월드 모델

월드 모델

화면을 작은 벡터로 압축하는 눈(V), 다음 장면을 예측하는 기억(M), 아주 작은 조종기(C)로 에이전트를 나눠, 자동차 경주 과제를 풀고 심지어 모델이 만든 '꿈속'에서만 연습한 정책을 실제 게임에 옮겨 쓴 논문입니다.

들어가기 전에

이 해설은 아래 내용을 안다고 보고 씁니다. 처음이라면 먼저 읽고 오면 수식과 손계산을 훨씬 쉽게 따라올 수 있습니다.

• 변분 오토인코더 (VAE) — 재매개변수화 트릭으로 인코더 · 디코더를 경사하강법으로 함께 학습하는 확률적 생성 모델
• 순환 신경망(RNN)과 LSTM 직관 — RNN 갱신식 · 기울기 소실 · LSTM 게이트를 손계산으로 익히는 기초 수업
• 강화학습 입문 (Sutton · Barto, 2판) — 슬롯머신 · MDP · 벨만 방정식 · 가치 반복을 3칸 격자 손계산으로 배우는 강화학습 첫 수업

한눈에 보기

• 에이전트를 '보는 부분 · 기억하고 예측하는 부분 · 행동을 고르는 부분'으로 나누면 왜 좋은지 알게 됩니다.
• VAE로 화면을 32차원(또는 64차원) 벡터로 줄이고, MDN-RNN으로 다음 벡터의 확률 분포를 예측하는 흐름을 따라갈 수 있습니다.
• 조종기 매개변수가 867개뿐인 이유를 직접 계산해 볼 수 있습니다.
• '꿈속에서 학습'이 무엇이고, 에이전트가 꿈의 허점을 파고드는 문제를 온도 $\tau$로 어떻게 누르는지 압니다.

초록 (번역)

이 논문은 널리 쓰이는 강화학습 환경에 대한 생성형 신경망 모델을 만드는 방법을 탐구합니다. 우리의 월드 모델은 비지도 방식으로 빠르게 학습해 환경의 공간적 · 시간적 표현을 압축해 배울 수 있습니다. 월드 모델에서 뽑은 특징을 에이전트의 입력으로 쓰면, 주어진 과제를 풀 수 있는 아주 작고 단순한 정책을 학습할 수 있습니다. 심지어 에이전트를 월드 모델이 만들어 낸 환각 같은 꿈속에서만 학습시킨 뒤, 그 정책을 실제 환경으로 옮겨 쓸 수도 있습니다. 이 논문의 대화형 버전은 이 https URL에서 볼 수 있습니다.

무엇을 풀었나

강화학습은 보상이라는 약한 신호 하나로 많은 매개변수를 학습해야 합니다. 이것을 신용 할당 문제(credit assignment problem) 라고 합니다. "지금 받은 보상이 몇 초 전 어떤 행동 덕분인가"를 알아내기 어렵기 때문입니다. 그래서 당시 모델 프리 강화학습은 작은 신경망을 쓰는 경우가 많았고, 큰 모델의 표현력을 살리지 못했습니다.

사람은 다르게 행동합니다. 눈에 들어온 장면을 추상적인 개념으로 줄여 기억하고, 다음에 무슨 일이 일어날지 머릿속 모델로 예측하며 반사적으로 반응합니다. 저자들은 이 생각을 그대로 옮겨, 큰 모델은 비지도 학습으로 세상을 배우게 하고, 보상으로 학습하는 부분은 아주 작게 만들자고 제안합니다.

핵심 아이디어

1234567891011flowchart LR
  O[화면 64x64x3] --> V[V: VAE 인코더]
  V --> Z[잠재 벡터 z_t]
  Z --> C[C: 선형 조종기]
  H[M의 숨은 상태 h_t] --> C
  C --> A[행동 a_t]
  A --> E[환경 또는 꿈]
  A --> M[M: MDN-RNN]
  Z --> M
  M --> H2[다음 h_t+1 · 다음 z 분포]
  E --> O

