Recurrent World Models Facilitate Policy Evolution

lastmod 2026-10-28created 2026-10-283 edits

World Models는 David Ha와 Jürgen Schmidhuber가 2018년에 발표한 논문이다. 강화학습 환경을 비지도 방식으로 학습한 생성형 월드 모델을 만들고, 에이전트를 그 모델이 만든 "꿈(dream)" 속에서만 학습시킨 뒤 정책을 실제 환경으로 옮길 수 있음을 보였다.

• 부제: "Can agents learn inside of their own dreams?"
• 저자: David Ha (Google Brain), Jürgen Schmidhuber (NNAISENSE, IDSIA)
• 공개: 2018-03-27 (arXiv:1803.10122)
• 학회: NeurIPS 2018 (발표 제목: Recurrent World Models Facilitate Policy Evolution)
• 형식: 브라우저에서 모델을 직접 실행하는 인터랙티브 웹 논문 (p5.js, deeplearn.js 사용)

Categories

• Model Works
• Reinforcement learning

Abstract

대표적인 강화학습 환경들에 대한 생성형 신경망 모델(World Model)을 구축한다. World Model은 비지도 방식으로 빠르게 학습되어 환경의 압축된 시공간(spatial and temporal) 표현을 익힌다. World Model에서 추출한 특징을 에이전트의 입력으로 쓰면, 매우 작고 단순한 정책으로도 과제를 해결할 수 있다. 나아가 에이전트를 World Model이 생성한 환각 속 꿈(hallucinated dream)에서만 학습시킨 뒤, 그 정책을 실제 환경으로 옮겨도 작동함을 보인다.

핵심 아이디어

• 인간은 제한된 감각으로 얻은 정보로 세계의 내부 모델(mental model)을 만들고, 이를 바탕으로 판단하고 행동한다. 뇌는 미래를 예측하며 지각하고, 야구 타자가 의식적으로 계획하기 전에 반사적으로 배트를 휘두르는 것도 이 예측 모델 덕분이다.
• 크고 표현력 있는 RNN이 유리하지만, model-free RL은 credit assignment 문제 때문에 큰 네트워크를 학습하기 어려워 보통 작은 네트워크를 쓴다.
• 그래서 에이전트를 둘로 나눈다.
  • 큰 World Model (V, M): 비지도 학습으로 환경을 압축하고 예측한다.
  • 작은 Controller (C): 강화학습으로 행동을 학습한다. 탐색 공간이 작아 credit assignment에 집중할 수 있다.
• 실제 환경 대신 World Model이 생성한 가상 환경에서만 Controller를 학습할 수 있다. 에이전트가 모델의 허점을 악용하는 문제는 온도(temperature) τ로 불확실성을 높여 완화한다.

Agent Model

에이전트는 세 개의 구성 요소로 이루어진다.

V: VAE

• 각 시점의 2D 이미지 프레임을 저차원 잠재 벡터 \(z\)로 압축하고 다시 복원하는 Convolutional Variational Autoencoder다.
• 입력은 64×64×3 RGB 이미지다. 4층 convolution 인코더와 4층 deconvolution 디코더를 쓰며, 모든 층은 stride 2이고 출력층을 제외하면 ReLU를 쓴다.
• 잠재 차원은 Car Racing 32, Doom 64이다.
• 손실 함수는 L2 재구성 손실과 KL 손실의 합이다. 가우시안 사전 분포는 정보 용량을 제한하지만, M이 만든 비현실적인 \(z\)에도 강건해지게 한다.

M: MDN-RNN

• V가 시간에 따라 보는 내용을 압축한다. 다음 \(z\)를 결정론적으로 예측하지 않고 확률 밀도 \(p(z)\)를 출력한다.
• LSTM에 Mixture Density Network 출력층을 붙인 구조다. 가우시안 5개를 혼합하고 대각 공분산을 쓴다.
• 은닉 유닛 수는 Car Racing 256, Doom 512이다.
• teacher forcing으로 20 epoch 학습한다. 저장해 둔 \(\mu, \sigma\)에서 배치마다 \(z\)를 새로 샘플링해 과적합을 줄인다.
• 샘플링할 때 온도 τ로 모델의 불확실성을 조절한다.
• SketchRNN 등 기존 MDN-RNN 계열 시퀀스 생성 모델과 같은 방식이다.

