World Models (Ha & Schmidhuber, 2018) 논문 정리

논문명: World Models
저자: David Ha, Jürgen Schmidhuber
발표: 2018 (arXiv: 1803.10122)
인터랙티브 버전: worldmodels.github.io

한 줄 요약

"에이전트가 환경의 축소 모형(World Model)을 스스로 학습하고, 그 꿈(Dream) 속에서 정책을 훈련하여 현실에 전이할 수 있다."

핵심 아이디어

인간은 세상을 직접 경험하기 전에 머릿속에서 시뮬레이션합니다. 야구공이 날아올 때 우리는 뇌의 내부 모델로 궤적을 예측하고 미리 반응하죠. 이 논문은 이 개념을 AI 에이전트에 그대로 적용합니다.

핵심 질문: 에이전트가 환경의 생성 모델(generative model)을 학습하면, 그 모델 안에서만 훈련해도 실제 환경에서 작동하는 정책을 얻을 수 있을까?

답: Yes. 단, 제약 조건이 있습니다.

아키텍처: V-M-C 3모듈 구조

관측(x) ──→ [V: VAE] ──→ z(잠재벡터)
                            │
                            ▼
            [M: MDN-RNN] ──→ h(은닉상태) + P(z_next)
                            │
                     z + h  │
                            ▼
                       [C: Controller]
                            │
                            ▼
                       행동(a)

Vision (V) — VAE

• 역할: 고차원 이미지를 저차원 잠재 벡터 z로 압축
• 입력: 환경 프레임 이미지 (64x64)
• 출력: 잠재 벡터 z ∈ ℝ³² (32차원)
• 비유: 우리 눈의 시각 피질이 시각 정보를 핵심 특징으로 압축하는 과정

Memory (M) — MDN-RNN

• 역할: "다음에 뭐가 일어날까?"를 확률적으로 예측
• LSTM + Mixture Density Network
• 하나의 답이 아닌 여러 가능성의 확률 분포를 출력
• 예: "몬스터가 공격할 확률 60%, 가만히 있을 확률 40%"
• 비유: 해마(hippocampus)의 시퀀스 예측, 전전두엽의 미래 시뮬레이션

Controller (C) — 선형 정책

• 역할: z와 h를 받아 행동을 결정
• 놀랍도록 단순한 선형 모델 (파라미터 수백 개 수준)
• 설계 철학: 복잡성은 World Model(V+M)에 집중, 정책은 단순하게
• CMA-ES(진화 전략)로 학습

훈련 파이프라인

1. 데이터 수집 — 랜덤 정책으로 10,000회 롤아웃
2. V 학습 — 수집한 프레임으로 VAE 훈련
3. M 학습 — 인코딩된 z 시퀀스 + 행동으로 MDN-RNN 훈련
4. C 학습 — 실제 환경 또는 Dream 환경에서 CMA-ES로 최적화

Dream 학습 — 꿈 속 훈련 🌙

이 논문의 가장 흥미로운 기여입니다.

학습된 M을 환경 그 자체로 사용하여 C를 훈련합니다. 실제 게임을 안 켜도 M이 "게임처럼" 작동하는 거예요.

[일반 학습]    실제 게임 → 결과 확인 → 학습
[Dream 학습]   M이 만든 꿈 → 결과 확인 → 학습 → 실제 게임에 전이

인간이 잠을 자면서 낮의 경험을 "재생(replay)"하고 학습하는 것과 같은 원리입니다.

실험 결과

Dream 해킹 문제

에이전트가 World Model의 약점을 이용하여 현실에서는 불가능한 방식으로 높은 점수를 얻는 현상. 해결책으로 Temperature τ를 높여서 Dream을 더 어렵게 만듦.

핵심 기여점

한계점

1. 제한된 모델 용량 — LSTM 기반이라 복잡한 환경 학습에 한계
2. 순차적 훈련 — V→M→C 분리 훈련, end-to-end 최적화 불가
3. 단순한 환경 — 2D 환경에서만 검증, 3D 실세계 확장 미검증
4. Dream 해킹 — 모델이 불완전할수록 심화

후속 연구 연결

• Dreamer v1~v3 → RSSM으로 M 개선, Actor-Critic으로 C 개선
• JEPA (LeCun) → 생성이 아닌 예측에 집중, 잠재 공간에서 직접 예측
• Sora (OpenAI) → 비디오 생성을 World Simulation으로
• Google Genie → 단일 이미지에서 플레이 가능한 환경 생성

참고 링크

• 논문 원문 (arXiv)
• 인터랙티브 버전
• 코드 구현 (PyTorch)