실습 #3: CartPole World Model — MLP 기반 상태 예측과 상상 기반 계획

시리즈: World Model 실습 | 난이도: 입문 | 소요시간: 3~5분 (CPU)

---

1. 목적 (Why)

이 실습에서 증명하려는 것

"AI가 환경의 물리법칙을 학습하면, 실제로 행동하지 않고도 머릿속 상상만으로 최적의 행동을 찾을 수 있다."

CartPole(막대 세우기) 환경에서 다음을 검증합니다:

• 랜덤으로 모은 경험 데이터만으로 물리 모델(World Model)을 학습할 수 있는가?
• 학습된 모델 안에서 미래를 시뮬레이션(상상)하여 더 나은 행동을 선택할 수 있는가?
• "상상 기반 계획"이 랜덤 행동 대비 얼마나 성능이 향상되는가?

논문과의 연결

이 실습은 Dreamer v3 (Hafner et al., 2023)의 핵심 원리를 극도로 단순화한 것입니다:

개념	Dreamer v3	이 실습
World Model	RSSM (GRU + Stochastic)	MLP 3-layer
상태 표현	확률적 잠재 변수 (z_t)	CartPole 상태 벡터 (4차원)
계획 방식	Actor-Critic (정책 경사)	Random Shooting MPC
입력	고차원 이미지 (64x64)	수치 상태 (위치, 속도, 각도, 각속도)
학습 방식	시퀀스 기반 (시간 의존)	단일 스텝 독립

핵심 공통점: 둘 다 "현실 세계를 모방하는 내부 모델을 만들고, 그 안에서 상상하여 행동을 결정한다"는 동일한 원리입니다.

---

2. 전체 절차 (Pipeline)