C: Controller

롤아웃 동안 기대 누적 보상을 최대화하도록 학습하는 단일 선형층이다.

\(a_t = W_c \; [z_t \; h_t] \; + \; b_c\)

• \(h_t\)에 미래 예측에 필요한 정보가 이미 담겨 있으므로, 예측된 \(z_{t+1}\)은 입력하지 않는다.
• 모델 복잡도 대부분을 V와 M에 두고 C를 작게 유지한다. 그래서 역전파 대신 CMA-ES(진화 전략)로 최적화할 수 있다. CMA-ES는 파라미터 수천 개 이하의 공간에서 잘 작동하고 병렬화가 쉽다.
  • 인구 크기는 64이고, 개체마다 서로 다른 시드로 16회 롤아웃한 평균 누적 보상을 적합도로 쓴다.
• 출력은 tanh로 행동 범위에 맞춘다.

V, M, C의 결합

def rollout(controller):
  ''' env, rnn, vae are '''
  ''' global variables  '''
  obs = env.reset()
  h = rnn.initial_state()
  done = False
  cumulative_reward = 0
  while not done:
    z = vae.encode(obs)
    a = controller.action([z, h])
    obs, reward, done = env.step(a)
    cumulative_reward += reward
    h = rnn.forward([a, z, h])
  return cumulative_reward

Car Racing 실험

OpenAI Gym의 CarRacing-v0에서 픽셀 입력만으로 주행한다. 트랙은 매번 무작위로 생성되며, 100회 연속 시행 평균 900점 이상이면 해결로 본다. 행동은 조향(−1~1), 가속(0~1), 브레이크(0~1)의 연속값이다.

학습 절차

1. 무작위 정책으로 롤아웃 10,000회를 수집한다.
2. V(VAE)를 학습해 각 프레임을 \(z \in \mathbb{R}^{32}\)로 인코딩한다.
3. M(MDN-RNN)을 학습해 \(P(z_{t+1} \mid a_t, z_t, h_t)\)를 모델링한다.
4. CMA-ES로 C를 진화시켜 기대 누적 보상을 최대화한다.

World Model(V, M)은 보상 신호를 전혀 보지 않고 프레임 압축과 예측만 학습한다. 보상에 접근하는 것은 C뿐이다.

결과

• V만 사용 (\(z_t\)만 입력): 주행이 흔들리고 급커브에서 트랙을 벗어난다. 선형 C는 632 ± 251점이고, 은닉층(tanh 40유닛, 파라미터 1,443개)을 추가해도 788 ± 141점이다.
• V와 M 사용 (\(z_t\), \(h_t\) 입력): 주행이 안정되고 급커브를 잘 처리한다. 미래를 명시적으로 계획하지 않고, \(h_t\)에 담긴 미래 확률 분포 정보로 '본능적으로' 반응한다.
• 1,800세대 후 1,024회 무작위 롤아웃 평균 900.46점에 도달했고, 100회 시행 결과는 906 ± 21점이다.

기존 Deep RL 방법은 엣지 검출 같은 전처리와 프레임 스택이 필요했다. World Models는 원시 RGB 픽셀에서 바로 시공간 표현을 학습했고, 저자들이 아는 한 이 과제의 해결 기준을 처음으로 달성했다.

Car Racing Dreams

M은 미래를 예측할 수 있으므로, 샘플링한 \(z_{t+1}\)을 실제 관측처럼 다시 입력해 가상의 주행 장면을 만들 수 있다. 학습된 C를 이 꿈 환경에 넣어 운전시킬 수 있으며, 웹 논문에서는 τ를 조절하며 직접 체험해 볼 수 있다.

VizDoom 실험: 꿈 속에서 학습하기

VizDoom의 DoomTakeCover-v0에서 방 반대편 몬스터가 쏘는 화염구를 피해 최대한 오래 살아남는다. 생존 시간이 곧 누적 보상이다. 한 에피소드는 최대 2,100 타임스텝(약 60초)이고, 100회 연속 평균 750 타임스텝(약 20초) 이상이면 해결로 본다.

