FastWAM

FastWAM 是一种用于机器人控制的世界动作模型策略。LeRobot 集成通过标准策略 API 暴露 FastWAM，使其可以通过 policy.type=fastwam 进行配置，使用 lerobot-train 进行训练，并通过 LeRobot 预训练策略接口进行加载。

模型概览

FastWAM 在训练阶段保留视频建模，但在推理阶段采用直接动作预测，而非迭代生成未来观测。该 LeRobot 策略封装了 FastWAM 动作模型，将 LeRobot 批次适配为 FastWAM 训练样本，并提供标准的处理器流水线以进行归一化和动作后处理。

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

LeRobot 集成涵盖的内容

• 通过 LeRobot 配置标准的 policy.type=fastwam
• 图像、状态、动作和语言任务批处理适配
• 通过 select_action 和 predict_action_chunk 进行动作块推理
• 通过 LeRobot 策略 API 保存/加载检查点
• 可配置的 LIBERO 夹爪动作后处理

安装要求

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

pip install -e ".[fastwam]"

这将从 pyproject.toml 中安装 FastWAM 策略的额外依赖项：transformers、diffusers、ftfy 和 regex，以及 LeRobot 的基础依赖项。

对于 LIBERO 评估，请同时安装基准测试的依赖项：

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

这将安装两个附加组件。除了上述 FastWAM 依赖项外，libero 附加组件还会安装 LeRobot 数据集依赖项、Linux 上的 hf-libero 以及 scipy。

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

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

数据要求

FastWAM 需要一个 LeRobot 数据集，其中包含：

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

默认的视觉配置是一个名为 observation.images.image 的图像特征，其形状为 (3, 224, 448)。如果数据集使用两个摄像头，请配置 policy.input_features，使它们的高度与 224 匹配，且宽度之和为 448。

用法

使用以下内容创建一个新的 FastWAM 策略：

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 格式检查点：

lerobot-eval \
  --policy.path=lerobot/fastwam_libero_uncond_2cam224 \
  --policy.device=cuda \
  --policy.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。

该检查点使用 VISUAL: IDENTITY 预处理。FastWAM 在模型内部将图像重新缩放到 [-1, 1]，因此在预处理器中应用图像归一化会导致图像被归一化两次。

对于真实机器人的 rollout，请使用相同的检查点路径：

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)，并且所有 camera_width 的总和必须等于配置的 width。

动作分块

policy.action_horizon 控制训练期间监督以及推理期间预测的未来动作数量。policy.n_action_steps 控制在策略预测新动作块之前消耗的动作数量。policy.n_action_steps 必须小于或等于 policy.action_horizon。

相对动作

FastWAM 默认预测绝对动作。设置 policy.use_relative_actions=true 后，它将改为预测相对于当前 observation.state 的偏移量，因此必须设置 policy.proprio_dim。在 policy.relative_exclude_joints 中指定的关节（默认为 ["gripper"]）仍保持为绝对动作。首先，请在相对空间中重新计算数据集统计信息，并将 --operation.chunk_size 设置为 policy.action_horizon（参见 π₀.₅ 相对动作）：

lerobot-train \
  --dataset.repo_id=your-org/your-dataset \
  --policy.type=fastwam \
  --policy.use_relative_actions=true \
  --policy.relative_exclude_joints='["gripper"]' \
  ...

在相对动作功能出现之前保存的检查点（例如 lerobot/fastwam_base），在使用此选项基于它们进行微调时，会添加相对动作处理器步骤。这仅适用于微调阶段：评估和 rollout 会直接使用随检查点保存的处理器。

文本编码器设备

冻结的 UMT5-XXL 文本编码器大约占用 11 GB。policy.text_encoder_device=cpu 可将其保留在 GPU 之外：提示词在 CPU 上进行编码，随后将嵌入向量移至 policy.device。在推理期间，编码后的任务会被缓存，直到调用 policy.reset()，因此较慢的 CPU 编码在每个 episode 中仅执行一次，而不是在每个 chunk 中执行一次。训练阶段会对每个 batch 进行编码，因此在该场景下，将编码器置于 CPU 会拖慢每一步的执行速度。

Wan 组件

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

注意力后端

FastWAM 的 DiT 在所有注意力计算中均使用 PyTorch 的 scaled_dot_product_attention（SDPA）。它不使用 FlashAttention：其 Mixture-of-Transformers (MoT) 路由需要任意布尔型 [query, key] 注意力掩码，而 FlashAttention 的 varlen API 无法表达这种掩码。因此，安装 flash-attn 包对 FastWAM 路径没有任何影响。（请注意，SDPA 本身在内部仍可能选择 PyTorch 自带的 flash / memory-efficient / math 内核——这与 flash-attn 包无关。）

编译后的 Action 推理

FastWAM 可以将其视频预填充和缓存的动作去噪图进行编译，以减少重复推理步骤中的 Python 和 CUDA 内核启动开销：

--policy.compile_action_infer=true

该选项默认处于禁用状态。首次动作预测会为当前的输入形状编译一个图，因此速度较慢；后续具有相同形状的预测会复用该图。形状变化可能会触发重新编译。在不支持的 PyTorch 或加速器配置上，eager 路径仍然可用。

LIBERO 操作切换

FastWAM LIBERO 检查点默认使用 policy.toggle_action_dimensions=[-1]，以匹配原始 FastWAM 评估流水线中使用的夹爪动作约定：

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

结果

在 LIBERO 上使用 lerobot/fastwam_libero_uncond_2cam224 进行评估：

复现：lerobot-eval --policy.path=lerobot/fastwam_libero_uncond_2cam224 --policy.device=cuda --policy.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 代码
• 上游发布的检查点
• 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}
}