[논문 리뷰] Dreamer v4: Training Agents Inside of Scalable World Models

https://danijar.com/project/dreamer4/

Training Agents Inside of Scalable World Models

Google Deepmind에서 발표한 월드 모델 agent 시리즈인 dreamer, 25년 9월에 버전 4가 나왔다. 
Full offline train만으로 다이아몬드를 채굴한 첫 에이전트라고 한다!

이전까지의 recurrent based dreamer 시리즈와는 약간 다른 구조(Diffusion + Transformer)를 가지며, 핵심 내용 요약은 다음과 같다.

 

 

요약

해결하려는 문제

• 기존 월드 모델은 복잡한 환경에서의 Objecct interaction과 게임 매커니즘의 물리학을 정확하게 예측하지 못함.
• diffusion model은 정확하지만 느림. 실시간 활용 불가능
• long rollout에서 에러가 누적됨. 길게 Imagination하기 어려움.

Main contribution

1. 정확한 world model 안에서 RL을 수행하는 구조를 실제로 성공시킴
   • 실제 게임의 mechanics까지 맞추는 world model을 구축
     → Purely offline으로 Minecraft diamond task를 해결
2. Shortcut forcing으로 diffusion world model을 빠르고 안정적으로 만듦
   • 적은 step (K=4)으로도 high-quality generation
   • real-time inference 가능
   • long horizon rollout 안정화
     → “빠르면서 정확한 world model”을 구현
3. Transformer 기반 video world model로 확장
   • video token + transformer
   • diffusion dynamics
     → 훨씬 큰 capacity & 복잡한 환경 (Minecraft) 모델링 가능
4. 적은 action + 많은 video로 학습 가능함을 입증
   • 대부분의 지식은 unlabeled video에서 학습
   • action은 일부만 있어도 충분
     → world model이 먼저 world를 이해하고, action은 나중에 grounding
5. 완전 offline setting에서 long-horizon task 해결
   • environment interaction 없음
   • dataset만 사용
     → robotics 등 현실 문제에 직접 연결

 

 

먼저 월드 모델과 reinforcement learning의 관계에 대해 설명하고, 이전 시리즈인 Dreamer v3와 비교하면서 Diffusion based method인 Dreamer v4를 설명하도록 하겠다.

---

Reinforcement Learning

행동을 통해 보상을 최대화하기.

1. Markov Decision Process (MDP)

구성 요소

• agent (행동 주체), environment (변화하는 환경)
• state $s_t$ (시간 t에서의 상황), action $a_t$(시간 t에서의 행동), reward $r_t$ (시간 t에서의 state와 action으로 발생한 보상)
• discount factor $\gamma$ (누적 보상에 곱하는 값. 1에 가까우면 미래가, 0은 현재가 중요하다)
• transition probability P (state에서 action이 주어졌을 때 다른 state로 넘어갈 확률)
• policy $\pi(a|s)=P(a_t|s_t)$ (state에서 특정 action을 선택할 확률.)
  state $s_t$에서 agent가 $a_t$를 취했을 때, environment의 영향을 받아 다음 state $s_{t+1}$이 되었을 때 보상 $r_t$를 받는다. 이때 환경은 $P(s_{t+1}|s_t,a_t)$의 확률로 움직이는데, 이는 현재 상태와 행동이 주어지면 다음 상태가 확률적으로 결정됨을 말한다.

Trajectory

이제까지 어떤 state와 action을 거쳤는가?
$τ=(s_0,a_0,s_1,a_1,…)$
특정 trajectory가 발생할 확률은 다음과 같다. 여기서 policy만 θ에 의존한다.
$P(\tau|\theta) = p(s_0) \prod_t \pi_\theta(a_t|s_t) P(s_{t+1}|s_t,a_t)$

Reward

누적 보상
$G_t(\tau) = r_t + γ r_{t+1} + γ² r_{t+2} + ...$

Value function

현재 상태가 얼마나 좋은가? 즉 이 상태에서 시작했을 때, 앞으로 받을 보상의 기댓값이다.
$V(s) = E[ G_t | s_t = s ]$

Markov property

State $s_0, ..., s_{t-1}, s_t$에서 $s_{t+1}$로 갈 확률은, 바로 이전 state $s_{t}$에서 $s_{t+1}$로 갈 확률과 동일하다는 성질이다. 이는 memoryless property이며 $s_0, ..., s_{t-1}, s_t$의 정보가 이미 $s_{t}$에 담겨 있기 때문에 가능하다. Action도 비슷하게 적용됨
$P(s_{t+1}| s_0, ..., s_{t-1}, s_t) = P(s_{t+1} | s_t)$

목표 & Train policy

Policy $\pi_\theta(a|s)$ 를 학습하는 것이 목표이다. θ는 policy network parameter
→ Objective: 누적 보상의 기댓값을 최대화하기.
maximize $J(\theta) = E[G] = E[\sum_{t=0}^∞ \gamma^t r_t] = \sum_\tau P(\tau|\theta) R(\tau)$
Timestep별 return을 사용해 표현한 Policy gradient.
$∇_θ J = E[ ∇_θ \log π_θ(a_t|s_t) G_t ]$

