ReMix: Optimizing Data Mixtures for Large Scale Imitation Learning

저자: Joey Hejna, Chethan Bhateja, Yichen Jiang, Karl Pertsch, Dorsa Sadigh · CoRL 2024

라이선스

오픈 액세스
논문: CC BY 4.0

아래 내용은 모두 논문의 라이선스(CC BY 4.0)에 따라 논문 전문을 바탕으로 합니다.

요약

ReMix는 정책의 초과 손실에 그룹 분포 강건 최적화를 적용하여 각 로봇 데이터셋을 얼마나 샘플링할지 학습하며, 실제 로봇에서 균등 가중치와 RT-X 전문가 가중치를 앞선다.

중요한 이유

통합된 로봇 데이터로 사전 학습하는 팀은 수작업으로 조정한 샘플링 가중치를 데이터에서 계산한 가중치로 대체할 수 있으며, 보고된 시험에서는 성능 손실을 제한하면서 데이터의 75%를 버릴 수 있다. 학습된 가중치는 또한 어떤 데이터셋의 가중치를 낮춰야 하는지와 그 이유를 보여 주는데, 예를 들어 자율적으로 수집된 Kuka 데이터가 그러하며, 이는 어떤 데이터를 수집하거나 유지할지 결정하는 데 도움이 된다.

문제

Open X-Embodiment 같은 대규모 로봇 데이터셋은 서로 다른 로봇, 카메라, 행동 공간, 제어기의 데이터를 모으지만, RT-X, Octo, OpenVLA 같은 범용 정책을 학습하는 데 쓰인 혼합 비율은 수작업으로 선택되었다. 비전과 언어의 데이터 필터링 방법은 잘 옮겨지지 않는데, n-gram은 적용되지 않고 시각 임베딩은 행동과 에피소드의 순차적 특성을 무시하기 때문이다.

제안 방법

ReMix는 큐레이션을 min-max 문제로 정식화한다. 이산 정책은 참조 모델 대비 초과 행동 복제 손실을 최소화하고, 도메인 가중치는 이를 최대화하도록 지수화 경사 상승으로 갱신되므로, 정책이 아직 개선할 수 있는 도메인에 가중치가 간다. 로봇 간에 손실을 비교할 수 있도록 행동은 도메인별로 가우시안 정규화되고 구간으로 이산화되며, 참조 체크포인트는 어떤 도메인에도 과적합되기 전에 조기 종료로 선택된다. 학습 전반에 걸쳐 평균한 가중치는 이후 별도의 확산 정책을 학습하거나 데이터를 부분 샘플링하는 데 재사용된다.

핵심 통찰

원시 손실 대신 참조 정책 대비 초과 손실을 사용하면, 이미 적합된 도메인과 적합하기 너무 어려운 도메인 모두의 가중치가 낮아진다.
행동 스케일의 차이는 강건 최적화를 왜곡한다. 경계 정규화를 사용한 장난감 시험에서는 가중치의 94.3%가 순수 무작위 행동 도메인으로 갔지만, 도메인별 가우시안 정규화를 사용하자 실제 데이터가 84.2%를 받았다.
연속 L1/L2 행동 손실은 이상치와 다봉 행동에 휘둘리므로, 최종 정책이 연속 확산 모델이더라도 행동을 이산화하면 더 신뢰할 수 있는 가중치를 얻는다.
로봇 데이터셋(약 10k-100k개 시연)은 참조 모델이 과적합될 만큼 작다. 조기 종료 시점보다 50K 스텝 뒤의 체크포인트를 사용하자 성능이 평균 15% 이상 떨어졌다.

주요 결과

