DriftWorld: Fast World Modeling through Drifting

기술 요약: DriftWorld: Drifting을 통한 빠른 월드 모델링

문제 정의

예측 가능한 월드 모델은 로봇 학습에 필수적이며, 에이전트가 실제 환경에서 실행하지 않고도 행동의 결과를 상상하여 계획을 세울 수 있게 한다. 최근 행동 조건부 비디오 생성 기술의 발전으로 이러한 시뮬레이션의 충실도가 향상되었으나, 제어 성능 측면에서의 실용성은 추론 속도에 의해 심각하게 제한되고 있다. 최첨단 모델들은 확산(diffusion) 프로세스에 의존하며, 이는 미래 프레임을 생성하기 위해 반복적인 다단계 디노이징(denoising) 과정을 요구한다. 이러한 다단계 샘플링은 계산 비용이 매우 높아, 생성적 예측 제어 시스템 전체 실행 시간의 90~95%를 소비한다. 결과적으로, 실시간 계획 과정에서 최적의 행동 후보를 식별하기 위해 필요한 수백 번의 롤아웃(rollout)을 시뮬레이션하는 것은 불가능에 가까우며, 이는 온라인 정책 개선 및 오프라인 정책 평가를 제한한다.

방법론

저자들은 반복적인 확산 샘플링 대신 **드리프팅 생성 모델(drifting generative models)**에 기반한 단일 단계 생성(single-step generation) 프로세스를 도입한 행동 조건부 월드 모델인 DriftWorld를 소개한다. DriftWorld는 추론 시 반복적으로 디노이징하는 대신, 노이즈 사전 분포(noise prior)를 데이터 분포로 직접 매핑하는 행동 조건부 드리프트 필드(drift field)를 학습 과정에서 학습한다.

핵심 아키텍처 및 학습

• 단일 단계 생성: 모델 $f_\theta$는 노이즈 사전 분포 $\epsilon$과 컨디셔닝 $c$(과거 관측치 및 미래 행동)로부터 단 한 번의 순전파(forward pass)를 통해 미래 비디오 프레임 $x$를 생성한다.
• 드리프팅 손실(Drifting Loss): 학습 목적 함수는 생성된 분포 $q$를 실제 데이터 분포 $p$로 유도하는 드리프팅 필드 $V_{p,q}(x)$를 활용한다. 여기서 드리프트는 양의 샘플(실제 데이터)과 음의 샘플(생성된 샘플)의 평균 이동 벡터 간의 차이로 정의된다:
  $$V_{p,q}(x) = V^+_p(x) - V^-_q(x)$$
  모델은 이 드리프트 필드에 의해 오프셋된 고정된 타겟과 예측값 사이의 거리를 최소화하도록 고정점 반복(fixed-point iteration)을 통해 최적화된다.
• 행동 강조 드리프팅(Action-Accentuated Drifting): 모델이 행동 입력을 엄격히 따르도록 하기 위해, 음의 샘플 분포를 생성된 샘플과 "무작위 행동(no-action)" 상태를 보여주는 실제 샘플의 혼합으로 수정한다. 이는 모델이 행동에 무관한 항등 매핑(identity mapping)을 학습하는 것을 방지한다.
• 특징 공간 드리프팅(Feature Space Drifting): 복잡하고 고해상도인 실제 세계 데이터셋(예: Bridge-V2, RT-1)의 경우, 드리프팅 손실은 원시 픽셀 공간이 아닌 DINOv2/v3 인코더를 사용한 특징 공간에서 계산된다. 이는 시각적 선명도를 유지하고 의미 있는 거리 척도를 제공한다.
• 모션 가중치(Motion Weighting): 배경 재구성 오차를 최소화하느라 행동 컨디셔닝을 무시하는 것을 방지하기 위해, 손실은 모션 크기에 따라 가중치가 부여된다. 모션이 큰 영역(예: 로봇 그리퍼)에 더 높은 가중치를 부여하여 모델이 역동적인 변화를 정확하게 예측하도록 강제한다.
• 아키텍처: 생성기는 과거 프레임과 행동 시퀀스를 FiLM(Feature-wise Linear Modulation)을 통해 컨디셔닝하는 U-Net 구조를 가지며, 이를 통해 각 미래 프레임이 해당 행동에 정밀하게 종속되도록 보장한다.

