[論文レビュー] EDGI: Equivariant Diffusion for Planning with Embodied Agents

Johann Brehmer, Joey Bose+2 more|arXiv (Cornell University)|Mar 22, 2023

Reinforcement Learning in Robotics被引用数 4

ひとこと要約

EDGIは、空間的（SE(3)）、時間的（Z）および物体置換（S_n）対称性の積群に対して不変である、身体的エージェントの計画に向けた新規な等長性拡散モデルを導入する。このモデルを条件付き生成計画フレームワークに統合し、分類器ガイドによる柔ららかな対称性破壊を施すことで、非等長性ベースラインと比較して顕著に高いサンプル効率と対称性変換における一般化性能を達成し、最大10倍少ない学習データで優れた性能を発揮する。

ABSTRACT

Embodied agents operate in a structured world, often solving tasks with spatial, temporal, and permutation symmetries. Most algorithms for planning and model-based reinforcement learning (MBRL) do not take this rich geometric structure into account, leading to sample inefficiency and poor generalization. We introduce the Equivariant Diffuser for Generating Interactions (EDGI), an algorithm for MBRL and planning that is equivariant with respect to the product of the spatial symmetry group SE(3), the discrete-time translation group Z, and the object permutation group Sn. EDGI follows the Diffuser framework (Janner et al., 2022) in treating both learning a world model and planning in it as a conditional generative modeling problem, training a diffusion model on an offline trajectory dataset. We introduce a new SE(3)xZxSn-equivariant diffusion model that supports multiple representations. We integrate this model in a planning loop, where conditioning and classifier guidance let us softly break the symmetry for specific tasks as needed. On object manipulation and navigation tasks, EDGI is substantially more sample efficient and generalizes better across the symmetry group than non-equivariant models.

研究の動機と目的

• 空間的、時間的、置換的対称性を明示的に符号化することで、身体的エージェントのモデルベース強化学習におけるサンプル効率と一般化性能を向上させること。
• SE(3)×Z×S_nに対して不変性を保ちつつ、テスト時に柔軟にタスク固有の対称性破壊が可能な、拡散ベースの計画フレームワークを開発すること。
• 3次元座標、速度、物体ラベルなどの複数の入力表現を、1つの等長性アーキテクチャ内で統合し、現実世界のロボットタスクにおける柔軟性を高めること。
• 等長性モデリングが、対称性変換および低データ環境下でも強固な性能を発揮することを実証すること。
• 積群上に等長性モデルを構築する一般的な枠組みを提供し、複数の入力表現をサポートすること。

提案手法

• EDGIは、オフライントラジェクトリ上で学習された拡散モデルを採用し、SE(3)×Z×S_nの全積群に対して等長性を示すノイズ除去ネットワークを用いる。
• グループ構造を尊重し、異なる対称性タイプに一貫して作用する、新規の時間的・物体的・置換層を導入する。
• 複数の入力表現（例：3次元座標、速度、物体インデックス）を統合する新しい埋め込み機構を提案し、等長性を保持する。
• 条件付きサンプリングにより計画を実行し、タスク固有の条件付けと分類器ガイドを用いて柔らかな対称性破壊を実現する。
• 等長性拡散モデルを計画ループに統合し、ワールドモデル学習と軌道生成の両方をサポートする。
• 等長性を損なわず、複数のデータ表現をサポートするアーキテクチャを実現し、複雑な環境への広範な適用性を高める。

Figure 1: Schematic of EDGI in a navigation task, where the agent (red square) plans the next actions (red dots) to reach the goal (green star) without touching obstacles (grey circles). Top : planning as conditional sampling from a diffusion model. Bottom : effect of a group action. Equivariance re

実験結果

リサーチクエスチョン

• RQ1 モデルアーキテクチャにSE(3)×Z×S_nの対称性を明示的に符号化することで、拡散ベースの計画アルゴリズムがサンプル効率を向上させられるか？
• RQ2 空間的・時間的・置換的対称性における等長性が、未学習の対称性変換を伴う環境への一般化性能に与える影響はいかほどか？
• RQ3 条件付けと分類器ガイドによる柔らかな対称性破壊は、等長性を損なわず、タスク固有の計画をどの程度可能にするか？
• RQ4 複数表現入力サポートの統合は、複雑なロボット操作およびナビゲーションタスクにおける性能向上に寄与するか？
• RQ5 等長性拡散モデルは、非等長性ベースラインと比較して、顕著に少ない学習データでも高い性能を維持できるか？

主な発見

• EDGIは、学習データの10％のみで学習した場合でも、ナビゲーションおよび物体操作タスクにおいて、最良の非等長性ベースラインと同等の性能を達成する。
• 低データ環境下では、EDGIが非等長性モデルを顕著に上回り、優れたサンプル効率を示す。
• EDGIは、回転や物体配置の置換を伴う未学習の対称性変換環境に対しても、強固な一般化性能を発揮する。
• Kuka操作ベンチマークの3つのタスクすべてにおいて、EDGIは対称性拡張済みのテスト条件下でも強力な性能を維持する。
• 分類器ガイドの統合により、効果的な柔らかな対称性破壊が実現され、モデルはタスク固有の解決を可能にしつつ、下位のダイナミクスにおける不変性を保持する。
• 実験的結果から、等長性がデータ効率性と耐障害性の向上に寄与することが確認され、本手法の核心仮説が妥当であることが裏付けられる。

Figure 2: Architecture of our $\mathrm{SE(3)}\times\mathbb{Z}\times\mathrm{S}_{n}$ -equivariant denoising network. Input trajectories (top left), which consist of features in different representations of the symmetry group, are first transformed into a single internal representation (green block). T