J를 maximize하는 θ를 찾아야 하므로 gradient 방향으로 움직이는 gradient ascent 수행해서 θ를 업데이트해서 policy를 학습한다. 여기서 $∇_θ \log π_θ(a_t|s_t)$는 특정 행동의 확률을 증가시키는 방향, G는 그 행동을 하고 얻은 return이다. 즉 얼마나 좋은 행동인지에 따라 그 확률을 증가시키는 정도가 커진다

• 좋은 행동 a, G ↑ → 확률 증가, $π(a|s)$ ↑
• 나쁜 행동 a, G ↓ → 확률 감소, $π(a|s)$ ↓

예시

• 로봇) state: 카메라 이미지, action: 팔 움직임, reward: 물체 집으면 +1
• 게임) state: 화면, action: 키보드 입력, reward: 점수

 

2. Model-Based RL vs Model-Free RL

1) Model-free RL

환경을 직접 경험하면서 policy을 학습한다. 환경에 대해 모델링을 따로 하지 않고 환경의 dynamics($P(s_{t+1} | s_t, a_t)$)도 배우지 않고, 단순이 이 행동이 좋았는지/나빴는지만 학습한다.

• 예시) Atari 게임: agent는 "화면 → 행동"을 반복하면서 어떤 행동이 점수를 올리는지만 학습하고, 게임의 물리 법칙을 배우지 않는다.

2) Model-based RL

환경 모델($P(s_{t+1} | s_t, a_t)$)을 먼저 학습한다. 환경을 실제로 실행하지 않아도 미래를 예측할 수 있게 된다. state $s_t$ 를 world model에 넣으면 다음 state $s_{t+1}$을 예측한다.

차이점

 

• Model-Free: environment → data → policy learning 환경을 계속 실행해야 한다. 문제는 sampling이 비효율적이라는 것이다. 로봇으로 치면 실험을 1억 번 해야 할 수도…
• Model-based: (environment → data) → world model -> imagination rollout → policy learning으로 간소화할 수 있다. 즉 environment → world model로 world model을 학습시키고, World Model의 imagination을 통해 current state → world model → fake data(future trajectory)를 만들어내서 agent를 학습시킬 수 있다.
  ⇨ Dreamer 계열의 핵심!

 

---

World Model

World Model 은 다양한 입력 데이터 를 받아서 물리학, 공간 역학, 인과 관계 (최소한 기본적 인과 관계) 등 실제 환경의 Internal Representation을 학습하는 generative AI system.
이렇게 학습된 표현을 사용해서 미래의 상태를 예측하고, 내부적으로 연속적으로 일어나는 행동을 시뮬레이션하고, 실제 환경에서의 실험을 하지 않고도 정교하게 계획하고 의사 결정을 지원할 수 있다.
AI 에이전트가 실세계에서 직접 시행착오를 겪는 방식으로 학습하는 대신 ‘월드 모델’ -환경에 대해서 학습된 시뮬레이션- 을 사용한다면, 이후 일어날 수 있는 일들과 할 수 있는 연속적인 행동과 사건의 순서를 상상하고 탐색한다. 이런 활동을 내부적으로 시뮬레이션한다면, AI가 원하는 목표, 원하는 결과를 가져오는 경로, 궤적을 찾아낼 수도 있다.

• 장점: 가능한 모든 행동을 실제로 물리적으로 해 보지 않아도’ 되도록 해 주기 때문에, 필요한 자원의 규모를 크게 줄여준다.

장점

1. Sample efficiency
   환경 interaction을 줄일 수 있다. 예를 들어 로봇을 학습할 때 실제 실험 없이 학습이 가능하다.
2. Offline learning
   dataset만으로 학습할 수 있다. 예를 들어, Dreamerv4에서는 Minecraft dataset만 사용하고 환경 interaction이 없다.
3. Planning 가능
   미래를 예측할 수 있다

유형 

credit: https://turingpost.co.kr/p/topic-35-world-models

• Implicit Models
  세계가 어떻게 변화하는지, 그리고 어떻게 관찰되는지 두 가지를 분리하지 않고, 미래의 결과를 예측하기 위해서 하나의 큰 신경망을 -마치 블랙박스처럼- 사용한다. 이런 프레임웍은 에이전트가 압축된 이미지와 예측을 사용해서 미래의 행동을 "꿈"처럼 그릴 수 있게 한다.
  예시) 많은 바둑 대국을 보고 다음 수를 패턴으로 예측하는 모델
• Explicit Models
  세계가 어떻게 변화하는지 (상태의 전이)와 에이전트가 무엇을 보고 있는지 (관찰 내용)를 명확하게 분리한다. 이런 구조는 시스템의 해석 가능성, 그리고 디버깅 용이성이 높아진다.
  예시) 자동차를 운전할 때, 물리 법칙(가속과 브레이크가 차량 속도에 미치는 영향)과 관찰(백미러와 사이드미러로 보이는 것)을 분리해서 이해하는 것.
