잔차 잠재 액션을 통한 시각적 특징 기반 월드 모델 학습

기존 월드 모델은 고비용의 비디오 생성에 의존하여 연산 부담이 크고 환각 현상에 취약했다. 이 논문은 픽셀 대신 DINO 특징의 잔차를 활용한 RLA 기법을 통해 연산 효율을 극대화하면서도 복잡한 3D 환경에서 정확한 물리적 예측과 정책 학습이 가능함을 입증했다.

기존의 비디오 생성 기반 월드 모델은 고차원의 픽셀 공간에서 미래를 예측하려다 보니 연산량이 과도하고 물리적으로 불가능한 장면을 생성하는 환각 문제가 잦았다. Transformer 기반의 시각적 특징인 DINO 토큰을 활용하면 의미론적 정보는 유지하면서 데이터 차원을 줄일 수 있지만, 단순히 미래 토큰을 직접 회귀(Regression) 방식으로 예측하면 복잡한 3D 상호작용에서 예측값이 흐릿해지는 평균으로의 회귀 문제가 발생한다.

이 논문은 현재 상태와 미래 상태의 차이인 '잔차(Residual)'에 주목한다. 두 프레임 사이의 DINO 토큰 변화량을 RLA라는 압축된 잠재 공간으로 인코딩하면, 모델은 절대적인 상태값이 아닌 '변화의 원인'을 학습하게 된다. 이는 물리 시뮬레이션에서 변위(Displacement)를 계산하는 것과 유사한 원리로, 훨씬 단순한 구조로도 정교한 물리 법칙을 모사할 수 있게 한다.

결과적으로 RLA 공간에서 Flow Matching을 통해 궤적을 생성함으로써, 모델은 고해상도 비디오를 직접 렌더링하지 않고도 로봇이 수행할 액션의 물리적 결과를 정확히 예측한다. 이는 로봇이 상상 속(월드 모델 내부)에서 수만 번의 시행착오를 거치며 학습할 수 있는 효율적인 가상 환경을 제공한다.

RLA Autoencoder는 f_enc와 f_dec로 구성된다. [현재 프레임 s_t와 미래 프레임 s_t+h의 DINO 토큰 차이값을 입력으로] → [Self-Attention 레이어를 거쳐 64차원의 잠재 벡터 z로 압축하고] → [다시 s_t와 결합하여 s_t+h를 재구성하는 연산을 수행해] → [물리적 변화를 담은 압축된 액션 표현 z를 얻는다].

RLA-WM은 Flow Matching을 사용하여 미래의 z를 예측한다. [현재 상태 s_t와 로봇 액션 a_t:t+h를 조건부 입력으로] → [가우시안 노이즈 z_0에서 목표 z_1로 가는 속도 벡터 v를 ODE Solver로 적분하여] → [미래의 잠재 액션 z를 생성하고] → [이를 f_dec에 통과시켜 미래의 DINO 토큰 s_t+h를 도출한다].

강화학습 프레임워크인 WMRL은 학습된 RLA-WM 내부에서 정책을 최적화한다. [정책 모델이 출력한 액션을 월드 모델에 입력으로] → [월드 모델이 예측한 미래 DINO 토큰과 참조 영상의 토큰 간 L1 거리를 계산하여] → [음의 거리값을 보상(Video Aligned Reward)으로 산출하고] → [PPO 알고리즘을 통해 실제 환경과의 상호작용 없이 정책 가중치를 갱신한다].

ManiSkill 시뮬레이션 벤치마크에서 RLA-WM은 LPIPS 0.071, SSIM 0.931을 기록하며 Vid2World(비디오 확산 모델) 및 DINO-WM(직접 회귀 모델)을 압도했다. 특히 연산 효율성 측면에서 Vid2World가 1.1P FLOPs를 소모할 때 RLA-WM은 3.5T FLOPs만 사용하여 약 300배 이상의 효율성을 보였다.

액션 없는 비디오를 활용한 모방 학습 실험에서 RLA 기반 모델은 PushT 작업 성공률 15.2%를 달성하여 기존 최고 수치인 AdaWorld(9.2%) 대비 60% 이상의 성능 향상을 보였다. 또한 실제 로봇 데이터셋인 IWS에서도 DINO L1 오차 0.053을 기록하며 실세계 복잡한 조작 환경에서도 높은 예측 정확도를 유지했다.

RLA-WM은 1M 차원에 달하는 DINOv3-L 토큰 공간에서 직접 생성 모델을 돌리는 대신, 이를 2048(32x64) 차원의 RLA 공간으로 투영하여 '차원의 저주'를 해결했다. Flow Matching은 확률 경로를 직접 학습하므로 확산 모델보다 적은 단계(30 Euler steps)로도 고품질의 샘플링이 가능하다.

학습 시 MSE Loss는 실제 RLA 벡터 z와 노이즈 ϵ 사이의 속도 벡터 v* = z - ϵ를 타겟으로 설정한다. 이는 이미지 재구성 손실 없이도 물리적 동역학을 학습할 수 있게 하며, 추론 시에는 학습된 f_dec를 통해 DINO 토큰을 RGB 이미지로 복원하기 위해 별도의 UNet 디코더를 사용한다.

배경의 무관한 움직임이 많은 경우 RLA가 물리적 액션이 아닌 배경 변화를 인코딩하여 효율이 저하될 수 있다. 또한 현재 모델은 단일 프레임 쌍(Pair)에 의존하므로 물체가 가려졌다가 다시 나타나는 폐색(Occlusion) 상황에서의 장기적 메모리 유지에 한계가 있다.