lingbot-vla-2.0代码框架解析
·
LingBot-VLA 2.0 端到端流程解析（输入 → MoE → 输出）
（本文档基本由ai生成，主要记录学习轨迹）
本文档基于仓库源码静态梳理，覆盖数据输入、模型前向（含 MoE 混合专家）、损失计算、训练编排、推理部署的完整链路。
所有结论均可按文末给出的文件/行号回溯。
目录
- 一句话总览
- 端到端闭环图
- 仓库文件职责地图（推荐阅读顺序）
- 核心设计概念
- 输入侧：数据管线
- 模型侧：双流架构与前向计算
- MoE 混合专家详解
- 流匹配（Flow Matching）动作生成
- 双查询蒸馏（Depth + Future Video）
- 损失函数全景
- 训练编排
- 输出侧：推理与部署
- 关键超参与配置对照
- 张量规格速查表
1. 一句话总览
LingBot-VLA 2.0 以 Qwen3-VL-4B 为感知/语言主干（VLM 流），并联一个 36 层 Qwen2 结构的 Action Expert（动作流，部分层为 32 专家 top-4 的稀疏 MoE）；图像/语言/状态/带噪动作块经双流联合注意力融合后，用流匹配（Flow Matching）迭代去噪，一次预测未来 50 步、55 维统一空间的机器人动作；训练时由冻结的 LingBot-Depth 与 DINO-Video 教师对"当前/未来感知查询 token"做特征蒸馏。
2. 端到端闭环图
3. 仓库文件职责地图（推荐阅读顺序）
4. 核心设计概念
4.1 统一 55 维动作/状态空间
异构机器人本体被映射到同一套 canonical 向量（README.md 模型设计节）：
| 分段 | 维度 | 含义 |
|---|---|---|
| arm.position | 14 | 双臂关节位置（单臂 7×2） |
| end.position | 14 | 双臂末端位姿（位置+旋转） |
| effector.position | 2 | 双夹爪 |
| hand.position | 12 | 灵巧手 |
| waist.position | 4 | 腰部 |
| head.position | 2 | 头部 |
| mobility | 3 | 移动底盘 |
| reserved | 4 | 预留 |
| 合计 | 55 | action_dim = max_action_dim = max_state_dim = 55 |
本体实际不存在的关节用 mask 屏蔽（joint_mask / state_joint_mask），padding 维度不计损失。
4.2 双流（Prefix / Suffix）架构
- VLM 流（Prefix）：Qwen3-VL-4B，36 层 Decoder，hidden 2560。承载图像 token、蒸馏 query token、语言 token，负责"看懂场景与指令"。
- Action Expert 流（Suffix）：Qwen2 结构，36 层，hidden 768（默认 16 Q-heads / 2 KV-heads，head_dim 128 → Q/K 投影宽度 2048，GQA；稠密 FFN intermediate 2752）。承载 1 个 state token + 50 个带噪动作 token，负责"输出动作"。头数可在配置中覆盖（
QwenvlWithExpertV2Config签名默认 32/8，实际取值来自LingbotVLAV2Config默认 16/2 或 YAML 显式设置）。 - 两条流在每一层做一次联合自注意力（KV 拼接），但注意力后的 FFN 各自独立（Expert 侧为 MoE），从而以 768 的小宽度做高频动作推理，同时完整复用 4B VLM 的语义先验。
4.3 动作块（Action Chunk）与流匹配
- 一次前向预测 50 步动作（
chunk_size = n_action_steps = 50）。 - 动作不是回归值，而是用 Flow Matching 从纯噪声经 10 步 Euler 积分得到（见第 8 节）。
4.4 双查询蒸馏（Dual-Query Distillation）
在 Prefix 中插入可学习的"感知查询 token"：当前深度 query、未来深度 query、未来视频 query。它们在联合注意力中向 VLM 提取几何/时序信息，再由轻量对齐头回归冻结教师的特征，迫使模型同时理解当前场景几何与未来场景演化。
5. 输入侧：数据管线
5.1 数据源与时间展开
入口 build_vla_dataset（lingbotvla/data/dataset.py）→ VLADataset（base_dataset.py:152）。
- 数据格式：LeRobot v2.1 / v3.0。
- 动作展开（
get_delta_timestamps，base_dataset.py:231）：对每个动作特征取当前帧起0/fps … 49/fps共 50 帧，形成动作块。 - 未来帧图像（
get_video_delta_timestamps，base_dataset.py:242）：开启use_future_image时，每个相机取[0, 49/fps]两帧，供深度/视频教师使用。 - 多数据集：
data_name: multi时用MultiVLADataset，按assets/training_data/robotwin.txt中每行<robot_config名> <LeRobot路径>拼接，并给样本注入rep_id用于分组统计 loss。
5.2 Robot Config：原始特征 → 统一空间
configs/robot_configs/robotwin.yaml 示例：
states:
- observation.state.arm.position: # 统一字段名
origin_keys: # 从原始 14 维 state 切片
- observation.state: {start: 0, end: 6}
- observation.state: {start: 7, end: 13} # 跳过第 6/13 维夹爪
- observation.state.effector.position:
origin_keys:
- observation.state: {start: 6, end: 7}
- observation.state: {start: 13, end: 14}
images:
- observation.images.camera_top: {origin_keys: observation.images.cam_high}
...
norm_stats: assets/norm_stats/robotwin.json
origin_keys支持切片、拼接、跨特征组合；subtract_state: true可把绝对动作转成相对状态的增量动作（末端位姿支持四元数相对变换）。- 顶层 YAML 中
data.joints/data.cameras/data.norm_type声明关节组顺序、相机顺序、每组归一化方式（robotwin 配置：14 arm + 14 end + 2 effector，三个相机，均bounds_99_woclip）。
5.3 FeatureTransform：单样本转换全流程
核心方法 FeatureTransform.apply（vla_data/utils.py:375）：
归一化统计量由 scripts/compute_norm_stats.py 预计算（RoboTwin 直接提供 assets/norm_stats/robotwin.json）；推理时用同一 FeatureTransform.unapply 做反归一化还原成真实关节量。
5.4 Collator 与 DataLoader
VLADataCollatorWithPacking（multimodal/data_collator.py:317）：对state/images/img_masks/lang_tokens/lang_masks/actions/joint_mask/state_joint_mask/image_grid_thw/…沿 batch 维cat，其余走default_collate。build_dataloader（data_loader.py:143）支持梯度累加（一个优化步吐出多个 micro-batch）与基于 token 预算的动态批大小（batching_strategy.py + dynamic_batching.py，含 warmup 与状态恢复）。
5.5 进入模型的字段
| 字段 | 形状（B=micro batch） | 说明 |
|---|---|---|
images | (B, N_cam, C, H, W) | N_cam=3（顶视+双腕），训练增强后输入 ViT |
image_grid_thw | (B, N_cam, 3) | Qwen3-VL 的 temporal/height/width 网格 |
img_masks | (B, N_cam) | 有效相机 mask |
state | (B, 55) | 当前状态（归一化、padding 后） |
actions | (B, 50, 55) | 未来 50 步动作真值（归一化） |
lang_tokens / lang_masks | (B, L) | 语言指令 token（上限 72） |
joint_mask | (B, 50, 55) | 动作损失掩码（屏蔽 pad/无效关节） |
state_joint_mask / action_joint_mask | (B, 55) | 关节有效性 |
pil_images / future_pil_images | List[PIL] | 不进模型，专供教师生成伪标签 |
6. 模型侧：双流架构与前向计算
顶层包装 LingbotVlaV2Policy（modeling_lingbot_vla_v2.py:1198）→ 内部 FlowMatchingV2（L429）→ QwenvlWithExpertV2Model（L121）。
6.1 模块构成
LingbotVlaV2Policy
└── FlowMatchingV2
├── state_proj Linear(55 → 768)
├── action_in_proj Linear(55 → 768)
├── action_out_proj Linear(768 → 55)
├── action_time_mlp_in/out 时间步注入 MLP（768×2 → 768）
├── depth/video align heads 蒸馏对齐头 + 可学习 query embeddings
└── qwenvl_with_expert QwenvlWithExpertV2Model
├── qwenvl Qwen3-VL-4B（ViT 视觉塔 + 36 层语言 Decoder, hidden 2560）
└── qwen_expert Qwen2 结构 36 层（hidden 768），指定层 MLP = Qwen2TokenMoeBlock
动作专家层数硬编码与 VLM 对齐为 36 层（modeling_lingbot_vla_v2.py:85），前向中有 assert action_num_layers == num_layers。MoE 层在 _install_moe_blocks（L170）按配置 token_moe_layers 逐层替换 MLP。
6.2 Prefix 构建（条件侧）
embed_prefix（L499）：
- 图像经 Qwen3-VL ViT 得到 patch embedding 与 deepstack 多层视觉特征（训练中可逐层回注语言 Decoder）；每张图外包
vision_start/end特殊 token。 - 语言 token 经 VLM embedding 表。
- 开启蒸馏时，在序列中插入 query 段（由
prefix_query_segments决定顺序）：
图像 → 当前深度 query(8 token) → 语言 → 未来深度 query(8) → 未来视频 query(8) …
query 是可学习参数（depth_align_embs等 256×2560，按num_task_tokens=8聚合），深度/视频共享 query 时还会过共享投影融合。 - 计算 Qwen3-VL 的 M-RoPE 多维位置编码。
6.3 Suffix 构建（动作侧）
embed_suffix（modeling_lingbot_vla.py:1126）：
state→state_proj→ 1 个 state token；- 带噪动作块
x_t (B,50,55)→action_in_proj；时间t经正弦位置编码后与动作 embedding 拼接 → SiLU MLP 融合 → 50 个 action-time token； - 序列：
[state, a₁, …, a₅₀]； - 注意力掩码语义：Prefix 看不到 Suffix；state 看不到动作；50 个动作 token 之间双向可见（动作块内部联合推理）。
6.4 逐层双流联合注意力（核心创新）
QwenvlWithExpertV2Model.forward（L311-L414）每层执行两阶段：
- 推理时 Prefix 只前向一次并缓存 KV（
fill_kv_cache=True），10 步去噪仅重复 Suffix 前向。 - Action Expert 的 RMSNorm 被替换为 AdaRMSNorm（
adanorm_time: true），让时间步t调制归一化，实现扩散式时间条件。 - 注意力后端支持
flex_cached（训练默认，预构建 BlockMask）/flex/eager（推理用）。
7. MoE 混合专家详解
MoE 只位于 Action Expert 的 FFN 位置（VLM 主干保持稠密），实现类 Qwen2TokenMoeBlock（qwen2_action_expert.py:222）。
7.1 前向路由
对每个动作 token 的隐状态 x (B*T, 768)：
- Gate 打分（强制 FP32）：
router_logits = Linear(x)→ (N, 32)。FP32 是为避免 bf16 下近邻分数翻转导致专家路由抖动。 - 激活函数：后训练配置
router_activation: sigmoid（各专家独立打分；默认也支持 softmax）。 - Loss-free 修正偏置：
scores = sigmoid(logits) + e_score_correction_bias，再 Top-4（token_top_k: 4）。 - 权重：取选中专家分数 → 归一化 → ×
routed_scaling_factor（后训练为 4.0，DeepSeek-V3 风格补偿）。 - 专家计算：
fused（默认）：32 个专家的权重存为 3D 张量（gate/up/down_proj: [E, I, H]，L113），训练走 fused_moe.py 的 GroupGEMM 融合 kernel；无梯度推理走更快的 robby_moe.py grouped kernel，失败自动回退。eager：所有专家全量前向 + one-hot 加权（调试用）。
- 共享专家：
shared_expert（FFN intermediate 704）处理全部 token；use_shared_expert_gate: false时直接与路由输出相加（后训练配置），为 true 时经 sigmoid 门控加权。 - 路由专家内部为标准 SwiGLU：
down(silu(gate(x)) * up(x))，MoE 专家 intermediate 512（稠密层 2752）。
每个 token 激活计算量 ≈ 4 个窄专家（512）+ 1 个共享专家（704），但可学习参数量覆盖 32 个专家，实现"固定激活算力、扩大容量"。
7.2 专家并行（EP）
lingbotvla/distributed/moe/moe_layer.py：按路由结果统计每专家 token 数 → all-to-all 派发 → EPGroupGemm 分组矩阵乘 → 按路由权重还原；FSDP2 下专家参数按 Shard(0) 切分（见 get_parallel_plan，L1204）。
7.3 三套负载均衡机制（可组合）
| 机制 | 代码位置 | 后训练取值 | 作用 |
|---|---|---|---|
| Loss-free bias（DeepSeek-V3 风格，无辅助损失） | moe_load_balance.py | bias_update_speed: 0（关闭，仅保留监控） | optimizer.step 前钩子：all-reduce 全局专家负载，bias -= speed·sign(load-mean)，过载专家减分、欠载加分；支持 bias_centering 防漂移与 bias_update_interval 累积更新 |
| Sequence-wise 辅助损失 | triton_sequence_wise_balance_loss（Triton kernel），调用见 modeling_lingbot_vla_v2.py:1084 | sequence_wise_loss_coeff: 1e-3，mode: per_sequence | 每个样本各自的动作 token 序列内均衡专家选择（DeepSeek-V3 意图）；global 则整批统一 |
| Router z-loss | L1077 | router_z_loss_coeff: 1e-4 | 惩罚 logits 绝对值 logsumexp²，稳定路由数值 |
预训练可用纯 loss-free（
bias_update_speed=0.00025并关掉两个辅助损失）；后训练采用辅助损失路线。两种模式在 README.md:154 有说明。
7.4 MoE 监控
模型每步产出 moe_summary/*（maxvio/minvio/entropy/topk_sigmoid/bias/死专家/负载CV），每 moe_monitor_interval（后训练 1000 步）额外写逐层指标与专家负载精确柱状图（训练脚本第 13 节 TB 写入逻辑）。环境变量 GATE_GRAD_PROBE=1 可在前 12 步打印 gate 梯度范数诊断塌缩。
7.5 MoE 专家学习率缩放
get_moe_param_groups（train_lingbotvla.py:75）：路由专家权重使用 lr × sqrt(num_experts/top_k) = lr × sqrt(32/4) ≈ 2.83× 的学习率，gate / 共享专家 / 稠密层保持基础 LR（后训练 use_moe_expert_lr: true）。
8. 流匹配（Flow Matching）动作生成
8.1 训练（modeling_lingbot_vla_v2.py:767）
- 采样时间：
t ~ Beta(1.5, 1)，线性映射到[0.001, 1]（modeling_lingbot_vla.py:1055）。 - 线性插值路径：
x_t = t·ε + (1−t)·a₀，其中ε ~ N(0,I)，a₀为真值动作块；目标速度场u_t = ε − a₀。 x_t与t、state组成 Suffix，与 Prefix 联合前向，取最后 50 个动作 token 经action_out_proj得预测速度v_t (B,50,55)。- 损失：
fm→ MSE；L1_fm→ L1（后训练用 L1_fm），并按joint_mask屏蔽无效关节/padding。
8.2 推理采样（modeling_lingbot_vla_v2.py:920）
- Prefix 前向一次，缓存 36 层 KV；
x_t = 噪声，t = 1，步长dt = −1/10；- 循环 10 步：
v_t = predict_velocity(...)（仅 Suffix 前向，复用缓存），x_t += dt·v_t，t += dt； t=0时x_t即去噪完成的动作块（归一化空间）。
predict_velocity 支持 torch.compile（部署时 --use_compile）。官方实测 RTX 4090D 单次推理约 130 ms（10 步）。
9. 双查询蒸馏（Depth + Future Video）
9.1 冻结教师（训练脚本第 3 节构建）
- MoGe-2（单目几何）+ MoRGBD / LingBot-Depth（
build_depth_model，module_utils.py:71）：对当前帧与未来帧输出深度/几何特征伪标签（get_depth_target，bf16、no_grad）。 - DINO-Video（
build_video_model，dino_video/teacher.py）：以[warmup当前帧, 当前帧, 未来帧]三帧输入，输出未来帧 patch 特征、CLS 特征及当前帧 patch 特征（get_video_target）。
9.2 学生对齐头
- 可学习 query：
depth_align_embs（256×2560，256 backbone token，按 8 个 task token 聚合）；未来深度与未来视频各有独立/共享 query。 TaskTokenDepthHead（align_heads/depth_head.py）是 resampler 式小网络（1 层 4 头，dim_head 32）：把"图像 token + query token"作为上下文，迭代读出与教师同维（1024）的特征。- 前向：
depth_emb_forward（modeling_lingbot_vla.py:1418）、video_emb_forward（L1445）；视频损失支持 patch MSE、cosine loss（0.2 权重）、CLS loss 开关组合；未来分支可detach当前图像特征以阻断梯度。 - 注意力掩码支持
block_future_depth_to_action/block_suffix_to_future_video，控制未来 query 是否向动作 token 泄漏未来信息。
后训练权重：depth 0.004、future_depth 0.004、future_video 0.004，前 5000 步为视觉 warmup（visual_steps: 5000）。
10. 损失函数全景
训练返回的总损失（LingbotVlaV2Policy.forward:1305）：
total_loss = vla_flowmatch_loss # L1/MSE 速度场回归，joint_mask 屏蔽
+ depth_loss_weight × depth # 当前深度蒸馏 (0.004)
+ future_depth_loss_weight × future_depth # 未来深度蒸馏 (0.004)
+ future_video_loss_weight × future_video # 未来视频 patch/CLS (0.004)
+ sequence_wise_loss_coeff × seq_wise # MoE 序列内均衡 (1e-3)
+ router_z_loss_coeff × router_z # MoE router z-loss (1e-4)
(+ loss-free bias 仅改路由偏置，不进 loss)
- 多 micro-batch 梯度累加时各项 loss 先除以累加份数；反向后在 DP 组 all-reduce 平均，日志中再除以权重还原"裸损失"。
- 逐样本 loss（
batch_mean_losses）按数据集名（rep_id）分组写 TensorBoard，便于观察多机器人数据的不均衡。
11. 训练编排
train_lingbotvla.py 共 15 个编号阶段：
- 配置与分布式：YAML + CLI 覆盖解析（
Arguments = model/data/train/eval）；init_process_group(nccl)+init_parallel_state（DP / HSDP replicate-shard / TP / EP / PP / CP / Ulysses）。 - 建模型：
LingbotVLAV2Config→build_foundation_model（可从预训练 hf_ckpt 加载，bf16 或 fp32-mixed）。 - 建教师：按
align_params加载冻结深度/视频教师（可选 compile）。 - 建数据：
build_processor（含 tokenizer）→ collator →build_vla_dataset→build_dataloader。 - 并行包装：
build_parallelize_model（后训练fsdp2+ HSDPenable_full_shard:false+vlm_fsdp:true+ 按 decoder layer 分片；可冻结 ViT、梯度检查点、激活 offload）。 - 优化器：默认 Muon（矩阵参数 Muon + 1D/embedding AdamW，optim/muon.py），可回退 AdamW；注入 MoE 专家 LR 分组。
- MoE 钩子：
build_moe_load_balance_hook注册到optimizer.step之前。 - LR 调度：warmup + cosine，
lr=1e-4 → lr_min=5e-5。 - 日志：rank0 的异步 TensorBoard（
AsyncTBWriter）+ 可选 W&B；profiler；首步落盘 model assets（config/processor）。 - 续训：扫描
checkpoints/global_step_<N>，从新到旧尝试恢复模型/优化器/调度器/DataLoader 位置/RNG。
11-13. 主循环：取 micro-batches → 教师伪标签 → 学生前向 → 累加反向 → 梯度裁剪（稳定期后改用更严的decayed_max_grad_norm）→optimizer.step（触发 MoE 钩子）→ 指标 all-reduce → 进度条/TB/W&B（含 MoE 专家负载图）。
14-15. 保存：每save_steps（10000）用 DCP 存分布式 ckpt，并由AsyncHFCheckpointSaver后台异步转换 HF 格式（hf_ckpt/，推理必须指向它）；epoch 末/达到 max_steps（50000）同样落盘。
启动方式：
bash train.sh tasks/vla/train_lingbotvla.py ./configs/vla/robotwin/robotwin.yaml \
--data.train_path assets/training_data/robotwin.txt \
--data.data_name multi --train.output_dir output/
12. 输出侧：推理与部署
deploy/lingbot_vla_v2_policy.py 的 LingbotVLAv2Server：
- 加载：读训练时落盘的
lingbotvla_cli.yaml重建LingbotVLAV2Config，从$QWEN3VL_PATH加载 Qwen3-VL-4B 基座，加载hf_ckpt/*.safetensors；构建FeatureTransform（robot config + norm stats）。 - 观测预处理（
_prepare_model_input）：图像 resize 256 → numpy 转 tensor →FeatureTransform.apply(policy_eval=True)（状态归一化、语言 tokenize、Qwen3-VL 图像网格）。 - 采样：
sample_actions（第 8.2 节）→ 归一化空间的(B, 50, 55)动作块。 - 反归一化：
unapply还原成机器人原始动作字段（关节/夹爪）。 - Chunk 执行策略：
use_length=25：每 25 个环境步重规划一次，期间顺序消费缓存动作块；chunk_ret=true：直接返回整 chunk（由调用方处理）。
- 服务化：
WebsocketPolicyServer对外提供动作服务，支持 batch（多仿真任务/GPU 常驻多实例）。
python -m deploy.lingbot_vla_v2_policy \
--model_path .../checkpoints/global_step_50000/hf_ckpt/ \
--use_compile --use_length 25 --port 8006
闭环评估：开环脚本 scripts/open_loop_eval.py；RoboTwin 50 任务见 experiment/robotwin/。注意官方发布指标基于 FP32 推理，BF16 显存更省但成功率可能有明显差异。
13. 关键超参与配置对照
以 configs/vla/robotwin/robotwin.yaml 为准：
| 类别 | 参数 | 值 | 含义 |
|---|---|---|---|
| 动作 | chunk_size / action_dim | 50 / 55 | 动作块长度 / 统一动作维度 |
| 动作 | loss_type | L1_fm | 流匹配 L1 速度回归 |
| 专家 | expert_hidden_size | 768 | Action Expert 宽度 |
| MoE | token_moe_layers | 0–35 全部 36 层 | 启用 MoE 的层 |
| MoE | token_num_experts / top_k | 32 / 4 | 专家数 / 每 token 激活数 |
| MoE | moe_intermediate / shared_intermediate | 512 / 704 | 路由专家 / 共享专家 FFN 宽度 |
| MoE | router_activation / scaling | sigmoid / 4.0 | 路由打分与权重补偿 |
| MoE | use_shared_expert_gate | false | 共享专家直接相加 |
| MoE | 均衡 | seq-wise 1e-3 + z 1e-4，bias speed 0 | 后训练辅助损失路线 |
| 蒸馏 | query token / backbone token | 8 / 256 | 感知查询粒度 |
| 蒸馏 | depth / future_depth / video 权重 | 0.004 / 0.004 / 0.004 | 三项蒸馏强度 |
| 蒸馏 | num_future_frames / use_warmup_frame | 1 / true | DINO-Video 三帧输入 |
| 训练 | micro/global batch | 32 / 1024（32 卡） | 梯度累加扩全局批 |
| 训练 | lr / lr_min / schedule | 1e-4 / 5e-5 / cosine | 学习率 |
| 训练 | max_steps / save_steps | 50000 / 10000 | 训练与保存 |
| 训练 | optimizer / enable_fp32 | muon / true | 优化器与动作头精度 |
| 训练 | attention_implementation | flex_cached | 块掩码缓存注意力 |
| 推理 | num_steps（config 内置） | 10 | Euler 去噪步数 |
14. 张量规格速查表
| 阶段 | 张量 | 形状 | 备注 |
|---|---|---|---|
| 数据 | 输入图像 | (B, 3, 3, 256, 256) | 3 相机 |
| 数据 | ViT patch token | (B·3, L_patch, 2560) | L_patch 由 grid_thw 决定 |
| 数据 | state / actions | (B,55) / (B,50,55) | 归一化后 |
| Prefix | 前缀隐状态 | (B, L_prefix, 2560) | 图像+query+语言 |
| Suffix | 后缀隐状态 | (B, 1+50, 768) | state + 动作块 |
| MoE | gate logits | (B·51, 32) | FP32 打分 |
| MoE | 专家权重 | 32×[512,768] 等 3D 存储 | fused GroupGEMM |
| 输出 | v_t / 最终动作 | (B, 50, 55) | 10 步积分 |
| 部署 | 反归一化后 | dict{各原始动作键: (50, d)} | 直接下发机器人 |
本文档为源码静态分析快照（2026-09），如后续代码调整请以对应文件最新实现为准。
有“AI”的1024 = 2048，欢迎大家加入2048 AI社区
更多推荐
0
0- 0
已为社区贡献1条内容
所有评论(0)