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 は Diffusion Policy のデノイザに SO(2) の回転対称性を組み込み、100 デモ時のシミュレーションの平均成功率を 21.9 ポイント高め、20-60 回のデモから実機タスクを学習する。

重要な理由

実機のデモンストレーションは高価であり、テーブル上の回転対称性を組み込むことで、ボクセル版 Diffusion Policy のベースラインがほぼ失敗したところを、長期的なベーグル焼きタスクを含む 6 つの実機タスクで 20 から 60 回のデモにより成功させた。解析は完全な 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% であった。

主な結果

絶対制御、100 デモの MimicGen の 12 タスクにおいて、ボクセル版の平均成功率は 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 の失敗はすべてロボットの関節可動域の限界によるものであり、ポリシーは関節可動域に阻まれた 1 つの姿勢を除き、未知のオーブンの回転にゼロショットで汎化した。

制限事項

著者らは、視覚システムにおける対称性の不一致のため同変性は部分的にしか活用されていないと述べている。ボクセルグリッドに写り込むアームやカメラのノイズが対称性を崩しうる。
モデルの対称性がデモンストレーションと矛盾する場合、「誤った同変性」が性能を損なうことがある。
SO(2) の行動理論は拡散ポリシーについてのみ示されており、相対制御で良好な性能を示した BC RNN など他のパイプラインについては示されていない。
本手法はまだナビゲーション、ロコモーション、モバイルマニピュレーションには拡張されていない。
RGB 入力の場合、agent-view カメラは真上からの視点ではないため、画像の回転は実際の状態の回転と一致しない。相対制御では、RGB 版はベースラインよりわずかに優れていたにすぎない。

分野への影響

CoRL 2024 Outstanding Paper Award のファイナリストであり、被引用数は 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/