A Controllable Generative World Model for Robot Manipulation — Stanford · Tsinghua (ICLR 2026)

 

 

Ctrl-World

Overview --> Ctrl-World is designed for policy-in-the-loop rollouts with generalist robot policies. It generates joint multi-view predictions (including wrist views), enforces fine-grained action control via frame-level conditioning, and sustains coherent

ctrl-world.github.io

 

Authors: Yanjiang Guo¹²*, Lucy Xiaoyang Shi¹*, Jianyu Chen², Chelsea Finn¹
 
Affiliation: ¹Stanford University · ²Tsinghua University
 
Venue: ICLR 2026 (International Conference on Learning Representations)
 
Preprint: arXiv:2510.10125v3 [cs.RO] (초판 Oct 2025, v3 개정 Mar 2026)

 

0. Abstract

 

Generalist 로봇 정책(Generalist robot policies)은 이제 광범위한 조작 기술을 수행할 수 있지만, 낯선 객체와 지시문에 대한 그것들의 능력을 평가하고 개선하는 것은 여전히 중대한 과제로 남아 있다.

엄밀한 평가(Rigorous evaluation)는 대량의 실세계 rollout을 요구하는 반면, 체계적 개선(systematic improvement)은 전문가 라벨이 붙은 추가적인 교정 데이터(corrective data)를 요구한다.

이 두 과정 모두 느리고, 비싸며, 확장하기 어렵다.

월드 모델은 정책이 상상 공간(imagination space) 안에서 rollout할 수 있게 함으로써 유망하고 확장 가능한 대안을 제공한다.

 

 

그러나 핵심 과제는 generalist 로봇 정책과의 다단계 상호작용(multi-step interactions)을 다룰 수 있는 제어 가능한(controllable) 월드 모델을 구축하는 것이다.

이는 멀티뷰 예측(multi-view prediction), 세밀한 액션 제어(fine-grained action control), 그리고 일관된 장기 지평 상호작용(consistent long-horizon interactions)을 지원함으로써 현대의 generalist 정책과 호환되는 월드 모델을 요구하는데, 이는 이전 연구들에서는 달성되지 않았다.

 

 

이 논문에서 generalist 로봇 정책의 지시 이행 능력(instruction-following ability)을 평가하고 개선하는 데 사용될 수 있는

제어 가능한 멀티뷰 월드 모델을 도입함으로써 한 걸음 나아간다.

본 논문의 모델은 pose 조건 메모리 검색 메커니즘(pose-conditioned memory retrieval mechanism)으로 장기 지평 일관성을

유지하고, 프레임 수준 액션 조건화(frame-level action conditioning)를 통해 정밀한 액션 제어를 달성한다.

 

DROID 데이터셋(95k 궤적, 564 장면)으로 학습된 본 연구의 모델은 새로운 시나리오와 새로운 카메라 배치 하에서 20초 이상 동안 공간적·시간적으로 일관된 궤적을 생성한다.

 

이 방법이 실세계 로봇 rollout 없이 정책 성능을 정확하게 순위 매길 수 있음을 보인다.

나아가, 상상 속에서 성공 궤적을 합성하고 이를 supervised fine-tuning에 사용함으로써, 본 접근법은 정책 성공률을 44.7% 개선할 수 있다.


 

1. Introduction 

 

비전-언어-액션(VLA) 모델의 최근 발전은 광범위한 조작 태스크와 시나리오에 걸쳐 역량을 입증했다.

그 유망함에도 불구하고, 현재의 정책들은 open-world 환경에서 시험될 때 여전히 취약(brittle)하다.

 

하나의 중심적 과제는 정책 평가(policy evaluation)이다.

 

Generalist 정책의 성능을 평가하는 것은 일반적으로 통계적 유의성을 달성하기 위해 태스크와 환경에 걸쳐 신중하게 반복되는 대량의

실세계 rollout을 요구한다.

이러한 프로토콜은 물류적으로 부담스럽고, 반복(iteration)을 늦추며, 현재 정책 능력에 대한 섬세한 이해를 저해한다.

 

 

마찬가지로 중요한 것이 정책 개선(policy improvement)이다:

 

약점이 드러나더라도, 기존 방법들은 더 많은 전문가 데이터를 수집하는 것 외에는 실패 사례에 대해 정책을 강화할 방법을 거의

제공하지 못한다.

대규모 사전학습이 어느 정도 강건성을 제공하지만, 정책들은 낯선 객체나 지시문을 마주칠 때 흔히 여전히 취약하다.

 

부족한 것은 generalist 모델을 정제하기 위한 빠르고 저렴한 피드백 주도(feedback-driven) 메커니즘이다:

 

실패 사례를 표면화하고, 교정 경험(corrective experiences)을 수집하며, 정책을 반복적으로 개선하는 방법이다.

 

예측 모델을 학습하고 상상 속에서 반복하는 것은 확장 가능하고 유망한 대안이다.

선행 연구가 액션 조건 월드 모델을 탐구했지만, 대부분의 접근은 수동적(passive) 비디오 예측 설정에 집중하며 진보된 generalist 정책과 능동적으로 상호작용하기에는 충분하지 않다.

 

policy-in-the-loop rollout을 지원하는 능력을 저해하는 몇 가지 중요한 한계를 관찰한다.

 

이러한 모델들은 일반적으로 단일 3인칭 카메라 시점만 시뮬레이션하는데, 이는 심각한 부분 관측성(partial observability)으로 이어지고, 그 결과 환각(hallucination)을 유발할 수 있다(예: 사전 물리적 접촉 없이 객체가 그리퍼로 순간이동하는 것).

 

이 단일 시점 입력은 또한 3인칭과 손목 시점(wrist-view) 카메라를 모두 입력으로 요구하는 많은 현대 VLA 정책과 비호환이다.

 

게다가, 기존 모델들은 일반적으로 고주파 액션의 인과적 효과(causal effects of high-frequency actions)를 포착하는 데 필요한 세밀한 제어가 부족하다.

 

마지막으로, 그것들은 장기 지평 비디오 생성에 걸쳐 시간적 일관성을 유지하는 데 어려움을 겪는다.

 

Ctrl-World는 generalist 로봇 정책과의 policy-in-the-loop rollout을 위해 설계되었다. 그것은 (손목 시점을 포함한) 공동 멀티뷰 예측을 생성하고, 프레임 수준 조건화를 통해 세밀한 액션 제어를 강제하며, pose 조건 메모리 검색을 통해 일관된 장기 지평 동역학을 지속한다. 이러한 구성 요소들은 (1) 실세계 rollout과의 정렬을 갖춘 상상 속에서의 정확한 정책 평가와, (2) 합성 궤적을 통한 표적 정책 개선을 가능하게 한다.

 

본 논문에서는 Figure 1에 예시된 것처럼, policy-in-the-loop 상호작용을 위해 설계되어 다단계 rollout을 전적으로 상상 공간 안에서 가능하게 하는, 제어 가능한(Controllable) 멀티뷰 생성형 월드 모델인 Ctrl-World를 도입한다.

 

설계는 세 가지 핵심 구성 요소에 의존한다:

 

(1) 공동 멀티뷰 예측(Joint multi-view prediction) : 장면의 더 포괄적인 시각 표현을 포착하고 현대 VLA 정책의 입력 형식을 충족한다.

특히, 손목 카메라 예측의 포함은 접촉이 많은(contact-rich) 객체 상호작용 동안의 환각을 유의미하게 감소시킨다.

 

(2) 프레임 수준 액션 조건화(Frame-level action conditioning) : 시각 동역학을 제어 신호와 긴밀히 정렬하여, 생성된 rollout이 각 액션의 인과적 효과를 반영하도록 보장한다.

 

(3) 메모리 검색(Memory retrieval) : 희소한 히스토리 프레임을 문맥에 추가하고 대응하는 pose 정보를 각 프레임에 투영하여,

모델이 유사한 과거 상태에 주의를 기울이고 관련 정보를 검색할 수 있게 한다.

이 메커니즘은 장기 지평 rollout을 안정화하고 시간적 일관성을 보존한다.

 

 

이러한 메커니즘들을 결합하면, 사전학습된 수동적 비디오 생성기(passive video generator)정책이 직접 제어·활용할 수 있는 상호작용형 시뮬레이터(policy-compatible interactive simulator)로 변환할 수 있게 된다.

 

이 연구의 핵심 기여는 로봇 조작을 위한 제어 가능한 월드 모델이다.

 

실험을 통해 이 모델이 정책이 평가될 수 있고 — 실세계 rollout과의 순위 정렬(ranking alignment)로 — 개선될 수 있는 — 성공률을 높이는 표적 합성 데이터(targeted synthetic data)를 통해 — 새로운 상상 기반 워크플로를 가능하게 함을 발견한다.

 

구체적으로, 본 연구에서는 Ctrl-World를 DROID 데이터셋(Khazatsky et al., 2024)으로 학습하고 그것이 새로운 장면과 카메라 배치로 일반화하여 20초 이상 동안 일관된 rollout을 지속함을 보인다.

 

나아가 Ctrl-World를 이용한 상상 기반 평가가 정책의 실세계 지시 이행 능력을 충실히 반영함을 보인다.

 

마지막으로, 월드 모델 안에서 성공 궤적을 합성하고 이러한 합성 rollout으로 지도 fine-tuning을 수행함으로써 미지의 객체와

새로운 지시문이 있는 downstream 태스크에서 π₀.₅-DROID(Intelligence et al., 2025)의 성능을 개선할 수 있음을 시연한다.


 

 

2. Related Works

 

 

Video Generation Models for Robotics.

 

비디오 생성 모델의 최근 발전(Agarwal et al., 2025; Wan et al., 2025; Blattmann et al., 2023a; Chi et al., 2025)은 사실적이고 시간적으로 일관된 콘텐츠의 생성을 가능하게 했으며, 이는 물리 세계에 대한 강한 이해를 반영한다.

 

일부 연구는 비디오 예측 모델을 활용해 가짜 액션 라벨(fake action labels)이 붙은 로봇 궤적을 합성하며, 이러한 합성 궤적은 그런 다음 정책 학습에 사용될 수 있다(Jang et al., 2025; Bharadhwaj et al., 2024).

 

다른 연구들은 비디오 모델을 정책 backbone으로 직접 사용하여, 트래킹(tracking)이나 역동역학(inverse dynamics)을 통해 액션을 디코딩한다(Black et al., 2023; Du et al., 2024; Yang et al., 2023; Hu et al., 2024; Liang et al., 2024; Liao et al., 2025; Tan et al., 2025; Feng et al., 2025).

 

이를 보완하는 연구 흐름은 co-training을 통해 미래 예측 목적함수(future-prediction objectives)를 generalist 정책에 통합하여, 물리 지식을 정책에 포함시킨다(Zhao et al., 2025; Li et al., 2025a; Zhu et al., 2025; Guo et al., 2024; Gao et al., 2024; Zhang et al., 2025; Zheng et al., 2025; Zhong et al., 2025).

 

이러한 연구들과 달리, 본 연구는 비디오 생성을 활용해 액션 조건 예측(action-conditioned prediction)을 수행하며, 이는 모델을 정책 평가 정책 개선 양쪽 모두에 사용할 수 있게 한다.

 

 

Action-Conditioned World Models.

 

사전학습된 비디오 모델은 강력하지만, 흔히 고수준 언어 지시(high-level language instructions)에만 조건화된다.

 

그럼에도 불구하고, 일부 선행 연구들은 저차원 상태 공간(low-dimensional state spaces)(Nagabandi et al., 2020)에서와 이미지 관측(image observations)으로(Hafner et al., 2019, 2020; Hansen et al., 2022; Wu et al., 2023; Oh et al., 2015; Huang et al., 2025) 액션 조건 예측 모델(action-conditioned predictive models)의 사용을 탐구했다.

이러한 접근들 중 다수는 태스크 특화(task-specific) 모델을 학습하는(Hafner et al., 2019) 반면, 본 연구는 generalist·다중 태스크(multi-task) 월드 모델을 학습하는 데 집중한다.

 

초기 연구들(Finn & Levine, 2017; Ebert et al., 2018; Xie et al., 2019; Dasari et al., 2019; Yang et al., 2023; Wu et al., 2024)뿐 아니라 diffusion을 활용하는 더 최근의 접근들(Quevedo et al., 2025; Chen et al., 2024; Ball et al., 2025; Gao et al., 2025; Ren et al., 2025; Hafner et al., 2025)과 프레임 수준 액션 조건화(frame-level action conditioning)(Zhu et al., 2024) 위에 구축하여, 본 연구는 멀티뷰 예측(multi-view prediction), 장기 지평 시간 일관성(long-horizon temporal coherence), 그리고 세밀한 제어 가능성(fine-grained controllability)을 통합하는 모델을 제안한다.

 

본 연구의 실험은 이러한 능력들이 최신 generalist VLA 정책의 효과적인 평가와 개선을 가능하게 함을 보인다.


 

 

3. Problem Formulation 

목표 — generalist 정책이 제안한 액션의 미래 결과를 예측하는 월드 모델 W 개발

 

본 연구는 generalist 로봇 정책이 제안한 액션의 미래 결과(future outcomes)를 예측할 수 있는 월드 모델을 개발하는 것을 목표로 한다.

 

현대의 generalist 정책 π는 일반적으로 멀티뷰 관측(multi-view observations)언어 지시(language instructions)

액션 시퀀스로 매핑한다(Zhao et al., 2023; Black et al., 2025).

 

구체적으로, 로봇 관측 o_t = [I¹_t, …, Iⁿ_t, q_t]는 n개의 카메라 시점 [I¹_t, …, Iⁿ_t]과 로봇 pose q_t를 포함하며, 정책은 지시 l이 주어졌을 때 H-스텝 액션 청크(action chunk)를 출력한다:

 

목표는 월드 모델 W를 사용해 A_t = [a_{t+1}, …, a_{t+H}]의 각 스텝을 실행한 결과를 예측하는 것이다.

상상 공간에서 정책과의 다단계 상호작용(multi-step interaction)을 가능하게 하려면, W는 미래의 멀티뷰 관측을 생성해야 한다:

 

그런 다음 최종 예측 o_{t+H}는 정책 π로 다시 보내져 다음 액션 청크 A_{t+H} ∼ π(· | o_{t+H}, l)를 생성할 수 있다. 이런 방식으로, 정책과 월드 모델은 자기회귀적으로(auto-regressively) 상호작용하며, 장기 지평 rollout을 전적으로 상상 공간 안에서 가능하게 한다.


 

 

 

4. Controllable World Model for Robot Manipulation

4.1 Learning World Model Ctrl-World

본 연구의 목표는 현대 VLA 정책을 평가하고 개선하는 데 사용될 수 있는 월드 모델을 학습하는 것이다.

이를 달성하려면, 모델은 먼저 그러한 정책들이 흔히 사용하는 멀티뷰 관측을 지원해야 한다.

또한 모델이 제어 가능(controllable)한 것 — 그러한 제어 능력이 없는 사전학습 backbone에서 초기화되었더라도 액션 입력을 신뢰성 있게 그리고 긴밀하게 따르는 것 — 이 중요하다.

 

마지막으로, 모델은 일관된 rollout을 생성하기 위해, 가림(occlusion)이 존재하는 상황에서도 장기 지평에 걸쳐

시간적 일관성을 유지해야 한다.

본 연구에서는 월드 모델을 spatial-temporal transformer를 가진 사전학습 비디오 확산 backbone(Blattmann et al., 2023b)에서 초기화하고, Figure 2에 예시된 세 가지 핵심 적응(key adaptations)을 도입한다.

Ctrl-World는 사전학습된 비디오 확산 모델에서 초기화되어, 다음으로 제어 가능하고 시간적으로 일관된 월드 모델로 적응된다: (1) 통합된 정보 이해를 위한 멀티뷰 입력과 공동 예측. (2) 희소한 히스토리 프레임을 문맥에 추가하고 프레임 단위 cross-attention을 통해 pose 정보를 각 프레임에 투영하여, 예측을 유사한 과거 상태에 재정박(re-anchoring)하는 메모리 검색 메커니즘. (3) 고주파 액션을 시각 동역학과 더 잘 정렬하기 위한 프레임 수준 액션 조건화.

 

Multi-View Joint Predictions.

 

최신 VLA 모델은 흔히 전역 문맥(global context)을 위한 다수의 3인칭 카메라와 정밀한 상호작용을 위한 손목 장착 카메라에 의존한다(Intelligence et al., 2025; Liu et al., 2024, 2025).

이에 맞추기 위해, 월드 모델은 각 스텝에서 모든 시점에 걸쳐 공간적으로 일관된 예측을 생성해야 한다(Jiang et al., 2025).

 

선행 연구는 feed-forward transformer멀티뷰 카메라 간의 공간 관계를 확장 가능한 방식으로 효과적으로 포착할 수 있음을 보였다(Wang et al., 2025).

선행 연구를 따라, 본 연구는 각각 H×W개의 토큰을 담은 N개의 입력 이미지를 토큰 차원을 따라 연결(concatenate)하고

모든 시점 o_{t:t+H}를 공동으로 예측한다.

 

실험을 통해 멀티뷰 공동 예측이 일관성도 개선하고 환각을 상당히 감소시킴을 발견한다.

 

 

Pose-conditioned Memory Retrieval Mechanism.

 

월드 모델의 예측 오차는 긴 rollout에 걸쳐 누적되는 경향이 있어, 드리프트(drift)와 비일관성으로 이어진다.

이를 완화하기 위해, 모델 입력을 과거 프레임(past frames)으로 증강한다.

문맥이 너무 길어지는 것을 방지하기 위해, stride m으로 k개의 히스토리 프레임을 샘플링하여, 모델이 o_{t+1:t+H} ∼ W(· | o_{t−km}, …, o_t, l)를 예측할 수 있게 한다.

 

추가로, 대응하는 로봇 팔 pose [q_{t−km}, …, q_t]를 spatial transformer 내부의 프레임 단위 cross-attention을 통해 프레임 [o_{t−km}, …, o_t]에 임베딩한다.

이는 모델이 팔 pose를 사용해 과거로부터 관련 프레임을 식별하게 하여, 미래 예측을 관련 히스토리에 효과적으로 정박(anchoring)시킨다.

 

 

Frame-level Action Conditioning.

 

사전학습 비디오 모델은 텍스트와 이미지에만 조건화되어, 그 제어 정밀도가 제한된다.

완전한 제어 가능성을 가능하게 하기 위해, 추가로 모델을 정책이 출력한 액션 시퀀스 [a_{t+1:t+H}]에 조건화한다.

 

또한 각 액션 시퀀스를 Cartesian 공간 로봇 팔 pose [a′{t+1:t+H}]로 변환하고 과거 pose [q{t−km}, …, q_{t−m}, q_t]와 연결한다.

그런 다음 프레임 단위 cross-attention(Zhu et al., 2024; He et al., 2025)이 spatial transformer 내부에 적용되어, 각 프레임의 시각 토큰이 그것에 연관된 pose 임베딩에 주의를 기울이게 한다.

히스토리 프레임의 경우, 이 pose는 [q_{t−km}, …, q_{t−m}, q_t]에 대응하고, 미래 프레임의 경우 [a′_{t+1:t+H}]에 대응한다.

 

 

Training Objective.

 

본 연구는 연구 모델을 사전학습된 1.5B Stable-Video-Diffusion(SVD) 모델(Blattmann et al., 2023a)로 초기화한다.

 

사전학습 비디오 모델의 지식과 구조를 계승하기 위해, 입력 액션을 위한 action-projection MLP만 새로 초기화하고 초기화 시점에 다른 파라미터는 변경하지 않고 유지한다.

그런 다음 이 액션 조건 월드 모델은 diffusion loss(Ho et al., 2020; Karras et al., 2022)로 fine-tuning된다.

 

학습 동안, 예측 목표 x₀ = o_{t+1:t+H}는 스케줄러 ᾱ_{t′}를 사용해 diffusion 스텝 t′ ∈ [0, T′]에서 가우시안 노이즈 ε ∼ 𝒩(0, I)로

교란되어, x_{t′} = √ᾱ_{t′} x₀ + √(1−ᾱ_{t′}) ε_{t′}가 된다.

 

모델 입력은 히스토리 토큰과 noised future의 연결이다: [o_{t−km}, …, o_{t−m}, o_t, x_{t′}].

 

전체 학습 목적함수는 다음과 같다:

 

여기서 x̂₀는 모델의 예측을 나타내고, c = [q_{t−km}, …, q_t, a′{t+1:t+H}, o{t−km}, …, o_t]는 모든 모델 입력에 대응한다.

모델의 더 자세한 내용은 Appendix A에서 찾을 수 있다.


4.2 Using Ctrl-World for Policy Evaluation and Improvement

Policy Evaluation within World Model.

 

제어 가능하고 일관된 월드 모델이 학습되면, 상상 공간에서 policy-in-the-loop rollout을 수행할 수 있다.

초기 관측 o₀와 지시 l이 주어지면, 정책 π는 월드 모델 W와 함께 합성 궤적 τ를 생성할 수 있다.

초기 관측은 검증 데이터셋에서 샘플링되거나 실세계 셋업의 스냅샷으로 기록될 수 있다.

 

본 연구의 실험에서, 각 궤적을 인간 선호 판단(human preference judgments)에 기반해 성공 또는 실패로 라벨링한다.

최근 연구들(Du et al., 2023)이 Vision-Language 모델을 범용 보상 모델로 사용하는 것을 탐구하지만, 그러한 확장을 향후 연구로 남겨둔다.

 

Policy Improvement with Synthetic Data.

 

평가를 넘어, 월드 모델은 정책 성능을 개선하기 위해 성공적인 합성 궤적을 탐색하는 것을 가능하게 한다.

 

본 연구는 고정된 초기 관측과 지시 하에서 정책 행동이 매우 결정적(deterministic)인 경향이 있음을 관찰한다.

 

예를 들어, 정책은 여러 시도에 걸쳐 다양한 객체로 확률적으로 손을 뻗기보다 같은 객체를 잡으려는 경향이 있다.

더 큰 탐색 공간을 탐색하기 위해, 본 연구는 rollout의 다양성을 장려하는 구조적 섭동(structured perturbations)을 도입한다.

 

구체적으로, (i) 지시문을 재표현(rephrase)할 수 있는데, VLA 정책이 조종 가능(steerable)한 경향이 있어 서로 다른 지시에 대해 서로 다른 행동을 보이기 때문이다;

또는 (ii) 월드 모델 안에서 정책을 무작위 초기 상태로 리셋할 수 있는데, 이는 다양한 초기 관측으로 이어진다.

 

언어 지시 [l⁰, …, l^M]가 있는 downstream 태스크 집합에서 시작하여, 합성 rollout을 수집하고 인간 선호에 기반해 그것들을

점수화한다.

정책 성능을 개선하기 위해, 성공 궤적에 대해 정책을 fine-tuning한다. 전체 절차는 Algorithm 1에 요약되어 있다.


 

 

5. Experiments

Ctrl-World를 평가하기 위한 실험을 수행한다.

 

다음 질문들에 답하는 것을 목표로 한다:

 

(1) Ctrl-World는 높은 제어 가능성을 유지하면서 공간적·시간적으로 모두 일관된 장기 지평 rollout을 생성할 수 있는가?

 

(2) Ctrl-World는 상상 공간에서 서로 다른 generalist 로봇 정책을 신뢰성 있게 평가하여, 그것들의 실세계 성능 순위를 충실히

재현할 수 있는가?

 

(3) Ctrl-World는 전적으로 그 상상 안에서 성공 궤적을 발견하고 합성함으로써 정책의 지시 이행을 개선할 수 있는가?

 

5.1 Experiment Setups

DROID Platform and Dataset.

 

실험은 DROID 플랫폼(Khazatsky et al., 2024)을 사용하는데, 이는 Robotiq Gripper가 장착된 Panda 로봇 팔을 특징으로 한다.

이 플랫폼은 하나의 손목 시점 카메라와 작업 공간을 관찰하는 무작위로 배치된 두 개의 3인칭 시점 카메라를 포함한다.

 

DROID 데이터셋(Khazatsky et al., 2024)은 564개 장면에서 수집된 95,599개의 다양한 궤적을 담고 있으며, 작업 공간의 조밀한

커버리지를 제공한다.

이는 약 76k개의 성공 궤적과 약 19k개의 실패 궤적을 포함한다.

 

다양한 액션과 실패 데이터의 포함은 결정적인데, 이것이 광범위한 미래 시나리오를 시뮬레이션할 수 있는 제어 가능한 월드 모델을 학습할 수 있게 하기 때문이다.

 

Training Details.

 

학습 동안, 본 연구 모델은 각각 192×320 해상도의 세 카메라 모두로부터의 출력을 공동으로 예측한다.

모델은 프레임 사이 1–2초 간격의 7개 프레임 히스토리에 조건화된다.

모델을 다음 15개 미래 액션에 조건화하는데, 이는 DROID에서 1초 액션 청크에 대응한다.

 

상호작용 동안, 정책의 출력이 15스텝보다 적으면, 액션 청크를 더미 액션(dummy actions)으로 패딩하고 유효한 액션에 대한

예측만 사용한다.

모델을 2×8 H100 GPU에서 총 배치 크기 64로 학습한다. 학습은 약 2–3일 걸린다.


5.2 World Model Quality Analysis (월드 모델 품질 분석)

Baseline and Evaluation Matrices.

 

본 연구는 Ctrl-World를 두 개의 선행 액션 조건 월드 모델 — World-model-based Policy Evaluation(WPE)(Quevedo et al., 2025)와 IRASim(Zhu et al., 2024) — 과 정량적으로 비교한다.

이 모델들은 단일 3인칭 카메라 시점만 예측하므로, 공정한 비교를 위해 단일 3인칭 카메라만 입력·예측하는

단일 시점 버전 Ctrl-World-third-view를 학습한다.

 

평가를 위해, 궤적의 2%를 검증 세트로 보류하고 각각 10초 길이의 256개 비디오 클립을 무작위로 샘플링한다.

rollout 동안, 월드 모델은 15스텝 액션 청크(1초에 대응)를 받아 다음 프레임을 10스텝 동안 자기회귀적으로 예측하여, 10초 길이

궤적을 생성한다.

 

그런 다음 예측 비디오를 계산 기반(PSNR(Hore & Ziou, 2010), SSIM(Wang et al., 2004))과 모델 기반(LPIPS(Zhang et al., 2018), FID(Heusel et al., 2017), FVD(Unterthiner et al., 2018)) 지표를 모두 사용해 ground truth와 비교한다.

 

Quantitative and Qualitative Results on Multi-step Interaction Trajectories.

검증 세트에 대한 인터랙티브 장기 궤적 생성의 정량 결과. 우리는 10초 궤적을 생성하여 월드 모델의 품질을 평가한다. 무작위로 샘플링된 초기 프레임이 주어지면, 모델은 각 상호작용에서 15스텝 액션 청크(1초 이상에 걸침)를 받아 10라운드 동안 자기회귀적으로 생성한다. 결과는 256개 클립에 대해 평균낸 것이다.

 

Table 1에서 보이듯이, Ctrl-World-third-view는 이러한 선행 모델들을 능가하며, 멀티뷰 공동 예측은 생성 품질을 더욱 개선한다.

선행 연구(Quevedo et al., 2025; Zhu et al., 2024)의 관찰과 일관되게, 본 연구는 또한 이러한 baseline들이 로봇-객체 상호작용을 포착하는 데 어려움을 겪고 흔히 환각적 예측을 생성함을 발견한다.

