[EAI-023] FAST，機器人動作專用的Tokenizer，提高VLA模型的能力和訓練效率

Paper Card

論文標題：FAST: Efficient Action Tokenization for Vision-Language-Action Models
論文作者：Karl Pertsch, Kyle Stachowicz, Brian Ichter, Danny Driess, Suraj Nair, Quan Vuong, Oier Mees, Chelsea Finn, Sergey Levine
論文鏈接：https://arxiv.org/abs/2501.09747
論文出處：/
論文被引：/
項目主頁：https://www.pi.website/research/fast

Abstract

基于Transformer的視覺-語言-動作（VLA）策略等自回歸序列模型，可以非常有效地捕捉復雜且可泛化的機器人行為。但是，此類模型對連續動作信號的進行 tokenization（離散標記化），這決定了模型預測的離散token （標記）如何映射到連續的機器人動作（action）。當前基于簡單逐維度、逐時間步長分箱方案的機器人動作離散標記化方法，在從高頻機器人數據中學習靈巧技能時，通常表現不佳。為了解決這一挑戰，本文提出了一種基于離散余弦變換的基于壓縮的機器人動作離散標記化方案——Frequency-space Action Sequence Tokenization (FAST)，能夠為高度靈巧和高頻的任務訓練自回歸VLA，而標準離散化方法在此類任務中完全失效。基于FAST，發布了一個通用的機器人動作離散標記化器 FAST+，它在100萬個真實的機器人動作軌跡上訓練。它可以用作各種機器人動作序列（具有不同的動作空間和控制頻率）的黑盒分詞器（tokenizer）。當與 VLA 模型 pi0 結合使用時，可以擴展到在1萬小時的機器人數據上進行訓練，并與擴散VLA模型的性能相當，但訓練時間減少了5倍。

Summary

研究背景

性能好的tokenizer對序列模型的性能至關重要。當前機器人策略通常使用基于每個維度、每個時間步長分箱方案的簡單 tokenization 方法，這類方法在學習具有高頻控制的靈巧技能時的表現不好（本文中涉及的測試全部失敗）。當預測未來動作序列時，簡單的tokenizer難以捕獲各個時間步之間的相關性，高度相關的action token降低了自回歸VLA模型對 next token prediction 建模的有效性。

方法介紹

在這里插入圖片描述

從第一性原理出發，開發新的 action tokenizer。關鍵想法：受到llama中使用的 BPE 編碼方法的啟發，在模型訓練之前壓縮機器人動作信號，減少連續token之間的相關性。考慮到機器人動作是連續信號，因此采用離散余弦變換編碼，由此產生的 tokenization 方法稱為 Frequency-space Action Sequence Tokenization (FAST)。在 DROID 數據集上高效訓練 VLA 模型，在未見的真實環境通過自然語言提示進行零樣本評測。
在這里插入圖片描述

構建了 FAST+ 通用的機器人動作 tokenizer，在100萬條真機數據上訓練。可以對各種機器人動作序列進行分詞，報錯單臂機器人、雙臂機器人和移動機器人。當pi0模型結合FAST進行訓練時，比原來的結合擴散思想的模型訓練時間縮短了 5 倍，但性能相當。

前置知識

問題描述：目標是訓練策略 $π(a_{1:H}|o)$ ，將觀測結果 $o$ 映射到未來機器人動作序列 $a_{1:H}$ 。假設策略輸出動作塊（Action Chunk），即序列長度為 $H$ 動作，這使得更容易產生時間上一致的動作并減少復合誤差。動作分詞（Action Tokenization）的目標是定義一個映射 $\mathcal{T}_{a}:a_{1:H} → [T_1,…,T_n]$ ，從維度為 $|\mathcal{A}|$ 的連續動作 $a_{1:H}$ 序列映射到來自大小為 $|\mathcal{V}|$ 的詞匯表中的 $n$ 個離散token序列 $T∈|\mathcal{V}|$ 。動作序列之間的token數量 n 可能不同，就像相同長度的句子可能被離散化為可變數量的文本token一樣。

基于分箱的動作token化：動作token化最常用的方法是簡單的分箱離散化。對于給定的動作 a，這種方法獨立地離散化每個維度，將訓練數據集中值的范圍劃分為 $N$ 個均勻的箱，最常用的是 $N = 256$ 。對于D維的動作序列 $a_{1:H}$ ，此token化方案將應用于每個時間步，從而產生最終的token序列 $\mathcal{T}a(a1:H)=[T_{1,1},…,T_{1,D},…,T_{H,1},…,T_{H,D}]$ 。對于高頻機器人數據，這種token化方案并非最優：很容易為每個動作片段生成數百個 token，使得難以訓練并且推理速度慢。

Tokenization 對 VLA 模型訓練的影響

在這里插入圖片描述

創建了一個簡單的合成時間序列數據集，其目標是預測一個插值四個隨機生成點的三次曲線，如圖3所示。這個問題反映了高頻動作片段上訓練的策略面臨的挑戰，即策略必須預測一系列連續的動作。訓練一個小型自回歸Transformer進行實驗，序列時間步H從25到800，以模擬不同頻率收集的動作數據。Navie方法表示將動作序列中的每一個元素進行256bins分箱操作。