• Simulator-based Models
  처음부터 학습을 진행하는 대신, 시뮬레이터나 실제 환경을 사용해서 행동과 결과를 테스트한다. 매우 정확한 방법이라고 할 수 있지만, 대신 진행이 느리고 비용이 많이 들 수 있다.
  예시) 비행기 조종사가 실제 비행 전에 비행 시뮬레이터에서 다양한 시나리오를 연습하는 것
• Hybrid and Instruction-driven Models
  학습된 모델과 외부의 규칙, 매뉴얼 또는 언어 모델을 결합. 신경망이 예측하는 것과 인간이 설계한 Rule을 조합하는 이 방식은 이전에 보지 못한 상황에서 모델을 더 유연하게 만들어 준다.
  예시) GPS 내비게이션이 도로 규칙(속도 제한, 일방통행)과 실시간 학습(교통 패턴)을 모두 사용해서 최적의 경로를 제안하는 것과 유사.

 

즉 월드 모델로 환경을 시뮬레이션하고, 
그 시뮬레이션한 환경 안에서 강화학습을 하면서 에이전트/로봇/모델 등을 학습시키면 데이터 효율성이 높아진다.

 

---

Diffusion

generative model의 대표격인 diffusion 모델은 높은 퀄리티로 이미지를 생성한다.

학습 방식: Flow Matching

Deep Unsupervised Learning using Nonequilibrium Thermodynamics, 2015
Denoising Diffusion Probabilistic Models, 2020
노이즈에서 실제 데이터로 가는 방향(vector field)를 학습하는 것

목표

• Train network $f_\theta (x_\tau, \tau)$ to restore data point $x_1$ given corrupted version $x_\tau$
  • Data $x_1 \sim \mathcal{D}$
  • Noise $x_0 \sim N(0, \mathbb{I})$ 
  • Corrupted data $x_\tau =(1-\tau)x_0 + \tau x_1$
    → data와 noise를 섞은 상태
  • Signal Level $\tau \in [0,1]$
    → $\tau=0$ means pure noise
• 모델은 noise to data의 velocity 벡터, $v=x_1 - x_0$을 예측하도록 train된다.

Loss

$$\mathcal{L}(\theta) = \left\| f_\theta(x_\tau, \tau) - (x_1 - x_0) \right\|^2$$

Inference time

Pure noise vector $x_0$에서 시작해서 K번을 반복해 clean data point $x_1$로 변화시킨다.
$x_{\tau+d} = x_{\tau} + f_{\theta}(x_{\tau}, \tau)d \quad x_0 \sim \text{N}(0, \mathbb{I})$

• step size d = 1/K

 

그러나 diffusion을 바로 사용하기에는 여러가지 문제가 있다. Dreamer v4는 여러 논문에서 제안된 해결법을 차용해서, 이미지 시퀀스를 다룰 수 있는 빠르면서도 long rollout에서 안정적으로 동작하는 모델로 만들었다.

1. 학습 속도: Shortcut Models

One Step Diffusion via Shortcut Models, 2024

문제

일반 diffusion은 한 번에 조금씩만 이동하기 때문에 큰 step (K=64)가 필요하고, 느림

 → 한 번에 크게 step을 밟을 수 있도록 학습하자! 그런데 큰 스텝은 비선형성을 예측하기 어려우니, 반반 나누어 이동해 보자.

방식

Network $f_\theta$가 step size d도 입력받도록 바꾼다.

$$\begin{align*} x_0 &\sim \mathcal{N}(0, I) & x_1 &\sim \mathcal{D} & \tau, d &\sim p(\tau, d) \\ \end{align*}$$

1. d가 기존에 설정해둔 최소 step size d_{min}보다 작을 때
   이전처럼 noise to data의 velocity 벡터, $v=x_1 - x_0$을 바로 예측
