FastWAM

FastWAM 是一个用于机器人控制的世界 action 模型（World Action Model）policy。LeRobot 集成通过标准 policy API 暴露 FastWAM，因此可以用 policy.type=fastwam 配置、用 lerobot-train 训练，并通过 LeRobot 预训练 policy 接口加载。

模型概览

FastWAM 在训练期间保留视频建模，但在 inference 时直接预测 action，而不是迭代生成未来 observation。这个 LeRobot policy 封装了 FastWAM action 模型，将 LeRobot 批次适配为 FastWAM 训练样本，并提供用于归一化和 action 后处理的标准处理器流水线。

该实现默认从 Wan-AI/Wan2.2-TI2V-5B 初始化视觉世界模型组件，并预测形状为 [batch, action_horizon, action_dim] 的 action chunk。

LeRobot 集成涵盖的内容

• 通过 LeRobot 进行标准 policy.type=fastwam 配置
• 图像、state、action 和语言任务批次适配
• 通过 select_action 和 predict_action_chunk 进行 action chunk inference
• 通过 LeRobot policy API 保存/加载 checkpoint
• 可配置的 LIBERO gripper action 后处理

安装要求

从源码安装 LeRobot，然后安装 FastWAM 依赖：

pip install -e ".[fastwam]"

这会从 pyproject.toml 安装 FastWAM policy 可选项：transformers、 diffusers、ftfy 和 regex，以及 LeRobot 的基础依赖。

对于 LIBERO 评估，还需安装 benchmark 依赖：

pip install -e ".[fastwam,libero]"

这会安装两个可选项。除上述 FastWAM 依赖外， libero 可选项还会安装 LeRobot dataset 依赖、Linux 上的 hf-libero，以及 scipy。

FastWAM 使用 Wan2.2 TI2V 主干。默认模型 id 为：

policy.model_id=Wan-AI/Wan2.2-TI2V-5B

数据要求

FastWAM 期望一个 LeRobot dataset，包含：

• 一个或多个视觉 observation，其宽度拼接后等于 policy.image_size[1]
• 当 policy.proprio_dim 不为 None 时的 observation.state
• action
• 通过 dataset task 字段提供的语言任务指令，或预计算的 context 和 context_mask 张量

默认视觉设置为一个名为 observation.images.image、形状为 (3, 224, 448) 的图像特征。如果 dataset 使用两个相机，请配置 policy.input_features，使其高度都匹配 224，宽度之和为 448。

用法

使用以下命令创建新的 FastWAM policy：

lerobot-train \
  --dataset.repo_id=your-org/your-dataset \
  --policy.type=fastwam \
  --policy.action_dim=7 \
  --policy.proprio_dim=8 \
  --policy.action_horizon=32 \
  --policy.n_action_steps=10 \
  --policy.image_size='[224,448]' \
  --output_dir=./outputs/fastwam_training \
  --job_name=fastwam_training \
  --steps=300000 \
  --batch_size=8 \
  --policy.device=cuda

使用以下命令在 LIBERO-10 上评估现有的 LeRobot 格式 checkpoint：

lerobot-eval \
  --policy.path=ZibinDong/fastwam_libero_uncond_2cam224 \
  --policy.device=cuda \
  --policy.torch_dtype=float32 \
  --policy.n_action_steps=10 \
  --env.type=libero \
  --env.task=libero_10 \
  --env.observation_height=224 \
  --env.observation_width=224 \
  --eval.batch_size=1 \
  --eval.n_episodes=50 \
  --seed=0 \
  --env.episode_length=600

对于 libero_goal、libero_spatial 和 libero_object，请使用 --env.episode_length=300。

对于真实机器人 rollout，使用相同的 checkpoint 路径：

lerobot-rollout \
  --robot.type=so101_follower \
  --robot.port=/dev/ttyACM0 \
  --policy.path=your-org/fastwam-real-robot

配置说明

图像特征

policy.image_size 是拼接后的 FastWAM 图像张量的尺寸，形如 (height, width)。每个配置的图像特征必须具有形状 (3, height, camera_width)，且所有相机宽度之和必须等于所配置的宽度。

action chunking

policy.action_horizon 控制训练期间监督、inference 期间预测的未来 action 数量。policy.n_action_steps 控制 policy 预测新 action chunk 之前会消费多少个 action。policy.n_action_steps 必须小于或等于 policy.action_horizon。

Wan 组件

FastWAM 从配置的 Wan 模型目录或 Hugging Face Hub 模型 id 加载 Wan VAE、视频 DiT、文本编码器和分词器。由 save_pretrained 保存的 LeRobot 格式 FastWAM checkpoint 还会复制 from_pretrained 所需的本地 Wan 组件文件。

注意力后端

FastWAM 的 DiT 对所有注意力都使用 PyTorch 的 scaled_dot_product_attention（SDPA）。它不使用 FlashAttention：其 Transformer 混合（MoT）路由需要任意的布尔 [query, key] 注意力掩码，而 FlashAttention 的 varlen API 无法表达。因此安装 flash-attn 包对 FastWAM 路径没有影响。（注意，SDPA 本身仍可能在内部选择 PyTorch 自带的 flash / 内存高效 / math 内核——这与 flash-attn 包无关。）

编译后的 action inference

FastWAM 可以编译其视频预填充和缓存的 action 去噪计算图，以减少 重复 inference 步骤中的 Python 和 CUDA 内核启动开销：

--policy.compile_action_infer=true

该选项默认禁用。首次 action prediction 会针对当前输入形状编译计算图，因此较慢； 后续相同形状的预测会复用它。形状变化可能触发重新编译。 在不支持的 PyTorch 或加速器配置上，仍可使用 eager 路径。

LIBERO action 开关

FastWAM LIBERO checkpoint 默认使用 policy.toggle_action_dimensions=[-1]， 以匹配原始 FastWAM 评估流水线所使用的 gripper action 约定：

--policy.toggle_action_dimensions='[-1]'

结果

使用 ZibinDong/fastwam_libero_uncond_2cam224 在 LIBERO 上评估：

复现：lerobot-eval --policy.path=ZibinDong/fastwam_libero_uncond_2cam224 --policy.device=cuda --policy.torch_dtype=float32 --policy.n_action_steps=10 --env.type=libero --env.task=libero_spatial --env.observation_height=256 --env.observation_width=256 --eval.batch_size=1 --eval.n_episodes=50 --seed=0 --env.episode_length=300（1x H20 140 GB）。

参考资料

• Fast-WAM 论文
• Fast-WAM 项目主页
• Fast-WAM 代码
• 发布的上游 checkpoint
• Wan2.2 TI2V 5B

引用

@article{yuan2026fastwam,
  title = {Fast-WAM: Do World Action Models Need Test-time Future Imagination?},
  author = {Tianyuan Yuan and Zibin Dong and Yicheng Liu and Hang Zhao},
  journal = {arXiv preprint arXiv:2603.16666},
  year = {2026},
  url = {https://arxiv.org/abs/2603.16666}
}