검증 세트로부터의 장기 지평 rollout에 대한 정성 결과. 선행 모델들은 단일 시점 예측에 의존하여, 부분 관측성과 환각(예: 초록 수건을 옮기지 못하거나 빨간 그릇을 잡지 못함)을 겪는다. 이와 대조적으로, Ctrl-World는 3인칭과 손목 시점 카메라로부터 공동 예측하여, ground truth와 정렬된 정밀한 미래 궤적을 산출한다.

 

예를 들어, Figure 3에서 보이듯이, 단일 시점 예측 방법 WPE, IRASim, Ctrl-World-third-view는 모두 초록 수건을 옮기거나 빨간 그릇을 잡는 데 실패한다.

이와 대조적으로, Ctrl-World는 접촉 이벤트와 객체 상태 변화에 대한 결정적이고 세밀한 정보를 제공하는 손목 카메라 시점의 공동 예측을 통해 로봇-객체 상호작용을 정밀하게 모델링한다.

 

Controllability of the World Model.

 

월드 모델의 핵심 요건은 서로 다른 액션에 조건화된 다양한 미래 결과를 시뮬레이션하는 능력이다.

Ctrl-World의 제어 가능성과 ablation. 서로 다른 액션 시퀀스는 Ctrl-World에서 센티미터 수준의 정밀도로 서로 구별되는 rollout을 생성할 수 있다. 메모리를 제거하면 흐릿한 예측(파란색)으로 이어지고, 프레임 수준 pose 조건화를 제거하면 제어 정밀도가 감소한다(보라색). t=4s 프레임을 예측할 때의 attention 시각화(왼쪽)는 같은 pose를 가진 t=0s 프레임에 강한 attention을 보이며, 메모리 검색의 효과를 예시한다. 명확성을 위해, 각 액션 청크는 자연어로 표현된다(예: "Z-axis -6 cm"). 공간 제약으로 인해, 중간 프레임은 손목 시점만 시각화된다.

 

본 연구 모델이 단 몇 센티미터만 다른 액션에 대해서도 정밀한 미래 예측을 생성하는 세밀한 제어 가능성을 보임을 발견한다(Figure 4 참조).

이 제어 가능성이 두 가지 주요 요인에서 발생한다고 가설을 세운다:

 

첫째, DROID 데이터셋의 조밀한 액션 공간 커버리지;

 

둘째, 본 논문의 멀티뷰 예측과 프레임 수준 액션 조건화 사용인데, 이는 ablation 연구로도 뒷받침된다.

 

Figure 4의 왼쪽에서, 본 연구는 t=4s 프레임을 예측할 때의 attention 가중치를 시각화하고 유사한 pose를 가진 t=0s 프레임에 강한 attention이 있음을 관찰하며, 이는 본 논문에서 제시하는 메모리 검색 메커니즘의 효과를 부각한다.

 

Consistency of the World Model.

 

손목 카메라의 경우, 카메라의 시야(field of view)가 단일 궤적 안에서 극적으로 변하기 때문에, 모델이 일관되고 장기적인 예측

생성하는 것이 어렵다.

Ctrl-World의 일관성. 손목 카메라의 시야가 단일 궤적 안에서 극적으로 변하기 때문에, 멀티뷰 정보와 메모리 검색을 활용하는 것이 일관된 손목 시점 예측 생성에 필수적이다. 초록 상자로 강조된 예측은 다른 카메라 시점으로부터 추론된 것이고, 빨간 상자의 예측은 메모리로부터 검색된 것이다.

 

Figure 5에서 보이듯이, 모델이 다른 카메라 시점과 히스토리 프레임 양쪽으로부터의 관련 정보를 효과적으로 활용하여, 일관된 손목 시점 예측을 생성할 수 있음을 발견한다.

메모리 구성 요소와 프레임 수준 조건에 대한 ablation은 Table 2에 있으며, 이는 각 구성 요소의 중요성을 확인한다.

Ctrl-World의 핵심 구성 요소에 대한 ablation. 메모리 메커니즘, 프레임 수준 액션 조건화, 또는 멀티뷰 공동 예측을 제거하는 것은 모두 성능 저하로 이어진다.

 


5.3 World Model for Policy Evaluation

 

Ctrl-World가 generalist 로봇 정책의 지시 이행 능력을 평가하고 실세계에서의 그것들의 성능 순위를 정확하게 반영하는 데 사용될 수

있는지 평가한다(Li et al., 2024).

본 연구는 DROID 플랫폼을 설정하고 작업 공간 주위에 두 개의 3인칭 카메라를 무작위로 배치한다.

 

선행 연구들이 DROID 정책이 새로운 셋업으로 일반화하는 것을 본 것(Pertsch et al., 2025)과 유사하게, 본 연구는 오픈소스 DROID 데이터셋으로만 사전학습된 Ctrl-World가 새로운 카메라 배치를 가진 새로 구성된 장면에서 zero-shot으로 정확한 미래 예측을 할 수 있음을 발견한다.

 

Policies and Tasks.

 

세 개의 공개된 정책 π₀(Black et al., 2023), π₀-FAST(Pertsch et al., 2025), π₀.₅(Intelligence et al., 2025)

본 연구의 DROID 플랫폼에서 Pick-and-Place, Towel-Folding, Drawer, Wipe-Table, Close-Laptop, Pull-tissue, Stack 태스크를 포함한 다양한 태스크에 걸쳐 평가한다.

 

실세계와 월드 모델 rollout을 동일한 초기 관측으로 초기화하고 Algorithm 1을 따라 각 정책을 실행한다.

실세계와 월드 모델에서의 π₀.₅ rollout 간 비교. 각 궤적은 π₀.₅와 Ctrl-World 간의 20회 상호작용을 담고 있다. 주목할 만하게, generalist 정책과 Ctrl-World 모두 우리의 새로운 DROID 셋업으로 zero-shot 전이한다.
실세계와 월드 모델 rollout 간의 정량적 상관관계. 월드 모델은 지시 이행 행동을 신뢰성 있게 포착하지만 실행 성공률을 과소평가하는 경향이 있다.

 

Figure 7에 지시 이행률(instruction following rates)과 성공률을 보고하고 Figure 6에 실제와 상상된 rollout 간의 정성적 비교를 시각화한다. 더 자세한 rollout 내용은 Appendix B에서 찾을 수 있다.

 

 

Comparison Between Real-World and World Model Rollouts.

 

위 결과는 정책의 고수준 지시 이행 행동이 월드 모델에서 실세계에서 관찰된 것과 밀접하게 상관됨을 보인다.

그러나 저수준 실행 평가에서 일부 격차를 알아차리는데, 구체적으로 충돌, 객체가 미끄러져 나가는 것, 회전 등 복잡한 물리 동역학의 정밀한 모델링에서이다 (예: Figure 6에서 노트북과의 상호작용이 부정확하다).

 

또한 generalist 정책들이 실세계에서 실패한 시도 후 계속 재시도하는 경향이 있는데, 월드 모델이 이를 때때로 포착하지 못함을 관찰한다. 일부 실패 궤적이 DROID 데이터셋에 포함되어 있지만, 데이터 분포 밖의 실패 모드가 여전히 많다.

 

추가적인 in-domain 정책 rollout 데이터를 수집하는 것이 학습된 동역학의 충실도를 개선하고 이 격차를 좁힐 것으로 기대한다(Team, 2025).


5.4 World Model for Policy Improvement (정책 개선을 위한 월드 모델)

 

Post-train Policy with Synthetic Data.

 

이제 Ctrl-World가 실세계 데이터 없이 VLA 모델을 개선하기 위한 합성 post-training 데이터를 생성하는 데 사용될 수 있는지

평가한다.

 

π₀.₅를 기본 정책으로 사용하고 Algorithm 1을 따른다.

4.2절에서 설명한 것처럼, (1) 태스크 지시를 재표현하거나 (2) 로봇 팔을 새로운 초기 상태로 리셋함으로써 rollout 다양성을 장려한다.

 

재표현의 경우, 본 연구는 LLM API(Team et al., 2023)를 호출해 지시를 paraphase한다(예: "place glove in box"를 "pick up the cloth and put it inside the box"로 변환).

 

리셋의 경우, 정책-상호작용이 시작되기 전에 새로운 목표 초기 위치를 무작위로 샘플링하고

선형 보간 모션 플래너(linear-interpolation motion planner)를 사용해 로봇 팔을 그곳으로 이동시킨다.

 

태스크당 400개의 궤적을 생성하고 인간 선호 판단에 기반해 25–50개의 성공 궤적을 유지한다.

이 선택 단계는 보상 모델로 자동화될 수 있는데, 이는 활발한 연구 영역이다(Ma, 2025).

 

마지막으로, 큐레이션된 합성 데이터셋에 대해 정책을 2k 스텝 동안 fine-tuning하여, 낯선 지시문과 객체에 대한 기본 모델의 능력을 개선한다.

 

상단 행은 post-training 태스크의 예시를 예시하고, 하단 행은 월드 모델 안에서 생성된 합성 궤적을 제시한다. 월드 모델은 성공과 실패 rollout을 모두 생성할 수 있다; 우리는 성공 궤적을 유지하고 정책 fine-tuning에 사용한다.

 

 

Results.

 

일부 대표적인 태스크 예시와 합성 궤적이 Figure 8에 시각화되어 있고, 정량 결과가 Figure 9에 보고되어 있다.

정책 개선. 합성 데이터에 대한 post-training은 정책 지시 이행을 평균 44.7% 개선한다.

 

사전학습된 π₀.₅ 정책이 낯선 객체와 새로운 지시문에서 낮은 성공률을 달성하는 반면, post-training모델을 새로운 지시와 정렬하고 이러한 downstream 태스크에서 성공률을 38.7%에서 83.4%로 높인다. 태스크 세부 사항을 Appendix C에 포함한다.


 

 

6. Conclusion

 

모델의 상상 안에서 전적으로 closed-loop 정책 평가와 개선을 지원하는, 로봇 조작을 위한 제어 가능한 월드 모델인

Ctrl-World를 제시했다.

 

Ctrl-World에서 평가된 정책들은 실세계의 것을 밀접하게 반영하는 지시 이행 행동을 보인다.

주목할 만하게, 생성된 데이터에 대한 post-training은 사전학습된 로봇 정책의 새로운 지시문에 대한 성공률을 38.7%에서 83.4%로 높인다.

이러한 유망한 결과에도 불구하고, 중요한 과제들이 남아 있다.

 

본 모델은 정밀한 상호작용이나 장기 지평 추론을 포함하는 태스크에서 실패할 수 있으며, 성능은 초기 관측에 민감하다.

이러한 한계들은 비디오 backbone이 시간이 지나며 더 물리적으로 정확하고 일관되어짐에 따라 줄어들 수 있다.

 

또한, 본 연구의 실험은 지시 이행 개선에 집중하며, 본 모델이 이전에 본 지시문에 대한 저수준 성공률 같은 다른 측면의 성능을 개선하기에는 충분히 정확하지 않다고 기대한다.

반복적 정책 rollout과 fine-tuning으로 모델을 개선하는 것은 흥미로운 향후 방향이다.


 

 

Appendix

A. MORE DETAILS FOR WORLD MODEL LEARNING

Model Architecture.

 

본 연구의 월드 모델은 Stable Video Diffusion(SVD)(Blattmann et al., 2023a)의 아키텍처를 밀접하게 따르며,

SVD 사전학습 체크포인트에서 초기화한다.

유일하게 새로 초기화되는 구성 요소는 7차원 Cartesian 공간 액션을 1024차원 latent 임베딩으로 투영하는 3-layer MLP이다.

 

입력 이미지는 먼저 8×8의 공간 다운샘플링 비율을 가진 VAE로 인코딩된다.

실제로, 본 연구는 k=7개의 히스토리 프레임을 사용하며, 각각 강건성을 개선하기 위해 독립적인 무작위 노이즈로 교란한다.

 

액션 조건화 윈도우를 1초, 즉 15개 액션 스텝에 대응하도록 설정한다.

GPU 메모리 소비를 줄이기 위해, 이 15개 액션을 Cartesian 공간에서 변환하고(B절 참조) 모델에 입력하기 전에

시간적으로 5스텝으로 다운샘플링한다.

각 프레임은 세 개의 192×320 이미지를 담고 있으며, 이는 24×40 형태의 latent feature로 인코딩된다.

 

결과적으로 전체 입력 토큰 형태는 B × (7+5) × (3×24×40)이며, 이는 그런 다음 spatial-temporal transformer backbone에 의해 처리된다.

 

Training Datasets.

 

DROID 데이터셋의 95k개 궤적 전부를 사용한다.

각 학습 스텝마다, 궤적을 무작위로 샘플링한 다음 그 궤적 안에서 프레임을 균등하게 샘플링하여 현재 프레임으로 삼는다.

그런 다음 시간을 거슬러 뒤로 샘플링하여 메모리 프레임을 검색하고 모델의 예측 목표를 후속 미래 프레임으로 설정한다.

 

Training Process.

 

모델을 2×8 H100 GPU에서 총 배치 크기 64로 학습한다.

학습률은 1e-5로 설정되며, 100k 스텝 동안 학습하는데, 이는 완료하는 데 약 2–3일 걸린다.


B. MORE DETAILS FOR POLICY EVALUATION

Details on interaction between policy and world model.

 

Ctrl-World와 상호작용하기 위해 https://github.com/Physical-Intelligence/openpi 의 공식

π₀-DROID, π₀-FAST-DROID, π₀.₅-DROID 정책을 직접 사용한다.

 

Ctrl-World는 최신 VLA 모델 간의 policy-in-the-loop 상호작용을 가능하게 하는 최초의 월드 모델이다.

 

이러한 오픈소스 정책들은 관절 각도(joint angles)와 두 개의 카메라 시점을 입력으로 받아 관절 속도(joint velocities)를 출력한다.

 

이와 대조적으로, 본 연구의 월드 모델은 Cartesian 공간의 end-effector pose에 조건화한다.

이 불일치를 잇기 위해, DROID 데이터셋에서 현재 관절 각도 q_t^joint와 예측된 관절 속도 a_{t+1:t+H}^jv를

미래 관절 구성 q_{t+1:t+H}^joint로 매핑하는 adapter를 학습한다.

 

그런 다음 Franka Panda 순기구학(forward kinematics, FK)을 적용해 이 관절 구성을 Cartesian 공간 pose q_{t+1:t+H}^cartesian로 변환한다. 이 adapter는 단순한 2-layer MLP로 구현된다.

 

전체 과정은 다음과 같다:

현재 관절 구성 q_t^joint, 멀티뷰 관측 o_t, 언어 지시 l이 주어지면, 정책은 H-스텝 관절 속도를 출력한다:

 

이것들은 adapter를 통과해 미래 관절 구성을 예측하고, 이어서 FK로 Cartesian pose를 계산한다:

 

마지막으로, 월드 모델은 현재 관측, 계산된 Cartesian pose, 그리고 히스토리 Cartesian pose에 조건화하여 다음 H개 프레임을 예측한다:

 

이 셋업은 완전한 자기회귀 rollout을 가능하게 하여, 공식 π₀-DROID, π₀-FAST-DROID, π₀.₅-DROID 정책과 Ctrl-World가 상상 공간에서 매끄럽게 상호작용할 수 있게 한다.

 

Breakdown for policy evaluation.

 

Table 3에 지시 이행률과 저수준 실행 성공률을 제시한다.

 

Task details and Criterion.

 

본 연구의 실험에서, 각 궤적이 성공인지 실패인지 평가하기 위해 인간 어노테이터를 사용한다.

이 평가 과정은 향후 대규모 비전-언어 보상 모델을 사용해 자동화될 수 있지만, 이 논문에서 초점은 월드 모델 자체이므로,

보상 신호로 인간 선호에 의존한다.

 

본 연구는 궤적이 단지 지시를 따르는지 아니면 완전한 태스크 성공을 달성하는지 판단하기 위한 명확한 기준을 제공한다:

  • Pick-place: 여러 객체와 용기(receptacle)가 탁상에 놓여 있다. 지시는 "A를 집어 B에 놓아라" 형태이다. 궤적은 정책이 올바른 객체 A를 잡으려 시도하면 지시를 따른 것으로 간주된다. 객체 A가 목표 용기 B에 성공적으로 놓이면 성공으로 간주된다.
  • Fold the Towel: 수건이 테이블에 평평하게 놓여 있고, 다른 객체가 있을 수도 있다. 지시는 "수건을 접어라"이다. 궤적은 그리퍼가 수건 가장자리로 이동해 들어 올려 접으려 시도하면 지시를 따른 것으로 간주된다. 궤적은 최종적으로 수건의 표면적이 절반이 되면 성공으로 간주된다.
  • Drawer: 지시는 "객체 A를 서랍에 넣어라"이다. 궤적은 로봇이 객체 A를 서랍 안에 넣으려 시도하면 지시를 따른 것이다. 객체 A가 최종적으로 서랍에 놓이면 성공이다.
  • Wipe Table: 지시는 테이블 표면을 닦는 것이다. 궤적은 그리퍼가 수건과 접촉해 쓸기 동작으로 움직이면 지시를 따른 것이다. 테이블의 큰 부분이 쓸기 동작으로 덮이면 성공으로 간주된다.
  • Close Laptop: 지시는 열린 노트북을 닫는 것이다. 궤적은 그리퍼가 노트북 덮개에 접근하면 지시를 따른 것이다. 덮개가 완전히 닫히면 성공으로 간주된다.
  • Pull Tissue: 지시는 티슈 상자에서 티슈를 뽑는 것이다. 궤적은 그리퍼가 티슈 슬롯에 접근해 티슈를 집으면 지시를 따른 것이다. 최소 하나의 티슈가 완전히 추출되면 성공으로 간주된다.
  • Stack: 지시는 한 객체를 다른 객체 위에 쌓는 것이다. 궤적은 그리퍼가 올바른 객체를 들어 올리면 지시를 따른 것이다. 객체가 목표 객체 위에 안정적으로 놓이면 성공이다.

C. MORE DETAILS FOR POLICY IMPROVEMENT

Fine-tuning Process.

 

공식 코드베이스 https://github.com/Physical-Intelligence/openpi 를 기반으로 π₀.₅-DROID 정책을 fine-tuning한다.

사전학습된 체크포인트를 본 연구의 합성 데이터셋에 대해 4개 H100 GPU에서 2k 스텝 동안 fine-tuning한다.

 

Task Descriptions :

  • 공간 이해(Spatial Understanding) 태스크: 2–6개의 무작위 객체가 테이블에 놓인다. 정책은 지정된 공간 위치의 객체를 집어 상자에 놓도록 지시받는다. 예시 지시는 "오른쪽 위(top-right)의 객체를 집어 상자에 놓아라" 또는 "가장 왼쪽(far-left)의 객체를 상자에 넣어라"를 포함한다.
    • 형태 이해(Shape Understanding) 태스크: 2–3개의 무작위 객체가 테이블에 놓이며, 일부는 같은 속성을 공유하지만 크기가 다르다. 정책은 크기에 기반해 객체를 구별해야 한다. 예시 지시: "더 큰 빨간 블록을 집어 상자에 놓아라."
  • 방향이 있는 수건 접기(Towel-Folding with Directions): 수건과 다른 방해물(distractor)이 테이블에 놓이며, 정책은 특정 접기 방향을 지정하는 지시를 받는다(예: "수건을 왼쪽에서 오른쪽으로 접어라").
  • 새로운 객체(Novel Objects): 사전학습된 정책이 잘 식별하지 못하는 장갑과 스테이플러 같은 미지의 객체를 도입한다.

Detailed success rate.

 

각 범주 내의 상세 태스크 성공률을 제공한다.

 

 

 

Interactive VideoGPTs are Scalable World Models — Tsinghua · Huawei Noah's Ark Lab (NeurIPS 2024)

https://arxiv.org/abs/2405.15223

 

iVideoGPT: Interactive VideoGPTs are Scalable World Models

World models empower model-based agents to interactively explore, reason, and plan within imagined environments for real-world decision-making. However, the high demand for interactivity poses challenges in harnessing recent advancements in video generativ

arxiv.org

 

Authors: Jialong Wu¹*, Shaofeng Yin¹²*, Ningya Feng¹, Xu He³, Dong Li³, Jianye Hao³⁴, Mingsheng Long¹
Affiliation: ¹School of Software · BNRist, Tsinghua University · ²Zhili College, Tsinghua University · ³Huawei Noah's Ark Lab · ⁴Tianjin University
Venue: NeurIPS 2024 (38th Conference on Neural Information Processing Systems)

 

Preprint: arXiv:2405.15223v3 [cs.CV] (초판 May 2024, v3 개정 Oct 2024)

0. Abstract

 

월드 모델(World models)은 모델 기반 에이전트(model-based agents)가 실세계 의사결정을 위해 상상된 환경(imagined environments) 안에서 상호작용적으로 탐색하고, 추론하고, 계획할 수 있도록 한다.

그러나 상호작용성(interactivity)에 대한 높은 요구는 비디오 생성 모델(video generative models)의 최근 발전을 대규모 월드 모델

개발에 활용하는 데 있어 어려움을 제기한다.

 

 

이 연구는 Interactive VideoGPT(iVideoGPT)를 소개하는데, 이는 멀티모달 신호시각 관측(visual observations), 액션(actions), 보상(rewards) — 을 하나의 토큰 시퀀스로 통합하여 next-token 예측을 통해 에이전트의 상호작용적 경험을 가능하게 하는 확장 가능한 자기회귀 transformer 프레임워크이다.

 

 

iVideoGPT는 고차원 시각 관측을 효율적으로 이산화(discretize)하는 새로운 "압축 토큰화(compressive tokenization)" 기법을

특징으로 한다.

확장 가능한 아키텍처를 활용하여, iVideoGPT를 수백만 개의 인간 및 로봇 조작 궤적으로 사전학습할 수 있으며, 이를 통해 다양한 downstream 태스크를 위한 상호작용적 월드 모델로서 기능하도록 적응 가능한 다목적(versatile) 기반(foundation)을 구축한다.

 

 

이러한 태스크에는 액션 조건 비디오 예측(action-conditioned video prediction), 시각 계획(visual planning), 그리고 모델 기반 강화학습(model-based reinforcement learning)이 포함되며, 여기서 iVideoGPT는 최신(state-of-the-art) 방법들과 비교해

경쟁력 있는 성능을 달성한다.

 

 

본 연구는 상호작용적 범용 월드 모델(interactive general world models)의 개발을 진전시키며, 생성형 비디오 모델(generative video models)과 실용적인 모델 기반 강화학습 응용 사이의 간극을 잇는다.

코드와 사전학습 모델은 https://thuml.github.io/iVideoGPT 에서 참고 가능하다.

 


 

1. Introduction 

 

 

최근 몇 년간 텍스트, 오디오, 이미지를 포함한 멀티모달 콘텐츠의 생성 모델에서 괄목할 만한 발전이 있었으며, 이제 비디오 생성이 새로운 프런티어(frontier)로 부상하고 있다.

 

다양한 인터넷 규모(Internet-scale) 데이터에 대해 비지도 방식으로 학습된 이러한 생성형 비디오 모델의 특히 중요한 응용 중 하나는

대규모의 예측형 월드 모델(predictive world models)을 구축하는 것이다.

 

이러한 월드 모델은 세계가 작동하는 방식에 대한 상식적 지식(commonsense knowledge)을 축적하여, 에이전트의 액션에 기반한 잠재적 미래 결과(예: 시각 관측 및 보상 신호)의 예측을 가능하게 할 것으로 기대된다.

 

이러한 월드 모델을 활용함으로써, 모델 기반 강화학습(model-based RL)을 사용하는 에이전트는 월드 모델 안에서 상상하고, 추론하고, 계획할 수 있으며, 따라서 실세계에서 소수의 시도만으로 더 안전하고 효율적으로 새로운 기술을 습득할 수 있다.

이러한 근본적 연결에도 불구하고, 비디오 생성을 위한 생성 모델과 에이전트 학습을 위한 시각 월드 모델 사이에는

여전히 상당한 간극이 존재한다.

 

 

한 가지 주요한 과제는 상호작용성(interactivity) 확장성(scalability)의 양쪽 모두를 최대한 달성하는 것이다.

모델 기반 RL에서 월드 모델은 주로 순환 신경망(recurrent network) 아키텍처를 사용한다.

이 설계는 각 스텝에서 액션에 조건화된 관측 또는 잠재 상태(latent states)의 전이를 자연스럽게 허용하여, 상호작용적 행동 학습을 용이하게 한다. 그러나 이러한 순환 모델들은 대부분 단순한 시각 요소를 가진 게임이나 시뮬레이션 환경에 집중되어 있으며, 복잡한

in-the-wild 데이터를 대규모로 모델링하는 능력이 제한적이다. 

 

 

한편, 인터넷 규모의 생성형 비디오 모델은 생성 시작 시점에 텍스트 설명이나 미래 액션 시퀀스를 통해 제어 가능한, 사실적인 긴 비디오를 합성할 수 있다. 이들은 고수준 계획(high-level planning)에는 적합하지만, 그것들의 궤적 수준(trajectory-level) 상호작용성은 에이전트가 정밀한 기초 기술을 효율적으로 학습하기 위해 시뮬레이션 도중 스텝 단위로 개입하는 데 필요한 세밀함(granularity)

충분히 제공하지 못한다.

 

이 딜레마는 자연스럽게 다음 질문을 제기한다:

 

확장 가능한 비디오 생성 모델의 발전을 상호작용적 시각 월드 모델 개발에 어떻게 활용할 수 있는가?

 

 

이 연구에서, GPT류(GPT-like) 자기회귀 transformer 프레임워크 안에서 상호작용적이면서도 확장 가능한 월드 모델을 탐구한다.

 

 

최근 diffusion 모델masked 생성 모델을 통해 선구적인 시도들이 이루어졌다.

그럼에도 불구하고, 자기회귀 transformer를 활용하는 것은 확립된 대규모 언어 모델(LLM) 생태계와의 매끄러운 통합, 그리고 adapter 모듈과 같은 특정 아키텍처 수정 없이도 다양한 조건을 처리하는 더 큰 유연성 같은 뚜렷한 이점을 제공한다.

 

 

본 연구는 시각 관측, 액션, 보상을 포함한 멀티모달 신호를 상호작용적으로 자기회귀적인(interactively autoregressive) 방식으로 통합하는 확장 가능한 월드 모델 아키텍처인 Interactive VideoGPT(iVideoGPT)를 제시한다.

 

 

이미지 토크나이저를 사용해 시각 관측을 프레임 단위로(frame-by-frame) 토큰으로 이산화하는 멀티모달 LLM과 달리, 확장성을 높이기 위한 iVideoGPT의 핵심 혁신은 풍부한 문맥적 관측(contextual observations)에 조건화하여 각 관측에 대한 압축 토큰화(compressive tokenization)를 학습하는 것으로, 토큰 시퀀스 길이를 점근적으로(asymptotic) 16배 감소시킨다.

 

 

더 압축적인(compact) 비디오 토큰화가 더 효율적인 학습과 생성을 용이하게 할 뿐만 아니라 비디오 품질도 향상시킬 수 있음을 강조한다. 이는 문맥(context)을 동역학(dynamics)으로부터 분리(decoupling)함으로써 달성되며, 모델이 장면 내 시간적 일관성을 유지하면서 객체의 움직임(motion) 예측에 집중할 수 있게 한다.

iVideoGPT의 실용적 응용. iVideoGPT는 확장 가능하도록 설계되어 수백만 개의 인간 및 로봇 조작 궤적에 대한 사전학습을 가능하게 한다. 이는 광범위한 downstream 태스크에 적응 가능한, 단일의 다목적(versatile) 상호작용적 월드 모델 기반(foundation)으로 귀결된다.

 

