2024年9月來自 MIT 何凱明團隊和 Meta 的論文“Scaling Proprioceptive-Visual Learning with Heterogeneous Pre-trained Transformers”。

當今訓練通用機器人模型的障礙之一是異構性。以前的機器人學習方法通常收集數據，用一個特定的具身來訓練一個任務，這種方法成本高昂，而且容易過擬合。這項工作研究通過對不同具身和任務的機器人數據進行大規模異構預訓練來學習策略表示的問題。提出異構預訓練 Transformer (HPT)，它預訓練一個大型、可共享的策略神經網絡主干，以學習與任務和具身無關的共享表示。這種通用架構將來自不同具身的特定本體感覺和視覺輸入對齊到一小段token序列，然后處理這些token以映射到控制機器人以執行不同的任務。利用最近的大規模多具身真實世界機器人數據集以及模擬數據集、部署的機器人數據集和人類視頻數據集，研究跨異構的預訓練策略。其實驗研究訓練目標的擴展行為，范圍達到 52 個數據集。

機器人領域的大規模策略學習分別利用來自真實機器人 [6, 72]、人類視頻 [53, 48] 和模擬域 [32, 62, 82, 79] 的各種數據。此外，還有多任務學習 [64, 65, 84, 22]、元學習 [78, 54, 18]、小樣本學習 [83] 和成隊學習 [81] 方面的研究。最近，RT-X、Octo、OpenVLA [6, 14, 55, 35] 在來自各種機器人實例的數據集上訓練通用的視覺-語言-動作機器人策略。

所謂異構預訓練概念如圖所示。它將不同的具身方案（每個都有自己的本體感覺和視覺傳感器）通過針對具身方案的token化器（“stem”）映射到共享的潛空間上。這將來自不同具身方案的異構數據對齊到聯合表示空間中。這樣就能夠在所有異構數據集的聯合上訓練共享的 Transformer 主干。預訓練的 Transformer 可以遷移到新的具身方案，并在遷移時學習一個小型的新token化器。

異構預訓練 Transformers (HPT)，是一類架構，旨在可擴展地從異構具身數據中學習。HPT 模塊化通用策略網絡架構（如圖所示），并使用監督學習預訓練潛 Transformer 的策略表示。受到從多模態數據中學習的啟發 [1, 73, 19, 30]，用針對具身的token化器（稱為“stem”）來對齊各種傳感器輸入，例如攝像機視圖和本體感覺輸入?！皌runk”在數據集之間共享和預訓練，并在適應預訓練期間未知的新具身和任務時進行傳輸。此外，使用針對任務的操作解碼器（稱為“頭”）來產生操作輸出。至關重要的是，在“對每個具身進行token化”之后，HPT 在一小段潛token 序列的共享空間上運行。這種層次結構的靈感來源于人類如何在脊髓神經回路層面處理特定運動反應和感知刺激之間的反饋回路 [68]。

在具有跨具身的異構機器人學習中，數據來自不同的領域（例如模擬和真實機器人），跨越 RGB 圖像、語言指令、深度圖、3D 點云和觸覺圖像等傳感模態。每個機器人都是一個獨特的硬件具身，具有不同的自由度、末端執行器、傳感器配置、控制器和動作空間以及特定于應用程序的物理設置。

網絡架構模塊化為具身特定的stem、共享chunk和任務特定的頭部。直觀地講，如圖所示的stem，是神經網絡的早期層，可將來自異構具身和模態的傳感輸入對齊到共享表示空間中。網絡的共享中間部分稱為trunk，它將傳感表征處理為可用于多項任務的潛表征。最后，網絡的最后一部分是頭部，它將潛表示映射到感興趣的單個任務動作空間。訓練過程稱為異構預訓練，根據采樣的方案和任務數據分配和對齊特定的stem/頭部對，同時仍然享有在共享trunk中進行聯合訓練的好處。這可以被認為是使用神經網絡對每個實施方案進行token化，并減輕在標準訓練過程中將具身方案統一為同質數據形式的需要。

HPT 中的stem（上圖所示）由本體感受token化器和視覺token化器組成。這些token化器將來自不同具身的異構輸入映射到固定數量、固定維度的tokens，這使得stem能夠以相同的方式處理它們，盡管異構性很大，并且能夠享受固定上下文長度的擴展和推理優勢。關鍵思想是利用注意機制 [77, 30, 9] 將固定數量的可學習token關注到觀察的特征上。雖然主要關注本體感受和視覺，但處理觸覺、3D 和動作輸入中的其他類型傳感器異構，可以在stem中靈活擴展。

作為預訓練的核心組件，trunk架構遵循一個transformer，參數化為維度為 d 潛空間參數。輸出token序列長度 L 與輸入token序列長度相同。輸出token序列只是池化作為觀察的最終組合特征。trunk在不同的具身和任務之間共享，以捕獲復雜的輸入輸出關系（即trunk參數的數量是固定的，與具身和任務的數量無關）。

策略頭獲取stem transformer的輸出并將其映射到每個數據集中的動作空間 A。對于每個具身和任務，策略頭可以是任意架構（例如 MLP），它將trunk的池化特征作為輸入并輸出規范化的動作軌跡。策略頭重新初始化以遷移到新的具身。

其訓練過程有兩個數據擴展軸：一個數據集的數量，和數據集的總數。在預訓練階段，每次迭代只更新trunk參數，而每個異構具身和任務的stem和頭部則基于訓練批次采樣進行更新。

