Cosmos Policy: 시각-운동 제어 및 계획을 위한 비디오 모델 미세 조정

Cosmos Policy: Fine-Tuning Video Models for Visuomotor Control and Planning

최근 비디오 생성 모델들은 복잡한 물리적 상호작용과 시간의 흐름에 따른 장면 변화를 포착하는 데 있어 놀라운 능력을 보여줍니다. 이러한 모델의 시공간적 사전 지식을 활용하기 위해, 로봇 공학 연구들은 비디오 모델을 정책 학습에 적용해 왔으나, 행동 생성을 위해 다단계의 사후 학습과 새로운 아키텍처 구성 요소를 필요로 하여 복잡성을 야기했습니다. 본 연구에서는 아키텍처 수정 없이 타겟 플랫폼에서 수집된 로봇 데모 데이터에 대한 단일 단계의 사후 학습만으로 대규모 사전 학습 비디오 모델(Cosmos-Predict2)을 효과적인 로봇 정책으로 조정하는 단순한 접근 방식인 Cosmos Policy를 소개합니다. Cosmos Policy는 로봇 행동을 비디오 모델의 잠재 확산(latent diffusion) 프로세스 내에서 잠재 프레임(latent frames)으로 인코딩하여 직접 생성하도록 학습하며, 모델의 사전 학습된 지식과 핵심 학습 알고리즘을 활용하여 복잡한 행동 분포를 포착합니다. 또한, Cosmos Policy는 미래 상태 이미지와 가치(기대 누적 보상)를 생성하는데, 이들 역시 잠재 프레임으로 인코딩되어 테스트 시(test-time)에 성공 확률이 더 높은 행동 궤적을 계획할 수 있게 해줍니다. 평가 결과, Cosmos Policy는 LIBERO 및 RoboCasa 시뮬레이션 벤치마크에서 각각 평균 성공률 98.5%와 67.1%로 최고 성능(SOTA)을 달성했으며, 까다로운 실제 양손 조작 작업에서도 가장 높은 평균 점수를 기록하여, 밑바닥부터 학습된 강력한 확산 정책, 비디오 모델 기반 정책, 그리고 동일한 로봇 데모로 미세 조정된 최신 시각-언어-행동(VLA) 모델들을 능가했습니다. 더 나아가, 정책 롤아웃(rollout) 데이터가 주어지면, Cosmos Policy는 경험을 통해 학습하여 세계 모델(world model)과 가치 함수를 정교화하고 모델 기반 계획을 활용함으로써 어려운 작업에서 더 높은 성공률을 달성할 수 있습니다. 우리는 코드, 모델 및 학습 데이터를 https://research.nvidia.com/labs/dir/cosmos-policy/ 에서 공개합니다.

Recent video generation models demonstrate remarkable ability to capture complex physical interactions and scene evolution over time. To leverage their spatiotemporal priors, robotics works have adapted video models for policy learning but introduce complexity by requiring multiple stages of post-training and new architectural components for action generation. In this work, we introduce Cosmos Policy, a simple approach for adapting a large pretrained video model (Cosmos-Predict2) into an effective robot policy through a single stage of post-training on the robot demonstration data collected on the target platform, with no architectural modifications. Cosmos Policy learns to directly generate robot actions encoded as latent frames within the video model's latent diffusion process, harnessing the model's pretrained priors and core learning algorithm to capture complex action distributions. Additionally, Cosmos Policy generates future state images and values (expected cumulative rewards), which are similarly encoded as latent frames, enabling test-time planning of action trajectories with higher likelihood of success. In our evaluations, Cosmos Policy achieves state-of-the-art performance on the LIBERO and RoboCasa simulation benchmarks (98.5% and 67.1% average success rates, respectively) and the highest average score in challenging real-world bimanual manipulation tasks, outperforming strong diffusion policies trained from scratch, video model-based policies, and state-of-the-art vision-language-action models fine-tuned on the same robot demonstrations. Furthermore, given policy rollout data, Cosmos Policy can learn from experience to refine its world model and value function and leverage model-based planning to achieve even higher success rates in challenging tasks. We release code, models, and training data at https://research.nvidia.com/labs/dir/cosmos-policy/