Equivariant Diffusion Policy

作者: Dian Wang, Stephen Hart, David Surovik, Tarik Kelestemur, Haojie Huang, Haibo Zhao, Mark Yeatman, Jiuguang Wang, Robin Walters, Robert Platt · CoRL 2024

授權條款

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

以下內容皆依據論文全文，依其 CC BY 4.0 授權條款使用；圖表已標示作者。

一句話摘要

Equivariant Diffusion Policy 把 SO(2) 旋轉對稱性內建到 Diffusion Policy 的去雜訊器中，在 100 次示範時讓模擬平均成功率提高 21.9 個百分點，並能從 20-60 次示範學會真實任務。

為何重要

真實示範的成本很高，而內建桌面旋轉對稱性讓策略在六項真實任務上成功，其中包括一項長時程的烤貝果任務，只需 20 到 60 次示範，而體素 Diffusion Policy 基準在這些任務上大多失敗。由於分析涵蓋完整的 SE(3) 夾爪動作，它適用於需要平面外旋轉的任務，而 SE(2) 版本無法解決這類任務。

問題

Diffusion Policy 能捕捉多模態的示範行為，但它必須在觀測、帶雜訊的動作與雜訊程度上學習去雜訊函數，這比學習從狀態直接對應到動作的策略更困難，也需要更多資料。先前的等變操作策略不是使用開迴路的 SE(3) 動作空間，就是使用閉迴路的 SE(2) 動作空間，而且都沒有把對稱性應用到擴散策略上。

提出的方法

作者證明，只要專家策略具有 SO(2) 等變性，雜訊預測函數就具有 SO(2) 等變性，並推導出在絕對與相對姿態控制下，繞重力軸的旋轉如何作用於完整的 6-DoF 夾爪動作。去雜訊器使用等變的觀測與動作編碼器（以 escnn 函式庫建構），為離散旋轉群產生嵌入，對每個群元素套用共用的 1D 時間 U-Net，並使用等變的雜訊解碼器。輸入可以是 RGB 影像或 64×64×64 的體素網格。

關鍵見解

• 若專家策略具有 SO(2) 等變性，擴散的雜訊預測函數也會具有等變性，這使等變網路成為去雜訊器的正確歸納偏置。
• 在示範很少、以及初始物體姿態隨機範圍很大的任務上，增益最大；在隨機性低的 Robomimic 任務上，此方法與 Diffusion Policy 相當或略好。
• 在消融實驗中，等變結構比體素輸入更重要：在 100 次示範時，移除等變性讓平均成功率下降 17.6 個百分點，移除體素則下降 10.3 個百分點。
• 在標準 CNN 上使用旋轉資料擴增，對 Stack 等簡單任務有幫助，但在較困難的任務上遠不如等變網路。
• 完整的 SE(3) 動作空間很重要：SE(2) 版本在 Coffee Preparation 上得分 0.0%，SE(3) 則為 85.3%；在 Threading 上為 12.7% 對 40.0%。

主要結果

• 在 12 個採用絕對控制、100 次示範的 MimicGen 任務上，體素版本平均成功率為 63.9%，Diffusion Policy（DiffPo-C）為 42.0%，提高 21.9 個百分點；RGB 版本平均為 53.7%。
• 在 200 次示範時，體素版本（72.6%）勝過所有以 1000 次示範訓練的基準（最佳為 71.4%）。
• 採用絕對控制時，體素版本在 12 個環境中的 11 個優於基準，唯一例外是 Hammer Cleanup D1。
• 在真實 Franka 手臂上，與體素 Diffusion Policy 基準相比的成功率為：Oven Opening 95% 對 60%（20 次示範）；Banana in Bowl 95% 對 30%、Letter Alignment 95% 對 0%、Trash Sweeping 90% 對 5%（各 40 次示範）；Hammer to Drawer 85% 對 5%（60 次示範）；Bagel Baking 80% 對 10%（58 次示範）。
• Bagel Baking 的所有失敗都來自機器人的關節極限，而且策略能零樣本泛化到未見過的烤箱旋轉角度，只有一個被關節極限擋住的姿態例外。

限制

• 作者指出，由於視覺系統中的對稱性不一致，等變性只被部分利用；出現在體素網格中的手臂與相機雜訊都可能破壞對稱性。
• 當模型的對稱性與示範互相衝突時，「錯誤的等變性」可能損害表現。
• SO(2) 動作理論只在擴散策略上展示，並未在其他流程上展示，例如在相對控制下表現良好的 BC RNN。
• 此方法尚未擴展到導航、行走運動或移動式操作（mobile manipulation）。
• 使用 RGB 輸入時，agent-view 相機並非由上往下拍攝，因此影像旋轉與真實狀態的旋轉並不一致；採用相對控制時，RGB 版本只比基準略好。

對領域的影響

入圍 CoRL 2024 傑出論文獎決選，被引用 107 次。論文指向專案頁面 equidiff.github.io，內文並未列出程式碼儲存庫。

摘要

近期研究顯示，在行為複製中，擴散模型是學習示範資料所呈現之多模態分布的有效方法。然而，這種方法的缺點之一是需要學習一個去雜訊函數，這比學習一個顯式策略複雜得多。在本研究中，我們提出 Equivariant Diffusion Policy，這是一種新穎的擴散策略學習方法，運用領域對稱性，讓去雜訊函數獲得更好的樣本效率與泛化能力。我們從理論上分析完整 6-DoF 控制的 SO(2) 對稱性，並刻畫擴散模型在何種條件下具有 SO(2) 等變性。我們進一步在 MimicGen 的一組 12 個模擬任務上實證評估此方法，結果顯示其成功率平均比基準 Diffusion Policy 高 21.9%。我們也在真實世界系統上評估此方法，證明它能以相對較少的訓練樣本學到有效的策略，而基準 Diffusion Policy 則無法做到。

連結

• arXiv 2407.01812: https://arxiv.org/abs/2407.01812
• DOI 10.48550/arXiv.2407.01812: https://doi.org/10.48550/arXiv.2407.01812
• PDF: https://raw.githubusercontent.com/mlresearch/v270/main/assets/wang25a/wang25a.pdf
• 專案頁面: https://equidiff.github.io/