1. V — 화면을 짧은 요약 벡터로 (VAE)

직관: 사진 한 장을 "왼쪽으로 굽은 도로, 차는 가운데" 같은 짧은 메모로 줄이는 것입니다(해설용 비유).

단계: 화면을 64×64 픽셀로 줄입니다 → 합성곱 층 4개로 평균 $\mu$와 표준편차 $\sigma$를 만듭니다 → 거기서 잠재 벡터 $z$를 뽑습니다 → 역합성곱 층 4개로 원래 화면을 다시 그려 보며, 원본과의 차이(L2)와 KL 항을 줄이도록 학습합니다.

(손실 식은 논문 부록 설명 'L2 재구성 손실 + KL 손실'을 식으로 적은 것입니다.)

왜 이렇게 했나: 보상과 상관없이 화면만으로 학습할 수 있어서, 무작위로 모은 데이터로도 충분히 배웁니다. 픽셀 수천 개 대신 숫자 32개만 넘기므로 뒤쪽 부분이 훨씬 가벼워집니다.

2. M — 다음 장면을 확률로 예측 (MDN-RNN)

직관: 운전하면서 "곧 커브가 나올 것 같다, 아니면 직진일 수도 있다"처럼 여러 가능성을 함께 떠올리는 기억 장치입니다.

단계: LSTM이 지금 행동 $a_t$, 지금 잠재 벡터 $z_t$, 자기 숨은 상태 $h_t$를 받습니다 → 출력층(혼합 밀도 네트워크, MDN)이 다음 $z_{t+1}$에 대한 가우시안 혼합 분포의 매개변수를 냅니다.

왜 혼합 분포인가: 게임은 갈림길이 많습니다. Doom에서 불덩이가 "생길 수도, 안 생길 수도" 있을 때 평균 하나로 예측하면 '반쯤 생긴 흐릿한 불덩이'가 됩니다. 혼합 분포는 서로 다른 미래를 따로따로 표현할 수 있습니다. 저자들은 온도를 아주 낮게 하면 꿈속에서 불덩이가 아예 안 생긴다고 보고합니다.

숫자로 해 보기(해설용 예): 1차원 $z$에서 혼합 2개가 $\pi=(0.7, 0.3)$, $\mu=(1.0, -2.0)$이면, 다음 $z$의 기댓값은 $0.7\times1.0 + 0.3\times(-2.0) = 0.1$입니다. 하지만 실제로는 0.1 근처가 아니라 1.0 근처(70%) 아니면 −2.0 근처(30%)가 나옵니다. 평균만 예측하면 이 갈림길이 사라집니다.

3. C — 아주 작은 선형 조종기와 진화 전략

직관: 복잡한 판단은 눈과 기억이 다 해 두고, 조종기는 그 요약을 보고 핸들을 얼마나 꺾을지만 정하는 단순한 규칙입니다.

숫자로 해 보기: 자동차 경주에서 $z$는 32차원, $h$는 256차원이라 입력은 288차원입니다. 행동이 3차원(핸들 · 가속 · 브레이크)이면 가중치 $288\times3=864$개, 편향 3개로 합계 867개입니다. 논문이 밝힌 조종기 매개변수 수 867과 같습니다. 반면 V는 4,348,547개, M은 422,368개입니다.

왜 이렇게 했나: 매개변수가 수천 개 이하이면, 기울기를 쓰지 않는 진화 전략인 CMA-ES(공분산 행렬 적응 진화 전략)로 바로 최적화할 수 있습니다. CMA-ES는 매개변수 수천 개 정도까지 잘 동작한다고 알려져 있습니다. 여러 CPU 코어에서 환경을 동시에 돌려 각 후보 조종기의 누적 보상만 보고 좋은 후보를 고릅니다. 신용 할당이 어려운 문제를 '작은 탐색 공간'으로 바꾼 셈입니다.

4. 꿈속에서 학습하기

직관: 운전 시뮬레이터에서 연습한 뒤 실제 도로에 나가는 것과 같습니다. 다만 시뮬레이터를 사람이 아니라 에이전트가 스스로 만들었습니다.