Car Racing과 달리 M은 다음 \(z\)뿐 아니라 에이전트의 사망 여부 \(d\)도 예측한다. 그래서 M을 Gym 환경 인터페이스로 감싸 완전한 가상 환경으로 쓸 수 있다. 이 가상 환경은 픽셀을 렌더링하지 않고 잠재 공간에서만 동작한다. 사망 확률이 50%를 넘으면 done을 True로 설정한다.

학습 절차

1. 무작위 정책으로 롤아웃 10,000회를 수집한다.
2. V(VAE)를 학습해 각 프레임을 \(z \in \mathbb{R}^{64}\)로 인코딩한다.
3. M(MDN-RNN)을 학습해 \(P(z_{t+1}, d_{t+1} \mid a_t, z_t, h_t)\)를 모델링한다.
4. CMA-ES로 C를 진화시켜 가상 환경 안에서 생존 시간을 최대화한다.
5. 학습된 정책을 실제 환경에 적용한다.

꿈 속 학습과 정책 이전

• M은 원시 이미지만으로 게임 로직, 몬스터의 행동, 물리, 3D 그래픽을 흉내 낸다. 벽에 막히는 동작과 여러 화염구의 움직임도 학습한다.
• 가상 환경에서는 약 900 타임스텝을 얻었다. 이 정책을 실제 환경에 옮기자 100회 평균 약 1,100 타임스텝을 기록했다. 해결 기준(750)을 크게 넘고, 더 어렵게 설정된 가상 환경에서보다도 높은 점수다.
• V가 몬스터 수 같은 세부를 정확히 재현하지 못해도, 에이전트는 잡음 많은 가상 환경에서 익힌 정책을 깨끗한 실제 환경에서 잘 사용한다.

World Model 속이기 (Cheating the World Model)

• 초기 버전의 에이전트는 몬스터가 화염구를 아예 쏘지 않게 만드는 적대적 정책(adversarial policy)을 찾아냈다.
• M은 실제 환경을 근사한 확률 모델이라 현실 법칙을 따르지 않는 궤적을 만들 때가 있다. C는 M의 은닉 상태에도 접근할 수 있어 이런 허점을 쉽게 악용한다. 결정론적 모델일수록 더 취약하다.
• 대책은 M을 확률 분포(가우시안 혼합)로 모델링하고, 온도 τ를 높여 가상 환경을 더 불확실하고 어렵게 만드는 것이다.
  • τ가 매우 낮으면(예: 0.1) mode collapse가 일어나 몬스터가 화염구를 쏘지 않는다. 그러면 가상 환경에서는 거의 만점이지만 실제 환경에서는 무작위 정책보다 못하다.
  • 단일 가우시안 대신 가우시안 혼합을 쓰는 이유는 "몬스터가 쏠지 말지" 같은 이산적 사건을 표현하기 위해서다.

• τ를 높이면 적대적 정책을 찾기 어려워지지만, 너무 높이면 에이전트가 아무것도 배우지 못한다. 따라서 τ는 튜닝해야 하는 하이퍼파라미터다.
• τ=1.15에서 실제 환경 최고 점수를 얻었다. τ=1.30에서는 평균은 낮지만 분산이 작은, 더 안전한 정책이 나왔다.

Iterative Training Procedure

무작위 정책으로 모은 데이터만으로는 복잡한 환경을 충분히 탐험할 수 없다. 그래서 다음 절차를 반복할 것을 제안한다.

1. M과 C를 무작위 파라미터로 초기화한다.
2. 실제 환경에서 N회 롤아웃하고, 모든 행동 \(a_t\)와 관측 \(x_t\)를 저장한다.
3. M이 \(P(x_{t+1}, r_{t+1}, a_{t+1}, d_{t+1} \mid x_t, a_t, h_t)\)를 모델링하도록 학습하고, C는 M 안에서 기대 보상을 최적화하도록 학습한다.
4. 과제가 끝나지 않았으면 2단계로 돌아간다.