2. d가 d_{min}보다 크고 최대 크기인 1보다 작을 때
   Bootstrap Loss를 사용한다. 큰 step을 두 개의 작은 step으로 분해하고 각 step에서 예측된 결과를 합쳐서 사용하는 방식이다.
   \begin{align*} b' &= f_\theta(x_\tau, \tau, \frac{d}{2}) & b'' &= f_\theta(x', \tau + \frac{d}{2}, \frac{d}{2}) & x' &= x_\tau + b'\frac{d}{2} \end{align*}
   • b': 모델 f_θ가 x_{\tau}에서 step size인 d/2만큼 이동했을 때의 예측 velocity 벡터.
   • b'': 시작점 x_{\tau}에서 b'방향으로 이동한 중간 위치 x'에서 d/2만큼 더 이동했을 때의 예측 velocity 벡터.
     Loss는 모델의 예측값과 목표값의 차이의 L2 norm.
     \begin{align*} \mathcal{L}(\theta) &= \|f_\theta(x_\tau, \tau, d) - v_{\text{target}}\|^2 & v_{\text{target}} &= \begin{cases} x_1 - x_0 & \text{if } d = d_{\min} \\ \frac{\text{sg}(b' + b'')}{2} & \text{else} \end{cases} \end{align*}

• $\text{sg}$: Stop gradient. Sg로 설정된 변수는 역전파 시 상수로 취급되어서 이 변수의 업데이트가 멈춘다. Distillation에서 자주 사용됨. 여기서 b', b'' 은 다른 입력을 넣은 f_θ의 결과물이므로, 현재 학습 중인 f_θ의 파라미터와 동시에 업데이트되면 복잡하고 불안정한 그래디언트 문제가 발생할 수 있다. 따라서 b', b''의 f_θ로 다시 역전파되지 않도록 멈춰 준다.

(참고) GPT의 stop gradient 설명

Dreamerv4에서 중요한 이유

Dreamerv4에서 world model은 rollout을 많이 해야 한다.

이 때 rollout이 느리면 RL 자체가 불가능하므로 빠르게 step을 밟아 학습하는 것이 중요함!

 

 

2. Diffusion Model Output: X-prediction vs V-prediction

고치원일수록 v-prediction에서는 망가지지만 x-prediction은 잘 보존됨

1) 문제: V-prediction

현재 상태에서 깨끗한 데이터로 가는 방향(velocity)을 예측

정의

• 모델이 예측하는 것: $v = x_1 - x_0$
  학습: $f_\theta(x_\tau, \tau) \approx (x_1 - x_0)$

특징

• diffusion에서 가장 흔한 방식
• 학습 안정적
• score matching과 유사한 성질

단점 (Dreamer 4에서 중요)

• high-frequency 신호를 학습함
  → 작은 오차가 누적되어 긴 rollout에서 붕괴

2) x-prediction

노이즈 상태에서 바로 clean data를 예측

정의

• 모델이 예측하는 것: $\hat{x}_1 = f_\theta(x_\tau, \tau)$

특징

• target이 더 직접적
• structured output (이미지/latent 자체)

장점

• 더 안정적
• long rollout(long horizon)에서 error 누적 감소
• world model에 적합

 

3. Sequence 처리: Diffusion forcing

Diffusion forcing: Next-token prediction meets full-sequence diffusion, 2024

문제

일반 diffusion은 single image를 대상으로 하지만, dreamerv4는 sequence를 다룬다. (z_1, z_2, ...)

해결 방법

각 time step마다 서로 다른 noise level ($\tau$)을 할당하여 전체 시퀀스를 corrupt시킨다.

장점

• 학습 안정성: 모든 time step에서 loss 계산 가능 → 각 time step에서 동시에 디노이징
• Conditioning: 앞 frame이 context로 사용된다. 즉 autoregressive + diffusion
• Inference 시 유연성: 모델이 이전 time step의 프레임들을 context로 삼아서 autoregressive하게 다음 프레임을 diffusion할 수 있게 된다. 이는 causal masking 기법을 통해 훈련되며, inference time에 이전 프레임들이 깨끗하든(실제 데이터의 프레임) 노이즈가 있든(모델이 추론한 프레임) 이를 받아들여서 다음 프레임을 디노이징하는 데 사용할 수 있도록 robustness를 제공한다.
  • Causal masking: 자연어 처리에서 미래의 단어를 보고 현재의 단어를 추론하면 안 되기 때문에, 나의 이전 단어들만 볼 수 있도록 미래의 단어들을 masking.
    → Dreamerv4에서는 block causal masking으로 구현되었으며, 이는 현재 time step 내에서는 자유롭게 어텐션을 할 수 있지만 서로 다른 time step에 대해서는 오직 이전 time step에만 어텐션할 수 있도록 마스킹하는 방법이다.

 

---

 

Dreamer v4의 구조는 Dreamer v3과 거의 동일하다. 
가장 큰 차이점: World Model의 형태가 Diffusion으로 완전히 바뀜! imagination training objective도 바뀌었다.

구조

 

 

(Phase 1) World Model Pretraining 

Tokenizer와 world model(dynamics)을 video + action으로 pretrain.

(Phase 2) Agent Finetuning

World model의 policy와 reward head를 finetune한다.

(Phase 3) Imagination Training

Imagined trajectory로 policy head를 optimize한다.

 

 

 

 

왜 월드 모델로 기존에 사용하던 RSSM 대신에 Diffusion + transformer를 선택했을까?

Why diffusion?

Diffusion의 핵심 아이디어는 다음 상태를 확률적으로 생성한다는 것

• 장점1: Multi-modality + Uncertainty
  • 현실에서는 같은 상태에서 같은 행동을 해도 다른 결과가 나올 수 있다.
    • e.g., 블록을 부수면 particle이 서로 다르게 튀김, 공은 다양한 방향으로 튐
  • Dreamer v3의 RSSM 구조는 categorical한 simple한 분포이므로 복잡한 multi-modal distribution을 표현하기에는 제한이 있음.
  • Diffusion은 노이즈에서 시작해서 매번 샘플링하면서 디노이징하므로, 샘플링할 때마다 다른 결과가 나올 수 있다. 즉 distribution 자체를 모델링하므로 여러 결과를 확률적으로 생성할 수 있으므로 현실의 다양한 분포와 복잡한 상호작용을 더 잘 모델링할 수 있음.
