跟着一次前向传播的数据流读 wall-x 源码（commit 97406f2，
仓库 ~/learn/02_wall-x/wall-x）。每个代码 cell 都标注了真实文件与行号，
可以直接跳过去继续下钻。配套的结果展示页在 这里。
📥 要真正跑代码? 下载可执行版
.ipynb —— 每个模块配一段纯 PyTorch 的等价复现(CPU 秒级跑出数字和图),
并用 show() 从仓库实时读取完整原始源码。本页是它的只读快照。
🔬 逐行精读版 —— 本页是按数据流走的导读（节选）。要看每一行代码的中文讲解 （做什么 / 为什么 / shape / 坑，4600 行全覆盖无跳行），进 源码逐行精读： vla_mixin.py · 987 行 · action_head.py · 808 行 · modeling_qwen2_5_vl_act.py · 2716 行 · train_libero_wrapper.py · 89 行。 每个行号可直链（如 vla_mixin.html#L34）。
模型只有三个文件是主战场：
wall_x/model/ ├── qwen2_5_based/ │ └── modeling_qwen2_5_vl_act.py # 2716 行主模型:MoE 化的 Qwen2.5-VL + 训练/推理入口 ├── vla_mixin.py # 987 行:路由/MoE block/embedding 注入等积木 └── action_head.py # 808 行:ActionProcessor(flow matching 的进出投影)
整个模型类叫 Qwen2_5_VLMoEForAction，加载方式和普通 HF 模型一样。
决定它"长成 VLA"的是 config 里三个开关（抄自权重的 config.json）：
from wall_x.model.qwen2_5_based.modeling_qwen2_5_vl_act import Qwen2_5_VLMoEForAction model = Qwen2_5_VLMoEForAction.from_pretrained("x-square-robot/wall-oss-flow") # config.json 的三个关键开关: # "num_experts": 2, # 每层 2 个 FFN 专家 # "attention_moe": false, # 注意力不分家 —— 图文/动作 token 同一个注意力 # "mlp_moe": true, # FFN 分家 —— 按 token 类型各走各的 # 以及动作空间定义(20 维跨本体布局): # "dof_config": { follow_left_ee_cartesian_pos: 3, follow_left_ee_rotation: 3, # follow_left_gripper: 1, ...right 同构 7 维..., head_actions: 2, # height: 1, car_pose: 3 } # 共 20
前向的第一步是把多模态输入拼成一条序列。文本先正常过词嵌入；然后三次
masked scatter——把图像 patch 特征、本体状态 embedding、带噪动作 embedding
分别替换进序列里对应的占位 token 位置（<|image_pad|>、proprio 占位、
<|action|> × chunk 长度）：
# 图像: ViT 输出替换 <|image_pad|> 占位符 image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw) inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) # 本体状态: 20 维 propri + 20 维 mask -> propri_proj -> 1 个 token inputs_embeds = self.scatter_proprioception_embeddings( input_ids, inputs_embeds, proprioception, dataset_names, agent_pos_mask ) # 动作: 训练时对 GT action chunk 加噪后投影成 N 个 token(见 §5), # 同时返回本次的监督目标 flow inputs_embeds, flow, adarms_cond = self.scatter_flow_action_embeddings( input_ids, inputs_embeds, action_chunk, dataset_names, dof_mask )
常规 MoE 的 router 是个可学习的打分网络；这里的 TokenTypeRouter
退化成一行取模——token 进序列时就已经带了类型标记（0=图文，1=动作/状态），
路由只是按类型分拣。没有负载均衡问题、没有路由抖动，因为分工是先验固定的：
class TokenTypeRouter(nn.Module): def __init__(self, num_experts: int): super().__init__() self.num_experts = num_experts def forward(self, token_types: torch.Tensor) -> torch.Tensor: # token_types: (batch, seq) 每个 token 的类型标记 experts_indices = token_types % self.num_experts return experts_indices
每层 transformer 的 FFN 被替换成 SparseMoeBlock：先把整个 batch 的 token
按专家序号 permute 重排（同专家的 token 变成连续内存，一次矩阵乘算完），
各专家分头过自己的 SwiGLU，再 unpermute 放回原位。专家 0（图文）中间维 11008，
专家 1（动作）中间维 2048——这就是"0.45B 动作专家"的全部主体：
class BlockSparseMLP(nn.Module): # 单个专家 = 标准 SwiGLU def _full_mlp(self, hidden_state): gate_out = self.gate_proj(hidden_state) up_out = self.up_proj(hidden_state) act_out = self.act_fn(gate_out) * up_out # SiLU(gate) ⊙ up return self.down_proj(act_out) class SparseMoeBlock(nn.Module): def forward(self, hidden_states, experts_indices, start_indices, end_indices): # 按专家序号重排:同专家 token 连续排布 permuted_inputs, row_id_map = ops.permute(flat_hidden, experts_indices) final_output = torch.zeros_like(permuted_inputs) for expert_idx, expert in enumerate(self.experts): start, end = start_indices[expert_idx], end_indices[expert_idx] if start == end: continue dim_input = self.dim_inputs[expert_idx] # 专家各自的输入维度 expert_input = permuted_inputs[start:end, :dim_input] partial_output = expert(expert_input) final_output[start:end, :dim_input] = partial_output[:, :dim_input] # 放回原顺序 final_output = ops.unpermute(final_output, row_id_map, probs) return final_output.reshape(batch_size, seq_length, hidden_dim)
注意力侧（attention_moe=false）完全不动——36 层里
图文 token 和动作 token 做的是同一场注意力，Q/K/V 权重共享。动作 token 因此能逐层
读到视觉表征，这是"深度融合"的实现位置。
ActionProcessor.forward 是 flow matching 的训练侧入口。四步：
采样噪声强度 t（Beta 分布偏向高噪端）→ 沿直线插值构造带噪动作 →
记下监督目标 flow = action − noise → 拼上 dof_mask 和时间步编码投影成 token：
def sample_time(self, batch_size, device, dtype): sample = self.beta_dist.sample([batch_size]) # Beta(1.5, 1) 偏向 1 time = (1 - sample) * self.s # 翻转 × 0.999 -> 偏向 0(高噪端) return time def forward(self, action_chunk, dataset_names, dof_mask=None): # 1. 加噪:t=0 纯噪声, t=1 真动作, 直线插值 noise = torch.randn_like(action_chunk) time = self.sample_time(batch_size, device, dtype) noisy_action = (1 - time_expanded) * noise + time_expanded * action_chunk flow = action_chunk - noise # 监督目标:插值直线的"速度" # 2. 时间步 -> 正弦位置编码 time_embed = self.time_embed(time) # 3. [带噪动作 20 维 | dof_mask 20 维] -> w1 -> 2048 维动作 token noisy_action = torch.cat([noisy_action, dof_mask], dim=-1) action_embed = self.w1(noisy_action) # 4. 与时间步编码融合(w2/w3 小 MLP) concat_embed = torch.cat([action_embed, time_embed], dim=-1) action_time_embed = self.w3(self.act_fn(self.w2(concat_embed))) return action_time_embed, flow, adarms_cond
动作 token 过完 36 层 transformer 后，最后一层隐状态被
action_proj_back 投回 20 维，与 flow 做 MSE。
dof_mask 在这里第二次出场：LIBERO 只占 20 维中的前 7 维，
后 13 个空槽的预测误差不产生任何梯度：
def flow_loss(self, action_hidden_states, flow, action_chunk, dof_mask=None, ...): action_pred = self.action_proj_back( action_hidden_states[:, : self.action_hidden_size] ) v_pred = action_pred loss = self.mse_loss(v_pred, flow) # 逐元素 MSE if dof_mask is not None: loss = loss * dof_mask # 空槽维度不计损失
汇合点在主模型 forward 的结尾：语言头照常算 next-token 交叉熵（VQA/CoT 能力靠它）， 动作侧算 flow loss，两个 loss 加权相加——这就是"语言和动作在同一个可微框架里"的字面意思：
outputs = self.model(..., moe_token_types=moe_token_types, adarms_conds=[None, adarms_cond]) hidden_states = outputs[0] logits = self.lm_head(hidden_states) loss, cross_entropy_loss, flow_loss, ... = self.compute_loss( hidden_states=hidden_states, logits=logits, input_ids=input_ids, labels=labels, action_chunk=action_chunk, dof_mask=dof_mask, flow=flow, flow_loss_mask=flow_loss_mask, ) # loss = CE(文本) + flow_loss_weight × MSE(动作流场)
generate_flow_action 是闭环控制时每 10 步调用一次的函数。
最值得学的是它的性能设计：序列 = [图文前缀 | 动作后缀]，去噪迭代只改动作 token，
图文前缀的 KV 第一次前向后就缓存复用——5 步去噪只有第 1 步付了全序列的钱：
@torch.no_grad() def generate_flow_action(self, input_ids, action_horizon, action_dim, num_inference_timesteps: int = 10, ...): # 从纯高斯噪声出发 noise = torch.randn(size=(batch_size, action_horizon, action_dim), ...) noisy_action = noise.clone() times = torch.linspace(0.0, 1.0, num_inference_timesteps + 1) # 部署配置 = 5 步 dt = times[1] - times[0] # —— 第一次前向:全序列,顺便建 KV cache —— prefetch_output = self.model(..., use_cache=True) prefix_kv_cache = prefetch_output.past_key_values # 只保留图文前缀的 KV(动作 token 每步都变,不能缓存) for layer_i in range(len(prefix_kv_cache.key_cache)): prefix_kv_cache.key_cache[layer_i] = \ prefix_kv_cache.key_cache[layer_i][:, :, :prefix_length, :] # —— 之后每步:只重算动作后缀,attend 到缓存的前缀 —— def step_with_kvcache(timestep, noisy_action): action_embed, adarms_cond = self.action_preprocessor.step( timestep=timestep, noisy_action=noisy_action, dof_mask=dof_mask) temp_inputs_embeds[action_mask] = action_embed transformer_outputs = self.model(..., past_key_values=prefix_kv_cache, inputs_embeds=temp_inputs_embeds) action_pred = self.action_preprocessor.action_proj_back( transformer_outputs.last_hidden_state[action_mask][:, :2048]) return action_pred.reshape(batch_size, action_horizon, action_dim) # 速度场 v_t # Euler 积分:噪声沿速度场流成动作 action_trajectory = odeint(step_with_kvcache, noisy_action, times[1:], method="euler") predict_action = action_trajectory[-1] predict_action = self.normalizer_action.unnormalize_data(predict_action, ...)
上面全是官方代码。这一段是本项目自己的关键改动（docker 挂载覆盖，不改仓库）： 微调时模型会走捷径——忽略图像、只从本体状态外推动作（官方 issue #58 同款，开环指标完美、 闭环全灭）。解法是在 Dataset 层随机隐藏本体状态，逼模型从图像找答案：
PROPRI_DROPOUT = float(os.environ.get("PROPRI_DROPOUT", "0")) # 训练用 0.5 def _padded_getitem(self, index): r = _orig_getitem(self, index) r["action"] = _pad_nan(r["action"], 7) # 7 维 pad 进 20 维布局前 7 槽 # 以概率 p 隐藏本体状态(全 NaN -> agent_pos_mask=0),对抗捷径学习: # 没有这一条,微调后模型忽略图像只回归 proprio,闭环必败 if PROPRI_DROPOUT > 0 and torch.rand(1).item() < PROPRI_DROPOUT: r["agent_pos"] = torch.full((*r["agent_pos"].shape[:-1], MAX_DOF), float("nan")) else: r["agent_pos"] = _pad_nan(r["agent_pos"], 7) return r PreprocessedDataset.__getitem__ = _padded_getitem # 猴子补丁进数据管线
vla_mixin.py:474-620 —— 动作 token 之间全连接
还是因果，由 causal_action_attention_mask 决定；padding 的行列如何置 False。action_head.py:21-141 Normalizer —— per-数据集 min/max
查表，这就是为什么权重不带 LIBERO 统计就没法直接评估。compute_loss 返回的 channel_loss_dict
把 20 维每维的 loss 分开记——排查"某一维学不动"时非常好用。use_fast_tokenizer=true 时动作走离散
token 自回归（generate），与文本生成共用采样代码。wall_x/infer/ 与 wall_x/serving/ 是已废弃的 legacy
栈（1.1.0 重写），本项目为了跑 LIBERO 闭环打了 5 个 patch 才修通——细节见
主页面的"踩坑记录"。数字与代码均核对自本地权重 config.json / safetensors header / 仓库源码 （commit 97406f2）。生成于 wall-x LIBERO 微调项目，结果展示页 → ../wallx-libero/