Figure 1에 예시된 것처럼, 시각 로봇 조작(visual robotic manipulation)을 위한 iVideoGPT의 일련의 실용적 응용을 시연한다.

 

 

LLM에 의해 대중화된 2단계 접근을 본떠, 본 방법은 사전학습(pre-training) 후 도메인 특화 적응(domain-specific adaptation)을 포함한다.

사전학습 동안, iVideoGPT는 100만 개가 넘는 로봇 및 인간 조작 궤적의 혼합에 대한 action-free 비디오 예측에 대해 확장 가능하다.

사전학습된 iVideoGPT는 액션 조건 비디오 예측, 시각 계획, 시각 모델 기반 RL과 같은 다양한 downstream 태스크를 위한 단일의, 적응 가능한 상호작용적 월드 모델 기반으로 기능한다.

 

 

또한, 본 연구는 미지의 도메인에 대해 토크나이저 적응만을 요구하는

"fine-tuning 없는 사전학습 transformer의 예비적(preliminary) zero-shot 비디오 생성 능력"을 보여준다.

나아가 시퀀스 모델링의 유연성을 부각하는, goal-conditioned 비디오 예측을 위한 iVideoGPT의 변형(variant)을 탐구한다.

 

 

이 연구의 주요 기여는 다음과 같이 요약할 수 있다:

 

1. 시각 관측을 위한 압축 토큰화를 특징으로 하는, 확장 가능한 월드 모델을 위한 자기회귀 transformer 아키텍처인 Interactive VideoGPT(iVideoGPT)를 소개한다.

 

2. 수백만 개의 로봇 및 인간 조작 궤적으로 구성된 대규모 데이터셋으로 iVideoGPT를 사전학습하고 이를 도메인 특화 태스크에 적응시킨다. 사전학습된 모델들은 추가 연구를 장려하기 위해 공개되었다.

 

3. 비디오 예측, 시각 계획, 시각 모델 기반 RL을 아우르는 광범위한 실험은 iVideoGPT가 정확하고 사실적인 경험을 시뮬레이션할 수 있으며 최신 방법들과 비교해 경쟁력 있는 성능을 제공함을 보여준다.

 


 

2. Problem Formulation 

 

월드 모델(world model)은 환경을 시뮬레이션하기 위해 에이전트가 학습한 내부 모델(internal model)이다.

 

이 환경은 일반적으로 부분 관측 마르코프 결정 과정(partially observable Markov decision process, POMDP)

(𝒮, 𝒪, φ, 𝒜, p, r, γ)로 모델링된다.

 

각 스텝에서, s_t ∈ 𝒮는 환경의 근본 상태(underlying state)를 나타내고,

o_t = φ(s_t)는 에이전트가 받는 관측(observation)으로, s_t에 대한 불완전한 정보만을 제공한다.

액션 a_t ∈ 𝒜를 취한 후, p(s_{t+1}|s_t, a_t)는 상태 s_t에서 s_{t+1}로의 전이 확률(transition probability)을 정의한다.

에이전트는 또한 즉각적 보상 r_{t+1} = r(s_t, a_t)를 받으며, γ-할인 누적 보상 𝔼_{p,π}[Σ_t γ^{t−1} r_t]를 최대화하는 정책 π를 학습하는 것을 목표로 하되, 이때 a_t ∼ π(o_{1:t})이다.

 

 

월드 모델은 많은 유형의 데이터로부터 학습될 수 있지만, 비디오는 태스크 불문(task-agnostic)이고, 널리 이용 가능하며, 자기지도 방식(self-supervised way)으로 학습될 수 있는 광범위한 지식을 담고 있는 하나의 모달리티이다.

 

따라서, 본 연구는 시각 제어(visual control)를 위한 월드 모델 학습을 인터랙티브 비디오 예측(interactive video prediction) 문제로 정식화하며, 여기서 𝒪 = R^{H×W×3}는 비디오 프레임의 공간이다.

 

구체적으로, T₀개 프레임의 짧은 히스토리 시각 관측 o_{1:T₀}가 주어지면, 각 스텝 t = T₀, …, T−1에서 에이전트는 자신의 정책이전의 상상된 관측(previous imagined observations)에 기반하여 액션 a_t를 취하고, 그런 다음 월드 모델은 에이전트에게 피드백하기 위해 전이 p(o_{t+1}, r_{t+1} | o_{1:t}, a_{T₀:t})를 근사하고 샘플링해야 한다.

 

 

아키텍처 간 개념적 비교, 단순화를 위해 단일 문맥 프레임(T₀ = 1)을 사용하여 예시함. (a) Dreamer[29]와 MuZero[80] 같은 월드 모델을 위한 순환 아키텍처는 스텝 수준(step-level) 상호작용성을 제공하지만 확장성은 제한적이다. (b) VideoGPT[101]와 Stable Video Diffusion[8, 7] 같은 최근 비디오 생성의 발전은 궤적 수준(trajectory-level) 상호작용성만 제공할 수 있는 non-causal 시간 모듈을 사용한다. (c) 우리 모델은 각 스텝을 개별적으로 토큰 시퀀스로 매핑하는 자기회귀 transformer를 활용하여, 확장성과 상호작용성을 모두 달성한다.

 

 

Figure 2에 묘사된 것처럼, VideoGPT를 포함한 대부분의 진보된 비디오 생성 모델[101, 8, 104]은 인터랙티브 비디오 예측 문제를 다룰 수 없는데, 이는 그것들이 시간 차원(temporal dimension)을 따라 정보를 융합하는 non-causal 모듈을 설계하여, 생성 도중의 인과적·중간적 액션 제어(causal, intermediate action control) 능력이 결여되어 있기 때문이다(확장된 논의는 Appendix C.2 참조).

MBRL 문헌의 기존 월드 모델들[29, 80], 예컨대 Dreamer는 순환 아키텍처를 활용하지만 확장성이 부족하다.

 


 

3. Interactive VideoGPT

 

시각 관측·액션·보상 및 기타 잠재적 감각 입력을 포함한 멀티모달 신호를 통합할 큰 유연성을 갖춘, 확장 가능한 월드 모델 아키텍처인 Interactive VideoGPT를 소개한다. 그 핵심에서, iVideoGPT는 비디오 프레임을 이산화하는 "압축 토크나이저(compressive tokenizer)"와 후속 토큰을 예측하는 "자기회귀 transformer"로 구성된다(3.1절).

이 모델은 다양한 인간 및 로봇 조작 비디오에 대한 사전학습을 통해 여러 장면 속 움직임과 상호작용에 대한 공통 지식을 습득할 수 있으며, 그런 다음 추가 모달리티를 통합하는 downstream 태스크로 효과적으로 전이할 수 있다.


3.1 Architecture

iVideoGPT의 아키텍처, 단일 문맥 프레임(T₀ = 1)만 보이도록 단순화함. (a) 압축 토큰화는 시간적 중복성을 처리하기 위해 문맥 프레임에 조건화하여 미래 프레임을 이산화하는 조건부 VQGAN을 활용하여, 비디오 토큰 수를 유의미하게 감소시킨다. (b) 자기회귀 transformer는 멀티모달 신호 — 시각 관측·액션·보상 — 을 토큰 시퀀스로 통합하여, next-token 예측을 통해 인터랙티브 에이전트 경험을 가능하게 한다. 액션과 보상은 선택적이며 action-free 비디오 사전학습에는 포함되지 않는다.

 

Compressive tokenization.

 

Transformer는 이산 토큰의 시퀀스에 대해 작동하는 데 특히 뛰어나다.

VQGAN은 원시 픽셀에서 이산 토큰으로 변환하는 데 흔히 사용되는 시각 토크나이저이다.

 

각 프레임을 독립적으로 이산화하여 시퀀스 길이가 급증하게 만드는 이미지 토크나이저를 사용하거나, 상호작용성을 희생하면서 비디오를 시공간적으로 압축하는 3D 토크나이저를 사용하는 대신, 본 연구는 "이중 인코더·디코더 {(E_c, D_c), (E_p, D_p)}"로 구성된 새로운

"조건부 VQGAN(conditional VQGAN)"으로 비디오를 토큰화할 것을 제안한다.

(²이 연결로 인해, 본 연구는 "비디오 프레임"과 "시각 관측"이라는 용어를 서로 바꾸어 사용한다.)

 

 

Figure 3a에 예시된 것처럼, 문맥 정보가 풍부한 초기 문맥 프레임 o_{1:T₀}는 N개 토큰을 통해 독립적으로 토큰화되고 재구성된다:

 

이와 대조적으로, 문맥 프레임과 미래 프레임 사이의 시간적 중복성(temporal redundancy) 때문에, 움직이는 객체의 위치와 자세 같은 필수 동역학 정보만 인코딩되면 된다.

이는 훨씬 적은 수의 n개 토큰(n ≪ N)을 요구하는 조건부 인코더와 디코더를 사용해 달성된다:

 

본 연구는 이 조건화 메커니즘을 다중 스케일 feature map 간의 cross-attention을 사용해 구현한다.

(자세한 내용은 Appendix A.1 참조).

 

전체적으로, 제안된 토크나이저는 다음 목적함수로 학습된다:

 

 

여기서 𝓛_VQGAN(o; E, D)는 L₁ 재구성 손실, commitment 손실, perceptual 손실, 그리고 선택적으로 adversarial 손실의 조합이다.

 

제안된 토큰화에는 주로 두 가지 이점이 있다:

 

첫째, 이는 토큰화된 비디오의 시퀀스 길이를 유의미하게 감소시키는데, 이 길이는 프레임 수에 따라 선형적으로 증가하되 훨씬 작은 비율 n으로 증가한다. 이 연구에서는 N = 16×16, n = 4×4로 설정하여, 점근적으로 16배의 감소를 가져오고, 모델 기반 계획과 강화학습을 위한 더 빠른 rollout을 용이하게 한다.

 

둘째, 조건부 인코딩을 통해, 후속 토큰을 예측하는 transformer는 문맥의 시간적 일관성을 훨씬 쉽게 유지하고 필수 동역학 정보 모델링에 집중할 수 있다. 본 연구의 토큰화의 가정과 한계를 6절에서 논의한다.

 

Interactive prediction with Transformers.

 

토큰화 후, 비디오는 토큰 시퀀스로 평탄화된다:

x = (z_1^{(1)}, …, z_1^{(N)}, [S], z_2^{(1)}, …, z_2^{(N)}, …, [S], z_{T₀+1}^{(1)}, …, z_{T₀+1}^{(n)}, …),

 

이때 길이는 L = (N+1)T₀ + (n+1)(T−T₀) − 1이다. 특수한 slot 토큰 [S]가 삽입되어 프레임 경계를 구분하고 액션 같은 추가 저차원

모달리티의 통합을 용이하게 한다(자세한 내용은 3.3절 참조).

 

 

위 Figure 3b처럼, GPT류 자기회귀 transformer가 프레임 단위 next-token 생성을 통한 인터랙티브 비디오 예측에 활용된다.

이 연구에서 GPT-2의 모델 크기를 취하되, LLM 아키텍처를 위한 최신 혁신을 수용하기 위해 LLaMA 아키텍처를 채택하여, RMSNorm을 사용한 pre-normalization, SwiGLU 활성화 함수, 그리고 rotary positional embedding을 적용한다.


3.2 Pre-Training 

대규모 언어 모델은 next-word 예측을 통해 자기지도 방식으로 인터넷 텍스트로부터 광범위한 지식을 얻을 수 있다.

유사하게, 월드 모델을 위한 action-free 비디오 사전학습 패러다임은 비디오 예측을 사전학습 목적으로 삼아, LLM에 없는 물리 세계 지식을 담은 인터넷 규모의 감독(supervision)을 제공한다.

 

이 일반적 목적함수로 iVideoGPT를 사전학습하며, 후속 비디오 토큰을 예측하기 위해 cross-entropy 손실을 적용한다:

 

 

여기서 L은 전체 시퀀스 길이이고, (N+1)T₀+1은 예측될 프레임의 첫 토큰 인덱스를 표시한다.

특히, 앞서 논의한 대로 iVideoGPT가 문맥 프레임을 생성하도록 학습시키지 않으며, 이를 통해 그것의 용량(capacity)이 동역학 정보에 집중하게 한다.

 

Pre-training data

 

인터넷에는 수많은 비디오가 이용 가능하지만, 계산상의 제약으로 인해 iVideoGPT를 로봇 조작 도메인에 특화하여 사전학습한다.

Open X-Embodiment(OXE) 데이터셋의 35개 데이터셋과 Something-Something v2(SSv2) 데이터셋의 혼합을 활용하며,

총 140만 궤적이다(자세한 내용은 Appendix A.2 참조).

OXE는 다양한 로봇 embodiment, 장면, 태스크에서 온 로봇 학습 데이터셋의 다채로운 집합이다.

이 데이터셋들은 매우 이질적(heterogeneous)이지만 action-free 비디오 예측 태스크로 쉽게 통일될 수 있다.

다양성을 더욱 높이기 위해 인간-객체 상호작용 비디오 데이터셋인 SSv2도 포함하는데, 이는 선행 연구가 이러한 인간 조작 비디오로부터 로봇 조작 태스크를 위한 월드 모델 학습으로의 지식 전이를 입증했기 때문이다.

 

Flexibility of sequence modeling.

 

토큰 시퀀스는 태스크, 입력, 출력을 명시하는 유연한 방법을 제공한다.

이 유연성을 예비적으로 보여주기 위해, 본 연구는 goal-conditioned 비디오 예측을 위한 iVideoGPT의 변형을 소개한다:

 

여기서 모델은 지정된 목표 관측 o_T에 도달하는 비디오 시퀀스를 예측한다.

 

이는 아키텍처와 학습 절차를 위와 동일하게 유지하면서 프레임 시퀀스를 õ_{1:T} = (o_T, o₁, o₂, …, o_{T−1})로 단순히 재배열함으로써 달성된다(자세한 내용은 Appendix A.2 참조).

goal-conditioned 예측의 정성적 결과는 Figure 4에 나와 있으며, 추가 탐구는 향후 연구로 남겨둔다³.

(³별도로 명시하지 않는 한, action-free 및 goal-free 비디오 예측이 모든 실험을 위한 사전학습 모델을 얻기 위한 기본(default) 사전학습 목적으로 사용된다.)


3.3 Fine-Tuning

 

Action conditioning & reward prediction

 

본 아키텍처는 또한 인터랙티브 월드 모델 학습을 위해 추가 모달리티를 유연하게 통합하도록 설계되었으며, 이는 위 Figure 3b에

드러나 있다.

액션은 linear projection을 통해 slot 토큰 임베딩에 더해짐으로써 통합된다.

보상 예측의 경우, 독립적인 보상 예측기를 학습하는 대신, 각 관측의 마지막 토큰의 hidden state에 linear head를 추가한다.

이 멀티태스크 학습 접근은 태스크 관련 정보에 대한 모델의 집중을 높여, 제어 태스크에 대한 예측 정확도를 향상시킬 수 있다.

 

본 연구에서는 위 식의 cross-entropy 손실에 더하여 보상 예측을 위한 평균제곱오차(MSE) 손실을 사용한다.

 

Tokenizer adaptation

 

본 연구는 downstream 태스크를 위해 토크나이저를 포함한 전체 모델을 갱신하기로 선택하며, 이 전략이

파라미터 효율적 fine-tuning 방법보다 더 효과적임을 발견한다.

이는 인터넷 규모 이미지에 비해 본 연구의 사전학습 데이터의 다양성이 제한적이기 때문일 가능성이 높은데, 이 데이터는 방대하긴 하지만 로봇공학 같은 특정 실세계 응용을 충분히 커버하지 못할 수도 있다.

 

VQGAN 토크나이저를 도메인 특화 데이터에 적응시키는 것을 탐구한 문헌은 거의 없다.

 

본 연구의 토큰화는 동역학 정보를 문맥 조건으로부터 분리하도록 설계되었으므로, 저자들은 모델이 downstream 태스크에서 서로 다른 로봇 유형 같은 미지의 객체를 마주칠 수는 있어도, transformer가 다양한 장면으로부터 학습한 물리의 근본 지식 — 움직임과 상호작용 같은 — 은 공통적으로 공유된다고 가설을 세운다.

이 가설은 iVideoGPT를 혼합 사전학습 데이터에서 미지의 BAIR 데이터셋으로 전이하는 실험에 의해 뒷받침되는데, 여기서 사전학습된 transformer는 미지의 로봇 그리퍼에 대해 토크나이저만 fine-tuning하면 되면서 자연스러운 움직임을 예측하도록 zero-shot 일반화할 수 있다(Figure 8 참조).

이 특성은 GPT류 transformer를 큰 크기로 확장하는 데 특히 중요한데, transformer를 그대로 유지하면서

도메인 간 경량 정렬(lightweight alignment)을 가능하게 하기 때문이다.

 


 

 

4. Experiments

 

저자들은 iVideoGPT를 세 가지 서로 다른 제어 관련(control-relevant) 설정에서 평가하고 그 성능을 기존 최신 방법들과 비교한다.

iVideoGPT가 다양한 태스크에 걸쳐 경쟁력 있는 성능을 제공할 만큼 다목적임을 시연하고, 토큰화 및 예측 능력, 데이터 효율성, 모델 스케일링, 계산 효율성을 이해하기 위한 심층 분석을 수행한다.

실험 세부 사항은 Appendix A에서 찾을 수 있다.


4.1 Video Prediction

Setup.

 

BAIR 로봇 pushing 데이터셋은 43k개의 학습 비디오와 256개의 테스트 비디오로 구성되며, 여기서 저자들은 선행 연구의 표준 프로토콜에 따라 단일 초기 프레임으로부터 15개 프레임을 예측한다.

RoboNet 데이터셋은 7개 로봇 팔에 걸친 162k개의 비디오를 담고 있다.

선행 연구를 따라, 본 연구는 테스트에 256개 비디오를 사용하여 두 개 프레임으로부터 10개 프레임을 예측한다.

특히, RoboNet은 본 연구의 사전학습 데이터인 OXE와 겹치며, 저자들은 여기서 테스트 비디오를 신중하게 필터링했다.

variational, diffusion, masked, autoregressive 모델을 포함한 다양한 비디오 예측 모델과 네 가지 지표 — FVD, PSNR, SSIM, LPIPS — 에 걸쳐 비교한다.

 

Results.

BAIR 로봇 pushing과 RoboNet 데이터셋에 대한 비디오 예측 결과. 우리는 세 번의 실행에 대해 계산된 각 지표의 평균과 표준편차를 보고한다. "-"는 원 논문에서 해당 값이 보고되지 않았음을 표시한다. LPIPS와 SSIM 점수는 편의상 표시를 위해 100배로 스케일링된다.

 

Table 1에서 보이듯이, iVideoGPT는 최신 방법들 — BAIR의 MAGVIT와 RoboNet의 FitVid — 과 비교해 경쟁력 있는 성능을 제공하면서, 그 아키텍처에서 상호작용성과 확장성을 모두 달성한다.

 

처음에 action-free로 사전학습된 본 연구의 모델은 액션 조건화를 유연하게 허용하며, 이는 BAIR의 FVD를 거의 20% 유의미하게 개선한다.

 

주 실험은 64×64의 낮은 해상도에서 수행되지만, iVideoGPT는 RoboNet의 경우 256×256으로 쉽게 확장될 수 있다.

프레임 단위 토큰화를 활용하는 선행 방법인 MaskViT가 VQGAN 재구성에서 시간적 비일관성깜빡임(flicker) 아티팩트

겪는다는 점을 강조한다.

일관된 문맥 정보에 조건화된 압축 토큰화를 사용하는 본 연구의 모델은 이를 개선하며 MaskViT를 유의미하게 능가한다.

정성적 평가: Open X-Embodiment, RoboNet, VP²에 대한 iVideoGPT의 비디오 예측 결과. 자세한 내용은 확대하여 보라. 확장된 예시는 Appendix B.1에서 찾을 수 있다.

 


4.2 Visual Planning

Setup.

 

VP는 네 개의 Robosuite 태스크와 일곱 개의 RoboDesk 태스크에 걸쳐 시각 모델-예측 제어(visual model-predictive control, MPC)를 위한 비디오 예측 모델을 평가하는 제어 중심(control-centric) 벤치마크이다.

각 환경의 학습 데이터셋에는 noisy scripted 상호작용 궤적이 포함된다.

원 벤치마크 논문의 프로토콜을 따라, Robosuite에 대해 5k개 궤적, RoboDesk에 대해 35k개 궤적으로 iVideoGPT를 학습하여, 모델을 확립된 baseline들과 비교했다.

 

Results.

VP² 벤치마크에 대한 시각 MPC 결과. 우리는 3번의 제어 실행에 대한 iVideoGPT의 평균과 최소/최대 성능을 보고한다. 오른쪽에서는, 낮은 시뮬레이터 성능으로 인해 flat block을 제외한 모든 태스크에 걸쳐 평균낸 평균 점수를 시뮬레이터의 성능으로 정규화하여 보인다.

 

Figure 5는 baseline 모델들과 비교한 iVideoGPT의 성공률을 제시한다.

Tian et al.은 뛰어난 지각 지표가 효과적인 제어 성능과 항상 상관관계가 있는 것은 아님을 관찰했지만, iVideoGPT는 두 개의 RoboDesk 태스크에서 큰 격차로 모든 baseline을 능가하고 가장 강력한 모델인 SVG′와 비슷한 평균 성능을 달성한다.

 

Appendix C.3에서, 저자들은 open slide 태스크에 대한 iVideoGPT의 최적이 아닌 성능을 분석하는데, 이는 모델의 이산화(discretization) 한계와 벤치마크의 불완전한 내장 보상 설계 양쪽 모두에 기인한다.


4.3 Visual Model-based Reinforcement Learning

Setup.

 

Meta-World에서 난이도가 다양한 여섯 개의 로봇 조작 태스크에 대해 실험을 수행한다.

iVideoGPT를 인터랙티브 월드 모델로 활용하여, MBPO에서 적응한 모델 기반 RL 방법을 개발했는데, 이는 표준 actor-critic RL 알고리즘을 학습하기 위해 replay buffer를 합성 rollout으로 증강한다(의사코드는 Appendix A.5 참조).

 

구현은 최신 시각 model-free RL 방법인 DrQ-v2 위에 구축된다.

또한 월드 모델 사전학습이 있는 경우와 없는 경우 모두에 대해, 최신 모델 기반 RL 알고리즘인 DreamerV3와 비교한다.

 

 

Results.

Meta-world에 대한 시각 모델 기반 RL. (왼쪽) 종합 결과는 여섯 개 태스크에 걸친 총 30번의 실행에 대한 사분위 평균(interquartile mean)과 95% 신뢰구간(CI)[2]을 보고한다. (오른쪽) 각 태스크에 대한 개별 결과로, 다섯 번의 실행에 걸친 평균과 95% CI를 보고하며, 20개의 평가 에피소드에 대한 성공률을 측정한다. PT는 사전학습(pre-training)을 나타낸다.

 

Figure 7은 본 연구의 모델 기반 알고리즘이 model-free 대응물 대비 샘플 효율을 현저히 개선할 뿐만 아니라 DreamerV3의 성능과

대등하거나 이를 능가함을 보여준다.

 

이는 MBPO를 시각 연속 제어(visual continuous control) 태스크에 성공적으로 적용한 최초의 사례를 보고하는 것이다.

이러한 결과는, 강력한 월드 모델이 있으면, 월드 모델 안에서 잠재 상태(latent states)의 rollout에 대해 정책을 학습하기 위해 진보된 MBRL 시스템에서 흔히 사용되는 전략인 latent imagination의 필요성을 제거할 기회를 부각한다. (Figure 6의 비교 참조).

강력한 iVideoGPT는 모델 rollout과 정책 학습을 분리하는, 단순하지만 성능 좋은 MBRL 알고리즘을 가능하게 한다.

성능 좋은 MBRL 알고리즘에 대한 본 연구의 개발은 모델 학습과 정책 학습을 분리하며, 여기서 iVideoGPT는 단순히 환경의 drop-in 대체물로 기능한다. 이는 설계 공간(design space)을 상당히 단순화하여, 실세계 응용에서 MBRL 알고리즘의 실용성과 효과를 크게 높일 수 있다.

 

Comparison to recurrent world models.

 

저자들은 순환 월드 모델이 실세계 데이터에 대한 대규모 사전학습 능력 — 현대 foundation 모델에 결정적인 능력 — 이 부족하다고 주장한다. 이를 검증하기 위해, DreamerV3 XL(200M 파라미터, iVideoGPT에 필적)을 동일한 데이터셋으로 사전학습한다.

Appendix의 Figure 11에서 보이듯이, DreamerV3는 자연스러운 로봇 동역학을 포착하는 데 실패하여, 저품질의 흐릿한 예측을 낸다. Figure 7의 Meta-World 벤치마크에 대한 추가 평가는 DreamerV3가 그러한 비효과적인 사전학습으로부터 이득을 얻지 못함을 드러낸다.


4.4 Model Analysis

 

Zero-shot Prediction.

 

먼저 미지의 BAIR 데이터셋에 대한 대규모 사전학습 iVideoGPT의 zero-shot 비디오 예측 능력을 분석한다.

iVideoGPT에서 사전학습 transformer에 의한 zero-shot 예측. fine-tuning 없이, transformer는 사전학습 토크나이저를 사용해 다른 로봇 그리퍼의 자연스러운 움직임을 예측하지만(두 번째 행), 적응된 토크나이저로는 올바른 그리퍼 유형에 대해 정확하게 예측한다(세 번째 행).

흥미로운 점은, Figure 8의 두 번째 행에서 iVideoGPT가 fine-tuning 없이 — 비록 본 연구의 사전학습 데이터셋과는 다른 것이지만 — 로봇 그리퍼의 자연스러운 움직임을 예측하는 것을 관찰했다는 것이다.

이는 사전학습 데이터의 불충분한 다양성으로 인해 모델이 완전히 미지의 로봇에 대한 zero-shot 일반화 능력은 제한적이지만, 장면 문맥(scene context)을 움직임 동역학(motion dynamics)으로부터 효과적으로 분리함을 나타낸다.

 

이와 대조적으로, 적응된 토크나이저를 사용하면, 그 자체로는 fine-tuning되지 않은 transformer가 사전학습 지식을 성공적으로 전이하여 세 번째 행에서 새로운 로봇 유형에 대한 움직임을 예측하며, 네 번째 행의 완전히 fine-tuning된 transformer와 유사한 지각 품질을 제공한다. 정량적 결과는 Figure 9a에서 찾을 수 있다.

 

Few-shot Adaption.

 

대규모 사전학습 모델은 특히 데이터가 부족한 시나리오에서 효과적임이 입증되었다.

모델 분석. (a) BAIR에 대한 다양한 fine-tuning 전략과 데이터 크기에 따른 비디오 예측 결과. (b) 사전학습 데이터셋에 대한 138M 및 436M transformer 모델의 validation loss. (c) 서로 다른 토크나이저의 계산 효율과 재구성 품질.

 

Figure 9a는 다양한 크기의 action-free BAIR 궤적으로 fine-tuning했을 때의 iVideoGPT 성능을 보여준다.

전체 downstream 데이터가 이용 가능할 때는 사전학습이 최소한의 이득만 제공하지만, 데이터 부족(100개 또는 1,000개 궤적) 하에서는 그 이점이 유의미해짐을 관찰한다.

또한 1,000개의 action-conditioned BAIR 궤적을 사용해 iVideoGPT를 적응시켜, FVD 82.3을 달성한다.

소수의 데이터로 빠르게 적응하는 능력은 모델 기반 RL에서 특히 결정적이다.

