기술 요약: WorldBagel

문제 정의

월드 모델(World models)은 인지, 추론 및 의사결정을 지원하기 위해 환경의 역학(dynamics)을 포착하는 것을 목표로 하며, 이는 체화된 지능(embodied intelligence)과 로보틱스의 핵심적인 목표 역할을 합니다. 최근의 시각-언어-행동(Vision-Language-Action, VLA) 모델들은 언어를 통합함으로써 멀티태스크 학습과 일반화에서 성공을 거두었으나, 이들은 주로 행동 예측과 태스크 수행에 중점을 두며 환경 역학의 명시적인 생성 모델링에는 집중하지 못했습니다. 반면, 전통적인 월드 모델링 연구는 언어나 대규모 멀티모달 지식을 깊이 있게 통합하지 않은 채 시각적 역학과 제어에 집중하는 경우가 많았습니다.

본 논문은 멀티모달 VLA 학습과 월드 모델링 사이에 자연스러운 수렴이 존재한다고 주장합니다. 통합된 모델은 시각적 관찰을 이해하고 언어 지시를 따를 수 있어야 할 뿐만 아니라, 공유된 표현 공간 내에서 미래의 시각적 상태를 생성하고 행동을 예측할 수 있어야 합니다. 그러나 통합된 멀티모달 모델을 어떻게 일반적인 시각-언어-행동-월드(Vision-Language-Action-World, VLAW) 모델로서 적응시키고 평가해야 하는지, 특히 행동 토큰화, 시퀀스 계획(sequence planning), 그리고 이질적인 데이터셋 간의 학습 안정성 측면에서 여전히 불분명한 점이 남아 있습니다.

방법론: WorldBagel

WorldBagel은 전용 이해(Understanding, UND) 및 생성(Generation, GEN) 전문가를 갖춘 투-타워(two-tower) 아키텍처를 특징으로 하는 현대적 멀티모달 통합 모델인 BAGEL을 기반으로 구축된 통합 VLAW 프레임워크입니다. WorldBagel은 이 아키텍처를 확장하여 멀티모달 이해, 구조화된 행동 모델링, 그리고 미래의 월드 예측을 공동으로 지원합니다.

1. 모델 아키텍처

이 프레임워크는 공유된 멀티모달 토큰 임베딩 공간 위에서 작동하며 두 개의 주요 브랜치로 구성됩니다:

• VLA 브랜치 (이해 및 행동 예측): 시각적 상태는 ViT 인코더를 통해 인코딩된 후, 교차 모달 접지(cross-modal grounding)를 전문으로 하는 UND 전문가를 통해 처리됩니다. 연속적인 제어 신호를 예측하기 위해 특화된 **푸리에 특징 행동 디코더(Fourier Feature Action Decoder, FFAD)**가 부착되어 있습니다.
• 월드 모델링 브랜치 (행동 조건부 생성): 이 브랜치는 과거의 시각적 토큰, 언어 지시, 그리고 구조화된 행동 토큰을 조건으로 합니다. 미래의 관찰을 예측하기 위해 플로우 매칭(flow matching)을 통한 GEN 전문가를 활용하며, 이후 VAE 디코더를 거칩니다.

2. 주요 기술 구성 요소

• 푸리에 특징 행동 토크나이저(FFAT) 및 디코더(FFAD): 표준 회귀(regression)나 이산적 빈닝(discrete binning, 불안정하거나 데이터 의존적일 수 있음) 대신, WorldBagel은 연속적인 행동 벡터를 구조화된 푸리에 특징 임베딩으로 매핑합니다.
  • 순방향 매핑(Forward Mapping): 행동은 조대(coarse) 및 미세(fine-grained) 변화를 모두 포착하기 위해 지수적으로 증가하는 스케일을 가진 주파수 대역으로 확장됩니다.
  • 예측: 모델은 가공되지 않은 행동 값 대신 잠재 공간(latent space)에서의 푸리에 특징을 예측하여 안정성을 향상시킵니다.
  • 역방향 매핑(Inverse Mapping): 행동은 주파수 전반에 걸친 위상 일관적 평균화(phase-consistent averaging)를 통해 재구성되며, 이를 통해 반복적인 최적화 없이도 연속성과 국소적 메트릭 구조를 보장합니다.
• 시퀀스 플랜을 통한 인터리브드(Interleaved) VLAW 모델링: 모델은 시각적 관찰, 언어 지시, 행동 토큰, 그리고 미래 상태 예측을 하나의 통합된 자기회귀(autoregressive) 시퀀스로 구성합니다. 유연한 "시퀀스 플랜"은 이러한 토큰들의 순서를 정의하며, 단일 순전파(forward pass) 내에서 다양한 목적(예: 미래 관찰 전에 행동을 예측하거나 그 반대의 경우)에 대해 모델을 학습할 수 있게 합니다.
• LLM 영감 데이터 샘플링: 다양한 규모와 감독 신호를 가진 이질적인 데이터셋을 처리하기 위해 저자들은 다음을 사용합니다:
  • 혼합 데이터셋 샘플링(Mixture Dataset Sampling): 대규모 데이터셋이 지배적이지 않도록 데이터셋 크기에 따른 평활화된 샘lu 확률을 적용합니다.
  • 우선순위 시퀀스-플랜 샘플링(Priority Sequence-Plan Sampling): "공동 VLAW 플랜"(행동과 월드 예측을 모두 감독함)에 정책 전용 플랜보다 높은 가중치를 부여하는 전략으로, 모델이 제어와 환경 역학 모두에 대한 감독을 빈번하게 관찰하도록 보장합니다.