使用 AdamW [46] 優化器訓練 HPT，權重衰減率為 0.05，基本學習率為 0.0002，采用帶有預熱和丟棄的余弦學習率調度。根據批次大小應用基本學習率的比例縮放。為了在遷移學習期間支持各種范圍，在訓練期間沿每個批次的時間維度應用隨機掩碼。由于動作軌跡在預測范圍內可能會出現不平衡的損失，使用 Huber 損失。預訓練階段在學習率調度和優化器中的各種超參以及驗證數據集的選擇方面都是穩定的。代碼是開源的，預發布的模型可以從 Huggingface 輕松下載。

在實踐中，由于訓練損失可能因不同的數據集而異，并且目標是在所有具身和任務上表現良好，因此應用加權采樣程序進行數據加載。對于每次訓練迭代，都會以數據集大小的指數作為溫度，逆概率方式對數據集進行采樣。具體來說，計算每個數據集大小的平方根，并對這些大小求和以計算歸一化常數。然后，對于每個批次，以相應的概率從這些數據集中采樣。這可以防止大型數據集主導整個訓練周期，這是多任務學習中的常見做法。

注：每個具身的stem和頭部的更新頻率與trunk不同，類似于混合專家 [71] 訓練程序。特別是在分布式訓練設置下，每個stem和頭部都使用來自特定具身和任務的數據進行訓練，trunk將從訓練工作器的所有批次中積累梯度。這些預訓練實驗的計算資源范圍從 8 V-100 到 128 V-100，訓練時間從半天到 1 個月不等?？倲祿疟P大小約為 10Tb，RAM 內存要求低于 50Gb。

對于模擬基準，用已發布的數據集作為專家演示 [89, 81, 13]?？傊?，Metaworld [89] 使用腕部攝像頭視圖，Robomimic [49] 以及 Simpler [42] 使用第三人稱視圖，并通過數據集使用自己的本體感受定義。Fleet-Tools [81] 使用兩個視圖作為輸入，并使用末端執行器姿勢作為本體感受輸入。用預訓練的凍結 ResNet 特征對圖像進行編碼，并在將本體感受輸入傳遞到 stem 之前對其進行規范化。為除 Metaworld 之外的所有這些模擬基準訓練單任務策略。

對于 Simpler[42] 基準，專注于 Close-Drawer、Move Near 和 Pick Coke Can 任務以及具有視覺匹配設置的 Google EDR 實例。測試 9 種不同的初始化，總共 218 個episodes。注：模擬任務側重于語言調節，不暴露本體感受輸入，這不是最合適的 HPT 測試平臺。為了解決這些問題，在具有 79 條軌跡的 RTX 監督數據集上微調 HPT 作為其他模擬基準。用 HPT-base 作為本次實驗的主干（backbone）。

默認情況下，用 20000 次迭代進行訓練，批量大小為 512，學習率較小，為 1e 5。圖像和狀態主干是隱維度為 128 的單層 MLP，頭部是雙層 MLP。只使用長度為 1 的觀察窗口和 MLP 作為策略頭。每個訓練數據集每個任務使用 10-100 條軌跡，每個測試涵蓋具有不同初始條件的 50 個episodes。模擬中的每個軌跡在場景初始化方面略有不同。為了減少方差，進行 5 次獨立訓練，并取每個基線的平均值。

在真實世界中，在兩種不同的機器人設置中試驗機器人工具使用的任務 Sweep Leftover, Fill Water, Scoop Food 和 Switch Insertion。雖然在這兩種設置中，都使用 Franka Panda 作為機器人，但傳感器位置以及動作空間截然不同。為每個任務收集大約 100 個演示，并對每個任務進行 15 次試驗評估，以測量平均成功率。

在評估過程中，一個人始終監督機器人。評估 episode 可以因安全問題、機器人故障、超時等原因終止。如果episode 完成了任務，則認為它成功了。在 Fill Water 任務中，成功分數為 1 表示將一些水倒入碗中。在 Sweep Leftover 任務中，成功分數為 1 表示將所有推入盤子中，成功分數為 0.5 表示將一些推入盤子中。在Scoop Food 任務中，成功分數為 1 表示舀出一些狗糧并全部倒入碗中，分數為 0.5 表示只舀出一些。在Switch Insertion 任務 [90] 中，成功分數為 1 表示開關精確插入 PCB 板上的三個引腳。機器人在嘗試插入之前會移動到預定義的姿勢。選擇這些具有挑戰性的任務，是因為它們需要與工具和顆粒體進行豐富的接觸交互，并且需要高精度和密集的接觸。確保機器人的初始條件相同。由于這些任務的復雜性和人為錯誤，目標設置的初始條件并不完全相同。

遷移學習。對于現實世界實驗中的策略頭，嘗試 MLP 和擴散策略 [13]。與其他方法中通常用于單個任務的更大模型（例如 100M）相比，微調的活動參數不超過 3Mb。用長度為 2 的觀察歷史窗和較小的學習率 2e 5。在單個 NVIDIA RTX 2080Ti GPU 上，以批量大小 256 進行 20000 次迭代（大約 4 小時時間）。

Scaling Proprioceptive-Visual Learning with Heterogeneous Pre-trained Transformers 2409.20537v1.pdf

END

轉載自CSDN-三谷秋水