• 장점2: long-horizon 안정성
  • RSSM(RNN)의 본질적 한계는 오래된 정보를 기억하기 어렵다는 점.
  • Dremaerv4는 diffusion에서 오차가 누적될 수 있는 v-prediction 대신 x-prediction을 사용했고, transformer 구조를 사용하여 과거 시점의 문맥을 직접 참조하도록 해서 장기 기억 문제를 극복해냄. Dreamer 4는 기존 모델(0.8~1.6초)에 비해 9.6초라는 훨씬 긴 context length를 가진다고 함.
• 장점3: video-level modeling
  • pixel/video를 직접 모델링하므로 더 rich한 representation
  • 기존에는 다음 state(latent)를 모델링했고 이를 바탕으로 latent space에서의 시나리오를 상상했음. dreamerv4는 실제 이미지/비디오 상에서 시나리오를 직접 상상함
  • 다음 사진을 보면 3에 비해 4가 고해상도로 이미지 상황을 직접 예측하므로, 훨씬 복잡한 상호작용을 잘 모델링함을 알 수 있음.

 

 

Phase 1: Train World Model

세계를 모델링하고, 태스크 기반 에이전트의 행동을 그 세계 안에서 이해할 수 있도록 만드는 모델.

Causal tokenizer + Transformer dynamics 두 개의 모듈로 구성됨

1) Causal Tokenizer

Pixel을 transformer에 친화적인 latent representation(token)으로 인코딩한다. 이 token은 World-model dynamics가 다음 token을 생성하기 위한 입력으로 사용됨.

Input

각 timestep에서, image는 patch token과 latent token으로 구성된다.

• Patch token: 이미지를 쪼갠 것
• Latent token: learnable token으로, 전체 정보를 담는 slot이다.

Encoder

Causal transformer encoder. Latent token과 patch token이 합쳐져 self attention을 하면서, patch → latent로 정보가 요약됨. 이후 latent token만 다음 bottleneck으로 전달한다.

• Causal: 이때 현재 frame은 미래 frame을 보면 안 되므로 미래 frame은 masking하고 전체 timestep에 대해 (patch + latent)의 self-attention 을 수행한다.
  • 효율성: Dreamer 4는 patch/latent token이 시간축까지 포함해 causal하게 상호작용할 수 있도록 설계되었지만, 실제로는 시간이 길어질수록 full attention의 비용이 너무 크기 때문에 space-only / time-only attention으로 분해한 효율적인 2D transformer를 사용한다.
• Modality별 Attention 범위: Latent는 모든 input modality(image, text …)와 어텐션을 할 수 있다. Image와 같은 개별 modality는 그 안에서만 어텐션할 수 있다. 이를 통해 latent는 여러 모달리티의 정보를 통합적으로 압축 요약할 수 있게 됨

Bottleneck (Latent Tokens)

Latent tokens만 읽어서 저차원으로 projection하여 압축한다.

Decoder

Causal transformer decoder. 이 압축된 latent를 보고 image를 복원한다.

• Modality별 Attention 범위: 인코더와 유사하지만 약간 다르다. Image와 같은 개별 modality는 자기 자신 + latent와 어텐션을 할 수 있다. Latent는 그 안에서만 어텐션할 수 있다. 이를 통해 latent는 자기 자신의 정보를 유지하고, 각 모달리티는 latent의 정보를 참고하여 복원할 수 있게 된다.

Train: Masked Autoencoding

MAE- Masked Autoencoders are Scaleable Vision Learners 참고
Patch를 dropout하고 복원하는 task로 학습한다. p ~ U(0, 0.9)
$$\mathcal{L}(\theta) = \mathcal{L}_{\text{MSE}}(\theta) + 0.2 \mathcal{L}_{\text{LPIPS}}(\theta)$$

• MSE Loss: 픽셀 단위 차이
• LPIP Loss: Feature level에서의 차이. 즉 perceptual한 high-level에서의 차이

 

2) Interactive Dynamics

Diffusion + transformer를 이용해 다음 상태(영상) 를 생성하는 world model

기존의 dreamer v3는 현재 state를 기반으로 미래 latent 하나를 예측했지만, dreamerv4는 noisy video를 기반으로 clean video latents 전체 (token sequence)를 생성한다.

Input

1. noise level τ
2. noisy latent tokens (현재 + 미래) $\tilde{z} = (1-\tau)z_0 +\tau z_1$
   • $z_0$: noise, $z_1$: clean latent tokens
3. action tokens a

Model Task

Clean latent $\hat{z}$를 출력하는 것이 목표이다. → X-prediction
즉 노이즈가 있는 latent $\tilde{z}$ 에서 실제 미래 $\hat{z}$를 복원한다.
$$f_\theta(\tilde{z}, \tau, a, d) = \hat{z}$$