3. 학습 목적 함수

모델은 사전 학습된 BAGEL 체크포인트에서 초기화되어 지도 미세 조정(Supervised Fine-Tuning, SFT)을 통해 미세 조정됩니다. 전체 손실 함수는 행동 예측 손실(푸리에 공간에서의 회귀)과 비전 생성 손실(미래 프레임을 위한 플로우 매칭)을 결합하며, 다음과 같이 가중치가 부여됩니다:
$$\mathcal{L} = \mathcal{L}_{action} + 0.1\mathcal{L}_{vision}$$

주요 기여

1. 통합 VLAW 프레임워크: 인지, 언어 접지, 행동 예측 및 미래 관찰 생성을 공동으로 모델링하기 위해 BAGEL의 투-타워 아키텍처를 확장한 WorldBagel을 도입했습니다.
2. 푸리에 특징 행동 모듈: 연속적인 제어 신호를 멀티모달 공간과 정렬된 구조화된 토큰으로 인코딩하는 FFAT 및 FFAD를 제안하여, 행동 표현의 품질과 안정성을 개선했습니다.
3. 통합성에 대한 체계적 연구: 멀티태스크 로봇 조작 및 교차 도메인 실험을 통해, 본 논문은 통합 자체가 효과적인 VLAW 모델을 학습하는 데 이득이 된다는 것을 입증하며, 일관된 경험적 이점과 멀티모달 월드 모델링에 대한 깊은 통찰을 제공합니다.

실험 결과

저자들은 세 가지 벤치마크인 LIBERO(시뮬레이션된 장기 호라이즌 조작), Language Table(실제 테이블탑 조작), Franka(물리 기반 시뮬레이션)에서 WorldBagel을 평가했습니다.

• 멀티태스크 성능: WorldBagel은 LIBERO 벤치마크에서 평균 성공률 **98.0%**를 달성하여, 강력한 베이스라인인 RynnVLA-002(+0.6%) 및 다양한 다른 VLA 베이스라인(예: OpenVLA-OFT, $\pi0.5$)을 능가했습니다.
• 월드 모델링 품질: 행동 조건부 월드 모델링 작업에서, WorldBagel은 모든 데이터셋에 대해 베이스라인 대비 일관되게 향상된 예측 충실도(낮은 FVD, 높은 PSNR/SSIM, 낮은 LPIPS)를 보여주었습니다. 예를 들어, LIBERO에서 FVD를 389.5에서 373.1로 감소시켰습니다.
• 행동 표현 품질: 푸리에 기반 방식(FFAD)은 Action MSE를 0.028로 줄였으며, 태스크 정체성(task identity)에 대한 선형 프로브 정확도 **68.7%**를 달성하여 회귀(49.8%) 및 기타 토큰화 방법보다 크게 앞섰습니다. 이는 학습된 임베딩이 더 구조적이고 의미론적으로 정렬되어 있음을 나타냅니다.
• 분포 변화 하의 안정성: 섭동(가우시안 노이즈, 행동 스케일링, 시간적 이동) 하에서도 WorldBagel은 RynnVLA-002에 비해 더 높은 예측 충실도(PSNR)와 낮은 LPIPS를 유지했습니다. 고유값 스펙트럼 분석 결과, WorldBagel은 더 높은 유효 랭크(33.7 대 16.3)와 더 낮은 지배적 고유값 비율을 보유하고 있어, 행동 역학에 대한 더 풍부하고 안정적인 표현을 가지고 있음을 시사했습니다.

의의 및 주장

본 논문은 통합 그 자체가 확장 가능한 월드 모델을 위한 핵심 요소라고 주장합니다. 단일 자기회귀 프레임워크 내에서 인지, 언어, 행동 및 미래 역학을 공동으로 모델링함으로써, WorldBagel은 통합된 아키텍처가 태스크 특화된 대안보다 뛰어날 수 있음을 입증합니다.

저자들은 다음과 같은 성과를 제시합니다:

1. 경험적 이점: 제어 성능(정책 성공률)과 예측적 환경 이해(월드 모델링 충실도) 모두에서 일관된 개선을 달성했습니다.
2. 방법론적 통찰: 푸리에 기반 행동 토큰화와 인터리브드 시퀀스 플래닝이 이질적인 멀티모달 학습을 안정화하고 고품질의 구조화된 행동 표현을 학습하는 데 결정적이라는 증거를 제공합니다.
3. 향후 연구를 위한 토대: 결과는 아키텍처의 통합이 정책 학습과 예측적 환경 모델링 모두를 개선함을 시사하며, 체화된 에이전트를 위한 확장 가능한 멀티모달 월드 모델의 유망한 기초를 제공합니다.

결론적으로, 최근의 통합 비전-언어 모델들이 전례 없는 생성 능력을 보여주고 있지만, 월드 모델로서의 역할은 아직 충분히 밝혀지지 않았습니다. WorldBagel은 이 간극을 체계적으로 해결하기 시작합니다.