實驗結果表明，分箱方案的模型在低采樣頻率的條件下預測效果較好，但是隨著采樣頻率增加，預測誤差急劇增加。為什么呢？因為自回歸模型的訓練目標是下一個token預測，因此，他們的學習信號在給定 $T_{1:i-1}$ 的情況下與 $T_i$ 的邊際信息內容成正比。分箱方案隨著采樣頻率的增加，邊際信息接近于零：對于平滑信號，隨著時間步長的縮短，每個時間步長的變化成比例地縮小。這極大地減慢了訓練收斂速度，并且難以你和復雜的高頻數據集。例如，OpenVLA在低頻的 BridgeV2 和 RT-1 數據集上運行良好，但是對于高頻的 DROID 數據集表現不佳。這說明為機器人動作設計更好的分詞器的重要性。

通過時間序列壓縮實現高效的動作分詞器

在這里插入圖片描述

為了解決前述提到的高頻動作軌跡中的冗余會導致每個action token的邊際信息量低進而導致訓練性能差的問題，需要一種能夠將高度冗余的動作信號壓縮減少成少量高信息量token的動作離散化方法。

本文使用基于離散余弦變換（DCT）的壓縮算法構建FAST。DCT是一種頻域變換，它將連續信號表示為各種頻率的余弦元素之和。低頻捕獲信號的整體形狀，而高頻份量反映尖銳的跳變（sharp jumps）。

圖4說明了FAST從原始機器人動作到action token的變換步驟。首先對輸入動作進行歸一化，然后對每個動作維度應用DCT，為了壓縮信號，忽略不重要的系數，得到量化后的稀疏的矩陣。然后將矩陣展平為一個一維整數向量，然后訓練一個BPE分詞器將其無損壓縮成稠密的token。
在這里插入圖片描述

通用機器人動作分詞器

FAST 中唯一需要學習的組件是 BPE 編碼器的詞匯表，這個詞匯表需要針對分詞器應用的每個新數據集進行循例那，雖然只需要幾分鐘，但也增加了使用FAST的難度。因此，使用100萬個1s的動作塊訓練了一個通用的機器人動作分詞器。已經開源并合入到Transformers庫了。
在這里插入圖片描述

訓練所需要的數據集：數據集有多種動作空間：聯合空間、末端執行器世界坐標系和末端執行器相機坐標系，以確保所得分詞器的通用性。Open X-Embodiment、DROID和Bridge?V2則以其原始形式包含在內。在分詞之前，所有動作都填充到32維，以適應不同維度的動作空間。在這里插入圖片描述

消融實驗

模型基線：pi0、OpenVLA

實驗目標：驗證FAST分詞器+自回歸VLA模型的有效性

評估任務

在這里插入圖片描述

如圖5所示，包含7個評估任務（6個真實機器人任務，1個模擬任務），旨在測試VLA在高度靈巧的任務（例如折疊衣物）和泛化任務（例如在未見環境中進行0樣本桌面操作）上的性能。

Libero：在Libero [43]模擬基準套件上進行測試。測量了Libero-Spatial、Libero-Object、Libero-Goal和Libero-10的平均性能。
餐桌清理 [7] (20 Hz)：一臺UR5單臂機器人需要清理桌子，將12個物體分類到垃圾桶（用于垃圾）和塑料容器（用于盤子、碗、杯子和餐具）中。此任務需要精確抓取各種物體。
折疊T恤 [7] (50 Hz)：一套雙臂ARX機器人需要在一個靜止的桌面上折疊各種襯衫。在任務開始時，襯衫平放在桌子上。成功完成此任務需要精確的抓取和移動才能折疊襯衫。
雜貨裝袋 [7] (20 Hz)：一臺UR5單臂機器人需要將 7 個物體從桌子上裝入雜貨袋中，注意不要弄倒或撕破袋子。此任務需要拾取各種各樣的物體并小心地將它們放入袋中。
從烤面包機中取出吐司 [7] (50 Hz)：一臺雙臂Trossen Viper-X機器人需要從烤面包機中取出兩片面包并將它們放在盤子上。此任務需要精確地抓取和放置面包片。
衣物折疊 [7] (50 Hz)：一臺雙臂ARX機器人需要從籃子里取出襯衫和短褲，將它們平放在桌子上，然后折疊并堆疊起來。這是我們測試中最靈巧的任務。它需要精確的抓取，動態的動作來使衣物平整，在衣物纏結時進行重試和糾正，以及將折疊好的衣物精確地放置在現有的衣物堆上。報告了單個服裝物品的成功率。
零樣本DROID桌面操作 [38] (15 Hz)：測試了一個在完整DROID數據集上訓練的策略，該策略涵蓋各種桌面操作任務，例如拾取和放置物體、擦拭、打開和關閉抽屜等。在未見的環境中測試該策略，該環境具有新的桌子設置、背景、新穎的物體、視角和桌子高度。這是第一次在完全未見的環境中對DROID策略進行“零樣本”評估，無需協同訓練或微調，只需使用自然語言提示預訓練模型即可。

