기술 요약: Cosmos Policy

문제 정의

최근 비디오 생성 모델의 발전은 복잡한 물리적 상호작작용, 시간적 인과관계, 그리고 장면의 진화를 포착하는 데 있어 인상적인 능력을 보여주었습니다. 그러나 이러한 모델을 로봇 제어에 적응시키는 과정은 전통적으로 복잡한 다단계 학습 파이프라인과 새로운 아키텍처 구성 요소(예: 별도의 액션 디퓨저 또는 역역학 모델)의 도입을 필요로 했습니다. 기존의 통합 비디오-액션 모델들은 처음부터 학습된 커스텀 설계를 사용하기 때문에, 대규모 사전 학습된 비디오 모델의 강력한 시공간적 사전 지식(spatiotemporal priors)을 효과적으로 활용하지 못하는 경우가 많습니다. 또한, 시각-언어-행동(VLA) 모델이 주요 백본으로 인기를 얻고 있지만, 이들은 주로 정적인 이미지-텍스트 쌍에 기반하여 학습되므로 비디오 모델에 비해 시간적 역동성이나 암시적 물리학을 파악하는 데 한계가 있을 수 있습니다.

본 연구가 다루는 핵심 과제는, 아키텍처의 수정 없이 단순한 단일 단계 프로세스를 통해 대규모 사전 학습된 비디오 파운데이션 모델을 로봇 정책으로 어떻게 효과적으로 적응시킬 것인가 하는 점이며, 동시에 학습된 월드 모델(world models)과 가치 함수(value functions)를 통해 모델 기반 계획(model-based planning)을 수행할 수 있도록 하는 것입니다.

방법론

Cosmos Policy 아키텍처

저자들은 NVIDIA Cosmos-Predict2-2B 비디오 파운데이션 모델로부터 미세 조정된 로봇 정책인 Cosmos Policy를 소개합니다. 이 접근 방식은 아키텍처의 변경이 없다는 점이 특징이며, 대신 모델 고유의 잠재 확산(latent diffusion) 프로세스를 활용하여 다중 모달 입출력을 처리합니다.

1. 잠재 프레임 주입 (Latent Frame Injection):
비디오 모델이 원래 지원하지 않는 모달리티(로봇 고유 수용 감각, 행동, 상태 가치, 다중 뷰 카메라)를 포함하기 위해, 저자들은 잠재 프레임 주입을 제안합니다.

