Unified World Models: Coupling Video and Action Diffusion for Pretraining on Large Robotic Datasets

저자: Chuning Zhu, Raymond Yu, Siyuan Feng, Benjamin Burchfiel, Paarth Shah, Abhishek Gupta · RSS 2025

라이선스

오픈 액세스
논문: CC BY 4.0

아래 내용은 모두 논문의 라이선스(CC BY 4.0)에 따라 논문 전문을 바탕으로 하며, 그림에는 저자를 표기했습니다.

요약

행동과 미래 이미지에 대해 별도의 노이즈 타임스텝을 갖는 단일 확산 트랜스포머가 정책, 동역학 모델 또는 비디오 예측기로 작동하며, 행동 없는 비디오로 학습할 수 있다.

중요한 이유

하나의 모델이 행동 레이블이 있는 로봇 데이터와 훨씬 구하기 쉬운 행동 없는 비디오 모두로 학습할 수 있다. DROID 방식의 Franka 장치에서 다섯 가지 작업으로 수행한 실제 로봇 시험에서 사전 학습된 UWM은 모든 작업에서 가장 높은 성공률을 보였고, 추가 비디오로 공동 학습하자 시각적 방해 요소가 있는 경우를 포함해 성공률이 더 높아졌다.

문제

대규모 로봇 모델을 위한 모방 학습은 고품질 전문가 시연에 의존하며, 이는 원격 조작으로 수집하기에 비용이 많이 들고 느리다. 실제 세계의 동역학을 보여 주는 비디오는 풍부하지만, 행동 레이블이 없어 현재 대부분의 모방 학습 방법은 이를 활용할 수 없으며, 행동 복제는 궤적에 존재하는 시간적 동역학을 명시적으로 포착하지 않는다.

제안 방법

UWM은 현재 관측을 조건으로 하는 하나의 트랜스포머에서 행동 확산 과정과 미래 이미지 확산 과정을 결합하며, 각 모달리티에 독립적인 확산 타임스텝을 둔다. 어떤 모달리티의 타임스텝을 T로 설정하면 그 모달리티가 주변화되고 0으로 설정하면 그것을 조건으로 삼으므로, 같은 모델이 정책 p(a|o), 순방향 동역학 p(o′|o,a), 역방향 동역학 p(a|o,o′) 또는 비디오 예측 p(o′|o)를 샘플링한다. 행동 없는 비디오는 행동 타임스텝을 T로 고정하고 누락된 행동을 가우시안 노이즈로 채워 학습에 사용한다.

핵심 통찰

모달리티별 독립 확산 타임스텝은 부분 마스킹처럼 작동하므로, 단일 네트워크가 행동과 미래 프레임에 대한 주변 분포와 조건부 분포를 모두 다룬다.
성능 향상을 이끄는 것은 현재 프레임을 단순히 재구성하는 것이 아니라 미래 프레임을 예측하는 것이다. Stack-Bowls/Block-Cabinet에서 성공률 0.86/0.76으로, 현재 프레임 재구성의 0.70/0.66, Diffusion Policy의 0.48/0.60과 비교된다.
레지스터 토큰(기본 8개)은 행동 토큰과 이미지 토큰 사이의 정보 교환을 돕고, LIBERO 절제 실험에서 AdaLN 조건화가 크로스 어텐션보다 우수했다.
처음부터 학습할 때 UWM과 Diffusion Policy는 성능이 비슷하며, UWM은 사전 학습으로부터 더 많은 이득을 얻는다.
도메인 내 로봇 비디오로 공동 학습하는 것이 인터넷의 인간 활동 비디오(Kinetics-400 및 Something-Something v2)로 공동 학습하는 것보다 더 도움이 되었지만, 둘 다 로봇 데이터만 사용하는 것보다는 나았다.

주요 결과

2,000개의 DROID 궤적으로 사전 학습한 실제 로봇의 분포 내 성공률: Stack-Bowls, Block-Cabinet, Paper-Towel, Hang-Towel, Rice-Cooker에서 UWM 0.86, 0.76, 0.78, 0.82, 0.60으로, Diffusion Policy의 0.48, 0.60, 0.52, 0.64, 0.35 및 GR1의 0.66, 0.66, 0.60, 0.66, 0.40과 비교된다.
행동 없는 비디오로 사용한 추가 2,000개의 DROID 궤적으로 공동 학습하자 UWM의 분포 내 성공률은 0.92, 0.84, 0.86, 0.86, 0.65로, 분포 밖 성공률(처음 네 작업)은 0.76/0.60/0.78/0.64에서 0.84/0.72/0.84/0.76으로 높아졌다.
범주화된 OOD 장면(조명, 배경, 어수선함): 공동 학습한 UWM은 Stack-Bowls와 Block-Cabinet에서 21/30과 15/30으로, 사전 학습만 한 UWM의 15/30과 8/30, Diffusion Policy의 12/30과 6/30보다 높다.
LIBERO(LIBERO-90으로 사전 학습하고 분포 이동이 있는 LIBERO-10 작업 다섯 개로 미세 조정): 평균 성공률은 UWM 0.79 ± 0.11, DP 0.71 ± 0.12, PAD 0.57 ± 0.19, GR1 0.58 ± 0.14이다.
궤적 길이 내에서 50개의 LIBERO 궤적을 추적하는 역방향 동역학 모드: Book-Caddy와 Soup-Cheese에서 성공률 0.73과 0.45로, 정책 모드의 0.37과 0.33보다 높다.
DROID로 UWM을 100K 스텝 학습하는 데 NVIDIA A100 GPU 4개로 24시간이 걸린다.