Meta-world에 대한 시각 모델 기반 RL. (왼쪽) 종합 결과는 여섯 개 태스크에 걸친 총 30번의 실행에 대한 사분위 평균(interquartile mean)과 95% 신뢰구간(CI)[2]을 보고한다. (오른쪽) 각 태스크에 대한 개별 결과로, 다섯 번의 실행에 걸친 평균과 95% CI를 보고하며, 20개의 평가 에피소드에 대한 성공률을 측정한다. PT는 사전학습(pre-training)을 나타낸다.

 

Figure 7에서 보이듯이, from scratch로 학습된 월드 모델은 부정확한 예측을 생성할 수 있으며, 이로써 모델 기반 에이전트에 필수적인 샘플 효율을 저하시킨다.

 

Model scaling.

 

이전의 모든 실험은 12개의 transformer layer와 768차원 hidden state(138M 파라미터)를 가진 iVideoGPT를 사용해 수행되었다.

모델의 스케일링 거동을 예비적으로 조사하기 위해, 24개 layer와 1024차원 hidden state(436M 파라미터)를 가진 더 큰 iVideoGPT를 학습했다.

Figure 9b는 사전학습 데이터셋에 대한 validation loss 곡선을 나타낸다.

이는 (1) validation loss(perplexity)가 모델 크기와 무관하게 계속 감소하며, (2) 모델 크기 증가가 loss 감소를 가속함을 보여준다.

이러한 결과는 더 큰 모델 크기와 증가된 계산이 더 강력한 iVideoGPT를 구축할 수 있다는 연구의 기대와 일치한다.

 

Tokenization efficiency.

 

각 프레임을 독립적으로 16×16 및 4×4 토큰으로 변환하는 표준 VQGAN 토크나이저와 비교하여 본 연구의 압축 토큰화의 효과를 평가한다. 세 개의 토크나이저를 동일한 스텝 수 동안 RoboNet에서 from scratch로 학습한다.

 

Figure 9c처럼, 4×4 토큰을 가진 토크나이저는 불충분한 용량으로 인해 낮은 재구성 품질을 겪는다.

본 연구에서 제안한 토큰화 방법은 표준 16×16 토크나이저에 비해 재구성 품질을 약간 희생하지만, 비디오 예측 태스크에 유익한

더 일관된 문맥 정보를 제공할 수 있다.

더 중요하게는, 이는 유의미하게 더 적은 양의 토큰으로 계산 효율을 크게 향상시켜, 시간과 메모리를 크게 절약하고, 더 적은 비용으로 모델 크기를 확장할 수 있게 한다 (정량적 결과는 Appendix B.5 참조).

 

Context-dynamics decoupling.

 

본 연구의 토크나이저는 훨씬 적은 토큰의 병목(bottleneck)으로 설계되어, 원시 픽셀을 재구성하기 위해

초기 프레임과 문맥 정보를 공유하면서 미래 프레임에 필요한 동역학 정보 포착에만 집중한다.

이 문맥과 동역학 정보의 분리를 명시적으로 시각화하기 위해, 본 연구에서는 미래 프레임을 재구성할 때 디코더에서

문맥 프레임으로의 cross-attention 블록을 제거한다.

우리 압축 토큰화에서의 문맥-동역학 분리. 미래 프레임에서 문맥 프레임으로의 cross-attention을 제거함으로써, 디코더는 여전히 원본과 같은 방식으로 움직이는 궤적을 재구성할 수 있으나, 시각적 문맥은 거의 완전히 사라진다.

 

Figure 10의 결과는 디코더가 여전히 움직임 궤적을 정확하게 재현할 수 있으나 최소한의 문맥 정보만 담고 있음을 보여준다. 이 시각화는 Figure 8에 나타난 모델의 일반화 능력에 대한 설명을 뒷받침한다.

 

Goal-conditioned Prediction.

 

Figure 4에서, 또한 대규모 인간 및 로봇 비디오로 사전학습된 goal-conditioned iVideoGPT가 생성한 비디오 예측을 보여준다(3.2절). ground truth에서 벗어나는 궤적을 자주 초래하는 action-free 예측과 달리, goal-conditioned 모델은 지정된 목표에 도달하는 더 정확한 경로를 생성한다. 저자들은 이것이 통합 시퀀스 모델링 패러다임의 유연성을 활용하는 잠재력을 부각한다고 믿는다.


 

5 Related Work 

 

World models for visual control.

 

시각 도메인에서 범용 월드 모델을 학습하는 것은 모델 기반 강화학습에서 여전히 중대한 과제이다.

 

직관적인 방법 하나는 액션 조건 비디오 예측 모델을 학습하는 것이다.

진보된 모델 기반 RL 알고리즘은 더 효율적이고 정확한 rollout을 위해 latent imagination을 활용하지만,

모델 학습과 정책 학습을 긴밀하게 결합함으로써 스스로를 복잡하게 만든다.

특정 태스크를 넘어서는 일반화 가능한 지식을 축적한 강력한 월드 모델로 이 복잡성을 줄일 수 있음을 보인다.

이를 용이하게 하기 위한 최근 노력에는 transformer 같은 확장 가능한 아키텍처를 활용하는 것 대규모 데이터로부터 사전학습하는 것이 포함된다.

본 연구와 특히 관련이 있는 것은 UniSimGenie인데,

이들은 각각 diffusion과 masked 모델로 광범위하게 학습된 월드 모델을 개발했으나 둘 다 공개되어 있지 않다.

본 연구는 범용 자기회귀 transformer 프레임워크를 활용함으로써 스스로를 차별화하며, 확장 가능한 월드 모델의 유연성을 진전시킨다.

 

Video generation and prediction.

 

인터넷 규모 비디오 생성 모델의 최근 발전은 이제 클래스 레이블, 텍스트 설명, 그리고 초기 프레임 — 마지막 것은 비디오 예측 문제로도 알려져 있다 — 에 조건화된 사실적 비디오의 합성을 가능하게 한다.

결정론적 RNN, variational autoencoder, diffusion, masked, autoregressive 모델을 포함한 다양한 모델이 개발되었다.

그러나 대부분의 최근 연구는 비디오 예측을 동역학 모델링 문제로 취급하지 않고 시공간 압축(spatiotemporal compression)을 수행하며, 이에 따라 월드 모델로 기능하기에는 제한적인 상호작용성만 제공한다.

 

본 연구는 최소한의 귀납적 편향(inductive bias)을 가진 cross-attention 메커니즘을 활용하는 문맥 인지(context-aware) 표현을 통해 압축 토큰화와 상호작용성을 모두 달성한다.

이 방법은 모션 벡터나 광류(optical flow)에 의존하는 이전 기법들에서 벗어나며, 더 범용적인 형태의 비디오 토큰화를 제공한다.


 

6 Discussion

 

본 연구는 확장 가능한 자기회귀 transformer를 활용해 멀티모달 신호를 토큰 시퀀스로 통합하는, 범용적이고 효율적인 월드 모델 아키텍처인 Interactive VideoGPT(iVideoGPT)를 소개했으며, 이는 next-token 예측을 통해 상호작용적 에이전트 경험을 제공한다.

iVideoGPT는 수백만 개의 인간 및 로봇 조작 궤적으로 사전학습되고 광범위한 downstream 태스크에 적응되었다.

월드 모델의 강력한 기반으로서, 이는 정확하고 일반화 가능한 비디오 예측뿐 아니라 단순하면서도 성능 좋은 모델 기반 계획 또는 강화학습을 가능하게 한다.

 

Limitations and future work.

 

iVideoGPT는 중대한 진전을 이루었지만, 개선의 여지는 상당하다.

대규모 Open X-Embodiment 데이터셋을 포함해 공개적으로 이용 가능한 로봇 데이터에서 제한적인 다양성을 발견했으며, 인간 비디오로부터 지식을 전이하려는 노력을 시작했다.

 

저자들은 인간과 로봇 사이의 지식을 잇기 위해 iVideoGPT가 더 광범위한 데이터로 사전학습되어야 한다고 믿는다.

 

이는 또한 iVideoGPT가 action-free 비디오 예측을 넘어, 다중 시점 관측(multi-view observations), 고유수용감각 로봇 상태(proprioceptive robot states), 그리고 액션 같은 더 많은 모달리티를 통합 정식화(unified formulation) 안에 포함할 것

요구한다.

 

구체적으로, 고차원 시각 관측을 처리하기 위해, 본 연구의 압축 토큰화는 초기 프레임이 미래 프레임에 대한 충분한 문맥을 제공한다고 가정하는데, 이는 모델 기반 에이전트가 흔히 수십 스텝 앞을 내다보는 저수준 제어 태스크에는 유효하지만, 긴 비디오와 상당한 카메라 움직임이 있는 시나리오에서는 흔들릴 수 있다.

이 문제는 키프레임 추출(keyframe extraction)로 완화될 수 있으나, 탐구할 중요한 향후 방향으로 남아 있다.

 

마지막으로, 더 복잡한 실기 로봇 태스크로 확장하는 것이 필수적인데, 이는 downstream 제어 태스크를 위한 시각적으로 단순한 시뮬레이션 안에서는 모델을 훨씬 더 큰 크기로 스케일링하는 이점이 이 연구에서 아직 관찰되지 않았기 때문이다.




 

Appendix

 

 

 

 

이상으로 본 논문의 리뷰를 마친다.

A Fine-Grained World Model for Robot Manipulation - HKUST · ByteDance Seed (ICCV 2025)

 

https://arxiv.org/abs/2406.14540

 

IRASim: A Fine-Grained World Model for Robot Manipulation

World models allow autonomous agents to plan and explore by predicting the visual outcomes of different actions. However, for robot manipulation, it is challenging to accurately model the fine-grained robot-object interaction within the visual space using

arxiv.org

Project Page: https://gen-irasim.github.io/

Authors: Fangqi Zhu¹², Hongtao Wu²†*, Song Guo¹*, Yuxiao Liu², Chilam Cheang², Tao Kong²

Affiliation: ¹Hong Kong University of Science and Technology, ²ByteDance Seed

Venue: arXiv Preprint (cs.RO) 

Preprint: arXiv:2406.14540v2 (초판 Jun 2024, v2 개정 Jul 2025)


0. Abstract

월드 모델은 자율 에이전트가 서로 다른 액션의 시각적 결과를 예측함으로써 계획하고 탐색할 수 있게 한다.

 

그러나 로봇 조작(manipulation)에서는, 시각 공간 안에서 미세한 로봇-물체 상호작용(fine-grained robot-object interaction)을 정확히 모델링하기가 어렵다. 기존 방법들은 각 액션과 그에 대응하는 프레임 사이의 "정밀한 정렬을 간과"하기 때문이다.

 

이 논문에서 IRASim을 제시한다.

이는 과거 관측과 로봇 액션 궤적에 조건화하여, 미세한 로봇-물체 상호작용 디테일을 담은 비디오를 생성할 수 있는 새로운 월드 모델이다.

 

디퓨전 트랜스포머를 학습하고, 각 트랜스포머 블록 안에 새로운 "프레임 단위 액션 조건화 모듈(frame-level action-conditioning module)"을 도입하여 액션-프레임 정렬을 명시적으로 모델링하고 강화한다.

 

 

광범위한 실험은 다음을 보인다.

 

(i) 이 방법이 생성한 비디오 품질이 모든 baseline을 능가하며, 모델 크기와 계산량이 늘수록 효과적으로 스케일된다.

 

(ii) IRASim으로 수행한 정책 평가정답 시뮬레이터(ground-truth simulator)로 수행한 평가와 강한 상관관계를 보여, 실제 세계 정책 평가를 가속할 잠재력을 부각한다.

 

(iii) IRASim을 이용한 모델 기반 계획(model-based planning)을 통한 테스트타임 스케일링(testing-time scaling)이 정책 성능을 크게 높인다. Push-T 벤치마크의 IoU 지표가 0.637 → 0.961로 향상된 것이 그 증거다.

 

(IV) IRASim은 유연한 액션 제어성을 제공하여, 데이터셋 안의 가상 로봇 팔을 키보드나 VR 컨트롤러로 조작할 수 있게 한다.


 

1. Introduction

월드 모델의 두 가지 응용

월드 모델은 세계의 근본적 동역학을 학습함으로써 에이전트가 자신의 행동 결과를 미리 내다볼 수 있게 한다.

 

이 능력은 로봇 조작에 두 가지 핵심 응용을 제공한다.

 

첫째, 정책 개선.

로봇이 모델 안에서 다양한 액션 제안(action proposal)을 탐색하고 그중 최적을 골라 실행함으로써 자율 정책을 개선할 수 있다.

 

둘째, 확장 가능한 정책 평가.

월드 모델은 사실적이고 합리적인 물리적 상호작용을 생성하여, 실제 세계 평가를 대체하는 효율적 방안을 제공한다.


왜 "미세한(fine-grained)" 월드 모델이 필요한가

로봇 조작용 월드 모델을 학습할 때, 로봇·물체·주변 환경 사이의 복잡한 상호작용을 정확히 시뮬레이션하는 것은 여전히 큰 난제다.

 

조작 과제는 본질적으로 섬세해서, 미세한 변화만으로도 과제가 실패할 수 있다.

 

따라서 이 정밀한 상호작용을 충실히 포착하려면 미세한(fine-grained) 월드 모델을 구축하는 것이 필수적이다.

 

또한 현대 로봇 조작 정책은 액션 청킹(action chunking) 기법을 활용하는데, 이는 복잡한 과제를 수행하기 위해 단일 액션이 아니라

액션 궤적(trajectory)을 생성한다.

(action chunking = 한 번에 한 스텝씩이 아니라 여러 스텝의 액션 묶음을 한꺼번에 내놓는 방식. 이것이 IRASim이 "궤적"을 조건으로 받는 이유다.)


trajectory-to-video 과제 정의

 

이 논문에서 과거 관측이 주어졌을 때 로봇이 액션 궤적을 실행하는 미세한 로봇-물체 상호작용 디테일의 비디오를 생성하는 월드 모델을 다룬다. 이 과제를 trajectory-to-video 과제라 부른다.

 

text-to-video 생성의 최근 발전에서 영감을 받아, 본 연구에서는 생성 모델을 활용해 세계의 복잡한 동역학(예: 로봇-물체 접촉, 물체의 관절 운동)을 표현하는 데 필수적인 시각 디테일을 포착한다.

 

그러나 text-to-video 모델은 상세한 프레임 단위 지시가 아니라 일반적 맥락 단서를 제공하는 "고수준 텍스트 서술"에 기반해

비디오를 생성하도록 학습된다.

 

이는 trajectory-to-video 과제와 다르다. 후자에서는 궤적의 각 액션이 예측 비디오의 각 프레임에서 로봇 움직임을 정확히 서술하기 때문이다.

(이 대비가 논문의 핵심 동기다. 텍스트 프롬프트는 "영상 전체를 대충" 서술하지만, 액션 궤적은 "프레임마다 정확히" 서술한다. 그래서 프레임 단위 조건화가 필요하다.)


IRASim 소개

이 간극을 메우기 위해, 본 연구에서는 복잡한 환경 동역학을 포착하도록 디퓨전 트랜스포머로 학습된 새로운 월드 모델 IRASim을 소개한다.

 

각 트랜스포머 블록 안에 새로운 "프레임 단위 액션 조건화 모듈"을 넣어, 각 액션과 대응 프레임 사이의 정렬을 명시적으로 모델링하고 강화한다.

 

IRASim은 미세한 로봇-물체 상호작용을 시뮬레이션하는 고품질 비디오를 생성한다.

과제 전체를 완수하는 긴 호라이즌(long-horizon) 비디오를 생성하려면, IRASim을 "자기회귀 방식으로 롤아웃"하여

각 클립 간 시간적 일관성을 유지할 수 있다.

(디퓨전 트랜스포머로 한 클립을 만들되, 긴 영상은 자기회귀로 이어붙인다 — 디퓨전 + 자기회귀의 하이브리드. 앞서 본 두 논문이 한쪽 계열만 택한 것과 대비된다.)


네 가지 과제와 데이터셋

네 과제로 방법을 검증한다.

1) 궤적 조건부 비디오 생성, 2) 정책 평가, 3) 모델 기반 계획, 4) 유연한 액션 제어성.

 

궤적 조건부 비디오 생성은 네 개의 실제 로봇 데이터셋(RT-1, Bridge, Language-Table, RoboNet)에서 검증한다.

IRASim은 고해상도(최대 288×512)와 긴 호라이즌(150프레임 이상) 비디오를 생성하며, 네 데이터셋 모두에서 baseline을 능가하고 인간 평가에서도 선호된다. 또한 모델 크기·계산량이 늘수록 효과적으로 스케일된다.

 

정책 평가는 IRASim과 LIBERO 시뮬레이션 환경 양쪽에서 수행한다.

IRASim의 평가 결과가 정답 시뮬레이터와 강하게 상관하여, 확장 가능한 정책 평가의 잠재력을 보인다.

 

또한 IRASim을 "시각 동역학 모델(visual dynamics model)"로 삼아 시뮬레이션과 실제 세계 양쪽에서 모델 기반 계획을 수행한다.

정책이 다양한 궤적 제안을 탐색하고 최적을 선택하게 함으로써 복잡한 과제 성능을 크게 높인다.

 

IRASim은 Push-T 벤치마크에서 바닐라 디퓨전 정책의 IoU를 0.637 → 0.961로 높인다.

더 중요한 건, 테스트타임 계산량이 늘수록 개선폭도 커진다는 것이며, 이는 로봇 조작을 위한 "테스트타임 스케일링(test-time scaling)"의 유망한 길을 보여준다.

 

마지막으로, 키보드나 VR 컨트롤러로 수집한 궤적으로 데이터셋 속 가상 로봇을 조작하는 비디오를 생성하여 유연한 액션 제어성을 보인다.


기여 요약 (threefold)

(1) trajectory-to-video 과제를 위해 미세한 상호작용 디테일의 고품질 비디오를 생성하는 IRASim을 제안한다.

프레임 단위 액션 조건화 모듈로 정밀한 액션-프레임 정렬을 달성한다.

 

(2) 궤적 조건부 비디오 생성에서 IRASim이 모든 baseline을 능가하며 모델 크기·계산량에 따라 효과적으로 스케일됨을 보인다.

 

(3) 정책 평가와 정책 개선 양면에서 IRASim의 유용성을 보인다.

IRASim 평가와 정답 시뮬레이터 평가가 강하게 상관하며, 모델 기반 계획과 결합하면 시뮬레이션과 실제 세계 모두에서 복잡한 과제 성능을 개선한다.


 

 

2.Related Work

World Models

현재 관측과 액션에 기반해 미래 관측을 예측하는 월드 모델(동역학 모델)을 학습하는 연구가 최근 점점 인기를 끌고 있다.

 

자율주행에서는 안전하고 강건한 주행을 위해 환경의 미래 상태를 추론하는 데 쓰였고, 게임에서는 안전하고 샘플 효율적인 RL 에이전트 학습에 활용되었다.

로봇 조작에서는 계획을 위해 액션 조건부 비디오 예측 모델을 학습해 왔다.

 

최근 iVideoGPT는 액션 조건부 예측을 위해 자기회귀 트랜스포머를 학습하는 방안을 제안했고, VLP·UniSim은 액션 정보를 담은 언어로 text-to-video 모델을 프롬프트해 비디오를 생성한다.

(iVideoGPT가 여기 등장한다. DWS가 자기회귀 베이스 모델로 쓴 바로 그 모델 — IRASim은 이걸 자기 방법의 선행 연구로 언급하고 있다.)

 

IRASim은 이들과 달리 초기 관측과 로봇 궤적이 주어졌을 때 고해상도(최대 288×512)·긴 호라이즌(150프레임 이상) 비디오를 생성하며, 미세한 로봇-물체 상호작용을 정확히 포착한다. 또한 모델 기반 계획을 통한 정책 개선과 확장 가능한 정책 평가에서 강한 능력을 보인다.


Video Models

비디오 모델은 프레임을 무조건적으로 생성하거나, 클래스·초기 프레임·텍스트·스트로크(strokes)·액션 등을 조건으로 생성한다.

 

최근 디퓨전 모델이 비디오 생성에서 점점 인기를 얻고 있다.

SoraDiffusion Transformer(DiT)로 놀라운 생성 능력을 보였다. IRASim도 DiT를 백본으로 활용한다.

(1X 리포트의 Wan, DWS의 Open-Sora, 그리고 IRASim 모두 DiT 계열 백본을 쓴다. Peebles & Xie의 DiT가 세 논문 모두의 공통 토대다.)

 

관련된 한 갈래는 모션으로 비디오 합성을 제어하는 연구다. 이들은 사용자가 지정한 "스트로크, 바운딩 박스, 인간 포즈"를 조건으로 쓴다.

반면 IRASim은 시간에 걸친 복잡한 2D·3D 액션프레임 단위 액션 조건화 모듈로 모델링한다.

(스트로크/박스/포즈는 대략적인 공간 힌트인 반면, IRASim은 정밀한 로봇 액션 벡터(2D·3D)를 다룬다는 차별점.)


Robot Learning with World Models

월드 모델은 로봇이 액션의 효과를 예측하고 복잡한 환경에서 해법을 계획할 가능성을 지닌다.

정책 학습을 위해 선행 연구들은 액션 조건부 비디오 예측을 "모델 예측 제어(MPC)"와 결합했다.

 

DreamerV3, DayDreamer는 "순환 상태 공간 모델(RSSM : Recurrent State Space Model)"을 활용해 RL용 월드 모델을

모델링함으로써 상태의 잠재 표현을 학습한다.

(RSSM 계열 — 1X와 DWS의 Related Work에도 나왔던 Dreamer 계보. 세 논문 모두 "픽셀/영상 공간, 그들은 잠재 상태 공간"이라는 대비를 공유한다.)

 

최근 FLIP흐름 예측(flow prediction)으로 보상을 최대화하는 비디오 계획을 생성한 뒤 역동역학(inverse dynamics)으로 액션을 만드는 방안을 제안했다. 본 연구의 모델 기반 계획과 다른 점은, 본 연구는 미래 비디오를 예측함으로써 액션의 보상을 예측하고 이를 통해 최적 액션을 선택한다는 것이다.

 

확장 가능한 정책 평가를 위해, 최근 연구는 물리 시뮬레이터에서의 평가와 실제 로봇에서의 평가 사이의 상관을 보였다.

물리 시뮬레이터를 쓰는 것과 대조적으로, 본 연구에서는 강력한 생성 모델로 정책의 롤아웃을 시뮬레이션하여 품질을 평가하고자 한다.

(핵심 차별점: 기존은 "물리 시뮬레이터 ↔ 실제"의 상관을 봤지만, IRASim은 물리 엔진 대신 생성 모델(영상)로 정책을 평가한다. 물리 규칙을 코딩하지 않고 데이터에서 배운다는 발상.)


 

 

3.Methods

3.1 Problem Statement

trajectory-to-video 과제를 과거 관측 이미지 I^{t−h:t}가 주어졌을 때, 궤적 a^{t:t+n}을 실행하는 로봇의 비디오를 예측하는 것으로 정의한다.

 

여기서 h는 과거 프레임 수, n은 비디오 내 액션 수, a^i ∈ ℝ^d는 i번째 시점의 액션이다.

이 논문은 로봇 팔 예측에 집중한다. 로봇 팔의 전형적 액션 공간은 7 자유도(DoF), 즉 a^i ∈ ℝ^7이다.

3 DoF는 3D 공간의 병진(translation), 3 DoF는 3D 회전, 1 DoF는 그리퍼(집게) 동작이다.

(액션이 25차원(1X)이나 게임 이산 액션(DWS)이 아니라 로봇 팔의 물리적 7DoF다. 이게 IRASim이 "fine-grained"를 말할 수 있는 근거 — 액션 자체가 정밀한 물리량이다.)


3.2 Preliminaries 

 

디퓨전 모델은 순방향 과정역방향 과정으로 구성된다.

순방향은 T스텝에 걸쳐 데이터 x_0에 가우시안 노이즈를 점진적으로 추가한다.

x_t는 t번째 diffusion timestep의 확산된 데이터, ᾱ_t 는 분산 스케줄로 정의되는 상수, ᾱ_t는 분산 스케줄로 정의되는 상수다.

 

역방향은 순수 노이즈 x_T ∼ 𝒩(0, I) 에서 시작해 노이즈를 점진적으로 제거하여 x_0를 복원한다.

 

이는 수학적으로 위와 같이 표현할 수 있으며, 여기서 μθ(·)와 Σθ(·)는 각각 평균 함수와 공분산 함수를 나타내고 신경망으로 파라미터화될 수 있다.

 

학습 단계에서는 timestep t를 [1, T] 구간에서 샘플링하고, reparameterization trick을 이용해 위 식으로 x_t를 얻는다. 여기서 ε_t ∈ 𝒩(0, I)이다.

 

DDPM처럼 노이즈 예측 모델 ε_θ를 학습한다. 손실은 예측 노이즈와 실제 노이즈의 L2 거리다.

 

추론 단계에서는 먼저 x_T를 𝒩(0, I)에서 샘플링한 뒤, t = 0이 될 때까지 반복적으로 계산하여 x₀를 생성한다.

 

조건부 디퓨전에서는 노이즈 예측 모델을 "ε_θ(x_t, t, c)"로 두며, c가 생성을 제어하는 조건이다.

논문 전반에서 위첨자(superscript)와 아래첨자(subscript)는 각각 입력 비디오 내 프레임의 timestep과 diffusion timestep을 나타낸다.

 

그러나 픽셀 공간에서 전체 비디오를 직접 확산하는 것은 시간이 오래 걸리고 상당한 계산을 요구한다.

Ma et al.에서 착안하여, 계산 효율을 위해 픽셀 공간 대신 저차원 latent 공간 z에서 확산 과정을 수행한다.

He et al.을 따라, "SDXL(Stable Diffusion XL)의 사전학습 VAE"를 활용해 각 프레임 I^t를 VAE 인코더로 latent 표현 z^t = Enc(I^t)로 압축한다.

이 latent 표현은 VAE 디코더로 다시 픽셀 공간으로 복원할 수 있다 (I^t = Dec(z^t)).


3.3 IRASim

 

IRASim은 VAE 잠재 공간에서 작동하는 조건부 디퓨전 모델이다.

조건 c = 과거 프레임의 잠재 표현 + 액션 궤적이고, 디퓨전 타깃 = 로봇이 궤적을 실행하는 이후 n프레임의 잠재 표현이다.

 

Sora에서 영감을 받아 DiT를 백본으로 채택한다.

설계 목표 세 가지: 1) 과거 프레임과의 일관성, 2) 주어진 액션 궤적의 준수, 3) 계산 효율.

 

표준 트랜스포머는 모든 토큰에 MHA를 적용해 이차 계산 비용이 든다.

IRASim은 "메모리 효율적 시공간 어텐션(spatial-temporal attention)"으로 이를 줄인다.

과거 프레임 조건화는 과거 프레임을 입력 시퀀스의 정답(ground-truth) 부분으로 취급해 달성한다.

즉 학습 중 예측 프레임에만 노이즈를 추가하고, 과거 프레임은 깨끗하게 유지하며(예측할 필요 없으니), 디퓨전 손실도 예측 프레임에만 계산한다.

(이것이 1X 리포트의 "마스킹으로 I2V→V2V" 트릭의 원본이다. 과거 프레임을 고정 슬롯으로 두고 미래만 생성하는 방식 — 세 논문이 공유하는 조건화 기법.)

 

이 방식은 예측 프레임이 어텐션으로 과거 프레임과 상호작용하게 하여 일관성을 보장한다.

 

궤적 조건을 주입하기 위해, DiT를 따라 "적응적 층 정규화(adaptive LayerNorm, adaLN)"를 사용한다.

