LeRobot documentation
LingBot-VA
LingBot-VA
LingBot-VA 是一个基于 Wan2.2 视频扩散技术栈构建的自回归视频-动作世界模型策略。它在同一个自回归序列中交替预测未来的视频潜在变量和机器人动作(“VA” = Video-Action)。LeRobot 集成将 LingBot-VA 接入标准的训练、评估和处理器接口。
模型概览
LingBot-VA 是一种双流“transformer 混合”架构:视频/潜在流(patch_embedding_mlp → blocks → proj_out)和动作流(action_embedder → blocks → action_proj_out)共享相同的 30 个 transformer 块以及文本条件。
| 组件 | 类 | 角色 |
|---|---|---|
| DiT 主干(可训练) | WanTransformer3DModel | 约 50 亿参数的双流 Transformer。 |
| VAE(冻结) | AutoencoderKLWan | Wan2.2 VAE,z_dim=48。从源仓库中按需拉取。 |
| 文本编码器(冻结) | UMT5EncoderModel | UMT5-XXL,d_model=4096。从源仓库中延迟拉取。 |
在推理阶段,策略针对每个块执行一个自回归循环:它使用两个独立的流匹配调度器对视频潜在流(CFG,约 20 步)和动作流(约 50 步)进行去噪,并在块之间维护 KV 缓存。随着块的执行,实际观测到的关键帧会被反馈到 KV 缓存中(闭环世界建模)。
LeRobot 集成涵盖的内容
- 通过 LeRobot 进行标准的
policy.type=lingbot_va配置。 - Hub 上可直接使用的 LeRobot 格式检查点(由上游发布的检查点转换而来)。
- 标准
select_action接口背后的自回归双流推理 (单环境评估,--eval.batch_size=1)。 - 在评估/训练期间,选择保存策略的预测(想象)视频。
- 在 LIBERO 和 RoboTwin 上使用
lerobot-eval进行评估。 - 通过双流流匹配损失(
policy.forward)进行训练 / 微调,详见下文。
安装
- 按照安装指南安装 LeRobot。
- 安装 LingBot-VA 附加组件:
pip install -e ".[lingbot_va]"检查点
已发布的上游检查点已转换为 LeRobot 格式并推送到 Hub:
| 变体 | LeRobot 检查点 |
|---|---|
| LIBERO-Long 后训练 | lerobot/lingbot_va_libero_long |
| RoboTwin 后训练 | lerobot/lingbot_va_robotwin |
| 预训练基础 | lerobot/lingbot_va_base |
LeRobot 的 model.safetensors 中仅存储可训练的 ~5B transformer。冻结的 VAE + UMT5 + tokenizer(约 20 GB)在加载时从 config.wan_pretrained_path 拉取(默认指向源 robbyant/* 仓库)。UMT5-XXL 文本编码器默认在 CPU 上运行(config.text_encoder_device),因此 5B transformer + VAE 可以适配单张 24–32 GB 的 GPU。
评估(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_liberoLingBot-VA 的流式推理(KV 缓存 + 观测关键帧反馈)是为单环境评估实现的;请使用 --eval.batch_size=1。
评估(RoboTwin)
RoboTwin 2.0 需要 SAPIEN + CuRobo 模拟器栈。您可以使用基准测试 Docker 镜像(docker/Dockerfile.benchmark.robotwin,该镜像还需要 warp-lang==1.3.1,并且 CuRobo 需使用 TORCH_CUDA_ARCH_LIST 中指定的 GPU 计算能力进行构建)。RoboTwin 使用 末端执行器位姿控制,因此请通过 --env.action_mode=ee 运行:策略预测每个机械臂的 xyz+quaternion+gripper 增量(robotwin_tshape 潜在布局),这些增量会叠加到该回合的初始末端执行器位姿上,并通过 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 解码,并在环境渲染的 eval_episode_*.mp4 视频旁边写入 pred_episode_*.mp4。
在 lerobot-train 期间的周期性评估中,相同的标志也适用。
实时观看想象中的视频
在运行 lerobot-rollout --display_data=true 时,设置 --policy.return_predicted_video=true,即可解码每个 chunk 的预测视频,并在每个动作执行时显示其所属的想象相机帧(prediction.images.<camera>,位于每个真实相机旁边)。这会导致每个 chunk 执行一次 VAE 解码。
训练 / 微调
LingBotVAPolicy.forward(batch) 实现了论文中的双流 flow-matching 损失
(latent_loss + action_loss,采用时间步加权并应用动作掩码):它将相机片段通过 VAE 编码为视频潜在变量,使用 UMT5 编码任务,对两个流添加噪声,执行 Transformer 的块因果训练前向传播,并返回 (loss, metrics)。优化器预设为 AdamW,采用线性预热后恒定的学习率调度(与上游实现保持一致)。
要求:
- 块因果掩码使用 PyTorch flex-attention,因此训练时请使用
--policy.attn_mode=flex构建策略(默认的torchSDPA 仅用于推理)。 - 完整的 5B DiT 在 AdamW 下无法装入单张 24–32 GB 的 GPU;请使用 LoRA(
--peft.method_type=LORA)和/或优化器卸载进行微调。get_optim_params仅返回可训练(例如适配器)参数;VAE + UMT5 文本编码器保持冻结。安装lerobot[peft]附加组件以启用 PEFT 支持(参见 PEFT 训练指南)。
lerobot-train \
--policy.path=lerobot/lingbot_va_libero_long --policy.attn_mode=flex \
--peft.method_type=LORA --peft.r=32 --peft.lora_alpha=32 \
--peft.target_modules='transformer\.blocks\.\d+\.attn[12]\.(to_q|to_v)' \
--dataset.repo_id=<your LeRobot-format dataset> \
--batch_size=1 --steps=... --output_dir=outputs/train/lingbot_va与 SmolVLA / π₀ 不同,LingBot-VA 未内置默认的 LoRA 目标,因此必须显式传入 --peft.target_modules。只有 self.transformer(双流 Wan transformer)是可训练的;上述示例适配了其自注意力(attn1)和交叉注意力(attn2)块中的 query/value 投影,这是标准的 LoRA 目标集合。如果需要更高容量的适配器,可以扩大范围(例如添加 to_k/to_out,或 ffn 层)。传入 --peft.method_type 即意味着启用 PEFT,因此无需再指定 --policy.use_peft=true。
数据集必须提供相机片段(每个相机的时间窗口,经 VAE 编码为 frame_chunk_size 个潜在帧)以及每个样本的 frame_chunk_size * action_per_frame 个动作步骤。
数据格式(动作通道 & 相机顺序)
LingBot-VA 是一种末端执行器(笛卡尔)位姿策略,它预测的是末端执行器位姿和夹爪状态,而非关节位置。动作位于一个固定的多本体30 维布局中;请将您机器人的动作维度映射到这些通道中,其余通道用0填充(used_action_channel_ids用于选择特定检查点实际使用的通道):
| channels | 含义 |
|---|---|
| 0–6 | 左臂末端执行器位姿 |
| 7–13 | 右臂末端执行器位姿 |
| 14–20 | 左臂关节(已发布的检查点未使用) |
| 21–27 | 右臂关节(已发布的检查点未使用) |
| 28 | 左夹爪 |
| 29 | 右夹爪 |
- LIBERO 使用通道
0–6:6 自由度 EEF 增量(xyz + 旋转)+ 夹爪(单臂)。 - RoboTwin 使用通道
[0–6, 28, 7–13, 29]:左侧 EEF(xyz + 四元数)+ 左侧夹爪 + 右侧 EEF + 右侧夹爪(16 维)。环境通过 CuRobo IK 将这些位姿转换为关节轨迹——从不预测关节。
在微调这些检查点之前,必须将关节空间数据集(或不同的末端执行器约定)重新映射到该架构中。
相机顺序是固定的且对顺序敏感,各相机的潜在变量会按照 obs_cam_keys 的顺序在空间上进行拼接,因此物理相机到槽位的映射必须与训练时保持一致:
| 基准测试 | obs_cam_keys(按顺序) | camera_layout |
|---|---|---|
| LIBERO | observation.images.image(agentview / 第三人称视角),observation.images.image2(眼在手腕视角) | width_concat(在宽度方向上拼接的潜在变量) |
| RoboTwin | observation.images.head_camera, observation.images.left_camera, observation.images.right_camera | robotwin_tshape(下方为全分辨率头部,上方为两个半分辨率手腕) |
第一个摄像头是外部/头部视角,其余的是手腕视角。
推理超参数(LIBERO)
| 密钥 | 值 |
|---|---|
| 高度 × 宽度 | 128 × 128 |
| cameras | observation.images.image(agentview),observation.images.image2(eye-in-hand) |
| 使用的动作通道 | 0–6(7 自由度机械臂 + 夹爪) |
| action_per_frame / frame_chunk_size | 4 / 4 |
| attn_window | 30 |
| 视频 / 动作去噪步骤 | 20 / 50 |
| guidance_scale / action_guidance_scale | 5 / 1 |
| snr_shift / action_snr_shift | 5.0 / 0.05 |
这些是 LingBotVAConfig 的默认值;您可以通过 --policy.<name>=... 覆盖其中任意一项。
备注
- 注意力后端:推理使用
torchSDPA 后端(始终可用)。flashattn和flex后端为可选项;flex仅在训练时需要。 - 模型大小:DiT 约为 50 亿参数,冻结的 VAE+UMT5 额外占用约 20 GB;推理大约需要 18–24 GB 的显存。
许可证
LingBot-VA 采用 Apache-2.0 许可证发布。请参阅 上游仓库。
Update on GitHub