한계

이 모델은 아직 대규모 인간 비디오로부터 학습하지 않으므로, 인간과 로봇 사이의 구현체 격차는 다루지 않는다.
순방향 동역학 예측에는 아티팩트가 포함될 수 있으며, 이는 계획에 대한 유용성을 떨어뜨릴 수 있다.
저자들은 이 버전에서 사용하지 않는 더 조밀한 비디오 예측으로부터 추가적인 향상을 기대한다.
LIBERO 분포 밖 설정에서 기준선 대비 향상 폭은 실제 로봇에서보다 훨씬 작다. 저자들은 시뮬레이션 동역학이 더 단순하기 때문이라고 본다.
Rice-Cooker는 분포에 가까운 초기화 20개에서만 평가되었고 OOD 설정에서는 평가되지 않았으며, 모든 실제 로봇 시험은 하나의 Franka 장치를 사용한다.

분야에 미친 영향

이 논문의 인용 횟수는 244회이다. 비디오와 코드를 프로젝트 페이지에서 볼 수 있다고 밝히며, 이 페이지는 GitHub의 코드 저장소 WEIRDLabUW/unified-world-model로 연결된다.

초록

모방 학습은 범용 로봇을 구축하기 위한 유망한 접근법으로 부상했다. 그러나 대규모 로봇 파운데이션 모델을 위해 모방 학습을 확장하는 것은 고품질 전문가 시연에 의존하기 때문에 여전히 어렵다. 한편, 다양한 환경과 다양한 행동을 담은 대량의 비디오 데이터는 쉽게 구할 수 있다. 이 데이터는 실제 세계의 동역학과 에이전트-환경 상호작용에 관한 풍부한 정보원이다. 그러나 대부분의 현대 방법에 필요한 행동 주석이 없기 때문에 이 데이터를 모방 학습에 직접 활용하는 것은 어려운 것으로 드러났다. 본 연구에서 우리는 정책 학습에 비디오 데이터와 행동 데이터를 모두 활용할 수 있게 하는 프레임워크인 Unified World Models(UWM)를 제시한다. 구체적으로 UWM은 행동 확산 과정과 비디오 확산 과정을 하나의 통합 트랜스포머 아키텍처 안에 통합하며, 각 모달리티는 독립적인 확산 타임스텝의 지배를 받는다. 우리는 각 확산 타임스텝을 제어하는 것만으로 UWM이 정책, 순방향 동역학, 역방향 동역학, 비디오 생성기를 유연하게 표현할 수 있음을 보인다. 시뮬레이션과 실제 실험을 통해 우리는 다음을 보인다. (1) UWM은 동역학 예측과 행동 예측을 모두 사용하여 대규모 다중 작업 로봇 데이터셋에서 효과적인 사전 학습을 가능하게 하며, 그 결과 모방 학습보다 일반화가 잘되고 강건한 정책을 얻는다. (2) UWM은 모달리티별 확산 타임스텝의 독립적인 제어를 통해 행동 없는 비디오 데이터로부터의 학습을 자연스럽게 지원하여, 미세 조정된 정책의 성능을 더욱 향상시킨다. 우리의 결과는 UWM이 확장 가능한 로봇 학습을 위해 대규모의 이질적인 데이터셋을 활용하는 방향으로 나아가는 유망한 한 걸음이며, 흔히 서로 동떨어진 모방 학습과 월드 모델링이라는 패러다임을 간단하게 통합한다는 것을 시사한다. 비디오와 코드는 https://weirdlabuw.github.io/uwm/ 에서 볼 수 있다

링크

arXiv 2504.02792: https://arxiv.org/abs/2504.02792
DOI https://doi.org/10.15607/rss.2025.xxi.015: https://doi.org/https://doi.org/10.15607/rss.2025.xxi.015
PDF: https://arxiv.org/pdf/2504.02792v1
코드: https://github.com/WEIRDLabUW/unified-world-model
프로젝트 페이지: https://weirdlabuw.github.io/uwm/