• 이때 action을 conditioning으로 넣어 z를 생성하므로 interactive dynamics라고 한다.

X-prediction: 프레임의 픽셀을 직접 예측하면 긴 시퀀스에서 안정적이다.

• V-prediction을 사용하면 noise to clean의 velocity vector $v=x_1-x_0$을 예측한다. 이는 빠르고 미세한 디테일이 많은 high-frequency인 경우가 많다. Video sequence를 autoregressive하게 생성하려고 할 때 시간이 길어질수록 오차가 누적되면 잘못된 방향으로 drift될 가능성이 높다. 즉 t-1에서 t로 가는 속도에 ϵ의 오차가 있었다면, t 위치= t-1 위치 + 정상 속도 + ϵ이 된다. 이 ϵ는 사라지지 않고, t+1을 예측할 때에도 누적되는 경향이 있다고 한다.
  $f_\theta(\tilde{z}, \tau)=z_1 - z_0$
• X-prediction은 corrupted data $x_\tau =(1-\tau)x_0 + \tau x_1$가 주어졌을 때 clean data $x_1$ 자체를 바로 예측하는 방법이다. 내가 t시점의 프레임을 약간 잘못 예측했더라도, t+1 프레임을 예측할 때는 내가 이제까지 이해한 1~t-1까지의 컨텍스트를 기반으로 어느 정도의 correction이 가능하다.
  $f_\theta(\tilde{z}, \tau)=\hat{z}_1$

Diffusion Forcing

각 timestep마다 (logit-normal 분포에서 샘플링해서 얻은) 서로 다른 $\tau$를 적용한다. 내가 디노이징한 clean한 과거 context(이전 프레임)를 이용해 미래를 예측할 수 있게 한다.

• Train time: Causal masking으로 과거의 프레임을 어텐션하도록 하면, 현재 타임스텝에서 손상된 $\tilde{z}$를 복원하기 위해 과거를 어텐션하면서 정보를 수집함.
  e.g., 과거 프레임 중 $\tau$가 높은 프레임은 원본 프레임과 거의 동일하므로 많이 참고할 만함.
• Inference time: 트랜스포머 모델은 전체 시퀀스(e.g., 192 frame)를 입력으로 한 번에 받고, 출력도 한 번에 전체 시퀀스를 낸다. 이때 각 타임스텝마다 spatial token, latent token, action token, agent token 등 여러 토큰이 있는데, 이 모든 토큰을 concat해서 한 개의 토큰 블록으로 만들고, 전체 192개 프레임의 192개 토큰 블록을 모두 합쳐서 입력으로 넣는다. 출력은 Causal masking때문에 나의 이전 프레임만 참조할 수 있으므로 autoregressive하게 순차적으로 생성하는 것처럼 보인다. 이 거대한 input 토큰 블록들끼리 self-attention하면서 각 타임스텝의 프레임을 생성한다. 이때 과거 정보만 볼 수 있도록 하며(Causal attention), 과거의 정보에는 $\tau_{ctx}=0.1$ 의 노이즈를 추가해 robust하게 만든다.
  • Train time에 각 타임스텝별로 다른 $\tau$를 넣으므로, 모델은 다양한 수준의 노이즈 레벨을 경험하면서 디노이징하는 방법을 배운다. 따라서 inference time에도 내가 이전에 생성한 프레임에 0.1의 노이즈를 넣어 주면, 훈련할 때 보았던 noisy image 상황을 다시 경험하므로 모델이 더욱 강건하고 안정적인 예측을 보인다. 실제로 내가 만든 프레임이 완전히 clean하거나 약간 noisy하거나의 차이를 줄여 주기도 한다.

Shortcut model

K=4의 매우 작은 step의 forward pass만으로 프레임을 생성한다.
$z_0 \sim \text{N}(0, 1) \quad z_1 \sim \mathcal{D} \quad \tau, d \sim p(\tau, d) \quad \tau, d \in [0, 1]^T$
$\hat{z}_1 = f_\theta(\tilde{z}, \tau, d, a) \quad \tilde{z} = (1 - \tau) z_0 + \tau z_1$

• 모델 f_θ는 corrupted latent $\tilde{z}$를 입력받아 clean latent \hat{z}를 예측한다.

$b' = \frac{\left(f_\theta\left(\tilde{z}, \tau, \frac{d}{2}, a\right) - z_\tau\right)}{1 - \tau} \qquad z' = \tilde{z} + b' \frac{d}{2}$

• $b'=f_θ​(\tilde{z},τ,d/2,a)$는 $\tilde{z}$에서 시작하여 τ 시그널 레벨에서 스텝 크기 d/2만큼 진행했을 때 도달해야 할 클린한 $\tilde{z}_1$을 예측하는 것이다. $v=(z_1 - z_\tau)/(1-\tau)$ 관계가 성립하므로, b'은 corrupted $z_\tau$에서 clean $z_1$로 d/2만큼의 step으로 가는 방향 벡터이다.
• z': corrupted $\tilde{z}$에서 b' 속도로 움직인 중간 지점