• 단순한 과제는 1회 반복으로 충분하지만, 복잡한 과제는 여러 번 반복해야 한다.
• M의 손실을 C의 보상으로 쓰면 인공 호기심(artificial curiosity)이 된다. 에이전트가 M이 잘 예측하지 못하는 상황을 찾아가며 새 데이터를 모은다.
• M이 행동 \(a_{t+1}\)까지 예측하면 C의 행동을 흡수할 수 있다. 그러면 걷기 같은 저수준 기술은 M에 맡기고 C는 고수준 기술에 집중할 수 있다.
• 신경과학적으로는, 동물이 쉬거나 잠잘 때 해마가 최근 경험을 재생(replay)해 기억을 공고화하는 현상과 연결된다.

Related Work

• PILCO: 가우시안 프로세스로 환경 동역학을 학습하는 확률적 model-based 방법이다. 저차원 관측에서만 효율적이다.
• 베이지안 신경망(BNN) 기반 동역학 모델, 오토인코더 잠재 공간에서 동역학을 학습해 픽셀 입력으로 진자를 제어한 연구 등이 선행 연구로 소개된다.
• Guzdial 등의 CNN 기반 게임 순방향 시뮬레이션 모델, Graves의 "Hallucination with RNNs" 강의, Finn의 Model-Based RL 강의가 언급된다.
• RNN 기반 World Model과 Controller의 결합은 Schmidhuber의 1990년 연구 "Making the World Differentiable" 등에서 이미 탐구되었고, "Learning to Think"가 이를 일반화했다.
• Controller 학습에 쓴 진화 전략(ES)은 최종 누적 보상만 있으면 되고 병렬화가 쉽다.

Discussion

• 장점
  • 게임 엔진의 무거운 렌더링과 물리 연산 대신, 가벼운 가상 환경에서 원하는 만큼 학습할 수 있다. GPU로 분산 학습하기도 쉽다.
  • World Model은 미분 가능하므로, 꿈 속 정책을 역전파로 직접 미세 조정할 수도 있다.
• 한계
  • VAE를 따로 학습하면 벽돌 무늬처럼 과제와 무관한 부분을 재현하고, 도로 타일처럼 중요한 부분은 놓칠 수 있다. 보상을 예측하는 M과 함께 학습하면 과제에 필요한 특징에 집중하지만, 다른 과제에 재사용하기는 어려워진다.
  • LSTM 기반 M은 용량이 제한적이고, 역전파로 학습할 때 catastrophic forgetting이 생긴다. 더 큰 모델이나 외부 메모리가 필요하다.
  • 타임스텝 단위로 미래를 시뮬레이션할 뿐, 계층적 계획이나 추상적 추론은 하지 않는다. "Learning to Think"처럼 C가 M을 서브루틴으로 질의하는 방식으로 확장할 수 있다.
• C와 M을 하나의 네트워크(One Big Net)로 합치는 방향도 향후 연구로 제시한다.

Citation

@incollection{ha2018worldmodels,
  title = {Recurrent World Models Facilitate Policy Evolution},
  author = {Ha, David and Schmidhuber, J{\"u}rgen},
  booktitle = {Advances in Neural Information Processing Systems 31},
  pages = {2451--2463},
  year = {2018},
  publisher = {Curran Associates, Inc.},
  url = {https://papers.nips.cc/paper/7512-recurrent-world-models-facilitate-policy-evolution},
  note = "\url{https://worldmodels.github.io}",
}

다이어그램과 텍스트는 CC-BY 4.0 라이선스다. 단, 다른 출처에서 가져온 그림은 제외된다.

Documentation

World Models

https://arxiv.org/abs/1803.10122

Recurrent World Models Facilitate Policy Evolution (NeurIPS 2018)

https://papers.nips.cc/paper/7512-recurrent-world-models-facilitate-policy-evolution

See also

• World Models - 월드 모델 일반 개념
• World Foundation Model
• Reinforcement learning
• Autoencoder
• VAE
• RNN
• LSTM
• Mixture Density Network
• Gaussian mixture model
• Physical AI

Favorite site

• World Models - Can agents learn inside of their own dreams?
• Github - worldmodels/worldmodels.github.io
• World Models Experiments - 실험 재현 가이드 (저자 블로그)
• World Models 발표 영상 (YouTube)