• Video-Level Condition

 

text-to-video에서 텍스트 임베딩으로 영상 전체를 조건화하듯, 선형 레이어로 궤적을 하나의 임베딩으로 인코딩한다. 이 임베딩을 diffusion timestep t의 임베딩에 더해, 각 공간·시간 어텐션 블록의 scale 파라미터 γ, α와 shift 파라미터 β를 생성한다.

 

 

• Frame-Level Condition

 

텍스트가 영상 전체를 서술하는 text-to-video와 달리, trajectory-to-video의 궤적은 더 미세한 서술이다.

궤적의 각 액션이 각 프레임에서 로봇이 어떻게 움직일지를 정의하므로, 생성된 각 프레임은 대응 액션과 일치해야 한다.

 

정밀한 프레임 단위 정렬을 위해, 각 프레임을 그 프레임의 액션으로 조건화한다.

궤적을 하나로 인코딩하는 대신, 선형 레이어로 각 액션을 개별 임베딩으로 인코딩한다.

diffusion timestep 임베딩을 각 액션 임베딩에 더해, 공간 블록에서 각 프레임마다 개별 scale·shift 파라미터를 생성한다.

 

시간 블록의 scale·shift 파라미터는 모든 프레임이 동일한 조건 임베딩을 공유하며, 이는 video-level과 유사하게 도출된다.

 

• 출력

 

출력 레이어는 "노이즈 예측 ε̂ = ε_θ(x_t, t, c)"을 내는 선형 레이어를 포함한다.

학습 중 정답 노이즈와의 L2 손실을 계산한다.

VAE는 전 학습 과정 동안 동결(frozen)된다.

추론 시 순수 노이즈에서 시작해 점진적으로 디노이징하여 예측 프레임의 잠재 표현을 얻고, VAE 디코더로 디코딩한다.


 

 

4.Experiments

세 질문에 답한다.

1) 미세 상호작용 모델링에 효과적인가? 2) 정책 평가에 쓸 수 있는가? 3) 모델 기반 계획으로 정책을 개선할 수 있는가?

4.1 Trajectory-Conditioned Video Prediction

Set up 

 

데이터셋 4개: RT-1·Bridge(7-DoF) · Language-Table(2-DoF) · RoboNet(최대 5-DoF)
컨텍스트·예측 구성: RT-1·Bridge·Language-Table은 과거 1프레임 + 15액션 → 다음 15프레임 예측 · RoboNet은 iVideoGPT를

따라 과거 2프레임 + 10액션 → 다음 10프레임
해상도: RT-1·Bridge 256×320 · Language-Table 288×512 · RoboNet 256×256
짧은 궤적(short): n개 액션 — 비디오를 단일 생성 패스로 생성 · 긴 궤적(long): 더 많은 액션 — 자기회귀로 여러 패스에 걸쳐 생성하며, 이전 패스의 마지막 생성 프레임이 다음 패스의 조건 과거 프레임이 된다.

 

Baselines 

VDM·LVDM — 둘 다 U-Net 기반 확산(IRASim의 Transformer와 대조). 

LVDM은 latent 공간, VDM은 픽셀 공간에서 확산.

궤적 조건은 원 논문의 텍스트 임베딩처럼 궤적을 임베딩으로 인코딩해 확산을 조건화
추가로 RoboNet에서 비확산 SOTA 2종 비교 — iVideoGPT(다음 visual token을 자기회귀 예측) · MaskViT(반복 정제로 visual token 생성)

 

Metrics

두 유형: 계산 기반(PSNR·SSIM) + 모델 기반(Latent L2 loss·FID·FVD).

하나의 텍스트 조건을 여러 비디오가 만족할 수 있는 text-to-video와 달리, trajectory-to-video는 변이가 훨씬 적다.

예측 비디오의 로봇이 입력 궤적을 엄격히 따라야 한다.

따라서 재구성 지표인 Latent L2 loss와 PSNR을 주 지표로 사용(부록 H에서 이 둘이 모든 평가 지표 중 인간 선호와 가장 잘 정렬됨을 확인)

 

(이 지표 선택이 중요하다. 1X에서 봤던 것 — PSNR은 흐릿함을 선호한다 — 이 여기선 오히려 타당해진다. 정답이 하나로 거의 정해진 과제이기 때문. 다양성을 벌하는 게 결함이 아니라 적합한 성질이 되는 상황이다.)

Video Generation of Short Trajectories.

Fig. 3(a)와 Fig. 10에서 보이듯이, IRASim-Frame-Ada는 세밀한(fine-grained) robot-object 상호작용을 효과적으로 모델링하고 ground truth와 밀접하게 정렬되는 고품질 비디오를 생성한다.

이는 주 평가 지표인 Latent L2와 PSNR뿐 아니라 Sec. 4.1의 인간 평가에서도 비교 대상 baseline 방법들을 모두 능가한다.

Appendix A.1과 A(Fig.)11에 예시되어 있듯이, baseline 방법들은 주어진 궤적을 따라 로봇 팔을 유도하는 데 어려움을 겪으며, 로봇과 객체 간의 상호작용을 사실적으로 시뮬레이션하는 데 실패한다.

 

Video Generation of Long Trajectories.

IRASim을 최고의 baseline 방법인 LVDM과 비교한다.

IRASim-Frame-Ada는 세 데이터셋 모두에서 Latent L2 loss 기준으로 비교 방법들을 일관되게 능가한다. Fig. 3(b)와 Fig. 11은 IRASim이 짧은 궤적 설정에서와 마찬가지로 시각적으로 사실적이고 정확한 비디오를 생성하는 강력한 능력을 유지함을 보여준다.

 

Human Preference Evaluation.

또한 IRASim-Frame-Ada와 다른 방법들 사이의 인간 선호를 이해하는 데 도움이 되도록 사용자 연구(user study)를 수행한다. IRASim-Frame-Ada가 예측한 비디오와 비교 방법이 예측한 비디오를 나란히 병치하고, 사람들에게 어느 쪽을 선호하는지 묻는다. Ground-truth도 참조용으로 함께 제공된다.

IRASim-Frame-Ada는 세 데이터셋 모두에서 비교 대상 방법들을 모두 이긴다(Fig. 4).

이 결과는 Latent L2 loss와 일치하며, Latent L2 loss를 주 평가 지표 중 하나로 사용하는 이유를 정당화한다.

Scaling

서로 다른 모델 크기로 IRASim-Frame-Ada를 학습한다.

모델들의 상세 파라미터는 Appendix E에 나와 있다. 결과는 Fig. 5에 나와 있다.

세 테스트 데이터셋 전반에 걸쳐, IRASim은 더 큰 모델 크기와 더 많은 학습 스텝에 따라 효과적으로 스케일링되며, 이는 계산량 증가를 통한 추가 성능 향상의 강력한 잠재력을 부각한다.

 


4.2 Policy Evaluation & 4.3 Model-based Planning

 

LIBERO 벤치마크를 쓴다. 평가자는 성공과 실패 롤아웃을 모두 시뮬레이션할 수 있어야 한다.

그래서 학습된 정책을 시뮬레이터에 배치해 성공·실패가 섞인 post-trained 롤아웃을 모은다.

데이터가 적으니 IRASim을 OpenSora 사전학습 가중치로 초기화한다.

 

두 평가 결과의 Pearson 상관계수는 0.99로, 정답 Mujoco 시뮬레이터와 IRASim 평가가 강하게 상관한다.


4.3 Model-based Planning

 

간단한 랭킹 알고리즘을 쓴다.

1) 정책에서 K개 궤적 샘플링 → 2) 각 궤적을 IRASim에서 롤아웃 → 3) 가장 높은 값(value)의 궤적을 실행.

 

Push-T에서 IoU를 value 함수로 쓴다.

IoU 예측용 ResNet50을 post-trained 롤아웃으로 학습한다. "post-trained 롤아웃 수(P)"를 변화시키며 ablation한다.

 

공정 비교를 위해 본 연구의 디퓨전 정책의 IoU를 GPC 논문의 baseline과 맞춘다(0.637 vs 0.642).

K=50, P=1000일 때 IRASim은 바닐라 정책 IoU를 0.637 → 0.961로 높인다.

 

P>0일 때 K가 커질수록 성능이 일관되게 향상된다. 즉 랭킹용 샘플 궤적을 늘려 test-time scaling을 할 수 있다.

P가 커지면 큰 K에서 성능이 계속 오른다. 데이터 크기(P)와 테스트타임 계산(K)이 동시에 스케일돼야 함을 시사한다.

 

이제 Push-T 표 — Abstract부터 세 번 나온 그 숫자를 뜯어볼게.

 

Real-Robot Experiments

 

실제 로봇 데이터셋으로 IRASim을 학습하고 세 과제(서랍 닫기, 귤을 초록/빨강 접시에 놓기)를 실험한다.

목표 이미지(goal image)로 과제를 지정하고, 현재 end-effector 위치 중심의 구(sphere)에서 50개 점을 샘플링해

K=50 궤적을 만든다.

 

유사도 함수 두 개(MSE, ResNet50 특징의 코사인 유사도)를 실험한다.

MSE가 ResNet보다 훨씬 우수했고, 두 변형 모두 계획 없이 무작위 선택하는 정책을 크게 능가한다

(Table 6: Random 0.20/0.07/0.13 → IRASim-MSE 0.87/0.80/0.87).


4.4 Flexible Action Controllability

"키보드(Language-Table의 2D)와 VR 컨트롤러(RT-1의 3D)"로 가상 로봇을 조작한다.

특히 이렇게 수집한 궤적은 원본 데이터셋의 분포에서 벗어난다.

 

IRASim은 서로 다른 입력 소스의 궤적을 따르고 사실적으로 상호작용을 시뮬레이션한다.

무엇보다 생성의 다중모드성(multimodality)을 강건하게 처리한다.

부록 A.4에서는 물리적으로 불가능한 궤적도 강건하게 처리함을 보인다.

(같은 초기 프레임 + 다른 궤적 → 다른 영상. 이게 월드 모델이 "궤적을 진짜 따른다"는 강한 증거다. 궤적을 무시하고 데이터셋 평균만 재생하는 게 아니라는 것.)


 

 

5.Conclusion

 

본 연구에서는 과거 관측이 주어졌을 때 로봇이 액션 궤적을 실행하는 미세한 상호작용 디테일의 비디오를 생성하는 월드 모델 IRASim을 제시했다. 프레임 단위 액션 조건화 모듈로 정밀한 액션-프레임 정렬을 달성한다.

 

실험은 IRASim이 로봇 궤적 롤아웃을 정확히 시뮬레이션하는 긴 호라이즌·고해상도 비디오를 생성함을 보인다.

또한 IRASim을 "정책 평가용 시뮬레이터"이자 "모델 기반 계획용 동역학 모델"로 활용할 수 있음을 보인다.

 

한계: 다른 생성 모델처럼 IRASim도 비디오 생성이 실시간이 아니다.

향후 diffusion distillation으로 생성 속도를 높이고, IRASim을 동역학 모델로 삼아 월드 모델 안에서 강화학습으로 정책을 개선하는 방향을 탐색할 계획이다.

(Limitation이 딱 하나 — "실시간이 아니다". Push-T에서 K=50 궤적을 매번 롤아웃하는 비용을 스스로 인정한 것.

그리고 "미래엔 RL로 정책 개선" = 이게 바로 DWS가 실제로 한 일. IRASim이 "future work"라 남긴 걸 DWS가 이어받은 셈이다.)

 

 

이상으로 본 논문의 리뷰를 마친다.

https://arxiv.org/abs/2502.07825

 

Pre-Trained Video Generative Models as World Simulators

Video generative models pre-trained on large-scale internet datasets have achieved remarkable success, excelling at producing realistic synthetic videos. However, they often generate clips based on static prompts (e.g., text or images), limiting their abil

arxiv.org

 

Authors: Haoran He¹, Yang Zhang², Liang Lin³, Zhongwen Xu⁴, Ling Pan¹†

Affiliation: ¹Hong Kong University of Science and Technology, ²Tsinghua University, ³Sun Yat-sen University, ⁴Tencent AI Lab

Venue: AAAI 2026 (40th Annual AAAI Conference on Artificial Intelligence)

Preprint: arXiv:2502.07825 (Feb 2025)


0.Abstract

 

대규모 인터넷 데이터셋으로 사전학습된 비디오 생성 모델들은 사실적인 합성 비디오를 만들어내는 데 뛰어난 성과를 거두었다.

그러나 이 모델들은 텍스트나 이미지 같은 정적 프롬프트에 기반해 클립을 생성하는 경우가 많아, 인터랙티브하고 동적인 시나리오를 모델링하는 능력에 한계가 있다.

 

이 논문에서 "Dynamic World Simulation(DWS)"을 제안한다.

사전학습된 비디오 생성 모델을, 지정된 액션 궤적을 실행할 수 있는 제어 가능한 월드 시뮬레이터로 전환하는 새로운 접근법이다.

 

조건으로 주어진 액션과 생성된 시각적 변화 사이의 정밀한 정렬을 달성하기 위해, 어떤 기존 모델에도 매끄럽게 통합되는 경량·범용 액션

조건화 모듈을 도입한다.

복잡한 시각적 디테일에 집중하는 대신, 일관된 동역학 전이 모델링(consistent dynamic transition modeling)이

강력한 월드 시뮬레이터를 만드는 열쇠임을 보인다.

이 통찰을 바탕으로, 모델이 동적 변화를 더 효과적으로 포착하도록 강제함으로써 액션 제어성을 높이는 motion-reinforced loss

추가로 도입한다.

 

실험 결과 DWS는 Diffusion, Autoregressive Transformer Model 양쪽 모두에 다목적으로 적용될 수 있으며, 게임과 로보틱스 도메인에서 액션 제어가 가능하고 동역학적으로 일관된 비디오를 생성하는 데 유의미한 개선을 달성한다.

 

나아가 학습된 월드 시뮬레이터를 모델 기반 강화학습(MBRL) 같은 다운스트림 과제에 활용하기 쉽도록, 샘플 효율을 개선하는 prioritized imagination을 제안하며 최신 기법들과 견줄 만한 성능을 입증한다.


 

1. Introduction

 

비디오 생성 분야는 최근 몇 년간 눈에 띄는 발전을 이뤘다.

Brooks et al.(2024), Polyak et al.(2024), Sharma et al.(2024), Yang et al.(2024c), Zheng et al.(2024) 같은 모델들은 다양한 입력, 특히 텍스트와 초기 프레임에 조건화하여 고품질이면서 시간적으로 일관된 비디오를 생성하는 탁월한 능력을 보여주었다.

 

그러나 이 모델들은 정적 프롬프트로 한 번에(one-shot) 생성하도록 학습되었기 때문에 인터랙티브 시뮬레이션 시나리오를 지원하는 데 한계가 있으며, 프레임 단위의 상호작용성프레임 간 동역학 모델링이 결여되어 있다.

 

이 간극을 메우기 위해, 학계는 액션 조건부(action-conditioned) 비디오 모델을 구축하는 데 점점 더 집중하고 있다.

 

이러한 액션 조건부 모델들은 사실상 인터랙티브 환경 시뮬레이터("월드 모델" 또는 "월드 시뮬레이터")로 기능하며, 발전된 트랜스포머나 디퓨전 모델 아키텍처를 활용해 에이전트의 행동에 기반한 미래의 시각적 결과를 예측한다.

 

이들의 목표는 근본적인 동역학 전이(dynamic transitions)에 대한 이해세계가 작동하는 방식에 관한 상식적 지식을 담아내어, 인간의 인지 과정과 유사한 "행동 기반 상상(action-driven imagination)"을 가능하게 하는 것이다.

 

이러한 월드 모델은 특히 "모델 기반 강화학습(MBRL)"에서 흥미로운 가능성을 연다.

에이전트가 월드 모델과 상호작용함으로써 새로운 기술을 더 효율적으로 학습할 수 있고, 실제 세계에서의 시도에 따르는 위험과 비용을

피할 수 있기 때문이다.

 

이 연구에서 인터랙티브 월드 시뮬레이터의 최근 발전을 검토하며, 현재 이들의 폭넓은 채택을 가로막는 핵심 난제들을 짚는다.

 

(i) 이 모델들은 "처음부터 학습(training from scratch)"하는 데 막대한 계산 자원을 필요로 하는 경우가 많다.

 

예를 들어 Genie(Bruce et al., 2024)는 비교적 단순한 Platformer game simulator를 학습하는 데

256개의 TPUv5p 코어(NVIDIA A100 GPU 약 226장에 상당)에서 12.5만 학습 스텝을 필요로 했다.

마찬가지로 GameNGen(Valevski et al., 2024)은 128개의 TPU-v5e 코어70만 스텝을 소모했다.

 

(ii) 사전학습된 비디오 생성 모델을 파인튜닝하는 것이 보다 효율적인 대안이 되긴 하지만, 기존 접근법들(Rigter et al., 2024; Yu et al., 2025)은 "본질적으로 아키텍처 의존적(architecture-dependent)"이다.

 

이는 비디오 생성 아키텍처의 빠른 발전을 활용하기 위해 이 방법들을 서로 다른 모델 아키텍처에 적용하는 데 걸림돌이 된다.

새로운 아키텍처가 나올 때마다 적응을 위한 상당한 엔지니어링 노력이 필요하기 때문이다.

 

(iii) 일반적인 비디오 생성 과제와 달리, 액션 조건부 월드 시뮬레이터는 미세한 동적 변화(fine-grained dynamic changes)를 정밀하게 포착해야 하며, 이를 위해서는 "프레임 단위의 액션 정렬(frame-level action alignment)"이 요구된다.

 

이 요구사항은 모델 기반 강화학습 응용에서 결정적이다. MBRL에서는 프레임 간 동역학·모션 변화를 포착하는 일이 정적인 시각 요소(예: 배경, 객체 디테일)를 모델링하는 것보다 우선하기 때문이다.

 

앞서 언급한 난제들을 해결하기 위해, "Dynamic World Simulation(DWS)"이라는 새로운 프레임워크를 제안한다.

이는 사전학습된 비디오 생성 모델을 월드 시뮬레이터로 효율적으로 전환하는 통합적이고 아키텍처 불가지론적(architecture-agnostic)인 접근법이다.

 

인터넷 규모 데이터셋에서 학습된 "사전학습 사전지식(pre-trained priors)"을 활용함으로써, 파인튜닝된 월드 시뮬레이터는 물리 법칙과 상식적 지식에 대한 기본적인 이해를 보일 수 있다.

그러나 이들은 본질적으로 인터랙티브 시뮬레이션에 맞게 설계되어 있지 않으며, 정밀한 프레임 단위 액션 조건화를 위한 핵심 메커니즘이 없다.

 

DWS는 아키텍처 유연성을 유지하면서 프레임 단위 액션 인식을 개선하는, 최소한이면서도 강력한 애드온(add-on) 액션 조건화 모듈을 도입한다.

이 모듈은 "단 두 개의 선형 레이어(linear layer)"로만 구성되며, 세심하게 설계된 scale과 shift 연산을 통해 어떤 네트워크 아키텍처에도 통합될 수 있다.

또한 예측된 시각적 변화와 조건으로 주어진 액션 사이의 정렬을 강화한다.

 

나아가 본 연구에서는 전통적인 "지도학습 손실 함수"가 비디오 모델로 하여금
"모든 시각 요소와 복잡한 디테일에 균일하게 집중하게 만든다"
는 점을 관찰했다.

여기에는 정적인 배경과 무관한 디테일까지 포함되며, 이는 월드 시뮬레이터 구축에 결정적인 프레임 간 동역학·모션 변화를 포착하는 능력을 훼손한다.

 

DWS는 이 동역학 모델링 난제를 해결하기 위해 "motion-reinforced loss"를 제시한다.

학습 중 프레임 간 변화의 모델링을 명시적으로 우선시하는 단순하지만 효과적인 방법으로, 그 결과 시간적 일관성이 크게 향상되고 더 신뢰할 수 있는 동역학 예측이 가능해진다.

파인튜닝된 비디오 생성 모델은  서로 다른 정책과 상호작용하며 다양한 궤적을 생성하는 효과적인 월드 시뮬레이터로 기능할 수 있다. 서로 다른 정책은 서로 다른 종료 상태(terminal state)로 이어지며, 이는 노란 원으로 표시되어 있다.

 

Fig. 1에서 보듯, DWS로 학습된 월드 시뮬레이터는 정확한 동역학 예측과 액션 반응성을 유지하면서 여러 도메인에 걸쳐 다양한 정책(policy)과 상호작용할 수 있다.

 

마지막으로, 모델 기반 강화학습에서 월드 시뮬레이터의 실용적 효용을 높이기 위해 "prioritized imagination"을 도입한다.

이는 이미 잘 이해된 상태 전이에 계산 자원을 낭비하는 대신 가장 정보량이 많은(informative) 전이에 집중하는 새로운 샘플링 전략으로, 에이전트-월드 모델 상호작용에서 샘플 효율을 개선한다.

 

 

이 논문의 기여를 다음과 같이 요약한다.

 

(i) "DWS"를 도입한다. 이는 사전학습된 사전지식을 "물리적 근거(physics grounding)"로 활용하여, 사전학습된 비디오 생성 모델을 낮은 학습 비용으로 월드 시뮬레이터로 효과적으로 전환하는 새롭고 아키텍처 불가지론적인 프레임워크다.

 

(ii) 단순하지만 효과적인 두 가지 기법을 도입한다.

하나는 정밀한 프레임 단위 제어를 가능하게 하고 액션 추종(action-following) 능력을 개선하는 "경량 액션 조건화 모듈"이고, 다른 하나는 모델의 주의를 정적인 시각 디테일에서 액션이 유발한 동적 변화로 돌려 시간적·동역학적 일관성을 개선하는 "motion-reinforced 학습"이다.

 

(iii) "prioritized imagination"을 통해 모델 기반 강화학습에서 월드 시뮬레이터의 실용적 효용을 끌어올린다.

이는 학습된 월드 시뮬레이터를 다운스트림 MBRL에 적용할 때 샘플 효율과 정책 성능을 개선한다.

 

(iv) 까다로운 게임 및 로보틱스 과제 전반에 걸친 포괄적 평가를 통해, DWS가 생성된 액션 조건부 비디오의 품질과 동역학적 일관성을

크게 개선함을 입증한다.

 

또한 prioritized imagination을 적용한 DWS 학습 월드 시뮬레이터는 기존 최신 MBRL 접근법들(Alonso et al., 2024 — DIAMOND; Hafner et al., 2023 — DreamerV3)과 비교해 더 효율적이고 효과적인 학습을 가능하게 한다.


 

 

2. Related Work

 

2.1 Video World Models

인터넷 규모 데이터셋(Bain et al., 2021; Chen et al., 2024)과 발전된 모델 아키텍처(Brooks et al., 2024; Peebles and Xie, 2023 — DiT)의 발전에 힘입어, 텍스트 설명과 초기 프레임에 조건화한 사실적 비디오 생성에서 상당한 진전이 이루어졌다.

이러한 토대 위에서, 현재 연구는 액션 제어 가능(action-controllable) 비디오 생성에 점점 더 집중하고 있으며, 물리적 동역학과 행동의 결과를 모두 효과적으로 모델링하는 범용(generalist) 월드 시뮬레이터를 개발하는 것을 목표로 한다.

그러나 이 모델들은 대체로 대규모 데이터셋에서 처음부터 학습해야 하고 수백만(또는 수십억) 개의 파라미터를 포함하므로, 상당한 계산 오버헤드와 느린 추론 속도를 초래한다.

 

이와 대조적으로, 공개된 사전학습 비디오 생성 모델행동 기반(action-driven) 월드 시뮬레이터로 적응시킬 것을 제안한다.

본 연구에서 제안하는 파인튜닝 접근법 DWS최소한의 계산 오버헤드로 효율적인 모델 적응을 달성하며, 학습 비용과 추론 지연(inference latency)을 모두 크게 줄인다.

 

동시기 연구들(Rigter et al., 2024; Yu et al., 2025)도 액션 조건부 비디오 생성을 위해 사전학습 모델을 활용하는 방안을 탐색했다.

그러나 이들의 연구는 디퓨전 기반 모델에 국한되어 있으며, 두 연구 모두 모델 기반 RL 같은 다운스트림 과제를 촉진하는 데 있어

"자기 접근법의 효과"를 검증하지 않았다.


2.2 Model-Based RL

모델 기반 RL은 실제 비용 없이 시행착오가 이루어질 수 있는 월드 모델을 구축하는 것을 목표로 한다.

충분히 정확한 월드 모델이 있으면, 에이전트는 상상(imagination) 능력을 발전시켜 상호작용을 시뮬레이션하고

"합성 경험 데이터(synthetic experience data)"를 생성할 수 있다.

 

이렇게 시뮬레이션된 데이터는 다양한 의사결정 과제에서 최적 정책을 학습하는 데 활용될 수 있으며, 실제 세계와의 상호작용 필요성을

효과적으로 줄인다.

 

Sutton(1991)은 모델 기반 RL의 첫 일반적 프레임워크를 도입하며, 추정된 동역학 모델가치 함수(value function)정책의 학습을 촉진하는 데 유용함을 강조했다(Sutton and Barto, 1998).

 

최근 몇 년간 비디오 게임이나 시각 제어 과제처럼 복잡한 환경 동역학을 학습하는 모델 기반 방법들이 눈에 띄게 발전했으며, 모델 프리(model-free) 방법들을 일관되게 능가해 왔다.

 

예를 들어, 환경 동역학의 결정론적 요소와 확률적 요소를 명시적으로 분리하는 순환 상태 공간 모델(RSSM)(Hafner et al., 2019)에 기반한 Dreamer 시리즈는 Atari 게임, DeepMind Control Suite, Minecraft 등 다양한 도메인에서 인상적인 성능을 보였다.

 

복잡한 패턴을 표현하는 데 있어 RNN의 한계를 해결하기 위해, 최근 연구들은 "향상된 시퀀스 모델링과 장기 의존성 포착"을 위해

"트랜스포머 모델을 활용"하거나(Micheli et al., 2023a; Robine et al., 2023; Zhang et al., 2023, 2024),

동역학 학습에서 "다중 모드(multi-modal) 분포를 더 잘 표현"하기 위해 "디퓨전 모델을 도입"하는 방안(Alonso et al., 2024 — DIAMOND; Ding et al., 2024)을 탐색해 왔다.

 

그러나 이 연구들 역시 월드 모델 학습에 트랜스포머나 디퓨전 모델을 사용하지만, 대부분 처음부터 학습에 의존하며 동역학 이해를 강화하기 위한 사전학습 지식을 활용하지 못한다. 이는 이들을 지나치게 과제 특화적(task-specific)으로 만들어, 다양한 과제에 걸쳐 일반화하는 능력을 제한한다.

 

나아가, 기존 방법들이 학습 중 상상된 모든 샘플을 균일한 중요도로 취급하는 반면, 본 연구에서 제안하는 DWS는 "중요한 샘플에 선택적으로 집중하는 새로운 우선순위화(prioritization) 메커니즘"을 도입하여 샘플 효율을 개선한다.


 

3. Preliminaries

 

3.1 Problem Formulation

조건부 비디오 생성 프레임워크는 월드 시뮬레이터(또는 월드 모델)를 구현하는 데 적용될 수 있다(Yang et al., 2023).

 

월드 모델은 어떤 액션을 입력으로 받아 그 액션의 시각적 결과를 출력으로 생성하며, 이를 통해 환경을 시뮬레이션하는 것을 목표로 한다.

 

이 환경은 "부분 관측 마르코프 결정 과정(POMDP : Partially Observable Markov Decision Process)"으로 표현될 수 있으며,