• 메커니즘: 추가적인 모달리티를 새로운 잠재 프레임으로 인코딩하여 비디오 모델의 잠재 확산 시퀀스에 직접 인터리빙(interleave)합니다.
• 시퀀스 구조: 다중 카메라와 고유 수용 감각을 갖춘 로봇 설정의 경우, 잠재 시퀀스는 $(s, a, s', V(s'))$로 구성됩니다:
  • $s$: 현재 상태 (로봇 고유 수용 감각 + 현재 카메라 이미지).
  • $a$: 액션 청크 (잠재 프레임으로 인코딩됨).
  • $s'$: 미래 상태 (미래 고유 수용 감각 + 미래 카메라 이미지).
  • $V(s')$: 미래 상태 가치 (기대 누적 보상, 잠재 프레임으로 인코딩됨).
• 유연성: 이 방법은 모델의 가중치나 구조를 변경하지 않고도 시퀀스에 잠재 프레임을 추가하거나 제거함으로써 임의의 로봇 구성(예: 다양한 카메라 수)을 허용합니다.

2. 결합 학습 목표 (Joint Training Objectives):
모델은 세 가지 함수를 공동으로 학습하는 단일 단계 목표를 사용하여 로봇 시연 데이터에 대해 미세 조정됩니다:

• 정책 (Policy): 현재 상태에 조건화되어 행동, 미래 상태 및 가치를 예측합니다 ($p(a, s', V(s') | s)$).
• 월드 모델 (World Model): 상태와 행동에 조건화되어 미래 상태와 가치를 예측합니다 ($p(s', V(s') | s, a)$).
• 가치 함수 (Value Function): 상태, 행동 및 미래 상태에 조건화되어 미래 가치를 예측합니다 ($p(V(s') | s, a, s')$).

학습 데이터는 전문가 시연(배치의 50%)과 정책 롤아웃(배치의 50%)이 혼합되어 구성됩니다. 롤아웃 데이터는 월드 모델과 가치 함수를 학습시키는 데 사용되어, 시스템이 성공적인 궤적과 실패한 궤적 모두로부터 학습할 수 있게 합니다.

3. 모델 기반 탐색을 통한 계획 (Planning with Model-Based Search):
Cosmos Policy는 두 가지 배포 모드를 지원합니다:

• 직접 정책 (Direct Policy): 병렬 디코딩을 사용하여 현재 상태로부터 행동을 직접 생성합니다.
• 모델 기반 계획 (Model-Based Planning): 학습된 월드 모델과 가치 함수를 활용하여 최적의 행동을 탐색합니다.
  • 프로세스: 시스템은 여러 액션 제안을 샘플링하고, "플래닝 모델"(롤아웃 데이터로 집중 학습된 정교한 체크포인트)을 사용하여 결과로 나타날 미래 상태와 가치를 예측한 뒤, 가장 높은 예측 가치로 이어지는 행동을 선택합니다.
  • 앙상블 전략: 다중 모달 분포와 분산을 처리하기 위해, 시스템은 액션당 월드 모델과 가치 함수를 여러 번 쿼리하여 예측을 앙상블하고, "다수 평균(majority mean)"을 통해 결과값을 집계하여 성공 가능성을 결정합니다.

주요 기여

1. 비디오 모델의 단순한 적응: 본 논문은 잠재 프레임 주입을 사용하여 액션과 가치를 인코딩함으로써, 아키텍처 수정 없이 단일 단계의 사후 학습(post-training)만으로 대규모 사전 학습된 비디오 모델을 최첨단 로봇 정책으로 적응시킬 수 있음을 보여줍니다.
2. 통합된 정책, 월드 모델 및 가치 함수: 이 접근 방식은 정책 학습, 월드 모델링, 가치 추정을 단일 잠재 확산 프레임워크 내로 통합하여, 복잡한 액션 분포와 시간적 역동성을 포착하기 위해 사전 학습된 비디오 모델의 우선 지식을 활용합니다.
3. 롤아웃 정제를 통한 계획: 저자들은 롤아웃 데이터를 수집하고 가치 및 월드 모델 구성 요소를 미세 조정함으로써, 시스템이 효과적인 모델 기반 계획(best-of-N 샘플링)을 수행하여 성공률을 높일 수 있음을 보여줍니다.
4. 최첨단 성능: 이 방법은 시뮬레이션 및 실제 환경에서 새로운 벤치마크를 달성하며, 처음부터 학습된 정책, 비디오 기반 정책, 그리고 미세 조정된 VLA 모델들을 능가합니다.

실험 결과

저자들은 세 가지 벤치마크를 통해 Cosmos Policy를 평가했습니다:

• LIBERO 시뮬레이션: Diffusion Policy, Dita, $\pi_0$, OpenVLA-OFT, CogVLA와 같은 강력한 베이스라인을 능가하는 **98.5%**의 평균 성공률을 달 기록했습니다.
• RoboCasa 시뮬레이션: 단 50개의 태스크당 인간 시연만을 사용하여 **67.1%**의 평균 성공률을 달성했으며(많은 경쟁 모델이 사용하는 300개 이상의 시연보다 현저히 적음), 이는 GR00T-N1.5 및 Video Policy를 앞서는 수치입니다.
• 실제 ALOHA 작업: 도전적인 양손 조작 작업(예: 셔츠 접기, 지퍼백 다루기)에서 Cosmos Policy는 가장 높은 평균 점수(93.5%)를 기록했습니다.
  • 특히 고정밀 파지(grasp)와 높은 액션 다중 모달리티를 요구하는 작업에서 미세 조정된 VLA($\pi_0.5$, OpenVLA-OFT+)를 눈에 띄게 능가했습니다.
  • 계획의 영향: 롤아웃 데이터를 사용한 모델 기반 계획으로 강화했을 때, 가장 어려운 ALOHA 작업에서 태스크 완료 점수가 평균 12.5% 증가했습니다.

절제 연구 (Ablation Studies):

• 보조 학습 목표(함께 $s'$와 $V$를 예측하는 것)를 제거했을 때 성공률이 1.5% 하락했습니다.
• 사전 학습된 비디오 가중치 없이 처음부터 모델을 학습시켰을 때 3.9%의 하락이 발생하여, 비디오 모델의 사전 지식의 중요성을 강조했습니다.
• $V(s')$ 정식화(월드 모델을 필요로 함)를 사용하는 모델 기반 계획이 모델 프리(model-free) $Q(s, a)$ 변형보다 우수한 성능을 보였으며, 이는 학습된 환경 역학의 가치를 시사합니다.

의의 및 주장

본 논문은 Cosmos Policy가 정적인 이미지 기반 VLA 모델보다 대규모 비디오 생성 모델의 시공간적 사전 지식이 저수준 제어(low-level control)를 위한 더 우수한 기초가 된다는 것을 입증함으로써 로봇 학습의 중요한 진전을 이루었다고 주장합니다.

주요 주장은 다음과 같습니다:

• 단순성 및 효율성: 아키텍처 변경이나 다단계 학습 파이프라인 없이 강력한 로봇 정책을 생성할 수 있는 능력.
• 데이터 효율성: 비디오 파운데이션 모델으로부터의 강력한 초기화 덕분에 다른 최첨단 방법들(300개 이상 대비 50개)보다 훨씬 적은 시연으로도 높은 성능을 달성합니다.
• 계획 능력: 통합된 아키텍처는 자연스럽게 모델 기반 계획을 지원하여, 로봇이 미래 상태와 가치를 "상상"하여 더 나은 행동을 선택할 수 있게 하며, 이는 직접적인 정책이 어려움을 겪는 고정밀 및 다중 모달리티 작업에서 성능을 향상시킵니다.
• 일반화: 이 접근 방식은 시뮬레이션 및 실제 환경 전반에 걸쳐 인-디스트리뷰션(in-distribution) 및 아웃-오브-디스트리뷰션(out-of-distribution) 시나리오 모두에서 강력한 일반화 능력을 보여줍니다.

저자들은 모델 기반 계획의 추론 속도(액션 청크당 약 5초)와 월드 모델을 정제하기 위해 상당한 롤아웃 데이터가 현재 필요하다는 점을 한계점으로 인정하며, 이를 향후 개선해야 할 영역으로 식별했습니다.