실습 #2: CartPole World Model 설계 — Dreamer를 극단적으로 단순화하기

시리즈: World Model 실습 | 난이도: 설계 분석 | 소요시간: 읽기 20분

1. 목적 (Why)

이 실습에서 하려는 것

"Dreamer v3의 복잡한 구조를 CartPole이라는 가장 단순한 환경에 맞게 단순화하여, World Model의 본질만 추출한다."

실습 #1에서 Dreamer v3의 설계도를 읽었습니다. 하지만 RSSM, Actor-Critic, CNN Encoder 등 복잡한 구조를 처음부터 구현하는 것은 비현실적입니다. 이 실습에서는:

• Dreamer의 어떤 부분을 남기고 어떤 부분을 버릴 것인가?
• 버려도 되는 이유는 무엇인가?
• CartPole 환경의 특성이 어떤 단순화를 가능하게 하는가?

논문과의 연결

Dreamer v3 논문의 핵심 아이디어를 추출하고, 구현 복잡성은 제거합니다:

Dreamer v3의 핵심 아이디어 (보존):
  1. 세상의 물리법칙을 데이터에서 학습한다
  2. 학습된 모델 안에서 미래를 상상한다
  3. 상상을 통해 좋은 행동을 찾는다

Dreamer v3의 구현 복잡성 (제거):
  1. RSSM (GRU + 확률적 변수) → 단순 MLP
  2. CNN Encoder/Decoder → 필요 없음 (수치 상태 사용)
  3. Actor-Critic → Random Shooting MPC
  4. 이산 카테고리컬 표현 → 연속 상태 그대로
  5. Symlog 예측 → 표준 MSE

2. 설계 결정: 무엇을 남기고, 무엇을 버리나

결정 1: 환경 선택 — CartPole-v1

왜 CartPole인가?
┌──────────────────────────────────────────────┐
│ 상태: 4차원 (위치, 속도, 각도, 각속도)        │
│ 행동: 2개 (왼쪽, 오른쪽)                     │
│ 물리: 뉴턴 역학 (결정적, 예측 가능)           │
│ 에피소드: 짧음 (평균 ~22스텝)                 │
│ 보상: 단순 (서있으면 +1)                      │
└──────────────────────────────────────────────┘

CartPole의 특성이 가능하게 하는 단순화:
  - 상태가 4차원 수치 → CNN Encoder 불필요
  - 결정적 물리 → 확률적 z 불필요
  - 짧은 에피소드 → GRU 시간 기억 불필요
  - 이산 행동(2개) → Random Shooting으로 충분

결정 2: World Model — RSSM → MLP

Dreamer v3:
  상태 = (h, z)
  h_t = GRU(h_{t-1}, z_{t-1}, a_{t-1})  ← 시간 의존
  z_t ~ p(z_t | h_t)                     ← 확률적

이 실습:
  ŝ_{t+1} = MLP(s_t, a_t)               ← 단일 스텝, 확정적

남긴 것: 핵심 함수

f(s_t, a_t) → s_{t+1}   # 전이 함수
# 이것이 World Model의 본질
# "현재 상태와 행동이 주어지면 다음 상태를 예측한다"

결정 3: 계획 — Actor-Critic → Random Shooting MPC

Dreamer v3:
  Actor network(h, z) → action
  - 상상 궤적에서 정책 경사로 학습
  - 한번 학습하면 빠른 추론
  - 연속 행동 공간 지원

이 실습:
  Random Shooting MPC
  - N개 랜덤 행동 시퀀스 생성
  - 각각 World Model로 시뮬레이션
  - 최고 점수 시퀀스의 첫 행동 선택

왜 Random Shooting으로 충분한가?

CartPole: 행동 = {왼쪽, 오른쪽} = 2가지

10스텝 시퀀스의 모든 경우의 수 = 2^10 = 1,024
50개 랜덤 샘플 = 전체의 ~5% 탐색

→ 이산 행동 2개 × 짧은 horizon이라 Random Shooting으로도
  충분히 좋은 행동을 찾을 수 있음

연속 행동(예: 로봇 관절 각도)에서는 경우의 수가 무한이므로
Actor-Critic이나 CEM이 필수 → 실습 2-3에서 다룸

3. 4-Step 파이프라인 설계

Dreamer v3 → CartPole 매핑

         Dreamer v3                    이 실습
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

