L1–160wall-x VLA 模型的 MoE 视觉语言模型定义：导入依赖(ODE求解/SafeTensors/LoRA/VLM基础组件)、输出数据结构、注意力类型枚举、核心 Qwen2_5_VLDecoderLayer_with_MoE 类初始化(支持 attention_moe/mlp_moe/norm_moe 三维开关、TokenTypeRouter 非学习路由、SparseMoeBlock 多专家动作处理)。
标准库和深度学习框架导入：os、torch、yaml、numpy、glob、torch.nn，为整个模型提供系统接口、张量计算和参数管理基础。
7
from torchdiffeq import odeint
从 torchdiffeq 导入 odeint 求解器，用于 flow matching 推理时的 Euler 积分反向过程(generate_flow_action)去噪动作。实战坑：部署时积分步数通常固定为 5 步(参考 action_head.py)，不是训练的任意 t 采样。
8
from dataclasses import dataclass
9
from torch.nn import CrossEntropyLoss
10
from safetensors.torch import load_file
11
from peft import LoraConfig, get_peft_model
12
from typing import Optional, List, Tuple, Any, Dict, Union
13
import time14
from transformers import AutoConfig, AutoProcessor
15
from transformers.utils import logging, is_torchdynamo_compiling
导入 dataclass、损失函数、类型提示、时间工具、transformers 核心 API(AutoConfig/AutoProcessor/logging/torch.compile检查)；为数据定义、训练loss、类型安全和推理优化做准备。
16
from transformers.cache_utils import (
17
Cache,
18
DynamicCache,
19
SlidingWindowCache,
20
StaticCache,
21
)
22
from transformers.modeling_attn_mask_utils import AttentionMaskConverter
导入 transformers 缓存管理(Cache/DynamicCache/SlidingWindowCache/StaticCache)和注意力掩码工具(AttentionMaskConverter)，支持 KV cache 加速和多种注意力掩码格式转换，为不同推理模式(batch/流式/滑动窗口)提供基础；第 23 行空行作为 import 块的分隔。
24
from wall_x.model.qwen2_5_based.modeling_qwen2_5_vl import (
25
Qwen2_5_VLMLP,
26
Qwen2_5_VLRotaryEmbedding,
27
Qwen2_5_VLPreTrainedModel,
28
Qwen2RMSNorm,
29
Qwen2_5_VLForConditionalGeneration,
30
)
31
from transformers.modeling_outputs import (
32
ModelOutput,
33
BaseModelOutputWithPast,
34
)
导入 Qwen2.5-VL 基础组件(VLMLP/RotaryEmbedding/PreTrainedModel/RMSNorm/ForConditionalGeneration)和输出数据结构(ModelOutput/BaseModelOutputWithPast)，这些是本模型的父类和基础模块，所有 VLA 扩展都建立在这些之上，shape 传递贯穿全网络。
36
from wall_x.fusions import ops
37
from wall_x.model.action_head import ActionProcessor
38
from wall_x.model.qwen2_5_based.configuration_qwen2_5_vl import Qwen2_5_VLConfig
39
from wall_x.model.model_utils import load_wallx_processors, update_model_config
第 35 行空行分隔导入段；第 36-39 行导入融合操作库 ops(custom CUDA 置换/反置换)、ActionProcessor、Qwen2_5_VLConfig、工具函数(load_wallx_processors/update_model_config)，为 MoE token 置换和动作归一化提供基础。
40
from wall_x.model.qwen2_5_based.modeling_qwen2_5_vl import (
41
Qwen2_5_VisionTransformerPretrainedModel,
42
Qwen2_5_VLAttention,
43
Qwen2_5_VLFlashAttention2,
44
Qwen2_5_VLSdpaAttention,
45
)
导入注意力实现类：Qwen2_5_VisionTransformerPretrainedModel、Qwen2_5_VLAttention(eager)、Qwen2_5_VLFlashAttention2(flash)、Qwen2_5_VLSdpaAttention(sdpa)，这些会根据 config._attn_implementation 在第 103/107 行动态选择。
46
from wall_x.model.vla_mixin import ActionGenerationMixin, ActionModelMixMin
47
from wall_x.model.vla_mixin import TokenTypeRouter, SparseMoeBlock
48
from wall_x.model.vla_mixin import (
49
ATTENTION_TYPES_WITH_2D_MASK,
50
)
51
from wall_x.model.joint_attention import JOINT_QWEN_ATTENTION_CLASSES
导入 VLA 关键组件：ActionGenerationMixin(生成动作的混入)、ActionModelMixMin(动作模型最小混入接口)、TokenTypeRouter(确定性按 token_type%num_experts 路由)、SparseMoeBlock(稀疏MoE块)、ATTENTION_TYPES_WITH_2D_MASK(注意力掩码类型枚举，见 vla_mixin.py)；第 52 行空行分隔。
53
from wall_x.data.utils import update_action_statistics
54
from wall_x.utils.constant import action_statistic_dof
55
from pprint import pprint
57
logger = logging.get_logger(__name__)
导入动作统计相关：update_action_statistics 函数、action_statistic_dof 常量(各数据集的动作 min/max)；导入 pprint；第 57 行创建 logger 实例用于调试和警告。实战坑：LIBERO 数据的动作统计在常量中预定义，训练前需确保 Normalizer 读到正确的数据集名称('libero_10'等)。
60
@dataclass61
class Qwen2_5_VLACausalLMOutputWithPast(ModelOutput):
62
loss: Optional[torch.FloatTensor] = None63
flow_loss: Optional[torch.FloatTensor] = None64
cross_entropy_loss: Optional[torch.FloatTensor] = None65
logits: Optional[torch.FloatTensor] = None66
past_key_values: Optional[List[torch.FloatTensor]] = None67
hidden_states: Optional[Tuple[torch.FloatTensor]] = None68
attentions: Optional[Tuple[torch.FloatTensor]] = None69
rope_deltas: Optional[torch.LongTensor] = None71
channel_loss_dict: Optional[dict[torch.FloatTensor]] = None
72
channel_loss_count_dict: Optional[dict[torch.FloatTensor]] = None
第 58-59 行空行；第 60 行 @dataclass 装饰器；第 61-73 定义 Qwen2_5_VLACausalLMOutputWithPast 输出结构(继承 ModelOutput)，包含：loss(总loss)、flow_loss(扩散loss，shape (B,A,D))、cross_entropy_loss(token预测loss)、logits(token logits)、past_key_values(KV缓存)、hidden_states(中间隐状态)、attentions(注意力权重)、rope_deltas(位置编码增量)、channel_loss_dict(按DoF拆分loss用于诊断不同动作维度的学习进度)。
75
QWEN2_5_VL_ATTENTION_CLASSES = {76
"eager": Qwen2_5_VLAttention,
77
"flash_attention_2": Qwen2_5_VLFlashAttention2,
78
"sdpa": Qwen2_5_VLSdpaAttention,
79
}
第 74 行空行分隔数据类和注意力字典；第 75-79 定义 QWEN2_5_VL_ATTENTION_CLASSES 映射字典，key 为注意力实现名('eager'/'flash_attention_2'/'sdpa')，value 为对应的类；在第 107 行根据 config._attn_implementation 选择具体实现。
82
class Qwen2_5_VLDecoderLayer_with_MoE(nn.Module, ActionModelMixMin):
83
def __init__(
84
self,85
config: Qwen2_5_VLConfig,
86
layer_idx: int,87
num_experts: int,88
use_selective_recompute: bool = True,
89
):
90
super().__init__()
91
self.hidden_size = config.hidden_size92
self.use_selective_recompute = use_selective_recompute第 80-81 行空行和类定义行；第 82-89 定义 Qwen2_5_VLDecoderLayer_with_MoE 类(继承 nn.Module 和 ActionModelMixMin)的 __init__ 签名：接收 config(Qwen2_5_VLConfig)、layer_idx(层编号)、num_experts(2)、use_selective_recompute(梯度checkpointing，默认 True)；第 90-92 初始化 hidden_size 和 recompute 标志。
94
if (95
config.use_sliding_window
96
and config._attn_implementation != "flash_attention_2"
97
):
98
logger.warning_once(
99
f"Sliding Window Attention is enabled but not implemented for `{config._attn_implementation}`; "
100
"unexpected results may be encountered."
101
)
第 93 行空行；第 94-101 检查 sliding window attention 的配置兼容性：若 config.use_sliding_window=true 但 config._attn_implementation!='flash_attention_2'(如 sdpa/eager)，则打印警告(logger.warning_once 全局仅一次)。实战意义：滑动窗口可减少长序列计算，但仅 flash_attention_2 支持。
102
if config.attention_moe:103
self.self_attn = JOINT_QWEN_ATTENTION_CLASSES[config._attn_implementation](104
config, layer_idx
105
)
106
else:107
self.self_attn = QWEN2_5_VL_ATTENTION_CLASSES[config._attn_implementation](108
config, layer_idx
109
)
条件选择自注意力实现：若 config.attention_moe=true，使用 JOINT_QWEN_ATTENTION_CLASSES(图文与动作token共享注意力，cross-expert)；否则使用 QWEN2_5_VL_ATTENTION_CLASSES(标准Qwen，per-expert独占)。wall-x 默认 attention_moe=false，故不做跨专家注意力融合，动作token只在专家 1 的独占注意力中计算。
111
if config.use_adarms:112
adarms_cond_dims = [None, config.adarms_cond_dim]113
else:114
adarms_cond_dims = [None, None]
第 110 行空行；第 111-114 根据 config.use_adarms(是否用 AdaRM-S 条件归一化)设置 adarms_cond_dims：若开启 use_adarms，则为 [None, config.adarms_cond_dim](图文专家不用条件、动作专家用动作条件维度，如来自状态的条件)；否则为 [None, None]。
116
if config.norm_moe:117
self.input_layernorms = nn.ModuleList(118
[
119
Qwen2RMSNorm(
120
config.dim_inputs[i],
121
eps=config.rms_norm_eps,
122
cond_dim=adarms_cond_dims[i],
123
)
124
for i in range(num_experts)
125
]
126
)
127
self.post_attention_layernorms = nn.ModuleList(128
[
129
Qwen2RMSNorm(
130
config.dim_inputs[i],
131
eps=config.rms_norm_eps,
132
cond_dim=adarms_cond_dims[i],
133
)
134
for i in range(num_experts)
135
]
136
)
137
self.input_layernorm, self.post_attention_layernorm = None, None
138
else:139
self.input_layernorm = Qwen2RMSNorm(140
config.hidden_size, eps=config.rms_norm_eps
141
)
142
self.post_attention_layernorm = Qwen2RMSNorm(143
config.hidden_size, eps=config.rms_norm_eps
144
)
145
self.input_layernorms, self.post_attention_layernorms = None, None
第 115 行空行；第 116 条件分支 if config.norm_moe：若为 true，第 117-126 为每个专家创建独立的 input_layernorm(ModuleList，长度 num_experts)，维度各为 config.dim_inputs[i](如 [2048, 2048])，支持 adarms 条件注入；第 127-136 同理创建 post_attention_layernorm；第 137 标记共享 LayerNorm 为 None。else 分支(第 138-145)：创建单个共享的 input_layernorm/post_attention_layernorm(都是 config.hidden_size=2048)，标记分组 norm 为 None。实战坑：norm_moe=true 允许专家有不同的规范化尺度，但 wall-x 通常 norm_moe=false 用共享归一化，避免专家间的文本归一化差异影响 token 混合后的稳定性。
147
if config.mlp_moe:148
self.router = TokenTypeRouter(num_experts=num_experts)149
self.moe = SparseMoeBlock(150
config,
151
num_experts=num_experts,
152
use_selective_recompute=use_selective_recompute,
153
)
154
self.mlp = None
155
else:156
self.mlp = Qwen2_5_VLMLP(config)157
self.moe, self.router = None, None
第 146 行空行；第 147 条件分支 if config.mlp_moe：若为 true，第 148 创建 TokenTypeRouter(非学习路由，token_type%num_experts)；第 149-153 创建 SparseMoeBlock(包含 num_experts 个 BlockSparseMLP 专家，各有独立的 gate_proj/up_proj/down_proj，中间维 config.experts[i]['intermediate_size']，如 [11008, 2048])，启用选择性 recompute；第 154 标记普通 mlp 为 None。else 分支(第 155-157)：创建单个 Qwen2_5_VLMLP(标准 FFN)，标记 moe 和 router 为 None。背景：wall-x 标配 mlp_moe=true，num_experts=2；专家0(图文，中间维11008，继承Qwen权重)和专家1(动作，中间维2048，0.45B新参数)；TokenTypeRouter 按 token_type 取模路由，不涉及学习的负载均衡(见 vla_mixin.py:48 实现)。
第 158 行空行；第 159-160 保存 config 到实例成员 self.config，供 forward() 和其他方法后续查询(如第 220 行检查 config.attention_moe)，以及任何派生类通过 self.config 访问全局配置。
L161–284Qwen2_5_VLDecoderLayer_with_MoE.forward 方法：实现图文动作融合的 MoE 变换器层，包含 LayerNorm MoE、共享注意力、MLP MoE 和门控残差等三大子模块；支持 norm_moe/attention_moe/mlp_moe 配置开关，以及 LIBERO 任务专用的 token_types 路由和 dof_mask 处理。
161
def forward(
162
self,163
hidden_states: torch.Tensor,
164
attention_mask: Optional[torch.Tensor] = None,165
position_ids: Optional[torch.LongTensor] = None,166
past_key_value: Optional[Tuple[torch.Tensor]] = None,167
output_attentions: Optional[bool] = False,
168
use_cache: Optional[bool] = False,
169
cache_position: Optional[torch.LongTensor] = None,170
position_embeddings: Optional[
171
Tuple[torch.Tensor, torch.Tensor]
172
] = None, # necessary, but kept here for BC
173
# for vla174
token_types: Optional[torch.LongTensor] = None,175
start_indices: Optional[torch.Tensor] = None,176
end_indices: Optional[torch.Tensor] = None,177
probs: Optional[torch.Tensor] = None,178
row_id_map: Optional[torch.Tensor] = None,179
orig_shape: Optional[Tuple[int, int, int]] = None,
180
adarms_conds: Optional[List[torch.Tensor]] = [None, None],
181
**kwargs,
182
) -> Tuple[
183
torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]
184
]:
Qwen2_5_VLDecoderLayer_with_MoE.forward() 方法签名。输入是 (B,S,H=2048) 的 hidden_states，token_types 用来区分图文(0)和动作/状态(1)两类 token；start_indices/end_indices 是 2 元组分别记录两专家在 flattened 序列中的起止位置；output_attentions/use_cache 控制是否返回注意力权重和 KV cache。VLA 扩展参数(token_types/start_indices 等)在 attention_moe=false 时也会透传给 attention 层做联合处理。
185
"""186
Args:187
hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`188
attention_mask (`torch.FloatTensor`, *optional*): attention mask of size189
`(batch, sequence_length)` where padding elements are indicated by 0.190
output_attentions (`bool`, *optional*):191
Whether or not to return the attentions tensors of all attention layers. See `attentions` under192
returned tensors for more detail.193
use_cache (`bool`, *optional*):194
If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding195
(see `past_key_values`).196
past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states197
cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):198
Indices depicting the position of the input sequence tokens in the sequence.199
position_embeddings (`Tuple[torch.FloatTensor, torch.FloatTensor]`, *optional*):200
Tuple containing the cosine and sine positional embeddings of shape `(batch_size, seq_len, head_dim)`,201
with `head_dim` being the embedding dimension of each attention head.202
kwargs (`dict`, *optional*):203
Arbitrary kwargs to be ignored, used for FSDP and other methods that injects code204
into the model205
"""Docstring 概括前向过程输入输出。hidden_states 是每层输入(B,S,H)；attention_mask 标记 padding 位置；position_embeddings 包含 RoPE 的 cos/sin 值；kwargs 用于 FSDP 兼容性。返回 tuple(output_hidden_states, optional(attentions), optional(present_key_value))。
206
residual = hidden_states
保存原始 hidden_states 为残差，后续用于 residual connection；形状 (B,S,H)。
空行。
208
hidden_states, gate, _ = self._apply_norm_moe(209
hidden_states,
210
token_types,
211
adarms_conds,
212
self.input_layernorms,213
self.input_layernorm,214
start_indices,
215
end_indices,
216
self.use_selective_recompute,217
)
_apply_norm_moe(hidden_states, token_types, adarms_conds, self.input_layernorms, self.input_layernorm, start_indices, end_indices, use_selective_recompute) 调用：对输入应用 LayerNorm，可选 norm_moe 和 AdaRMS，返回 (norm_output, gate_for_expert1, gate_mask)。若 norm_moe=true，按 token_types 或 start/end 分配给专家 0/1 各自的 RMSNorm；若 adarms=true，动作专家的 Norm 会接收 condition 张量调节。gate 用于后续 gated_residual。形状 hidden_states (B,S,H)→(B,S,H) 保持不变。
空行。
219
# Self Attention注释行，标记自注意力块开始。
220
if self.config.attention_moe:
条件分支：if self.config.attention_moe，检查是否启用注意力 MoE；当前项目 false，图文和动作 token 共享注意力。
221
hidden_states, self_attn_weights, present_key_value = self.self_attn(222
hidden_states=hidden_states,
223
attention_mask=attention_mask,
224
position_ids=position_ids,
225
past_key_value=past_key_value,
226
output_attentions=output_attentions,
227
use_cache=use_cache,
228
cache_position=cache_position,
229
token_types=token_types,
230
start_indices=start_indices,
231
end_indices=end_indices,
232
probs=probs,
233
row_id_map=row_id_map,
234
orig_shape=orig_shape,
235
position_embeddings=position_embeddings,
236
)
attention_moe=true 分支(未在 LIBERO 项目使用)：调用 self.self_attn 并传递 token_types/start_indices/end_indices/probs/row_id_map/orig_shape，支持专家级别的注意力路由。注意力输出 (B,S,H)、权重、KV cache。
237
else:else 分支：attention_moe=false，图文和动作 token 在同一注意力头中交互。
238
hidden_states, self_attn_weights, present_key_value = self.self_attn(239
hidden_states=hidden_states,
240
attention_mask=attention_mask,
241
position_ids=position_ids,
242
past_key_value=past_key_value,
243
output_attentions=output_attentions,
244
use_cache=use_cache,
245
cache_position=cache_position,
246
position_embeddings=position_embeddings,
247
)
attention_moe=false 时的自注意力调用：hidden_states (B,S,H)→(B,S,H)；不传递 VLA 特定参数，仅使用标准 transformer 接口。返回 hidden_states(B,S,H)、self_attn_weights、present_key_value。此为 LIBERO 实际路径：所有 token 类型混在一起做多头注意，实现图文-动作的深层融合。
空行。
249
hidden_states = self._gated_residual(250
residual, hidden_states, gate, start_indices, end_indices
251
)
_gated_residual(residual, hidden_states, gate, start_indices, end_indices)：残差连接，若 gate=none 则返回 residual+hidden_states，否则仅对动作专家区间内的 token 应用 gate 调制；形状 (B,S,H)。实战中对动作路径的 attention 输出进行选择性缩放，避免动作 token 对图文的过度干扰。
空行。
253
# Fully Connected注释行，标记全连接(MLP)块开始。
254
residual = hidden_states
再次保存残差(此时已过注意力层)，用于 MLP 后的残差连接；形状仍为 (B,S,H)。
空行。
256
hidden_states, gate, gate_mask = self._apply_norm_moe(257
hidden_states,
258
token_types,
259
adarms_conds,
260
self.post_attention_layernorms,261
self.post_attention_layernorm,262
start_indices,
263
end_indices,
264
self.use_selective_recompute,265
)
_apply_norm_moe 第二次调用：在 MLP 前对 attention 输出再做一次 LayerNorm，返回 (norm_output, gate_for_expert1, gate_mask)。与第一次不同的是使用 post_attention_layernorms 而非 input_layernorms，各专家独立 Norm 参数。
空行。
267
hidden_states = self._apply_mlp_moe(268
hidden_states, token_types, start_indices, end_indices
269
)
_apply_mlp_moe(hidden_states, token_types, start_indices, end_indices)：若 mlp_moe=true(项目采用)，调用 SparseMoeBlock 将 token 按 token_types 路由给图文专家或动作专家，每专家处理对应维度；否则调用共享 MLP。输出形状 (B,S,H)。LIBERO 实战中专家 0 占 95% 参数，专家 1 只有 0.45B，通过 sparse 路由减少动作专家的计算。
空行。
271
hidden_states = self._gated_residual(272
residual, hidden_states, gate, start_indices, end_indices
273
)
_gated_residual 第二次调用：MLP 输出与之前保存的 residual 做门控融合，形式同第一次，动作区间内应用 gate。结果 (B,S,H)。
空行。
275
outputs = (hidden_states,)
初始化返回值元组，包含最终的 hidden_states(B,S,H)。
空行。
条件判断：若 output_attentions=true，则在 outputs 元组后追加自注意力权重张量。
条件判断：若 use_cache=true，则在 outputs 元组后追加当前层的 KV cache，供下一次 forward 重用以加速推理。
281
return outputsreturn outputs：返回包含 hidden_states 及可选 attentions/cache 的元组，供上层模型收集。
空行。
空行。
284
class Qwen2_5_VLMoEModel(Qwen2_5_VLPreTrainedModel, ActionModelMixMin):
Qwen2_5_VLMoEModel 类定义开始，继承 Qwen2_5_VLPreTrainedModel(huggingface 标准接口)和 ActionModelMixMin(VLA 动作前处理)；为整个模型的对外包装，负责 embedding、layer stack 和 action head 的端到端前向。
L285–355Qwen2_5_VLMoEModel 核心类定义：from_pretrained 动态配置专家数、__init__ 初始化 token embedding/MoE 解码层/规范化层、get/set_input_embeddings
285
@classmethod286
def from_pretrained(
287
cls, pretrained_model_name_or_path, num_experts=None, *args, **kwargs
288
):
@classmethod 装饰器 + from_pretrained 类方法签名：支持在加载预训练模型时动态传入 num_experts 参数，覆盖配置中的专家数。被外部加载器（如 .from_pretrained('path', num_experts=2)）调用以支持弹性专家配置。
289
# If `num_experts` is provided, ensure it is added to the config.290
config = kwargs.get("config", None)
291
if config is None:
292
config = AutoConfig.from_pretrained(pretrained_model_name_or_path)
从 kwargs 获取现有 config 或从本地路径 / HuggingFace Hub 自动加载配置。如果 config 为 None，调用 AutoConfig.from_pretrained 从模型名或本地路径加载 Qwen2_5_VLConfig。
若调用时显式提供 num_experts，则覆盖加载的 config.num_experts，允许同一模型快速切换专家数量。实战：LIBERO 微调时可通过该机制从预训练的 2 专家扩展到 4 专家。
297
kwargs["config"] = config
298
return super().from_pretrained(pretrained_model_name_or_path, *args, **kwargs)
将修改后的 config 写回 kwargs，调用父类（Qwen2_5_VLPreTrainedModel）的 from_pretrained 加载权重，完成模型权重与配置同步。
__init__ 方法签名：接收 Qwen2_5_VLConfig 配置对象和 use_selective_recompute 布尔标志（梯度检查点）。调用父类初始化。
301
super().__init__(config)
302
self.config = config303
self.use_selective_recompute = use_selective_recompute304
self.padding_idx = config.pad_token_id305
self.vocab_size = config.vocab_size缓存 config、use_selective_recompute 标志、padding_idx（来自 pad_token_id，用于遮蔽 pad token 梯度）和 vocab_size 到实例属性，供后续各层和评估使用。
307
self.embed_tokens = nn.Embedding(308
config.vocab_size, config.hidden_size, self.padding_idx309
)
创建 token embedding 层：nn.Embedding(vocab_size, hidden_size, padding_idx)。输入 input_ids (B,S) 经此层得 (B,S,H) 的浮点向量，其中 H=2048。padding_idx 避免对 pad token 的嵌入反向传播梯度。
310
self.layers = nn.ModuleList(311
[
312
Qwen2_5_VLDecoderLayer_with_MoE(
313
config,
314
layer_idx,
315
config.num_experts,
316
use_selective_recompute=use_selective_recompute,
317
)
318
for layer_idx in range(config.num_hidden_layers)
319
]
320
)
为 config.num_hidden_layers 层各创建一个 Qwen2_5_VLDecoderLayer_with_MoE 对象并存入 ModuleList。每层包含：TokenTypeRouter（按 token_type 路由到 2 个专家）+ MoE 注意力/FFN（Expert0=图文特征处理、Expert1=动作/状态处理）+ 选择性重计算。完整堆叠形成 (B,S,H)→(B,S,H)。
321
self._attn_implementation = config._attn_implementation缓存 config._attn_implementation 字符串（"flash_attention_2"/"sdpa"/"eager"），指导各层选择具体注意力实现以优化效率。
323
if config.use_adarms:324
adarms_cond_dims = [None, config.adarms_cond_dim]325
else:326
adarms_cond_dims = [None, None]
根据 use_adarms 布尔标志决定各专家的规范化层是否接收条件向量。若 True：adarms_cond_dims=[None, config.adarms_cond_dim]（Expert0 无条件、Expert1 接受动作条件）；否则 [None, None]，两专家规范化均无条件。
328
if config.norm_moe:329
self.norms = nn.ModuleList(330
[
331
Qwen2RMSNorm(
332
config.dim_inputs[i],
333
eps=config.rms_norm_eps,
334
cond_dim=adarms_cond_dims[i],
335
)
336
for i in range(config.num_experts)
337
]
338
)
339
self.norm = None
340
else:341
self.norm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)342
self.norms = None
根据 norm_moe 配置选择规范化策略：若 True，为每个专家创建独立 Qwen2RMSNorm（输入维度取自 config.dim_inputs[i]，如 (1536, 1536)），支持 AdaRMS 条件；self.norm=None。若 False，创建全局单一 RMSNorm（维度=hidden_size），两专家共享；self.norms=None。LIBERO 实战：通常 norm_moe=False，两专家共享最终规范化层。
初始化旋转位置编码对象 Qwen2_5_VLRotaryEmbedding(config=config)，支持 RoPE（Rotary Position Embedding）在 attention 层对 Q/K 进行位置信息融合。
346
self.gradient_checkpointing = False
347
# Initialize weights and apply final processing348
self.post_init()初始化梯度检查点标志为 False（可由外部设置 model.gradient_checkpointing=True 动态启用以降低显存占用）；调用 post_init() 执行权重初始化和 transformers 框架的最终设置（如权重绑定等）。
get_input_embeddings 方法：简单返回 self.embed_tokens，允许外部访问 token embedding 层进行权重查看、权重共享或计算 token logits。
set_input_embeddings 方法：允许外部更新 self.embed_tokens，支持权重共享（如与输出层绑定）或动态扩展 vocab_size。末尾空行标记方法结束。
L356–583Qwen2_5_VLMoEModel 核心前向传递方法：包括参数验证、KV缓存初始化、位置编码计算、流过每层 MoE decoder、应用最终归一化和可选的 mot_opt 排列优化，返回编码后的隐状态和可选的缓存/注意力。
356
def forward(
357
self,358
input_ids: torch.LongTensor = None,359
attention_mask: Optional[torch.Tensor] = None,360
position_ids: Optional[torch.LongTensor] = None,361
past_key_values: Optional[List[torch.FloatTensor]] = None,362
inputs_embeds: Optional[torch.FloatTensor] = None,363
moe_token_types: Optional[torch.LongTensor] = None,364
start_indices: Optional[torch.Tensor] = None,365
end_indices: Optional[torch.Tensor] = None,366
positional_masks: Optional[dict] = None,
367
use_cache: Optional[bool] = None,
368
output_attentions: Optional[bool] = None,
369
output_hidden_states: Optional[bool] = None,
370
return_dict: Optional[bool] = None,
371
cache_position: Optional[torch.LongTensor] = None,372
adarms_conds: Optional[List[torch.Tensor]] = [None, None],
373
**kwargs,
374
) -> Union[Tuple, BaseModelOutputWithPast]:
375
output_attentions = (
376
output_attentions
377
if output_attentions is not None
378
else self.config.output_attentions
379
)
380
output_hidden_states = (
381
output_hidden_states
382
if output_hidden_states is not None
383
else self.config.output_hidden_states
384
)
385
use_cache = use_cache if use_cache is not None else self.config.use_cache
387
return_dict = (
388
return_dict if return_dict is not None else self.config.use_return_dict
389
)
forward() 函数签名、参数列表（包括 input_ids、attention_mask、position_ids、past_key_values、inputs_embeds、moe_token_types 用于 MoE 路由、start_indices/end_indices 为 mot_opt 段指标）及返回类型 Union[Tuple, BaseModelOutputWithPast]。包含空行 390。
391
if (input_ids is None) ^ (inputs_embeds is not None):
392
raise ValueError(393
"You must specify exactly one of input_ids or inputs_embeds"
394
)
395
if moe_token_types is None:
396
raise ValueError("moe_token_types must be provided for MoE routing.")
397
if start_indices is None or end_indices is None:
398
raise ValueError(399
"start_indices and end_indices must be provided for MoE routing"
400
)
设置 output_attentions/output_hidden_states/use_cache/return_dict 的默认值：如果调用端没传，就使用 config 的默认配置。包含空行 401。
402
if self.gradient_checkpointing and self.training:
403
if use_cache:404
logger.warning_once(
405
"`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
406
)
407
use_cache = False409
# torch.jit.trace() doesn't support cache objects in the output输入验证三项：(1) input_ids 和 inputs_embeds 必须恰好提供一个；(2) moe_token_types 必须提供（用于 MoE 路由判断每个 token 的专家类型）；(3) start_indices/end_indices 必须提供（用于 mot_opt 排列或非 mot_opt 的专家分配）。包含空行 408-409。
410
if use_cache and past_key_values is None and not torch.jit.is_tracing():
411
past_key_values = DynamicCache()
如果启用梯度检查点（gradient_checkpointing）且处于训练模式，禁用 use_cache（两者不兼容），梯度检查点通过重计算中间激活来节省显存。包含空行 412。
如果需要 KV 缓存但未提供，初始化 DynamicCache()。DynamicCache 会在每层的 forward 中累积 key/value，加速自回归生成。包含空行 415。
416
if cache_position is None:
417
past_seen_tokens = (
418
past_key_values.get_seq_length() if past_key_values is not None else 0
419
)
420
cache_position = torch.arange(
421
past_seen_tokens,
422
past_seen_tokens + inputs_embeds.shape[1],423
device=inputs_embeds.device,
424
)
426
# the hard coded `3` is for temporal, height and width.如果没传 inputs_embeds，用 vocab embedding 层把 input_ids(B,S) 转为嵌入 (B,S,2048)。初始化 cache_position（当前生成位置）：如果是首次或无缓存，范围从 0 到 inputs_embeds.shape[1]；否则从 past_key_values.get_seq_length() 开始追加。行 426 的注释说硬编码 3 是时间/高度/宽度维度（后续用于旋转位置编码）。
427
if position_ids is None:
428
position_ids = cache_position.view(1, 1, -1).expand(
429
3, inputs_embeds.shape[0], -1
430
)
431
elif position_ids.dim() == 2:
432
position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1)
处理位置 id。如果位置 id 为 None，从 cache_position 广播到 (3, B, S)；如果已提供的是 2D (B,S)，扩展为 (3,B,S)。用于三维旋转位置编码。包含空行 433。
434
if not self.config.attention_moe:
435
causal_mask = self._update_causal_mask(436
attention_mask,
437
inputs_embeds,
438
cache_position,
439
past_key_values,
440
output_attentions,
441
moe_token_types,
442
)
443
else:444
causal_mask = attention_mask
根据 config.attention_moe 决定是否在不同专家间共享注意力。若 False（默认），调用 _update_causal_mask() 生成专家间的对齐因果掩码；若 True（图文和动作共一套 attention），直接用 attention_mask。实战：LIBERO 用 attention_moe=False。包含空行 445。
初始化 hidden_states = inputs_embeds，维度 (B,S,2048)。之后通过 36 层 decoder 逐层更新。包含空行 447。
448
if (449
self.config._attn_implementation != "flash_attention_2"
450
and self.config.attention_moe is True
451
):
452
position_ids = self._update_position_ids(453
position_ids, moe_token_types, positional_masks
454
)
如果非 flash_attention_2 且启用 attention_moe，调用 _update_position_ids() 根据 token_types 和 positional_masks 调整位置 id。flash_attention_2 内部已处理位置信息。包含空行 455。
456
# create position embeddings to be shared across the decoder layers457
position_embeddings = self.rotary_emb(hidden_states, position_ids)通过 Qwen2_5_VLRotaryEmbedding 计算旋转位置编码。输入 hidden_states (B,S,2048) 和 position_ids (3,B,S)，输出同维 position_embeddings。旋转编码会在每层 attention 中与 Q/K 融合。包含空行 458。
459
# If `mot_opt` is enabled, the tokens from different experts will be permuted first, resulting in a dimension of [Tokens, HiddenSize].460
orig_shape = hidden_states.shape
461
if self.config.mot_opt:
462
hidden_states = hidden_states.view(-1, hidden_states.size(-1))
463
hidden_states, row_id_map = ops.permute(
464
hidden_states, moe_token_types.view(-1)465
)
466
else:467
row_id_map = None468
probs = torch.ones_like(moe_token_types.view(-1), dtype=torch.float32).view(469
-1, 1
470
)
mot_opt（"mixture of tokens" 优化）：若启用，按 token_types 对 tokens 进行排列（gather），使同专家的 tokens 连续，加快 specialized MLPMoE 计算。(1) 记录原始形状；(2) 展平为 (total_tokens, 2048)；(3) 调用 ops.permute() 按 moe_token_types 排列，返回排列后的 hidden_states 和 row_id_map（逆向映射）；(4) 初始化 probs=(1,1) 全 1，用于后续 unpermute。包含空行 471。
472
# decoder layers473
all_hidden_states = () if output_hidden_states else None
474
all_self_attns = () if output_attentions else None
475
next_decoder_cache = None初始化输出容器：all_hidden_states（若 output_hidden_states=True）、all_self_attns（若 output_attentions=True）、next_decoder_cache（若 use_cache=True）。包含空行 476。
477
# generate 2d attention mask if needed478
if (479
self.config._attn_implementation in ATTENTION_TYPES_WITH_2D_MASK
480
and self.config.attention_moe is True
481
):
482
if causal_mask is not None and inputs_embeds.shape[1] > 1:
483
causal_mask = self._update_joint_attention_mask_2d(484
attention_mask=causal_mask,
485
moe_token_types=moe_token_types,
486
positional_masks=positional_masks,
487
)
如果需要 2D attention mask（用于 attention_moe=True 且非 flash_attention_2）且 sequence_length>1，调用 _update_joint_attention_mask_2d() 生成联合注意力掩码。此时图文和动作 tokens 在同一 attention 中，需要专用的 2D 掩码矩阵。包含空行 488。
489
for decoder_layer in self.layers:
490
if output_hidden_states:491
assert (492
self.config.mot_opt is False
493
), "When using mot_opt, output_hidden_states is not supported yet."
494
all_hidden_states += (hidden_states,)
遍历 36 层 decoder。若 output_hidden_states=True，先把当前 hidden_states 保存到 all_hidden_states 元组（但 mot_opt 下不支持，会断言失败）。维度仍为 (B,S,2048) 或 mot_opt 下的 (total_tokens,2048)。包含空行 495。
496
if self.gradient_checkpointing and self.training:
497
layer_outputs = self._gradient_checkpointing_func(498
decoder_layer.__call__,499
hidden_states,
500
causal_mask,
501
position_ids,
502
past_key_values,
503
output_attentions,
504
use_cache,
505
cache_position,
506
position_embeddings,
507
# for vla508
moe_token_types,
509
start_indices,
510
end_indices,
511
probs,
512
row_id_map,
513
orig_shape,
514
adarms_conds,
515
)
516
else:若启用梯度检查点且训练中，调用 _gradient_checkpointing_func() 包装 decoder_layer 的 forward。避免保存 decoder 层的中间激活，而是在反向时重计算。传入的 moe_token_types, start_indices, end_indices, probs, row_id_map, orig_shape 用于 MoE 路由和排列。包含空行 516。
517
layer_outputs = decoder_layer(
518
hidden_states,
519
attention_mask=causal_mask,
520
position_ids=position_ids,
521
past_key_value=past_key_values,
522
output_attentions=output_attentions,
523
use_cache=use_cache,
524
cache_position=cache_position,
525
position_embeddings=position_embeddings,
526
# for vla527
token_types=moe_token_types,
528
start_indices=start_indices,
529
end_indices=end_indices,
530
probs=probs,
531
row_id_map=row_id_map,
532
orig_shape=orig_shape,
533
adarms_conds=adarms_conds,
534
)
535
hidden_states = layer_outputs[0]正常前向（非梯度检查点）或梯度检查点返回后：提取当前层的输出隐状态 layer_outputs[0]，维度同前 (B,S,2048) 或 mot_opt 下 (total_tokens,2048)。decoder_layer 是 Qwen2_5_VLDecoderLayer_with_MoE，关键参数：token_types, start_indices, end_indices 用于层内 MoE FFN 路由；probs, row_id_map 用于 mot_opt unpermute。包含空行和 cache 更新逻辑。
如果启用 KV 缓存，从 layer_outputs 中提取 cache（位置取决于是否有 output_attentions）。对于每层，都会追加新的 key/value 到 DynamicCache。包含空行 536 和 539。
540
if output_attentions:541
assert (542
self.config.mot_opt is False
543
), "When using mot_opt, output_hidden_states is not supported yet."
544
all_self_attns += (layer_outputs[1],)如果启用 output_attentions，从 layer_outputs[1] 提取注意力权重。mot_opt 下不支持此选项，会断言失败。通常用于可视化或调试。包含空行 545。
546
hidden_states, _, _ = self._apply_norm_moe(547
hidden_states,
548
moe_token_types,
549
adarms_conds,
550
self.norms,551
self.norm,552
start_indices,
553
end_indices,
554
self.use_selective_recompute,555
)
经过全部 36 层后，应用最终的层归一化。_apply_norm_moe() 根据 config.norm_moe 决定是专家间独立归一化还是共享单个 LN。传入 moe_token_types, adarms_conds（用于 AdaRMS 条件）、start_indices/end_indices（mot_opt 下的段指标）。返回 (hidden_states, gate, gate_mask)，其中 gate 用于后续的 gated residual（只在 action 专家有非 None 的 gate）。包含空行 556。
557
# add hidden states from the last decoder layer558
if output_hidden_states:559
assert (560
self.config.mot_opt is False
561
), "When using mot_opt, output_hidden_states is not supported yet."
562
all_hidden_states += (hidden_states,)
若 output_hidden_states=True，把最终的 hidden_states 附加到 all_hidden_states。mot_opt 下断言失败。包含空行 563。
确定返回的 cache。如果 use_cache=True，返回 next_decoder_cache（每层累积的 KV）；否则返回 None。包含空行 565。
566
if self.config.mot_opt:
567
hidden_states = ops.unpermute(hidden_states, row_id_map, probs)
568
hidden_states = hidden_states.view(orig_shape)
若启用 mot_opt，反向排列 hidden_states 回原始顺序。(1) ops.unpermute(hidden_states, row_id_map, probs) 用 row_id_map 和全 1 的 probs 恢复顺序；(2) 重塑回 (B,S,2048)。实战：LIBERO 使用 mot_opt=True 加速 MoE 路由。包含空行 569。
570
if not return_dict:
571
return tuple(
572
v
573
for v in [hidden_states, next_cache, all_hidden_states, all_self_attns]
574
if v is not None
575
)
若 return_dict=False，返回一个元组，过滤掉 None 值。顺序通常为 (hidden_states, cache, all_hidden_states, all_attentions)。包含空行 576。
577
return BaseModelOutputWithPast(578
last_hidden_state=hidden_states,
579
past_key_values=next_cache,
580
hidden_states=all_hidden_states,
581
attentions=all_self_attns,
582
)
若 return_dict=True（默认），返回 BaseModelOutputWithPast 对象，包含 last_hidden_state (B,S,2048)、past_key_values（DynamicCache 或 None）、hidden_states（所有层中间激活的元组，若 output_hidden_states=True）、attentions（所有层注意力权重的元组，若 output_attentions=True）。forward 方法结束。
L584–692attention mask 工厂函数：处理 Flash Attention 2/SDPA 两条路径下的因果掩码生成与修改，支持 MoE token type 1(动作/状态 token)之间的双向注意力，处理各类缓存(DynamicCache/SlidingWindowCache/StaticCache)与左填充场景。
584
def _update_causal_mask(
585
self,586
attention_mask: torch.Tensor,
587
input_tensor: torch.Tensor,
588
cache_position: torch.Tensor,
589
past_key_values: Cache,
590
output_attentions: bool,591
moe_token_types: Optional[torch.LongTensor] = None,592
):
函数签名：_update_causal_mask 是 Qwen2_5_VLModel 的方法，被 forward() 调用生成训练/推理时的 attention mask。参数包括原始 attention_mask(batch 掩码)、input_tensor、cache_position、past_key_values(KV 缓存)、output_attentions、moe_token_types(MoE 路由标记)。返回修改后的 4D 因果掩码或 None。
593
if self.config._attn_implementation == "flash_attention_2":
594
if attention_mask is not None and past_key_values is not None:
595
is_padding_right = (
596
attention_mask[:, -1].sum().item() != input_tensor.size()[0]
597
)
598
if is_padding_right:599
raise ValueError(600
"You are attempting to perform batched generation with padding_side='right'"
601
" this may lead to unexpected behaviour for Flash Attention version of Qwen2_5_VL. Make sure to "
602
" call `tokenizer.padding_side = 'left'` before tokenizing the input. "
603
)
604
if attention_mask is not None and 0.0 in attention_mask:
605
return attention_mask606
return None
Flash Attention 2 分支(line 593)：若使用 FA2 且启用 KV 缓存，检查是否为右填充(line 595-597)——对比 attention_mask 最后一列的有效 token 数与 batch_size，若不匹配则报错(padding_side 必须='left')；若 attention_mask 中有 0.0 掩码值则直接返回(FA2 自动处理因果性)，否则返回 None(FA2 用 is_causal=True 替代掩码)。LIBERO 场景：生成推理时批量解码若启用了右填充会导致生成错误。607 为空行。
608
# For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in609
# order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail610
# to infer the attention mask.611
past_seen_tokens = (
612
past_key_values.get_seq_length() if past_key_values is not None else 0
613
)
614
using_static_cache = isinstance(past_key_values, StaticCache)615
using_sliding_window_cache = isinstance(past_key_values, SlidingWindowCache)注释说明 SDPA 路由策略(line 608-610)，随后计算 past_seen_tokens(line 611-613)：若 KV 缓存存在则调用 get_seq_length() 得到已处理的序列长度，否则为 0；检查缓存类型是否为 StaticCache 或 SlidingWindowCache(line 614-615)，用于后续条件分支。616 为空行。
617
# When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward618
if (619
self.config._attn_implementation == "sdpa"
620
and not (using_static_cache or using_sliding_window_cache)
621
and not output_attentions
622
):
623
if attention_mask.ndim == 2:
624
if AttentionMaskConverter._ignore_causal_mask_sdpa(625
attention_mask,
626
inputs_embeds=input_tensor,
627
past_key_values_length=past_seen_tokens,
628
sliding_window=self.config.sliding_window,629
is_training=self.training,630
):
631
return None
632
elif attention_mask.ndim == 3:
633
return attention_maskSDPA 优化路径判断(line 618-622)：满足 SDPA 实现、非静态/滑动窗口缓存、非输出注意力权重时进入；若 attention_mask 为 2D 掩码，调用 Transformers 内置的 _ignore_causal_mask_sdpa 检查是否可忽略因果掩码(line 624-631)——当掩码全 1 且满足其他条件则返回 None；若为 3D 掩码(已经是 4D 的降维形式)则直接返回。634 为空行。
635
dtype, device = input_tensor.dtype, input_tensor.device
636
min_dtype = torch.finfo(dtype).min
637
sequence_length = input_tensor.shape[1]638
# SlidingWindowCache or StaticCache639
if using_sliding_window_cache or using_static_cache:
640
target_length = past_key_values.get_max_cache_shape()
641
# DynamicCache or no cache642
else:643
target_length = (
644
attention_mask.shape[-1]645
if isinstance(attention_mask, torch.Tensor)
646
else past_seen_tokens + sequence_length + 1
647
)
提取 dtype/device(line 635)，获取 min_dtype(line 636，用于掩码填充值)，计算 sequence_length=input_tensor 的第 1 维(line 637)；计算 target_length(line 639-647)：若使用静态/滑动窗口缓存则用缓存的最大形状，否则用 attention_mask 的最后一维或推断值(past_seen_tokens+sequence_length+1)。target_length 是生成 4D 掩码时的 key 序列长度。648 为空行。
649
# In case the provided `attention` mask is 2D, we generate a causal mask here (4D).650
causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position(651
attention_mask,
652
sequence_length=sequence_length,
653
target_length=target_length,
654
dtype=dtype,
655
device=device,
656
cache_position=cache_position,
657
batch_size=input_tensor.shape[0],658
config=self.config,659
past_key_values=past_key_values,
660
)
调用静态方法 _prepare_4d_causal_attention_mask_with_cache_position(line 650-660)生成 4D 因果掩码，参数包括原始 2D 掩码、sequence_length、target_length、dtype、device、cache_position、batch_size、config、past_key_values；返回形状 (B, num_heads, S_query, S_key) 的因果掩码(满足当前位置只能看到历史位置)。
661
# Modify the mask to support bidirectional attention.662
if moe_token_types is not None:
663
# Find the positions of all tokens of type 1.664
type1_tokens = (
665
(moe_token_types == 1).unsqueeze(1).unsqueeze(2)
666
) # [B, 1, 1, S]668
# Create a square mask for the type1 region.669
type1_mask = torch.zeros_like(causal_mask) # [B, num_heads, S, S]670
type1_region = type1_tokens & type1_tokens.transpose(-1, -2) # [B, 1, S, S]
671
type1_mask = type1_mask.masked_fill(type1_region, 1.0).to(torch.bool)672
# Set the original causal_mask to zero in the type1 region, and then add the type1_mask.673
causal_mask = torch.where(
674
type1_mask,
675
torch.zeros_like(causal_mask),
676
causal_mask,
677
)
MoE 修改掩码分支(line 662)：若 moe_token_types 不为 None(即有动作/状态 token)；line 664-666 生成 type1_tokens 布尔张量表示类型为 1 的位置，unsqueeze 到 (B,1,1,S)；line 669 初始化 type1_mask 与 causal_mask 同形(B,num_heads,S,S)；line 670 通过转置运算找出双向注意力区域 type1_region(B,1,S,S)，标记所有 type1_token 之间的注意力对；line 671 将 type1_region 处的 type1_mask 填为 1.0；line 673-677 通过 torch.where 把原因果掩码中的 type1 区域替换为全 0(变成双向)，其余区域保持原因果性。LIBERO 实战：动作 token 之间允许互相看到上文，解决因果约束导致的当前 action 看不到 proprioception 问题。
678
if (679
self.config._attn_implementation == "sdpa"
680
and attention_mask is not None
681
and attention_mask.device.type in ["cuda", "xpu"]
682
and not output_attentions
683
):
684
# Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when685
# using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path.686
# Details: https://github.com/pytorch/pytorch/issues/110213687
causal_mask = AttentionMaskConverter._unmask_unattended(
688
causal_mask, min_dtype
689
)
SDPA 内存高效路径补偿(line 678-683)：当使用 SDPA、有原始掩码、在 CUDA/XPU 设备、非输出注意力时，调用 _unmask_unattended(line 687-689)处理因果掩码中的完全掩码行(如左填充的前几行)——F.scaled_dot_product_attention 的高效路由要求这些行可以注意到所有位置，否则会报错(PyTorch issue #110213)。690 为空行。
返回修改后的 causal_mask(4D 布尔张量)，或之前提前返回的 None；mask 值为 True 表示可注意(未掩码)，False 表示掩码。
L693–828因果注意力掩码构建与 Qwen2.5 VL MoE 动作模型类定义、自定义机器人配置初始化（包含 normalizer 统计数据更新）
693
@staticmethod694
def _prepare_4d_causal_attention_mask_with_cache_position(
695
attention_mask: torch.Tensor,
696
sequence_length: int,697
target_length: int,698
dtype: torch.dtype,
699
device: torch.device,
700
cache_position: torch.Tensor,
701
batch_size: int,702
config: Qwen2_5_VLConfig,
703
past_key_values: Cache,
704
):
@staticmethod 装饰器 + 函数签名。_prepare_4d_causal_attention_mask_with_cache_position 根据缓存位置和目标长度构建 4D 因果掩码，是 _prepare_causal_attention_mask() 内部调用的辅助方法。参数包括 attention_mask(2D 或 4D Tensor)、sequence_length(int)、target_length(int)、dtype、device、cache_position、batch_size、config、past_key_values(Cache 对象)。
705
"""706
Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape707
`(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing.709
Args:710
attention_mask (`torch.Tensor`):711
A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape `(batch_size, 1, query_length, key_value_length)`.712
sequence_length (`int`):713
The sequence length being processed.714
target_length (`int`):715
The target length: when generating with static cache, the mask should be as long as the static cache, to account for the 0 padding, the part of the cache that is not filled yet.716
dtype (`torch.dtype`):717
The dtype to use for the 4D attention mask.718
device (`torch.device`):719
The device to plcae the 4D attention mask on.720
cache_position (`torch.Tensor`):721
Indices depicting the position of the input sequence tokens in the sequence.722
batch_size (`torch.Tensor`):723
Batch size.724
config (`Qwen2_5_VLConfig`):725
The model's configuration class726
past_key_values (`Cache`):727
The cache class that is being used currently to generate728
"""函数 docstring：详细说明输入输出形态、处理逻辑。输入可以是 (batch_size, key_value_length) 的 2D mask 或 (batch_size, 1, query_length, key_value_length) 的 4D mask；输出总是 4D 掩码。target_length 用于静态 cache 长度对齐，cache_position 记录输入序列在完整序列中的位置索引。
729
if attention_mask is not None and attention_mask.dim() == 4:
730
# In this case we assume that the mask comes already in inverted form and requires no inversion or slicing.731
causal_mask = attention_mask
条件判断：若输入 attention_mask 已是 4D 形式（mask.dim()==4），则直接使用，跳过后续构建逻辑（因为前向计算中已处理过逆序和切片）。
732
else:733
min_dtype = torch.finfo(dtype).min
734
causal_mask = torch.full(
735
(sequence_length, target_length),
736
fill_value=min_dtype,
737
dtype=dtype,
738
device=device,
739
)
740
diagonal_attend_mask = torch.arange(
741
target_length, device=device
742
) > cache_position.reshape(-1, 1)
初始化因果掩码。(1) min_dtype = torch.finfo(dtype).min 取 dtype 最小值（float32 约 -3.4e38）；(2) torch.full 创建 (sequence_length, target_length) 的张量，全填 min_dtype；(3) diagonal_attend_mask 通过 arange 和比较生成对角线掩码，target_length > cache_position 的位置为 True（可参与注意）。形态转换：(seq_len, tgt_len) → (seq_len, tgt_len)，仅逻辑值。
743
if config.sliding_window is not None:
744
# if we have sliding window, we should not attend to tokens beyond sliding window length, so we mask them out also745
# the check is needed to verify is current checkpoint was trained with sliding window or not746
if (747
not isinstance(past_key_values, SlidingWindowCache)
748
or sequence_length > target_length749
):
750
sliding_attend_mask = torch.arange(
751
target_length, device=device
752
) <= (cache_position.reshape(-1, 1) - config.sliding_window)
753
diagonal_attend_mask.bitwise_or_(sliding_attend_mask)
滑动窗口处理。若 config.sliding_window 非空且模型训练时用了滑动窗口，则额外屏蔽超出窗口范围的位置：sliding_attend_mask = target_arange <= (cache_position - sliding_window)，用 bitwise_or_ 并入 diagonal_attend_mask。防止模型关注距离过远的 token（因果性 + 窗口大小约束）。
754
causal_mask *= diagonal_attend_mask
755
causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1)
causal_mask *= diagonal_attend_mask：将布尔掩码转换为浮点掩码。被掩蔽位置保留 min_dtype，非掩蔽位置变为 min_dtype*1=min_dtype（实际逻辑需配合后续 masked_fill）。形态：(seq_len, tgt_len)。然后通过 [None, None, :, :] 添加 batch_size 和 num_heads 维，expand(batch_size, 1, -1, -1) 广播到 (B, 1, seq_len, tgt_len)，对应 multi-head attention 的掩码形态。
756
if attention_mask is not None:
757
causal_mask = (
758
causal_mask.clone()
759
) # copy to contiguous memory for in-place edit760
if attention_mask.shape[-1] > target_length:
761
attention_mask = attention_mask[:, :target_length]
762
mask_length = attention_mask.shape[-1]763
padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[
764
:, None, None, :
765
].to(causal_mask.device)
766
padding_mask = padding_mask == 0767
causal_mask[:, :, :, :mask_length] = causal_mask[
768
:, :, :, :mask_length
769
].masked_fill(padding_mask, min_dtype)
770
return causal_mask处理 padding mask（若有）。(1) clone() 确保 contiguous 内存供 masked_fill 修改；(2) 若 attention_mask 长度超过 target_length 则截断；(3) padding_mask = causal_mask[:,:,:,:mask_length] + attention_mask[:, None, None, :] 相加，两个掩蔽相叠；(4) 仅 == 0 处保留（无遮挡），其余用 min_dtype 填充。(5) return causal_mask 返回构建好的 4D 因果掩码 (B, 1, seq_len, tgt_len)。(6) 空行分隔。在 LIBERO 场景中，若使用 sequence packing 或变长序列，attention_mask 会标记真实 token 位置。
空行与类定义开始。类定义：Qwen2_5_VLMoEForAction 继承自 Qwen2_5_VLForConditionalGeneration(基础 VL 模型)、ActionGenerationMixin(动作生成)、ActionModelMixMin(动作模型最小接口)。是用于 LIBERO 等动作学习任务的完整 VLA 模型。
774
Qwen2_5_VLForConditionalGeneration, ActionGenerationMixin, ActionModelMixMin
775
):
776
"""777
Qwen2.5 Vision-Language Mixture of Experts model for action processing.779
This model extends the base Qwen2.5 VL model with action token processing capabilities780
and optional LoRA fine-tuning support.781
"""类 docstring（多行）：说明该模型扩展基础 Qwen2.5 VL 模型，添加动作 token 处理与可选 LoRA 微调能力。是整套 wall-x VLA 架构的核心类，三个基类协同提供完整的 VLA 功能。
783
_tied_weights_keys = ["lm_head.weight"]
784
config_class = Qwen2_5_VLConfig
785
_no_split_modules = ["Qwen2_5_VLDecoderLayer_with_MoE", "Qwen2_5_VLVisionBlock"]
类属性。_tied_weights_keys = ["lm_head.weight"]：权重共享配置（lm_head 与 embedding 层共享，节省参数）；config_class = Qwen2_5_VLConfig；_no_split_modules 声明不应被拆分到不同设备的模块（FSDP 相关，层内的 MoE 和 ViT Block 保持完整）。
@classmethod 装饰器与方法签名。_set_customized_config 是类方法，用于从 config dict 读取并处理 norm_stats.json，重建机器人 DoF 映射。
789
"""790
Processing norm_stats.json and reconstruct the DoF mapping791
"""792
dataload_config = config["data"]
793
if not dataload_config.get("use_lerobot", False):
794
raise NotImplementedError(795
"Not implemented for non-lerobot dataset currently"
796
)
方法 docstring：说明功能是处理 norm_stats.json 并重建 DoF（自由度）映射。该方法在初始化时调用，用于将归一化统计量与机器人配置关联。
读取 config["data"] 中的 use_lerobot 标志。若为 False，抛出 NotImplementedError 表示当前仅支持 LeRobot dataset。实战约束：LIBERO 数据集必须经过 lerobot_config 包装。
801
assert (802
enable_customized_robot_config
803
), "enable_customized_robot_config must be true when use lerobot dataset"
验证 enable_customized_robot_config 必须为 True。当使用 lerobot 数据集时，自定义机器人配置强制启用（否则 normalizer 无法初始化）。
805
customized_dof_config = config["customized_robot_config"][
806
"customized_dof_config"
807
]
808
customized_agent_pos_config = config["customized_robot_config"][
809
"customized_agent_pos_config"
810
]
从 config 中提取三个关键配置：(1) customized_dof_config：DoF 维数映射表（如 {"left_arm": 7, "right_arm": 7, ...}）；(2) customized_agent_pos_config：关键点位置配置；(3) norm_stats_path：归一化统计文件路径。这些是后续 update_action_statistics 的必需参数。
811
norm_stats_path = config["norm_stats_path"]
813
# Use the compute_action_statistics function from utils读取 robot_name 和注释行。name = config["customized_robot_config"]["name"] 获取机器人名（如 "libero_object_tabletop"），用于在 action_statistic_dof 全局字典中查表。注释说明接下来调用 compute_action_statistics 等函数。
赋值与空行。name = config["customized_robot_config"]["name"] 获取机器人名称作为 action_statistic_dof 的索引键。
817
update_action_statistics(
818
action_statistic_dof=action_statistic_dof, # Assuming this is a global variable819
norm_stats_path=norm_stats_path,
820
repo_id=config["data"]["lerobot_config"]["repo_id"],
821
robot_name=name,
822
customized_dof_config=customized_dof_config,
823
customized_agent_pos_config=customized_agent_pos_config,
824
)
调用 update_action_statistics 函数，将归一化统计数据从 norm_stats.json 加载到全局 action_statistic_dof dict。参数包括 norm_stats_path、lerobot dataset repo_id、robot_name、customized_dof_config、customized_agent_pos_config。加载后，Normalizer 实例可根据 robot_name 查表获得 min/delta。实战：LIBERO 数据集名为 "libero_object_tabletop" 等，对应的统计量应提前计算并存在 norm_stats.json。
打印完成与空行。print("Customized robot config added") 输出初始化完毕信号；pprint(action_statistic_dof) 格式化打印全局字典内容（包含该 robot_name 的所有 DoF 归一化统计）；最后空行结束方法。此时 action_statistic_dof 已持久化，供后续 Normalizer 实例使用。
L829–988from_pretrained 类方法加载预训练模型，处理 LIBERO 配置/processor/safetensors 权重；__init__ 初始化 VLMoE 模型、action_preprocessor、LoRA、损失函数等完整训练推理前向图。
829
@classmethod830
def from_pretrained(
831
cls,832
pretrained_model_path,
833
train_config=None,834
config_path=None,835
processor_path=None,836
action_tokenizer_path=None,837
is_train=False,838
**kwargs,
839
):
@classmethod 和 from_pretrained 方法签名。这是类方法，用于从预训练模型路径加载完整模型实例，支持 train_config/config_path/processor_path/action_tokenizer_path 等多种输入方式；是 from_pretrained 的自定义实现，override 了 HF 默认行为以支持 action_tokenizer/processor 等 VLA 特有的组件。
840
"""841
Load model from pretrained model path.843
Args:844
pretrained_model_path (str): Model directory path containing model.safetensors file845
config_path (str, optional): Configuration file path, if None will look for qwen25_config.json in pretrained_model_path846
processor_path (str, optional): Processor path, if None will load from default config847
action_tokenizer_path (str, optional): Action tokenizer path, if None will load from default config848
**kwargs: Additional arguments850
Returns:851
Qwen2_5_VLMoEForAction: Loaded model instancedocstring 前 12 行：描述加载流程和所有参数含义。pretrained_model_path 是核心输入，指向包含 model.safetensors 的目录；config_path/processor_path/action_tokenizer_path 若不提供则从默认位置推导；train_config 为 None 时会尝试自动加载 config.yml。返回 Qwen2_5_VLMoEForAction 实例，已初始化所有权重。
852
"""docstring 结束符 """。
853
# Load model components from pretrained path# 注释：声明下面逻辑的意图是从预训练路径加载模型各组件。
855
if train_config is None:
856
try:857
with open(os.path.join(pretrained_model_path, "config.yml"), "r") as f:
858
train_config = yaml.load(f, Loader=yaml.FullLoader)
859
except Exception as e:
860
print(f"load train_config.yml fail: {e}")
861
train_config = Nonetry-except 块：尝试从 pretrained_model_path 目录加载 config.yml 作为 train_config（通常包含 data/model/robot 配置）。若加载失败（文件不存在或格式错误），打印错误信息后令 train_config=None；后续根据 train_config 是否为 None 来决定是否调用 update_model_config/load_wallx_processors（LIBERO 场景）或使用默认 AutoProcessor（推理场景）。
863
model_config_path = os.path.join(pretrained_model_path, "config.json")
864
model_config = cls.config_class.from_pretrained(model_config_path)866
if train_config is not None:
867
model_config = update_model_config(train_config, model_config)
868
processors_dict = load_wallx_processors(train_config)
869
processor = processors_dict["processor"]
870
else:871
processor = AutoProcessor.from_pretrained(
872
pretrained_model_path, use_fast=True873
)
加载模型配置和处理器。从 pretrained_model_path/config.json 加载 Qwen2_5_VLConfig；若有 train_config 则用 update_model_config 融合训练参数，并通过 load_wallx_processors 获取自定义的 processor（包含 LIBERO-specific tokenizer/action_processor 等）；否则用 AutoProcessor.from_pretrained 的默认 processor；这里 processor 包含文本/图像/动作 tokenizer 和 action_processor。
875
if not is_train:
876
model_config._attn_implementation = "sdpa"
878
if action_tokenizer_path is not None:
879
processor.action_processor = AutoProcessor.from_pretrained(
880
action_tokenizer_path, trust_remote_code=True881
)
推理优化和 action_tokenizer 加载。若 is_train=False（推理模式），设置 attn_implementation='sdpa'（使用 Flash Attention 加速）；若显式提供 action_tokenizer_path，从该路径加载 action_processor 覆盖 processor 中默认的（支持加载微调后的 action tokenizer）；这里 AutoProcessor 调用 trust_remote_code=True 允许加载自定义代码。
883
# Set the customized robot configuration to ensure consistency between cross-embodiment884
# representations and the Wall-X action dimensionality.885
# if not train_config:886
# cls._set_customized_config(train_config)887
# customized_dof_config = train_config["customized_robot_config"][888
# "customized_dof_config"889
# ]890
# customized_agent_pos_config = train_config["customized_robot_config"][891
# "customized_agent_pos_config"892
# ]893
# setattr(model_config, "customized_dof_config", customized_dof_config)已注释掉的自定义机器人配置代码块。原本用于从 train_config 提取 customized_dof_config/customized_agent_pos_config 并挂到 model_config，保证跨本体表示和 Wall-X 20 维动作空间一致；目前代码注释掉这部分，改用 ActionProcessor 内部的 dof_mask 机制来处理（LIBERO 只用前 7 维，后 13 维 dof_mask=0）。
894
# setattr(model_config, "customized_agent_pos_config", customized_agent_pos_config)896
# Initialize model with configuration and processor897
model = cls(model_config, processor=processor, **kwargs)899
# Resize token embeddings to match processor tokenizer vocabulary size900
model.resize_token_embeddings(len(processor.tokenizer))初始化模型实例。先空行，注释声明下面创建模型；cls(model_config, processor=processor, **kwargs) 调用 __init__，实例化 Qwen2_5_VLMoEForAction；然后调用 resize_token_embeddings 调整 embed_tokens.weight 大小到 processor.tokenizer 的词表大小（支持自定义 token 扩展）；processor 中包含 action_processor 会在 __init__ 中被设置到模型属性。
902
# Load model state dict from safetensors file903
safetensor_files = glob.glob(
904
os.path.join(pretrained_model_path, "*.safetensors")
905
)
906
state_dict = {}加载 safetensors 权重准备。glob 扫描 pretrained_model_path 下所有 *.safetensors 文件（支持模型被拆成多个 shards 的场景）；初始化空 state_dict 字典；记录 embed_tokens_size 初值为 processor.tokenizer 词表大小（后续若权重中 embed_tokens 更大，会覆盖此值）。
907
embed_tokens_size = len(processor.tokenizer)908
for file in safetensor_files:
909
sd = load_file(file, device="cpu")
910
# filter normalizer statistic params911
del_keys = []
912
for key in sd.keys():
913
if "action_preprocessor.normalizer" in key:
914
print(f"filter load model weight {key}")
循环遍历每个 safetensors 文件，使用 load_file(..., device='cpu') 加载到 CPU（避免显存溢出）。内部初始化 del_keys=[] 列表用于记录需删除的 key；遍历 sd（当前 shard 的 state_dict）中所有 key，过滤掉 'action_preprocessor.normalizer' 开头的权重（因为 normalizer 的统计量会在后续数据加载时重新计算）；若 key 包含 'embed_tokens.weight'，更新 embed_tokens_size 到该权重的词表维度，这允许适应扩展后的词表大小。
915
del_keys.append(key)
916
if "embed_tokens.weight" in key:
917
embed_tokens_size = sd[key].shape[0]918
# if train_config is not None:919
for key in del_keys:
920
del sd[key]921
state_dict.update(sd)
删除被过滤的 key 并合并权重。for 循环逐个删除 del_keys 中的 key（从当前 shard state_dict 中移除 normalizer 权重）；state_dict.update(sd) 将当前 shard 的权重合并到总的 state_dict 中；多个 shard 的权重会依次追加到同一个字典，形成完整的模型权重。
922
if embed_tokens_size != len(processor.tokenizer):
923
model.resize_token_embeddings(embed_tokens_size)
924
model.load_state_dict(state_dict, strict=False)926
return model校验并再次调整 embed_tokens 大小，然后加载权重。若权重中 embed_tokens 的实际大小不等于 processor.tokenizer 的词表大小，再次调用 resize_token_embeddings 同步；model.load_state_dict(state_dict, strict=False) 加载全部权重，strict=False 允许权重中有 unused key（如 normalizer 已被过滤）或模型有额外 module 未在权重中（如新增的 LoRA adapter）；返回加载完成的 model 实例。
from_pretrained 方法返回加载完成的模型实例。
928
def __init__(
929
self,930
config,
931
use_fast_tokenizer=False,932
processor=None,933
action_tokenizer=None,934
action_mapper=None,935
flow_loss_weight=1.0,936
use_selective_recompute=False,937
):
def __init__(...) 方法签名，初始化 Qwen2_5_VLMoEForAction。主要参数：config 是模型配置；processor 包含 text/image/action tokenizer；use_selective_recompute 为 True 时启用梯度重计算（省显存）；action_mapper 是可选的动作映射工具；flow_loss_weight 是 flow matching 目标函数的权重（用于加权多任务损失）；被 from_pretrained 或直接实例化时调用。
938
"""939
Initialize the Qwen2.5 VLMoE model for action processing.941
Args:942
config: Model configuration943
use_fast_tokenizer (bool): Whether to use fast tokenizer944
processor: Text and image processor945
action_tokenizer: Action-specific tokenizer946
action_mapper: Action mapping utility947
flow_loss_weight (float): Weight for flow loss computation948
"""__init__ 的 docstring（11 行）：描述该构造器初始化 VLMoE 模型的目的和参数。config 包含所有模型超参（hidden_size/vocab_size/dof_config 等）；processor 是多模态处理器；flow_loss_weight 用于 flow matching 损失加权；action_tokenizer/action_mapper 为可选的 VLA 扩展组件。
949
super().__init__(config)
调用父类 PreTrainedModel 的 __init__(config)，初始化 HF 框架的基础属性（如 config/generation_config 等）。
951
# Initialize vision transformer and language model components952
self.visual = Qwen2_5_VisionTransformerPretrainedModel._from_config(953
config.vision_config
954
)
955
self.model = Qwen2_5_VLMoEModel(956
config, use_selective_recompute=use_selective_recompute
957
)
958
self.vocab_size = config.vocab_size959
self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
初始化 VL 模型的核心组件。注释声明下面初始化 vision transformer 和语言模型；self.visual = Qwen2_5_VisionTransformerPretrainedModel._from_config(config.vision_config) 创建视觉编码器（处理图像输入）；self.model = Qwen2_5_VLMoEModel(...) 创建多专家 VL decoder（包含 MoE 路由、双专家 FFN、cross-attention 等）；self.vocab_size/self.lm_head 初始化文本输出投影头，将隐状态 (B,S,H) 投到 (B,S,vocab_size) 用于语言建模。
961
# Initialize loss function without reduction for channel-wise loss computation962
self.loss_fct = CrossEntropyLoss(reduction="none")
963
self.flow_loss_weight = flow_loss_weight964
self.use_fast_tokenizer = use_fast_tokenizer965
self.processor = processor初始化训练损失函数和处理器。CrossEntropyLoss(reduction='none') 不做归约，保持 (B,S) 形状以支持逐 token 的加权（LIBERO 中按 dof_mask 逐维加权）；flow_loss_weight 存储 flow matching 损失权重；use_fast_tokenizer/processor 分别记录是否使用快速 tokenizer 和多模态处理器（后续 forward 时调用 processor.action_processor 处理动作输入）。
定义动作 token ID 映射和时间步缓存。self.define_action_token_id() 调用方法（见 989 行）初始化 fast_action_token_list/propri_token_id/action_token_id 映射，用于在解码时识别动作/本体感觉 token；self.times_cache = {} 用于缓存不同 num_inference_timesteps 对应的时间步 linspace，避免 generate_flow_action 推理时重复计算。
969
self.times_cache = {} # cache times linspace for each num_inference_timesteps
971
# Cache for rope deltas972
self.rope_deltas = None
974
# Initialize action preprocessor975
self.action_preprocessor = ActionProcessor(config)初始化绳索（RoPE）缓存和动作处理器。空行+注释；self.rope_deltas = None 初始化（后续 forward 中计算）；self.action_preprocessor = ActionProcessor(config) 创建动作处理器，内含 normalizer（查 dof_mask/min/max 表归一化）、时间步正弦编码、action_head（w1/w2/w3 投影）等（详见 action_head.py）；LIBERO 场景下 ActionProcessor 会按 dataset_name 查表获取 dof_mask 和 min/max 统计量。
977
# Apply LoRA if specified in configuration978
if hasattr(config, "use_lora") and config.use_lora:
979
self.add_lora(980
r=config.lora_r,
981
lora_alpha=config.lora_alpha,
982
target_modules=config.lora_target_modules,
983
lora_dropout=config.lora_dropout,
984
)
条件化应用 LoRA 适配器。注释声明下面是 LoRA 初始化；if hasattr(config, 'use_lora') and config.use_lora: 检查 config 中是否启用了 LoRA；若启用，调用 self.add_lora(r/lora_alpha/target_modules/lora_dropout) 给模型的 q_proj/v_proj/etc 线性层插入低秩适配器（常用于高效微调，3090 上 freeze_vlm 微调 LIBERO 时使用）；LoRA 相关参数从 config 中读取（config.lora_r/lora_alpha/lora_target_modules/lora_dropout）。
最后的初始化和权重设置。空行+注释；self.post_init() 调用 PreTrainedModel 的 post_init hook，初始化所有权重（使用 config.initializer_range 等超参做高斯初始化），用于新增的层（action_head/lora adapter 等）；确保模型可以立即开始训练或推理，无需手动初始化。
L989–1065定义 action token 映射、LoRA 微调接口、和 HuggingFace 模型标准 embedding/decoder getter/setter 方法，完成模型初始化后的 token ID 配置和可选微调
989
def define_action_token_id(self):
990
"""991
Define action token IDs based on tokenizer configuration.993
Creates mappings for fast action tokens, proprioception tokens, and general action tokens.994
"""define_action_token_id() 方法定义+docstring，被 __init__():968 调用；作用是为推理时快速查询做准备，初始化 self.action_token_id_set 字典，存储 3 类 token ID 映射（快速 action token、proprioception token、通用 action token）
初始化空列表 fast_action_token_list，后续按需填充快速 action token IDs
997
if self.use_fast_tokenizer:
998
for i in range(
999
self.processor.tokenizer.init_kwargs["action_token_vocab_size"]
1000
):
1001
action_token_id = self.processor.tokenizer.convert_tokens_to_ids(1002
f"<|action_token_{i}|>"
1003
)
1004
fast_action_token_list.append(action_token_id)
if self.use_fast_tokenizer 块（997-1004）+空行：当启用快速 tokenizer 时，遍历 action_token_vocab_size 个 token，每个调用 convert_tokens_to_ids() 将 '<|action_token_{i}|>' 特殊符号转换为词表 ID，追加到列表；该列表用于前向推理时直接索引预定义的 action token，避免每次都转换字符串
1006
# Get special action token IDs1007
action_token_id = self.processor.tokenizer.convert_tokens_to_ids("<|action|>")
1008
propri_token_id = self.processor.tokenizer.convert_tokens_to_ids("<|propri|>")
注释行 + 获取两个固定特殊 token 的词表 ID：'<|action|>' 用于标记 action 序列段，'<|propri|>' 用于标记本体状态（proprioception）段；这两个是 tokenizer 初始化时预注册的，ID 值固定且全数据集一致 + 空行
1010
# Store action token ID mappings1011
self.action_token_id_set = {1012
"fast_action_token_list": fast_action_token_list,
1013
"propri_token_id": propri_token_id,
1014
"action_token_id": action_token_id,
1015
}
注释行 + 将 3 种 token ID 存入 self.action_token_id_set 字典：fast_action_token_list(若启用快速 tokenizer 则为长度 = action_token_vocab_size 的列表，否则空列表），propri_token_id（int），action_token_id（int）；这个字典在 forward() 和 generate_flow_action() 中用于识别 token 类型，路由到不同专家 + 空行
1017
def add_lora(
1018
self, r=8, lora_alpha=32, target_modules=["q_proj", "v_proj"], lora_dropout=0.1
1019
):
1020
"""1021
Add LoRA (Low-Rank Adaptation) adapters to the model.1023
Args:1024
r (int): Rank of adaptation1025
lora_alpha (int): LoRA scaling parameter1026
target_modules (list): List of module names to apply LoRA to1027
lora_dropout (float): Dropout probability for LoRA layers1028
"""add_lora() 方法签名 + docstring：被 __init__():978-984 可选调用；作用是如果 config.use_lora==True，给模型添加 LoRA 低秩适应层进行高效微调；参数 r(秩，默认 8)、lora_alpha(缩放因子，默认 32)、target_modules(应用 LoRA 的模块名列表，默认 q_proj/v_proj 注意力投影)、lora_dropout(LoRA 层 dropout 比例，默认 0.1)
1029
config = LoraConfig(
1030
r=r,
1031
lora_alpha=lora_alpha,
1032
target_modules=target_modules,
1033
lora_dropout=lora_dropout,
1034
bias="none",
1035
task_type="CAUSAL_LM",
1036
)
1037
self.model = get_peft_model(self.model, config)
1039
# Print information about trainable parameters1040
self.model.print_trainable_parameters()构建 LoraConfig 对象（1029-1036），指定秩、缩放因子、目标模块、dropout、无偏置、任务类型为 CAUSAL_LM；然后调用 get_peft_model() 用 PEFT 库包装 self.model，将指定模块替换为 LoRA 适配器（1037）；最后打印可训练参数数量（1039-1040） + 空行；这样只需训练少量参数（约 0.01-0.1 倍原模型），大幅降低显存和计算成本，实战中 LIBERO 微调通常 trainable_params 占 0.5-1%
1042
def get_input_embeddings(self):
1043
"""Get input embeddings layer."""1044
return self.model.embed_tokens
get_input_embeddings() 方法：HuggingFace 标准接口，返回 self.model.embed_tokens，用于外部访问输入 embedding 层；在推理/可视化时被上游调用 + 空行
1046
def set_input_embeddings(self, value):
1047
"""Set input embeddings layer."""1048
self.model.embed_tokens = valueset_input_embeddings() 方法：HuggingFace 标准接口，更新 self.model.embed_tokens 为新权重；用于共享 embedding（如 tied weights）或初始化时替换 + 空行
1050
def get_output_embeddings(self):
1051
"""Get output embeddings layer."""1052
return self.lm_head
get_output_embeddings() 方法：HuggingFace 标准接口，返回 self.lm_head，用于外部访问输出投影层（vocab projection）；在解码评估/可视化时被上游调用 + 空行
1054
def set_output_embeddings(self, new_embeddings):
1055
"""Set output embeddings layer."""1056
self.lm_head = new_embeddingsset_output_embeddings() 方法：HuggingFace 标准接口，更新 self.lm_head 为新权重；用于初始化时替换或多任务微调时切换输出投影 + 空行
set_decoder() 方法：HuggingFace 标准接口，直接替换 self.model 为新的 decoder；在模型重组件/蒸馏时被调用，为了保持接口兼容性而设置 + 空行
get_decoder() 方法：HuggingFace 标准接口，返回 self.model，用于外部访问 decoder 模块；在模型恢复/导出时被上游调用，返回的是底层 Qwen2_5_VLMoEModel
空行，方法定义结束；下一个方法 get_rope_index() 在 1066 开始
L1066–1286get_rope_index 方法：为混合视觉-文本序列计算 3D RoPE 位置索引。处理 image/video 的时空 3D 坐标（时间、高、宽）与纯文本的 1D 坐标，输出 (3,B,S) 位置张量和批次偏移量，支持 Qwen2.5-VL 的多模态旋转位置编码。
1066
def get_rope_index(
1067
self,1068
input_ids: Optional[torch.LongTensor] = None,1069
image_grid_thw: Optional[torch.LongTensor] = None,1070
video_grid_thw: Optional[torch.LongTensor] = None,1071
second_per_grid_ts: Optional[torch.Tensor] = None,1072
attention_mask: Optional[torch.Tensor] = None,1073
) -> Tuple[torch.Tensor, torch.Tensor]:
函数签名及类型注解。get_rope_index 被 forward():1427 和 train_step_forward():1287 调用，输入包括 input_ids (B,S)、image_grid_thw 和 video_grid_thw 分别为各 image/video 的 (T,H,W)、second_per_grid_ts 为视频时间间隔、attention_mask (B,S)；返回 position_ids (3,B,S) 与 mrope_position_deltas (B,1)。
1074
"""1075
Calculate the 3D rope index based on image and video's temporal, height and width in LLM.docstring 摘要：计算 3D RoPE 索引。针对混合视觉（image/video）与文本的序列，视觉部分生成 3D 坐标（时间-高-宽），文本部分生成 1D 坐标。
1077
Explanation:1078
Each embedding sequence contains vision embedding and text embedding or just contains text embedding.1080
For pure text embedding sequence, the rotary position embedding has no difference with modern LLMs.1081
Examples:1082
input_ids: [T T T T T], here T is for text.1083
temporal position_ids: [0, 1, 2, 3, 4]1084
height position_ids: [0, 1, 2, 3, 4]1085
width position_ids: [0, 1, 2, 3, 4]1087
For vision and text embedding sequence, we calculate 3D rotary position embedding for vision part1088
and 1D rotary position embeddin for text part.1089
Examples:1090
Temporal (Time): 3 patches, representing different segments of the video in time.1091
Height: 2 patches, dividing each frame vertically.1092
Width: 2 patches, dividing each frame horizontally.1093
We also have some important parameters:1094
fps (Frames Per Second): The video's frame rate, set to 1. This means one frame is processed each second.1095
tokens_per_second: This is a crucial parameter. It dictates how many "time-steps" or "temporal tokens" are conceptually packed into a one-second interval of the video. In this case, we have 25 tokens per second. So each second of the video will be represented with 25 separate time points. It essentially defines the temporal granularity.1096
temporal_patch_size: The number of frames that compose one temporal patch. Here, it's 2 frames.1097
interval: The step size for the temporal position IDs, calculated as tokens_per_second * temporal_patch_size / fps. In this case, 25 * 2 / 1 = 50. This means that each temporal patch will be have a difference of 50 in the temporal position IDs.1098
input_ids: [V V V V V V V V V V V V T T T T T], here V is for vision.1099
vision temporal position_ids: [0, 0, 0, 0, 50, 50, 50, 50, 100, 100, 100, 100]1100
vision height position_ids: [0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1]1101
vision width position_ids: [0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1]1102
text temporal position_ids: [101, 102, 103, 104, 105]1103
text height position_ids: [101, 102, 103, 104, 105]1104
text width position_ids: [101, 102, 103, 104, 105]1105
Here we calculate the text start position_ids as the max vision position_ids plus 1.docstring Explanation 段落（30 行）：详细说明纯文本与视觉-文本混合序列的 RoPE 编码差异。纯文本为 1D 线性位置；视觉部分按 (T,H,W) 三维网格展开，其中 tokens_per_second=4 定义时间粒度、spatial_merge_size 影响高宽分割、second_per_grid_t 为视频每个时间片段占用秒数，最后计算间隔 interval=tokens_per_second*temporal_patch_size/fps；文本部分续接在视觉部分之后，位置 ID 从视觉最大值+1 开始。
1107
Args:1108
input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):1109
Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide1110
it.1111
image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):1112
The temporal, height and width of feature shape of each image in LLM.1113
video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):1114
The temporal, height and width of feature shape of each video in LLM.1115
second_per_grid_ts (`torch.Tensor` of shape `(num_videos)`, *optional*):1116
The time interval (in seconds) for each grid along the temporal dimension in the 3D position IDs.1117
attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):1118
Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:1120
- 1 for tokens that are **not masked**,1121
- 0 for tokens that are **masked**.docstring Args 段落（16 行）：参数说明。input_ids (B,S) 输入 token IDs，padding 被 attention_mask 屏蔽；image_grid_thw (num_images,3) 各 image 的 LLM 特征空间尺寸；video_grid_thw (num_videos,3) 各 video 的 LLM 特征空间尺寸；second_per_grid_ts (num_videos,) 视频帧率调整参数；attention_mask (B,S) 标记有效位置为 1、填充为 0。
1123
Returns:1124
position_ids (`torch.LongTensor` of shape `(3, batch_size, sequence_length)`)1125
mrope_position_deltas (`torch.Tensor` of shape `(batch_size)`)1126
"""docstring Returns 段落：返回两个张量。position_ids (3,B,S) 包含时间/高/宽三维坐标；mrope_position_deltas (B,) 为每条序列的最大坐标与原序列长度的偏移，用于调整 RoPE 的缩放。
1127
spatial_merge_size = self.config.vision_config.spatial_merge_size1128
image_token_id = self.config.image_token_id1129
video_token_id = self.config.video_token_id1130
vision_start_token_id = self.config.vision_start_token_id1131
mrope_position_deltas = []
1132
if input_ids is not None and (
1133
image_grid_thw is not None or video_grid_thw is not None
1134
):
1135
total_input_ids = input_ids
1136
if attention_mask is None:
1137
attention_mask = torch.ones_like(total_input_ids)
初始化阶段。读取 config 中的 spatial_merge_size（通常=2，用来合并特征的高宽）、image_token_id/video_token_id（特殊标记）、vision_start_token_id（视觉起始标记）；初始化 mrope_position_deltas 列表、检查输入是否同时包含 input_ids 和视觉信息。
1138
position_ids = torch.ones(
1139
3,1140
input_ids.shape[0],1141
input_ids.shape[1],1142
dtype=input_ids.dtype,
1143
device=input_ids.device,
1144
)
1145
image_index, video_index = 0, 0
创建 position_ids 张量 (3, B, S)，shape=(3, input_ids.shape[0], input_ids.shape[1])，初值全 1，以便后续修改；初始化 image_index=0、video_index=0 用于遍历 image_grid_thw/video_grid_thw 数组。
1146
attention_mask = attention_mask.to(total_input_ids.device)
1147
for i, input_ids in enumerate(total_input_ids):
1148
input_ids = input_ids[attention_mask[i] == 1]1149
image_nums, video_nums = 0, 0
attention_mask 转移到 input_ids 所在设备；开始对 batch 中每条序列遍历（i 为批次索引，input_ids 为该序列的 token IDs）。
1150
vision_start_indices = torch.argwhere(
1151
input_ids == vision_start_token_id
1152
).squeeze(1)1153
vision_tokens = input_ids[vision_start_indices + 1]1154
image_nums = (vision_tokens == image_token_id).sum()
1155
video_nums = (vision_tokens == video_token_id).sum()
使用 attention_mask 过滤出有效 token（mask==1 的位置），计算该序列中实际的 image/video 数量：找到所有 vision_start_token_id，取其后一个 token，统计其中 image_token_id 和 video_token_id 的个数，用于后续逐个处理每个视觉模块。
1156
input_tokens = input_ids.tolist()
1157
llm_pos_ids_list: list = []1158
st = 01159
remain_images, remain_videos = image_nums, video_nums
将 token IDs 转换为 Python list 便于索引查找；llm_pos_ids_list 存储该序列各段落的 position IDs 张量；st=0 为当前扫描位置，remain_images/remain_videos 为还未处理的视觉模块计数。
1160
for _ in range(image_nums + video_nums):
1161
if image_token_id in input_tokens and remain_images > 0:
1162
ed_image = input_tokens.index(image_token_id, st)
1163
else:1164
ed_image = len(input_tokens) + 1
1165
if video_token_id in input_tokens and remain_videos > 0:
1166
ed_video = input_tokens.index(video_token_id, st)
1167
else:1168
ed_video = len(input_tokens) + 1
外层循环遍历 image_nums+video_nums 次（总视觉模块数）；内部分别在 input_tokens 中查找下一个 image_token_id 和 video_token_id 的位置（从 st 开始搜索），若找不到则返回 len(input_tokens)+1（哨兵值），以确定哪个类型先出现。
1169
if ed_image < ed_video:1170
t, h, w = (
1171
image_grid_thw[image_index][0],1172
image_grid_thw[image_index][1],1173
image_grid_thw[image_index][2],1174
)
1175
second_per_grid_t = 01176
image_index += 11177
remain_images -= 11178
ed = ed_image
若 image 更早出现（ed_image < ed_video），从 image_grid_thw[image_index] 读取该 image 的 (T,H,W)，second_per_grid_t=0（image 无时间维），更新 image_index 和 remain_images，设 ed=ed_image 为该模块的结束位置。
1180
else:1181
t, h, w = (
1182
video_grid_thw[video_index][0],1183
video_grid_thw[video_index][1],1184
video_grid_thw[video_index][2],1185
)
1186
if second_per_grid_ts is not None:
1187
second_per_grid_t = second_per_grid_ts[video_index]
1188
else:1189
second_per_grid_t = 1.01190
video_index += 11191
remain_videos -= 11192
ed = ed_video
否则 video 先出现或同时出现，从 video_grid_thw[video_index] 读取 (T,H,W)；根据 second_per_grid_ts[video_index] 计算该 video 每时间片占用秒数（若 second_per_grid_ts=None 则默认 1.0，配合 tokens_per_second 计算实际时间间隔），更新 video_index 和 remain_videos，设 ed=ed_video。
1193
llm_grid_t, llm_grid_h, llm_grid_w = (
1194
t.item(),
1195
h.item() // spatial_merge_size,
1196
w.item() // spatial_merge_size,
1197
)
计算 LLM 空间中的 grid 维度：T 不变、H_llm=H//spatial_merge_size、W_llm=W//spatial_merge_size（spatial_merge_size=2 表示相邻 4 个特征被合并为 1）；text_len=ed-st 为该视觉模块前方的 text token 数。
1198
text_len = ed - st
1200
st_idx = (
1201
llm_pos_ids_list[-1].max() + 1
1202
if len(llm_pos_ids_list) > 0
1203
else 0
1204
)
1205
llm_pos_ids_list.append(
1206
torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx
1207
)
计算该视觉模块前的纯 text 部分的 position IDs：st_idx 为该部分的起始位置（等于上一段落的最大坐标+1，若无前序段落则=0）；创建 (3, text_len) 张量，三维都是 [0,1,...,text_len-1]+st_idx 的线性序列（文本无空间-时间结构，只有序列位置）；第 1208 行空行并入此组。
1209
range_tensor = torch.arange(llm_grid_t).view(-1, 1)
1210
expanded_range = range_tensor.expand(-1, llm_grid_h * llm_grid_w)1212
time_tensor = (
1213
expanded_range
1214
* second_per_grid_t
1215
* self.config.vision_config.tokens_per_second1216
)
1218
time_tensor_long = time_tensor.long()
1219
t_index = time_tensor_long.flatten()
计算视觉部分的时间索引：range_tensor (llm_grid_t,1) 扩展为 (llm_grid_t, llm_grid_h*llm_grid_w)；time_tensor 将时间索引乘以 second_per_grid_t 和 tokens_per_second（得到每个时间片对应的时间坐标），转为 long 类型并展平为 t_index；第 1220 行空行并入此组。
1221
h_index = (
1222
torch.arange(llm_grid_h)
1223
.view(1, -1, 1)
1224
.expand(llm_grid_t, -1, llm_grid_w)1225
.flatten()
1226
)
1227
w_index = (
1228
torch.arange(llm_grid_w)
1229
.view(1, 1, -1)
1230
.expand(llm_grid_t, llm_grid_h, -1)1231
.flatten()
1232
)
1233
llm_pos_ids_list.append(
1234
torch.stack([t_index, h_index, w_index]) + text_len + st_idx
1235
)
1236
st = ed + llm_grid_t * llm_grid_h * llm_grid_w
计算高度和宽度索引：h_index 用 torch.arange(llm_grid_h) 沿第二维展开，再跨 llm_grid_t 倍复制，最后展平；w_index 类似，用 torch.arange(llm_grid_w) 沿第三维展开，跨 llm_grid_t*llm_grid_h 倍复制，最后展平。将 t_index、h_index、w_index 堆叠为 (3, llm_grid_t*llm_grid_h*llm_grid_w) 张量，加上 (text_len+st_idx)；st 更新为下一循环做准备；第 1237 行空行并入此组。
1238
if st < len(input_tokens):
1239
st_idx = (
1240
llm_pos_ids_list[-1].max() + 1
1241
if len(llm_pos_ids_list) > 0
1242
else 0
1243
)
1244
text_len = len(input_tokens) - st1245
llm_pos_ids_list.append(
1246
torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx
1247
)
若循环后还有剩余 token（st < len(input_tokens)），说明序列末尾是纯文本，对其计算 1D position IDs（方法同 1205-1207），从 st 到 len(input_tokens) 的部分视为 text_len=len(input_tokens)-st，st_idx 为上一段落最大值+1；第 1248 行空行并入此组。
1249
llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)
1250
position_ids[..., i, attention_mask[i] == 1] = llm_positions.to(1251
position_ids.device
1252
)
1253
mrope_position_deltas.append(
1254
llm_positions.max() + 1 - len(total_input_ids[i])
1255
)
将 llm_pos_ids_list（包含该序列各段落的 position IDs）沿 dim=1 拼接，得到 (3,S') 张量（S' 为有效 token 数）；赋值回 position_ids[i, attention_mask[i]==1]（只更新有效位置）；计算 mrope_position_deltas 为该序列最大坐标+1 与原 token 数的差。
1256
mrope_position_deltas = torch.tensor(
1257
mrope_position_deltas, device=input_ids.device
1258
).unsqueeze(1)1259
return position_ids, mrope_position_deltas将 mrope_position_deltas 列表转为张量 (B,)、移到 input_ids 所在设备、reshape 为 (B,1)，准备返回；position_ids 此时已是 (3,B,S) 完整张量。
1260
else:1261
if attention_mask is not None:
1262
position_ids = attention_mask.long().cumsum(-1) - 1
1263
position_ids.masked_fill_(attention_mask == 0, 1)
1264
position_ids = (
1265
position_ids.unsqueeze(0)1266
.expand(3, -1, -1)
1267
.to(attention_mask.device)
1268
)
1269
max_position_ids = position_ids.max(0, keepdim=False)[0].max(
1270
-1, keepdim=True
1271
)[0]1272
mrope_position_deltas = max_position_ids + 1 - attention_mask.shape[-1]
else 分支：无 image/video 信息时的降级处理。若有 attention_mask，用其生成位置 IDs：cumsum(-1)-1 计算掩码后的位置（累积和），masked_fill_ 把 mask==0 的位置填充为 1（保证全正）；unsqueeze 和 expand 到 (3,B,S) 三维；求 position_ids 的最大值以计算 mrope_position_deltas=最大值+1-原序列长度。
1273
else:1274
position_ids = (
1275
torch.arange(input_ids.shape[1], device=input_ids.device)1276
.view(1, 1, -1)
1277
.expand(3, input_ids.shape[0], -1)
1278
)
1279
mrope_position_deltas = torch.zeros(
1280
[input_ids.shape[0], 1],
1281
device=input_ids.device,
1282
dtype=input_ids.dtype,
1283
)
else 分支：若无 attention_mask，直接用 torch.arange(input_ids.shape[1]) 创建 (1,1,S) 位置序列，expand 到 (3,B,S)；mrope_position_deltas 全 0（无有效数据偏移），保证类型与设备一致；第 1284 行空行并入此组。
返回 position_ids (3,B,S) 和 mrope_position_deltas (B,1)。这两个张量后续被传入 self.model(...) 的前向传播，供 RoPE 旋转操作使用。
L1287–1484train_step_forward 函数：VLA 训练主前向传播，集成图文/视频/本体感觉嵌入、MoE token 路由、flow matching 损失计算，支持 KV cache 和 RoPE 增量位置编码。
1287
def train_step_forward(
1288
self,1289
input_ids: torch.LongTensor = None,1290
attention_mask: Optional[torch.Tensor] = None,1291
position_ids: Optional[torch.LongTensor] = None,1292
past_key_values: Optional[List[torch.FloatTensor]] = None,1293
inputs_embeds: Optional[torch.FloatTensor] = None,1294
labels: Optional[torch.LongTensor] = None,1295
use_cache: Optional[bool] = None,
1296
output_attentions: Optional[bool] = None,
1297
output_hidden_states: Optional[bool] = None,
1298
return_dict: Optional[bool] = None,
1299
pixel_values: Optional[torch.Tensor] = None,1300
pixel_values_videos: Optional[torch.FloatTensor] = None,1301
image_grid_thw: Optional[torch.LongTensor] = None,1302
video_grid_thw: Optional[torch.LongTensor] = None,1303
rope_deltas: Optional[torch.LongTensor] = None,1304
cache_position: Optional[torch.LongTensor] = None,1305
second_per_grid_ts: Optional[torch.Tensor] = None,1306
# for vla1307
moe_token_types: Optional[torch.LongTensor] = None,1308
start_indices: Optional[torch.Tensor] = None,1309
end_indices: Optional[torch.Tensor] = None,1310
positional_masks: Optional[dict] = None,
1311
action_chunk: Optional[torch.FloatTensor] = None,1312
proprioception: Optional[torch.FloatTensor] = None,1313
dataset_names: Optional[str] = None,
1314
dof_mask: Optional[torch.FloatTensor] = None,1315
agent_pos_mask: Optional[torch.FloatTensor] = None,1316
flow_loss_mask: Optional[torch.FloatTensor] = None,1317
**kwargs,
1318
) -> Union[Tuple, Qwen2_5_VLACausalLMOutputWithPast]:
train_step_forward 函数签名：被 Trainer 在每个训练 step 调用。包含 32 个参数：标准 LLM 参数（input_ids、attention_mask 等）+ VLA 特定参数（moe_token_types、action_chunk、dof_mask、flow_loss_mask 等）+ 视觉参数（pixel_values、image_grid_thw 等）。返回 Qwen2_5_VLACausalLMOutputWithPast 对象，包含 loss、flow_loss、logits、rope_deltas 等。
从 input_ids shape 提取 batch_size 和 seq_length，为后续位置编码和 reshape 操作提供维度信息。
1322
output_attentions = (
1323
output_attentions
1324
if output_attentions is not None
1325
else self.config.output_attentions
1326
)
1327
output_hidden_states = (
1328
output_hidden_states
1329
if output_hidden_states is not None
1330
else self.config.output_hidden_states
1331
)
1332
return_dict = (
1333
return_dict if return_dict is not None else self.config.use_return_dict
1334
)
处理 output_attentions/output_hidden_states/return_dict 的默认值。若为 None，使用 self.config 中的配置；否则使用传入的值，允许动态控制模型输出内容。
1336
if start_indices is None or end_indices is None:
1337
# Calculate the start and end positions of each expert group's tokens after permutation1338
group_size = torch.zeros(
1339
self.config.num_experts, dtype=torch.long, device="cpu"
1340
)
1341
for i in range(self.config.num_experts):
1342
group_size[i] = (moe_token_types == i).sum()
1344
# Calculate start and end indices for each expert group1345
start_indices = torch.cumsum(group_size, dim=0) - group_size1346
end_indices = torch.cumsum(group_size, dim=0)若未提供 start_indices/end_indices，则从 moe_token_types 计算专家分组位置。对每个专家（num_experts=2：专家0=图文、专家1=动作/状态），统计其 token 数量 (B=batch, num_experts=2)；cumsum 得 start/end indices，用于后续稀疏 MoE 路由的 token 排序。
1348
# if we get 4D attention mask we cannot calculate rope deltas anymore. TODO @raushan fixme1349
if position_ids is None and (
1350
attention_mask is None or attention_mask.ndim == 2
1351
):
1352
# calculate RoPE index once per generation in the pre-fill stage only1353
if (1354
(cache_position is not None and cache_position[0] == 0)
1355
or self.rope_deltas is None
1356
or (past_key_values is None or past_key_values.get_seq_length() == 0)
1357
):
1358
position_ids, rope_deltas = self.get_rope_index(1359
input_ids=input_ids,
1360
image_grid_thw=image_grid_thw,
1361
video_grid_thw=video_grid_thw,
1362
second_per_grid_ts=second_per_grid_ts,
1363
attention_mask=attention_mask,
1364
)
1365
self.rope_deltas = rope_deltas首次计算 position_ids 和 rope_deltas（预填充阶段）：当 cache_position[0]==0 或 rope_deltas 不存在或 KV cache 为空时触发。调用 get_rope_index 根据图文/视频 token 网格重新计算 RoPE 位置索引，考虑多模态内容的特殊位置对齐；rope_deltas 保存在 self.rope_deltas 以供推理阶段复用。
1366
# then use the prev pre-calculated rope-deltas to get the correct position ids1367
else:1368
# batch_size, seq_length, _ = inputs_embeds.shape1369
delta = (
1370
(cache_position[0] + self.rope_deltas).to(cache_position.device)
1371
if cache_position is not None
1372
else 0
1373
)
1374
position_ids = torch.arange(seq_length, device=cache_position.device)
1375
position_ids = position_ids.view(1, -1).expand(batch_size, -1)
1376
if cache_position is not None: # otherwise `deltas` is an int `0`
1377
delta = delta.repeat_interleave(batch_size // delta.shape[0], dim=0)
1378
position_ids = position_ids.add(delta)
1379
position_ids = position_ids.unsqueeze(0).expand(3, -1, -1)
推理阶段使用预计算的 rope_deltas：增量调整 position_ids。首先计算 delta = cache_position[0] + self.rope_deltas，按 batch_size 重复扩展；然后为当前 seq 生成 arange(seq_length) 并加上 delta；最后 unsqueeze(0) 扩展到 3 行（多头配置），形状 (3, B, S)。shape：(B,S)→(1,1,S)→(3,B,S)，位置编码支持 KV cache 增量。
若 inputs_embeds 未提供，先用 embed_tokens 将 input_ids 转换为词嵌入 (B,S)→(B,S,H)，其中 H=hidden_size=2048。
1383
if pixel_values is not None:
1384
pixel_values = pixel_values.type(self.visual.dtype)1385
image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw)1386
mask = input_ids == self.config.image_token_id1387
mask_unsqueezed = mask.unsqueeze(-1)1388
mask_expanded = mask_unsqueezed.expand_as(inputs_embeds)
1389
image_mask = mask_expanded.to(inputs_embeds.device)
1391
image_embeds = image_embeds.to(
1392
inputs_embeds.device, inputs_embeds.dtype
1393
)
1394
inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds)
若存在图像输入 pixel_values，经过视觉编码器得 image_embeds，再用 masked_scatter 将对应位置的词嵌入替换为图像特征。逻辑：mask=input_ids==image_token_id，然后 scatter；image_embeds 已在视觉编码器中处理网格对齐，无需额外处理。
1396
if pixel_values_videos is not None:
1397
pixel_values_videos = pixel_values_videos.type(self.visual.dtype)1398
video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw)1399
n_video_tokens = (input_ids == self.config.video_token_id).sum().item()1400
n_video_features = video_embeds.shape[0]1401
if n_video_tokens != n_video_features:1402
raise ValueError(1403
f"Video features and video tokens do not match: tokens: {n_video_tokens}, features {n_video_features}"
1404
)
1406
mask = input_ids == self.config.video_token_id1407
mask_unsqueezed = mask.unsqueeze(-1)1408
mask_expanded = mask_unsqueezed.expand_as(inputs_embeds)
1409
video_mask = mask_expanded.to(inputs_embeds.device)
1411
video_embeds = video_embeds.to(
1412
inputs_embeds.device, inputs_embeds.dtype
1413
)
1414
inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds)
若存在视频输入 pixel_values_videos，类似图像处理但先校验 token 数与特征数匹配。视频通过 visual(pixel_values_videos, grid_thw=video_grid_thw) 编码，n_video_tokens/n_video_features 不一致则报错（防止维度错配）；最后 masked_scatter 替换词嵌入。
1416
inputs_embeds = self.scatter_proprioception_embeddings(1417
input_ids, inputs_embeds, proprioception, dataset_names, agent_pos_mask
1418
)
调用 scatter_proprioception_embeddings 将本体感觉特征（关节位置、速度等）嵌入到对应 token 位置。该方法在 vla_mixin.py:387 定义，内部按 dataset_name 查 normalizer 表，投影到 H 维后 scatter；LIBERO 场景下 proprioception_proj 处理 7 个 dof 加上 action_horizon 维度。
1420
inputs_embeds, flow, adarms_cond = self.scatter_flow_action_embeddings(1421
input_ids, inputs_embeds, action_chunk, dataset_names, dof_mask
1422
)
调用 scatter_flow_action_embeddings 获得加噪动作嵌入 + flow 目标 + adarms_cond。返回修改后的 inputs_embeds、flow（形状 (B,A,D) 其中 A=action_horizon=10、D=dof=20，LIBERO 时前 7 维有意义）以及 adarms_cond（动作条件，仅在需要时计算）。
若 attention_mask 非空，转移到 inputs_embeds 的设备，确保后续 model() 调用时 attention_mask 与 inputs 在同一设备。
1427
outputs = self.model(1428
input_ids=None,1429
position_ids=position_ids,
1430
attention_mask=attention_mask,
1431
past_key_values=past_key_values,
1432
inputs_embeds=inputs_embeds,
1433
moe_token_types=moe_token_types,
1434
start_indices=start_indices,
1435
end_indices=end_indices,
1436
positional_masks=positional_masks,
1437
use_cache=use_cache,
1438
output_attentions=output_attentions,
1439
output_hidden_states=output_hidden_states,
1440
return_dict=return_dict,
1441
adarms_conds=[None, adarms_cond],1442
# cache_position=cache_position,1443
)
调用 self.model 前向传播（Qwen2.5-VL with MoE）：传入已嵌入的输入序列 inputs_embeds、位置编码 position_ids、MoE 参数 (moe_token_types、start_indices、end_indices)、KV cache、adarms_cond（动作条件投入 MoE 专家选择）。模型内部：注意力跨图文和动作 token（attention_moe=false），FFN 按 token_type 路由到专家0（图文）或专家1（动作/状态）。返回 outputs=(hidden_states, past_key_values, hidden_states_tuple, attentions_tuple)。
1446
logits = self.lm_head(hidden_states)将 hidden_states 经过 lm_head（线性投影层）生成 logits，形状 (B,S,vocab_size)，用于语言建模损失。
1448
(
1449
loss,
1450
cross_entropy_loss,
1451
flow_loss,
1452
channel_loss_dict,
1453
channel_loss_count_dict,
1454
) = self.compute_loss(1455
hidden_states=hidden_states,
1456
logits=logits,
1457
input_ids=input_ids,
1458
dataset_names=dataset_names,
1459
labels=labels,
1460
action_chunk=action_chunk,
1461
dof_mask=dof_mask,
1462
flow=flow,
1463
flow_loss_mask=flow_loss_mask,
1464
)
调用 compute_loss 计算多目标损失（定义在 vla_mixin.py:763）。传入 hidden_states、logits、input_ids、dataset_names、labels（language token）、action_chunk（目标动作）、dof_mask（维度掩码，LIBERO 仅前 7 维=1）、flow（flow matching 监督目标）、flow_loss_mask（可选的动作 token 选择）。返回：总 loss、cross_entropy_loss（语言建模）、flow_loss（动作 MSE）、channel_loss_dict/count_dict（按数据集统计，目前未启用）。
1466
if not return_dict:
1467
output = (logits,) + outputs[1:]1468
return (loss,) + output if loss is not None else output
若不返回字典（return_dict=False），则返回元组 (loss, logits, past_key_values, ...)；否则构造 Qwen2_5_VLACausalLMOutputWithPast 对象。
1470
return Qwen2_5_VLACausalLMOutputWithPast(1471
loss=loss,
1472
cross_entropy_loss=(
1473
cross_entropy_loss.clone() if cross_entropy_loss is not None else None
1474
),
1475
flow_loss=flow_loss,
1476
logits=logits,
1477
past_key_values=outputs.past_key_values,
1478
hidden_states=outputs.hidden_states,
1479
attentions=outputs.attentions,
1480
rope_deltas=self.rope_deltas,1481
channel_loss_dict=channel_loss_dict,
1482
channel_loss_count_dict=channel_loss_count_dict,
1483
)
返回完整输出对象 Qwen2_5_VLACausalLMOutputWithPast：包含 loss（总损失）、cross_entropy_loss（NLP 部分）、flow_loss（动作 MSE，仅 action_chunk 非空时计算）、logits（(B,S,vocab_size)）、past_key_values（KV cache）、hidden_states 和 attentions（若配置启用）、rope_deltas（保存用于下一 token 推理）、channel_loss_dict 和 channel_loss_count_dict（按数据集分类统计，当前实现中为 None）。1484 为空行，函数结束。
L1485–1501predict_action 方法是推理入口的简化包装器，处理 fast/diffusion 两种预测模式，调用内部 predict 方法后解包预测动作和真实标注。
1485
def predict_action(self, predict_mode: str, **kwargs):
Qwen2_5_VLMoEForAction 类的 predict_action 方法签名：接受 predict_mode 字符串和灵活的关键字参数。被推理脚本调用作为 VLA 模型推理的顶层入口。
1486
"""1487
Predict actions using specified prediction mode.1489
Args:1490
predict_mode (str): Prediction mode, either "fast" or "diffusion"1491
**kwargs: Additional arguments passed to the predict method1493
Returns:1494
tuple: (predicted_action, ground_truth_action) where ground_truth_action may be None1495
"""函数 docstring：说明 predict_action 是高级 API，支持 fast（离散动作）或 diffusion（flow matching 连续动作）两种预测模式，返回 (预测动作, 真实标注动作) 元组，真实标注可为 None。
1496
assert predict_mode in ["fast", "diffusion"]
断言校验：predict_mode 必须是 'fast' 或 'diffusion' 中的一个，否则抛出 AssertionError。这是防御性编程，确保调用者不会传入无效的预测模式。
调用核心推理方法 self.predict()，传递 predict_mode 和 **kwargs 中的所有超参数（含 input_ids、pixel_values、action_chunk、dof_mask、num_inference_timesteps 等）。predict 方法完整处理前向传播，返回字典 output，包含 'predict_action'、'gt_action'、'input_text'、'predict_output_text' 等多个键。
返回元组：(1) output['predict_action'] = 模型预测的动作序列，shape (B, A, D) 其中 A=pred_horizon、D=dof_dim(20)；(2) output.get('gt_action', None) = 真实标注动作或 None。LIBERO 实战中，predict_action 作为控制环闭包的输入，需通过 denormalizer 反转回原始动作空间。
函数体结束的空行，为下一个方法 @torch.no_grad() def predict() 分隔。
L1502–1631predict() 方法的核心推理入口：支持文本/快速/扩散三种模式，处理多模态输入（图像/视频/本体感觉），执行输入嵌入初始化和多模态 token 替换（masked_scatter）。关键实战细节：dof_mask 控制 LIBERO 场景的动作维度激活，num_inference_timesteps=10 是扩散去噪步数，dataset_names 用于 normalizer 查表
1502
@torch.no_grad()@torch.no_grad() 装饰器：禁用梯度计算，因为 predict() 是推理方法，不需要反向传播
1503
def predict(
1504
self,1505
predict_mode: str,1506
pred_horizon: Optional[int] = None,
1507
action_dim: Optional[int] = None,
1508
input_ids: torch.LongTensor = None,1509
attention_mask: Optional[torch.Tensor] = None,1510
position_ids: Optional[torch.LongTensor] = None,predict() 方法签名开始。predict_mode 参数指定推理模式（'text'/'fast'/'diffusion'），pred_horizon/action_dim 用于动作预测的地平线和维度。input_ids/attention_mask/position_ids 是标准语言模型输入参数
1511
past_key_values: Optional[List[torch.FloatTensor]] = None,1512
inputs_embeds: Optional[torch.FloatTensor] = None,1513
moe_token_types: Optional[torch.LongTensor] = None,1514
labels: Optional[torch.LongTensor] = None,1515
use_cache: Optional[bool] = None,
1516
output_attentions: Optional[bool] = None,
1517
output_hidden_states: Optional[bool] = None,
1518
return_dict: Optional[bool] = None,
1519
pixel_values: Optional[torch.Tensor] = None,1520
pixel_values_videos: Optional[torch.FloatTensor] = None,1521
image_grid_thw: Optional[torch.LongTensor] = None,多模态输入参数：past_key_values 用于 KV cache 加速推理；inputs_embeds 是预计算的输入嵌入；moe_token_types 用于 MoE 路由（按 token 类型）；pixel_values/pixel_values_videos/image_grid_thw/video_grid_thw 是图像/视频相关参数
1522
video_grid_thw: Optional[torch.LongTensor] = None,1523
action_chunk: Optional[torch.FloatTensor] = None,1524
proprioception: Optional[torch.FloatTensor] = None,1525
rope_deltas: Optional[torch.LongTensor] = None,1526
cache_position: Optional[torch.LongTensor] = None,1527
second_per_grid_ts: Optional[torch.Tensor] = None,1528
num_inference_timesteps: Optional[int] = 10,
1529
dataset_names: Optional[str] = None,
1530
dof_mask: Optional[torch.FloatTensor] = None,1531
agent_pos_mask: Optional[torch.FloatTensor] = None,1532
re_generate: bool = False,
1533
**kwargs,
1534
):
动作相关参数：action_chunk 是 ground truth 动作序列；proprioception 是本体感觉数据；rope_deltas 是 RoPE 位置增量；num_inference_timesteps 是扩散推理步数（默认 10）；dof_mask 掩码动作自由度（LIBERO 场景只用前 7 维）；agent_pos_mask 和 re_generate 用于本体感觉和采样控制
1535
"""1536
Multi-modal prediction method supporting text generation, fast action prediction, and diffusion-based action prediction.1538
This method handles three prediction modes:1539
1. "text": Pure text generation using autoregressive decoding1540
2. "fast": Fast action prediction using discrete action tokens1541
3. "diffusion": Continuous action prediction using diffusion/flow matching1543
Args:1544
predict_mode (str): Prediction mode ("text", "fast", or "diffusion")1545
pred_horizon (int, optional): Prediction horizon for action sequencesDocstring 主要说明：predict() 支持三种模式：text（纯文本）、fast（离散动作tokens）、diffusion（连续动作，使用 flow matching）。这是整个 VLA 推理的核心入口
1546
action_dim (int, optional): Dimensionality of action space1547
input_ids (torch.LongTensor, optional): Input token IDs1548
attention_mask (torch.Tensor, optional): Attention mask for input tokens1549
position_ids (torch.LongTensor, optional): Position IDs for tokens1550
past_key_values (List[torch.FloatTensor], optional): Cached key-value pairs1551
inputs_embeds (torch.FloatTensor, optional): Pre-computed input embeddings1552
moe_token_types (torch.LongTensor, optional): Token type assignments for MoE routing1553
labels (torch.LongTensor, optional): Target labels for evaluation1554
use_cache (bool, optional): Whether to use key-value caching1555
output_attentions (bool, optional): Whether to return attention weights1556
output_hidden_states (bool, optional): Whether to return hidden states1557
return_dict (bool, optional): Whether to return structured output1558
pixel_values (torch.Tensor, optional): Image pixel values1559
pixel_values_videos (torch.FloatTensor, optional): Video pixel values1560
image_grid_thw (torch.LongTensor, optional): Image grid dimensions1561
video_grid_thw (torch.LongTensor, optional): Video grid dimensions1562
action_chunk (torch.FloatTensor, optional): Ground truth action sequences1563
proprioception (torch.FloatTensor, optional): Proprioceptive sensor data1564
rope_deltas (torch.LongTensor, optional): RoPE position deltas1565
cache_position (torch.LongTensor, optional): Cache position indices1566
second_per_grid_ts (torch.Tensor, optional): Time interval per temporal grid1567
num_inference_timesteps (int, optional): Number of diffusion inference steps1568
dataset_names (str, optional): Dataset names for normalization1569
dof_mask (torch.FloatTensor, optional): Degrees of freedom mask1570
agent_pos_mask (torch.FloatTensor, optional): Agent position mask1571
re_generate (bool, optional): Whether to use sampling for regeneration1572
**kwargs: Additional keyword argumentsDocstring Args 部分：详细列举 30+ 个输入参数及其类型/默认值。关键实战细节：dataset_names 用于 normalizer 查表（如 LIBERO）；dof_mask 在 LIBERO 上是 [1]*7+[0]*13（20 维中只激活前 7 个）；num_inference_timesteps=10 是部署时扩散去噪步数
1574
Returns:1575
dict: Dictionary containing prediction results with keys like:1576
- 'predict_action': Predicted action sequences1577
- 'gt_action': Ground truth actions (if available)1578
- 'input_text': Input text (for text/fast modes)1579
- 'predict_output_text': Generated text (for text/fast modes)1580
- 'gt_output_text': Ground truth text (for text/fast modes)1581
"""Docstring Returns 部分：返回 dict，包含 'predict_action'（预测动作），'gt_action'（真值，可选），'input_text'/'predict_output_text'/'gt_output_text'（文本/快速模式的生成结果）。多模态返回支持不同预测模式的多种输出格式
1582
batch_size = (
1583
input_ids.shape[0] if input_ids is not None else inputs_embeds.shape[0]
1584
)
确定批大小：从 input_ids（如果存在）的第 0 维或 inputs_embeds 的第 0 维获取，典型值 B=1（文本/快速模式约束）
空行
1586
# Text and fast modes require batch size 1 for autoregressive generation1587
if predict_mode in ["text", "fast"]:
1588
assert (1589
batch_size == 11590
), "predict only support batch size 1 for ar generation"
文本/快速模式约束：assert batch_size==1，因为这两种模式使用自回归生成，需要逐 token 采样，不支持批处理
空行
1592
# Set output configuration from model config if not specified1593
output_attentions = (
1594
output_attentions
1595
if output_attentions is not None
1596
else self.config.output_attentions
1597
)
output_attentions 参数初始化：如果未显式指定，则从 self.config.output_attentions 读取。用于控制是否返回注意力权重
1598
output_hidden_states = (
1599
output_hidden_states
1600
if output_hidden_states is not None
1601
else self.config.output_hidden_states
1602
)
output_hidden_states 参数初始化：如果未显式指定，则从 self.config.output_hidden_states 读取。用于控制是否返回各层隐状态
1603
return_dict = (
1604
return_dict if return_dict is not None else self.config.use_return_dict
1605
)
return_dict 参数初始化：如果未显式指定，则从 self.config.use_return_dict 读取。控制是否返回结构化 output dict 还是元组
空行
1607
# Process input embeddings with multi-modal data注释行：'Process input embeddings with multi-modal data'，标记后续处理多模态输入嵌入的代码块
inputs_embeds 初始化：如果调用方未提供预计算的嵌入，则通过 self.model.embed_tokens(input_ids) 从 token IDs 获取初始嵌入。这是后续替换多模态嵌入（图像/视频/本体）的基础
1611
# Process image embeddings注释行：'Process image embeddings'，标记图像处理子块
1612
if pixel_values is not None:
1613
pixel_values = pixel_values.type(self.visual.dtype)1614
image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw)图像处理开始：if pixel_values is not None 检查是否有图像输入；pixel_values.type(self.visual.dtype) 类型转换到视觉编码器所需格式；self.visual() 调用视觉编码器编码图像
1615
n_image_tokens = (input_ids == self.config.image_token_id).sum().item()1616
n_image_features = image_embeds.shape[0]Token/Feature 计数：n_image_tokens=(input_ids==self.config.image_token_id).sum().item() 统计输入中图像占位符 token 的数量；n_image_features=image_embeds.shape[0] 获取视觉编码器输出的特征数量（通常一图多patch，形状 (N_patches, hidden_dim)）
空行
1618
# Validate image token and feature count match注释行：'Validate image token and feature count match'，标记验证逻辑
1619
if n_image_tokens != n_image_features:1620
raise ValueError(1621
f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}"
1622
)
验证图像 token 和特征数量匹配：如果 n_image_tokens != n_image_features，抛出 ValueError。这确保每个图像占位符都有对应的视觉特征，是多模态对齐的关键检查
空行
1624
mask = input_ids == self.config.image_token_id创建图像 token 掩码：mask=(input_ids==self.config.image_token_id)，布尔张量形状 (B,S)，标记每个位置是否是图像占位符
1625
mask_unsqueezed = mask.unsqueeze(-1)1626
mask_expanded = mask_unsqueezed.expand_as(inputs_embeds)
1627
image_mask = mask_expanded.to(inputs_embeds.device)
掩码形状扩展：mask_unsqueezed=mask.unsqueeze(-1) 扩展为 (B,S,1)；mask_expanded=mask_unsqueezed.expand_as(inputs_embeds) 扩展为 (B,S,H)，与 inputs_embeds 形状相同；image_mask=mask_expanded.to(inputs_embeds.device) 转移到正确设备
空行
图像嵌入类型转换和设备转移：image_embeds.to(inputs_embeds.device, inputs_embeds.dtype) 确保视觉特征与输入嵌入在同一设备和数据类型（如 float16/float32）。准备用于后续的 masked_scatter 替换操作
L1632–1761多模态输入处理与 position ID 计算：图像/视频嵌入注入、本体感觉投影、RoPE 位置编码、文本/快速模式的 autoregressive 生成前置处理
使用 masked_scatter 把已处理的 image_embeds (N_img, H) 按 image_mask 位置填充回 inputs_embeds (B, S, H)；(B,S,H) 中只有图像 token 位置被替换，其余文本 token 保持不变
1634
# Process video embeddings1635
if pixel_values_videos is not None:
1636
pixel_values_videos = pixel_values_videos.type(self.visual.dtype)1637
video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw)视频嵌入处理分支：同样的模式处理 pixel_values_videos 和 video_grid_thw（视频时空网格分辨率），通过 self.visual encoder 得到视频特征序列
1638
n_video_tokens = (input_ids == self.config.video_token_id).sum().item()1639
n_video_features = video_embeds.shape[0]统计视频 token 和特征数量以验证一致性：n_video_tokens 计数 input_ids 中 video_token_id 的个数，n_video_features 从 video_embeds 的第一维（序列长度）读取
1641
# Validate video token and feature count match1642
if n_video_tokens != n_video_features:1643
raise ValueError(1644
f"Video features and video tokens do not match: tokens: {n_video_tokens}, features {n_video_features}"
1645
)
视频 token 数与特征数量验证逻辑：若不匹配则抛异常，防止编码器输出的视频帧数与模板中预留的 token 位置不对齐；这在多尺度视频编码或数据预处理错误时容易触发
1647
mask = input_ids == self.config.video_token_id1648
mask_unsqueezed = mask.unsqueeze(-1)1649
mask_expanded = mask_unsqueezed.expand_as(inputs_embeds)
1650
video_mask = mask_expanded.to(inputs_embeds.device)
构建视频 mask 与 inputs_embeds 对齐：(1) mask=(input_ids==video_token_id) 形状 (B,S)；(2) unsqueeze(-1) 扩展为 (B,S,1)；(3) expand_as 广播到 (B,S,H) 以标记待替换位置
1652
video_embeds = video_embeds.to(
1653
inputs_embeds.device, inputs_embeds.dtype
1654
)
1655
inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds)
视频特征转移设备/dtype 后，用 masked_scatter 填充：video_embeds (N_vid,H) 按视频 mask 位置逐个填入 inputs_embeds，无缝替换预留的 video_token 位置
1657
# Process proprioceptive data1658
if proprioception is not None:
1659
proprioception = proprioception.to(inputs_embeds.device).to(
1660
inputs_embeds.dtype
1661
)
本体感觉数据处理分支：proprioception 是机器人关节角度/位置反馈，形状通常 (B,T,D_proprio)；先转移到 inputs_embeds 的设备和 dtype，为后续投影做准备
本体感觉掩码处理：agent_pos_mask 标记 LIBERO 任务中哪些 DOF 有效(形状 B,D)；同样转移设备/dtype，后续与噪声本体感觉拼接供网络学习(proj_with_mask=True 时生效)
1665
proprio_embed = self.action_preprocessor.proprioception_proj(1666
proprioception,
1667
dataset_names,
1668
agent_pos_mask,
1669
use_history=proprioception.shape[1] > 1,
1670
)
本体感觉嵌入投影：调用 action_preprocessor.proprioception_proj()，将 (B,1,D_proprio) 通过 MLP 投影到 (B,1,H_hidden)；use_history=True 若本体感觉有多步历史，否则只用单步；产出的 proprio_embed 用于替换 input_ids 中 propri_token_id 的位置
1671
proprioception_mask = (
1672
input_ids == self.action_token_id_set["propri_token_id"]
1673
)
1674
inputs_embeds[proprioception_mask] = proprio_embed.reshape(
1675
-1, inputs_embeds.shape[-1]
1676
).to(inputs_embeds.dtype)
本体感觉 token 掩码与注入：(1) proprioception_mask 定位 input_ids 中 propri_token_id 位置；(2) reshape(-1, H) 把 (B,1,H) 打平为 (B,H)；(3) inputs_embeds[mask]=proprio_embed_flat 原位替换对应 token 的嵌入；LIBERO 实战中这步在 PROPRI_DROPOUT 概率下可被关闭以测试失明情况
attention_mask 转移设备：防止跨设备计算时的设备不匹配错误；mask 形状通常 (B,S) 或 (B,1,S,S)，标记 padding/无效位置应被 mask 掉
1681
# Calculate RoPE position IDs if not provided1682
# Note: Cannot calculate rope deltas with 4D attention mask. TODO: Fix this limitation1683
if position_ids is None and (
1684
attention_mask is None or attention_mask.ndim == 2
1685
):
1686
# Calculate RoPE index once per generation in the pre-fill stage only1687
if (1688
(cache_position is not None and cache_position[0] == 0)
1689
or self.rope_deltas is None
1690
or (past_key_values is None or past_key_values.get_seq_length() == 0)
1691
):
RoPE position ID 缓存与计算逻辑：当 position_ids 未提供且 attention_mask 为 2D 时，检查是否需要重算 rope_deltas（仅在 prefill 阶段首次计算）；条件包括：cache 首位置、本类 rope_deltas 缓存为空、KV cache 首次初始化
1692
position_ids, rope_deltas = ops.get_rope_index(
1693
input_ids=input_ids,
1694
image_grid_thw=image_grid_thw,
1695
video_grid_thw=video_grid_thw,
1696
second_per_grid_ts=second_per_grid_ts,
1697
attention_mask=attention_mask,
1698
spatial_merge_size=self.config.vision_config.spatial_merge_size,1699
image_token_id=self.config.image_token_id,1700
video_token_id=self.config.video_token_id,1701
vision_start_token_id=self.config.vision_start_token_id,1702
tokens_per_second=self.config.vision_config.tokens_per_second,1703
)
调用 ops.get_rope_index() 计算 RoPE 编码的位置索引与增量向量：输入图像/视频网格分辨率(grid_thw)、时间分辨率(second_per_grid_ts)、spatial_merge_size；输出 position_ids (B,S,3) 和 rope_deltas (rope_len,) 用于后续自注意力；该步处理视图和文本 token 的混合位置编码
1704
self.rope_deltas = rope_deltas1705
# Use previously calculated rope deltas to get correct position IDs1706
else:1707
batch_size, seq_length, _ = inputs_embeds.shape
1708
delta = (
1709
(cache_position[0] + self.rope_deltas).to(inputs_embeds.device)
1710
if cache_position is not None
1711
else 0
1712
)
缓存 rope_deltas 并进入 else 分支（KV cache 已有内容时）：batch_size, seq_length 从 inputs_embeds 读取；delta = cache_position[0] + rope_deltas 计算当前窗口的绝对位置偏移；若无 cache_position 则 delta=0 保持相对位置编码
1713
position_ids = torch.arange(seq_length, device=inputs_embeds.device)
1714
position_ids = position_ids.view(1, -1).expand(batch_size, -1)
1715
if cache_position is not None: # otherwise `deltas` is an int `0`
1716
delta = delta.repeat_interleave(batch_size // delta.shape[0], dim=0)
1717
position_ids = position_ids.add(delta)
1718
position_ids = position_ids.unsqueeze(0).expand(3, -1, -1)
使用缓存的 rope_deltas 重建 position_ids：(1) torch.arange(seq_length) 生成相对位置 (S,)；(2) view(1,-1).expand(batch_size,-1) 广播为 (B,S)；(3) repeat_interleave 重复 delta 以对齐 batch；(4) add(delta) 加上绝对偏移；(5) unsqueeze(0).expand(3,-1,-1) 扩展为 (3,B,S) 以供三组注意力头使用
1720
# Prepare action chunk data if provided1721
if action_chunk is not None:
1722
action_chunk = action_chunk.to(inputs_embeds.device).to(inputs_embeds.dtype)
action_chunk（动作序列）类型转换：形状 (B, A_horizon, D_action) 表示未来 action_horizon 步的 20 维连续动作；转移到同一设备和精度，为后续流匹配损失计算做准备
1726
# Split input sequence for text and fast modes (not needed for diffusion)1727
if predict_mode == "text" or predict_mode == "fast":
1728
# Look for generation prompt tokens: <|im_start|>assistant1729
generation_prompt_ids = torch.tensor(
1730
[151644, 77091], device=input_ids.device, dtype=input_ids.dtype
1731
)
1732
matches = (input_ids[0, :-1] == generation_prompt_ids[0]) & (
1733
input_ids[0, 1:] == generation_prompt_ids[1]
1734
)
文本和快速模式的输入序列分割：定位生成起点标记 <|im_start|>assistant (token_id=[151644, 77091])；matches 标记两个 token 连续出现的位置，用于后续找到生成提示与 GT 输出的分界点
1736
if matches.any():1737
split_pos = torch.nonzero(matches, as_tuple=True)[0][0].item()
1738
# Extract ground truth output tokens (including newline)1739
gt_output_ids = input_ids[:, split_pos + 3 :]1740
# Remove output part from input, keeping prompt1741
input_ids = input_ids[:, : split_pos + 3]1742
inputs_embeds = inputs_embeds[:, : split_pos + 3, :]1743
if attention_mask is not None:
1744
attention_mask = attention_mask[:, : split_pos + 3]1745
if labels is not None:
1746
labels = labels[:, split_pos + 3 :]序列分割逻辑：若找到生成起点 (matches.any())，nonzero 获取首个匹配位置，split_pos+3 跳过标记后提取 gt_output_ids；input_ids/inputs_embeds 截断至 split_pos+3，保留提示部分；attention_mask 和 labels 同步截断
1747
else:1748
raise Warning(1749
"input_ids does not contain the generation prompt tokens <|im_start|>assistant"
1750
)
缺失生成起点标记时抛警告：正常情况下不应触发，表示输入格式不符合预期（多轮对话中可能丢失 assistant 提示）
1752
# Decode input text for output1753
input_text = self.processor.batch_decode(1754
input_ids, skip_special_tokens=False, clean_up_tokenization_spaces=True
1755
)
1756
output["input_text"] = input_text
对截断后的 input_ids（纯提示部分）调用 tokenizer.batch_decode() 还原为文本字符串；output["input_text"] 存储解码结果，用于后续生成日志或对比
1758
# Handle text and fast prediction modes using autoregressive generation1759
if predict_mode == "text" or predict_mode == "fast":
1760
# Initialize MoE token types for generation1761
moe_token_types = torch.zeros_like(input_ids)
为 autoregressive 生成初始化 MoE 路由信息：moe_token_types = torch.zeros_like(input_ids) 表示所有 token 默认属于图文专家(type=0)；后续在 forward/generate 中按 token 类型动态更新以指导双专家混合
L1762–1957三种预测模式的统一出口：文本生成模式 text 进行自回归解码并保存 ground truth 与预测文本；快速模式 fast 离散 token→连续动作→反归一化；扩散模式 diffusion 定义 step 函数接收扩散时间步与噪声动作，迭代调用 transformer 并更新隐状态，最后用 odeint(method=euler) 从高斯噪声一路去噪到清晰动作。所有模式都支持可选的 ground truth 动作以供评估。
1762
batch = {1763
"input_ids": input_ids,
1764
"attention_mask": attention_mask,
1765
"pixel_values": pixel_values,
1766
"moe_token_types": moe_token_types,
1767
"image_grid_thw": image_grid_thw,
1768
"dof_mask": dof_mask,
1769
"agent_pos_mask": agent_pos_mask,
1770
"proprioception": proprioception,
1771
"dataset_names": dataset_names,
1772
}
batch = {...} 构造输入数据字典，包含 input_ids, attention_mask, pixel_values, moe_token_types(全为 0 表示文本/图像 token), image_grid_thw, dof_mask, agent_pos_mask, proprioception, dataset_names；这个字典将被展开传递给 generate() 方法进行自回归文本生成。
空行
1774
# Generate output tokens1775
predict_output_ids = self.generate(1776
**batch,
1777
max_new_tokens=100,1778
eos_token_id=[self.processor.tokenizer.eos_token_id],1779
use_cache=True,1780
pad_token_id=self.processor.tokenizer.pad_token_id,1781
temperature=(
1782
1.0 if not re_generate else 0.7
1783
), # Higher temperature for regeneration1784
do_sample=(
1785
False if not re_generate else True
1786
), # Enable sampling for regeneration1787
)
调用 self.generate() 执行自回归文本生成；max_new_tokens=100 限制输出长度；use_cache=True 使用 KV 缓存加速；temperature 根据 re_generate 标志设置（False 时为 1.0，True 时为 0.7），do_sample 也相应调整（False 时贪心解码，True 时启用采样），实现重新生成时的温度降低和采样启用；返回 predict_output_ids 是生成的 token ID 序列 shape (B, new_len)。
空行
1789
# Decode generated and ground truth text1790
gt_output_text = self.processor.batch_decode(1791
gt_output_ids,
1792
skip_special_tokens=False,1793
clean_up_tokenization_spaces=True,1794
)
1795
predict_output_text = self.processor.batch_decode(1796
predict_output_ids,
1797
skip_special_tokens=False,1798
clean_up_tokenization_spaces=True,1799
)
1800
output["gt_output_text"] = gt_output_text
1801
output["predict_output_text"] = predict_output_text
分别解码 ground truth（gt_output_ids）和预测输出（predict_output_ids）为文本字符串；batch_decode() 处理一批 token 序列，skip_special_tokens=False 保留特殊 token（如 <|im_start|>），clean_up_tokenization_spaces=True 清理分词空格；结果存入 output 字典作为 gt_output_text 和 predict_output_text，供后续评估或展示。
空行
1803
# Convert tokens to actions for fast prediction mode1804
if predict_mode == "fast":
1805
action_id = []
1806
# Extract action tokens from generated sequence1807
for token_id_i in predict_output_ids[0]:
1808
if (进入 fast 预测模式分支；检查 predict_mode == "fast"；初始化 action_id 列表，用于收集从生成序列中提取的动作 token 的相对索引；遍历第一个样本的预测 token 序列 predict_output_ids[0]。
1809
token_id_i.item()
1810
>= self.processor.tokenizer.init_kwargs["action_token_start_index"]
1811
):
1812
action_id.append(
1813
token_id_i.item()
1814
- self.processor.tokenizer.init_kwargs[1815
"action_token_start_index"
1816
]
1817
)
if 条件判断：token_id_i.item() >= action_token_start_index 表示该 token 是动作 token（ID 在动作 token 范围内），则从原始 ID 减去 action_token_start_index 得到该动作 token 的相对索引，累加到 action_id 列表；这样可以从混合的文本+动作 token 序列中分离并重建动作索引序列。
1819
predict_action = self.processor.action_processor.decode(1820
[action_id], time_horizon=pred_horizon, action_dim=action_dim
1821
)
调用 action_processor.decode() 将离散动作 token 索引序列转换为连续动作张量；输入 [action_id] 是列表的列表（batch_size=1），time_horizon=pred_horizon 和 action_dim 指定动作的时间维度和动作维度；返回 predict_action 是 numpy 数组 shape (1, pred_horizon, action_dim)。
1822
# Handle action decoding errors1823
if np.sum(predict_action) == 0:
1824
print("Error in decoding action, predict_action is None")
1825
output["predict_action"] = None
检查解码后的 predict_action 是否有效（通过检查所有元素之和是否非零）；如果解码失败或返回全 0 数组（意味着没有有效的动作 token 被解析出来），打印错误信息并将 output['predict_action'] 设为 None；否则进入下一步处理。
1826
else:1827
# Convert discrete tokens to continuous actions1828
predict_action = torch.tensor(predict_action, device=self.device)1829
dof_mask = dof_mask.to(self.device).to(pixel_values.dtype)1830
predict_action = (
1831
self.action_preprocessor.normalizer_action.unnormalize_data(1832
predict_action, dataset_names, dof_mask
1833
)
1834
)
1835
output["predict_action"] = predict_action
将 predict_action 转换为 torch.Tensor 并移到模型设备上；应用 dof_mask 的类型转换（to self.device 和 pixel_values.dtype）；调用 normalizer_action.unnormalize_data() 反归一化（从标准化的 [-1,1] 空间恢复到原始值域），根据 dataset_names 查表获取该数据集的 min/max 统计参数，进行逆线性变换 actual = normalized * (max-min) + min；结果存入 output['predict_action']。在 LIBERO 实战中，dof_mask 有 13 个非零位（前 7 维 DOF + 3 维 gripper + 2 维 head + 1 维 height），反归一化后的动作可直接用于控制机器人。
空行
1837
# Process ground truth actions if available1838
if action_chunk is not None:
1839
# Apply DOF mask and unnormalize action chunk to get ground truth actions处理 ground truth 动作：检查 action_chunk 是否不为 None（即是否有标注的动作标签）；如有，取该 batch 中第一个样本的动作（[0, :, :]），并仅保留 dof_mask==True 的维度，shape 从 (B,A,D) 变为 (B,A,dof_kept)。
1840
action_chunk = action_chunk[:, :, dof_mask[0, 0, :].bool()]
1841
output["gt_action"] = (
1842
self.action_preprocessor.normalizer_action.unnormalize_data(1843
action_chunk, dataset_names, dof_mask
1844
)
1845
)
调用 normalizer_action.unnormalize_data() 对 dof_mask 过的 action_chunk 进行反归一化，恢复原始值域；使用 dataset_names 查表获得该数据集的统计参数；注意这里传入的 dof_mask 仍然是完整的 20 维 mask（反归一化器内部会正确处理）；结果存入 output['gt_action']，用于与 predict_action 对比评估。在 LIBERO 中，gt_action 是从人类示范 trajectory 中提取的标准答案。
当 action_chunk 为 None 时（无 ground truth 标签），output['gt_action'] 显式设为 None。
空行
1849
# Handle diffusion-based action prediction进入扩散预测模式分支；检查 predict_mode == "diffusion"；后续执行基于 flow matching 的去噪采样过程。
1850
if predict_mode == "diffusion":
1851
# Initialize with random noise1852
noisy_action = torch.randn(
1853
size=(batch_size, pred_horizon, action_dim),
1854
dtype=inputs_embeds.dtype,
1855
device=inputs_embeds.device,
1856
)
1857
dof_mask = dof_mask.to(inputs_embeds.device).to(inputs_embeds.dtype)
初始化 noisy_action：从标准高斯分布采样，shape (batch_size, pred_horizon, action_dim)；dtype 和 device 与 inputs_embeds 保持一致，确保与后续计算的兼容性；dof_mask 转换到 inputs_embeds 的设备和 dtype，用于掩蔽不需要的动作维度。
空行
1859
# Calculate token distribution across MoE expert groups1860
group_size = torch.zeros(
1861
self.config.num_experts, dtype=torch.long, device="cpu"
1862
)
1863
for i in range(self.config.num_experts):
1864
group_size[i] = (moe_token_types == i).sum()
初始化 group_size 张量统计每个 MoE expert 的 token 数量；遍历 num_experts（通常为 2），计算有多少个 token 被路由到各个专家；group_size[i] = (moe_token_types == i).sum()；在本模型中，expert 0 处理文本/图像 token，expert 1 处理动作 token。在 LIBERO 实战中，通常大多数 token 被路由到 expert 0，只有少数动作 token 被路由到 expert 1。
1866
# Calculate start and end indices for each expert group1867
start_indices = torch.cumsum(group_size, dim=0) - group_size1868
end_indices = torch.cumsum(group_size, dim=0)计算各 expert group 在 token 序列中的起止索引；start_indices[i] = cumsum(group_size)[:i] = cumsum[:i] - group_size[i]；end_indices[i] = cumsum[:i+1]；这两个索引数组将被传入 transformer 的前向传播中，用于定位各 expert 的输入 token 范围，使 MoE 路由能找到正确的 token 块。
空行
1870
def step(timestep, noisy_action):
1871
"""1872
Single denoising step for diffusion process.1874
Args:1875
timestep: Current diffusion timestep1876
noisy_action: Current noisy action estimate1878
Returns:1879
torch.Tensor: Predicted clean action1880
"""定义 step(timestep, noisy_action) 内部函数，被 odeint() 调用执行单次去噪迭代；该函数签名 (t, y) -> dydt 符合 ODE 求解器接口；timestep 是当前扩散步的时间参数（0-1 区间，0 对应初始噪声，1 对应清晰动作），noisy_action 是当前的噪声动作张量 shape (B,A,D)；返回值是预测的清晰动作 shape (B,A,D)，作为 ODE 积分的梯度（velocity）。
1881
action_mask = input_ids == self.action_token_id_set["action_token_id"]
1882
assert action_mask.any(), "No action token found in input_ids"
在 step 函数内部定位所有动作 token 在 input_ids 中的位置（值等于 self.action_token_id_set['action_token_id']）；action_mask shape (1, seq_len) 为布尔张量；assert 确保至少存在一个动作 token，否则无法进行动作预测，抛出 AssertionError。
1884
# Prepare timestep for batch processing1885
timestep = timestep.unsqueeze(0).repeat(noisy_action.shape[0])
1886
action_embed, _ = self.action_preprocessor.step(1887
timestep=timestep, noisy_action=noisy_action, dof_mask=dof_mask
1888
)
准备时间步嵌入：timestep 是 scalar，用 unsqueeze(0) 变为 shape (1,)，再 repeat(noisy_action.shape[0]) 扩展为 (B,)；调用 action_preprocessor.step() 进行时间步的正弦位置编码和噪声动作的融合，输出 action_embed shape (B,A,H) 其中 H=inputs_embeds.shape[-1]；reshape 为 (-1, H) 即 (B*A, H) 以匹配后续 token 替换的需求。
1889
action_embed = action_embed.reshape(-1, inputs_embeds.shape[-1])
1891
# Create temporary copy of embeddings for thread safety1892
temp_inputs_embeds = inputs_embeds.clone()
1893
temp_inputs_embeds[action_mask] = action_embed.to(
1894
temp_inputs_embeds.dtype
1895
)
创建 inputs_embeds 的临时副本 temp_inputs_embeds，进行 clone() 深拷贝，避免直接修改原始嵌入（线程安全，防止并发访问冲突）；将动作位置的嵌入 temp_inputs_embeds[action_mask] 替换为当前步骤的 action_embed，并类型转换到 temp_inputs_embeds.dtype；这样就把当前去噪步的动作编码融入到输入序列中，准备输入 transformer。
空行
1897
# Forward pass through transformer1898
transformer_outputs = self.model(1899
input_ids=None,1900
attention_mask=attention_mask,
1901
position_ids=position_ids,
1902
past_key_values=past_key_values,
1903
inputs_embeds=temp_inputs_embeds,
1904
moe_token_types=moe_token_types,
1905
start_indices=start_indices,
1906
end_indices=end_indices,
1907
use_cache=True,1908
output_attentions=False,1909
output_hidden_states=False,1910
return_dict=True,1911
)
执行 transformer 前向传播：input_ids=None 表示直接使用 inputs_embeds（避免重复嵌入）；attention_mask 和 position_ids 与原始输入相同；use_cache=True 会保存和更新 KV cache 供下一步使用；moe_token_types, start_indices, end_indices 用于 MoE 路由（expert 0 处理 type==0 的文本/图像 token，expert 1 处理 type==1 的动作 token，两类 token 共享同一注意力层）；返回 dict 包含 last_hidden_state 等信息。在扩散迭代中，transformer 多次被调用，但文本 KV cache 只在第一步计算，后续可复用（虽然这里每次都重新计算了）。
1913
# Extract action predictions from hidden states1914
hidden_states = transformer_outputs.last_hidden_state
1915
action_mask = input_ids == self.action_token_id_set["action_token_id"]
1916
action_hidden_states = hidden_states[action_mask]
提取动作位置的隐状态：hidden_states = transformer_outputs.last_hidden_state shape (B,S,H)；用 action_mask 进行布尔索引选取动作 token 对应的隐状态 action_hidden_states shape (B*A, H)；只截取前 action_hidden_size 维的隐状态（丢弃超出部分，如果 H > action_hidden_size），为投影器准备输入。
1917
pred = self.action_preprocessor.action_proj_back(1918
action_hidden_states[
1919
:, : self.action_preprocessor.action_hidden_size1920
]
1921
)
1922
return pred.reshape(batch_size, pred_horizon, action_dim)调用 action_preprocessor.action_proj_back() 将隐状态投影回动作空间；输入 shape (B*A, action_hidden_size)，输出 shape (B*A, D=20)；最后 reshape 回 (B, A, D) 的标准形状，得到当前去噪步的预测速度场（flow）；该 flow 将被 ODE 求解器用于积分更新 noisy_action。
空行
1924
# Perform ODE integration for diffusion sampling1925
times = torch.linspace(
1926
0,1927
1,1928
num_inference_timesteps + 1,1929
device=inputs_embeds.device,
1930
dtype=inputs_embeds.dtype,
1931
)
生成去噪时间轴：torch.linspace(0, 1, num_inference_timesteps+1) 从 0 均匀采样到 1，共 num_inference_timesteps+1 个点（e.g., 11 个点对应 10 步积分）；device 和 dtype 与 inputs_embeds 保持一致；times 张量定义了 ODE 求解器的积分路径和时间点。
1932
action_trajectory = odeint(
1933
step,
1934
noisy_action.to(torch.float32),
1935
times.to(torch.float32),
1936
method="euler",
1937
)
调用 odeint() 执行 ODE 数值积分（method="euler" 一阶 Euler 方法）；初值 noisy_action.to(torch.float32) 是纯高斯噪声 shape (B,A,D)；积分函数 step 在每个时间点 t ∈ times 被调用，逐步去噪，每步更新 noisy_action += step(t, noisy_action) * dt；返回 action_trajectory shape (num_timesteps+1, B, A, D)，第 0 行是初始高斯噪声，最后一行是完全去噪的清晰动作。
空行
1939
# Extract final predicted action and unnormalize1940
predict_action = action_trajectory[-1]1941
predict_action = (
1942
self.action_preprocessor.normalizer_action.unnormalize_data(1943
predict_action, dataset_names
1944
)
1945
)
提取 ODE 积分的最终输出：action_trajectory[-1] 是时间轴最后一刻（t=1）的去噪结果，shape (B,A,D)；调用 normalizer_action.unnormalize_data() 反归一化（从 [-1,1] 标准空间恢复到原始值域），使用 dataset_names 查表获取该数据集的统计参数（min/max），进行逆线性变换 actual = normalized * (max-min) + min；在 LIBERO 实战中，每个数据集都有独立的归一化统计。
1946
output["predict_action"] = predict_action
将反归一化后的预测动作存入 output['predict_action']，shape (B,A,D)，这是最终的控制命令张量。
空行
1948
# Process ground truth actions if available1949
if action_chunk is not None:
1950
output["gt_action"] = (
1951
self.action_preprocessor.normalizer_action.unnormalize_data(1952
action_chunk, dataset_names
1953
)
1954
)
处理 ground truth 动作（扩散模式）：如果 action_chunk 不为 None（有标注的动作标签），则直接（无需 dof_mask 处理，因为 action_chunk 已在数据管道中处理完毕）调用 normalizer_action.unnormalize_data() 进行反归一化；返回值存入 output['gt_action']，用于与 predict_action 对比评估性能；注意 diffusion 模式不像 fast 模式那样需要预先提取 dof_mask 维度。
空行
返回 output 字典，包含在三种预测模式中累积的所有输出（input_text, gt_output_text, predict_output_text, predict_action, gt_action 等）；调用方可根据 predict_mode 查取相应的预测结果；字典为后续的评估、可视化、机器人控制等流程提供完整的信息。
L1958–2087generate_flow_action 推理入口(1958-1995) + 初始化计时与配置参数(1997-2015) + 图文/视频/本体状态嵌入融合管线(2017-2082) + RoPE 位置编码条件检查(2084-2087)
1958
@torch.no_grad()@torch.no_grad() 装饰器关闭梯度计算，因推理不需反向传播，节省显存
1959
def generate_flow_action(
1960
self,1961
input_ids,
1962
action_horizon,
1963
action_dim,
1964
num_inference_timesteps: int = 10,
1965
padding_action: Optional[torch.Tensor] = None,1966
prefix_length: Optional[int] = None,
1967
attention_mask: Optional[torch.Tensor] = None,1968
position_ids: Optional[torch.LongTensor] = None,1969
past_key_values: Optional[List[torch.FloatTensor]] = None,1970
inputs_embeds: Optional[torch.FloatTensor] = None,1971
moe_token_types: Optional[torch.LongTensor] = None,1972
start_indices: Optional[torch.Tensor] = None,1973
end_indices: Optional[torch.Tensor] = None,1974
positional_masks: Optional[torch.LongTensor] = None,1975
labels: Optional[torch.LongTensor] = None,1976
use_cache: Optional[bool] = None,
1977
output_attentions: Optional[bool] = None,
1978
output_hidden_states: Optional[bool] = None,
1979
return_dict: Optional[bool] = None,
1980
pixel_values: Optional[torch.Tensor] = None,1981
pixel_values_videos: Optional[torch.FloatTensor] = None,1982
image_grid_thw: Optional[torch.LongTensor] = None,1983
video_grid_thw: Optional[torch.LongTensor] = None,1984
action_chunk: Optional[torch.FloatTensor] = None,1985
proprioception: Optional[torch.FloatTensor] = None,1986
unnorm_proprioception: Optional[torch.FloatTensor] = None,1987
rope_deltas: Optional[torch.LongTensor] = None,1988
cache_position: Optional[torch.LongTensor] = None,1989
second_per_grid_ts: Optional[torch.Tensor] = None,1990
dataset_names: Optional[str] = None,
1991
dof_mask: Optional[torch.FloatTensor] = None,1992
agent_pos_mask: Optional[torch.FloatTensor] = None,1993
unnorm: Optional[bool] = True,
1994
**kwargs,
1995
):
generate_flow_action 函数签名：VLA 推理主入口，参数包括 input_ids、action_horizon/action_dim(动作空间)、num_inference_timesteps=10(部署推理 ODE 积分步数)、pixel_values/pixel_values_videos(视觉)、proprioception/agent_pos_mask(本体状态)、dataset_names/dof_mask(数据集和自由度掩码，LIBERO 只用前 7 维 pad 进 20 维)、padding_action(可选初始噪声)、position_ids/attention_mask/moe_token_types(位置编码和 MoE 路由) 等，由外侧 forward() 或 generate() 调用
空行
记录推理总耗时起点和 timing_results 字典，用于后续分段计时(embed/position/action_init)
空行
2000
batch_size = (
2001
input_ids.shape[0] if input_ids is not None else inputs_embeds.shape[0]
2002
)
提取 batch_size，优先从 input_ids 推导，若为 None 则从 inputs_embeds 推导；确保推理时总能确定批次大小
2003
output_attentions = (
2004
output_attentions
2005
if output_attentions is not None
2006
else self.config.output_attentions
2007
)
output_attentions 配置参数赋值：若函数参数传入非 None，用传入值；否则用 self.config 默认值，防止 None 漏透到后续模块
2008
output_hidden_states = (
2009
output_hidden_states
2010
if output_hidden_states is not None
2011
else self.config.output_hidden_states
2012
)
output_hidden_states 配置参数赋值：同上逻辑，确保推理输出格式一致
2013
return_dict = (
2014
return_dict if return_dict is not None else self.config.use_return_dict
2015
)
return_dict 配置参数赋值：同上逻辑，决定后续返回值是字典还是元组
空行
2017
embed_start_time = time.time()
记录嵌入处理模块起始时刻，计算 embed_processing 耗时
若外侧未预先计算 inputs_embeds，则通过 self.model.embed_tokens(input_ids) 将 token ID 映射到嵌入空间 (B,S,H) 其中 H=hidden_size 通常 2048
2020
if pixel_values is not None:
2021
pixel_values = pixel_values.type(self.visual.dtype)2022
image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw)2023
n_image_tokens = (input_ids == self.config.image_token_id).sum().item()2024
n_image_features = image_embeds.shape[0]2025
if n_image_tokens != n_image_features:2026
raise ValueError(2027
f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}"
2028
)
图像嵌入融合流程：若 pixel_values 非空(形状 (N_img,C,H,W))，先转为 visual 模块的 dtype，然后调 self.visual(pixel_values, grid_thw=image_grid_thw) 生成图像特征 (N_image_tokens, hidden_size)，检验数量匹配否则报错
空行
2030
mask = input_ids == self.config.image_token_id2031
mask_unsqueezed = mask.unsqueeze(-1)2032
mask_expanded = mask_unsqueezed.expand_as(inputs_embeds)
2033
image_mask = mask_expanded.to(inputs_embeds.device)
创建图像位置掩码：input_ids == image_token_id 得布尔掩码 (B,S)，unsqueeze(-1) 扩到 (B,S,1)，expand_as 扩展到输入嵌入形状 (B,S,H)，再转到正确设备；用于后续 masked_scatter
空行
图像特征设备和数据类型对齐至 inputs_embeds 所在设备和 dtype，防止异构计算
2038
inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds)
inputs_embeds.masked_scatter(image_mask, image_embeds)：按掩码把展平的 image_embeds 填回 inputs_embeds 对应位置，实现图像 token 的嵌入替换 (B,S,H) 中的对应位，LIBERO 实战中注意图像 token 数必须严格等于视觉编码器输出的特征数
空行
2040
if pixel_values_videos is not None:
2041
pixel_values_videos = pixel_values_videos.type(self.visual.dtype)2042
video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw)2043
n_video_tokens = (input_ids == self.config.video_token_id).sum().item()2044
n_video_features = video_embeds.shape[0]2045
if n_video_tokens != n_video_features:2046
raise ValueError(2047
f"Video features and video tokens do not match: tokens: {n_video_tokens}, features {n_video_features}"
2048
)
视频嵌入融合流程：同图像逻辑(2020-2028)，若 pixel_values_videos 非空，生成视频特征并验证数量匹配
空行
2050
mask = input_ids == self.config.video_token_id2051
mask_unsqueezed = mask.unsqueeze(-1)2052
mask_expanded = mask_unsqueezed.expand_as(inputs_embeds)
2053
video_mask = mask_expanded.to(inputs_embeds.device)
创建视频位置掩码：同图像掩码逻辑，标记 input_ids 中 video_token_id 对应位置
空行
视频特征设备和 dtype 对齐
2058
inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds)
inputs_embeds.masked_scatter(video_mask, video_embeds)：按掩码融合视频特征
空行
2060
if (2061
proprioception is not None
2062
and not self.config.use_state_string_representation
2063
):
条件检查：若 proprioception 非空且 config 未启用 use_state_string_representation(即本体状态用向量不用字符串表示)，则进行本体状态嵌入处理，这是多体状态融合的关键分支
2064
proprioception = proprioception.to(inputs_embeds.device)
2065
agent_pos_mask = agent_pos_mask.to(inputs_embeds.device)
proprioception (B,T,num_proprio) 和 agent_pos_mask (B,T,num_agents) 转到与 inputs_embeds 同一设备，通常为 GPU；LIBERO 实战中，proprioception 包含关节角/速度，agent_pos_mask 标记有效智能体
2066
proprio_embed = self.action_preprocessor.proprioception_proj(2067
proprioception,
2068
dataset_names,
2069
agent_pos_mask,
2070
use_history=proprioception.shape[1] > 1,
2071
)
调 self.action_preprocessor.proprioception_proj 投影本体状态向量：输入 proprioception (B,T,num_proprio_dim) 和 agent_pos_mask (B,T,num_agents)，dataset_names 用于查多数据集 normalizer，use_history=proprioception.shape[1]>1 表示是否传入历史帧；返回投影后嵌入 (B*T, H) 或类似形状，可选项：若为 None 表示所有智能体都无效
创建本体状态位置掩码：在 input_ids 中查找 propri_token_id(本体状态特殊 token)位置，布尔掩码形状同 input_ids (B,S)
2075
inputs_embeds[proprioception_mask] = proprio_embed.reshape(
2076
-1, inputs_embeds.shape[-1]
2077
).to(inputs_embeds.dtype)
inputs_embeds[proprioception_mask] = ... 按掩码把本体状态嵌入填回输入嵌入对应位置；reshape(-1, inputs_embeds.shape[-1]) 确保形状匹配 (被选中的位数, H)，再 .to(inputs_embeds.dtype) 确保数据类型一致，防止 float32 和 float16 混用；LIBERO 闭环中若本体状态被随机隐藏(PROPRI_DROPOUT)则此处嵌入为全 0
空行
若 attention_mask 非空(通常 (B,S) 或 (B,1,S,S))，转到与 inputs_embeds 同一设备，避免设备不匹配
空行
2082
timing_results["embed_processing"] = time.time() - embed_start_time
记录嵌入处理耗时到 timing_results['embed_processing']，用于性能分析和瓶颈定位
空行
2084
position_start_time = time.time()
2085
# if we get 4D attention mask we cannot calculate rope deltas anymore. TODO @raushan fixme2086
if position_ids is None and (
2087
attention_mask is None or attention_mask.ndim == 2
记录位置编码起始时刻；注释指出若 attention_mask 为 4D(如 (B,1,S,S))不能计算 RoPE 增量(TODO 待修复)；条件 position_ids is None and (attention_mask is None or attention_mask.ndim==2) 判断是否需计算新的 RoPE 位置编码，仅在 attention_mask 为 None 或 2D 时才能进行增量计算
L2088–2217RoPE位置编码计算与初始化、MoE专家索引计算、Flow Matching去噪过程初始化（生成噪声、嵌入初始动作）、预填充前向传播获得KV缓存、动作速度场计算与Euler积分更新、前缀长度计算用于KV缓存分割
2088
):为前置if语句闭括号。2089为注释，说明RoPE位置索引仅在预填充阶段计算一次，后续自回归阶段复用缓存以提高效率，这是KV缓存优化的关键2090
if (2091
(cache_position is not None and cache_position[0] == 0)
2092
or self.rope_deltas is None
2093
or (past_key_values is None or past_key_values.get_seq_length() == 0)
2094
):
if条件判断（5行）：检查三个触发重新计算RoPE的条件：(1) cache_position[0]==0即预填充阶段开始 (2) rope_deltas未被计算 (3) KV缓存为空。满足任一条件都重新计算，保证位置编码的正确性
2095
position_ids, rope_deltas = self.get_rope_index(2096
input_ids,
2097
image_grid_thw,
2098
video_grid_thw,
2099
second_per_grid_ts,
2100
attention_mask,
2101
)
调用self.get_rope_index()获取position_ids和rope_deltas（7行）。position_ids用于RoPE旋转位置编码，rope_deltas存储位置偏移供后续KV缓存阶段使用。参数包括input_ids、图像/视频网格、注意力掩码
2102
self.rope_deltas = rope_deltasself.rope_deltas=rope_deltas缓存位置增量到实例变量，供自回归生成阶段复用，避免重复计算
2103
# then use the prev pre-calculated rope-deltas to get the correct position ids注释行说明：进入else分支表示已计算过rope_deltas，现在复用缓存的rope_deltas来生成position_ids
else分支处理KV缓存阶段位置编码。2105从inputs_embeds形状解包出batch_size和seq_length，seq_length为当前生成步的序列长度
2106
delta = (
2107
(cache_position[0] + self.rope_deltas).to(inputs_embeds.device)
2108
if cache_position is not None
2109
else 0
2110
)
计算delta：位置偏移量（5行三元表达式）。若有cache_position，则delta=cache_position[0]+缓存的rope_deltas（移到当前device）；否则delta=0。这个偏移确保KV缓存阶段的位置编码与预填充阶段连续
2111
position_ids = torch.arange(seq_length, device=inputs_embeds.device)
2112
position_ids = position_ids.view(1, -1).expand(batch_size, -1)
初始化position_ids基础值：(1) 2111创建0到seq_length-1的整数序列 (2) 2112展开为(1,seq_length)再扩展成(batch_size,seq_length)，shape变化：(S,)→(1,S)→(B,S)
2113
if cache_position is not None: # otherwise `deltas` is an int `0`
2114
delta = delta.repeat_interleave(batch_size // delta.shape[0], dim=0)
2115
position_ids = position_ids.add(delta)
2116
position_ids = position_ids.unsqueeze(0).expand(3, -1, -1)
若有cache_position则应用位置偏移。2114重复delta使其与batch_size匹配，2115加法应用偏移，2116扩展position_ids为(3,B,S)形状（可能对应多头或多表示法），完整shape变化：(B,S)→(3,B,S)
空行，分隔位置编码逻辑与MoE索引计算逻辑
2118
if start_indices is None or end_indices is None:
if条件：检查start_indices和end_indices是否都为None，若为None则需根据moe_token_types计算各专家的起始/结束位置。这些索引用于TokenTypeRouter在MoE前向时分割和路由token
2119
# Calculate the start and end positions of each expert group's tokens after permutation (the dataset does not contain `num_expert` information, so this calculation must be done here).2120
group_size = torch.zeros(
2121
self.config.num_experts, dtype=torch.long, device="cpu"
2122
)
初始化group_size张量：shape(num_experts,)，所有值初始为0。为每个专家统计其负责的token数，这是计算MoE路由索引的第一步
遍历每个专家，统计其token数：group_size[i]=(moe_token_types==i).sum()。moe_token_types中0表示图文token，1表示动作/状态token，分别路由到两个专家
空行，分隔统计逻辑与索引计算逻辑
2126
# Calculate start and end indices for each expert group注释说明：计算累积和以得到起始和结束索引，用于后续前向传播时高效地对token进行排列和分组处理
2127
start_indices = torch.cumsum(group_size, dim=0) - group_size2128
end_indices = torch.cumsum(group_size, dim=0)cumsum计算起始和结束索引（2行）。start_indices=cumsum(group_size)-group_size为各专家的起始位置，end_indices=cumsum(group_size)为结束位置，支持高效的token排列操作，避免python循环
空行，分隔MoE索引计算与动作初始化前的时间记录
2130
timing_results["position_encoding"] = time.time() - position_start_time
timing_results['position_encoding']=time.time()-position_start_time记录位置编码阶段耗时，用于性能分析和瓶颈定位
空行
2132
action_init_start_time = time.time()
2133
if action_chunk is not None:
2134
action_chunk = action_chunk.to(inputs_embeds.device).to(torch.float32)
2132记录动作初始化开始时间。2133-2134若action_chunk非None，转移到inputs_embeds设备并转为float32，为后续处理做准备
空行
2136
output = {}初始化output字典，用于存储该函数返回的所有结果（如隐藏状态、损失等）
2137'# Reproduce'注释说明这段噪声生成用于可复现性。2138被注释掉的torch.manual_seed(0)显示曾用于调试确保随机性可控
2139
noise = torch.randn(
2140
size=(batch_size, action_horizon, action_dim),
2141
dtype=torch.float32,
2142
device=inputs_embeds.device,
2143
)
生成标准高斯噪声张量，shape(B,action_horizon,action_dim)。在Flow Matching中，噪声作为ODE初始条件，将通过去噪网络逐步积分演化为动作预测，tensor shape：(B,A,D)其中A=action_horizon(如10),D=action_dim(20)
2144
noisy_action = noise.clone()
clone()创建噪声的独立副本作为初始noisy_action，该变量在Euler积分循环中迭代更新，在LIBERO实战中A通常为10
2145
dof_mask = dof_mask.to(inputs_embeds.device).to(torch.float32)
将dof_mask转移到inputs_embeds设备并转为float32。dof_mask用于mask掉LIBERO任务中不使用的自由度（后13维），实战中只有前7维active
空行
2147
if num_inference_timesteps not in self.times_cache:
if条件：检查num_inference_timesteps是否已在self.times_cache中，times_cache缓存不同推理步数下的时间步序列，避免重复计算
2148
self.times_cache[num_inference_timesteps] = torch.linspace(2149
0.0,2150
1.0,2151
num_inference_timesteps + 1,2152
device=inputs_embeds.device,
2153
dtype=torch.float32,
2154
)
若cache未命中，创建时间步缓存（7行）。torch.linspace(0,1,num_inference_timesteps+1)生成从0到1均匀分布的时间步，shape(num_inference_timesteps+1,)。在Flow Matching中，时间步t∈[0,1]，0表示纯噪声，1表示纯动作
2155
times = self.times_cache[num_inference_timesteps]从times_cache获取当前配置的时间步张量，shape(num_inference_timesteps+1,)
2156
dt = times[1] - times[0]
计算时间步间隔dt=times[1]-times[0]，用于后续Euler积分的步长。部署时通常为1/num_inference_timesteps（如1/5=0.2）
2157
time_0 = times[0].unsqueeze(0).repeat(noisy_action.shape[0])
获取初始时间步time_0=times[0]（通常为0），unsqueeze(0)添加batch维，repeat()扩展为shape(B,)，broadcast到batch中每个样本
2158
action_embed, adarms_cond = self.action_preprocessor.step(2159
timestep=time_0, noisy_action=noisy_action, dof_mask=dof_mask
2160
)
调用self.action_preprocessor.step()生成初始动作嵌入（3行）。step()返回(action_embed,adarms_cond)：前者是初始动作对应的隐藏表示，后者是时间和dof_mask的条件编码，用于融合到后续前向中
2161
action_embed = action_embed.reshape(-1, inputs_embeds.shape[-1]).to(
2162
inputs_embeds.dtype
2163
)
重塑action_embed为(B*A,H)维度并转为inputs_embeds的dtype（3行）。H=2048为模型隐藏维度，该过程将动作嵌入转换为可直接代入LLM的格式，shape变化：(B,A,*)→(B*A,H)
2164
flow_action_mask = input_ids == self.action_token_id_set["action_token_id"]
生成flow_action_mask，标记input_ids中所有action_token的位置（通常是特殊token），shape(B,S)其中S为序列长度
空行
2166
inputs_embeds[flow_action_mask] = action_embed
masked_scatter()将action_embed按flow_action_mask代入inputs_embeds对应位置，inputs_embeds在forward过程中会逐步叠加不同token的嵌入（如图像token、动作token）
空行
2168
timing_results["action_initialization"] = time.time() - action_init_start_time
timing_results['action_initialization']=time.time()-action_init_start_time记录动作初始化耗时，包括噪声生成、step计算、嵌入代入等步骤
空行
2170
prefetch_start_time = time.time()
记录prefetch_start_time，prefetch是指预填充（pre-fill）阶段的前向传播，该阶段一次性处理整个图文前缀和初始动作
2171
prefetch_output = self.model(2172
input_ids=None,2173
attention_mask=attention_mask,
2174
position_ids=position_ids,
2175
past_key_values=None,2176
inputs_embeds=inputs_embeds,
调用self.model()执行预填充前向传播（6行）。input_ids=None因为已转为inputs_embeds，position_ids和inputs_embeds包含图文+初始动作，use_cache=True保存KV缓存供后续自回归阶段使用，past_key_values=None表示首次无缓存
2177
moe_token_types=moe_token_types,
2178
start_indices=start_indices,
2179
end_indices=end_indices,
2180
positional_masks=positional_masks,
传递MoE路由所需参数（4行）：moe_token_types标记token类型，start/end_indices指定各专家token边界，positional_masks用于attention计算，adarms_conds[None,adarms_cond]是时间条件（第0个专家用None，第1个用adarms_cond）
2181
use_cache=True,2182
output_attentions=False,2183
output_hidden_states=False,2184
return_dict=True,2185
adarms_conds=[None, adarms_cond],2186
)
模型配置参数（6行）：use_cache=True保存KV缓存，output_attentions/output_hidden_states=False不保存中间态以省显存，return_dict=True返回字典格式便于访问属性
2187
hidden_states = prefetch_output.last_hidden_state
提取prefetch_output的last_hidden_state，shape(B,S,H)，包含图文前缀和初始动作在该处的LLM隐藏表示，后续用于提取动作token的隐状态进行反投影
2188
prefix_kv_cache = prefetch_output.past_key_values
提取past_key_values（KV缓存），包含所有LLM层的key和value张量，在自回归生成阶段复用该缓存以加速计算（前缀部分不重新算attention）
按flow_action_mask从hidden_states中提取动作token对应的隐藏状态，转为float32，shape(B*A,H)，是后续反投影的输入
2190
action_hidden_states = hidden_states[flow_action_mask].to(torch.float32)
2191
action_pred = self.action_preprocessor.action_proj_back(2192
action_hidden_states[:, : self.action_preprocessor.action_hidden_size]调用action_proj_back()反投影：将隐状态的前action_hidden_size维（通常<<H）投影回action_dim(20)维速度场（3行）。action_hidden_size为action head的瓶颈维度，shape变化：(B*A,H)→(B*A,D)其中D=20
2193
)
if条件：检查config.use_x_pred标志，决定反投影输出的解释：若True则输出为x_pred（去噪后图像），需减去噪声得到速度；若False则直接为速度场
2194
if getattr(self.config, "use_x_pred", False):
2195
v_0 = action_pred - noise.reshape(-1, noise.shape[-1])
若use_x_pred=True，v_0=action_pred-noise.reshape()用x_pred减去噪声得到速度场，这是v_prediction方式的反演，shape：(B*A,D)
else分支：v_0=action_pred直接为速度场，该配置常用于epsilon_prediction（噪声预测）或v_prediction（速度预测）方式
if条件：检查是否需要填充动作处理。(not dof_mask.all())表示有些自由度被masked，padding_action非None表示有预定义的填充轨迹
2199
if (not dof_mask.all()) and (padding_action is not None):
打印日志'use padding action'，调试信息说明启用了填充动作补偿，该特性在LIBERO中用于处理未激活的维度
计算填充部分的速度v_padding=padding_action-noisy_action（2行），即预定义轨迹与当前噪声动作的差异
2202
v_0 = (v_padding) * (1 - dof_mask) + v_0 * dof_maskmask融合：v_0=(v_padding)*(1-dof_mask)+v_0*dof_mask。active维度(dof_mask=1)用预测速度，inactive维度(dof_mask=0)用填充速度，这是LIBERO任务的关键补丁处理
2204
noisy_action = noisy_action + dt * v_0.reshape(
2205
batch_size, action_horizon, action_dim
2206
)
Euler积分更新noisy_action（4行）。noisy_action=noisy_action+dt*v_0对速度进行单步积分，dt为时间步长，reshape确保shape恢复为(B,A,D)，这是ODE求解中离散化的体现
timing_results['prefetch_forward']=time.time()-prefetch_start_time记录预填充前向和动作更新的总耗时
2208
timing_results["prefetch_forward"] = time.time() - prefetch_start_time
空行
记录cache_prep_start_time，后续需要对KV缓存进行处理（截断到prefix_length）为自回归生成准备
2210
cache_prep_start_time = time.time()
空行
if条件：若prefix_length为None则自动计算。has_true检查各batch是否存在action_token，argmax找最早的action_token位置（3行），前缀长度即图文部分的长度
2214
prefix_length = torch.argmax(flow_action_mask.float(), dim=1, keepdim=True)
2215
prefix_length[~has_true] = flow_action_mask.shape[1]2216
prefix_length = prefix_length[0]特殊处理无action_token的样本：prefix_length[~has_true]设为序列总长度，确保没有动作的样本仍然被正确处理。2216取[0]是因为argmax返回的是(B,1)，需降维为标量或1D张量用于后续切分
注释行说明下面代码处理不同版本transformer库的兼容性（key_cache/value_cache vs layers），对KV缓存进行prefix截断处理以减少自回归阶段的计算量
L2218–2362KV cache 版本兼容、postfix 阶段 attention mask 构建、ODE 积分推理循环、动作反归一化后处理。支持两种 transformers API（key_cache 新版/layers 老版），为 postfix action token 预测阶段准备 KV cache（截断到 prefix_length）和注意力掩码，通过 step_with_kvcache 闭包逐步积分去噪，最后反归一化 action 空间。
2218
# support different transformers version注释行：标记 KV cache 的版本兼容处理段，支持 transformers 中 key_cache/value_cache 和 DynamicCache.layers 两种存储格式。
2219
if hasattr(prefix_kv_cache, "key_cache"):
2220
for layer_i in range(len(prefix_kv_cache.key_cache)):
2221
prefix_kv_cache.key_cache[layer_i] = prefix_kv_cache.key_cache[layer_i][
2222
:, :, :prefix_length, :
2223
]
2224
prefix_kv_cache.value_cache[layer_i] = prefix_kv_cache.value_cache[
2225
layer_i
2226
][:, :, :prefix_length, :]
新版 transformers API 检查和截断：若 prefix_kv_cache 有 key_cache 属性，则逐层截断 key 和 value 缓存到 prefix_length（只保留图文前缀的 KV）；shape (B,num_heads,seq,head_dim) → (B,num_heads,prefix_len,head_dim)；消除之前 ODE 积分时积累的 postfix KV，后续只用 prefix KV。
2227
else:2228
for layer_i in range(len(prefix_kv_cache.layers)):
2229
prefix_kv_cache.layers[layer_i].keys = prefix_kv_cache.layers[
2230
layer_i
2231
].keys[:, :, :prefix_length, :]
2232
prefix_kv_cache.layers[layer_i].values = prefix_kv_cache.layers[
2233
layer_i
2234
].values[:, :, :prefix_length, :]
老版本 transformers API 处理：若无 key_cache，使用 DynamicCache.layers[i].keys/values 格式，同样逐层截断各 layer 的 keys 和 values 到 prefix_length；LIBERO 可能使用较旧版本（<=4.35）走此分支；语义与 key_cache 版本一致。
空行。
2236
postfix_position_ids = position_ids[:, :, prefix_length:]
2237
postfix_inputs_embeds = inputs_embeds[:, prefix_length:, :]
2238
postfix_attention_mask = attention_mask[:, prefix_length:]
2239
postfix_moe_token_types = moe_token_types[:, prefix_length:]
2240
postfix_input_ids = input_ids[:, prefix_length:]
从完整序列中提取 postfix（动作预测）段的各张量：position_ids/inputs_embeds/attention_mask/moe_token_types/input_ids 都从 prefix_length 到末尾切片；这些是即将参与 action token 预测的前向传播的输入；shape 保持 (B,...,postfix_len,...)。
空行。
2242
group_size = torch.zeros(
2243
self.config.num_experts, dtype=torch.long, device="cpu"
2244
)
2245
for i in range(self.config.num_experts):
2246
group_size[i] = (postfix_moe_token_types == i).sum()
计算 postfix 中每个 MoE 专家的 token 数量：初始化 group_size 全 0 张量 shape (num_experts,)，遍历 num_experts，统计 postfix_moe_token_types==i 的 token 个数；用于后续 expert routing 时确定各专家负责的 token 范围。
空行。
2248
# Calculate start and end indices for each expert group2249
postfix_start_indices = torch.cumsum(group_size, dim=0) - group_size2250
postfix_end_indices = torch.cumsum(group_size, dim=0)计算 postfix 各专家的起始和结束索引：postfix_start_indices = cumsum(group_size) - group_size，postfix_end_indices = cumsum(group_size)；供 TokenTypeRouter 查表时快速定位各专家负责的 token 区间 [start_i, end_i)。
空行。
2252
pad_token_id = self.processor.tokenizer.pad_token_id2253
padding_mask = input_ids == pad_token_id
获取 pad_token_id 并构建全局 padding_mask：pad_token_id 从 tokenizer 取值，padding_mask = (input_ids==pad_token_id) shape (B,seq_len,bool)；用于屏蔽 padding 位置在 attention 中的影响；覆盖 prefix 和 postfix 全部。
空行。
2255
# prefix_length, postfix_length = prefix_indices.shape[0], postfix_indices.shape[0]注释行：此行注释出的代码（prefix_length, postfix_length = prefix_indices.shape[0]...）已被后续单独计算 postfix_length 替代，代码演进中的遗留注释。
空行。
2257
postfix_length = input_ids.shape[-1] - prefix_length2258
_postfix_attention_mask = torch.ones(
2259
(batch_size, postfix_length, prefix_length + postfix_length),
2260
dtype=torch.bool,
2261
device=postfix_attention_mask.device,
2262
)
计算 postfix_length 并初始化 postfix attention mask：postfix_length = 总长 - prefix_length，创建全 True 的 _postfix_attention_mask shape (B,postfix_len,prefix_len+postfix_len,bool)；初始允许 postfix 所有 token 看到 prefix 和 postfix 全部位置（稍后被 causal/padding mask 修改）。
空行。
2264
# Use a padding mask to set the corresponding rows and columns to false.2265
# Get the padding mask for the postfix portion.2266
postfix_padding_mask = padding_mask[
2267
:, prefix_length:
2268
] # [batch_size, postfix_length]2269
full_padding_mask = padding_mask # [batch_size, prefix_length + postfix_length]准备 padding mask 用于 postfix attention：postfix_padding_mask 为 postfix 段的 padding 标记 shape (B,postfix_len)，full_padding_mask 为全序列的 padding 标记 shape (B,prefix_len+postfix_len)；后续逐 batch 用这两个 mask 将 postfix attention 矩阵的对应行列置 False。
空行。
2271
# causal mask for postfix attention2272
if self.config.causal_action_attention_mask:
2273
_postfix_attention_mask[:, :, prefix_length:] = torch.tril(
2274
torch.ones(
2275
(postfix_length, postfix_length),
2276
dtype=torch.bool,
2277
device=postfix_attention_mask.device,
2278
)
2279
)
条件分支应用因果掩码：若 self.config.causal_action_attention_mask=True，对 _postfix_attention_mask 的 postfix-to-postfix 部分应用下三角 mask（torch.tril）；这样 action token i 只能看到 i 及之前的 action token，遵循因果约束（每步输入应该依赖前面步的预测）。
空行。
2281
for batch_idx in range(padding_mask.shape[0]):
2282
# Set the rows corresponding to the padding positions to False (where the query position is padding).2283
_postfix_attention_mask[batch_idx, postfix_padding_mask[batch_idx], :] = (
2284
False2285
)
2286
# Set the columns corresponding to the padding positions to False (the key position is the padding).2287
_postfix_attention_mask[batch_idx, :, full_padding_mask[batch_idx]] = False逐 batch 应用 padding mask 到 postfix attention：循环遍历 batch，将 postfix_padding_mask[batch_idx] 位置的行置 False（padding query 无效），将 full_padding_mask[batch_idx] 位置的列置 False（padding key 无效）；LIBERO 中 batch 内可能有不同长度序列，此处处理各自的 padding。
空行。
2289
timing_results["cache_preprocessing"] = time.time() - cache_prep_start_time
记录 cache 预处理耗时到 timing_results['cache_preprocessing']。
空行。
2291
ode_start_time = time.time()
启动 ODE 积分计时，ode_start_time = time.time()。
空行。
2293
def step_with_kvcache(timestep, noisy_action):
2294
action_mask = (
2295
postfix_input_ids == self.action_token_id_set["action_token_id"]
2296
)
2297
assert action_mask.any(), "No action token found in input_ids"
2298
timestep = timestep.unsqueeze(0).repeat(noisy_action.shape[0])
2299
action_embed, adarms_cond = self.action_preprocessor.step(2300
timestep=timestep, noisy_action=noisy_action, dof_mask=dof_mask
2301
)
2302
action_embed = action_embed.reshape(-1, postfix_inputs_embeds.shape[-1])
2304
temp_inputs_embeds = postfix_inputs_embeds.clone()
2305
temp_inputs_embeds[action_mask] = action_embed.to(temp_inputs_embeds.dtype)
2306
transformer_outputs = self.model(2307
input_ids=None,2308
attention_mask=_postfix_attention_mask,
2309
position_ids=postfix_position_ids,
2310
past_key_values=prefix_kv_cache,
2311
inputs_embeds=temp_inputs_embeds,
2312
moe_token_types=postfix_moe_token_types,
2313
start_indices=postfix_start_indices,
2314
end_indices=postfix_end_indices,
2315
use_cache=False,2316
output_attentions=False,2317
output_hidden_states=False,2318
return_dict=True,2319
adarms_conds=[None, adarms_cond],2320
)
2322
hidden_states = transformer_outputs.last_hidden_state
2323
action_hidden_states = hidden_states[action_mask].to(torch.float32)
2324
action_pred = self.action_preprocessor.action_proj_back(2325
action_hidden_states[:, : self.action_preprocessor.action_hidden_size]2326
)
2327
if getattr(self.config, "use_x_pred", False):
2328
v_t = action_pred - noise.reshape(-1, noise.shape[-1])
2329
else:2330
v_t = action_pred
2331
return v_t.reshape(batch_size, action_horizon, action_dim)定义 step_with_kvcache(timestep, noisy_action) 闭包函数：ODE 积分的核心单步器，接收当前时间步和噪声动作，返回 (B,action_horizon,action_dim) 的速度场。内部：(1) 用 postfix_input_ids 定位 action token 掩码，确保至少有一个 action token；(2) 将 timestep 扩展到 batch 维度，调用 action_preprocessor.step() 进行时间嵌入和 action 投影（融入 dof_mask 和 adarms 条件）；(3) reshape action_embed 到隐层维度，复制 postfix_inputs_embeds 并替换 action token 位置的嵌入；(4) 调用 self.model() 前向传播（复用 prefix KV cache，只计算 postfix，use_cache=False）；(5) 提取 action token 对应的隐状态，投影回动作空间（action_proj_back），根据 use_x_pred 配置决定速度场 v_t = pred - noise 或 v_t = pred；(6) reshape 为 (B,action_horizon,action_dim) 返回；与 odeint() 的签名兼容。
空行。
2333
action_trajectory = odeint(
2334
step_with_kvcache, noisy_action, times[1:], method="euler"
2335
)
调用 torchdiffeq.odeint() 执行 ODE 积分：输入 step_with_kvcache 函数作为动力系统，初始状态 noisy_action，积分时间点 times[1:]（跳过 t=0，从 t_1 开始），method='euler' 使用 Euler 单步法；输出 action_trajectory shape (num_timesteps,B,action_horizon,action_dim)；整个从高噪声到完全去噪的过程由此完成，部署中通常 5~10 步。
空行。
2337
timing_results["ode_integration"] = time.time() - ode_start_time
记录 ODE 积分耗时到 timing_results['ode_integration']。
空行。
2339
postprocess_start_time = time.time()
启动后处理计时，postprocess_start_time = time.time()。
2340
predict_action = action_trajectory[-1]2341
if unnorm:2342
predict_action = (
2343
self.action_preprocessor.normalizer_action.unnormalize_data(2344
predict_action, dataset_names
2345
)
2346
)
后处理第一步：提取最终预测动作并反归一化。predict_action = action_trajectory[-1] 取积分轨迹的最后一步（完全去噪，t=1.0 对应真实动作）；若 unnorm=True（推理时默认 True），调用 normalizer_action.unnormalize_data() 将 [-1,1] 的归一化空间反归到原始动作空间（LIBERO 的绝对位置/速度），dataset_names 用于查表获得该 dataset 的 min/delta 统计，shape (B,action_horizon,action_dim)；将反归一化后的动作存入 output['predict_action']。
2347
output["predict_action"] = predict_action
注释行：说明后续处理 ground-truth action 的目的是评估对比。
2348
# normalize action chunk to get gt_action2349
if action_chunk is not None:
2350
output["gt_action"] = (
2351
self.action_preprocessor.normalizer_action.unnormalize_data(2352
action_chunk, dataset_names
2353
)
2354
)
条件分支处理 ground-truth action：若在推理时提供了 action_chunk（如 LIBERO 闭环中一步的真实动作），也进行同样的反归一化处理（但这里代码有注释写反了，应该是 unnormalize 而非 normalize），将反归一化后的真实动作存入 output['gt_action']；用于离线评估预测 vs 真实的性能对比（通常在评估脚本中用到）。
空行。
2356
timing_results["postprocessing"] = time.time() - postprocess_start_time
2357
timing_results["total_time"] = time.time() - total_start_time
2359
output["timing_results"] = timing_results
2361
return output统计各阶段耗时并组装返回值：timing_results['postprocessing'] 记录后处理耗时，timing_results['total_time'] 记录从函数入口到此的总耗时；将 timing_results 存入 output['timing_results']；返回 output 字典，包含 predict_action（必有）、gt_action（可选）、timing_results（性能分析用）；这是 generate_flow_action() 的完整输出，可直接用于闭环控制或评估。
空行或函数结束。
L2363–2398Qwen2_5_VLMoEForAction 类的 forward 方法签名与多模式路由逻辑，实现训练/验证/推理三大执行路径分发。
2363
def forward(
2364
self, mode: Optional[str] = None, predict_mode: Optional[str] = "text", **kwargs
2365
):
forward 方法签名：mode=None(训练)/predict(推理)/train/validate 进行路由分发，predict_mode 控制推理策略。本方法是 Qwen2_5_VLMoEForAction 的主入口，被训练循环和推理框架直接调用，位于模型最外层。
2366
"""2367
Main forward pass dispatcher for different execution modes.2369
This method routes execution to appropriate forward functions based on the specified mode:2370
- No mode (None): Training step with gradient disabled2371
- 'predict': Prediction/inference mode2372
- 'train': Training mode with gradients enabled2373
- 'validate': Validation mode with gradients disabled2375
Args:2376
mode (str, optional): Execution mode. If None, defaults to training step without gradients2377
predict_mode (str, optional): Prediction mode for 'predict' mode ("text", "fast", or "diffusion")Docstring 第一部分：说明 forward 为模式分发器，四种执行模式分别对应训练步骤无梯度、推理、训练有梯度、验证无梯度。Args 部分说明 mode 和 predict_mode 参数含义。实战视角：mode=None 用于数据加载 dry-run、mode='train' 用于微调时反向传播、mode='validate' 用于验证集评估、mode='predict' 用于闭环控制器推理。
2378
**kwargs: Additional arguments passed to the selected forward function2380
Returns:2381
Model outputs appropriate for the selected mode2383
Todo:2384
- Add support for distinguishing multi-modal data types in prediction mode2385
"""Docstring 第二部分（Args 续+Returns+Todo）：说明 **kwargs 透传给下游 forward 函数，Returns 为模型输出，Todo 提示多模态推理还需扩展。文档与实现对齐：目前 predict_mode 参数在 generate_flow_action 中被消费，控制去噪策略（Euler 固定为 5 步）。
mode 为 None 或 falsy 时进入第一分支（默认训练步骤）：with torch.no_grad() 显式禁梯度，调用 train_step_forward 进行前向传播但不计算梯度。用途：DDP 验证、dry-run 测试、推理前置步骤。不会产生损失反向传播。
2389
elif mode == "predict":
2390
return self.generate_flow_action(predict_mode=predict_mode, **kwargs)
mode == 'predict' 分支（推理/生成模式）：调用 generate_flow_action（line:1959）进行 flow matching 去噪推理，返回预测的 20 维或 7 维（LIBERO mask 后）动作。predict_mode 透传控制去噪算法选择（当前固定 Euler 5 步）。LIBERO 实战：此分支在闭环中使用，返回值需经 dof_mask[7:]=0 遮挡后馈给机械臂。
mode == 'train' 分支（显式训练模式）：调用 train_step_forward 并禁用 KV cache(use_cache=False)，保留梯度用于反向传播。区别于 mode=None：此分支显式保梯度，用于微调循环中的训练步骤。
2393
elif mode == "validate":
2394
with torch.no_grad():2395
return self.train_step_forward(use_cache=False, **kwargs)
mode == 'validate' 分支（验证模式）：with torch.no_grad() 禁梯度，use_cache=False 禁用 KV 缓存，逐样本评估不更新参数。用于验证集损失/指标计算。与 mode=None 的区别：validate 显式命名意图清晰，同时禁 cache 和梯度。
异常处理：mode 取值不在 [None, 'predict', 'train', 'validate'] 时抛出 NotImplementedError。起防御性编程作用，捕获调用方传入无效 mode 的 bug。
空行，函数体结束，下一个方法 prepare_inputs_for_generation(line:2399) 开始。
L2399–2552prepare_inputs_for_generation 方法：多模态生成阶段的输入准备，包括 MoE token 类型填充、KV缓存位置切片、视觉输入管理、attention mask 预处理、传感器数据传递，为 Euler flow matching 去噪生成提供规范化输入。
2399
def prepare_inputs_for_generation(
2400
self,2401
input_ids,
2402
past_key_values=None,2403
attention_mask=None,函数签名开始：def prepare_inputs_for_generation(self, input_ids, past_key_values, attention_mask, inputs_embeds)。生成阶段输入预处理入口，由 transformers.GenerationMixin 的 generate() 调用。这里声明的前4个主要参数用于传统语言模型的生成流程。
2404
inputs_embeds=None,2405
moe_token_types=None,2406
cache_position=None,2407
position_ids=None,2408
use_cache=True,MoE 路由相关参数(moe_token_types, cache_position, position_ids, use_cache)。moe_token_types 是关键的 token 类型标记(0=文本/图像, 1=动作/状态)，TokenTypeRouter 后续根据这个查表路由到对应专家。cache_position 指示生成当前步的位置(用于 KV 缓存索引)。
2409
pixel_values=None,2410
pixel_values_videos=None,2411
image_grid_thw=None,2412
video_grid_thw=None,2413
second_per_grid_ts=None,视觉多模态参数(pixel_values, pixel_values_videos, image_grid_thw, video_grid_thw)。这些参数在生成的初始步(cache_position[0]==0)被转发到视觉编码器，后续步骤为空(减少计算)。image/video_grid_thw 记录了图像/视频序列的时空网格维度，用于拼接到 KV 缓存。
2414
dataset_names=None,2415
proprioception=None,2416
dof_mask=None,2417
agent_pos_mask=None,2418
**kwargs,
时间、数据集、传感器参数(second_per_grid_ts, dataset_names, proprioception, dof_mask, agent_pos_mask, **kwargs)。dataset_names 标识数据源，用于查表获取 action normalizer；proprioception 是本体感觉(臂关节角度等)；dof_mask 掩码哪些自由度参与(LIBERO 只用前7维，其余mask=0)。
2419
):
函数签名结束，右括号和冒号。进入 docstring 说明区。
2420
"""2421
Prepare inputs for autoregressive generation with multi-modal support.2423
This method handles input preparation for generation, including proper slicing of inputs2424
based on cache position, MoE token type management, and multi-modal data handling.2425
Vision inputs are selectively forwarded only when needed during generation.2427
Args:2428
input_ids: Input token IDs2429
past_key_values: Cached key-value pairs from previous generation stepsdocstring 主说明部分：该方法为自回归生成的输入准备阶段，处理多模态数据、KV缓存切片、token 类型管理。特别说明了视觉输入的选择性转发策略(初始步才用，后续步跳过)。注明这是对原 Qwen2.5-VL 的重写以支持 MoE 和 action token。
2430
attention_mask: Attention mask for input tokens2431
inputs_embeds: Pre-computed input embeddings2432
moe_token_types: Token type assignments for MoE routing2433
cache_position: Current cache position for generation2434
position_ids: Position IDs for tokens2435
use_cache: Whether to use key-value caching2436
pixel_values: Image pixel values2437
pixel_values_videos: Video pixel values2438
image_grid_thw: Image grid dimensions2439
video_grid_thw: Video grid dimensions2440
second_per_grid_ts: Time interval per temporal griddocstring Args 前半部分：input_ids 是 token 序列(包含图像占位符)；past_key_values 是前面层累积的 KV 缓存；attention_mask 掩码；inputs_embeds 预计算的嵌入；moe_token_types token 类型标记；cache_position 缓存读取位置；position_ids 绝对位置编码。
2441
dataset_names: Dataset names for processing2442
proprioception: Proprioceptive sensor data2443
dof_mask: Degrees of freedom mask2444
agent_pos_mask: Agent position mask2445
**kwargs: Additional arguments2447
Returns:2448
dict: Prepared model inputs for generation step2450
Todo:docstring Args 后半部分及 Returns：pixel_values_videos/image_grid_thw/video_grid_thw 用于视觉多模态；second_per_grid_ts 是时间采样间隔；dataset_names/proprioception/dof_mask/agent_pos_mask 是 action head 所需的传感器和控制信息。Returns 是 dict，包含上述所有处理过的输入，直接喂给模型前向传递。
2451
- Test this function thoroughly with various input configurations2453
Note:2454
This is an overridden method that handles specific cases for multi-modal generation:2455
- Slices input_ids through cache_position to keep only unprocessed tokens2456
- Handles special cases for input_embeds, generation methods, and GPU synchronization2457
- Manages vision inputs to avoid unnecessary forward passes2458
"""docstring Note 部分：列举了该方法的关键处理逻辑——(1)通过 cache_position 切片 input_ids 仅保留待处理 token；(2)处理 input_embeds 特殊情况(不含 input_ids)；(3)静态缓存时 GPU 同步边界处理；(4)视觉输入管理避免多余前向。该方法是原 Qwen generate 流程的多模态扩展。
2459
# Initialize MoE token types if not provided2460
if moe_token_types is None:
2461
moe_token_types = torch.zeros_like(
初始化 moe_token_types：若未传入则创建全 0 tensor (torch.zeros_like(input_ids))，shape=(B, S)，代表所有 token 默认为文本/图像类型。FIXME 注释指出当 input_embeds 代替 input_ids 时这里处理有问题，实战中 embedding-only 路径较少见。
赋值 moe_token_types = torch.zeros_like(input_ids)。shape: input_ids(B,S) → moe_token_types(B,S)，dtype=int64。当生成开始且无预设 token 类型时，假定所有 token 初始都走文本/图像专家(类型0)路由。
2464
else:2465
# Ensure moe_token_types length matches input_ids2466
if moe_token_types.shape[1] < input_ids.shape[1]:
2467
# Calculate required padding length2468
pad_length = input_ids.shape[1] - moe_token_types.shape[1]
else 分支：若已传入 moe_token_types，检查其长度是否匹配 input_ids。内层 if 条件判断 moe_token_types.shape[1] < input_ids.shape[1]，即 token 类型序列比 input_ids 短(可能前向过程中动态添加了新 token)，需要填充到相同长度。
2469
# Create padding tensor with default token type (0)2470
pad_tensor = torch.zeros(
2471
(moe_token_types.shape[0], pad_length),2472
dtype=moe_token_types.dtype,
计算填充长度并创建填充 tensor：pad_length = input_ids.shape[1] - moe_token_types.shape[1]。然后创建 pad_tensor = torch.zeros((B, pad_length), dtype=..., device=...)，shape=(B, pad_length)，用全 0 填充(默认文本类型)。
2473
device=moe_token_types.device,
2474
)
2475
# Concatenate padding to existing moe_token_types2476
moe_token_types = torch.cat([moe_token_types, pad_tensor], dim=1)pad_tensor 指定 dtype 和 device：确保填充 tensor 与原 moe_token_types 的数据类型和设备一致。然后 torch.cat([moe_token_types, pad_tensor], dim=1) 沿序列维连接，结果 shape=(B, input_ids.shape[1])。LIBERO 实战中这通常在第一步生成时触发(追加 action token 时)。
空行及注释开始："Handle input slicing based on cache state and special cases"。接下来的大 if 块处理不同的缓存状态和特殊边界情况，确保生成步骤中 input_ids/moe_token_types 的切片正确对应当前要处理的 token。
2479
if past_key_values is not None:
2480
if (2481
inputs_embeds is not None and input_ids.shape[1] == 0
2482
): # Exception 4: input_embeds case2483
inputs_embeds = inputs_embeds[:, -cache_position.shape[0] :]大 if 块：if past_key_values is not None，即有累积 KV 缓存(非首步生成)。此时需要从 input_ids 中提取「仅本步新增的 token」，避免重复编码。内层嵌套的 if/elif/elif 处理 4 种边界情况和正常情况。
2484
moe_token_types = moe_token_types[:, -cache_position.shape[0] :]2485
elif inputs_embeds is not None or ( # Exception 1: input_embeds provided
第一种异常情况(Exception 4)：input_embeds is not None and input_ids.shape[1] == 0，即用预计算嵌入代替 token IDs。此时不能依赖 input_ids，而是切片 inputs_embeds 和 moe_token_types 至最后 cache_position.shape[0] 个(本步)。shape: (B, seq) → (B, cache_position.shape[0])。
2486
is_torchdynamo_compiling() or cache_position[-1] >= input_ids.shape[1]
2487
): # Exception 3: GPU sync edge case2488
input_ids = input_ids[:, -cache_position.shape[0] :]2489
moe_token_types = moe_token_types[:, -cache_position.shape[0] :]2490
elif (第一种正常/特殊情况(Exception 1/3)：if inputs_embeds is not None or (is_torchdynamo_compiling() or cache_position[-1] >= input_ids.shape[1])。检测到使用 embedding-only 模式或 TorchDynamo 编译中或 cache_position 越界(GPU 同步边界)。此时也用 input_ids[:, -cache_position.shape[0]:] 切片，保留末尾 cache_position.shape[0] 个 token。
2491
input_ids.shape[1] != cache_position.shape[0]
2492
): # Default case (Exception 2 is no-op)2493
cache_pos = cache_position.clone()
第二种正常情况(Exception 2 是 no-op，即 elif 条件为真且符合注释说法)：elif input_ids.shape[1] != cache_position.shape[0]，即 input_ids 长度与 cache_position 长度不匹配。这是生成流程中常见的边界情况。
用 cache_position 的具体索引数组切片：cache_pos = cache_position.clone()，然后 input_ids = input_ids[:, cache_pos]，moe_token_types = moe_token_types[:, cache_pos]。这是最通用的做法，根据缓存位置数组直接索引当前步的 token。shape: (B, S) → (B, cache_position.shape[0])。
2497
# Skip vision inputs for continuation steps (not initial generation)2498
if cache_position[0] != 0:
2499
pixel_values = None2500
pixel_values_videos = None跳过非初始步的视觉输入：if cache_position[0] != 0(生成不是从步数 0 开始)，则将 pixel_values 和 pixel_values_videos 置为 None。因为视觉编码器计算(patch embedding + 时空聚合)只在首步进行，后续步可复用 KV 缓存。这是重要的推理优化。空行后接后续逻辑。
2502
# Determine whether to use inputs_embeds or input_ids for this generation step2503
if inputs_embeds is not None and len(cache_position) == inputs_embeds.shape[1]:
2504
model_inputs = {"inputs_embeds": inputs_embeds, "input_ids": None}确定用 inputs_embeds 还是 input_ids：注释说明逻辑。if inputs_embeds is not None and len(cache_position) == inputs_embeds.shape[1]，即预计算嵌入的长度与缓存位置长度相同，此时用嵌入而非 token。
构造 model_inputs dict，赋值相应的编码方式。else 分支：使用 input_ids，model_inputs = {"input_ids": input_ids, "inputs_embeds": None}。shape: input_ids(B, cache_position.shape[0])。空行分隔开后续的 attention mask 处理。
2508
# Prepare 4D causal attention mask for static cache2509
if isinstance(past_key_values, StaticCache) and attention_mask.ndim == 2:
2510
if model_inputs["inputs_embeds"] is not None:
2511
batch_size, sequence_length, _ = inputs_embeds.shape
2512
device = inputs_embeds.device
静态缓存时的 4D attention mask 预处理：if isinstance(past_key_values, StaticCache) and attention_mask.ndim == 2，即有静态缓存且原 mask 是 2D(B, S)。需要转换为 4D(B, num_heads, seq, max_cache_shape) 才能与静态缓存兼容。
初始化 batch_size 和 sequence_length：从 inputs_embeds 或 input_ids 推导(含 device 信息)。batch_size, sequence_length = input_ids.shape。这些是构造 4D mask 所需的形状参数。
2517
attention_mask = (
2518
self.model._prepare_4d_causal_attention_mask_with_cache_position(2519
attention_mask,
2520
sequence_length=sequence_length,
2521
target_length=past_key_values.get_max_cache_shape(),
调用模型的 mask 转换方法：attention_mask = self.model._prepare_4d_causal_attention_mask_with_cache_position(...)。这是 Qwen2.5-VL 的内部方法，处理因果掩码与静态缓存的兼容性。输入 2D mask，输出 4D mask，shape: (B, 1, seq, max_cache) ≈ (B, 1, 1, past_len+seq)。
2522
dtype=self.lm_head.weight.dtype,2523
device=device,
2524
cache_position=cache_position,
2525
batch_size=batch_size,
2526
config=self.config,4D mask 方法的主要参数：attention_mask(原 2D)、sequence_length(当前序列长)、target_length(缓存最大形状)、dtype(输出数据类型)、device(输出设备)。这些共同定义转换后 mask 的精确形状。
4D mask 方法的缓存相关参数：cache_position(当前生成步在缓存中的位置)、batch_size(批大小)、config(模型配置)、past_key_values(静态缓存对象)。空行后进入最终的模型输入组装。
2531
# Assemble all model inputs for generation2532
model_inputs.update(
2533
{2534
"position_ids": position_ids,
2535
"past_key_values": past_key_values,
注释和 model_inputs.update() 开始："Assemble all model inputs for generation"。用字典 update() 方法追加所有必需的前向参数。前两个字段是 position_ids(位置编码)和 past_key_values(KV 缓存)。
2536
"moe_token_types": moe_token_types,
2537
"use_cache": use_cache,
2538
"attention_mask": attention_mask,
2539
"pixel_values": pixel_values,
2540
"pixel_values_videos": pixel_values_videos,
传递 transformer 核心参数：moe_token_types(路由索引)、use_cache(开启缓存写)、attention_mask(4D 因果掩码)、pixel_values(图像张量，首步非空)、pixel_values_videos(视频张量，首步非空)。
2541
"image_grid_thw": image_grid_thw,
2542
"video_grid_thw": video_grid_thw,
2543
"cache_position": cache_position,
2544
"second_per_grid_ts": second_per_grid_ts,
2545
"proprioception": proprioception,
传递多模态视觉参数：image_grid_thw(图像网格时空维度)、video_grid_thw(视频网格时空维度)、cache_position(缓存位置数组)、second_per_grid_ts(视频时间采样率)、proprioception(关节状态，LIBERO 实战中通常 shape=(B,7) 或用 PROPRI_DROPOUT 隐藏)。
2546
"dataset_names": dataset_names,
2547
"dof_mask": dof_mask,
2548
"agent_pos_mask": agent_pos_mask,
2549
}
2550
)
传递控制和数据集参数：dataset_names(数据来源，用于 action normalizer 查表)、dof_mask(掩码有效 dof，LIBERO 只前 7 维有效，其余 dof_mask=0)、agent_pos_mask(代理位置掩码，可选)。这些字段全部到达 forward() 的各个子模块。
返回 model_inputs dict 给 transformers.generate() 的下一个迭代。dict 包含所有必需的前向参数、缓存、mask、多模态数据和控制信号，完整支持 Qwen2.5-VLMoE 的多模态生成流程(视觉编码→注意力融合→action flow matching→Euler 去噪)。
L2553–2716两个生成辅助方法：_get_image_nums_and_video_nums() 通过扫描 input_ids 中的视觉 token 计数图像/视频数，用于后续张量分割；_expand_inputs_for_generation() 在束搜索等扩展场景重复扩展输入，特殊处理非标准 batch 维的视觉张量（pixel_values 按样本内 grid 长度分割+重复）。
2553
def _get_image_nums_and_video_nums(
2554
self,2555
input_ids: Optional[torch.LongTensor],
2556
) -> Tuple[torch.Tensor, torch.Tensor]:
2557
"""def _get_image_nums_and_video_nums(...) 函数签名：输入 input_ids (B,S)，返回图像数和视频数张量，用于后续计算每个样本的视觉内容长度。被 _expand_inputs_for_generation():2628 调用。
2558
Get the number of images and videos for each sample to calculate tensor separation lengths.2560
These parameters are computed directly from input_ids rather than being passed through2561
the processor to avoid unpredictable impacts from interface modifications.2563
Args:2564
input_ids (torch.LongTensor): Input token IDs of shape (batch_size, sequence_length)2566
Returns:2567
tuple:2568
- image_nums (torch.LongTensor): Number of images per sample2569
- video_nums (torch.LongTensor): Number of videos per sample2570
"""docstring 完整说明：通过查找 input_ids 中的特定 token ID（而非依赖 processor）直接计算图像/视频数，避免接口改动的不可预期影响。返回两个 (B,) 张量，每元素为该样本内的图像/视频计数。
2571
image_token_id = self.config.image_token_id2572
video_token_id = self.config.video_token_id2573
vision_start_token_id = self.config.vision_start_token_id从 config 中取出 image_token_id、video_token_id、vision_start_token_id 三个特殊 token ID，用于后续的 input_ids 扫描和掩码匹配。
2575
# Find vision start tokens and their following tokens2576
vision_start_mask = input_ids == vision_start_token_id
2577
vision_first_mask = torch.roll(vision_start_mask, shifts=1, dims=1)
2578
image_mask = input_ids == image_token_id
2579
video_mask = input_ids == video_token_id
构造掩码：vision_start_mask 找出所有视觉起始 token，vision_first_mask 是 start 右移一位（找紧跟 start 后的 token），image_mask/video_mask 找所有图像/视频 token。形状均为 (B,S)。
2581
# Count images and videos following vision start tokens2582
image_nums = torch.sum(vision_first_mask & image_mask, dim=1)2583
video_nums = torch.sum(vision_first_mask & video_mask, dim=1)2585
return image_nums, video_nums联合掩码逻辑与后沿求和：(vision_first_mask & image_mask).sum(dim=1) 统计每样本的'紧跟 vision_start 后出现的图像 token 数'，返回 (B,) 的 image_nums 和 video_nums（2586 行空行）。实战：LIBERO 使用这两个计数来反向 split pixel_values 和 video 特征。
2587
def _expand_inputs_for_generation(
2588
self,2589
expand_size: int = 1,
2590
is_encoder_decoder: bool = False,
2591
input_ids: Optional[torch.LongTensor] = None,2592
**model_kwargs,
2593
) -> Tuple[torch.LongTensor, Dict[str, Any]]:def _expand_inputs_for_generation(...) 函数签名：重写 HF GenerationMixin 的扩展方法，支持非标准 batch 维的多模态张量（如 pixel_values 在多图像样本时无统一 batch 维）。expand_size 用于束搜索、多候选生成等策略的扩展倍数。
2594
"""2595
Expand inputs for generation with support for multi-modal tensors.2597
This is an overridden method that supports expanding tensors without a standard batch2598
size dimension, specifically for vision-related tensors:2599
- pixel_values.shape[0] = sum(sequence_lengths for all image samples)2600
- image_grid_thw.shape[0] = sum(num_images for all samples)2601
- Similar patterns for video tensors2603
Args:2604
expand_size (int): Factor by which to expand inputs (for beam search, etc.)2605
is_encoder_decoder (bool): Whether using encoder-decoder architecture2606
input_ids (torch.LongTensor, optional): Input token IDs2607
**model_kwargs: Additional model arguments to expand2609
Returns:2610
tuple: (expanded_input_ids, expanded_model_kwargs)2611
"""docstring：说明重写必要性及处理的视觉张量特殊形状（pixel_values.shape[0]=所有样本的图像序列长之和，image_grid_thw.shape[0]=所有样本的图像数之和）。这些张量没有统一 batch 维，需要先按每样本的图像/视频数 split 再逐样本 repeat。
快速路径：expand_size==1 时不需扩展，直接返回原 input_ids 和 model_kwargs（2614 行空行）。
2615
# Define keys for vision-related tensors that need special handling2616
visual_keys = [
2617
"pixel_values",
2618
"image_grid_thw",
2619
"pixel_values_videos",
2620
"video_grid_thw",
2621
"second_per_grid_ts",
2622
]
定义 visual_keys 列表，列出需要特殊处理的视觉张量键名（pixel_values、image_grid_thw、pixel_values_videos、video_grid_thw、second_per_grid_ts），与其他标准 Tensor 参数区分，因为它们没有统一 batch 维（2623 行空行）。
2624
def _expand_dict_for_generation_visual(dict_to_expand):
2625
"""Expand vision-related tensors based on image/video counts per sample."""2626
image_grid_thw = model_kwargs.get("image_grid_thw", None)
2627
video_grid_thw = model_kwargs.get("video_grid_thw", None)
2628
image_nums, video_nums = self._get_image_nums_and_video_nums(input_ids)2630
def _repeat_interleave_samples(x, lengths, repeat_times):
2631
"""Split tensor by lengths and repeat each sample."""2632
samples = torch.split(x, lengths)
2633
repeat_args = [repeat_times] + [1] * (x.dim() - 1)
2634
result = torch.cat(
2635
[sample.repeat(*repeat_args) for sample in samples], dim=0
2636
)
2637
return result内层函数 _expand_dict_for_generation_visual 定义及 docstring：处理视觉张量的扩展，先获取 image_grid_thw 和 video_grid_thw，调用 _get_image_nums_and_video_nums() 计数；内嵌 _repeat_interleave_samples 辅助函数用于按 lengths 分割后逐子张量 repeat。
2639
for key in dict_to_expand:
2640
if key == "pixel_values":
2641
# Split images into samples and compute sequence lengths2642
samples = torch.split(image_grid_thw, list(image_nums))2643
lengths = [torch.prod(sample, dim=1).sum() for sample in samples]
2644
dict_to_expand[key] = _repeat_interleave_samples(
2645
dict_to_expand[key], lengths=lengths, repeat_times=expand_size
2646
)
处理 pixel_values：for 循环遍历 dict_to_expand，key=='pixel_values' 时先按 image_nums 分割 image_grid_thw，计算每样本实际序列长度（grid 沿 dim=1 乘积求和），作为 lengths 传给 _repeat_interleave_samples。(B*img_seq, C, H, W)→(B*expand_size*img_seq, C, H, W)。
2647
elif key == "image_grid_thw":
2648
# Expand based on number of images per sample2649
lengths = list(image_nums)2650
dict_to_expand[key] = _repeat_interleave_samples(
2651
dict_to_expand[key], lengths=lengths, repeat_times=expand_size
2652
)
处理 image_grid_thw：直接以 image_nums 作为 lengths（每样本的图像数），调用 _repeat_interleave_samples 扩展。(sum(num_images), 3)→(sum(num_images)*expand_size, 3)。
2653
elif key == "pixel_values_videos":
2654
# Split videos into samples and compute sequence lengths2655
samples = torch.split(video_grid_thw, list(video_nums))2656
lengths = [torch.prod(sample, dim=1).sum() for sample in samples]
2657
dict_to_expand[key] = _repeat_interleave_samples(
2658
dict_to_expand[key], lengths=lengths, repeat_times=expand_size
2659
)
处理 pixel_values_videos：与 pixel_values 逻辑相同，按 video_nums 分割 video_grid_thw 得每样本的视频 grid，计算序列长度再用 _repeat_interleave_samples 扩展。(B*video_seq, C, T, H, W)→(B*expand_size*video_seq, C, T, H, W)。
2660
elif key == "video_grid_thw":
2661
# Expand based on number of videos per sample2662
lengths = list(video_nums)2663
dict_to_expand[key] = _repeat_interleave_samples(
2664
dict_to_expand[key], lengths=lengths, repeat_times=expand_size
2665
)
处理 video_grid_thw：直接以 video_nums 作为 lengths（每样本的视频数），调用 _repeat_interleave_samples 扩展。(sum(num_videos), 3)→(sum(num_videos)*expand_size, 3)。
2666
elif key == "second_per_grid_ts":
2667
# Handle list-type temporal grid data2668
if not isinstance(dict_to_expand[key], list):
2669
raise TypeError(2670
f"Expected value for key '{key}' to be a list, but got {type(dict_to_expand[key])} instead."
2671
)
2672
tensor = torch.tensor(dict_to_expand[key])
2673
lengths = list(video_nums)2674
tensor = _repeat_interleave_samples(
2675
tensor, lengths=lengths, repeat_times=expand_size
2676
)
2677
dict_to_expand[key] = tensor.tolist()
2678
return dict_to_expand处理 second_per_grid_ts（列表格式的时间网格）：检查类型是否为 list，转换为 tensor，按 video_nums split 和 repeat，最后转回列表（保持原格式）。实战：LIBERO 的视频时间戳辅助数据，扩展时需保持列表化以适配生成流程。
2680
def _expand_dict_for_generation(dict_to_expand):
2681
"""Expand standard tensors using repeat_interleave."""2682
for key in dict_to_expand:
2683
if (2684
key != "cache_position"
2685
and dict_to_expand[key] is not None
2686
and isinstance(dict_to_expand[key], torch.Tensor)
2687
and key not in visual_keys
2688
):
2689
dict_to_expand[key] = dict_to_expand[key].repeat_interleave(
2690
expand_size, dim=02691
)
2692
return dict_to_expand内层函数 _expand_dict_for_generation 定义及 docstring（2679 行空行）：标准 Tensor 扩展逻辑，遍历 dict_to_expand，对所有 Tensor（除 cache_position 和 visual_keys）沿 dim=0 用 repeat_interleave 扩展。处理 attention_mask、position_ids、past_key_values 等非视觉参数。
2694
# Expand visual inputs only if input_ids is available for counting images/videos2695
# If input_ids is unavailable, visual inputs won't be used, so no expansion needed2696
if input_ids is not None and input_ids.numel() != 0:
2697
model_kwargs = _expand_dict_for_generation_visual(model_kwargs)
条件判断：仅当 input_ids 可用且非空（input_ids.numel()!=0）时才调用 _expand_dict_for_generation_visual 扩展视觉张量（2693 行空行）。若 input_ids 不可用（如某些编码器-解码器情景），视觉输入不需扩展。
2699
# Expand input_ids using standard repeat_interleave2700
if input_ids is not None:
2701
input_ids = input_ids.repeat_interleave(expand_size, dim=0)扩展 input_ids：若存在则沿 dim=0 用 repeat_interleave 扩展 expand_size 倍。(B,S)→(B*expand_size,S)。
2703
# Expand all other model arguments2704
model_kwargs = _expand_dict_for_generation(model_kwargs)
调用标准扩展函数处理剩余 model_kwargs 中的非视觉 Tensor（已排除 visual_keys）（2705 行空行）。
2706
# Handle encoder-decoder specific expansion2707
if is_encoder_decoder:2708
if model_kwargs.get("encoder_outputs") is None:
2709
raise ValueError(2710
"If `is_encoder_decoder` is True, make sure that `encoder_outputs` is defined."
2711
)
2712
model_kwargs["encoder_outputs"] = _expand_dict_for_generation(
2713
model_kwargs["encoder_outputs"]
2714
)
编码器-解码器特殊路径：若是 encoder-decoder 模型，需确保 encoder_outputs 存在，并对其内部 Tensor 也应用标准扩展（encoder_outputs 是字典，可能包含 last_hidden_state 等）。
返回扩展后的 (input_ids, model_kwargs) 元组，供后续生成循环使用。完成了对所有多模态输入的一致性扩展，用于束搜索、多候选生成等场景。
源码零改写,与仓库 wall-x/wall_x/model/qwen2_5_based/modeling_qwen2_5_vl_act.py 逐字节一致(commit 97406f2)。
生成于 wall-x LIBERO 微调项目 · ← 返回导读