튜플 (𝒮, 𝒪, φ, 𝒜, p, r, γ)로 정리된다.

여기서 "𝒮는 상태 공간(state space)"이고, "𝒪는 관측 공간(observation space)"으로 𝒮에 대한 불완전한 정보만 제공한다.

 

각 시점 t에서 에이전트는 "정책 π : 𝒪 → Δ(𝒜)"를 따라 액션 aₜ를 선택하고, 환경은 "동역학 sₜ₊₁ ∼ p(sₜ₊₁ | sₜ, aₜ)"에 따라 상태를 갱신하며, "다음 관측 oₜ₊₁ = φ(sₜ₊₁)"을 받고, "스칼라 보상 rₜ = R(sₜ, aₜ, sₜ₊₁)"을 계산한다.

(주: Δ(𝒜)는 액션 공간 위의 확률분포로, 정책이 결정론적 액션이 아니라 액션에 대한 분포를 출력함을 뜻한다.)

 

에이전트의 목표는 γ로 할인된 누적 보상을 최대화함으로써 최적 정책을 학습하는 것이다.

 

(γ는 할인율(discount factor). 먼 미래의 보상일수록 γᵗ로 가중치를 줄여, 가까운 보상을 더 중시한다.)

 

잘 학습된 월드 모델은 환경을 대체하여 에이전트와 상호작용할 수 있으며, 따라서 무한한 경험을 제공함으로써 다운스트림 정책 학습에

도움을 준다.

 

구체적으로, "과거 관측 o(T₀)"가 주어졌을 때, 각 시점 t = T₀, ⋯, T−1에서 에이전트는 자신의 정책과 이전에 상상된(imagined) 관측에 기반해 "액션 aₜ"를 취하고, 월드 모델은 "전이 p(oₜ₊₁, rₜ₊₁ | oₜ, aₜ)"를 예측하여 에이전트에게 피드백한다.

(핵심은 "imagined observations"다. 에이전트가 실제 환경이 아니라 월드 모델이 상상해낸 관측을 다시 입력으로 받아 다음 행동을 정한다. 이것이 MBRL의 상상 롤아웃(imagination rollout)이다.)


3.2 Pre-Trained Video Generative Models

시각 제어(visual control)를 위한 월드 모델 학습을 인터랙티브 비디오 생성 문제로 정식화함으로써, 널리 존재하는 비디오 데이터를 활용할 수 있다. 이 데이터에는 여러 도메인에 걸쳐 일반화 가능한 폭넓은 지식이 담겨 있다.

 

비디오 데이터는 의미적인 시각 디테일뿐 아니라, 물리 세계의 동역학 법칙을 포착하는 모션 움직임까지 담고 있다.

 

그러나 이러한 비디오 월드 모델을 인터넷 규모 데이터셋에서 처음부터 학습하는 것은 비싸고 시간이 오래 걸린다.

 

본 연구에서는 사전학습된 발전된 비디오 모델을 파인튜닝하여 상호작용을 시뮬레이션할 수 있게 만들 것을 제안한다.

 

구체적으로, 서로 다른 두 개의 사전학습 비디오 생성 모델을 베이스 모델로 채택한다.

하나는 디퓨전 모델인 "Open-Sora(Zheng et al., 2024)"이고, 다른 하나는 자기회귀 트랜스포머 모델인 "iVideoGPT(Wu et al., 2024)"이다.

 

Open-Sora는 rectified flow 기반 디퓨전 모델로, 완전히 오픈소스이며 수백만 개의 인터넷 비디오로 사전학습되었다.

이를 채택한 이유는 flow-matching 학습 덕분에 샘플링 스텝이 적게 필요하기 때문이다.

 

iVideoGPT는 LLaMA 아키텍처 위에 구축된 자기회귀 트랜스포머다.

이는 서로 다른 모달리티의 학습 데이터(예: 시각적 결과, 액션, 보상)를 하나의 토큰 시퀀스로 압축하여 인터랙티브 비디오 예측을 수행한다.


 

4. Method

 

이 절에서는 먼저 §4.1에서 제안하는 "액션 조건화 모듈(action-conditioned module)"을, §4.2에서 동역학 모델링을 강화하는 "motion-reinforced loss"를 소개한다. 일반 비디오 생성 모델을 월드 시뮬레이터로 파인튜닝하는 방법을 제시한 뒤, §4.3에서 모델 기반 강화학습 성능을 개선하는 "prioritized imagination" 기법을 소개한다.

 

4.1 Action-Conditioned Module

4.1 Action-Conditioned Module

최근 비디오 생성 모델들은 조건으로 주어진 텍스트 프롬프트나 초기 프레임과 연관된 사실적 비디오를 생성하는 데 상당한 성공을 거두었다. 그러나 이러한 text-to-video 과제는 다음 프레임이 무엇이어야 하는지 명시하지 않고, 전체 비디오를 전역적으로 서술하는 정적 프롬프트로 작동한다.

 

이 설계 패러다임은 일반적인 비디오 생성에는 효과적이지만, 조건이 프레임 단위의 미세한 액션 궤적액션이 풍부한 상호작용을 시뮬레이션하는 월드 모델로 적응시킬 때는 근본적인 어려움을 낳는다.

 

이 요구를 해결하고 생성된 각 프레임이 궤적 상의 대응하는 액션과 일치하도록, 각 프레임의 생성을 그에 대응하는 액션으로 개별적으로 조건화하는 액션 조건화 모듈을 활용한다.

 

전체 액션 궤적을 하나의 임베딩으로 압축하는 기존 text-to-video 모델과 달리, 접근법은 "IRASim(Zhu et al., 2024)"과 유사하게 더 세분화된(granular) 액션 인코딩 메커니즘을 구현한다.

 

각 트랜스포머 블록 안에 두 개의 선형 레이어로 구성된 경량 애드온 모듈을 도입하여, 개별 액션을 각각 따로 인코딩한다.

 

이 설계는 각 프레임의 내용이 전역적 서술이 아니라 대응하는 액션에 의해 직접 변조되도록 보장하며, 액션과 생성된 프레임 사이의 직접적 대응관계로 이어진다.


Action Representation (액션 표현)

비디오 생성 모델을 액션 기반 제어에 적응시킬 때의 "핵심 난제는 액션의 표현"에 있다.

이산화된 액션 공간에서 액션은 보통 정수값으로 표현되는데, 이는 기존 text-to-video 모델의

텍스트 프롬프트가 가진 "풍부한 의미적 맥락이 결여"되어 있다.

 

의미적 간극(semantic gap)은 모델이 서로 다른 액션을 효과적으로 해석하고 반응하는 능력을 제한할 수 있다.

이를 메우기 위해, 본 연구는 액션의 의미를 서술하는 "언어 템플릿(language templates)"을 사용해 액션을 표현할 것을 제안한다.

 

구체적으로, 액션 궤적 y = {a_t, a_{t+1}, ⋯, a_{t+H−1}}(H는 궤적의 호라이즌)이 주어지면, 추상적인 액션 정수를 의미 있는 언어로 번역하는 "매핑 함수 ψ"를 만든다. 즉 ψ : 𝒜 → L(L은 언어 공간)이다.

 

이 매핑을 통해 사전학습 모델의 텍스트 인코더를 활용하여 풍부한 특징 임베딩 c ∈ ℝ^{n_H × n_d} 를 얻는다.

(여기서 n_H는 호라이즌, n_d는 각 토큰의 차원이다.)

 

연속 액션 공간의 경우, 언어 번역 없이 "학습 가능한 선형 액션 임베더(linear action embedder)"를 사용해

특징 임베딩 c를 직접 생성한다.

(핵심: 이산 액션은 정수 → 언어로 번역해 사전학습 텍스트 인코더를 재활용하고, 연속 액션은 선형 임베더로 직접 처리한다.

사전학습 모델이 이미 이해하는 "언어"라는 인터페이스로 액션을 밀어넣는다.)


Frame-Level Condition (프레임 단위 조건화)

비디오 생성 모델이 월드 시뮬레이터로 기능할 때, 각 액션이 바로 다음 프레임의 시각적 내용을 직접 변조해야 하므로 정밀한 시간적 제어가 중요하다.

 

액션-프레임 대응관계를 명시적으로 모델링하고 강화하기 위해, IRASim에서 영감을 받아 "각 트랜스포머 블록 안에 프레임 단위 액션 조건화 모듈을 삽입"한다.

 

IRASim의 구현은 시공간 트랜스포머 구조를 가진 특정 디퓨전 모델에 국한되었으나, 저자들은 이 개념을 크게 확장하여

디퓨전 기반과 트랜스포머 기반을 모두 아우르며 서로 다른 아키텍처에 일반화되는 다목적 애드온 모듈을 개발했다.

 

따라서 본 연구의 설계는 향상된 아키텍처 유연성과 더 넓은 적용성을 제공한다. 단 두 개의 선형 레이어로 구현된 미니멀한 구조는 경량 통합과 최소 오버헤드의 효율적 파인튜닝을 가능하게 한다.

 

구체적으로, 각 비디오 임베딩 x ∈ ℝ^{T×C×H×W} 에 대해, 트랜스포머 블록에 넣기 전 다음과 같이 처리한다 :

 

 

여기서 α^i와 β^i는 i번째 프레임에 대한 scale과 shift 파라미터이며, "액션 임베딩 c^i로부터 회귀(regress)"되어 나온다.

 

이 모듈은 Fig. 2에 나타냈으며, 어떤 네트워크 블록(예: 어텐션 블록)에도 매끄럽게 통합될 수 있다.

(수식 (1)은 1X 리포트의 adaLN-Zero와 사실상 같은 형태다. LayerNorm(x) × (1+α) + β가 scale·shift 변조이고, x = x + FFN(...)가 residual이다. 다만 여기서는 α, β가 프레임 i마다 따로 나온다 — 이것이 "frame-level"의 실체다. 1X에서 우리가 "국소(local) 조건"이라 부른 것과 정확히 같은 아이디어다.)


4.2 Motioin-Reinforced Loss

 

 

전통적인 비디오 생성 모델은 학습 목적함수로 연속 공간에서는 제곱 l₂ 거리(rectified flow의 경우)를, 이산 공간에서는 cross-entropy 손실(next-token-prediction 트랜스포머의 경우)을 흔히 사용한다.

 

목적함수들은 일반 비디오 생성에는 효과적이지만, "모든 픽셀을 동등하게 취급"하기 때문에 액션에 의존하는 상태 변화를 포착하는 능력을 훼손할 수 있다.

 

따라서 이들은 RL 에이전트를 위한 월드 시뮬레이터로 효과적으로 기능하기에는 비효율적이다.

RL에서는 고품질 배경 디테일을 유지하는 것보다 동역학 전이를 정확히 모델링하는 것이 학습에 더 결정적이기 때문이다.

이 한계는 RL 에이전트가 정적인 시각 요소가 아니라 액션이 유발한 상태 변화로부터 주로 학습하기 때문에 발생한다.

 

이 문제를 해결하고 더 정밀한 동역학 모델링을 가능하게 하기 위해, 비디오 모델의 액션 추종 능력을 개선하는 새로운 motion-reinforced loss를 도입한다. 각 학습 스텝에서 정답 비디오 임베딩 배치 **x = {x⁰, ⋯, x^k, ⋯, x^j, x^{H−1}}**를 샘플링하고(H는 클립의 호라이즌), 연속된 프레임 사이의 차이 ω를 다음과 같이 계산한다.

 

여기서 ω_i → [1, c] 범위를 가지며, 초기 프레임 x⁰은 조건 프레임 역할을 하므로 ω_0 = 1로 둔다.

여기서 c는 motion-reinforced 강도를 조절하는 하이퍼파라미터다.

 

ω를 얻은 뒤, 이를 픽셀 단위 가중치로 지도학습 손실에 통합한다. 최종 motion-reinforced 손실은 다음과 같다.

 

여기서 ℒ_prev는 디퓨전 모델의 원래 MSE 손실이거나, 트랜스포머 기반 구조의 cross-entropy 손실이다.

 

이 정식화를 통해, 프레임 간에 변하는 픽셀이 손실 역전파에 더 큰 영향을 미치게 된다.

이 픽셀들은 보통 연속적으로 변하는 모션 관련 요소에 해당하며, 안정적으로 유지되는 배경은 학습 중 영향이 작아진다.

이 메커니즘은 액션 조건부 예측에 거의 기여하지 않는 정적 배경 요소의 영향을 본질적으로 약화시킨다.

 

동역학 전이를 강조함으로써, 본 연구의 접근법은 월드 모델이 "액션-상태 인과관계(action-state causal relationships)"를 포착하는 능력을 높이고, 이로써 강화학습에서 더 효과적인 정책 학습을 촉진한다.


Figure 3 — 실험 관찰

Fig. 3(a)에서 보듯, 프레임 간 차이는 주로 모션 관련·동적 요소에 해당하므로, ω가 학습 중 이 픽셀들에 더 높은 가중치를 부여하게 된다.

 

Fig. 3(b)는 BAIR 데이터셋에서 iVideoGPT를 파인튜닝한 SFT 손실(ℒ_prev) 곡선을 제시하며, 본 연구 방법의 여러 변형을 비교한다.

실험 결과는 motion-reinforced loss와 action-conditioned module 둘 다 결정적 구성요소임을 보여준다.

어느 하나라도 없으면 프레임 간 동역학을 가진 액션 조건부 비디오 예측 성능이 크게 저하되기 때문이다.

(이것이 이 논문의 ablation이다. Fig. 3b에서 DWS(전체) < DWS w/o motion-reinforced loss, DWS w/o action-conditioned module 순으로 손실이 낮다. 1X 리포트와 결정적으로 다른 지점 — DWS는 자기 두 기여를 각각 뺀 baseline을 실제로 그렸다.)

 

따라서 ℒ_motion으로 파인튜닝된 비디오 모델은 학습하기 어려운 복잡한 시각 디테일 대신 동역학·모션 예측에 더 집중하게 된다.

나아가 월드 시뮬레이터에게는 정적 배경 디테일보다 동역학·모션 일관성이 더 중요하다.

시뮬레이터는 액션에 조건화된 시각적 결과를 예측해야 하기 때문이다.

 


 

4.3 Model-Based Reinforcement Learning

 

사전학습 비디오 생성 모델에서 파인튜닝된 비디오 기반 월드 시뮬레이터의 가장 유망한 응용 중 하나는,

이를 MBRL의 정책 학습을 위한 월드 모델로 활용하는 것이다.

MBRL에서 에이전트는 학습된 월드 모델과 상호작용하여 누적 보상을 최대화하는 정책을 최적화하며, 이는 샘플 효율을 개선한다.

 

MBRL에서 효과적인 정책 학습을 위해, 월드 모델은 전이 동역학과 보상을 모두 예측해야 한다.

보상 예측 모델로 월드 모델을 완성한다.

보상 추정은 스칼라 예측 문제이므로, 선형 레이어, self-attention Block, cross-attention Block 으로 구성된 별도의 모델

R_ψ를 도입하여 과거 관측과 액션이 주어졌을 때의 보상을 추정한다.

 

RL 에이전트는 공유 CNN 백본에서 갈라져 나온 별도의 "정책 헤드(policy head)"와 "가치 헤드(value head)"로 이루어진

"actor-critic Network"를 사용한다.

MBPO 프레임워크(Janner et al., 2019b) 위에서, 합성 롤아웃(synthetic rollouts)으로 리플레이 버퍼를 증강하여

"표준 actor-critic RL 알고리즘"을 학습한다. 베이스 알고리즘으로는 "PPO(Schulman et al., 2017)"를 채택한다.


Prioritized Imagination

월드 모델의 상상(imagination)은 초기 관측(initial observations)에서 시작해야 하며, 이는 환경에서 수집된 경험에서 샘플링된다.

이 초기 상태들은 월드 모델이 합성 궤적을 생성하는 출발점이 된다.

 

기존 MBRL 방법들은 상상을 위한 초기 관측을 "균일하게 샘플링(uniform sampling)"한다.

그러나 이 전략은 정책 학습에서 서로 다른 상태가 가진 중요도의 차이를 무시하여 학습 비효율을 초래한다.

 

본 연구에서는 "서로 다른 초기 관측에서 비롯된 상상 전이들이 MBRL 정책 최적화에 대한 중요도와 과제 관련성에서 상당한 이질성(heterogeneity)"을 보인다는 점을 강조한다.

 

파인튜닝된 비디오 생성 모델의 월드 시뮬레이션 능력을 더 잘 끌어내기 위해, 본 연구는 더 가치 있는 전이에 선택적으로 집중하는 "prioritized imagination" 방법을 제안한다.

 

핵심 통찰은 더 높은 학습 잠재력(learning potential)과 학습 가능성(learn-ability)을 가진 전이로 이어지는

"초기 관측을 더 자주 샘플링해야 한다"는 것이다.

 

환경과의 상호작용 중 마주친 관측을 저장하는 버퍼 𝔅를 유지하고, 기대 학습 진전(expected learning progress)이 높은 초기 관측을 우선시한다. 이 학습 진전은 "TD 손실(TD loss)"의 크기로 측정된다.

 

이 우선순위화 메커니즘은 정책 학습에 더 큰 기여를 하는 상태에 상상 자원을 집중시켜 월드 모델을 더 효율적으로 활용하게 한다.

전체 의사코드는 다음 Algorithm 1에 제시된다.

(핵심: 상상 롤아웃의 출발점을 무작위로 고르지 말고, TD 손실이 큰 상태 = 아직 에이전트가 잘 모르는 상태를 더 자주 고르자는 것. 강화학습의 Prioritized Experience Replay(PER)를 상상 롤아웃의 시작점 선택에 옮겨온 발상이다.)


 

5. Experiments

이 절에서는 본 연구의 방법으로 파인튜닝된 비디오 월드 모델을 두 가지 핵심 관점에서 평가한다.

(1) 액션 시뮬레이션 능력 — 액션 조건부 비디오 예측의 품질로 평가,

(2) 온라인·오프라인 모델 기반 강화학습에서의 효과 — 과제별 누적 리턴(cumulative return)으로 정량화.

 

5.1 Action-Conditioned Simulation

월드 시뮬레이션을 위한 액션 조건부 비디오 예측 개선에서 DWS의 효과를 평가하기 위해, 아키텍처가 뚜렷이 다른 두 사전학습 모델로 실험한다. 디퓨전 기반 Open-Sora와 자기회귀 기반 iVideoGPT다.

 

계산 효율을 위해, 원래 1.1B 모델 대신 12층(약 280M 파라미터)으로 압축한 Open-Sora를 사용한다.

 

평가는 세 가지 데이터셋에서 수행한다.

연속 액션 공간BAIR, 그리고 이산 액션 공간Procgen과 Atari다.

 

Experiment Setup

BAIR 로봇 푸싱 데이터셋(로봇 팔이 다양한 물체를 조작)은 학습 43k, 테스트 256개 비디오로 구성된다.

선행 연구를 따라 단일 초기 프레임에서 15프레임을 예측한다.

이 벤치마크에서 디퓨전·마스크·자기회귀 계열의 다양한 비디오 예측 모델과 비교한다.

 

Procgen에서는 두 플랫포머 게임 Coinrun, Ninja를, Atari에서는 두 게임 Breakout, Battle Zone을 평가한다.

각 게임마다 RGB 관측·액션·보상을 포함한 100만 전이 스텝을 수집한다. 해상도는 BAIR·Procgen이 64×64, Atari가 84×84이다.

 

Metrics

예측 비디오의 품질을 재는 데 널리 쓰이는 네 지표를 채택한다.

 

1.FVD — 사전학습 네트워크 특징 간 Fréchet 거리참조·생성 비디오 분포의 통계적 유사도를 측정 (↓)

2. PSNR — 최대 픽셀값과 MSE의 로그 비율로 계산되는 픽셀 단위 충실도 (↑)

3.SSIM — 휘도·대비·구조 비교를 포함한 구조 정보 보존도 (↑)

4.LPIPS — 사전학습 딥 네트워크의 특징 공간 거리로 재는 지각적 이미지 유사도 (↓)

게임·로보틱스 등 여러 도메인의 생성 비디오 정성 결과. 초기 관측과 조건 액션이 주어졌을 때, 본 연구의 방법으로 파인튜닝된 Open-Sora와 iVideoGPT가  동적 월드 모델링 능력을 크게 개선 함을 관찰한다.

 

사전학습 모델과 DWS 파인튜닝 모델의 비교.  DWS 없는 베이스 모델은 일관성 없는 픽셀을 생성하고 조건 액션과 일치하지 못한다 (빨간 원으로 표시).

Qualitative Results Analysis

보지 못한(unseen) 초기 프레임으로 두 모델을 정성 평가한다. Fig. 4는 처음 보는 입력에 대한 생성 예시를 보여준다.

DWS로 파인튜닝된 두 베이스 모델 모두 일관된 시간적 동역학과 배경 보존을 갖춘 고품질·제어 가능 비디오를 성공적으로 생성한다.

또한 배경과 텍스처가 다양한 처음 보는 입력을 처리할 때 강건한 일반화 능력을 보인다.

Fig. 5에서 보듯, 베이스 모델은 시각적 왜곡이 있는 비디오를 생성하는 경향이 있는 반면, DWS는 객체 일관성을 유지하고 정밀한 시각 예측을 만들어 출력 품질을 크게 높인다.

 

Quantitative Results Analysis

Table 1, 2에서 다음 관찰이 나온다.

(i) DWS는 서로 다른 베이스 모델 전반에서 액션 조건부 비디오 예측 성능이 우수하다.

 

연속 액션 공간의 BAIR에서, DWS는 디퓨전 기반 Open-Sora와 자기회귀 기반 iVideoGPT 둘 다의 성능을 크게 높여, 아키텍처를 넘나드는 일반화 가능성과 다목적성을 입증한다.

 

정적·전역 텍스트 프롬프트에 조건화하는 전통적 text-to-video 모델인 Open-Sora조차 DWS 파인튜닝 후 액션 제어성에서 큰 개선을 얻는다. 구체적으로 FVD 11.7% 감소, LPIPS 27.9% 감소를 관찰한다.

 

액션 조건부 생성을 위해 특별히 설계된 iVideoGPT에서도 DWS는 눈에 띄는 개선을 달성한다.

이산 액션 공간의 Procgen·Atari에서도 DWS는 베이스 모델 성능을 뚜렷한 폭으로 일관되게 개선한다.

 

이 개선은 두 핵심 구성요소 덕분이다.

각 반응을 대응 프레임에 효율적으로 변조하는 액션 조건화 모듈과 프레임 간 픽셀 동역학을 효과적으로 포착하는 motion-reinforced loss.

 

(ii) DWS는 범용 방법으로서의 다목적성을 보인다.

연속 액션의 로보틱스부터 이산 액션의 게임까지 다양한 데이터셋에 효율적으로 적용되며, 디퓨전 기반과 자기회귀 트랜스포머 기반이라는 두 대중적 아키텍처에 매끄럽게 통합된다.

 


5.2 Model-Based Reinforcement Learning

DWS로 파인튜닝된 월드 시뮬레이터의 실용적 효용을 검증하기 위해, 이를 MBRL 정책 학습의 월드 모델로 활용하여 성능을 평가한다.

 

Benchmarks and Baselines

Procgen의 Coinrun, Ninja와 Atari의 Breakout, Battle Zone에서 실험한다.

모두 RGB 이미지를 관측으로, 이산 액션을 제어로 사용한다. 액션 차원은 Procgen 15, Breakout 4, Battle Zone 18이다.

prioritized imagination을 적용한 MBRL을 다음 baseline과 비교한다.

 

1) PPO — 널리 쓰이는 모델 프리 RL(본 연구 방법의 토대).

2) DreamerV3 — 순환 네트워크로 동역학을 예측하고 actor-critic으로 정책을 학습하는 SOTA 모델 기반 RL.

Procgen에서는 경쟁력 있는 PPG를 추가 비교한다.

 

 

Results Analysis

Fig. 6은 DWS 학습 월드 모델을 단순한 PPO와 결합했을 때, 바닐라 PPO와 SOTA 모델 기반 RL인 DreamerV3를 모두 크게 능가함을 보여준다.

 

Procgen에서 DWS는 PPO·PPG 같은 모델 프리 접근을 상당히 능가한다.

Procgen 특유의 에피소드마다 배경·객체가 달라지는 특성 때문에 DWS 월드 모델은 MBRL 정책 학습 중 파인튜닝이 필요하지만, 기존 모델 기반 방법과 견줄 만한 성능을 유지한다.

 

Atari에서 DWS는 기존 방법을 크게 능가하는데, 이는 월드 모델이 액션 시뮬레이션을 위한 포괄적 동역학 지식을 습득했기 때문이다.

특히 Breakout에서는 SOTA 대비 7배 성능 향상을 달성한다.

 

이 우수한 성능, 특히 샘플 효율prioritized imagination 덕분이다.

이 기여를 Procgen에서의 ablation으로 검증하며, 결과는 Fig. 7에 나타난다.

 


5.3 Offline Model-Based Reinforcement Learning

나아가 학습된 비디오 월드 모델로 오프라인 데이터셋을 증강하여 정책을 개선하는 가능성을 탐색한다.

CoinRun, Ninja를 사례로, 잘 학습된 PPO 에이전트로 환경마다 100만 전문가 궤적을 수집해 baseline 데이터셋을 만든다.

 

각 환경에서 Open-Sora+DWS로 추가 100만 상태-액션 전이를 합성하여, 원본 데이터셋 크기를 사실상 두 배로 늘린다.

 

이 데이터 증강의 효과를 두 오프라인 RL 알고리즘 CQL(Conservative Q-Learning), IQL(Implicit Q-Learning)로 검증한다.

Table 3에서 보듯, 월드 모델 생성 데이터로 오프라인 RL을 증강하면 알고리즘·환경 전반에서 성능이 크게 향상된다.

이 결과는 DWS 월드 시뮬레이터가 의미 있는 상태-액션-보상 전이를 생성하여 오프라인 데이터셋을 효과적으로 보완함을 보인다.

이는 사전학습 비디오 모델을 RL용 월드 시뮬레이터로 쓰는 가능성을 보여주며, DWS 모델이 정확한 동역학 전이를 생성할 수 있음을 검증한다.


 

6. Conclusion and Limitation

이 논문에서 DWS를 제시했다.

이는 대규모 데이터셋에서 학습된 풍부한 사전지식을 다운스트림 액션 조건부 시뮬레이션에 활용하여, 사전학습된 비디오 생성 모델을 월드 시뮬레이터로 효율적으로 적응시키는 새로운 접근법이다.

 

본 연구의 프레임워크는 두 가지를 도입한다.

액션-프레임 정렬을 가능하게 하고 다양한 아키텍처에 매끄럽게 통합되는 경량 액션 조건화 모듈, 그리고 정확한 월드 시뮬레이션에 결정적인 프레임 간 픽셀 동역학을 모델링하도록 특별히 설계된 motion-reinforced loss다.

 

광범위한 실험 결과는 DWS가 비디오 예측 품질과 MBRL 성능을 모두 크게 개선하여 의미 있는 응용으로 이어짐을 보여준다.

 

그러나 DWS는 현재 긴 시간 호라이즌(extended temporal horizons)이나 높은 공간 해상도(high spatial resolutions)의

비디오를 모델링하는 데 한계가 있다.

이를 향후 연구로 남긴다.


Appendix A. Details of Dataset Collection

A.1 Atari

두 Atari 게임(Breakout, Battle Zone) 실험에서, 오프라인 데이터셋 수집에 잘 정립된 baseline인 DQN 에이전트를 사용했다.

