RoboDreamer: Learning Compositional World Models for Robot Imagination

Text-to-video models have demonstrated substantial potential in robotic decision-making, enabling the imagination of realistic plans of future actions as well as accurate environment simulation. However, one major issue in such models is generalization -- models are limited to synthesizing videos subject to language instructions similar to those seen at training time. This is heavily limiting in decision-making, where we seek a powerful world model to synthesize plans of unseen combinations of objects and actions in order to solve previously unseen tasks in new environments. To resolve this issue, we introduce RoboDreamer, an innovative approach for learning a compositional world model by factorizing the video generation. We leverage the natural compositionality of language to parse instructions into a set of lower-level primitives, which we condition a set of models on to generate videos. We illustrate how this factorization naturally enables compositional generalization, by allowing us to formulate a new natural language instruction as a combination of previously seen components. We further show how such a factorization enables us to add additional multimodal goals, allowing us to specify a video we wish to generate given both natural language instructions and a goal image. Our approach can successfully synthesize video plans on unseen goals in the RT-X, enables successful robot execution in simulation, and substantially outperforms monolithic baseline approaches to video generation.

한 줄 요약

RoboDreamer는 언어 조건을 구성 요소로 분해하여 비전-언어 세계 모델을 학습하여 로봇의 제어 계획을 개선하는 새로운 접근법이다.

핵심 기여도

RoboDreamer는 자연어 조건을 원시적 구성 요소(primitive components)로 분해하여 비전-언어 세계 모델을 학습함으로써 구성적 일반화(compositional generalization)를 달성.
다중 모달 조건(예: 목표 이미지, 목표 스케치)을 결합하여 공간 관계를 정확히 표현.
RT-X 데이터셋에서 기존 모델 대비 15% 높은 성공률을 기록하며, 단일 카메라 입력에서도 우수한 성능 보임.
T5-XXL, VQVAE, PerceiverSampler 등 사전 학습된 모듈을 활용하여 효율적인 조건 결합을 구현.

핵심 아이디어

기존 텍스트-비디오 생성 모델은 학습 시 보지 못한 언어 지시에 대한 일반화 능력이 제한적이었다. RoboDreamer는 자연어의 구성성(compositionality)을 활용해 텍스트를 원시적 행동 및 객체 간 관계로 분해하고, 이를 각각의 디퓨전 모델에 조건으로 전달함으로써 구성적 생성을 가능하게 한다. 예를 들어, "pepsi can을 plastic bottle 옆에 옮기라"는 지시는 "pepsi can", "plastic bottle", "near"라는 세 가지 원시 요소로 분해되어 각각의 디퓨전 모델에 조건으로 전달된다. 이는 새로운 언어 지시가 기존 구성 요소의 조합으로 표현될 수 있음을 의미하며, 학습되지 않은 조합에도 일반화할 수 있게 한다. 또한, RoboDreamer는 텍스트 외에도 목표 이미지와 스케치를 조건으로 결합하여 공간 정보를 보완적으로 제공함으로써, 더 정확한 비디오 생성과 로봇 실행을 가능하게 한다.

기술적 접근법

**모델 구조**: RoboDreamer는 AVDC와 Imagen 기반의 디퓨전 모델을 사용하며, ResNet 블록에 시간적 주의 계층(temporal attention layer)을 추가하여 시간 일관성을 강화.
**텍스트 조건**: T5-XXL을 사용하여 자연어 지시를 문맥 임베딩(contextual embedding)으로 변환.
**이미지/스케치 조건**: Stable Diffusion의 VQVAE 이미지 인코더를 활용하여 목표 이미지와 스케치를 인코딩.
**다중 모달 결합**: PerceiverSampler 아키텍처를 사용하여 텍스트, 이미지, 스케치 조건을 통합.
**하이퍼파라미터**: 3단계 캐스케이드 디퓨전 모델을 사용하여 초해상도(super-resolution) 비디오 생성.
**학습 데이터**: 약 70,000개의 로봇 조작 데모와 500개의 다른 작업을 사용하여 학습.

주요 결과

**RT-X 데이터셋**: RoboDreamer는 기존 모델 AVDC, HiP 대비 15% 높은 작업 성공률을 기록.
**Zero-shot 일반화**: 학습되지 않은 언어 지시에서도 70% 이상의 작업 완료율을 달성.
**다중 모달 조건**: RoboDreamer(t+i)와 RoboDreamer(t+s)는 각각 85%와 82%의 작업 완료율을 보이며, 단일 텍스트 조건 대비 10~15% 개선.
**RLBench 환경**: 단일 카메라 입력에서도 15%의 성공률을 달성하며, Image-BC, Hiveformer, UniPi 대비 우수한 성능.

의의 및 한계

RoboDreamer는 로봇이 학습되지 않은 작업을 해결할 수 있도록 구성적 추론을 가능하게 하며, 이는 데이터 부족한 환경에서 특히 유용하다. 또한, 텍스트, 이미지, 스케치를 결합하는 방식은 공간 관계를 명확히 표현할 수 있어, 로봇 실행의 정확도를 높인다. 그러나, RoboDreamer는 여전히 학습된 구성 요소의 조합에만 제한되며, 완전히 새로운 객체나 행동에 대한 일반화는 제한적이다. 또한, 실제 로봇 환경에서의 성능 검증은 아직 충분히 이루어지지 않았으며, 더 많은 실험과 데이터가 필요하다.

실용적 활용

RoboDreamer는 로봇이 복잡한 작업을 스스로 계획하고 실행하는 데 활용될 수 있으며, 특히 데이터가 제한된 환경에서 유용하다. 예를 들어, 물류 시스템에서 다양한 물체를 조합해 배치하거나, 제조 환경에서 정밀한 작업을 수행하는 데 적용 가능하다. 또한, 사용자가 스케치나 자연어로 작업을 지시할 수 있어, 인간-로봇 상호작용을 유연하게 지원한다.