LeRobot 中的 LingBot-VA：基于 Wan2.2 的自回归视频-动作世界模型策略集成指南

LeRobot 中的 LingBot-VA：基于 Wan2.2 的自回归视频-动作世界模型策略集成指南

【免费下载链接】lerobot🤗 LeRobot: Making AI for Robotics more accessible with end-to-end learning项目地址: https://gitcode.com/GitHub_Trending/le/lerobot

LingBot-VA（Video-Action）是一个构建在Wan2.2 视频扩散技术栈之上的自回归视频-动作世界模型策略：它在同一条自回归序列中交错预测未来的视频潜变量与机器人动作。本文以 src/lerobot/policies/lingbot_va/README.md 为核心，结合 LeRobot 仓库中的配置、模型与处理器源码，系统讲解其架构原理、安装方式、Checkpoint 使用、LIBERO/RoboTwin 双基准评测、LoRA 微调、数据格式规范与推理超参数，帮助你直接在 LeRobot 标准select_action/lerobot-eval/lerobot-train接口下复现与二次开发这一世界模型策略。

一、模型总览：双流 "Mixture-of-Transformers"

LingBot-VA 是一个双流 transformer：视频/潜变量流（patch_embedding_mlp → blocks → proj_out）与动作流（action_embedder → blocks → action_proj_out）共享同一组 30 层 transformer block和同一份文本条件（text conditioning）。两条流在训练时以块因果（block-causal）方式交织在一条序列中，在推理时则由两个独立的 flow-matching 调度器分别去噪。

从源码结构看，核心实现位于 modeling_lingbot_va.py 与 utils.py：后者内置了 Wan2.2 模型代码（patch 重组、注意力后端、VAE 归一化与冻结组件加载器、flow-matching 调度器、WanTransformer3DModel及其子模块），前者只保留面向 LeRobot 的LingBotVAPolicy编排器，二者保持单向依赖。

推理时的闭环世界建模：每个 chunk 内，策略先对视频潜变量流做约 20 步的 CFG 去噪，再对动作流做约 50 步去噪；KV cache 跨 chunk 保持。当 chunk 内的动作被执行时，真实观测到的关键帧会被反馈回 KV cache，形成 closed-loop 的世界模型——这是它区别于"开环生成"策略的关键设计。

LeRobot 集成覆盖的能力

• 通过 LeRobot 标准policy.type=lingbot_va配置使用；
• 官方上游 checkpoint 已转换为 LeRobot 格式并发布在 Hub；
• 自回归双流推理封装在标准select_action接口之后（单环境评测，--eval.batch_size=1）；
• 可选择保存策略预测（想象）的视频（--policy.save_predicted_video=true）；
• 支持用lerobot-eval在LIBERO与RoboTwin上评测；
• 通过双流 flow-matching 损失（policy.forward）支持训练/微调。

注册机制上，LingBotVAConfig通过@PreTrainedConfig.register_subclass("lingbot_va")接入 LeRobot 策略工厂（见 configuration_lingbot_va.py），工厂层仅需按类型字符串分发即可，无需改动其他模块。

二、安装

1. 先按照 安装指南 安装 LeRobot 本体；
2. 再安装 LingBot-VA 的扩展依赖：

pip install -e ".[lingbot_va]"

lingbot_vaextra 在 pyproject.toml 中定义为transformers-dep + diffusers-dep + accelerate-dep的组合（对应 UMT5 文本编码器、Wan VAE 与分布式训练所需的三类依赖）。由于模型代码在导入时对diffusers/transformers做惰性导入（见 utils.py 中的TYPE_CHECKING or _diffusers_available分支），未安装这些库时模块仍可被安全导入，但实例化策略会通过require_package("diffusers", extra="lingbot_va")提示安装对应 extra。

三、Checkpoint：瘦身存储 + 冻结组件懒加载

官方上游 checkpoint 已转换为 LeRobot 格式并发布到 Hub：

