FastWAM

FastWAM adalah kebijakan World Action Model untuk pengendalian robot. Integrasi LeRobot mengekspos FastWAM melalui API kebijakan standar sehingga dapat dikonfigurasi dengan policy.type=fastwam, dilatih dengan lerobot-train, dan dimuat melalui antarmuka kebijakan pralatih LeRobot.

Ringkasan Model

FastWAM mempertahankan pemodelan video selama pelatihan, tetapi menggunakan prediksi aksi langsung pada waktu inferensi alih-alih menghasilkan observasi masa depan secara iteratif. Kebijakan LeRobot ini membungkus model aksi FastWAM, menyesuaikan batch LeRobot menjadi sampel pelatihan FastWAM, dan menyediakan pipeline prosesor standar untuk normalisasi dan pascapemrosesan aksi.

Implementasi ini menginisialisasi komponen world-model visual dari Wan-AI/Wan2.2-TI2V-5B secara default dan memprediksi chunk aksi dengan bentuk [batch, action_horizon, action_dim].

Cakupan Integrasi LeRobot

• Konfigurasi policy.type=fastwam standar melalui LeRobot
• Adaptasi batch untuk tugas gambar, state, aksi, dan bahasa
• Inferensi chunk aksi melalui select_action dan predict_action_chunk
• Menyimpan/memuat checkpoint melalui API kebijakan LeRobot
• Pemrosesan pasca aksi gripper LIBERO yang dapat dikonfigurasi

Persyaratan Instalasi

Instal LeRobot dari sumber, lalu instal dependensi FastWAM:

pip install -e ".[fastwam]"

Ini menginstal ekstra kebijakan FastWAM dari pyproject.toml: transformers, diffusers, ftfy, dan regex, serta dependensi dasar LeRobot.

Untuk evaluasi LIBERO, instal juga dependensi benchmark:

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

Ini menginstal kedua extras tersebut. Selain dependensi FastWAM di atas, extra libero menginstal dependensi dataset LeRobot, hf-libero di Linux, dan scipy.

FastWAM menggunakan backbone Wan2.2 TI2V. ID model default adalah:

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

Kebutuhan Data

FastWAM mengharapkan dataset LeRobot dengan:

• satu atau lebih pengamatan visual yang lebarnya digabungkan menjadi policy.image_size[1]
• observation.state ketika policy.proprio_dim bukan None
• action
• instruksi tugas bahasa melalui field tugas dataset, atau tensor context dan context_mask yang sudah dihitung sebelumnya

Setup visual default adalah satu fitur gambar bernama observation.images.image dengan bentuk (3, 224, 448). Jika dataset menggunakan dua kamera, konfigurasikan policy.input_features sehingga tinggi mereka sama dengan 224 dan lebar mereka berjumlah 448.

Penggunaan

Buat kebijakan FastWAM baru dengan:

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

Evaluasi checkpoint format LeRobot yang sudah ada pada LIBERO-10 dengan:

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

Untuk libero_goal, libero_spatial, dan libero_object, gunakan --env.episode_length=300.

Checkpoint ini menggunakan prapemrosesan VISUAL: IDENTITY. FastWAM melakukan penskalaan ulang gambar ke [-1, 1] di dalam model, sehingga menerapkan normalisasi gambar di preprocessor akan menormalisasi mereka dua kali.

Untuk rollout robot nyata, gunakan jalur checkpoint yang sama:

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

Catatan Konfigurasi

Fitur Gambar

policy.image_size adalah ukuran tensor gambar FastWAM yang digabungkan sebagai (height, width). Setiap fitur gambar yang dikonfigurasi harus memiliki bentuk (3, height, camera_width), dan semua lebar kamera harus berjumlah sama dengan lebar yang dikonfigurasi.

Chunking Aksi

policy.action_horizon mengontrol jumlah aksi masa depan yang diawasi selama pelatihan dan diprediksi selama inferensi. policy.n_action_steps mengontrol berapa banyak aksi yang dikonsumsi sebelum policy memprediksi chunk baru. policy.n_action_steps harus kurang dari atau sama dengan policy.action_horizon.

