擴散 Transformer 策略:大規模預訓練驅動跨具身 VLA 泛化,Calvin 基準 3.6 任務突破


2025年2月來自上海AI實驗室、浙大、香港中文大學、北大、商湯科技、清華和中科院香港科學創新研究院的論文“Diffusion Transformer Policy: Scaling Diffusion Transformer for Generalist Vision-Language-Action Learning”。
最近,在多樣化的機器人數據集上進行預訓練的大型視覺-語言-動作模型,已展示出利用少量域內數據泛化到新環境的潛力。然而,這些方法通常通過小型動作頭預測單個離散或連續動作,這限制處理多樣化動作空間的能力。相比之下,用大型多模態擴散Transformer(稱為擴散Transformer策略)對連續動作序列進行建模,其中直接通過大型Transformer模型對動作塊進行去噪,而不是使用小型動作頭進行動作嵌入。通過利用Transformer的擴展能力,所提出的方法可以有效地在大型多樣化機器人數據集上對連續末端執行器動作進行建模,并獲得更好的泛化性能。
大量實驗證明,擴散Transformer策略在 Maniskill2、Libero、Calvin 和 SimplerEnv 以及真實世界的 Franka Arm 上的有效性和泛化性,與 OpenVLA 和 Octo 相比,在 Real-to-Sim 基準 SimplerEnv、真實世界的 Franka Arm 和 Libero 上實現一致更好的性能。具體來說,所提出的方法無需任何花哨的手段,在 Calvin 任務 ABC→D 中僅使用單個第三視角攝像機流就實現最佳性能,將一行 5 個任務中連續完成的平均數量提高到 3.6 個,并且預訓練階段顯著提高 Calvin 上的成功序列長度超過 1.2。
傳統的機器人學習范式通常依賴于為特定機器人和任務收集的大規模數據,但由于現實世界中機器人硬件的限制,為通用任務收集機器人數據既耗時又昂貴。如今,自然語言處理和計算機視覺中的基礎模型 [37, 46– 48, 58] 在廣泛、多樣、與任務無關的數據集上進行預訓練,已展示出強大的能力,無論是零樣本還是使用少量任務特定樣本解決下游任務。原則上,暴露于大規模多樣化機器人數據集的通用機器人策略,可能會提高下游任務的泛化和性能 [6, 7]。然而,在具有多樣化傳感器、動作空間、任務、相機視圖和環境的大規模跨具身數據集上訓練通用機器人策略是一項挑戰。
為了實現統一的機器人策略,現有的研究直接將視覺觀察和語言指令映射到具有大型視覺-語言-動作模型的動作上,用于機器人導航 [62, 63] 或操控 [6, 7, 28, 69],并展示對新環境的零樣本或少樣本泛化。機器人 Transformers [6, 7, 50] 提出基于 Transformer 架構的機器人策略,并通過在大規模 Open X-Embodiment 數據集 [50] 上進行訓練展示穩健的泛化能力。Octo [69] 遵循帶有擴散動作頭的自回歸 Transformer 架構,而 OpenVLA [28] 則將動作空間離散化,并利用預訓練的視覺-語言模型來構建暴露于 Open X-Embodiment 數據集 [50] 的 VLA 模型。盡管那些視覺-語言-動作 (VLA) 模型 [28, 69] 已經顯示出從大型跨具身數據集 [50] 中學習機器人策略的潛力,但跨具身數據集中機器人空間的多樣性仍然限制泛化。
最近的擴散策略 [12, 26, 57, 77] 已經展示其在機器人策略學習中的穩定能力,用于使用 UNet 或交叉注意架構進行單任務模仿學習,而擴散Transformer則展示其在多模態圖像生成中的可擴展性 [52]。具體而言,Octo [69] 提出一種通用策略,該策略使用以自回歸多模態 Transformer 單嵌入為條件的小型 MLP 網絡,對動作進行去噪。然而,大規模跨具身數據集的機器人空間,包含各種攝像機視圖和不同的動作空間,這對小型 MLP 分別對單動作頭嵌入為條件的每個連續動作進行去噪提出重大挑戰。同時,先前的擴散策略 [26, 57, 69] 首先將歷史圖像觀察和指令融合到嵌入中,然后再進行去噪過程,這可能會限制動作去噪學習,因為動作預期通常直接依賴于詳細的歷史觀察,而不是融合嵌入。
去噪擴散技術 [8、15、21、52、59] 是圖像生成的先驅,而最近的擴散策略 [9、10、12、26、35、57、71、72、77] 與之前的機器人策略相比,表現出強大的多模態動作建模能力。當前的擴散策略方法通常遵循 Unet 結構或淺層交叉注意網絡來完成單個操作任務,而大規模多模態擴散策略研究甚少。例如,3D 擴散策略 [77] 對以 3D 點云為條件的策略進行去噪,而 3D 3D diffuser actor [26] 提出一種基于點云的交叉注意 3D 擴散策略。最近的通用策略 Octo [69] 使用小型 MLP 擴散器對 Transformer 模型的嵌入進行降噪處理。
語言條件策略 [11, 19, 39, 45, 56, 78] 更適合實際應用,并且具身社區對通才機器人策略的興趣日益濃厚,其基礎多模態模型可用于機器人導航 [4, 23, 62, 63, 67, 76] 和操控 [2, 5– 7, 16, 17, 25, 28, 40, 42, 50, 51, 56, 60, 64, 65, 69, 74, 75]。最近的方法 [6, 7, 28, 50, 69] 旨在通過可擴展的視覺-語言-動作模型實現通才策略。[6, 7, 28, 50] 通過將機器人動作的每個維度分別離散化為 256 bins來構建動作 tokens。然而,這種離散化策略會導致機器人執行的內部偏差。
DiT,是一個用于通才機器人策略學習的擴散 Transformer 架構。與之前的機器人 transformer 模型 [6、7、28、50、69] 類似,利用 transformer 作為基礎模塊,保留大規模跨具身數據集的可擴展性。
架構
如圖所示不同的機器人策略架構:

指令token化。語言指令由 frozon CLIP [55] 模型token化。
圖像觀察表示。圖像觀察首先傳遞到 DINOv2 [49] 以獲得圖像塊特征。注:DINOv2 是在與機器人數據不同的網絡數據上訓練的,因此通過端到端的方式與 Transformers 一起聯合優化 DINOv2 參數。
Q-Former。為了降低計算成本,結合 Q-Former [30] 和 FiLM [53] 條件,根據指令上下文從 DINOv2 [49] 的塊特征中選擇圖像特征。
動作預處理。用末端執行器動作,并用 7D 向量表示每個動作,包括 3 個平移向量維度、3 個旋轉向量維度和 1 個夾持器位置維度。為了將維度與圖像和語言token對齊,只需用零填充連續動作向量即可構建動作表示。僅在去噪擴散優化期間將噪聲添加到 7D 動作向量中。
架構。核心設計是擴散Transformer結構 [52],它對動作token塊進行去噪,而不是對每個單個動作token進行去噪,直接以圖像觀察和指令token為條件,通過具有因果Transformer的上下文條件風格進行條件化,即,只需將圖像特征、語言token和時間步嵌入連接到序列的前面,平等地處理來自指令token的嘈雜動作,如圖所示。這種設計保留了Transformer網絡的縮放屬性,并允許直接以圖像塊為條件進行去噪學習,從而有助于模型捕捉歷史觀察中的詳細動作變化。該模型以語言指令和圖像觀察為條件,具有因果Transformer結構,由添加到連續動作中的噪聲進行監督。換句話說,用大型Transformer模型直接在動作塊空間中執行擴散目標,這與具有幾個 MLP 層的共享擴散動作頭不同 [69]。