단계(Doom 'Take Cover' 과제):

1. 무작위 정책으로 10,000번 플레이해 화면을 모읍니다.
2. V가 각 화면을 $z \in \mathbb{R}^{64}$로 바꿉니다.
3. M이 다음 $z$와 함께 '이번 칸에서 죽었는가'($d_{t+1}$)까지 예측하도록 학습합니다. 죽을 확률이 50%를 넘으면 에피소드를 끝냅니다.
4. 실제 게임 대신 M이 만든 가상 환경 안에서 CMA-ES로 C를 학습합니다. 이때 V는 픽셀을 인코딩할 필요가 없어, 학습이 전부 잠재 공간에서 일어납니다.
5. 학습한 C를 실제 VizDoom 환경에 그대로 넣어 평가합니다.

꿈의 허점 파고들기: 에이전트가 꿈속에서 불덩이가 생기지 않게 하는 이상한 움직임을 찾아내, 실제로는 통하지 않는 '치트'로 높은 점수를 내는 일이 생겼습니다. 조종기가 M의 숨은 상태까지 볼 수 있어서, 모델의 결함을 이용할 수 있기 때문입니다.

왜 온도를 올리나: 온도 $\tau$를 높여 꿈을 더 불확실하고 어렵게 만들면 치트가 잘 통하지 않습니다. 너무 높이면 꿈이 너무 어려워 아무것도 못 배웁니다. 그래서 $\tau$는 조절해야 하는 하이퍼파라미터입니다.

왜 결정적 모델이 아니라 확률 모델인가: 다음 장면을 한 가지로만 예측하는 모델이라면 에이전트는 그 한 가지 미래에 딱 맞춘 꼼수를 찾기 쉽습니다. 확률 모델에서 샘플을 뽑아 매번 조금씩 다른 꿈을 꾸게 하면, 여러 미래에 두루 통하는 정책을 찾도록 압박할 수 있습니다. 온도는 그 압박의 세기를 정하는 손잡이입니다.

숫자로 해 보기(해설용 예): 같은 정책을 꿈속에서 3번 돌렸더니 점수가 900 · 950 · 1000이고, 실제 환경에서 3번 돌렸더니 200 · 180 · 210이라면, 이 정책은 꿈의 허점에 기대고 있다는 신호입니다. 온도를 올려 꿈속 점수가 600 · 700 · 800으로 떨어져도 실제 점수가 올라간다면 그쪽이 더 좋은 정책입니다. 위 '결과'의 표 2가 바로 이런 모양입니다.

5. 더 어려운 과제를 위한 반복 학습 절차

직관: 지도 없이 처음 간 도시는 동네 한 바퀴로는 다 알 수 없습니다. 돌아다니며 지도를 고치고, 고친 지도로 다시 돌아다니는 일을 되풀이해야 합니다.

단계(논문 5절):

1. M과 C를 무작위 매개변수로 시작합니다.
2. 실제 환경에서 $N$번 플레이하고 행동 · 관측을 모두 저장합니다.
3. M이 다음 관측 · 보상 · 행동 · 종료 여부를 예측하도록 학습하고, M 안에서 C가 기대 보상을 최대화하도록 학습합니다.
4. 과제를 아직 못 풀었으면 2로 돌아갑니다.

저자들은 이 논문의 단순한 과제는 한 번의 반복으로 충분했다고 밝힙니다. 또한 M의 예측 손실이 크면 "에이전트가 익숙하지 않은 곳에 왔다"는 뜻이므로, 실제 환경에서 M의 손실 부호를 뒤집어 보상으로 쓰면 낯선 곳을 탐험하도록 호기심을 줄 수 있다는 방향을 제안합니다(실험으로 보인 것은 아닙니다).

결과

자동차 경주(CarRacing-v0) — 100회 연속 평균 900점을 넘으면 '해결'로 봅니다.