機器人動作分詞器對比

在這里插入圖片描述

使用1秒的動作片段。FAST分詞器對所有數據集都實現了有效的壓縮，高頻動作上效果更顯著，token為20-53。
在這里插入圖片描述

先前工作中應用的naive分詞方法難以在高頻機器人數據上學習有效的策略。最高頻的任務中尤為明顯：餐桌整理 (20Hz) 和 T 恤折疊 (50Hz)。

在這里插入圖片描述

FAST 分詞技術實現了在 DROID 數據集上成功訓練強大的通用策略，該策略可以通過自然語言提示，在未見過的環境中進行零樣本評估，無需微調。所有先前的工作都沒有顯示零樣本結果，而是完全專注于聯合訓練或微調評估。在三個大學的校園中對各種桌面操作任務進行測試，證明了策略的通用性（圖 7）。無需額外訓練，該策略能夠熟練地執行簡單的操作任務，例如在各種場景和攝像機視角下拾取和放置物體、打開和關閉櫥柜以及打開水龍頭。即使是不成功的嘗試也表現出合理的行為，例如靠近微波爐和洗碗機門的把手，即使最終未能打開它們。

消融研究

回答兩個問題：

FAST分詞方法是否獨立于底層VLA主干？
BPE壓縮步驟有多重要？

在這里插入圖片描述

為了回答第一個問題，在高頻T恤折疊數據集上訓練了一個OpenVLA策略，修改了OpenVLA模型代碼以接受多個輸入圖像并預測1秒的動作塊。結果表明，FAST能夠顯著提高OpenVLA的性能，使其能夠有效地訓練高頻機器人操作數據。這表明，分詞方法獨立于底層模型主干，并且可以輕松應用于各種預訓練的自回歸Transformer模型。

在這里插入圖片描述

在桌面整理和T恤折疊任務上消融了BPE編碼步驟。結果表明，沒有BPE編碼的策略獲得了更差的性能（但仍然優于樸素分詞）。直觀地說，DCT變換仍然將大部分信號信息集中在少數幾個token中，從而改善了學習信號。然而，如果沒有BPE，就會出現大量重復的值為0的token，這會稀釋學習信號，并顯著減慢推理速度，因為模型需要自回歸地預測數百個動作token，最終導致策略性能下降。

自回歸VLA與擴散VLA對比

在這里插入圖片描述
圖 9 所示，在小型數據集（Libero，折疊T恤；<50小時）上，兩種VLA的性能相當。但是在像搬運桌子這樣的大型數據集上，基于FAST的VLA收斂速度明顯更快，在訓練步驟比π0的擴散變體少3倍的情況下達到了高性能。使用FAST分詞訓練的自回歸 π0 模型更嚴格地遵循語言指令：在DROID評估中，擴散π0模型經常忽略語言指令，導致分數較低。未來會繼續研究擴散和自回歸VLA的語言遵循能力。

自回歸VLA的一個當前局限性在于其推理速度：π0使用擴散模型通常可以在NVIDIA 4090 GPU上在100毫秒內預測一秒鐘的動作片段，但π0模型使用FAST分詞需要大約750毫秒的每個片段推理時間，因為它必須執行更多自回歸解碼步驟（通常需要解碼30-60個動作token，而擴散模型π0需要10個擴散步驟）并使用完整的20億參數語言模型主干進行自回歸解碼（而擴散模型π0使用3億參數的“動作專家”）。未來會繼續研究離散token自回歸Transformer模型的推理提速。

實驗結論

本文介紹了 FAST，一種用于高頻機器人控制數據的動作分詞器。FAST使用離散余弦變換（DCT）和字節對編碼（BPE）來壓縮動作塊，使得其具有更好的壓縮效果。實驗表明，與以前使用的簡單動作離散化方法相比，FAST帶來了顯著的性能提升，并且優于基于矢量量化的更復雜的學習分詞方法。

未來工作：
動作分詞器。FAST是朝著通用機器人動作分詞器邁出的重要一步，但仍有很多問題有待解決。在這項工作對靜態機器人機械臂測試了 FAST。FAST+在其他機器人形態（如移動機器人、靈巧手和人形機器人）上具有良好的壓縮能力。在這些平臺上測試實際策略性能是未來工作的一個令人興奮的方向。探索替代壓縮方案，以及測試基于壓縮的動作編碼與非自回歸解碼方法（如擴散[7]）的組合，是未來研究的有趣方向。

VLA 架構。本文初步探索了兩種主要類型的 VLA 架構（自回歸和擴散解碼 VLA）之間的權衡，但最佳 VLA 架構仍未確定。未來工作應仔細研究訓練速度、語言基礎能力和任一方法的表達能力之間的權衡。

推理速度。雖然π0-FAST 的整體性能與擴散π0匹配，但在推理時間上較慢。未來的工作應該探索加快自回歸 VLA 模型推理速度的方法，以使它們能夠解決高度動態的任務。