┌─────────────┐     ┌─────────────┐     ┌─────────────┐     ┌─────────────┐
│  Step 1      │ ──→ │  Step 2      │ ──→ │  Step 3      │ ──→ │  Step 4      │
│  데이터 수집  │     │  모델 학습   │     │  상상+계획   │     │  시각화      │
│  (경험 축적)  │     │  (패턴 학습) │     │  (꿈속 탐색) │     │  (결과 분석) │
└─────────────┘     └─────────────┘     └─────────────┘     └─────────────┘
     ↓                    ↓                    ↓                    ↓
 22,263개 전이       World Model         랜덤 vs WM 비교      4종 차트 생성
 (s, a, s') 쌍      MSE 0.000033        11.4배 성능향상

---

3. 환경 설정 (Setup)

필수 패키지

Python 3.8+
gymnasium==0.29.1    # CartPole 환경
numpy>=1.24          # 수치 연산
torch>=2.0           # 신경망 학습
matplotlib>=3.7      # 시각화

설치 명령

pip install gymnasium numpy matplotlib torch

하드웨어

• CPU만으로 충분 (GPU 불필요)
• 메모리: 2GB 이상
• 디스크: 약 10MB (데이터 + 모델 + 차트)

---

4. 각 Step 상세

---

Step 1: 데이터 수집 (step1_collect_data.py)

목적

AI에게 "세상 경험"을 제공합니다. 아기가 세상을 마구잡이로 탐색하면서 인과관계를 기록하는 것과 같습니다.

Input → Process → Output

Input:  CartPole-v1 환경 (OpenAI Gymnasium)
        - 상태: [카트위치, 카트속도, 막대각도, 각속도] (4차원 실수)
        - 행동: 0(왼쪽) 또는 1(오른쪽)

Process: 1,000 에피소드 × 랜덤 행동
         매 스텝마다 (현재상태, 행동, 다음상태, 보상, 종료여부) 기록

Output:  collected_data.pkl (약 1.1MB)
         - states: (22263, 4) ndarray
         - actions: (22263,) ndarray
         - next_states: (22263, 4) ndarray
         - rewards: (22263,) ndarray
         - dones: (22263,) ndarray

핵심 로직

# 매 스텝의 기록 구조 (하나의 "전이")
data = {
    'states':      obs,           # s_t = [0.01, 0.03, -0.02, 0.01]
    'actions':     action,        # a_t = 1 (오른쪽)
    'next_states': next_obs,      # s_{t+1} = [0.02, -0.15, -0.01, 0.21]
}
# 이것이 World Model이 배울 핵심 데이터:
# f(s_t, a_t) → s_{t+1}

논문 연결

Dreamer v3에서는 환경과 상호작용하며 Replay Buffer에 경험을 저장합니다. 이 실습에서는 한 번에 모든 데이터를 수집하지만, 원리는 동일합니다. 모델 학습의 원재료가 되는 (state, action, next_state) 삼쌍입니다.

실행 & 예상 출력

$ python step1_collect_data.py

Step 1: CartPole 데이터 수집
환경: CartPole-v1
  상태 차원: 4 (위치, 속도, 각도, 각속도)
  행동 수: 2 (왼쪽, 오른쪽)

수집 결과:
  에피소드 수: 1000
  전이(transition) 수: 22263
  평균 에피소드 길이: 22.3 스텝

---

Step 2: World Model 학습 (step2_train_world_model.py)

목적

수집된 경험에서 환경의 물리법칙을 학습합니다. "공을 밀면 굴러간다"는 인과관계를 신경망이 스스로 발견하는 과정입니다.

Input → Process → Output

Input:  collected_data.pkl
        - 22,263개의 (s_t, a_t, s_{t+1}) 삼쌍

Process: MLP 신경망 학습
         입력: s_t(4) + a_t(2, one-hot) = 6차원
         출력: s_{t+1}(4차원) 예측
         손실함수: MSE(predicted, actual)
         학습: Adam lr=1e-3, 200 에폭, batch=512

Output:  world_model.pt (약 140KB)
         - model_state: 학습된 가중치
         - state_mean/std: 정규화 통계량
         - losses: 에폭별 손실 기록

모델 아키텍처

WorldModel (35,204 parameters)
┌──────────────────────────────────┐
│ Input: [s_t(4) + a_t_onehot(2)] │  6차원
│            ↓                     │
│    Linear(6 → 128) + ReLU       │  128차원
│            ↓                     │
│    Linear(128 → 128) + ReLU     │  128차원
│            ↓                     │
│    Linear(128 → 128) + ReLU     │  128차원
│            ↓                     │
│    Linear(128 → 4)              │  4차원
│            ↓                     │
│ Output: ŝ_{t+1} (예측 다음 상태) │
└──────────────────────────────────┘

핵심 로직: 왜 정규화하는가?

# 상태값 범위가 제각각이라 정규화 필수
# 카트 위치: -2.4 ~ 2.4
# 각속도:   -3.0 ~ 3.0 (범위가 훨씬 큼)

state_mean = states.mean(dim=0)     # 각 차원의 평균
state_std  = states.std(dim=0)      # 각 차원의 표준편차
states_norm = (states - state_mean) / state_std  # z-score 정규화

논문 연결

Dreamer v3의 RSSM(Recurrent State-Space Model)은 GRU + 확률적 잠재 변수를 사용하여 시간 의존성과 불확실성을 모두 표현합니다. 이 실습의 MLP는 그 극단적 단순화로, 시간 의존성도 확률성도 없지만 "상태 전이 예측"이라는 핵심 기능은 동일합니다.

Dreamer: h_t, z_t = RSSM(h_{t-1}, z_{t-1}, a_{t-1})
이 실습: ŝ_{t+1} = MLP(s_t, a_t)

실행 & 예상 출력

$ python step2_train_world_model.py

학습 시작: 200 에폭, 배치=512
  에폭  20/200: 손실 = 0.006543
  에폭  40/200: 손실 = 0.001234
  ...
  에폭 200/200: 손실 = 0.000033
최종 손실: 0.000033
손실 감소: 0.428451 → 0.000033 (99.99% 감소)
모델 저장: output/world_model.pt
  파라미터 수: 35,204개

---

Step 3: 상상 기반 계획 (step3_imagine_and_plan.py)

목적

학습된 World Model 안에서 미래를 상상하고, 최적의 행동을 찾습니다. 체스 고수가 머릿속으로 수를 읽는 것과 같은 원리입니다. Dreamer의 "Dream Learning" 핵심에 해당합니다.

Input → Process → Output

Input:  world_model.pt (학습된 모델)
        CartPole-v1 환경 (성능 테스트용)

Process: Random Shooting MPC (Model Predictive Control)
         매 스텝:
           1. 50개 랜덤 행동 시퀀스 생성 (horizon=10)
           2. 각 시퀀스를 World Model로 상상 시뮬레이션
           3. 막대가 가장 오래 서있는 시퀀스의 첫 행동 선택
           4. 실제 환경에서 해당 행동 실행
           → 반복

Output:  test_results.pt
         - random_scores: 랜덤 정책 50회 결과
         - wm_scores: World Model 정책 10회 결과

핵심 로직: Random Shooting MPC

def plan_action(model, current_state, horizon=10, num_candidates=50):
    best_score = -1
    best_first_action = 0

    for _ in range(num_candidates):  # 50개 후보
        actions = random(horizon=10)  # 랜덤 10스텝 행동 시퀀스

        # === 이 부분이 "상상" ===
        state = current_state
        score = 0
        for t in range(horizon):
            state = world_model(state, actions[t])  # 실제 환경 아님!
            if 막대가_서있는가(state):
                score += 1
            else:
                break
        # ========================

        if score > best_score:
            best_score = score
            best_first_action = actions[0]  # 첫 행동만 사용

    return best_first_action

이 과정에서 단 한 번도 실제 환경과 상호작용하지 않습니다. 50개 × 10스텝 = 500번의 "상상"을 통해 최적의 첫 행동 1개를 선택합니다.

논문 연결

측면	Dreamer v3	이 실습
상상	Learned policy가 RSSM 안에서 rollout	랜덤 시퀀스로 MLP 안에서 rollout
행동 선택	Actor network가 직접 행동 출력	50개 후보 중 최고 점수 선택
장점	연속 행동 공간, 빠른 추론	이해하기 쉬움
한계	복잡한 학습 필요	후보 수에 성능 의존

Dreamer의 Actor-Critic은 이 Random Shooting을 학습된 정책으로 대체하여 효율을 극적으로 높인 것입니다.

실행 & 예상 출력

$ python step3_imagine_and_plan.py

테스트: 랜덤 정책 vs World Model 계획
랜덤 정책 (50회 평균):
  평균 점수: 22.8 스텝
  최고 점수: 68 스텝

World Model 계획 정책 테스트 중... (10회)
  시도 1: 500 스텝 (상상 점수: 10)
  시도 2: 245 스텝 (상상 점수: 10)
  ...

World Model 계획 정책 결과:
  평균 점수: 260.7 스텝
  성능 향상: 11.4배

---

Step 4: 시각화 (step4_visualize.py)

목적

실험 결과를 4종 차트로 시각화하여 World Model의 학습 품질과 계획 효과를 정량적으로 분석합니다.

Output 차트 4종

차트 1 — Training Loss (학습 곡선)

X축: Epoch (0~200)
Y축: MSE Loss (log scale)
의미: 모델이 물리법칙을 얼마나 잘 학습했는가
결과: 0.428 → 0.000033 (99.99% 감소)
논문연결: Dreamer의 representation loss + dynamics loss에 대응

차트 2 — 1-Step Prediction Accuracy (예측 정확도)

X축: 실제 다음 상태값
Y축: 모델 예측값
이상적: 대각선 위에 모든 점이 위치 (predicted = actual)
결과: R² = 0.97~0.99 (4개 상태 차원 모두)
의미: 1스텝 전이를 거의 완벽하게 예측
논문연결: Dreamer의 decoder가 정확히 재구성하는 것에 대응

차트 3 — Multi-Step Imagination (상상 궤적)

실선: 실제 환경 궤적 (ground truth)
점선: World Model 상상 궤적 (40스텝)
빨간 영역: 오차 누적 구간
의미: 스텝이 길어질수록 오차 누적 (Dreamer가 짧은 horizon 사용하는 이유)
논문연결: Dreamer의 imagination horizon = 15 선택의 근거

차트 4 — Performance Comparison (성능 비교)

히스토그램: 랜덤 vs WM의 점수 분포
막대그래프: 평균 성능 비교 (22.8 vs 260.7 = 11.4배)
의미: "상상하는 능력"이 실제 성능으로 직결됨을 증명
논문연결: Dreamer의 zero-shot 일반화 성능 근거

실행

$ python step4_visualize.py

그래프 1 저장: 1_training_loss.png
그래프 2 저장: 2_prediction_accuracy.png
그래프 3 저장: 3_imagination_trajectory.png
그래프 4 저장: 4_performance_comparison.png

---

5. 실험 결과 종합

지표	수치	의미
데이터 크기	22,263 전이	약 1,000 에피소드 분량
모델 파라미터	35,204개	매우 작은 모델
학습 손실	0.000033	거의 완벽한 물리 학습
1-step R²	0.97~0.99	다음 상태를 97%+ 정확도로 예측
랜덤 정책	22.8 스텝	기준선 (baseline)
WM 계획 정책	260.7 스텝	11.4배 향상
상상 비용	500회/스텝	50후보 × 10스텝 시뮬레이션

---

6. 한계점과 다음 단계

이 실습의 한계

1. 수치 상태만 사용: 실제 환경은 이미지(픽셀)를 관측함 → 2-2에서 해결
2. 확정적(deterministic) 모델: 불확실성을 표현할 수 없음 → Dreamer의 RSSM이 필요
3. 시간 의존성 없음: 각 전이를 독립으로 취급 → RNN/GRU가 필요
4. 이산 행동 공간: 왼쪽/오른쪽 2개만 → 2-3에서 연속 행동으로 확장

학습 경로

이 실습 (2-1)          다음 실습 (2-2)            심화 (2-3)
MLP + 수치 상태   →    VAE + 이미지 입력   →    MuJoCo + 연속 행동
Random Shooting   →    Latent WM           →    CEM Planning
이산 행동(2개)    →    이산 행동(2개)       →    연속 행동(다차원)