RT-X 데이터 혼합에서 ReMix 가중치로 학습한 정책은 네 가지 WidowX 및 두 가지 Franka 실제 과제 전반에 걸쳐 평균적으로 균등 가중치보다 38%, RT-X 인간 전문가 가중치보다 32% 뛰어났다. 전문가 가중치는 균등 가중치보다 6% 나았다.
RT-X 혼합의 25% 부분집합에서 ReMix는 성능을 유지한 반면 인간 가중치는 10% 넘게 떨어졌다. ReMix는 UR5 데이터의 72%와 Kuka의 12%를 유지했고, 인간 가중치는 각각 30%와 24%였다.
ReMix는 작은 Toto 데이터셋을 3.42%(균등)에서 16.3%로 올리고 Kuka를 24.9%에서 12.1%로 줄였다.
전체 Bridge V2 데이터셋에서는 ReMix와 균등 가중치 사이에 눈에 띄는 차이가 없었다. 25% Bridge 부분집합에서는 ReMix가 균등 가중치와 SSP 이미지 임베딩 기준선보다 뛰어났다.
방해 물체가 있는 WidowX 큐브 과제에서 ReMix 가중치는 7/10 성공한 반면, 인간 가중치와 균등 가중치는 모두 0/10이었다.
시뮬레이션(조작자별로 나눈 Robomimic NutAssemblySquare)에서 ReMix는 50% 부분 샘플링 시 77/100 대 균등 53/100, 25%에서는 59/100 대 39/100을 기록했다.

한계

저자들은 실제 환경 평가가 WidowX와 Franka 두 팔만 다루었다고 언급하며, 더 많은 형태의 로봇은 향후 과제로 남긴다.
ReMix는 Toto처럼 비정상적이고 다봉적인 행동 분포를 가진 데이터셋의 가중치를 높이는데, 저자들은 이것이 반드시 바람직하지는 않다고 말한다.
가중치를 계산하려면 전체 데이터로 두 번 학습해야 한다. 한 번은 참조 모델을 위해, 한 번은 강건 최적화를 위해서이다.
Bridge V2와 OpenX의 RT-X 부분집합에서 시험했으며, 2M개가 넘는 에피소드로 이루어진 전체 OpenX 데이터셋에서는 시험하지 않았다. 4.1절은 Bridge V2를 50k개 시연으로 기술하는 반면, 부록 B.2는 전체 데이터셋의 60K개 중 약 45,000개 궤적을 사용했다고 말한다.
사전 학습 데이터로만 학습한 정책은 저자들의 구성에 제로샷으로 일반화되지 않았으므로, 모든 평가는 과제당 도메인 내 시연 25개와의 공동 학습에 의존한다.

분야에 미친 영향

CoRL 2024 Outstanding Paper Award 최종 후보였으며 74회 인용되었다. 저자들은 github.com/jhejna/remix 에서 코드를 공개한다. 논문은 또한 RT-X 혼합(표 2)과 32개 Bridge V2 도메인(표 5)에 대해 학습된 도메인 가중치를 공개한다.

초록

로보틱스를 위한 파운데이션 모델을 학습시킬 목적으로 점점 더 큰 모방 학습 데이터셋이 수집되고 있다. 그러나 데이터 선택이 비전과 자연어 처리에서 매우 중요했음에도, 로보틱스에서는 이러한 모델이 실제로 어떤 데이터로 학습되어야 하는지를 묻는 연구가 거의 없었다. 본 연구에서 우리는 로봇 파운데이션 모델 사전 학습을 위해 로보틱스 데이터셋의 서로 다른 부분집합, 즉 “도메인”에 어떻게 가중치를 줄지 조사한다. 구체적으로, 가능한 모든 다운스트림 도메인에 걸쳐 최악의 경우 성능을 최대화하기 위해 분포 강건 최적화(DRO)를 사용한다. 우리 방법 Re-Mix는 데이터셋 간 행동 공간과 동역학의 가변성을 포함하여, 로보틱스 데이터셋에 DRO를 적용할 때 생기는 광범위한 문제를 다룬다. Re-Mix는 이러한 문제에 대응하기 위해 조기 종료, 행동 정규화, 이산화를 사용한다. 최대 규모의 오픈 소스 로봇 조작 데이터셋인 Open X-Embodiment 데이터셋에 대한 광범위한 실험을 통해, 데이터 큐레이션이 다운스트림 성능에 매우 큰 영향을 미칠 수 있음을 입증한다. 구체적으로, 기존 범용 로봇 정책, 특히 RT-X 모델을 학습하는 데 사용된 데이터셋에서 Re-Mix가 학습한 도메인 가중치는 균등 가중치보다 평균 38%, 사람이 선택한 가중치보다 32% 뛰어나다.

링크

arXiv 2408.14037: https://arxiv.org/abs/2408.14037
DOI 10.48550/arXiv.2408.14037: https://doi.org/10.48550/arXiv.2408.14037
PDF: https://raw.githubusercontent.com/mlresearch/v270/main/assets/hejna25a/hejna25a.pdf
코드: https://github.com/jhejna/remix