Aksi Relatif

FastWAM secara default memprediksi aksi absolut. Dengan policy.use_relative_actions=true, model memprediksi offset dari observation.state saat ini, sehingga policy.proprio_dim harus diatur. Sendi yang disebutkan dalam policy.relative_exclude_joints (default ["gripper"]) tetap absolut. Hitung ulang statistik dataset dalam ruang relatif terlebih dahulu, dengan --operation.chunk_size diatur ke policy.action_horizon (lihat aksi relatif π₀.₅):

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

Checkpoint yang disimpan sebelum aksi relatif ada, seperti lerobot/fastwam_base, akan mendapatkan langkah pemroses aksi relatif saat fine-tuning dari checkpoint tersebut dengan opsi ini. Hal ini hanya berlaku pada saat fine-tuning: evaluasi dan rollout menggunakan pemroses yang disimpan bersama checkpoint apa adanya.

Perangkat Text Encoder

Teks encoder UMT5-XXL yang dibekukan membutuhkan sekitar 11 GB. policy.text_encoder_device=cpu menjaganya agar tidak berada di GPU: prompt dikodekan di CPU dan embedding dipindahkan ke policy.device. Selama inferensi, tugas yang telah dikodekan di-cache hingga policy.reset(), sehingga pengkodean CPU yang lebih lambat hanya berjalan sekali per episode, bukan sekali per chunk. Pelatihan mengodekan setiap batch, sehingga penempatan CPU di sana memperlambat setiap langkah.

Komponen Wan

FastWAM memuat Wan VAE, video DiT, text encoder, dan tokenizer dari direktori model Wan yang dikonfigurasi atau model id Hugging Face Hub. Checkpoint FastWAM format LeRobot yang disimpan oleh save_pretrained juga menyalin file komponen Wan lokal yang diperlukan oleh from_pretrained.

Backend Atensi

DiT dari FastWAM menggunakan scaled_dot_product_attention (SDPA) dari PyTorch untuk semua perhatian. Ia tidak menggunakan FlashAttention: routing Mixture-of-Transformers (MoT)-nya membutuhkan mask perhatian boolean [query, key] yang arbitrer, yang tidak dapat diekspresikan oleh API varlen FlashAttention. Oleh karena itu, menginstal paket flash-attn tidak berdampak pada jalur FastWAM. (Perhatikan bahwa SDPA itu sendiri masih mungkin memilih kernel flash / memory-efficient / math milik PyTorch secara internal — ini tidak terkait dengan paket flash-attn.)

Inferensi Aksi Terkompilasi

FastWAM dapat mengompilasi graf video-prefill dan action-denoising yang di-cache untuk mengurangi overhead peluncuran kernel Python dan CUDA di seluruh langkah inferensi yang berulang:

--policy.compile_action_infer=true

Opsi ini nonaktif secara default. Prediksi aksi pertama akan mengompilasi graf untuk bentuk input saat ini dan karenanya lebih lambat; prediksi berikutnya dengan bentuk yang sama akan memakainya kembali. Perubahan bentuk dapat memicu kompilasi ulang. Jalur eager tetap tersedia pada konfigurasi PyTorch atau akselerator yang tidak didukung.

LIBERO Action Toggle

Checkpoint FastWAM LIBERO menggunakan policy.toggle_action_dimensions=[-1] secara default untuk menyesuaikan konvensi aksi gripper yang digunakan oleh pipeline evaluasi FastWAM asli:

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

Hasil

Dievaluasi pada LIBERO dengan lerobot/fastwam_libero_uncond_2cam224:

Reproduksi: 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).

Referensi

• makalah Fast-WAM
• Halaman proyek Fast-WAM
• Kode Fast-WAM
• Checkpoint yang dirilis upstream
• Wan2.2 TI2V 5B

Sitasi

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