[AI논문리뷰-강화학습] DreamerV2 - 이산 잠재 변수와 KL 균형으로 Atari 55개 게임을 인간 수준으로 정복한 세계 모델 RL

논문  Mastering Atari with Discrete World Models 저자  Hafner et al. (Google Research / DeepMind) 학회/연도  ICLR 2021 arXiv  2010.02193

📄 한 문단 요약

DreamerV2는 ICLR 2021에 발표된 Dreamer의 직접 후속으로, 연속 잠재 변수(Gaussian) 대신 32×32 범주형(categorical) 잠재 변수와 KL 균형(KL balancing)을 도입해 이산 행동 공간인 Atari로 확장했다. 단일 GPU에서 200M 환경 스텝, 468억 개의 상상 상태를 학습해 55개 Atari 게임의 게이머 정규화 중간값에서 Rainbow를 능가하고 인간 수준 성능을 달성했다. 연속 제어 전용이었던 DreamerV1의 한계를 한 번에 뛰어넘어, 단일 세계 모델 RL 프레임워크가 연속·이산 환경을 모두 커버할 수 있음을 증명한다.

이 논문은 무엇을 해결하려고 했나?

DreamerV1(ICLR 2020)은 RSSM 세계 모델과 잠재 공간 역전파 Actor-Critic으로 DeepMind Control Suite 연속 제어에서 뛰어난 샘플 효율을 보였다. 그러나 Atari처럼 이산 행동이 필요하고 게임 상태가 비연속적으로 전환되는 환경에서는 두 가지 구조적 한계가 있었다.

첫째, Gaussian 잠재 변수의 부적합성. 연속 Gaussian 분포는 고차원 게임 상태를 표현할 때 불필요하게 부드러운 잠재 공간을 만든다. 체스판의 말 위치, 아타리의 블록 배치 같은 "명확히 구분되는 상태"는 이산 표현이 훨씬 자연스럽다.

둘째, KL 손실의 불균형. Gaussian 잠재 변수에서 사전 분포(prior)와 사후 분포(posterior)를 같은 속도로 학습하면 하나가 지배하는 문제가 생긴다. 사후 분포가 관측을 무시하고 사전 분포에 수렴하는 "KL 붕괴(KL collapse)"가 일어나 세계 모델이 관측 정보를 제대로 활용하지 못한다.

DreamerV2는 이 두 문제를 각각 범주형 잠재 변수와 KL 균형 계수 α로 해결하고, Straight-through 추정기로 이산 변수를 역전파한다.

배경지식과 핵심 키워드

배경지식과 핵심 키워드