$b'' = \frac{\left(f_\theta\left(z', \tau + \frac{d}{2}, \frac{d}{2}, a\right) - z'\right)}{1 - \left(\tau + \frac{d}{2}\right)}$

• b'': 중간 지점 z'에서 시작해서 τ+d/2의 시그널 레벨(기존보다 더 clean)에서 스텝 크기 d/2만큼 진행했을 때 도달할 클린한 $\tilde{z}$를 예측하고, 그 방향 벡터를 b''이라고 한다. 즉 더 클린한 상태에서 예측된 속도 벡터임.

$\begin{align*} \mathcal{L}(\theta) &= \begin{cases} \left\|\hat{z}_1 - z_1\right\|_2^2 & \text{if } d = d_{\min} \\ (1 - \tau)^2\left\|\frac{\hat{z}_1 - \tilde{z}}{1 - \tau} - \frac{\text{sg}(b_1 + b_2)}{2}\right\|_2^2 & \text{else} \end{cases} \end{align*}$

• Loss
  • 가장 작은 step 크기일 때(e.g., 1/K_max), clean과 corrupted latent의 MSE 에러를 직접 최소화한다.
  • 조금 더 큰 step 크기일 때, 두 개의 작은 스텝으로 나누어 예측한 b', b''의 bootstrap loss를 사용한다. 이때 $(\hat{z}_1 - \tilde{z})/(1-\tau)$는 큰 스텝 d에 대한 예측된 속도 벡터이고, 이를 $sg(b_1+b_2)/2$와 비교한 차이를 loss로 사용한다. $sg(b_1+b_2)/2$는 작은 스텝으로 움직인 목표 속도 벡터임. 여기에 $(1-\tau)^2$로 scale을 맞춰 주어 v-space → x-space로 바꿔줄 수 있다.
    → 한 번의 큰 스텝은 그 비선형성 등을 한 번에 예측하기 어렵지만, 작은 스텝으로 두 번 움직인 것을 목표 벡터로 사용하면 그 비선형성을 비교적 잘 반영할 수 있음.

 

월드 모델을 이용해서 task 기반으로 에이전트를 학습시킨다.
Action-conditioned video prediction으로 pretrain한 world model에 agent token(=task)을 넣어 policy와 reward를 예측하게 함
1) 먼저 dataset으로 behavior cloning + reward modeling.
2) 그 다음 실제 환경이 아니라 world model 안에서 imagined rollout을 하고, 그 위에서 RL로 policy를 더 개선한다.

Phase 2: Agent Finetuning

Agent Fintuning

• Agent token
  Diffusion으로 pretrain한 world model에 agent token을 새로운 modality로 삽입한다. 다이아몬드 얻기와 같은 task embedding이 Agent token이 되며, 나중에 reward MLP와 policy MLP를 통과시켜 각각 reward과 action distribution을 예측할 수 있다.
• Attention of agent token
  Agent token은 월드 모델의 다른 모든 토큰과 자기 자신을 어텐션할 수 있지만, 다른 모달리티는 agent에 어텐션할 수 없다. 다른 모달리티가 agent(task)를 참고하는 순간 월드 모델의 dynamics가 task 기반으로 생성될 수 있기 때문이다. 환경은 다이아몬드를 캐는 task이든, 나무를 얻는 task이든 현재 나의 state와 action 기반으로만 결정되어야 한다.
  → 따라서 agent는 월드 모델을 참고할 수 있지만 월드 모델은 agent에 관계없이 동작하도록 한다.

Behavior cloning and reward model

Policy와 reward의 task-conditioned learning.

현재 상태에서 dataset에서는 이런 행동을 했고 이런 reward가 나왔다는 것을 학습하는 단계이다.
데이터셋의 video를 tokenizer가 $z_t$로 바꾸고, action token a, tasks q, scalar rewards r이 있다. agent token의 월드 모델을 통과하여 생성된 embedding $h_t$는 현재 state와 task를 반영한 representation이다. 이 $h_t$를 이용해서 policy $p_θ(a_t|h_t)$와 reward $p_θ(r_t|h_t)$ head를 학습시켜서 task를 알 때 action과 reward가 예측 가능해진다.

• $h_t$ → $a_t$: dataset의 action을 맞추는 behavior cloning
• $h_t$ → $r_t$: dataset의 reward를 맞추는 supervised learning
  이때 video prediction loss도 다시 걸어서, noisy representation을 그대로 사용하면서 dynamics 학습 세팅을 유지한다. Agent를 붙여도 월드 모델이 망가지지 않게 됨!

Phase 3

Imagination Training (Reinforcement Learning)

Policy를 dataset에 있는 행동 수준에서 멈추지 않고, world model이 생성한 imagined rollout 안에서 reinforcement learning으로 더 향상시킨다. 이때 Dreamerv3처럼 environment interaction은 전혀 없고, policy는 purely inside the world model에서 학습된다.

