[논문 리뷰] RDT-1B: a Diffusion Foundation Model for Bimanual Manipulation

Songming Liu, Lingxuan Wu외 7명|arXiv (Cornell University)|2024. 10. 10.

Mechanics and Biomechanics Studies인용 수 5

한 줄 요약

RDT-1B는 Robotics Diffusion Transformer (RDT)로, 언어 조건의 양손 조작을 위한 1.2B-parameter diffusion-based foundation model이며, 대규모 다중 로봇 데이터로 사전 학습하고 다중 태스크 양손 데이터셋에서 미세 조정되어 실제 로봇에서 강한 제로샷 및 소샷 일반화를 달성합니다.

ABSTRACT

Bimanual manipulation is essential in robotics, yet developing foundation models is extremely challenging due to the inherent complexity of coordinating two robot arms (leading to multi-modal action distributions) and the scarcity of training data. In this paper, we present the Robotics Diffusion Transformer (RDT), a pioneering diffusion foundation model for bimanual manipulation. RDT builds on diffusion models to effectively represent multi-modality, with innovative designs of a scalable Transformer to deal with the heterogeneity of multi-modal inputs and to capture the nonlinearity and high frequency of robotic data. To address data scarcity, we further introduce a Physically Interpretable Unified Action Space, which can unify the action representations of various robots while preserving the physical meanings of original actions, facilitating learning transferrable physical knowledge. With these designs, we managed to pre-train RDT on the largest collection of multi-robot datasets to date and scaled it up to 1.2B parameters, which is the largest diffusion-based foundation model for robotic manipulation. We finally fine-tuned RDT on a self-created multi-task bimanual dataset with over 6K+ episodes to refine its manipulation capabilities. Experiments on real robots demonstrate that RDT significantly outperforms existing methods. It exhibits zero-shot generalization to unseen objects and scenes, understands and follows language instructions, learns new skills with just 1~5 demonstrations, and effectively handles complex, dexterous tasks. We refer to https://rdt-robotics.github.io/rdt-robotics/ for the code and videos.

연구 동기 및 목표

• 대형 규모의 다중 로봇 데이터를 사전 학습에 활용하고 대상 로봇 데이터로 미세 조정하여 양손 조작의 데이터 부족 문제를 해소한다.
• 텍스트, 이미지, 및 고유감각으로부터의 이질적 입력을 처리하고 양손 작용의 다중 모달리티를 모델링할 수 있는 확장 가능한 diffusion-based 아키텍처를 개발한다.
• 물리적 의미를 보존하면서 로봇 간 행동 표현을 통일하기 위한 물리적으로 해석 가능한 통합 행동 공간을 도입한다.
• 실제 이중 팔 로봇에서 제로샷 및 소샷 역량, 언어 지시 따르기, 정교한 조작을 포함한 강한 일반화를 시연한다.

제안 방법

• 다중 모달리티를 포착하기 위해 denoising diffusion 프로세스를 이용한 연속적 조건 분포 p(a_t|l,o_t)로 행동을 모델링한다.
• DiT(Diffusion Transformer) 백본을 사용하고 로봇 데이터 특성에 맞춘 아키텍처 적응(MLP 디코더, QKNorm, RMSNorm, 교대 조건 주입)을 적용한다.
• 저차원 고유감각을 갖는 MLP와 Fourier 특징을 통해 이질적 입력을 인코딩하고, 이미지는 비전 인코더(SigLIP)로, 언어는 사전 학습된 Transformer(T5-XXL)로 인코딩한다; 모달리티 의존 과다를 방지하기 위해 입력 마스킹을 적용한다.
• 동일한 시퀀스의 시간적 일관성을 촉진하고 오차 누적을 줄이기 위해 action chunks (a_t:t+T_a)에 대한 확산 기반 디노이징 손실을 사용하여 학습한다.
• 이질적인 로봇 액션 공간을 Physically Interpretable Unified Action Space로 통합하여 46개 데이터셋(≈1M trajectories, 21TB)에서 다중 로봇 사전 학습을 가능하게 한다.
• 대규모 다중 로봇 데이터에서 RDT를 1.2B 매개변수로 사전 학습하고, 대상 이중 팔 조작을 위한 자체 수집 다중 태스크 양손 데이터셋(>6K trajectories)에서 미세 조정한다.
• 샘플링 가속을 위해 DPM-Solver++를 활용하여 6 Hz의 action-chunk 추론과 하드웨어에서 초당 높은 액션 처리량을 가능하게 한다.

실험 결과

연구 질문

• RQ1 RDT가 미지의 물체와 장면에 대해 제로샷으로 일반화할 수 있는가?
• RQ2 미지의 모달리티에 대한 RDT의 제로샷 지시 수행 능력은 얼마나 효과적인가?
• RQ3 이전에 보지 못한 기술에 대해 소샷 학습을 가능하게 하는가?
• RQ4 정교하고 섬세한 동작을 요구하는 작업을 수행할 수 있는가?
• RQ5 모델 크기, 데이터 규모, 확산 모델링이 성능 향상에 기여하는가?

주요 결과

• RDT는 다수의 양손 작업에서 최첨단 성능을 달성하고 베이스라인 대비 큰 격차로 우수한 성능을 보인다(예: 성공률 56% 향상).
• RDT는 미지의 물체, 장면, 지시 및 기술에 대해 제로샷 및 소샷(1–5샷) 일반화를 보여준다.
• 대형 모델 크기와 광범위한 사전 학습 데이터가 확산 모델링과 함께 우수한 성능에 기여한다.
• RDT는 실제 로봇에서 섬세하고 정교한 작업을 수행할 수 있으며 언어 지시를 효과적으로 따른다.