범주형 잠재 변수 (Categorical Latents) DreamerV2는 32개의 범주형 분포를 각각 32개 클래스로 구성해 32×32 = 1024차원 원-핫 잠재 벡터를 사용한다: s_t = [cat_1⊕cat_2⊕...⊕cat_32], cat_i ~ Categorical(32). 연속 Gaussian 대비 게임 상태처럼 명확히 구분된 상황을 표현하기 유리하다. "적이 오른쪽에 있다 vs 왼쪽에 있다" 같은 이진 정보가 부드러운 연속값보다 범주형으로 더 자연스럽게 표현된다. Gumbel-Softmax 또는 Straight-through 추정기로 이산 샘플링을 역전파 가능하게 만든다. ablation에서 Gaussian 대비 Atari 점수가 대폭 향상됨을 확인했다.	KL 균형 (KL Balancing, α=0.8) 표준 ELBO의 KL 항: KL[posterior || prior]를 두 항으로 분리한다: α·KL[sg(posterior) || prior] + (1−α)·KL[posterior || sg(prior)], α=0.8. sg()는 stop-gradient. α=0.8이면 prior 학습 속도가 posterior 대비 4배(0.8/0.2) 빠르다. 세계 모델이 관측을 활용하는 posterior보다 미래를 예측하는 prior를 우선 학습하도록 유도한다. 이 균형이 없으면 KL 붕괴가 발생해 잠재 상태가 관측과 무관한 사전 분포로 수렴한다. ablation에서 α=0.8 제거 시 성능이 급락했다.
Straight-Through 추정기 (STE) 이산 샘플링 s_t ~ argmax(logits)은 미분 불가능하다. STE는 순전파에서는 원-핫 샘플을, 역전파에서는 softmax 확률을 그래디언트로 통과시키는 트릭이다. 순전파: z = one_hot(argmax(logits)). 역전파: ∂L/∂logits ← ∂L/∂softmax(logits). 불연속점의 그래디언트를 근사적으로 계산한다. DreamerV2에서는 세계 모델의 잠재 상태(범주형)와 행동 모델(이산 행동) 모두에 STE를 적용한다.	REINFORCE (이산 Actor) DreamerV1의 연속 Actor는 재파라미터화(reparameterization)로 역전파가 됐다. Atari의 이산 행동은 재파라미터화가 불가능해 REINFORCE를 사용한다. REINFORCE gradient: ∇J(φ) = E[∇log π_φ(a_τ|s_τ) · (V^λ_τ − baseline)]. Baseline은 가치 함수 V_ψ(s_τ)를 사용해 분산을 줄인다. STE와 REINFORCE를 함께 사용하는 것이 DreamerV2 행동 학습의 핵심이며, V1의 reparameterization만 쓰는 방식과 가장 큰 구조적 차이다.
게이머 정규화 점수 (Gamer-Normalized Score) 정규화 점수 = (에이전트 점수 − 무작위 점수) / (전문 게이머 점수 − 무작위 점수). 값이 1.0이면 인간 수준, >1.0이면 초월. 55개 Atari 게임 중간값으로 집계하면 DreamerV2 0.217, Rainbow 0.153, IQN 0.143. DreamerV2가 모델 프리 SOTA를 40% 이상 앞선다. 단, 중간값과 평균값은 크게 다를 수 있다. 일부 게임에서 초인적 성능이 평균을 왜곡할 수 있어 중간값이 더 공정한 지표다.	Rainbow / IQN (비교 대상) Rainbow(2017)는 DQN에 6가지 개선(Double DQN, Prioritized Replay, Dueling, Multi-step, Distributional, Noisy)을 통합한 모델 프리 Atari SOTA. 게이머 정규화 중간값 0.153. IQN(Implicit Quantile Networks, 2018)은 분포 기반 Q-learning 방법. 중간값 0.143. DreamerV2(0.217)는 이 두 방법을 단일 GPU 학습 환경에서 명확히 앞선다. 세계 모델 기반 방법이 처음으로 Atari 모델 프리 SOTA를 넘어선 결과다.
상상 상태 (Imagined States) DreamerV2는 200M 환경 스텝에서 468억 개의 상상 상태를 생성해 Actor-Critic을 학습했다. 실제 환경 상호작용 대비 234배 많은 학습 신호다. 상상 지평선 H=15. 각 실제 경험 1회당 세계 모델 안에서 234번의 상상 경험이 이뤄지는 셈이다. 이것이 샘플 효율의 근본이다. 상상이 많을수록 세계 모델 오차가 누적된다. DreamerV2는 범주형 잠재 변수로 세계 모델 정확도를 높여 이 문제를 완화했다.	DreamerV1 → V2 주요 변경 V1: 연속 Gaussian 잠재 변수 → V2: 32×32 범주형 잠재 변수. V1: 재파라미터화 그래디언트만 사용 → V2: STE(잠재 상태) + REINFORCE(이산 행동). V1: 표준 KL 손실 → V2: KL 균형(α=0.8 stop-gradient 비대칭). 이 세 가지 변경만으로 DM Control은 유지하면서 Atari 인간 수준까지 확장됐다.

핵심 인사이트

DreamerV2의 핵심은 "세 가지 외과적 변경"이다: ① Gaussian → 범주형 잠재 변수, ② 표준 KL → KL 균형(α=0.8), ③ 재파라미터화 → STE+REINFORCE. 세계 모델 구조(RSSM)는 V1과 동일하다. 이 세 변경만으로 연속 제어 전용 → 연속+이산(Atari) 동시 커버가 가능해졌다.

기존 방법의 한계

DreamerV1 (연속 환경 전용): Gaussian 잠재 변수와 재파라미터화 Actor로 DM Control에서 탁월했지만, Atari의 이산 행동과 비연속적 게임 상태 전환에 맞지 않았다. 적용 시 세계 모델 학습이 불안정하고 성능이 낮았다.

Rainbow / IQN (Atari 모델 프리 SOTA): 모델 프리로 Atari에서 강력했지만 200M 스텝을 전부 실제 환경과 상호작용해야 한다. 상상 학습이 없어 샘플 효율이 낮다. 단일 GPU로는 학습 시간이 Dreamer보다 훨씬 길다.

SimPLe (모델 기반 Atari 선행): 픽셀 레벨 비디오 예측 세계 모델로 Atari를 시도했지만 고차원 픽셀 예측 오차가 계획에 누적되어 성능이 제한됐다. DreamerV2처럼 잠재 공간에서 계획하는 방식이 아니었다.

MuZero (MCTS 기반): 학습된 역학 모델로 Atari를 정복했지만 MCTS에 매 스텝 50 시뮬레이션이 필요하고 대규모 컴퓨팅이 필요하다. DreamerV2는 단일 GPU로 경쟁 수준의 성능을 낸다.

제안 방법의 핵심 아이디어

DreamerV2는 DreamerV1 구조를 유지하면서 세 가지를 바꾼다.

변경 1 — 범주형 잠재 변수:

// DreamerV1 (Gaussian)
s_t ~ Normal(μ_t, σ_t)  →  s_t ∈ ℝ^200

// DreamerV2 (Categorical)
s_t = [cat_1 ⊕ cat_2 ⊕ ... ⊕ cat_32]
cat_i ~ Categorical(32)  // 각 슬롯: 32개 클래스 중 하나 선택
s_t ∈ {0,1}^(32×32)     // 원-핫 1024차원

32개 범주형 분포 × 각 32개 클래스 = 총 32^32 ≈ 10^48가지 가능한 상태 공간. Gaussian 200차원보다 훨씬 표현력이 크다.

변경 2 — KL 균형:

// 표준 KL (V1)
L_KL = KL[ posterior || prior ]

// KL 균형 (V2, α=0.8)
L_KL = α·KL[ sg(posterior) || prior ]      // prior 학습 (α=0.8)
     + (1-α)·KL[ posterior || sg(prior) ]  // posterior 학습 (1-α=0.2)

α=0.8이면 prior가 posterior 대비 4배 빠르게 학습된다. 세계 모델이 미래를 예측하는 능력(prior)을 관측을 인코딩하는 능력(posterior)보다 우선 향상시킨다.

변경 3 — 이산 행동 Actor:

// Atari 이산 행동
a_t ~ Categorical(π_φ(s_t))    // 18개 버튼 중 하나

// REINFORCE gradient (STE 불가 — argmax 뒤에 환경이 옴)
∇J(φ) = E[ ∇log π_φ(a_τ|s_τ) · (V^λ_τ - V_ψ(s_τ)) ]

// 잠재 상태 그래디언트는 STE로 통과
s_t = one_hot(argmax(logits)),  역전파: ∂L/∂logits ← ∂L/∂softmax(logits)

논문 그림/표로 이해하기

원문 Figure 2: DreamerV2 세계 모델 — CNN, RSSM, 보상·종료·이미지 예측 헤드. 출처: Hafner et al., "Mastering Atari with Discrete World Models", ICLR 2021. arXiv:2010.02193 — arXiv 라이선스 하에 인용.

Figure 2는 세계 모델 구조다. 왼쪽 CNN이 84×84 그레이스케일 이미지를 인코딩하고, RSSM의 GRU가 결정론적 상태 h_t를 유지하며, 범주형 분포에서 s_t를 샘플링한다. 오른쪽으로 이미지 재구성, 보상 예측, 에피소드 종료 예측 헤드가 달린다. DreamerV1과 구조는 동일하고 잠재 변수 타입만 범주형으로 바뀌었다.

원문 Figure 3: 잠재 공간에서 상상된 H=15 롤아웃으로 Actor-Critic 학습. 출처: Hafner et al., ICLR 2021. arXiv:2010.02193 — arXiv 라이선스 하에 인용.

Figure 3은 행동 학습 단계다. 현재 잠재 상태 s_t에서 H=15 스텝의 가상 궤적을 RSSM 전이 모델로 롤아웃하고, 각 스텝에서 이산 행동을 샘플링(REINFORCE), λ-return으로 가치를 추정해 Actor-Critic을 업데이트한다. 이 전체 과정이 실제 환경 없이 잠재 공간에서만 이루어진다.

원문 Figure 1: Atari 55개 게임 200M 스텝에서 게이머 정규화 중간값 점수 비교. 출처: Hafner et al., ICLR 2021. arXiv:2010.02193 — arXiv 라이선스 하에 인용.

Figure 1은 200M 스텝 학습 과정에서 DreamerV2(파란선)가 Rainbow와 IQN을 뚜렷이 앞서는 것을 보여준다. 특히 100M 스텝 이후부터 격차가 크게 벌어진다. 세계 모델이 충분히 학습된 후 상상 학습의 효율이 극대화되기 때문이다.

실험 설정과 결과 해석

Atari 55개 게임 주요 결과 (200M 환경 스텝):

알고리즘	유형	중간값 (gamer norm.)	인간 수준 게임 수
IQN	모델 프리	0.143	39/55
Rainbow	모델 프리	0.153	40/55
DreamerV1 (Gaussian)	모델 기반	0.057	—
✦ DreamerV2 (범주형)	모델 기반	0.217	45/55

논문 Table 1 재구성. DreamerV2는 Rainbow 대비 게이머 정규화 중간값 42% 향상, 인간 수준 게임 수 12.5% 증가.

ablation 결과 (Figure 5): 범주형 잠재 변수, KL 균형, REINFORCE 중 하나라도 제거하면 성능이 크게 떨어진다. 특히 KL 균형 제거(α=0.5 사용) 시 중간값이 0.093까지 급락 — DreamerV1 수준으로 회귀한다. 세 기법이 상호보완적으로 작용한다는 것을 보여준다.

