ReMix: Optimizing Data Mixtures for Large Scale Imitation Learning

作者: Joey Hejna, Chethan Bhateja, Yichen Jiang, Karl Pertsch, Dorsa Sadigh · CoRL 2024

授權條款

• 開放取用
• 論文: CC BY 4.0

以下內容皆依據論文全文，依其 CC BY 4.0 授權條款使用。

一句話摘要

ReMix 對策略的超額損失使用群組分布強健最佳化，學習每個機器人資料集該取樣多少，在真實機器人上勝過均勻權重與 RT-X 專家權重。

為何重要

以彙整的機器人資料進行預訓練的團隊，可以用從資料計算出的權重取代人工調整的取樣權重，而且在報告的測試中，可以捨棄 75% 的資料而表現損失有限。學到的權重也能顯示哪些資料集應該降低權重及其原因，例如自主蒐集的 Kuka 資料，這有助於決定要蒐集或保留哪些資料。

問題

Open X-Embodiment 等大型機器人資料集彙整了來自不同機器人、相機、動作空間與控制器的資料，但用來訓練 RT-X、Octo 與 OpenVLA 等通用策略的資料配比都是人工挑選的。來自視覺與語言領域的資料篩選方法難以轉用，因為 n-gram 並不適用，而視覺嵌入會忽略動作以及回合的序列性質。

提出的方法

ReMix 把資料整理表述為一個最小-最大問題：離散策略最小化其相對於參考模型的超額行為複製損失，而領域權重則以指數梯度上升更新以最大化該損失，使權重流向策略仍可改進的領域。為了讓不同機器人之間的損失可以比較，動作會按領域做高斯正規化並離散化成數個區間，參考檢查點則以提前停止在它對任何領域過度擬合之前選出。接著，訓練過程中平均的權重會被重新用來訓練另外的擴散策略，或對資料進行子取樣。

關鍵見解

• 使用相對於參考策略的超額損失而非原始損失，能同時降低已經擬合好的領域與難以擬合的領域的權重。
• 動作尺度的差異會讓強健最佳化產生偏差：在一項使用邊界正規化的簡化測試中，94.3% 的權重落到一個純隨機動作的領域；改用逐領域的高斯正規化時，真實資料獲得 84.2% 的權重。
• 連續的 L1/L2 動作損失容易受離群值與多模態動作左右，因此即使最終策略是連續的擴散模型，將動作離散化仍能得到更可靠的權重。
• 機器人資料集（約 10k-100k 次示範）小到足以讓參考模型過度擬合；使用提前停止點之後 50K 步的檢查點，會讓表現平均下降超過 15%。

主要結果

• 在 RT-X 資料配比上，以 ReMix 權重訓練的策略，在四項 WidowX 與兩項 Franka 真實世界任務上，平均比均勻權重高 38%，比 RT-X 人類專家權重高 32%；專家權重則比均勻權重高 6%。
• 在 RT-X 配比的 25% 子集上，ReMix 維持了表現，而人工權重下降超過 10%；ReMix 保留了 72% 的 UR5 資料與 12% 的 Kuka 資料，人工權重則分別為 30% 與 24%。
• ReMix 把小型 Toto 資料集的比例從 3.42%（均勻）提高到 16.3%，並把 Kuka 從 24.9% 降到 12.1%。
• 在完整的 Bridge V2 資料集上，ReMix 與均勻權重之間沒有明顯差異；在 25% 的 Bridge 子集上，ReMix 優於均勻權重與 SSP 影像嵌入基準。
• 在帶有干擾物體的 WidowX 方塊任務上，ReMix 權重成功 7/10 次，人工與均勻權重皆為 0/10。
• 在模擬中（依操作員劃分的 Robomimic NutAssemblySquare），在 50% 子取樣下 ReMix 得分 77/100，均勻權重為 53/100；在 25% 下為 59/100 對 39/100。

限制

• 作者指出，真實世界評估只涵蓋 WidowX 與 Franka 兩種手臂，更多機器人形態留待未來研究。
• ReMix 會提高動作分布異常、多模態之資料集（例如 Toto）的權重，作者表示這不一定是理想的結果。
• 計算權重需要在完整資料上訓練兩次，一次用於參考模型，一次用於強健最佳化。
• 只在 Bridge V2 與 OpenX 的 RT-X 子集上測試，並未在超過 2M 個回合的完整 OpenX 資料集上測試。第 4.1 節描述 Bridge V2 有 50k 次示範，而附錄 B.2 則說從完整資料集的 60K 條軌跡中使用了約 45,000 條。
• 只以預訓練資料訓練的策略無法零樣本泛化到作者的設置，因此每項評估都仰賴每項任務以 25 次領域內示範進行共同訓練。

對領域的影響

入圍 CoRL 2024 傑出論文獎決選，被引用 74 次。作者在 github.com/jhejna/remix 釋出程式碼。論文也公布了為 RT-X 配比學到的領域權重（表 2），以及 32 個 Bridge V2 領域的權重（表 5）。

摘要

為了訓練機器人基礎模型，人們正在蒐集規模越來越大的模仿學習資料集。然而，儘管資料選擇在視覺與自然語言處理中極為重要，機器人領域卻很少有研究探討這類模型實際上應該以什麼資料來訓練。在本研究中，我們探討在機器人基礎模型預訓練中，應如何為機器人資料集的不同子集或「領域」分配權重。具體而言，我們使用分布強健最佳化（DRO）來最大化所有可能下游領域中最差情況的表現。我們的方法 Re-Mix 處理了將 DRO 應用於機器人資料集時出現的各種挑戰，包括不同資料集之間動作空間與動力學的差異。Re-Mix 採用提前停止、動作正規化與離散化來因應這些問題。透過在最大的開源機器人操作資料集 Open X-Embodiment 上進行大量實驗，我們證明資料整理會對下游表現造成極大影響。具體而言，在用來訓練現有通用機器人策略（也就是 RT-X 模型）的資料集上，Re-Mix 學到的領域權重平均比均勻權重高 38%，比人工選擇的權重高 32%。

連結

• arXiv 2408.14037: https://arxiv.org/abs/2408.14037
• DOI 10.48550/arXiv.2408.14037: https://doi.org/10.48550/arXiv.2408.14037
• PDF: https://raw.githubusercontent.com/mlresearch/v270/main/assets/hejna25a/hejna25a.pdf
• 程式碼: https://github.com/jhejna/remix