오픈소스 DQN 구현(cleanrl)을 채택해 100만 환경 스텝 학습하고 리플레이 버퍼를 오프라인 데이터셋으로 저장했다.

 

본 연구는 더 정확한 월드 모델 개발에 초점을 두므로, 두 게임 환경을 random sticky action 없이 의도적으로 초기화했다.

이는 수집된 궤적이 환경의 정답 상태 전이를 정확히 포착하도록 보장한다.

(sticky action = 이전 액션이 확률적으로 반복되는 무작위성. 이를 끄면 "이 액션 → 이 결과"의 인과가 깨끗해져 월드 모델이 배우기 쉬워진다. 다만 실제 Atari 벤치마크는 보통 sticky action을 켜서 평가하므로, 이 세팅은 월드 모델 학습에 유리하게 단순화한 것이다.)

 

또한 Frame Skip 4의 FrameSkipping Wrapper와

RGB 관측을 (84, 84)로 리사이즈하는 ResizingObservation Wrapper를 적용했다.

DQN 하이퍼파라미터는 기본 설정을 따랐다.

 

A.2 Procgen

Atari와 유사하게, 두 Procgen 게임(Coinrun, Ninja)에서 리플레이 데이터셋을 수집했다.

공개 PPO 구현을 사용해, 게임마다 기본 해상도 64×64로 100만 상태 전이를 수집했다.


Appendix B. Implementation Details

B.1 Details of Base Models

Open-Sora

 

Open-Sora는 rectified-flow와 시공간 트랜스포머 기반의 텍스트·프레임 조건부 비디오 생성 모델이다.

DWS 구현에는 Open-Sora v1.2를 베이스로 썼다.

원래 아키텍처와의 일관성을 위해 동일한 VAE를 사용했으나, 계산 효율을 위해 원래 텍스트 인코더 대신 T5-small을 채택했다.

 

원래 Open-Sora가 11억 파라미터로 추론·다운스트림에 큰 제약이 되므로, 앞 12개 레이어만 사용해 모델을 초기화했다.

이로써 계산 부담을 크게 줄이면서 핵심 생성 능력은 보존했다. 액션 조건화 모듈은 Open-Sora의 모든 레이어에 통합된다.

 

iVideoGPT

 

iVideoGPT는 LLaMA 기반의 자기회귀 트랜스포머다. 전용 보상 예측 헤드 레이어를 추가해 비디오 생성과 보상 추정을 모두 수행한다.

특징은 효율적 예측을 위한 "압축 토큰화를 구현한 문맥 인식 토크나이저(context-aware tokenizer)"다.

 

토크나이저 114M + 트랜스포머 138M 크기의 사전학습 가중치로 초기화한다. 액션 조건화 모듈은 각 트랜스포머 블록에 통합된다.

 

B.2 Implementation of Motion-Reinforced Loss

디퓨전 모델은 손실 계산과 그래디언트 갱신에 제곱 l₂ 거리를 쓰며, 원래 손실은 다음과 같다.

 

y는 조건, g는 추정할 타깃 신호로, 아키텍처마다 형태가 다르다. **DDPM 스타일에서는 정답 노이즈, flow-matching에서는 속도(velocity)**에 해당한다.

 

p_θ의 출력과 비디오 임베딩 x_t의 차원이 동일하므로, 모션 기반 가중치 ω를 손실에 통합할 수 있다.

 

자기회귀 트랜스포머는 cross-entropy 손실을 쓰며 원래 손실은 다음과 같다.

 

x_i는 이산 토큰, L은 시퀀스 길이다.

 

이 경우 ω는 다음 수식으로 계산하며, c = e(자연상수), 𝕀는 지시함수다.

 

최종 손실은 다음과 같다.


B.3 Implementation of the MBRL Algorithm

MBPO 프레임워크 안에서 DWS 월드 모델을 쓰는 단순한 MBRL 알고리즘을 만들었고, 베이스 actor-critic으로 PPO를 쓴다.

cleanrl 구현 위에 코드를 짰고, actor-critic 학습에는 동일한 하이퍼파라미터·아키텍처를 쓴다.

MBRL 전용 하이퍼파라미터는 Table 4에 있고, DreamerV3를 따라 보상 예측에 symlog 변환을 쓴다.

 

이상으로 본 논문의 리뷰를 마친다.

1x World Model Challenge Techinical Report 

 

0.Abstract

 

월드 모델은 미래를 예측하는 내부 시뮬레이터다.

그런데 무엇을 예측할 것인가에서 갈래가 둘로 나뉜다 — 시각적 관측(픽셀), 압축된 잠재 상태(토큰)

1X World Model Challenge의 두 트랙이 정확히 이 두 갈래에 대응한다.

 

Sampling 트랙은 미래의 이미지 프레임을 픽셀 공간에서 직접 예측하는 과제다.

저자들은 이 트랙을 처음부터 학습하는 대신, Video Generation Foundation Model인 Wan-2.2 TI2V-5B

video-state-conditioned future frame prediction에 맞게 개조했다.

원래 이 모델은 텍스트와 이미지 한 장을 받아 비디오를 만드는 TI2V 모델인데, 이를 과거 영상과 로봇 상태를 조건으로 미래 프레임을 예측하도록 바꾼 것이다.

로봇 상태를 주입하는 데는 AdaLN-Zero를 썼고, 5B 규모의 모델을 통째로 학습할 수는 없으니 LoRA로 일부만 추가 학습했다.

두 기법 모두 사전학습된 가중치를 최대한 건드리지 않겠다는 하나의 원칙에서 나온 선택이다.

그 결과 PSNR 23.0 dB를 기록했다.

 

Compression 트랙은 미래의 이산 잠재 코드, 즉 토큰을 예측하는 과제인데, 여기서는 시공간 트랜스포머(Spatio-Temporal Transformer)를 아무런 사전학습 없이 처음부터 학습시켰다.

픽셀이 아니라 토큰을 다루는 순간 문제가 GPT의 다음 토큰 예측과 같은 형태로 바뀌기 때문에, 검증된 언어모델 레시피를 그대로 이식할 수 있었다. 그 결과 Top-500 cross-entropy 6.6386을 기록했다.

 

그리고 두 트랙 모두에서 SOTA를 차지했다.

 

 

1.Introduction

Figure 1은 가로로 입력, 생성, 평가 세 단계로 나뉘고, 위가 픽셀 공간(Sampling), 아래가 토큰 공간(Compression)이다. 두 줄을 잇는 "Tokenizer" 화살표가 의미하는 바는, 같은 원본 영상을 쓰되 표현 공간만 다르다는 뜻이다.

 

Sampling 트랙에서 모델은 과거 17프레임을 관측하고, 여기에 더해 로봇의 상태를 함께 받아 미래 프레임을 픽셀 공간에서 생성한다.

여기서 짚고 넘어갈 비대칭이 하나 있다. 영상은 과거만 주어지는데, 로봇 상태는 과거와 미래가 전부 주어진다.

즉 "로봇이 앞으로 이렇게 움직일 테니, 그때 카메라에 무엇이 보일지 그려봐라"에 가깝다.

순수한 미래 예측이 아니라 행동 조건부 예측이고, 이 설정이 문제 난이도를 상당히 낮춘다.

 

평가는 예측한 77번째 프레임과 정답 프레임 사이의 PSNR로 이뤄진다.

여러 프레임을 생성하지만 점수는 마지막 한 장으로만 결정된다는 점이 중요하다.

중간 과정이 아무리 엉망이어도 마지막 한 장만 맞으면 되기 때문에, 시간적 일관성은 사실상 평가되지 않는다.

뒤에서 저자들이 후처리와 앙상블 트릭으로 점수를 크게 끌어올릴 수 있었던 구조적 배경이 여기에 있다.

 

Compression 트랙에서는 Cosmos 8×8×8 토크나이저가 17개의 RGB 프레임을 3 × 32 × 32 형태의 잠재 토큰 그리드로 압축한다.

공간은 256을 8로 나눠 32가 되고, 시간은 (17−1)/8 + 1 = 3이 된다.

시간축 공식이 단순한 나눗셈이 아닌 이유는 첫 프레임을 따로 인코딩하는 causal 구조 때문이다.

프레임 0은 앞에 묶을 과거가 없으니 혼자 남고, 나머지 16프레임이 8개씩 두 덩어리로 묶인다.

모델은 다음 17프레임에 대응하는 다음 세 개의 잠재 토큰 그리드를 예측해야 한다.

 

평가는 예측 토큰과 정답 토큰 사이의 Top-500 cross-entropy에 기반한다.

어휘 전체가 아니라 토큰당 상위 500개 로짓만 보고 계산하는 방식인데, 이 때문에 확률을 소수의 후보에 집중시키는 모델이 유리해진다.

나중에 저자들이 greedy decoding을 택하는 이유가 바로 여기서 나온다.

 

 

 

월드 모델은 에이전트에게 자신이 놓인 환경의 내부 시뮬레이터를 쥐여준다.

로봇이 머릿속에 세계의 복제본을 하나 갖고 있는 셈이다.

행동의 결과를 미리 "상상"함으로써, 에이전트는 실제 세계와 직접 부딪히지 않고도 계획을 세우고 결과를 예측하고 의사결정을 개선할 수 있다.

이 마지막 조건이 중요한데, 실제 로봇으로 시행착오를 반복하는 일은 비싸고 위험하기 때문에 월드 모델의 존재 이유는

결국 샘플 효율성으로 귀결된다.

 

월드 모델링의 핵심 난제는 아키텍처 설계다.

충분히 표현력이 높으면서 동시에 계산적으로 다룰 수 있어야 한다는 것인데, 이 둘은 보통 서로 상충한다.

표현력을 키우면 계산이 무거워지고, 계산을 가볍게 하면 표현력을 잃는다.

이 논문이 뒤에서 시공간 트랜스포머를 도입하는 이유도 결국 이 트레이드오프에 대한 대응이다.

 

초기 접근법들은 대체로 순환 신경망이나 다층 퍼셉트론에 의존했다.

Dreamer 계열이 전자를, PETS나 TD-MPC2 같은 연구가 후자를 대표한다.

그러다 최근 들어 생성 모델링의 발전이 아키텍처 선택의 새로운 흐름을 만들었다.

 

이 새로운 흐름은 다시 두 갈래로 나뉜다.

한쪽은 이산 잠재 공간 위에서 자기회귀 트랜스포머를 활용하는 계열로, DeepMind의 Genie가 대표적이다.

다른 한쪽은 디퓨전과 플로우 기반 접근을 탐색하는 계열로, DIAMOND나 Oasis가 여기 속한다.

그리고 이 방법들을 대규모로 확장하면, 사실적이고 정확한 비디오 예측을 만들어내는 강력한 파운데이션 모델이 된다.

Wan과 NVIDIA Cosmos가 그 결과물이다.

 

2.Sampling Challenge

 

 

2.0 Problem Statement

1. Sampling 과제 : 모델은 로봇이 2초 뒤에 보게 될 512×512 프레임 한 장을 예측해야 함

2. 제공되는 조건 : 첫 17프레임과 로봇 상태의 전체 시퀀스

3. 평가 : 예측한 마지막 프레임과 정답 프레임 사이의 PSNR

 

PSNR은 픽셀 단위 MSE를 로그 스케일로 바꾼 값인데, MSE 기반이라는 점이 중요하다.

MSE를 최소화하는 최적해는 가능한 모든 미래를 픽셀별로 평균낸 것이고, 그건 필연적으로 흐릿한 이미지다.

뒤에서 저자들이 앙상블 평균으로 점수를 크게 끌어올리는 트릭도, CFG가 끝내 도움이 되지 않는 것도 전부 이 한 가지 성질에서 파생된다.

 

2.1 Data Pre-processing 

원본 77프레임 클립을 4배 downsampling해 21개짜리 짧은 클립으로 만든다.

그러면 조건 프레임(conditioning frame)이 다섯 장(x₀, x₄, …, x₁₆) 남고, 나머지 16개가 예측 대상이 된다.

이렇게 줄인 클립을 Wan2.2-VAE가 다시 압축하는데, 첫 프레임에는 공간 압축(spatial compression)만 적용하고

나머지 프레임에는 시간 압축(temporal compression) 4배를 적용한다.

결과적으로 길이 L의 클립은 (1 + (L−1)/4) 길이의 잠재 시퀀스가 된다. L = 21이면 여섯 개의 latent 프레임이 나온다.

 

첫 프레임만 예외로 취급하는 이유는 시간 압축이 여러 프레임을 묶어야 하는데 프레임 0에는 묶을 과거가 없기 때문이다.

Cosmos 토크나이저의 (17−1)/8 + 1 = 3과 정확히 같은 구조이고, 비디오 VAE와 토크나이저의 공통 설계 패턴이다.

다만 숫자를 직접 검산해보면 조금 어긋난다.

77프레임을 4배 다운샘플하면 인덱스 0, 4, …, 76으로 스무 개가 나오는데 논문은 스물한 개라고 쓴다.

아마 L = 21이어야 VAE 공식이 1 + 20/4 = 6으로 딱 떨어지기 때문에 역으로 맞춘 것으로 보인다.

재현하려면 코드를 확인해야 하는 지점이다.

 

2.2 Base Model

https://huggingface.co/Wan-AI/Wan2.2-TI2V-5B

모델은 Wan 2.2 TI2V-5B를 가져다 쓴다.

30-Layer DiT backbone을 가진 flow-matching 기반 비디오 생성 모델로, 50억 개의 파라미터를 가진다.

원래 이 모델은 텍스트와 이미지 한 장을 받아 비디오를 만드는 TI2V(Text-Image-to-Video) 구조인데, 저자들은 이를 과거 영상과 로봇 상태를 조건으로 받도록 개조했다.

그리고 모델은 512×512 픽셀을 직접 다루는 게 아니라 Wan2.2-VAE(which compresses clips to a size (1 + (L − 1)/4) × 16 × 16.)가 압축한 잠재 표현 위에서만 동작한다.

비디오는 데이터가 워낙 커서 픽셀 공간에서 디퓨전을 직접 돌리는 건 계산적으로 불가능하기 때문이다.

 

 

2.3 Video-State Conditioning 

영상 조건화를 통합하기 위해 입력 latent의 마스킹을 수정했다.

저자들은 가중치도, 레이어도, 파라미터도 하나 건드리지 않았다. 그저 입력 latent의 마스킹만 바꿨을 뿐이다.

표준 image-to-video 모델에서는 시간 차원의 첫 번째 latent가 마스킹되어 있다.

그 슬롯은 노이즈를 씌우지도 생성하지도 않고 깨끗한 상태로 고정되며, 이것이 곧 조건 역할을 한다.

나머지 슬롯들만 노이즈에서 시작해 생성된다.

 

즉 마스크는 "어디까지가 주어진 것이고 어디부터가 만들 것인가"를 정의하는 스위치다.

저자들은 이 아이디어를 그대로 확장해, 고정할 슬롯을 하나가 아니라 여러 개로 늘렸다.

그것만으로 모델의 과제가 image-to-video에서 video-to-video로 바뀐다.

 

영리한 지점은, 사전학습된 5B 모델이 이미 "고정된 슬롯을 조건으로 나머지를 채우는" 능력을 학습했다는 것이다.

고정 슬롯을 늘리는 건 모델 입장에서 이미 아는 과제의 자연스러운 확장이다.

추가 학습 비용이 거의 없고 사전학습 지식이 하나도 파괴되지 않는다. 비용이 0인 트릭이다.

 

남은 문제가 하나 있다. 원래 Wan 2.2는 텍스트 프롬프트도 조건으로 받는다.

그런데 1X 데이터셋에는 텍스트 설명이 없다. 저자들은 cross-attention 레이어를 그대로 남겨둔 채 빈 문자열을 프롬프트로 넣는 방식을 택했고, 향후 연구가 텍스트 조건화를 활용할 수 있도록 하기 위해서라고 언급한다.

 

Wan2.2 TI2V-5B의 DiT 아키텍처는  adaLN-Zero 로 상태 조건화를 가능하게 하고, 이를  Flow Matching 스케줄러의 타임스텝과 결합 하도록 수정되었다.

 

로봇 상태를 모델에 주입하는 방식은 adaLN-Zero다. Wan의 DiT 블록 안에서 작동한다.

 

과정은 네 단계다.

① 시간 해상도 맞추기 : 먼저 상태 시퀀스를 다운샘플링해 다운샘플된 영상과 시간 해상도를 맞춘다.

 

② 특징 투영 : 그다음 연속값인 각도와 속도에 사인파 특징을 덧붙이고, 모든 상태를 MLP로 은닉 차원 256에 투영한다.

각도에 사인파를 쓰는 이유는 각도가 주기적이기 때문이다. 0도와 360도가 같은 값인데 raw 스칼라로 넣으면 모델이 이 순환성을 학습하기 어렵다.

 

③ 시간 차원 압축 : 투영된 특징은 2층 1D 합성곱을 통과하며 시간 차원이 압축된다. Wan-VAE가 영상 프레임에 가한 압축과 정확히 같은 비율로 맞추기 위해서다.