DM Control 성능 유지: 새로운 구성이 연속 제어 환경에서도 V1과 동등하거나 더 나은 성능을 유지했다. 이산 환경을 위한 변경이 연속 환경을 해치지 않았다.

핵심 인사이트

DreamerV1에 Gaussian 잠재 변수를 그대로 두고 Atari를 시도하면 중간값 0.057에 불과하다. V2의 범주형 잠재 변수로만 바꿔도 0.171로 세 배 이상 뛴다. KL 균형까지 더하면 0.217. "어떤 종류의 잠재 변수를 쓰는가"가 세계 모델 성능의 결정적 요인이라는 강력한 증거다.

한계와 비판적 관점

1. 일부 Atari 게임에서의 낮은 성능: Montezuma's Revenge, Pitfall 같은 희소 보상·탐색 중심 게임에서는 DreamerV2도 무작위 에이전트 수준에 머무른다. 호기심 기반 탐색(intrinsic motivation)이 없어 새로운 방을 방문할 동기가 부족하다.

2. 긴 학습 시간: 단일 GPU에서 55개 Atari 게임을 각각 학습하는 데 ~10일이 걸린다. Rainbow가 유사한 하드웨어에서 더 빠르게 수렴하는 게임도 있다. "단일 GPU"라는 점은 긍정적이지만 10일은 여전히 길다.

3. 세계 모델 표현력 한계: 32×32 범주형이 표현할 수 있는 상태 수는 방대하지만, 장기 의존성이 있거나 매우 세밀한 연속 정보가 중요한 게임(Skiing, Solaris 등)에서는 한계를 보인다.

4. 단일 환경 학습: V2도 게임 하나당 독립적으로 학습한다. 다중 게임에서 단일 세계 모델을 공유하는 일반화는 DreamerV3에서야 시도된다.

5. 클리핑 정규화의 모호성: 논문의 "clipped record-normalized score" 지표가 표준적이지 않아 다른 논문 결과와 직접 비교하기 어렵다는 비판이 있었다.

구현하거나 응용한다면 무엇을 봐야 하나?

DreamerV1 → V2 마이그레이션 체크리스트:

1. 잠재 변수 교체: Normal(μ, σ) 샘플링을 Categorical(32)×32 샘플링으로 교체. 원-핫 인코딩 후 연결(concat)해 1024차원 벡터 사용.
2. KL 손실 수정: 기존 KL[q||p]를 0.8·KL[sg(q)||p] + 0.2·KL[q||sg(p)]로 변경. PyTorch에서는 q.detach()로 stop-gradient 구현.
3. Straight-Through 구현:

# 순전파: one-hot, 역전파: softmax 통과
logits = self.encoder(h_t)
probs = F.softmax(logits, dim=-1)
indices = probs.argmax(dim=-1)
one_hot = F.one_hot(indices, num_classes=32).float()
# STE: 순전파는 one_hot, 역전파는 probs
s_t = one_hot + probs - probs.detach()

4. 행동 손실 변경: 연속 환경은 재파라미터화 Actor 유지, Atari는 REINFORCE = log_prob(a) * (V_lambda - baseline).

오픈소스 참고:

• danijar/dreamerv2 (GitHub): 저자 공식 TensorFlow 구현.
• zhaoyi11/dreamerv2-pytorch: PyTorch 포팅 버전. 범주형 잠재 변수와 KL 균형 구현 참고용.

현실적 고려: Atari에서 처음부터 DreamerV2를 구현하기보다는 저자의 DreamerV3를 쓰는 것이 훨씬 빠르다. V3는 단일 코드베이스로 Atari·DM Control·Minecraft를 모두 커버하며 하이퍼파라미터도 최적화됐다.

한 줄 결론과 다음에 읽을 논문

한 줄 결론: DreamerV2는 Gaussian → 범주형 잠재 변수, KL 균형, STE+REINFORCE 세 가지 변경만으로 Atari 55개 게임에서 모델 프리 SOTA(Rainbow)를 넘어섰다 — 세계 모델 RL이 이산 환경에서도 경쟁력 있음을 처음 증명한 논문이다.

다음에 읽을 논문:

• DreamerV3 (2023, arXiv:2301.04104): 하나의 하이퍼파라미터 세트로 Atari·DM Control·Minecraft·로보틱스를 모두 학습. V2의 한계를 정규화(symlog) 등으로 해소.
• EfficientZero (2021): MuZero에 RSSM 유사 자기지도 일관성 손실 추가. Atari 100K 벤치마크에서 DreamerV2와 경쟁.
• IRIS (2022): Transformer 기반 세계 모델로 DreamerV2보다 낮은 샘플 수(100K)에서 우수한 성능.

---

참고 자료:

• 논문 원문: arXiv:2010.02193
• 저자 공식 구현: github.com/danijar/dreamerv2
• DreamerV3 (최신): github.com/danijar/dreamerv3