存储策略：LeRobot 的model.safetensors只保存可训练的约 5B transformer；冻结的 VAE + UMT5 + tokenizer（约 20 GB）在加载时从config.wan_pretrained_path拉取（默认指向源robbyant/*仓库，也支持本地目录）。

这一设计在代码中有清晰体现（modeling_lingbot_va.py）：

• self.transformer是唯一注册进nn.Module的模块，随 checkpoint 往返保存；
• 冻结模块存放在注册表之外的普通字典self._frozen中，既不会写进model.safetensors，也不会被.to()移动，而是由_ensure_frozen_modules()在首次推理时从wan_pretrained_path下的vae/、text_encoder/、tokenizer/子目录懒加载；
• RoboTwin 的 T 形布局会额外加载第二个流式 VAE（streaming_vae_half），用于处理半分辨率的左右腕部相机。

显存规划：UMT5-XXL 文本编码器默认跑在 CPU 上（config.text_encoder_device="cpu"，每个 episode 只编码一次任务描述），从而把约 11 GB 的显存让出来，使 5B transformer + VAE 可以放进单张 24–32 GB GPU。推理总体约需18–24 GB 显存。

四、LIBERO 评测

lerobot-eval \ --policy.path=lerobot/lingbot_va_libero_long \ --policy.device=cuda \ --env.type=libero --env.task=libero_10 \ --env.observation_height=128 --env.observation_width=128 \ --eval.n_episodes=50 --eval.batch_size=1 \ --output_dir=outputs/eval/lingbot_va_libero

LingBot-VA 的流式推理（KV cache + 观测关键帧反馈）目前只实现了单环境评测，因此必须使用--eval.batch_size=1。这是select_action内部维护deque动作队列、按 substep 缓冲关键帧所决定的约束，与上游 LIBERO 客户端循环（evaluation/libero/client.py）保持一致。

五、RoboTwin 评测

RoboTwin 2.0 需要 SAPIEN + CuRobo 仿真器栈。可以直接使用基准 Docker 镜像 docker/Dockerfile.benchmark.robotwin，同时还需要warp-lang==1.3.1，并且 CuRobo 需要按 GPU 计算能力在TORCH_CUDA_ARCH_LIST中编译。

RoboTwin 使用末端执行器位姿（end-effector-pose）控制，因此要加--env.action_mode=ee：策略预测每只手臂的xyz + quaternion + gripper增量（对应robotwin_tshape潜变量布局），由环境把这些增量合成到 episode 初始 EEF 位姿上，再经CuRobo IK转成关节轨迹执行——关节位置从不被直接预测。

lerobot-eval \ --policy.path=lerobot/lingbot_va_robotwin \ --policy.device=cuda \ --env.type=robotwin --env.task=beat_block_hammer --env.action_mode=ee \ --eval.n_episodes=10 --eval.batch_size=1 \ --output_dir=outputs/eval/lingbot_va_robotwin

保存预测（想象）视频

设置--policy.save_predicted_video=true后，策略会把预测的视频潜变量做VAE 解码（decode_predicted_latents，先denormalize_latents反归一化再vae.decode，最终输出[F, H, W, C]的 uint8 帧栈），并写出pred_episode_*.mp4，与仿真渲染的eval_episode_*.mp4并存。lerobot-train周期评测中同样适用。解码后的帧缓存在self.last_predicted_frames/self.last_predicted_latents中供日志与可视化使用。

六、训练与微调

LingBotVAPolicy.forward(batch)实现了论文中的双流 flow-matching 损失：latent_loss + action_loss（按时间步加权、动作掩码）。流程为：VAE 编码相机片段为视频潜变量 → UMT5 编码任务 → 对两条流加噪 → 执行 transformer 的块因果训练前向 → 返回(loss, metrics)。优化器预设为 AdamW（默认lr=1e-5, betas=(0.9, 0.95), eps=1e-8, weight_decay=1e-4, grad_clip_norm=1.0），调度器为线性 warmup 后恒定（ConstantWithWarmupSchedulerConfig，默认 1000 步），与上游train.py一致（见 configuration_lingbot_va.py 的get_optimizer_preset/get_scheduler_preset）。

硬性要求：

• 块因果掩码使用 PyTorchflex-attention，因此训练必须用--policy.attn_mode=flex构建策略（默认torchSDPA 仅供推理）。在 utils.py 的FlexAttnFunc中，flex 后端通过torch.compile编译flex_attention与create_block_mask，并在init_mask里构造 block-causal / window / noise-vs-clean 组合掩码（clean2clean、noise2clean、noise2noise的与/或组合），且要求半精度 dtype；
• 完整的 5B DiT 在 AdamW 下放不进单张 24–32 GB GPU，因此微调需要LoRA（--policy.use_peft=true）和/或优化器 offload。get_optim_params只返回requires_grad的 transformer 参数（启用 LoRA 后自然只剩 adapter 参数），VAE 与 UMT5 文本编码器保持冻结。

lerobot-train \ --policy.path=lerobot/lingbot_va_libero_long --policy.attn_mode=flex \ --policy.use_peft=true \ --dataset.repo_id=<your LeRobot-format dataset> \ --batch_size=1 --steps=... --output_dir=outputs/train/lingbot_va

数据集必须提供每台相机的时序片段（VAE 编码后得到frame_chunk_size个潜变量帧），且每个 item 需要frame_chunk_size * action_per_frame步动作。训练损失的具体计算在_flow_matching_loss与training_loss_from_streams中：latent 与 action 分别采样时间步、加噪、用linear_timesteps_weights加权，action 侧乘以actions_mask只对used_action_channel_ids对应的通道求损失，最终loss = latent_loss + action_loss。

七、数据格式：动作通道与相机顺序

LingBot-VA 是末端执行器（笛卡尔）位姿策略，预测 EEF 位姿 + 夹爪，而不是关节位置。动作固定存放在多具身通用的30 维布局中；你需要把机器人的动作维度映射进这些通道，其余通道填0（used_action_channel_ids用于选定某个 checkpoint 实际使用的通道）：

• LIBERO使用通道0–6：6 自由度 EEF 增量（xyz + 旋转）+ 夹爪（单臂，共 7 维）。
• RoboTwin使用通道[0–6, 28, 7–13, 29]：左 EEF（xyz + 四元数）+ 左夹爪 + 右 EEF + 右夹爪（16 维）。环境经 CuRobo IK 把这些位姿转为关节轨迹。

关节空间数据集（或不同的 EEF 约定）在微调这些 checkpoint 之前，必须先重映射进上述 schema。这一点在_build_training_streams中同样成立：训练批次的动作[B, F*apf, n_used]会被scatter到完整的action_dim=30空间，未被选中的通道全部置零并由actions_mask屏蔽。

相机顺序是固定且敏感的：每台相机的潜变量按obs_cam_keys的顺序做空间拼接，因此物理相机→槽位的映射必须与训练时一致：

第一个相机是外部/头部视角，其余是腕部视角。从源码看，两种布局的实现分别在_encode_frames（逐相机堆叠为[num_cam, C, F, H, W]后一次性流式编码、按宽度 concat）与_encode_frames_tshape（头部全分辨率、左右腕部各半分辨率，腕部先横向拼接再叠在头部潜变量上方，得到约 1.5 倍高的潜变量网格）中。

八、推理超参数（LIBERO 默认）

| 关键项 | 值 | | ------ | -- | | height × width | 128 × 128 | | cameras |observation.images.image（agentview）、observation.images.image2（手眼） | | action channels used | 0–6（7 维：机械臂 + 夹爪） | | action_per_frame / frame_chunk_size | 4 / 4 | | attn_window | 30 | | video / action denoising steps | 20 / 50 | | guidance_scale / action_guidance_scale | 5 / 1 | | snr_shift / action_snr_shift | 5.0 / 0.05 |

以上均为LingBotVAConfig的默认值（configuration_lingbot_va.py），可用--policy.<name>=...覆盖任意一项。值得补充的源码细节：

• chunk_size = frame_chunk_size * action_per_frame（默认 16），即每个自回归 chunk 产生 16 步动作；
• attn_window决定 KV cache 容量：create_empty_cache按(attn_window // 2) * latent_token_per_chunk + (attn_window // 2) * action_token_per_chunk分配槽位，并在槽位耗尽时按 cache id 淘汰最旧的 token；
• 两条去噪流各自使用独立FlowMatchScheduler（视频流shift=5.0、动作流shift=0.05），_infer中视频流逐timesteps步进时把首帧替换为真实观测条件帧，动作流同样用全零条件帧锚定 chunk 起点；
• CFG 通过_repeat_input_for_cfg将输入沿 batch 复制为 2 份（prompt / 空负向 prompt），再按guidance_scale或action_guidance_scale做pred = uncond + scale * (cond - uncond)外推；
• 处理器层面（processor_lingbot_va.py）：预处理管线为rename_observations → add_batch_dim → normalize → to_device（归一化映射全部为 IDENTITY，图像缩放 + VAE 编码在策略内部完成），后处理管线用内置UnnormalizerProcessorStep按QUANTILES模式，用 checkpoint 中恢复的逐通道 q01/q99 把策略输出的[-1, 1]动作映射回物理单位。

九、注意事项

• 注意力后端：推理使用torchSDPA（始终可用）；flashattn与flex为可选，其中flex仅训练需要。WanAttention在构造时按attn_mode选择算子：torch走custom_sdpa，flashattn懒加载flash_attn（优先flash_attn_interface），flex走FlexAttnFunc。
• 模型规模：DiT 约 5B 参数，冻结的 VAE+UMT5 约 20 GB；推理大约需要18–24 GB 显存。
• 归一化设计：策略与 checkpoint 采用 IDENTITY 归一化（图像缩放 + VAE 编码、动作在策略内做分位数归一化/去归一化），动作分位数统计存放在 checkpoint 的policy_postprocessor.json中而非配置里。
• 文本编码：任务描述经clean_prompt（HTML 反转义 + 空白折叠，等价于 diffusers Wan 的prompt_clean去掉ftfy）后送入 UMT5，max_sequence_length=512，训练与推理都会 padding 到定长。

十、License

LingBot-VA 以 Apache-2.0 协议发布（上游为 Robbyant/lingbot-va）。LeRobot 集成代码（configuration / modeling / processor / utils 四件套）采用 Apache-2.0 许可，其中utils.py保留了上游 Robbyant Team 的版权声明，属于从 Wan2.2 / LingBot-VA 上游 vendor 的模型代码，使用与分发时请保留相应版权与许可信息。

【免费下载链接】lerobot🤗 LeRobot: Making AI for Robotics more accessible with end-to-end learning项目地址: https://gitcode.com/GitHub_Trending/le/lerobot