그래야 상태 특징과 프레임별 latent가 1대 1로 정렬된다. 결과는 ((1 + L//4), r_dim) 형태가 된다.

 

④ 변조값 생성 : 마지막으로 이 특징을 MLP에 넣어 adaLN-Zero가 사용할 변조값을 얻고, 이렇게 얻은 로봇 변조를 flow matching 타임스텝의 변조에 더한다.

 

adaLN-Zero는 DiT 논문에서 제안된 조건 주입 기법이다.

조건 벡터로부터 γ와 β(scale, shift), 그리고 α(residual gate)를 만들어 LayerNorm과 잔차 연결을 변조한다.

이름의 "Zero"는 α를 생성하는 마지막 레이어를 0으로 초기화한다는 뜻이다.

학습 초기에는 조건의 영향이 정확히 0이라서 사전학습된 5B 모델이 항등함수처럼 동작하고, 학습이 진행되면서 조건이 서서히 개입한다. 사전학습 가중치를 파괴하지 않고 새로운 조건을 주입하는 데 적합한 구조다.

 

이 절에서 가장 중요한 문장은 마지막에 나온다.

로봇 변조는 latent에 다르게 작용한다는 것이다. 타임스텝 임베딩은 latent 전체에 동일하게 적용되지만, 상태는 각 프레임에 대응하는 latent 슬라이스만 변조한다.

 

두 조건의 작용 범위가 다르다는 뜻이다.

 

flow matching 타임스텝 t는 전역(global) 조건으로, "지금 노이즈 제거가 몇 퍼센트 진행됐는가"를 알려준다.

로봇 상태 s는 국소(local) 조건으로, "이 시점에 팔이 어디 있는가"를 해당 프레임에만 알려준다.

 

원래 DiT의 adaLN은 시퀀스 전체에 단일 변조를 적용하는데, 이를 시간축을 따라 달라지도록 확장한 것이 이 논문의 실질적인 아키텍처 기여다.

 

2.4 Training

학습은 AdamW로 23k 스텝, 학습률은 4×10⁻⁴ 고정이다.

Wan 2.2 DiT 백본에 rank 32의 LoRA 파인튜닝을 적용했다. 

저차원 부분공간만 업데이트하므로 큰 학습률을 견딜 수 있고, 스케줄러 없이 constant LR을 쓰는 것도 관행이다.

rank 32면 학습되는 파라미터는 전체의 1퍼센트에도 못 미친다.

 

학습 중 classifier-free guidance를 적용한 경우와 적용하지 않은 경우를 모두 실험했지만 PSNR에서 유의미한 개선은 없었다.

논문은 여기서 멈추지만 이유는 명확하다.

CFG는 조건부 예측과 무조건부 예측을 외삽해 분포의 mode 쪽으로 밀어붙이는 기법이고, 그래서 이미지가 선명해진다.

그런데 PSNR은 mean을 원한다.

목적이 정반대이므로 CFG가 통하지 않는 것은 우연이 아니라 필연이다.

 

실제로 Table 2를 보면 CFG scale을 2.0, 1.5, 1.0으로 낮출수록 PSNR이 24.20, 24.59, 25.53으로 단조 증가한다.

 

학습 자원은 DataCrunch 클러스터의 4개 노드, 노드당 NVIDIA B200 8장으로 총 32장이다.

B200의 184GB VRAM 덕분에 메모리를 많이 쓰는 비디오 생성 모델을 효율적으로 학습할 수 있었다고 밝힌다.

유효 배치 크기는 1024다.

저자 여덟 명 중 세 명이 DataCrunch 소속이라는 점을 감안하면 이 문단의 성격은 짐작할 수 있다.

다만 연구적으로 중요한 사실은 따로 있다.

이 결과는 B200 32장 없이는 재현되지 않는다. "LoRA라서 가볍다"는 인상과 달리 실제로는 상당한 자원이 투입됐다.

 

 

2.5 Inference

출발점은 대회 규칙이 추론 연산 시간을 제한하지 않는다는 관찰이다.

초기 시도에서 저자들은 2위 팀(Duke)의 후처리 파이프라인을 따라

예측 프레임에 Gaussian Blur를 적용하고 Histogram Matching을 수행했다. 이 후처리만으로 PSNR이 1.2dB 올랐다.

모델을 개선한 것이 아니라 예측 결과를 흐리게 만든 것뿐인데도 그렇다.

참고로 1등과 2등의 최종 격차가 1.44dB이니, 후처리 하나가 순위를 뒤집을 수 있는 크기다.

 

저자들은 그 이유도 정확히 진단한다.

PSNR은 타깃으로부터의 이상치 편차를 제곱으로 처벌하기 때문에, 약간의 오차를 가진 선명한 이미지가 비슷한 오차를 가진

흐릿한 이미지보다 낮은 점수를 받는다.

선명한 예측은 "팔이 5픽셀 어긋남" 같은 큰 국소 오차를 만들고, 흐릿한 예측은 오차를 넓게 퍼뜨려 제곱합을 줄인다.

 

여기서 저자들의 핵심 아이디어가 나온다.

예측 불확실성을 앙상블로 착취하는 편이 Gaussian Blur보다 낫다는 것이다.

같은 조건으로 여러 번 샘플링하면 모델이 확신하는 영역(배경, 정적 물체)은 매번 거의 같게 나오고,

불확실한 영역(빠르게 움직이는 팔)만 매번 다르게 나온다.

이 샘플들을 픽셀 단위로 평균내면 확신하는 영역은 선명하게 남고 불확실한 영역만 자동으로 흐려진다.

균일하게 뭉개는 Gaussian Blur와 달리 Spatially Adaptive Blur가 만들어진다.

실제로 논문은 휴머노이드의 팔처럼 움직임이 큰 영역에서 주로 블러가 발생한다고 보고한다.

 

이론적으로도 근거가 있다.

여러 샘플의 평균은 조건부 기댓값에 대한 Monte Carlo Approximation이고, 조건부 기댓값은 MSE의 최적 추정량이다.

PSNR이 MSE 기반이므로, 앙상블 평균은 PSNR을 직접 최적화하는 원리적으로 올바른 방법이다.

앙상블 샘플 수를 늘리자 검증셋과 공개 리더보드 양쪽에서 PSNR이 올랐다.

검증셋에서만 올랐다면 과적합을 의심할 수 있지만, 저자들이 볼 수 없는 리더보드에서도 같은 방향으로 움직였다는 것은

이 효과가 지표 자체의 성질에서 나온다는 뜻이다. 추론 스텝 수와 CFG 가중치를 조정한 결과도 함께 Table 2에 보고돼 있다.

 

 

2.6 Results

Table 2는 검증셋 PSNR과 함께 SSIM, LPIPS, FID를 추가로 보고한다.

대회가 요구하지 않은 지표를 자발적으로 덧붙였다.

 

표는 세 블록이다.

첫 블록은 CFG 없이 학습한 모델들로, 최종 예측에 평균낼 샘플 수를 1개에서 20개까지 바꿔가며 ablation을 수행했다.

논문은 샘플 수를 늘리면 검증 PSNR은 오르지만 나머지 지표에 드러나듯 시각적 품질은 저하된다고 서술한다.

 

숫자를 보면 이야기가 조금 다르다.

LPIPS와 FID는 확실히 나빠졌고, 특히 FID는 40.23에서 90.71로 2.3배 악화됐다.

PSNR 2.25dB를 얻는 대가로 이미지의 사실성을 절반 이상 포기한 셈이다. 그런데 SSIM은 오히려 개선됐다.

SSIM은 국소 평균·분산·공분산 기반이라 블러에 관대하고, LPIPS와 FID는 지각적·분포적 거리를 재기 때문에 블러를 잡아낸다.

 

첫 블록의 마지막 행은 검증셋으로 추가 학습한 모델이다.

검증셋 수치가 비교 불가능해진다는 점을 저자들이 명시했고, 리더보드에서는 PSNR이 0.58dB 올랐다고 보고한다.

이 행이 1위를 확정지은 모델(TEST 23.00)이다.

 

두 번째와 세 번째 블록은 학습 중 CFG를 적용한 모델들이다.

cfg_scale을 일정 수준 이상 올려도 PSNR이 개선되지 않았지만 해당 실행을 2순위 제출작으로 유지했다고 밝힌다.

 

마지막으로 같은 체크포인트에서 샘플링 스텝을 20에서 100으로 늘리자 TEST PSNR 23.04로 표 전체 최고값이 나왔다.

다만 여기엔 †가 붙어 있고, 캡션에 따르면 마감 이후 측정된 값이라 순위에 반영되지 않았다.

 

 

3. Compression Challnge

3 × 32 × 32 형태의 과거 그리드 3개와 64 × 25 형태의 로봇 상태를 문맥으로 받아, 미래 3개 그리드를 예측한다.

ST-Transformer는 L개 층의 시공간 블록으로 구성되며, 각 블록은 시간 스텝별 공간 어텐션, 동일 공간 좌표의 causal 시간 어텐션, feed-forward 네트워크를 포함한다. 공간·시간 어텐션의 각 색상은 하나의 self-attention map을 나타낸다.

 

픽셀 공간에서 직접 예측을 측정하는 Sampling Challenge와 달리, Compression Challenge는 이산 잠재 공간에서 모델을 평가한다. 각 비디오 시퀀스는 먼저 Cosmos 8×8×8 토크나이저를 사용해 이산 토큰 그리드로 압축되며, 이렇게 만들어진 압축 시퀀스는 시퀀스 아키텍처로 모델링할 수 있다.

 

3.0 Problem Statement

 

32 × 32 토큰 그리드 H = 3개와 과거·미래 양쪽 시간 스텝의 로봇 상태가 문맥으로 주어질 때, 과제는 32 × 32 토큰의 다음 M = 3개 그리드를 예측하는 것이다. 여기서 _{H:H+M−1}은 미래 프레임에 대해 예측된 토큰 그리드다.

 

 

토큰화된 학습 데이터셋 𝒟는 약 306,000개의 샘플을 포함한다.

각 샘플은 다음으로 구성된다.

  • 토큰화된 비디오: 6개의 연속된 토큰 그리드(과거 3, 미래 3), 각각 32 × 32 크기로, 샘플당 6,144개 토큰이며 전체로는 약 18.8억 토큰이다.
  • 로봇 상태: 대응하는 원본 비디오 프레임에 정렬된 시퀀스 s ∈ ℝ⁶⁴ˣ²⁵.

 

3개의 32 × 32 토큰 그리드 블록은 256×256 해상도의 17개 RGB 프레임에 대응하므로, 토큰 공간에서의 예측은 원본 비디오와 정렬을 유지한다. 성능은 top-500 cross-entropy로 평가되며, 이는 토큰당 상위 500개 로짓만 고려한다.

 

 

3.1 Spatio-temporal Transforemer

Genie를 따라, 이 월드 모델은 Vision Transformer(ViT) 위에 구축된다.
표준 트랜스포머의 이차(quadratic) 메모리 비용을 줄이기 위해, 공간 어텐션 블록과 시간 어텐션 블록을 번갈아 배치하고 그 뒤에 feed-forward 레이어를 두는 시공간(Spatio-Temporal) 트랜스포머를 사용한다.

 

공간 어텐션은 프레임당 1 × 32 × 32 토큰에 대해 어텐션하고, 시간 어텐션은 causal mask를 적용해 시간축을 따라 T × 1 × 1 토큰에 대해 어텐션한다.
이 설계는 주요 병목인 공간 어텐션이 프레임 수에 대해 선형으로 확장되게 한다.

 

안정성을 위해 pre-LayerNormQKNorm을 적용한다.
위치 정보는 공간 토큰과 시간 토큰 모두에 대해 학습 가능한 절대 위치 임베딩으로 추가된다.
트랜스포머는 24층, 8헤드, 임베딩 차원 512, 시퀀스 길이 T = 5를 사용했고, dropout 0.1을 어텐션·MLP·잔차 연결 전부에 적용했다.

 

3.2 State Conditioning

로봇 상태는 Bruce et al.6을 따라 덧셈 임베딩(additive embedding)으로 인코딩된다.
상태 벡터는 MLP로 투영되고, 1D 합성곱(커널 크기 3, 패딩 1)으로 처리되며, 절대 위치 임베딩으로 보강된 뒤 비디오 토큰과 결합된다.

 

3.3 Training

모델은 PyTorch로 구현했으며, β₁ = 0.9, β₂ = 0.95의 fused AdamW 옵티마이저로 80 에폭 학습했다.

 

weight decay 0.05는 파라미터 행렬에만 적용되었고, 편향(bias), 정규화 파라미터, gain, 위치 임베딩은 제외되었다.

 

GPT-2 및 Bertolotti and Cazzola, Press and Wolf를 따라 입력 임베딩과 출력 임베딩을 묶었다(weight tying).
이는 가장 큰 두 개의 가중치 행렬 중 하나를 제거하여 메모리 사용량을 줄이고, 일반적으로 학습 속도와 최종 성능을 모두 향상시킨다.

 

3.4 Training Objective

모델은 미래 시간 스텝에서 예측 토큰과 정답 토큰 사이의 Cross-Entropy Loss를 최소화하도록 학습되었다.

 

 

여기서 ẑ_t는 시간 t에서의 모델 출력이고, CE는 그리드 내 모든 토큰에 대한 cross-entropy 손실이며, 𝒟는 토큰화된 비디오 및 상태 시퀀스로 이루어진 데이터셋이다.

 

학습에는 시간축 전반의 병렬 계산을 가능하게 하기 위해 teacher forcing을 사용했으며, 2000 스텝의 warmup 이후 최고 8×10⁻⁴에서 0까지 선형으로 감소하는 학습률 스케줄을 적용했다.

 

3.5 Implementation

학습에는 bfloat16을 이용한 자동 혼합 정밀도(AMP : Automatic Mixed Precision)를 사용했지만,

추론에서는 bfloat16에서 성능이 저하되어 float32를 사용했다.

 

선형 레이어의 편향은 0으로 초기화되었고, 가중치(임베딩 포함)는 𝒩(0, 0.02)에서 추출되었다.

학습은 Sampling과 동일한 B200 DataCrunch 클러스터에서 유효 배치 크기 160으로 수행되었다.

 

 

3.6 Inference

 

자기회귀 모델은 위 식을 통해 시퀀스를 생성하며, 각 스텝은 각 공간 토큰에 대한 범주형(categorical) 분포를 출력한다.

 

샘플링은 z_t ∼ f_θ(·)를 추출하여 다양성을 도입하지만, 일반적으로 더 낮은 확률의 궤적과 더 높은 손실을 낳는다. 

 

 

반면 greedy decodingargmax를 선택하여 결정론적이고 확률이 높은 시퀀스를 생성하며, 이것이 효과적이면서도 효율적임을 확인했다.

 

 

3.7 Results

학습 시에는 teacher forcing(파랑)을 사용한다. 검증셋에서 비현실적인 teacher forcing(주황)으로 평가하고, 추론 시 실제로 사용될 greedy 자기회귀 생성(초록)으로도 평가한다.

 

 

teacher-forced 검증 손실은 정답 입력에 조건화하기 때문에 낙관적이지만, 달성 가능한 손실의 하한(lower bound)으로 해석될 수 있으며, 이는 완벽한 추론을 하는 이상화된 자기회귀 모델의 성능을 나타낸다.

 

teacher-forced 성능과 자기회귀 성능 사이의 간극을 줄이기 위해 scheduled sampling을 실험했다.
그러나 이는 유의미한 개선으로 이어지지 않았다.

 

 

 

 

 

 

 

Conclusion

이 리포트에서 우리는 두 1X World Model Challenge 전반에서 강력한 성능을 달성한 두 가지 상호보완적 접근법을 제시했다.

 

첫째, 사전학습된 image-text-to-video 파운데이션 모델을 파인튜닝함으로써 인터넷 규모의 데이터를 어떻게 활용할 수 있는지 보였다.

 

DataCrunch 인스턴트 클러스터에서의 다중 노드 학습을 통해, 단 36시간 만에 리더보드 1위에 도달했으며, 이는 약 한 달이 걸린 2위보다 한 자릿수(약 10배) 빠른 것이다.

 

추론을 더욱 개선하기 위해, 샘플들을 평균하여 예측 불확실성이 높은 영역을 선택적으로 흐리게 처리했다.

 

이것이 PSNR 최적화에는 효과적임이 입증되었으나, 다운스트림 의사결정에 가장 적합한 추론 전략은 여전히 열린 질문으로 남아 있다.

 

둘째, 시공간 트랜스포머 월드 모델이 토큰화된 데이터셋에서 17시간 이내에 학습될 수 있음을 입증했다.

또한, greedy 자기회귀 추론이 속도와 정확도의 실용적 균형을 제공함을 확인했다.

 

단순함에도 불구하고, 이 모델은 다른 리더보드 참가작들보다 상당히 낮은 손실 값을 달성했다.

 

두 트랙 모두에서 1위를 만든 결정적 요인은 아키텍처가 아니라 추론 전략(앙상블 평균 / greedy decoding)이었고, 두 전략 모두 "평가지표가 불확실성 표현을 벌한다"는 사실을 이용한 것이다. 이는 저자들의 잘못이라기보다 대회 평가 설계의 구조적 한계이며, 저자들도 Conclusion에서 이를 "open question"으로 인정했다.

이 분야의 진짜 다음 연구 주제가 바로 그것이다.

https://arxiv.org/abs/2310.03744

 

Improved Baselines with Visual Instruction Tuning

Large multimodal models (LMM) have recently shown encouraging progress with visual instruction tuning. In this note, we show that the fully-connected vision-language cross-modal connector in LLaVA is surprisingly powerful and data-efficient. With simple mo

arxiv.org

Authors:  Haotian Liu, Chunyuan Li, Yuheng Li, Yong Jae Lee

Affiliation:  University of Wisconsin–Madison (Computer Sciences) · Microsoft Research, Redmond

Venue:  CVPR 2024 (IEEE/CVF Conference on Computer Vision and Pattern Recognition)

Preprint:  arXiv:2310.03744 (Oct 2023)

 

 

Abstract

핵심은 세 가지.

첫째, LLaVA라는 동일한 틀 위에서 설계 선택지를 조건을 맞춰 비교한 첫 체계적 연구라는 점.

둘째, 이미지와 언어를 잇는 단순한 완전연결(fully-connected) 커넥터가 의외로 강력하고 데이터 효율적이라는 발견.

셋째, 두 가지 간단한 수정 — MLP 커넥터, 그리고 응답 형식 프롬프트를 붙인 학술 VQA 데이터 — 만으로 11개 벤치마크에서 최고 성능을 냈다는 것이다. 최종 13B 모델은 약 1.2M 공개 데이터만 쓰고 단일 8×A100에서 하루면 학습이 끝난다.

Abstract의 메시지는 "복잡하게 만들지 않아도 된다"로 요약된다. 그런데 이 주장은 "그럼 기존엔 뭐가 문제였길래?"라는 의문을 부른다.

 

Figure 1.  세 패널로 논문 전체를 압축한다. (상단 레이더) 11~12개 벤치마크에서 빨간 선(LLaVA-1.5)이 거의 모든 축의 가장 바깥 — 전방위 우위. (좌측 로그 스케일 막대) 학습 샘플 수로, LLaVA-1.5는 ≈1.2M인 반면 Qwen-VL-Chat 1400M+50M, InstructBLIP 129M+1.2M —  수백~수천 배 적은 데이터 로 더 높은 성능. (우측 구조도) CLIP ViT-L/336px → MLP 커넥터 → Vicuna v1.5 13B.

 

1. Introduction

최근 멀티모달 연구는 visual instruction tuning으로 수렴했고, LLaVA와 MiniGPT-4가 좋은 성과를 냈다.

그런데 핵심 문제가 있었다 — 벤치마크와 모델은 쏟아지는데 "어떻게 학습하는 것이 최선인가"에는 아무도 명확히 답하지 못했다.

예를 들어 LLaVA는 자유로운 대화형 추론에 강하고 InstructBLIP은 단답형 VQA에 강한데, 둘은 구조도 데이터도 완전히 달라서 무엇이 그 차이를 만들었는지 분리해낼 수가 없었다.

그래서 저자들은 변수를 하나씩 통제하는 실험을 설계한다.

LLaVA를 고정된 출발점으로 두고 입력·모델·데이터를 한 번에 하나씩만 바꿔가며 효과를 측정했다. 이 통제 실험이 두 기여로 이어진다.

  • 기여 1 — 더 강한 baseline. 단순한 완전연결 커넥터가 실제로 강력하고 데이터 효율적임을 규명했다. InstructBLIP·Qwen-VL이 수억~수십억 쌍으로 특수 resampler를 학습하는 반면, LLaVA는 projection 한 층을 약 60만 쌍으로만 학습하고도 8×A100으로 하루 만에, 공개 데이터만으로 최고 성능에 도달한다.
  • 기여 2 — 미해결 4문제 탐구. ① 고해상도 입력 확장, ② 조합적 능력, ③ 데이터 효율(최대 75% 줄여도 성능 유지), ④ 데이터 입도(granularity) 스케일링을 처음으로 함께 다뤘다.

 

2. Related Work

Instruction-following large multimodal models (LMMs).

멀티모달 모델은 보통 세 부품으로 구성된다 — 이미지를 인코딩하는 vision backbone, 지시를 이해하고 답을 생성하는 LLM, 둘을 이어 맞추는 cross-modal connector.

학습은 2단계다.

 alignment pretraining: 이미지-텍스트 쌍으로 시각 특징을 LLM의 단어 임베딩 공간에 정렬.

 visual instruction tuning: 시각 지시로 튜닝해 다양한 요청을 따르게 함. Q-Former 같은 리샘플러로 시각 토큰 수를 줄이는 방법도 있지만, LLaVA는 그런 복잡한 장치 없이 가장 단순한 축에 속한다.

Multimodal instruction-following data.

NLP의 교훈은 "지시 데이터의 품질이 모델 능력을 좌우한다"는 것이다.

LLaVA는 텍스트 전용 GPT-4로 COCO의 박스·캡션을 대화형 QA·상세 설명·복잡 추론 세 유형의 멀티모달 데이터로 확장했다.

다만 알려진 문제가 있었다 — 단순히 학술 VQA를 합쳐 넣으면 모델이 VQA에 과적합되어 자연스러운 대화 능력을 잃는다. 

이 '대화 ↔ 학술 균형' 문제가 다음 Approach의 핵심 질문이다.

 

3. Approach

3.1. Preliminaries

Approach는 먼저 LLaVA와 InstructBLIP의 장단점을 정밀하게 대비한다. LLaVA는 단일 선형층(single linear layer)으로 시각 특징을 언어 공간에 투영하고 LLM 전체를 학습한다. 대화엔 강하지만 단답형 학술 벤치마크엔 약하고, 예/아니오 질문에 'yes'로 치우친다 — 학습 분포에 단답 데이터가 거의 없었기 때문이다. 반대로 InstructBLIP은 학술 데이터를 넣었지만 대화가 약하다.

3.2. Response Format Prompting

이 약점을 Table 1(a)가 한 장면으로 보여준다.

Table 1(a).  택시 뒤에서 다림질하는 사진에 "이거 특이한가요? 자세히 설명해 주세요(explain in detail)"라고 묻는데, InstructBLIP(Vicuna-13B)은 딱 한 단어  "yes" 만 답한다. 상세 설명 요청에도 단답을 내놓는 — 단답형 VQA 과적합의 전형적 실패.

 

 

원인은 둘이다. 

(논리적 원인) 모호한 프롬프트  Q: {질문} A: {답} 형태는 "얼마나 길게 답하라"는 신호가 없다. 학습 신호가 모호하니 모델은 자연 대화에서도 짧게 답하도록 굳어진다. 

(구조적 원인) LLM을 학습하지 않음 — InstructBLIP은 LLM을 얼려두고 용량이 작은 Q-Former만 튜닝한다. 출력 길이·형식 제어를 시각 토큰이 떠맡아야 하는데, LLaMA급 LLM보다 용량이 작아 제어를 제대로 못 한다.

원인을 알면 해법은 단순하다. 출력 형식을 직접 지시하는 한 줄, "Answer the question using a single word or phrase."를 단답이 필요한 VQA 질문 끝에 붙인다. 핵심은 LLaVA는 LLM 전체를 학습하므로 이 지시를 실제로 학습해 따를 수 있다는 점이다(InstructBLIP의 구조적 한계를 정면으로 비껴간다). 부수 효과로 ① 형식을 자유롭게 조정 ② ChatGPT 후처리 불필요 ③ 다양한 데이터로 확장이 쉬워진다.

Table 1(b).  같은 "셔츠 색이 뭐죠?" 질문에 Normal·Ambiguous( Q:…A: ) 프롬프트는 "The man is wearing a yellow shirt."처럼 길게 답하지만, Formatting 프롬프트를 붙이면 정확히  "Yellow."  한 단어로 답한다. 2단계 튜닝 전 1단계만 거친 zero-shot 결과라 더 의미가 크다. 정량적으로 VQAv2를 넣는 것만으로 MME가  809.6 → 1323.8 로 뛰고 InstructBLIP을 111점 앞선다.

 

 

3.3. Scaling the Data and Model

형식 프롬프트가 한 일은 "학술 데이터를 섞어도 대화 능력이 망가지지 않게" 안전장치를 만든 것이다. 그래서 이제 마음 놓고 구성 요소를 하나씩 키운다.

MLP vision-language connector.

단일 선형층은 회전·스케일·이동만 하는 아핀 변환이라, CLIP과 LLM이라는 두 이질적 공간의 비선형 대응을 못 잡는다. 그래서 GELU 비선형을 가진 2-layer MLP로 바꾼다(H=W·Z  H=W₂·σ(W₁Z+b₁)+b₂, σ=GELU는 보편 근사기). 자기지도학습 SimCLR에서 projection head를 선형→비선형으로 바꿔 성능이 오른 것을 차용한 것이다.

Academic task oriented data.

LLaVA가 단답 학술셋에 약했던 건 그 분포가 학습에 없었기 때문이니, OKVQA·A-OKVQA·OCRVQA·TextCaps로 공백을 직접 메운다(A-OKVQA는 객관식+형식 프롬프트로 변환). 일부만 써도 InstructBLIP을 능가한다.

Additional scaling.

영역 데이터(Visual Genome·RefCOCO)로 세밀한 위치 인지를 강화하고, 입력 해상도를 336²로, 데이터에 GQA·ShareGPT를, 그리고 LLM을 13B로 키운다. 특히 MM-Vet은 13B에서 가장 크게 오른다 — 시각 대화의 품질 상한은 결국 기반 LLM의 추론·언어 능력이 정하기 때문이다.

LLaVA-1.5.

이 모든 것을 합친 최종 모델이 LLaVA-1.5다 — CLIP-ViT-L-336px → MLP 커넥터 → Vicuna v1.5 13B. Table 2는 개선을 하나씩 누적하며 GQA·MME·MM-Vet을 추적한다.

Table 2.  개선을 하나씩 누적(파랑=데이터, 분홍=모델, 주황=해상도). 원판 LLaVA(MME 809.6) → +형식 프롬프트(1323.8) → +MLP(1355.2) → +학술데이터 → 영역·336·GQA·ShareGPT →  +13B에서 GQA 63.3 / MME 1531.3 / MM-Vet 36.1 의 LLaVA-1.5. 가장 큰 점프 두 곳이  형식 프롬프트 와  13B 다.

 

Computational cost.

비용은 해상도 증가로 원판 LLaVA 대비 학습이 약 2배다 — 사전학습 약 6시간 + 튜닝 약 20시간(8×A100). 그럼에도 하루 남짓이면 끝난다.

3.4. Scaling to Higher Resolutions

해상도를 올릴수록 좋지만, 공개 CLIP은 336²가 상한이다. 기존처럼 ViT 위치 임베딩을 보간·재학습하려면 대규모 데이터가 필요하고 추론 해상도도 고정된다. 그래서 발상을 바꾼다 — 큰 이미지를 인코더가 원래 학습한 해상도의 작은 타일로 쪼개 각각 독립 인코딩하고 합친다. 인코더를 재학습 없이 그대로 쓰면서 임의 해상도를 처리할 수 있다. 타일 경계의 맥락 단절은 축소한 전체 이미지를 함께 넣어 global context로 보완한다.

 

 

 

 

4. Empirical Evaluation

4.1. Benchmarks

평가는 총 12개 벤치마크다.

  • 학술용 — VQA-v2·GQA(자유 단답 시각 인지), VizWiz(시각장애인 질문 zero-shot, 8,000장), ScienceQA-IMG(과학 객관식 zero-shot), TextVQA(글자 많은 이미지).
  • 지시수행용 — POPE(COCO의 random/common/adversarial로 환각을 F1로 측정), MME-Perception(예/아니오 인지), MMBench/MMBench-CN(보기 순서를 섞어 답변 견고성 측정, 중국어판 포함), SEED-Bench(이미지+비디오 객관식), LLaVA-Wild·MM-Vet(GPT-4가 채점하는 시각 대화).

4.2. Results

결과는 명확하다. LLaVA-1.5는 다른 방법 대비 수십~수천 배 적은 사전학습·튜닝 데이터를 쓰고도 12개 벤치마크에서 종합 최고이고, 지시수행 LMM 전부에서 원판 LLaVA를 크게 앞선다(원판 LLaVA는 VQA-v2처럼 단답을 요구하는 학술셋에서 평가 자체가 어렵다). 해상도를 448²로 키운 LLaVA-1.5-HD는 모든 벤치마크를, 특히 디테일이 필요한 과제(MM-Vet의 OCR, LLaVA-Wild의 상세 설명)에서 더 끌어올린다. 또한 global context를 더하면 split-and-merge 아티팩트가 회복되어 모델이 고해상도 특징에서 관련 영역을 더 쉽게 찾는다.

가장 고무적인 점은 가장 단순한 구조·학술 수준 장비·공개 데이터만으로 최고 성능을 냈다는 것이다. 이는 완전히 재현 가능하고 저렴한 baseline이며, "멀티모달엔 대규모 vision-language 정렬 사전학습이 필수"라는 통념에 의문을 던진다 — 비전 인코더(CLIP·OpenCLIP·EVA-CLIP 등)가 이미 web-scale로 사전학습돼 있기 때문이다. 실제로 7B 모델조차 연결에만 수십억 파라미터를 쓰는 80B IDEFICS(Flamingo 계열)를 능가한다.

Table 3.  VQAv2·GQA·VizWiz·SciQA-IMG·TextVQA 5개 학술 벤치마크 비교. LLaVA-1.5는 4개 1위·1개 2위이고, Qwen-VL(1.4B+50M)보다 수백~수천 배 적은 데이터(558K+665K)로 이를 달성한다. ( * =학습 중 해당 데이터를 봄,  † =비공개 자체 데이터 포함.)
Table 4.  POPE·MME·MMBench·SEED·LLaVA-Wild·MM-Vet 비교. LLaVA-1.5-13B가 대부분 최고이고, 환각 지표 POPE는  87.x 로 원판 LLaVA의 76.3을 크게 웃돈다. 원판 7B의 MME 809.6이 1.5-7B에서  1510.7 로 뛰고, 그 7B가 80B IDEFICS를 능가한다.

 

 

Global context.

고해상도를 위해 이미지를 패딩·축소해 단일 224² 이미지로 만든 뒤 고해상도 특징과 concat하여 전역 맥락을 제공한다. 7B 모델 ablation에서 이 전역 맥락은 세 벤치마크 모두를 확실히 끌어올린다.

Global context ablation (7B).  패치만 쓸 때(high-res patch only) 대비 전역 맥락을 더하면  GQA +0.9, MME +71, MM-Vet +3.2  — 분할로 생긴 이음새(artifact)를 전체 이미지가 보정한다.

 

 

4.3. Emerging Properties

Format instruction generalization.

학습한 형식은 소수인데 처음 보는 형식도 따른다. VizWiz의 'Unanswerable' 출력이 11.1% → 67.8%로 오른다.

 

Table 5. 사막이 없는 해변 사진에 "이 사막에서 무슨 일이?"라고 함정을 거는데, GPT-4V는 "어느 사막인지 모르겠다", 구판 LLaVA는 "사막에 있다"고 틀리는 반면, LLaVA-1.5는 "이미지에 사막은 없다"고 전제를 바로잡는다.
Table 6. 운전면허증 이미지를 주고 지정된 JSON 스키마(class·DLN·DOB·Name…)로 추출하게 하면, 일부 필드 오차는 있어도 요구된 구조에 맞춰 정확히 채운다(GPT-4V 대비 일부 오차).

 

 

Multilingual multimodal capability.

모든 시각 지시와 VQA가 영어인데도 사용자가 다른 언어로 물으면 그 언어로 답한다. ShareGPT의 다국어 텍스트에서 "사용자 언어에 맞춰 답하는 습관"을 배웠고 그게 시각 대화로 전이된 것이다. 정량적으로 MMBench-CN에서 LLaVA-1.5가 63.6%로, 중국어로 직접 학습한 Qwen-VL-Chat의 56.7%보다 +7.3% 높다.

Figure 5.  영어로만 학습했지만 다른 언어로 물으면 그 언어로 답하는 다국어 창발 능력.

 

4.4. Ablation on LLM Choices

기반 LLM을 바꿔 비교한다 — LLaMA-1 계열(Vicuna-v1.1·v1.3) vs LLaMA-2 계열(Vicuna-v1.5·LLaMA-2-Chat). v1.3과 v1.5는 동일한 약 150K ShareGPT를 쓰고, Vicuna는 SFT만, LLaMA-2-Chat은 거기에 RLHF까지 받았다.

Figure 3.  Vicuna-v1.5가 전 축에서 가장 바깥, 특히 MMBench-CN에서 격차가 크다 —  같은 데이터로 튜닝해도 더 좋은 기반 LLM이 더 멀리 간다.

 

 

 

5. Open Problems in LMMs

5.1. Data Efficiency

학습 데이터를 10~100%로 줄여도 50%면 성능의 98% 이상이 유지되고 30%까지 안정적이라, 멀티모달에도 'less-is-more'가 성립한다. 특히 50%에서 MMBench·ScienceQA·POPE는 손실이 전혀 없고 MMBench는 오히려 오른다.

Figure 4.  학습 데이터를 10~100%로 줄여가며 측정 — 50%만 써도 대부분 성능이 유지된다.

 

5.2. Rethinking Hallucination in LMMs

환각 재고가 가장 통찰적이다. 흔히 "데이터에 오류가 있어서"로 봤지만, 해상도를 448²로 키우면 환각이 크게 준다 → 즉 LMM은 데이터의 소수 오류엔 의외로 견고하다. 다만 핵심 균형이 있다 — 해상도가 디테일을 분간하기엔 부족한데 모델 능력을 넘는 세밀한 데이터(granularity)만 늘리면, 모델은 오히려 환각을 '학습'한다. "주석을 얼마나 세밀하게 하느냐"와 "모델이 그 세밀함을 처리할 수 있느냐" 사이의 균형이 필요하다.

5.3. Compositional Capabilities

따로 배운 능력을 합동 학습 없이 조합 과제로 일반화한다(예: 학술 데이터의 시각 지식이 대화 답변의 'visual groundness', 즉 이미지 충실도를 높임). 단, VQA에서 속성을 맞혀도 상세 설명에서 정확히 묘사한다는 보장은 없고, 한국어 같은 일부 언어는 아직 약하다.

 

 

6. Conclusion

정리하면, 단순·효과적·재현 가능한 baseline LLaVA-1.5를 제안했고, 고해상도·환각·조합 능력에 관한 발견을 남겼다. 결국 메시지는 하나다 — "성능을 위해 반드시 복잡한 구조나 거대한 비공개 데이터가 필요한 건 아니다." 단순함을 유지한 채 커넥터·프롬프트·데이터·해상도·기반 LLM만 적절히 조정하면, 공개 자원만으로도 강력하고 재현 가능한 멀티모달 모델을 만들 수 있다.

Limitations.

한계는 ① 전체 이미지 패치를 쓰기에 고해상도 학습 비용이 크고, ② 다중 이미지(multiple-image) 이해가 부족하며, ③ 환각이 잔존하고, ④ 의료 같은 critical 응용에서는 주의가 필요하다는 점이다.

 

 

Appendix

부록은 본문 주장을 실제 설정과 사례로 뒷받침한다.

A. 구현 세부

LLaVA-1.5-HD는 CLIP-ViT-L-14의 224²를 기준 타일로 삼아 패딩·분할·독립 인코딩 후 전역 맥락을 더한다. 총 665K 데이터에 7가지 효율화 기법(예: 한 배치를 단일 모달리티로만 구성해 학습을 25% 가속)을 쓴다.

Table 7.  총 665K 데이터 구성표. VQA·OCR·영역·대화 데이터마다  서로 다른 형식 프롬프트 를 부여한 게 보인다(단답 VQA엔 "single word or phrase", 객관식엔 "option's letter", 영역엔 bounding box 형식).
Table 8.  평가용 형식 프롬프트. 벤치마크마다(단답·객관식·VizWiz의 'Unanswerable') 출력 형식을 통일해  공정 비교 를 보장한다.
Table 9.  하이퍼파라미터. MLP 사용 때문에 사전학습 learning rate를 절반으로 낮추고, AdamW·cosine·epoch 1·DeepSpeed를 쓴다. 평가는 재현성을 위해 greedy decoding을 사용한다.

 

 

 

B. 정성 예시

Table 10.  같은 대화 안에서 "한 단어로 → 자세히"로 요청을 바꾸면 답의 상세도를 자유자재로 조절한다(다림질 사진을 단답·중간·장문으로).
Table 11.  시각 글쓰기 — 하와이 사진에 "여행 블로그를 써 달라"고 하면, 구판 LLaVA보다 훨씬 길고 와이키키·진주만 등  이미지에 충실한  글을 쓴다(ShareGPT 추가 효과).
Table 12.  Stable Diffusion 프롬프트 생성 —  [형용사][주체][헤어][의상]…  지정 양식에 맞춰 "단발머리 소녀가 흰 기모노를 입고…" 식 프롬프트를 정확히 채운다.

 

C. 추가 한계

전체 패치를 쓰기에 학습이 길고, 리샘플러는 패치를 줄이지만 아직 LLaVA만큼 샘플 효율적으로 수렴하지 못한다 — 샘플 효율적 리샘플러 개발이 향후 과제다.

 

+ Recent posts