yuantianyuan01/FastWAM

解決的問題

FastWAM 解決了機器人領域中世界動作模型（WAMs）所帶來的計算開銷與延遲問題。具體而言，它探討了在高表現力下是否必須進行「測試時未來想像」（在預測動作前先預測未來影片幀），並提供一種直接從當前觀測預測動作的方法，從而實現顯著更快的推論速度。

工作原理

FastWAM 採用基於去噪核心（ActionDiT）與 VAE 編碼的模型架構。它提供兩種主要運作模式：

• 首幀模式（Fast-WAM）： 跳過未來影片想像步驟，直接從當前觀測預測動作，以降低延遲。
• IDM 模式： 遵循傳統方法，先想像未來影片幀，再基於這些想像預測動作。

為優化效能，本專案實作了編譯後的去噪核心、批量 VAE 編碼以及輕量級 CUDA Graph 後端，三者共同減少推論時間與訓練時長。

適用對象

本專案專為從事具身智能研究的機器人研究人員與開發者設計，特別是關注動作預測、世界模型以及提升機器人控制策略即時性能的群體。

主要亮點

• 加速推論： 相較於先前版本，實現約 2 倍的端到端推論速度提升。
• 雙模式能力： 「可選 IDM」變體允許單一模型在不重新訓練的情況下，自由切換未來想像（IDM）與直接動作預測（首幀）模式。
• 資料集支援： 原生支援 LeRobot 2.1 與 3.0 資料集，提升大規模資料的可擴展性。
• 高成功率： 在 LIBERO 基準測試中，於空間、目標、物件及長時序任務上均表現出色。
• 優化之流程： 包含預計算的 T5 文字嵌入與編譯後的訓練路徑，顯著提升吞吐量。