提出的擴散Transformer策略是一種通用設計,可以擴展到不同的數據集,并表現出色。同時,還可以向Transformer 結構中添加其他觀察token和輸入。
訓練目標
在架構中,去噪網絡 ε_θ(x^t, c_obs, c_instru, t) 是整個因果Transformer,其中 c_obs 是圖像觀察,c_instru 是語言指令,t ∈ 1, 2, …T 是實驗中的步進索引。在訓練階段,在時間步 t 處采樣一個高斯噪聲向量 x^t ∈ N (0, I),其中 T 是去噪時間步數,并將其作為 a? 添加到動作 a 中以構造噪聲動作token,最后根據去噪網絡 ε_θ(a?, c_obs, c_instru, t) 預測噪聲向量 x?,其中 t 是在訓練期間隨機采樣的。利用 xt 和 ?x^t 之間的 MSE 損失優化網絡。
為了生成一個動作,使用優化的 Transformer 架構 ε_θ 從采樣的高斯噪聲向量 x^T 應用 T 步去噪,如下所示,
xt?1 = α(xt?γε_θ(xt, c_obs, c_instru, t)+N(0, σ^2I))
其中 α、γ、σ 是噪聲調度器 [21]。在實驗中,ε_θ 用于預測添加到操作中的噪聲。
預訓練數據
為了評估所提出的 Diffusion Transformer 策略,選擇 Open X-Embodiment 數據集 [50] 對模型進行預訓練。主要遵循 [28, 69] 來選擇數據集并為每個數據集設置權重。對類似于 [50] 的動作進行規范化,并過濾掉數據集中的異常動作。
預訓練細節
設計所提出的 Diffusion Transformer 架構,并在大型跨具身數據集 [50] 中評估預訓練方法。在 Open X-Embodiment 數據集 [50] 的預訓練階段使用 DDPM [21] 擴散目標,T = 1000,而本文用 DDIM [66] 將 T = 100 設置為零樣本評估以加速推理。根據 Maniskill2 [18] 的初步實驗,用 2 張觀察圖像并預測 32 個動作塊。用 AdamW [38]將 casual transformer 和 Q-Former 的學習率設置為 0.0001,DINOv2 的學習率設置為 0.00001,批量大小設置為 8902。
實驗
SimplerEnv [34] 是一個 Real-to-Sim 平臺,用于使用模擬平臺評估從真實機器人數據中學習的策略。CALVIN(從語言和視覺組合動作)[43] 是一個開源模擬基準,用于學習長期語言條件任務。
LIBERO 是一個通用基準,專注于多任務和終身機器人學習問題中的知識轉移[36]。Maniskill2 [18] 是 SAPIEN Maniskill 基準 [44] 的下一代版本,被廣泛用于評估具體模型的廣義操縱能力。
將 Franka 設置在黑色背景的桌子上。同時,用距 Franka Arm 約 1.5 米的單個第三人稱 RGB 相機。考慮到 Franka 設置環境與預訓練數據 Open X-Effirmation [68] 中的場景不同,主要在開箱即用生成和小樣本生成方面評估所提出的方法。實驗中,用相同的場景評估每個模型,并且將物體放置在 Franka Arm 前面的 9 個相似位置中,形成 9 宮格格式。同時,在放置物體的位置上保持較小的差異以進行評估。
為了評估使用少量訓練樣本進行快速微調的能力,設置兩個拾取和放置任務,并進一步收集五個復雜的操作任務,即“拿起碗并將球倒入盒子中”,“打開盒子并拿起盒子內的積木”,“拿起香蕉并將其移入筆筒”,“疊起碗”和“拿起杯子并將咖啡豆從杯中倒入碗中”。為每個任務收集10個樣本以進行10次微調泛化。如圖展示場景和任務。圖中的圖像是模型輸入。現實環境具有挑戰性,因為物體(小于20毫米)與整個場景相比很小。還對域內評估進行簡單的評估,用于比較域內泛化下的不同動作策略,其中5個拾取和放置任務包括“綠色積木”,“獼猴桃”,“香蕉在前三個任務中為每個任務收集了 50 條軌跡,而將剩余兩個任務(“拾起綠色小塊”和“拾起粉色塊”)留作分布外評估。

微調細節。在實驗中,用 Lora [22] 和 AdamW [38] 在真實的 Franka Arm 上對所提出的方法進行 10,000 步的微調。將 DDPM [21] 的時間步數設置為 100,將批處理大小設置為 512。
END
轉載自CSDN-三谷秋水

