기술 요약: WEAVER - 로봇 조작을 위한 효과적인 월드 모델

문제 정의

월드 모델(World Models, WMs) 또는 학습된 시뮬레이터는 제한된 실제 환경 상호작용을 통해 정책 평가, 개선 및 테스트 타임 플래닝(test-time planning)을 가능하게 함으로써 로봇 공학 분야에서 큰 가능성을 품고 있습니다. 그러나 이러한 다운스트림 역량을 실현하기 위해, 로봇 월드 모델은 종종 상충하는 세 가지 요구 사항을 동시에 충족해야 합니다:

1. 충실도(Fidelity): 실제 세계의 결과와 상관관계를 갖는 물리적으로 정확한 예측을 생성해야 합니다.
2. 일관성(Consistency): 특히 폐쇄(occlusion) 및 다단계 상호작용이 존재하는 상황에서 긴 예측 지평(prediction horizons) 동안 일관된 궤적을 생성해야 합니다.
3. 효율성(Efficiency): 실시간 플래닝 및 반복적인 정책 개선을 지원할 수 있을 만큼 빠르게 예측을 생성해야 합니다.

기존의 접근 방식들은 이 세 가지 기준을 동시에 만족시키지 못합니다. 비디오 생성 모델은 효율성이 부족한 경우가 많고, 잠재 모델(예: JEPA 스타일)은 임의의 시각-운동 정책 평가에 필요한 디코딩 능력이 부족할 수 있으며, 이전의 조작 월드 모델(예: Ctrl-World)은 테스트 타임 플래닝을 수행하기에는 너무 느리게 작동합니다. 더욱이, 로봇 조작은 다중 뷰 관찰, 가려진 객체 추론, 그리고 단순한 시각적 미학이 아닌 고충실도 상태 예측의 필요성과 같은 독특한 복잡성을 수반합니다.

방법론: WEAVER 아키텍처

저자들은 충실도, 일관성, 효율성이라는 세 가지 요소를 모두 달성하기 위해 설계된 다중 뷰 월드 모델인 WEAVER(World Estimation Across Views for Embodied Reasoning)를 제안합니다. 이 아키텍처는 비디오 생성, 잠재 월드 모델, 그리고 JEPA 스타일 아키텍처의 설계 결정들을 통합합니다.

핵심 아키텍처

• 다중 뷰 인코딩(Multi-View Encoding): WEAVER는 손목 카메라와 외부 카메라 등 여러 카메라 뷰를 프로프리오셉티브(proprioceptive, 고유 수용성) 상태와 함께 처리합니다. 사전 학습된 Stable Diffusion 3 VAE 인코더를 사용하여 이미지 패치와 프로프리오셉티브 토큰을 잠재 공간으로 매핑합니다.
• 잠재 역학 모델(Latent Dynamics Model): 핵심 예측기는 다음 조건들에 기반하여 미래의 잠재 상태($\hat{z}_t$)를 예측하는 2D 트랜스포머입니다:
  • 희소 메모리(Sparse Memory): 매 $k$번째 이전 잠재 상태로 구성된 장기 문맥($z_{mem}$).
  • 단기 이력(Short-Term History): 즉각적인 행동의 결과를 포착하는 최근 잠재 상태($z_{hist}$).
  • 행동 계획(Action Plan): 미래 행동의 덩어리(chunk)인 $a_t$.
• 예측 목적 함수(Prediction Objectives):
  • 잠재 예측(Latent Prediction): 미래 잠재 상태를 예측하기 위해 **플로우 매칭(Flow Matching)**을 통해 학습됩니다. 긴 지평의 일관성을 높이기 위해, 모델은 미래의 각 타임스텝에서 독립적으로 노이즈 레벨을 샘플링하는 **디퓨전 포싱(Diffusion Forcing)**을 채택합니다.
  • 프로프리오셉티브 예측(Proprioceptive Prediction): 일부 선행 연구들과 달리, WEAVER는 미래의 프로프리오셉티브 상태(관절 위치 및 그리퍼 너비)를 명시적으로 예측하며, 이는 변형 가능한 객체 조작에 매우 중요합니다.
  • 보상 및 크리틱 헤드(Reward & Critic Heads): 경량화된 보상 헤드($R$)와 크리틱 네트워크($V$)가 잠재 상태와 언어 지시문에 직접 작용합니다. 이를 통해 이미지를 디코딩하거나 외부 시각-언어 모델(VLM)을 호출하지 않고도 행동 덩어리를 효율적으로 점수화할 수 있습니다.

학습 및 추론 최적화

• 학습(Training): 모델은 DROID 데이터셋에서 사전 학습된 후 특정 작업 데이터에 대해 미세 조정됩니다. 효율성을 극대화하기 위해 패치 토큰을 공격적으로 제거하는 SPRINT 블록을 활용합니다.
• 추론 가속화(Inference Acceleration):
  • KV 캐싱(KV Caching): 디노이징 단계 전반에 걸쳐 메모리 및 이력 토큰에 적용되어 순전파 비용을 줄입니다.
  • 정류된 흐름(Rectified Flow): 적은 수의 함수 평가(NFE) 내에서 고품질 생성을 가능하게 하기 위해 정류된 흐름 목적 함수로 사후 학습됩니다.
  • 점진적 노이즈(Progressive Noise): 코사인 노이즈 스케줄을 사용하여 저노이즈 영역에 더 많은 예산을 할당함으로써 충실도를 높입니다.

