UC伯克利等提出具身智能「動作Tokenizer」,效率飆升5倍!
來自π,伯克利和斯坦福的研究者近期提出了FAST,一種為VLA模型設計的高效動作Tokenizer。
FAST旨在解決傳統動作Tokenization方法在處理精細任務時面臨的挑戰,為VLA模型在更復雜、高頻的機器人控制任務帶來了顯著改進。
將FAST與π0 VLA結合,在1萬小時的機器人數據上對π0-FAST進行訓練,效果與最先進的擴散VLA相當,訓練時間最多縮短5倍。
論文鏈接:https://arxiv.org/abs/2501.09747
想訓練Transformer來控制機器人,具體該怎么做?
Transformer輸出的動作塊是一小段機器人動作序列,例如手臂關節角度。對于不太精細的系統,其動作序列可能包含3-5個動作;而對于高頻靈巧機器人,動作序列可能會有20-50個動作。
用合適的方法表示這些動作,對于高效學習非常重要。
現有的VLA模型常用簡單的離散劃分方法,這種方法對于簡單的任務是可行的,但對于需要高精度和高頻控制的復雜精細技能,就不太管用了。
擴散方法或流匹配通常展現出更好的性能,如π0模型。然而,擴散方法的訓練時間會更長。
那么,如何表示這些動作,才能在保持靈活性和精確性的同時,快速地訓練Transformer呢?
為解決這個問題,Physical Intelligence提出一款專為機器人動作設計的新型Tokenizer「FAST」。
借助FAST,研究者開發出了一種高效的機器人動作Tokenization方法,讓機器人技術能夠與自回歸Transformer訓練流程無縫銜接。
FAST能把動作像語言一樣,用離散Token表示。FAST提升了從大規模互聯網數據預訓練的遷移能力,增強了機器人執行語言指令的能力。
借助FAST,首次實現在DROID數據集上訓練出只需通過簡單的自然語言命令,機器人就能在新環境中零樣本執行多種操作任務。
FAST Tokenizer
用離散余弦變換(DCT)壓縮
FAST通過在訓練前壓縮原始動作塊,能顯著提升基于靈巧機器人數據的策略訓練和推理效率。
FAST使用了一種基于離散余弦變換(DCT)的壓縮算法,來提高VLA模型的訓練速度。DCT是一種頻域變換,因簡潔和計算高效,常用于壓縮算法,如JPEG圖像壓縮、MP3音頻的編解碼。
FAST Tokenization算法
研究者利用離散余弦變換(DCT)設計了FAST,它是一種快速且高效的機器人動作Tokenization方法。下圖展示了從原始機器人動作到動作token的步驟。
首先,對輸入的動作進行歸一化。完成歸一化后,每個動作維度分別應用離散余弦變換(DCT)。
研究者用字節對編碼(BPE)壓縮DCT矩陣,BPE是訓練LLM常用的壓縮算法。將DCT和字節對編碼(BPE)結合,就能把原始動作塊壓縮成數量少但更密集的動作Token。
通常每個動作塊包含30-60個Token,和以前的動作Tokenization方法相比,壓縮率提高了10倍。
通用機器人動作Tokenizer
基于FAST,研究者發布了FAST+,這是通用的機器人動作Tokenizer,用100萬個真實機器人動作序列上訓練而成。
FAST+ Tokenizer能高效地標記各種機器人動作序列,包括單臂、雙臂和移動機器人。FAST+能當黑盒Tokenizer,對各種機器人的動作序列編碼,涵蓋多樣的動作空間和控制頻率。
借助這個Tokenizer訓練自己的策略,只需三行代碼即可實現:
π0-FAST:開箱即用
研究者將FAST和π0模型結合進行測試。
以往的離散化VLA模型只能處理簡單操作任務。但有了FAST,就能借助它訓練自回歸Transformer策略,完成像折疊衣物、收拾餐桌、打包購物袋這類高難度的精細任務。同時,模型的訓練速度快了5倍。
下面是一些能用FAST策略解決的任務。
借助FAST,研究者首次在DROID數據集上訓練出通用策略,能在新環境中對多種指令實現泛化執行。
DROID數據集是一個由全球大型機器人研究聯盟歷經兩年收集的開源數據集,包含多樣化的場景和機器人操作任務,涵蓋從大學建筑到真實家庭等多種場景。
到目前為止,尚未有一種方法能在DROID數據集上訓練出通用策略,在新環境中零樣本執行語言指令。
但借助FAST,研究者實現了這一目標。快來看看它的實際表現!
在美國三所大學(加州大學伯克利分校、斯坦福大學、華盛頓大學)的測試環境中,都能直接完成簡單的操作任務。
即便在執行任務失敗,它也會基于常識去嘗試解決問題。
雖然這一策略尚不完善,卻為我們描繪了未來的圖景:未來能直接下載并使用通用機器人模型,就像今天使用語言模型一樣。
實驗結果
表中展示了FAST Tokenizer與先前方案在壓縮率方面的比較。
FAST對所有數據集的輸入動作序列都實現了顯著壓縮,對高頻動作數據的數據集壓縮效果更明顯。
基于100萬條真實機器人動作序列訓練的通用動作Tokenizer FAST+與之前的VLA相比顯著提高了訓練效率。
FAST+ Tokenizer在各類機器人數據集上有良好的壓縮性能,且在各種機器人形態、動作空間和控制頻率范圍中都展現出有效性。
研究者借助FAST Tokenizer成功訓練出首個自回歸通用策略π0-FAST。它基于之前的π0模型,沿用相同的模型架構和訓練數據集。
研究者把π0-FAST與最先進的擴散π0 VLA進行對比。
總體而言,自回歸π0-FAST模型的表現與擴散π0模型相當,即便在最具挑戰性的洗衣折疊任務上也是如此,且所需計算量顯著更少。
因其采用簡單的自回歸離散化方法,π0-FAST的收斂速度比擴散π0模型快得多,所需的訓練GPU小時數減少5倍。
對于最先進的VLA訓練,通常需數千個GPU小時,計算資源減少5倍至關重要。
基于FAST的方法可在復雜的機器人任務上訓練自回歸VLA,而先前的方法在這些任務上效果不佳。
研究者對于自回歸VLA訓練的前景感到十分興奮。然而,當前的模型存在一個較為顯著的缺點,即推理速度較慢。π0-FAST采用自回歸解碼的方式,其速度明顯比不上π0所使用的流匹配解碼。
未來,LLM中廣泛應用的加速離散自回歸Transformer模型推理的技術,或許能為提高VLA推理速度帶來啟發。