주요 기여

1. DriftWorld 모델: 단일 순전파로 롤아웃을 생성하는 단일 단계 행동 조건부 월드 모델을 도입하여, 높은 시각적 품질을 유지하면서도 확산 기반 베이스라인 대비 평균 17배의 속도 향상을 달환했다.
2. 로보틱스를 위한 드리프팅 적응: 조건부 비디오 예측에 드리프팅 생성 모델을 처음으로 적용하였다. 여기에는 행동 강조 드리프팅 필드, 특징 공간 드리프팅(DINOv2/v3 활용), 그리고 U-Net 아키텍처 내의 프레임 단위 행동 컨디셔닝의 혁신적인 사용이 포함된다.
3. 실험적 검증: 다섯 가지 벤치마크(Bridge-V2, RT-1, Language Table, Push-T, Robomimic)에 대한 종합적인 평가를 통해, DriftWorld가 시각적 지표(SSIM, PSNR, LPIPS, FID, FVD)에서 확산 베이스라인과 대등하거나 우수한 롤아웃 품질을 보이면서도 추론 시간을 크게 단축함을 입증했다.

결과

• 속도: DriftWorld는 다양한 환경에서 30+ fps 이상의 생성 속도를 달성한다. 반면, 확산 기반 베이스라인(예: GPC, Ctrl-World)은 훨씬 낮은 프레임 레이트(예: 0.3–3.5 fps)로 동작하며, DriftWorld는 이보다 평균적으로 약 17배 더 빠르다.
• 시각적 품질: Push-T 벤치마크에서 DriftWorld는 확산 베이스라인 및 MSE 학습 기반 단일 단계 베이스라인보다 우수한 시각적 지표(예: GPC의 SSim 0.9717 대비 0.9925)를 기록했다. 또한, 다른 모델들이 긴 롤아웃에서 정확하게 렌더링하지 못하는 그리퍼-물체 접촉과 같은 미세한 상호작용을 성공적으로 시뮬레이션한다.
• 정책 개선 (온라인 제어): GPC-RANK 알고리즘의 행동 제안 순위를 매길 때, DriftWorld는 정책 성능을 유의미하게 향상시킨다. 예를 들어, Push-T에서 IoU 점수를 기본 정책의 0.635에서 0.781로 높였으며, 이는 GPC 확산 월드 모델(0.698)보다 높은 수치인 동시에 훨씬 짧은 시간 내에 탐색을 수행한 결과이다.
• 오프라인 정책 평가: DriftWorld는 매우 정확한 오프라인 시뮬레이터 역할을 한다. DriftWorld에서 롤아웃된 정책의 성공률 및 IoU 점수는 실제 환경의 성능과 강력한 상관관계를 보이며, 피어슨 상관 계수(Pearson correlation coefficient)에서 Robomimic Lift의 경우 0.9916, Robomimic Can의 경우 0.9250, Push-T의 경우 0.9515를 달성했다. 이는 Ctrl-World나 GPC와 같은 베이스라인 모델들의 상관계수를 상회하는 수치이다.

의의 및 주장

본 논문은 드리프팅 모델이 로봇 월드 모델링에 매우 적합하다고 주장하며, 생성 속도와 품질 사이의 결정적인 트레이드오프 문제를 해결한다. DriftWorld는 빠른 고품질의 "상상"을 가능하게 함으로써 다음을 직접적으로 지원한다:

1. 효율적인 계획: 확산 모델로는 불가능했던 수백 개의 행동 후보를 단일 제어 루프 내에서 평가할 수 있게 한다.
2. 신뢰할 수 있는 정책 평가: 롤아웃 기반 점수가 실제 성능과 높은 상관관계를 갖는 시뮬레이터를 제공함으로써, 정책 순위 선정 시 비용이 많이 드는 하드웨어 롤아웃의 필요성을 줄여준다.

저자들은 결론적으로 DriftWorld가 상상의 속도가 더 이상 계획 및 정책 평가의 병목 현상이 되지 않는 새로운 로봇 월드 모델링의 토대를 마련했다고 밝힌다.