월드 가이던스: 액션 생성을 위한 컨디션 공간에서의 월드 모델링

로봇이 정밀하게 움직이려면 미래의 변화를 예측해야 하지만, 고해상도 영상을 직접 생성하는 방식은 연산량이 너무 많고 비효율적이다. 이 논문은 미래 정보를 행동 제어에 꼭 필요한 핵심 데이터로만 압축해 예측함으로써, 복잡한 환경에서도 정밀하고 안정적인 로봇 조작을 가능하게 한다.

기존 VLA(Vision-Language-Action) 모델은 현재의 시각 정보와 명령어를 임베딩하여 행동을 결정한다. 하지만 정밀한 조작을 위해서는 미래에 일어날 변화를 미리 고려하는 월드 모델링이 필수적이다. 기존 방식은 미래의 고해상도 영상을 직접 생성하거나 너무 단순한 잠재 변수만을 사용해, 연산이 너무 무겁거나 정보가 부족해 정밀 제어가 어렵다는 한계가 있었다.

WoG는 미래의 시각 정보를 행동 생성 헤드에 직접 주입되는 압축된 컨디션으로 변환한다. 이는 마치 운전자가 앞 차의 움직임을 예상하며 핸들을 꺾는 것처럼, 모델이 미래의 핵심적인 물리적 변화(물체의 궤적, 충돌 가능성 등)를 저차원 벡터로 요약하여 행동 결정의 근거로 삼게 한다.

특히 2단계 학습을 통해 모델은 현재 상태만 보고도 미래에 필요한 컨디션을 내부적으로 미리 계산할 수 있게 된다. 결과적으로 추론 시에는 미래 영상을 실제로 보지 않고도, 마치 미래를 내다보는 듯한 정밀한 궤적 계획과 장애물 회피가 가능해진다.

전체 아키텍처는 Prismatic VLM 백본과 DiT(Diffusion Transformer) 기반의 액션 헤드로 구성된다. 미래 관측 정보를 처리하기 위해 DINOv2와 Wan VAE를 결합한 고정된 비전 인코더를 사용하며, 이를 Q-Former 기반의 인코더로 압축하여 저차원 컨디션($O^c$)을 생성한다.

1단계 학습에서는 현재 관측값의 임베딩($z$)과 미래 컨디션($O^c$)을 모두 입력받아 행동($A$)을 예측한다. Rectified Flow 손실 함수를 사용하여 예측 속도와 타겟 속도의 차이를 최소화한다. [행동 궤적 $A_\tau$, 타임스텝 $\tau$, 현재 임베딩 $z$, 미래 컨디션 $O^c$를 입력으로] → [신경망 $v_\theta$를 통해 속도 필드를 연산] → [예측 속도와 타겟 속도 $v^*$ 사이의 MSE 값을 얻고] → [이 값이 작을수록 미래 정보를 활용한 행동 생성이 정확해짐을 의미함]

2단계 학습에서는 미래 관측 정보에 대한 접근을 차단하고, VLM 백본이 현재 관측값($z$)만으로 미래 컨디션($O^c$)을 예측하도록 정렬한다. [현재 관측값 $O$와 명령어 $l$을 입력으로] → [쿼리 메커니즘 $f_q$를 통해 내부 표현을 연산] → [미래 컨디션 $O^c$와의 코사인 유사도 $S$를 얻고] → [이 유사도가 높을수록 모델이 미래를 정확히 내다보고 있음을 의미함]

SIMPLER 시뮬레이션의 Google Robot 작업에서 WoG는 평균 69.4%의 성공률을 기록하여, 기존 SOTA인 $\pi_0$-FAST(60.5%)와 OpenVLA(33.8%)를 크게 앞질렀다. 특히 장애물 회피가 중요한 'Move Near' 작업에서 82.5%의 높은 성능을 보였다.

WidowX Robot 작업에서도 평균 성공률 63.5%를 달성하며 ViPRA(62.5%)와 UniVLA(45.6%) 대비 우수한 성능이 확인됐다. 특히 'Eggplant in Basket' 작업에서는 91.7%의 성공률로 정밀한 조작 능력이 나타났다.

실제 로봇 실험(UR5)에서 Microwave 닫기(100%), Pick and Place(60%), 수건 접기(60%) 등 난이도 높은 작업에서 우수한 성능을 보였으며, 배경 변화나 조명 변화가 있는 OOD 환경에서도 성능 저하가 가장 적어 강력한 일반화 능력이 확인됐다.

WoG의 핵심은 미래 관측치를 직접 재구성하는 대신, 행동 생성에 최적화된 컨디션 공간으로 투영하는 것이다. 이는 비전 파운데이션 모델(DINOv2, SigLIP)의 고수준 시맨틱 정보와 VAE의 시공간적 특징을 결합하여 풍부한 표현력을 확보한다.

Q-Former 기반 인코더는 $N=16$개의 학습 가능한 쿼리 토큰을 사용하여 미래 프레임들로부터 행동 관련 특징을 집계하며, 최종적으로 $D=32$ 차원의 콤팩트한 벡터로 압축한다. 이 저차원 설계는 월드 모델링의 복잡도를 낮추면서도 제어에 필요한 핵심 정보를 보존한다.

학습 과정에서 인간 조작 영상을 활용하는 전략은 두 가지다. 첫째는 행동 라벨이 있는 소량의 영상을 1단계에 포함하는 것이고, 둘째는 행동 라벨이 없는 대량의 영상을 2단계의 미래 컨디션 예측 학습에만 사용하여 일반화 성능을 높이는 방식이다.

추론 시에는 VLM 백본의 마지막 4개 토큰을 쿼리 임베딩과 교차 어텐션시켜 미래 컨디션을 예측한다. 이는 VLM이 단순한 시각 인코더를 넘어, 환경의 동역학을 이해하는 내부 월드 모델로서 기능하게 함을 의미한다.

Stack Green on Yellow나 Drawer 작업과 같이 매우 미세한 기하학적 정밀도가 요구되는 작업에서는 성능 향상이 상대적으로 작다. 이는 현재 사용된 백본 모델의 공간 해상도 한계와 동적 예측만으로는 정밀한 기하학적 제약을 완벽히 모델링하기 어렵기 때문이다.