FastWAM

FastWAMは、ロボット制御のためのWorld Action Modelポリシーです。LeRobotの統合により、FastWAMは標準的なポリシーAPIを通じて公開され、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]"

これにより、両方のextrasがインストールされます。上記のFastWAMの依存関係に加えて、libero extraは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]に連結される幅を持つ1つ以上の視覚的観測
• policy.proprio_dim が None ではない場合の observation.state
• action
• データセットのタスクフィールドを介した言語タスクの指示、または事前計算済みのcontextおよびcontext_maskテンソル

デフォルトのビジュアル設定は、(3, 224, 448) の形状を持つ observation.images.image という名前の1つの画像特徴です。データセットが2つのカメラを使用している場合は、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

…

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]にリスケールするため、プリプロセッサで画像正規化を適用すると、正規化が2回行われてしまいます。

実機でのロールアウトには、同じチェックポイントパスを使用します:

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 に設定して、相対空間でのデータセット統計を再計算してください（π₀.₅ relative actions を参照）：

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 など）から、このオプションを指定してファインチューニングを行う場合、相対アクションのプロセッサのステップが追加されます。これはファインチューニング時のみ適用されます。評価やロールアウトでは、チェックポイントに保存されたプロセッサがそのまま使用されます。

テキストエンコーダーデバイス

凍結された UMT5-XXL テキストエンコーダーは約 11 GB を占有します。policy.text_encoder_device=cpu を設定すると GPU 上に配置されず、プロンプトは CPU でエンコードされ、埋め込みベクトルが policy.device に移動されます。推論時には、エンコードされたタスクが policy.reset() が呼ばれるまでキャッシュされるため、低速な CPU でのエンコードはエピソードごとに 1 回のみ実行され、チャンクごとに実行されることはありません。一方、学習ではすべてのバッチがエンコードされるため、CPU への配置は各ステップを遅くします。

Wan コンポーネント

FastWAMは、設定されたWanモデルディレクトリまたはHugging Face HubのモデルIDから、Wan VAE、video 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パッケージとは無関係です。）

コンパイル済みアクション推論

FastWAM は、繰り返し推論ステップにわたる Python および CUDA カーネル起動のオーバーヘッドを削減するために、ビデオプリフィルとキャッシュされたアクションデノイジングのグラフをコンパイルできます:

--policy.compile_action_infer=true

このオプションはデフォルトで無効です。最初のアクション予測では、現在の入力形状用にグラフがコンパイルされるため、処理が遅くなります。その後、同じ形状での予測ではそのグラフが再利用されます。形状が変更されると、再コンパイルがトリガーされる場合があります。サポートされていないPyTorchまたはアクセラレータの構成では、eagerパスが引き続き利用可能です。

LIBERO アクションの切り替え

FastWAMのLIBEROチェックポイントは、元のFastWAM評価パイプラインで使用されているグリッパーアクションの慣例に合わせるため、デフォルトでpolicy.toggle_action_dimensions=[-1]を使用します:

--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}
}