Step 1   환경 상호작용              데이터 수집
(데이터)  + Replay Buffer           (랜덤 정책 1000 에피소드)
         → 시퀀스 저장              → (s, a, s') 삼쌍 저장

Step 2   RSSM 학습                 MLP 학습
(모델)    + Encoder/Decoder         (6→128→128→128→4)
         + KL balancing            + MSE loss
         → 잠재공간 전이 학습        → 상태공간 전이 학습

Step 3   Imagination               Random Shooting
(계획)    + Actor-Critic 학습       (50후보 × 10스텝)
         → 상상 속 정책 최적화      → 상상 속 최고 행동 선택

Step 4   (논문에는 없음)            시각화
(분석)                              (학습곡선, 예측정확도,
                                    상상궤적, 성능비교)

각 Step의 Input/Output 설계

Step 1: 데이터 수집
  Input:  CartPole-v1 환경
  Output: collected_data.pkl
          {states(N,4), actions(N), next_states(N,4)}

Step 2: 모델 학습
  Input:  collected_data.pkl
  Output: world_model.pt
          {model_weights, normalization_stats, loss_history}

Step 3: 상상 + 계획
  Input:  world_model.pt + CartPole-v1 (테스트용)
  Output: test_results.pt
          {random_scores[50], wm_scores[10]}

Step 4: 시각화
  Input:  world_model.pt + test_results.pt
  Output: 4 PNG charts

4. 모델 아키텍처 상세 설계

WorldModel 클래스

class WorldModel(nn.Module):
    """
    Dreamer의 RSSM을 극단적으로 단순화.

    Dreamer RSSM:
      h_t = GRU(h_{t-1}, concat(z_{t-1}, a_{t-1}))
      z_t ~ p(z_t | h_t)
      loss = recon + KL(posterior || prior)

    이 실습:
      s_{t+1} = MLP(s_t, a_t)
      loss = MSE(predicted, actual)
    """

    Architecture:
      Linear(6 → 128) → ReLU
      Linear(128 → 128) → ReLU
      Linear(128 → 128) → ReLU
      Linear(128 → 4)

    Total: 35,204 parameters
    (Dreamer v3 Small: 22,000,000 parameters → 625배 작음)

설계 근거

5. Planning 알고리즘 설계

Random Shooting MPC 의사 코드

function PLAN_ACTION(world_model, current_state):
    best_score = -infinity
    best_action = 0

    for candidate in 1..50:                    # 50개 후보
        actions = RANDOM_SEQUENCE(length=10)    # 10스텝 행동 시퀀스
        state = current_state
        score = 0

        for t in 0..9:                          # 상상 시뮬레이션
            state = world_model.predict(state, actions[t])
            if IS_ALIVE(state):                 # 막대가 서있는가?
                score += 1
            else:
                break

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

    return best_action

하이퍼파라미터 선택 근거

6. 예상 결과와 검증 기준

성공 기준

실패 시 디버깅 가이드

문제: 학습 Loss가 안 줄어듦
→ 원인: 정규화 안 함, learning rate 너무 큼
→ 해결: Z-score 정규화 확인, lr를 1e-4로 낮춤

문제: WM 정책이 랜덤과 비슷함
→ 원인: horizon 너무 짧음, candidates 너무 적음
→ 해결: horizon 15, candidates 100으로 증가 (단, 시간 증가)

문제: 상상 궤적이 실제와 크게 다름
→ 원인: multi-step 오차 누적
→ 해결: 이것은 정상! Dreamer가 짧은 horizon을 사용하는 이유

7. 이 설계에서 배우는 것 → 다음 실습으로

실습 #1 (이론)        실습 #2 (설계) ← 현재       실습 #3 (구현)
Dreamer 설계도     →  CartPole 단순화 설계    →   실제 코드 실행
"어떻게 만들었나"     "무엇을 남기고 버릴까"       "돌려보고 검증"

                    이 실습에서 결정한 것:
                    ├── MLP 3-layer (128 hidden)
                    ├── MSE loss + Z-score 정규화
                    ├── Random Shooting (50×10)
                    └── 4종 시각화로 검증

                         ↓ 실습 #3에서 실행 ↓
                    결과: 11.4배 성능 향상 확인!

전체 실습 시리즈 흐름

#1 코드 해부 (이론)         → Dreamer의 설계도를 읽다
#2 CartPole 설계 (설계)     → 무엇을 남기고 버릴지 결정
#3 CartPole 실행 (구현)     → 수치 상태 World Model 검증 (11.4배)
#4 Vision 실행 (확장)       → 이미지 입력, VAE, 잠재공간 (R²=0.978)
#5 Embodied AI (심화)       → 연속 행동, 로봇 제어 (예정)