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개의 시연으로 실제 과제를 학습한다.

중요한 이유

실제 시연은 비용이 크며, 탁상 회전 대칭성을 내장하자 정책은 긴 호라이즌의 베이글 굽기 과제를 포함한 여섯 가지 실제 과제에서 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) 등변이면 확산 노이즈 예측 함수도 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의 실패는 모두 로봇 관절 한계에서 비롯되었으며, 정책은 관절 한계로 막힌 자세 하나를 제외하고 처음 보는 오븐 회전에 제로샷으로 일반화되었다.

한계

저자들은 시각 시스템의 대칭 불일치 때문에 등변성이 부분적으로만 활용된다고 언급한다. 복셀 격자에 나타나는 팔과 카메라 노이즈가 대칭성을 깨뜨릴 수 있다.
"잘못된 등변성"은 모델의 대칭성이 시연과 충돌할 때 성능을 떨어뜨릴 수 있다.
SO(2) 행동 이론은 확산 정책에서만 제시되었으며, 상대 제어에서 좋은 성능을 보인 BC RNN 같은 다른 파이프라인에서는 제시되지 않았다.
방법은 아직 내비게이션, 보행, 모바일 매니퓰레이션으로 확장되지 않았다.
RGB 입력의 경우 에이전트 시점 카메라가 위에서 내려다보는 방향이 아니므로 이미지 회전이 실제 상태 회전과 일치하지 않는다. 상대 제어에서 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/