조종기가 $z_t$만 보고 $h_t$를 못 보면 운전이 흔들리고 급커브에서 길을 벗어났습니다. M의 기억을 함께 주자 안정적으로 커브를 공략했습니다.

Doom Take Cover, 꿈속 학습 — 750 타임스텝 이상 살아남으면 해결입니다.

온도가 낮으면 꿈속 점수는 높지만 실제로는 무작위 정책 수준입니다(치트에 기댄 정책). $\tau = 1.15$에서 실제 점수가 가장 높았고, $\tau = 1.30$은 점수는 낮지만 변동이 작은 덜 위험한 전략을 찾았습니다. 본문에는 꿈에서 학습한 정책을 실제 환경에 넣었을 때 100회 평균 약 1100 타임스텝을 기록했다는 서술도 있습니다.

왜 중요한가 · 이후 영향

• '꿈에서 연습'을 실제로 보여 줬습니다. 학습한 모델 안에서만 정책을 만들고 실제 환경에 옮긴 사례로, 모델 기반 강화학습의 대중적 출발점이 됐습니다.
• 역할 분담 설계. 큰 모델은 비지도 학습으로, 보상 학습은 작게 — 이 분리는 이후 잠재 공간 모델 연구(Dreamer 등)로 이어졌습니다.
• 모델의 허점 문제를 일찍 짚었습니다. 학습한 시뮬레이터는 완벽하지 않고, 최적화가 그 허점을 찾아낸다는 관찰은 오프라인 강화학습 · 보상 모델 과최적화(RLHF의 '보상 해킹')를 이해할 때도 그대로 통합니다.

한계 · 주의할 점

• 과제가 단순합니다. 무작위 정책으로 모은 데이터만으로 세상을 충분히 볼 수 있는 과제였습니다. 저자들도 복잡한 환경에서는 탐험하며 데이터를 다시 모으고 모델을 다시 학습하는 반복 학습 절차가 필요하다고 씁니다.
• VAE는 과제에 중요한 것을 모릅니다. 보상과 상관없이 화면을 재현하도록 배우므로, 작지만 중요한 물체를 놓칠 수 있습니다(저자들도 Doom 화면의 일부 세부를 V가 제대로 담지 못한다고 언급합니다).
• 꿈속 점수를 믿으면 안 됩니다. 위 표처럼 꿈속 점수와 실제 점수가 크게 다를 수 있습니다.
• 조종기가 선형이라 표현력이 작습니다. 복잡한 결정은 V · M의 표현이 좋아야만 가능합니다.

이어서 읽을 것

• 뒤: Dreamer (잠재 공간 상상으로 행동 학습) · MuZero (관측을 복원하지 않는 모델)
• 비교 기준: A3C · DQN
• 배경: 변분 오토인코더 (VAE)

용어 정리

스스로 점검

1. 조종기 C를 일부러 작게 만든 이유는?

   — 보상 신호만으로 배워야 하는 부분을 작게 해, CMA-ES 같은 진화 전략으로 쉽게 최적화하고 신용 할당 문제를 줄이기 위해서입니다. 복잡한 표현은 비지도 학습하는 V와 M이 맡습니다.

2. 온도 $\tau = 0.10$에서 꿈속 점수는 높은데 실제 점수가 낮은 이유는?

   — 꿈이 너무 예측 가능해, 에이전트가 모델의 허점(불덩이가 안 생기는 움직임)을 이용한 정책을 찾았기 때문입니다. 실제 게임에는 그런 허점이 없습니다.

3. 조종기 매개변수가 867개가 되는 계산은?

   — 입력 $32 + 256 = 288$차원 × 행동 3차원 $= 864$개 가중치에 편향 3개를 더해 867개입니다.

다음에 볼 것

이 글을 이해했다면 이어서 보면 좋은 항목입니다.

• 잠재 공간 상상으로 행동 학습 (Dreamer) — 월드 모델의 잠재 공간에서 상상한 궤적으로 가치 기울기를 역전파해 행동을 배우는 에이전트

출처

• 원문: https://arxiv.org/abs/1803.10122 (2018 · David Ha 외)