Updating head: value, policy

Value (현재 state의 future reward 기댓값) head를 새로 초기화하고, policy head의 frozen copy를 behavioral prior로 둔다. 그리고 transformer 본체는 freeze 한 상태로, policy와 value head만 업데이트한다. 즉 환경은 냅두고 decision-making head만 업데이트한다.
Reward는 GT 데이터셋으로 배운 것으로 끝나며, 고정된 기준(이 상태에서 실제로 몇 점이었는가)이 되어 value랑 policy를 업데이트하는 데 사용된다.

Imagination rollout

Context는 dataset의 짧은 시퀀스로, 이전 dreamer에서는 한 context에서 여러 개의 trajectory를 만들었다. Dreamerv4는 이전 dreamer들과 다르게 한 context당 rollout을 한 개만 시작한다. 이는 data diversity를 높이면서 더 적은 rollout을 하므로 memory consumption을 줄인다고 한다.
Rollout은 transformer를 반복적으로 호출하며 생성된다. 생성된 trajectory $(z_t, a_t, r_t), (z_{t+1}, a_{t+1}, r_{t+1}), ...$ 로 Value head와 Policy head를 학습한다.

• dataset의 마지막 state $z_t$ → policy로 action head에서 샘플링한 $a_t$와 여기서 계산한 reward $r_t$
• → dynamics로 샘플링한 $z_{t+1}$로 다음 action과 state $a_{t+1}, z_{t+1}$ 를 샘플링하고 reward $r_{t+1}$계산 → …

Train Value head

이 상태에서 시작하면 앞으로 총 reward가 얼마나 될까?
Learning to predict by the methods of temporal differences., 1988의 TD-learning을 사용한다. 이는 미래 시퀀스의 예측된 reward와 value들을 조합해서 계산된 λ-return (R^\lambda_t)을 예측하도록 학습시킨다. *𝛾 = 0.997
$$\mathcal{L}(\theta) = - \sum_{t=1}^{T} \ln p_\theta(R_t^\lambda | s_t) \quad R_t^\lambda = r_t + \gamma c_t (1 - \lambda) v_t + \lambda R_{t+1}^\lambda \quad R_T^\lambda = v_T$$

Train Policy head

어떤 action이 더 좋은가?
Policy는 현재 state에서의 value를 최대화하는 action의 확률을 높이는 것이 목표이다. Preference optimization as probabilistic inference, 2024의 PMPO라는 RL objective funtion을 사용한다.
기존에는 얼마나 더 좋았는가 (magnitude)를 사용했다. 그러나 dreamerv4는 reward 도 모델이 예측한 값이고 완전히 offline이므로 어떤 trajectory는 reward 가 과하게 크고, 어떤 건 작다. magnitude를 사용하면 학습이 몇 개의 큰 값에 의해 망가진다고 한다. Dreamerv4는 모든 return을 normalize하고 scale을 맞추는 대신, 좋았는지 나빴는지의 방향만 보기로 했다.
Advantage $A_t = R^\lambda_t - v_t$가 양수라면, 이 trajectory에서 한 행동의 λ-return이 value head에서 계산한 value 기댓값보다 크다는 뜻이다. 즉 기대한 것보다 잘했다는 것이므로 이 행동을 선택할 확률을 높이는 방향으로 policy를 학습한다.

$$\mathcal{L}(\theta) = \frac{1-\alpha}{|\mathcal{D}^-|} \sum_{i \in \mathcal{D}^-} \ln \pi_\theta(a_i | s_i) - \frac{\alpha}{|\mathcal{D}^+|} \sum_{i \in \mathcal{D}^+} \ln \pi_\theta(a_i | s_i) + \frac{\beta}{N} \sum_{i=1}^N \text{KL}[\pi_\theta(a_i | s_i) \| \pi_{\text{prior}}]$$

• $\mathcal{D}^-$는 A_t<0인 state의 set, $\mathcal{D}^+$는 0 이상인 경우의 set.

 

---

 

Experiment

policy로 나온 액션을 주거나, 인간이 직접 액션을 했을 때의 성능을 비교하기도 하고 robot simulation에서 잘 되는지 확인하기도 한다. 신기한 점은 action 라벨링이 된 데이터의 양이 적어도 학습이 잘 된다는 것이다. 이는 모델이 대부분의 학습을 비디오에서 한다는 뜻으로, 환경(게임 dynamics)을 이해하는 것이 모델의 인지 능력을 발전시키는 데 가장 큰 도움이 된다는 것이다!

 

action이 10시간일 때 50~75%, 100시간일 때 full 성능이 거의 달성되는 것을 볼 수 있다.

따라서 실제 행동 데이터를 수집하기 어려운 로봇학습 태스크에서 유용하게 사용될 수 있겠다.

 

 

Limitation

여전히 길지 않은 context length, 그리고 게임 환경과 완벽하게는 동일하지 않은 상황이 한계점이다.

e.g., Inventory를 아주 정확하게 모사하지는 못하고 가끔 오류가 난다고 한다