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/lerobotLingBot-VAVideo-Action是一个构建在Wan2.2 视频扩散技术栈之上的自回归视频-动作世界模型策略它在同一条自回归序列中交错预测未来的视频潜变量与机器人动作。本文以 src/lerobot/policies/lingbot_va/README.md 为核心结合 LeRobot 仓库中的配置、模型与处理器源码系统讲解其架构原理、安装方式、Checkpoint 使用、LIBERO/RoboTwin 双基准评测、LoRA 微调、数据格式规范与推理超参数帮助你直接在 LeRobot 标准select_action/lerobot-eval/lerobot-train接口下复现与二次开发这一世界模型策略。一、模型总览双流 Mixture-of-TransformersLingBot-VA 是一个双流 transformer视频/潜变量流patch_embedding_mlp → blocks → proj_out与动作流action_embedder → blocks → action_proj_out共享同一组 30 层 transformer block和同一份文本条件text conditioning。两条流在训练时以块因果block-causal方式交织在一条序列中在推理时则由两个独立的 flow-matching 调度器分别去噪。组件类作用DiT 主干可训练WanTransformer3DModel约 5B 参数的双流 transformerVAE冻结AutoencoderKLWanWan2.2 VAEz_dim48从源仓库懒加载lazy-pulled文本编码器冻结UMT5EncoderModelUMT5-XXLd_model4096同样懒加载从源码结构看核心实现位于 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.typelingbot_va配置使用官方上游 checkpoint 已转换为 LeRobot 格式并发布在 Hub自回归双流推理封装在标准select_action接口之后单环境评测--eval.batch_size1可选择保存策略预测想象的视频--policy.save_predicted_videotrue支持用lerobot-eval在LIBERO与RoboTwin上评测通过双流 flow-matching 损失policy.forward支持训练/微调。注册机制上LingBotVAConfig通过PreTrainedConfig.register_subclass(lingbot_va)接入 LeRobot 策略工厂见 configuration_lingbot_va.py工厂层仅需按类型字符串分发即可无需改动其他模块。二、安装先按照 安装指南 安装 LeRobot 本体再安装 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, extralingbot_va)提示安装对应 extra。三、Checkpoint瘦身存储 冻结组件懒加载官方上游 checkpoint 已转换为 LeRobot 格式并发布到 Hub变体LeRobot CheckpointLIBERO-Long 后训练lerobot/lingbot_va_libero_longRoboTwin 后训练lerobot/lingbot_va_robotwin预训练基础版lerobot/lingbot_va_base存储策略LeRobot 的model.safetensors只保存可训练的约 5B transformer冻结的 VAE UMT5 tokenizer约 20 GB在加载时从config.wan_pretrained_path拉取默认指向源robbyant/*仓库也支持本地目录。这一设计在代码中有清晰体现modeling_lingbot_va.pyself.transformer是唯一注册进nn.Module的模块随 checkpoint 往返保存冻结模块存放在注册表之外的普通字典self._frozen中既不会写进model.safetensors也不会被.to()移动而是由_ensure_frozen_modules()在首次推理时从wan_pretrained_path下的vae/、text_encoder/、tokenizer/子目录懒加载RoboTwin 的 T 形布局会额外加载第二个流式 VAEstreaming_vae_half用于处理半分辨率的左右腕部相机。显存规划UMT5-XXL 文本编码器默认跑在 CPU 上config.text_encoder_devicecpu每个 episode 只编码一次任务描述从而把约 11 GB 的显存让出来使 5B transformer VAE 可以放进单张 24–32 GB GPU。推理总体约需18–24 GB 显存。四、LIBERO 评测lerobot-eval \ --policy.pathlerobot/lingbot_va_libero_long \ --policy.devicecuda \ --env.typelibero --env.tasklibero_10 \ --env.observation_height128 --env.observation_width128 \ --eval.n_episodes50 --eval.batch_size1 \ --output_diroutputs/eval/lingbot_va_liberoLingBot-VA 的流式推理KV cache 观测关键帧反馈目前只实现了单环境评测因此必须使用--eval.batch_size1。这是select_action内部维护deque动作队列、按 substep 缓冲关键帧所决定的约束与上游 LIBERO 客户端循环evaluation/libero/client.py保持一致。五、RoboTwin 评测RoboTwin 2.0 需要 SAPIEN CuRobo 仿真器栈。可以直接使用基准 Docker 镜像 docker/Dockerfile.benchmark.robotwin同时还需要warp-lang1.3.1并且 CuRobo 需要按 GPU 计算能力在TORCH_CUDA_ARCH_LIST中编译。RoboTwin 使用末端执行器位姿end-effector-pose控制因此要加--env.action_modeee策略预测每只手臂的xyz quaternion gripper增量对应robotwin_tshape潜变量布局由环境把这些增量合成到 episode 初始 EEF 位姿上再经CuRobo IK转成关节轨迹执行——关节位置从不被直接预测。lerobot-eval \ --policy.pathlerobot/lingbot_va_robotwin \ --policy.devicecuda \ --env.typerobotwin --env.taskbeat_block_hammer --env.action_modeee \ --eval.n_episodes10 --eval.batch_size1 \ --output_diroutputs/eval/lingbot_va_robotwin保存预测想象视频设置--policy.save_predicted_videotrue后策略会把预测的视频潜变量做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默认lr1e-5, betas(0.9, 0.95), eps1e-8, weight_decay1e-4, grad_clip_norm1.0调度器为线性 warmup 后恒定ConstantWithWarmupSchedulerConfig默认 1000 步与上游train.py一致见 configuration_lingbot_va.py 的get_optimizer_preset/get_scheduler_preset。硬性要求块因果掩码使用 PyTorchflex-attention因此训练必须用--policy.attn_modeflex构建策略默认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_pefttrue和/或优化器 offload。get_optim_params只返回requires_grad的 transformer 参数启用 LoRA 后自然只剩 adapter 参数VAE 与 UMT5 文本编码器保持冻结。lerobot-train \ --policy.pathlerobot/lingbot_va_libero_long --policy.attn_modeflex \ --policy.use_pefttrue \ --dataset.repo_idyour LeRobot-format dataset \ --batch_size1 --steps... --output_diroutputs/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 维布局中你需要把机器人的动作维度映射进这些通道其余通道填0used_action_channel_ids用于选定某个 checkpoint 实际使用的通道通道含义0–6左臂末端执行器位姿7–13右臂末端执行器位姿14–20左臂关节发布版 checkpoint 未使用21–27右臂关节发布版 checkpoint 未使用28左夹爪29右夹爪LIBERO使用通道0–66 自由度 EEF 增量xyz 旋转 夹爪单臂共 7 维。RoboTwin使用通道[0–6, 28, 7–13, 29]左 EEFxyz 四元数 左夹爪 右 EEF 右夹爪16 维。环境经 CuRobo IK 把这些位姿转为关节轨迹。关节空间数据集或不同的 EEF 约定在微调这些 checkpoint 之前必须先重映射进上述 schema。这一点在_build_training_streams中同样成立训练批次的动作[B, F*apf, n_used]会被scatter到完整的action_dim30空间未被选中的通道全部置零并由actions_mask屏蔽。相机顺序是固定且敏感的每台相机的潜变量按obs_cam_keys的顺序做空间拼接因此物理相机→槽位的映射必须与训练时一致benchmarkobs_cam_keys按顺序camera_layoutLIBEROobservation.images.imageagentview / 第三人称、observation.images.image2手眼/腕部width_concat潜变量沿宽度拼接RoboTwinobservation.images.head_camera、observation.images.left_camera、observation.images.right_camerarobotwin_tshape全分辨率头部在下两个半分辨率腕部在上第一个相机是外部/头部视角其余是腕部视角。从源码看两种布局的实现分别在_encode_frames逐相机堆叠为[num_cam, C, F, H, W]后一次性流式编码、按宽度 concat与_encode_frames_tshape头部全分辨率、左右腕部各半分辨率腕部先横向拼接再叠在头部潜变量上方得到约 1.5 倍高的潜变量网格中。八、推理超参数LIBERO 默认| 关键项 | 值 | | ------ | -- | | height × width | 128 × 128 | | cameras |observation.images.imageagentview、observation.images.image2手眼 | | action channels used | 0–67 维机械臂 夹爪 | | 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视频流shift5.0、动作流shift0.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_sdpaflashattn懒加载flash_attn优先flash_attn_interfaceflex走FlexAttnFunc。模型规模DiT 约 5B 参数冻结的 VAEUMT5 约 20 GB推理大约需要18–24 GB 显存。归一化设计策略与 checkpoint 采用 IDENTITY 归一化图像缩放 VAE 编码、动作在策略内做分位数归一化/去归一化动作分位数统计存放在 checkpoint 的policy_postprocessor.json中而非配置里。文本编码任务描述经clean_promptHTML 反转义 空白折叠等价于 diffusers Wan 的prompt_clean去掉ftfy后送入 UMT5max_sequence_length512训练与推理都会 padding 到定长。十、LicenseLingBot-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创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考