[논문 리뷰] EDGI: Equivariant Diffusion for Planning with Embodied Agents

Johann Brehmer, Joey Bose외 2명|arXiv (Cornell University)|2023. 03. 22.

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에 대해 불변성을 유지하면서도 테스트 시에 유연하고 작업 기반의 대칭성 파괴를 허용하는 확산 기반 계획 프레임워크를 개발하는 것.
• 더 큰 유연성을 확보하기 위해 단일 등변성 아키텍처 내에서 다수의 입력 표현(예: 위치, 속도, 물체 레이블)을 지원하는 것.
• 등변성 모델링이 대칭 변환과 저자료 환경에서 강건한 성능을 낼 수 있음을 입증하는 것.
• 다중 표현 입력 지원 기능을 갖춘 제품군 위에서의 등변성 모델을 구축하는 일반적인 블루프린트를 제공하는 것.

제안 방법

• EDGI는 오프라인 트레이젝터리에서 훈련된 확산 모델을 사용하며, 이 모델의 노이즈 제거 네트워크는 전체 곱군 SE(3)×Z×S_n에 대해 등변성이다.
• 모델은 그룹의 구조를 존중하고 다양한 대칭 유형 간에 일관되게 작용하는 새로운 시간, 물체, 순열 레이어를 사용한다.
• 다양한 입력 표현(예: 3차원 좌표, 속도, 물체 인덱스)을 공통 내부 표현으로 통합하는 새로운 임bedding 메커니즘이 등변성을 유지한다.
• 계획은 조건부 샘플링을 통해 확산 모델에서 수행되며, 작업 기반 조건화와 분류기 가이던스를 통해 대칭성을 부드럽게 파괴한다.
• 등변성 확산 모델을 계획 루프에 통합하여 월드 모델 학습과 궤적 생성을 모두 지원한다.
• 등변성을 포기하지 않으면서도 다수의 데이터 표현을 지원하여, 더 복잡한 환경에서의 광범위한 적용 가능성을 확보한다.

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 조작 벤치마크의 세 가지 모든 작업에서 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