다운스트림 응용 분야

WEAVER는 세 가지 특정 다운스트림 작업을 지원하도록 설계되었습니다:

1. 정책 평가(Policy Evaluation): 기록된 실제 세계의 궤적을 월드 모델 내에서 오픈 루프로 실행하여 성공률을 예측합니다.
2. 정책 개선(Policy Improvement): 행동 덩어리를 샘플링하고, 이를 월드 모델 내에서 시뮬레이션하며, 높은 이점을 가진 궤적을 필터링하여 기본 정책($\pi_{0.5}$)으로 증류(distill)합니다.
3. 테스트 타임 플래닝(Test-Time Planning): 여러 행동 덩어리를 샘플링하고 상상한 뒤, 잠재 보상/크리틱 헤드를 통해 평가하여 실행할 최적의 행동을 선택하는 "Best-of-N" 접근 방식을 사용합니다.

주요 결과

저자들은 Franka Emika Panda 로봇을 사용하여 다섯 가지 로봇 조작 작업(Stack Bowls, PnP Bag, PnP Marker, PnP Towel, Pour Beans)에 대해 WEAVER를 평가합니다.

• 충실도 대 효율성(Fidelity vs. Efficiency): WEAVER는 주요 베이스라인인 Ctrl-World를 **파레토 지배(Pareto-dominates)**합니다. WEAVER는 훨씬 적은 추론 시간(예: 10초 롤아웃에 대해 1030초 소요, Ctrl-World는 3050초 소요)을 요구하면서도 더 낮은 FID 및 FVD 점수(높은 시각적 충실도 의미)를 달-성합니다. WEAVER는 낮은 NFE(8–16)에서도 높은 품질을 유지하는 반면, Ctrl-World의 품질은 NFE가 감소함에 따라 크게 저하됩니다.
• 장기 지평 일관성(Long-Horizon Consistency): WEAVER는 긴 지평의 롤아웃(최대 10초)에서 우수한 성능을 보이며, 특히 폐쇄 및 변형 가능한 객체가 포함된 작업에서 베이스라인보다 낮은 에러율을 유지합니다.
• 정책 평가(Policy Evaluation): 정책을 평가할 때, WEAVER-FT(미세 조정된 모델)는 실제 성공률과 $\rho = 0.870$의 피어슨 상관계수를 달성하며, 사전 학습된 베이스라인을 크게 앞섭니다.
• 정책 개선(Policy Improvement): WEAVER가 생성한 합성 데이터(실제 상호작용 없이)를 사용하여 $\pi_{0.5}$ 파운데이션 모델을 미세 조정하면 실제 성공률이 38% 향상됩니다. 합성 데이터와 실제 데이터를 결 조합하면 실제 데이터만 사용했을 때보다 추가로 11% 더 향상됩니다.
• 테스트 타임 플래닝(Test-Time Planning): WEAVER는 Ctrl-World 대비 5~10배 빠른 속도로 테스트 타임 플래닝을 가능하게 합니다. 이 플래닝 방식은 높은 이점을 가진 행동 덩어리를 선택함으로써 기본 정책의 실제 성공률을 14%(특정 작업에서는 최대 20%) 향상시킵니다.
• 분포 외(OOD) 강건성: WEAVER는 사전 학습된 인코더를 사용하기 때문에 OOD 작업 데이터셋에서도 강력한 성능을 유지합니다.

의의 및 주장

본 논문은 WEAVER가 로봇 조작을 위한 단일 월드 모델로서 고충실도, 장기 일관성, 추론 효율성을 성공적으로 통합함으로써 중요한 진전을 이루었다고 주장합니다.

• 다운스트림 역량의 실현: 저자들은 기존의 월드 모델들이 (높은 충실도와 느린 추론 사이의) 트레이드오프 때문에 실시간 플래닝이나 효율적인 정책 개선을 가로막는 제약을 받았다고 주장합니다. 세 가지 요구 사항을 동시에 만족하는 WEAVER의 능력은 이러한 이전에 접근 불가능했던 역량들을 열어줍니다.
• 데이터 효율성: 결과는 고충실도 월드 모델이 정책 개선을 위해 값비싼 실제 데이터만큼 효과적인 합성 데이터를 생성할 수 있음을 시사하며, 이는 광범위한 실제 상호작용에 대한 의존도를 줄일 수 있음을 보여줍니다.
• 실용성: 정류된 흐름(rectified flow)과 KV 캐싱을 통해 실시간 제약 조건에 부합하는 추론 속도를 달성함으로써, WEAVER는 월드 모델 기반의 테스트 타임 플래닝을 이론적 개념이 아닌 물리적 로봇을 위한 실질적인 현실로 만듭니다.

저자들은 부분 관측성(촉각 센싱의 부재), 복잡한 변형 역학의 어려움, 현재의 짧은 지평 플래닝 제한과 같은 한계점을 인정합니다. 그러나 WEAVER가 로봇 공학을 위한 확장 가능하고 효율적이며 효과적인 월드 모델링의 새로운 기준을 세웠다고 봅니다.