FastWAM

FastWAM là một chính sách World Action Model để điều khiển robot. Tích hợp LeRobot phơi bày FastWAM thông qua API chính sách tiêu chuẩn, cho phép cấu hình với policy.type=fastwam, huấn luyện với lerobot-train và nạp thông qua giao diện chính sách được huấn luyện trước của LeRobot.

Tổng quan về Mô hình

FastWAM duy trì mô hình hóa video trong quá trình huấn luyện, nhưng sử dụng dự đoán hành động trực tiếp tại thời điểm suy luận thay vì tạo ra các quan sát tương lai theo cách lặp lại. Chính sách LeRobot này bọc mô hình hành động FastWAM, điều chỉnh các lô LeRobot thành các mẫu huấn luyện FastWAM, và cung cấp đường ống xử lý chuẩn cho việc chuẩn hóa và xử lý hậu hành động.

Cài đặt khởi tạo các thành phần mô hình thế giới thị giác từ Wan-AI/Wan2.2-TI2V-5B theo mặc định và dự đoán các khối hành động có hình dạng [batch, action_horizon, action_dim].

Những gì Tích hợp LeRobot Bao Gồm

• Cấu hình policy.type=fastwam chuẩn thông qua LeRobot
• Điều chỉnh theo lô cho tác vụ hình ảnh, trạng thái, hành động và ngôn ngữ
• Suy luận chunk hành động thông qua select_action và predict_action_chunk
• Lưu/tải checkpoint thông qua các API chính sách của LeRobot
• Xử lý hậu kỳ hành động kẹp LIBERO có thể cấu hình

Yêu cầu cài đặt

Cài đặt LeRobot từ mã nguồn, sau đó cài đặt các phụ thuộc của FastWAM:

pip install -e ".[fastwam]"

Điều này cài đặt phần mở rộng chính sách FastWAM từ pyproject.toml: transformers, diffusers, ftfy và regex, cùng với các phụ thuộc cơ bản của LeRobot.

Để đánh giá LIBERO, hãy cài đặt các phụ thuộc của benchmark:

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

Điều này cài đặt cả hai gói bổ sung. Ngoài các phụ thuộc của FastWAM ở trên, gói bổ sung libero sẽ cài đặt các phụ thuộc của bộ dữ liệu LeRobot, hf-libero trên Linux, và scipy.

FastWAM sử dụng backbone Wan2.2 TI2V. Model id mặc định là:

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

Yêu cầu về dữ liệu

FastWAM mong đợi một tập dữ liệu LeRobot có:

• một hoặc nhiều quan sát thị giác có độ rộng được nối lại thành policy.image_size[1]
• observation.state khi policy.proprio_dim không phải là None
• action
• một chỉ thị tác vụ ngôn ngữ thông qua trường tác vụ của bộ dữ liệu, hoặc các tensor context và context_mask đã được tính trước

Cấu hình hình ảnh mặc định là một đặc trưng hình ảnh có tên observation.images.image với hình dạng (3, 224, 448). Nếu bộ dữ liệu sử dụng hai camera, hãy cấu hình policy.input_features sao cho chiều cao của chúng khớp với 224 và tổng chiều rộng bằng 448.

Cách sử dụng

Tạo một chính sách FastWAM mới với:

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

Đánh giá một checkpoint có sẵn ở định dạng LeRobot trên LIBERO-10 bằng cách:

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

Đối với libero_goal, libero_spatial và libero_object, hãy sử dụng --env.episode_length=300.

Checkpoint này sử dụng tiền xử lý VISUAL: IDENTITY. FastWAM tự động điều chỉnh lại thang giá trị ảnh về [-1, 1] bên trong mô hình, do đó việc áp dụng chuẩn hóa ảnh trong bộ tiền xử lý sẽ khiến ảnh bị chuẩn hóa hai lần.

Để triển khai trên robot thực, hãy sử dụng cùng đường dẫn checkpoint:

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

Ghi chú cấu hình

Đặc điểm hình ảnh

policy.image_size là kích thước của tensor ảnh FastWAM đã ghép nối dưới dạng (height, width). Mỗi đặc trưng ảnh được cấu hình phải có shape (3, height, camera_width), và tổng tất cả các camera_width phải bằng width đã cấu hình.

Chia nhỏ hành động

policy.action_horizon kiểm soát số lượng hành động tương lai được giám sát trong quá trình huấn luyện và được dự đoán trong quá trình suy luận. policy.n_action_steps kiểm soát số lượng hành động được tiêu thụ trước khi chính sách dự đoán một khối mới. policy.n_action_steps phải nhỏ hơn hoặc bằng policy.action_horizon.

Các thành phần của Wan

FastWAM nạp Wan VAE, video DiT, text encoder và tokenizer từ thư mục mô hình Wan được cấu hình hoặc model id trên Hugging Face Hub. Các checkpoint FastWAM định dạng LeRobot được lưu bởi save_pretrained cũng sao chép các tệp thành phần Wan cục bộ cần thiết cho from_pretrained.

Backend chú ý

DiT của FastWAM sử dụng scaled_dot_product_attention (SDPA) của PyTorch cho tất cả các phần chú ý. Nó không sử dụng FlashAttention: việc định tuyến Mixture-of-Transformers (MoT) của nó cần các mặt nạ chú ý boolean tùy ý [query, key], mà API varlen của FlashAttention không thể biểu diễn. Do đó, việc cài đặt gói flash-attn không có tác dụng gì đối với đường dẫn FastWAM. (Lưu ý rằng bản thân SDPA vẫn có thể chọn các kernel flash / memory-efficient / math của riêng PyTorch bên trong — điều này không liên quan đến gói flash-attn.)

Suy luận Hành động Đã Biên dịch

FastWAM có thể biên dịch các đồ thị video-prefill và action-denoising đã được cache để giảm thiểu overhead khởi tạo kernel Python và CUDA qua các bước suy luận lặp lại:

--policy.compile_action_infer=true

Tùy chọn này bị vô hiệu hóa theo mặc định. Dự đoán hành động đầu tiên sẽ biên dịch một đồ thị cho các hình dạng đầu vào hiện tại và do đó sẽ chậm hơn; các dự đoán sau với cùng các hình dạng sẽ tái sử dụng nó. Sự thay đổi hình dạng có thể kích hoạt việc biên dịch lại. Lối đi eager vẫn khả dụng trên các cấu hình PyTorch hoặc bộ tăng tốc không được hỗ trợ.

LIBERO Action Toggle

Checkpoint LIBERO của FastWAM sử dụng policy.toggle_action_dimensions=[-1] theo mặc định để khớp với quy ước hành động kẹp (gripper action) được sử dụng bởi pipeline đánh giá FastWAM gốc:

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

Kết quả

Đánh giá trên LIBERO với lerobot/fastwam_libero_uncond_2cam224:

Tái hiện: 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).

Tham khảo

• Bài báo Fast-WAM
• Trang dự án Fast-WAM
• Mã Fast-WAM
• Điểm kiểm tra đã phát hành ở hạ nguồn
• Wan2.2 TI2V 5B

Trích dẫn

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