실습 #4: Vision-Based World Model — 이미지에서 잠재 공간까지, Dreamer의 눈을 만들다

시리즈: World Model 실습 | 난이도: 중급 | 소요시간: 10~15분 (CPU)

---

1. 목적 (Why)

이 실습에서 증명하려는 것

"AI가 이미지(픽셀)로부터 세상의 핵심 구조를 추출하고, 그 압축된 공간에서 미래를 예측할 수 있다."

2-1에서는 환경이 친절하게 숫자 상태(위치, 속도, 각도, 각속도)를 제공했습니다. 하지만 현실 세계의 AI는 카메라 이미지(픽셀)만 볼 수 있습니다. 이 실습에서는:

• 5,000개 픽셀(50x100 이미지)을 32개 숫자로 압축할 수 있는가? (156배 압축)
• 압축된 공간(latent space)에서 "다음 장면"을 예측할 수 있는가?
• 여러 스텝을 상상하면 오차가 어떻게 누적되는가?

논문과의 연결

이 실습은 Dreamer v3의 인코더-디코더 파이프라인을 재현합니다:

        Dreamer v3                    이 실습
┌───────────────────────┐   ┌───────────────────────┐
│ 64x64 RGB 이미지      │   │ 50x100 Grayscale 이미지│
│       ↓ CNN Encoder   │   │       ↓ CNN Encoder    │
│ 잠재 변수 z_t (RSSM)  │   │ 잠재 벡터 z_t (VAE)    │
│       ↓ RSSM          │   │       ↓ MLP            │
│ 예측 z_{t+1}          │   │ 예측 z_{t+1}           │
│       ↓ CNN Decoder   │   │       ↓ CNN Decoder     │
│ 예측 이미지            │   │ 예측 이미지             │
└───────────────────────┘   └───────────────────────┘

핵심 차이점:

개념	Dreamer v3	이 실습
인코더	CNN → RSSM (확률적+결정적)	CNN → VAE (확률적)
잠재 공간	이산 카테고리컬 (32 × 32)	연속 가우시안 (32차원)
전이 모델	GRU + MLP (시간 의존)	MLP only (단일 스텝)
디코더	TransposedCNN	TransposedCNN
학습 신호	Recon + KL + Reward + Continue	Recon + KL

핵심 공통점: 둘 다 "고차원 이미지를 저차원 잠재 공간으로 압축 → 잠재 공간에서 전이 예측 → 디코딩으로 복원"이라는 동일한 3단계 구조입니다.

---

2. 전체 절차 (Pipeline)

┌──────────────┐     ┌──────────────┐     ┌──────────────┐     ┌──────────────┐
│  Step 1       │ ──→ │  Step 2       │ ──→ │  Step 3       │ ──→ │  Step 4       │
│  이미지 수집   │     │  VAE 학습     │     │  Latent WM    │     │  시각화       │
│  (카메라 촬영) │     │  (압축 학습)  │     │  (잠재 예측)  │     │  (결과 분석)  │
└──────────────┘     └──────────────┘     └──────────────┘     └──────────────┘
     ↓                     ↓                     ↓                     ↓
 6,318 프레임         5000px → 32dim       z+action → z'        4종 차트 생성
 (50x100 gray)        156배 압축             R²=0.9779

---

3. 환경 설정 (Setup)

필수 패키지

Python 3.8+
gymnasium==0.29.1    # CartPole 환경
numpy>=1.24          # 수치 연산
torch>=2.0           # 신경망 (CNN, VAE)
matplotlib>=3.7      # 시각화
pygame>=2.5          # CartPole 렌더링 (이미지 캡처용)

설치 명령

pip install gymnasium numpy matplotlib torch pygame

하드웨어

• CPU 가능, GPU 있으면 VAE 학습 2~3배 빨라짐
• 메모리: 4GB 이상 (이미지 데이터 ~242MB)
• 디스크: 약 250MB (이미지 데이터 + 모델 + 차트)

2-1과의 차이

항목	2-1 CartPole	2-2 Vision
pygame 필요	X	O (렌더링용)
메모리 사용	~2MB	~250MB
학습 시간	1~2분	8~12분
모델 복잡도	MLP 35K params	VAE 530K + WM 42K params

---

4. 각 Step 상세

---

Step 1: 이미지 데이터 수집 (vision_step1_collect_images.py)

목적

CartPole 환경을 카메라로 촬영하여 이미지 데이터셋을 구축합니다. 2-1에서는 환경이 숫자를 제공했지만, 여기서는 AI가 "눈"으로 보는 것처럼 픽셀만 제공합니다.

Input → Process → Output

Input:  CartPole-v1 환경 (render_mode='rgb_array')
        → 매 스텝 400x600x3 RGB 이미지 생성

Process:
  1. 환경 렌더링 → 원본 이미지 (400x600 RGB)
  2. 중앙 크롭 (150:350, 100:500) → 200x400 영역
  3. 평균 풀링 다운샘플링 → 50x100
  4. Grayscale 변환 → 단일 채널
  5. 0~1 정규화
  6. (현재이미지, 행동, 다음이미지) 삼쌍으로 저장

  300 에피소드 × 랜덤 행동 → 총 6,318개 프레임

Output:  vision_data.pkl (약 242MB)
         - frames: (6318, 50, 100) float32
         - actions: (6318,) int
         - next_frames: (6318, 50, 100) float32

이미지 전처리 파이프라인

원본 렌더링           중앙 크롭            다운샘플             그레이스케일
400×600×3 RGB  →  200×400×3 RGB  →  50×100×3 RGB  →  50×100×1 Gray
                   (핵심 영역)       (평균 풀링)        (0~1 정규화)

총 픽셀: 5,000개/프레임

핵심 로직: 왜 이미지 전처리를 하는가?

# 원본 400x600x3 = 720,000 픽셀 → 너무 큼
# 1. 카트+막대가 있는 중앙만 크롭 (정보 밀도 높임)
img = frame[150:350, 100:500, :]  # 200x400

# 2. 다운샘플링 (계산량 감소, 핵심 구조 보존)
img_small = 평균풀링(img, 50, 100)  # 50x100

# 3. 그레이스케일 (색 정보 불필요, 메모리 절약)
gray = 0.299*R + 0.587*G + 0.114*B  # ITU-R 601 표준

# 최종: 5,000 픽셀/프레임 (원본 대비 144배 축소)

논문 연결

Dreamer v3는 64x64 RGB(12,288 픽셀)를 사용합니다. 이 실습은 50x100 Grayscale(5,000 픽셀)로 더 작지만, 핵심 원리는 동일합니다: 고차원 이미지를 효율적으로 처리할 수 있는 크기로 전처리하는 것이 첫 번째 단계입니다.

실행 & 예상 출력

$ python vision_step1_collect_images.py

Vision Step 1: 이미지 데이터 수집
  에피소드  50/300 완료 — 누적 1035개 프레임
  에피소드 100/300 완료 — 누적 2090개 프레임
  ...
총 6318개 (image, action, next_image) 쌍 수집
이미지 크기: (50, 100) (50x100 grayscale)
저장: output/vision_data.pkl (242.1 MB)

---

Step 2: VAE 학습 (vision_step2_train_vae.py)

목적

5,000개 픽셀을 32개 숫자(잠재 벡터)로 압축하는 "인코더"와, 32개 숫자에서 이미지를 복원하는 "디코더"를 동시에 학습합니다. 이것이 Dreamer의 "눈"에 해당합니다.

Input → Process → Output

Input:  vision_data.pkl
        - 6,318개 프레임 (50×100 grayscale)

Process: VAE (Variational Autoencoder) 학습
         Encoder: CNN 3층 → mu, logvar (32차원)
         Reparameterize: z = mu + sigma * epsilon
         Decoder: ConvTranspose 3층 → 재구성 이미지
         손실 = 재구성 MSE + 0.5 × KL Divergence
         학습: Adam lr=1e-3, 50 에폭, batch=128

Output:  vae_model.pt (약 3.6MB)
         - model_state: 전체 VAE 가중치
         - all_z: (6318, 32) 전체 프레임의 잠재 벡터
         - all_mu: (6318, 32) 전체 프레임의 평균 벡터
         - losses: 에폭별 (recon, kl) 기록

VAE 아키텍처 상세

                    Encoder (CNN)                    Decoder (TransposedCNN)
               ┌─────────────────────┐          ┌─────────────────────────┐
 50×100×1 ──→ │ Conv(1→16, k=4, s=2)│          │ Linear(32 → flat)       │
               │ ReLU                │          │ Reshape → 6×12×64       │
               │ Conv(16→32, k=4,s=2)│          │ ConvT(64→32, k=4, s=2) │
               │ ReLU                │          │ ReLU                    │
               │ Conv(32→64, k=4,s=2)│          │ ConvT(32→16, k=4, s=2) │
               │ ReLU                │          │ ReLU                    │
               │ Flatten → Linear    │          │ ConvT(16→1, k=4, s=2)  │
               └────────┬────────────┘          │ Sigmoid                │
                        ↓                       └───────────┬─────────────┘
                    ┌────┴────┐                             ↑
                    │ mu (32) │                             │
                    │logvar(32)│                             │
                    └────┬────┘                             │
                         ↓                                  │
               z = mu + σ × ε  ─────────────────────────→ z (32차원)
               (reparameterization trick)

총 파라미터: 529,633개
압축률: 5,000 픽셀 → 32 차원 = 156배

핵심 로직: Reparameterization Trick

def reparameterize(mu, logvar):
    """
    왜 이렇게 복잡하게?
    단순히 z = encoder(image) 하면 "미분 가능"하지 않아서 역전파가 안 됨.
    z = mu + std * random_noise 로 하면:
      - mu와 logvar에 대해 미분 가능 (학습 가능)
      - random_noise가 확률성을 제공 (같은 이미지도 약간 다른 z)
      - 이것이 VAE의 핵심 아이디어

    Dreamer v3와의 관계:
      Dreamer는 이산 카테고리컬 분포 + straight-through gradient 사용
      VAE는 연속 가우시안 분포 + reparameterization trick 사용
      둘 다 "확률적 잠재 변수를 미분 가능하게 만드는" 방법
    """
    std = exp(0.5 * logvar)   # 표준편차
    eps = randn_like(std)      # 랜덤 노이즈
    return mu + eps * std      # 샘플링이지만 미분 가능!

핵심 로직: VAE 손실 함수

# 두 가지 목표를 동시에 최적화:
loss = reconstruction_loss + 0.5 * kl_loss

# 1. 재구성 손실: "원본과 복원이 얼마나 비슷한가?"
recon_loss = MSE(재구성_이미지, 원본_이미지)
# → 이것만 쓰면 압축 품질은 좋지만, 잠재 공간이 엉망

# 2. KL 발산: "잠재 분포가 정규분포와 얼마나 가까운가?"
kl_loss = -0.5 * sum(1 + logvar - mu² - exp(logvar))
# → 잠재 공간을 부드럽고 연속적으로 만듦 (보간, 샘플링 가능)
# → beta=0.5로 가중치 낮춤 (재구성 품질 우선)

논문 연결

Dreamer v3의 인코더도 CNN → 잠재 변수 구조이지만, 가우시안 대신 이산 카테고리컬 분포(32 classes × 32 categories = 1024 가능한 상태)를 사용합니다. 이유는 이산 표현이 더 표현력이 높고, 모드 붕괴(mode collapse)가 적기 때문입니다. 이 실습의 VAE는 그 단순화 버전입니다.

실행 & 예상 출력

$ python vision_step2_train_vae.py

VAE 파라미터: 529,633개
잠재 공간: 32차원
압축률: 5000 → 32 (156배 압축)

학습 시작 (50 에폭)...
  에폭  10 | 총 Loss: 8.1234 | 재구성: 6.2345 | KL: 3.7778
  에폭  20 | 총 Loss: 6.5432 | 재구성: 5.1234 | KL: 2.8396
  ...
  에폭  50 | 총 Loss: 6.9843 | 재구성: 4.5523 | KL: 4.9320

모델 저장: output/vae_model.pt
잠재 벡터: torch.Size([6318, 32])

---

Step 3: Latent World Model (vision_step3_latent_world_model.py)

목적

VAE가 압축한 잠재 공간(32차원)에서 "다음 상태"를 예측합니다. 5,000 픽셀에서 직접 예측하는 것이 아니라, 32개 숫자에서 예측하므로 156배 효율적입니다.

Input → Process → Output

Input:  vae_model.pt (VAE 인코더 + 잠재 벡터)
        vision_data.pkl (행동 데이터)

Process:
  1. 현재 프레임의 z_t (VAE 인코더로 추출)
  2. 다음 프레임의 z_{t+1} (VAE 인코더로 추출)
  3. Latent WM 학습: f(z_t, a_t) → z_{t+1}
     MLP: (32+2) → 128 → 128 → 128 → 32
     학습: Adam lr=1e-3, 100 에폭, StepLR 스케줄러

Output:  latent_world_model.pt (약 167KB)
         - model_state: Latent WM 가중치
         - losses: 에폭별 MSE
         - r2: 1-step 예측 R² 값

모델 아키텍처

Latent World Model (41,632 parameters)
┌────────────────────────────────────┐
│ Input: [z_t(32) + a_t_onehot(2)]  │  34차원
│            ↓                       │
│    Linear(34 → 128) + ReLU        │  128차원
│            ↓                       │
│    Linear(128 → 128) + ReLU       │  128차원
│            ↓                       │
│    Linear(128 → 128) + ReLU       │  128차원
│            ↓                       │
│    Linear(128 → 32)               │  32차원
│            ↓                       │
│ Output: ẑ_{t+1} (예측 다음 잠재벡터)│
└────────────────────────────────────┘

핵심 로직: 왜 잠재 공간에서 예측하는가?

방법 A (직접 예측):  50×100 이미지 → MLP → 50×100 이미지
  - 입력 5,000차원, 출력 5,000차원
  - 파라미터 수: 수천만 개 필요
  - 학습 어려움, 과적합 위험

방법 B (잠재 예측):  이미지 → Encoder → z(32) → MLP → z'(32) → Decoder → 이미지
  - 입력 34차원, 출력 32차원
  - 파라미터 수: 41,632개로 충분
  - 핵심 구조만 예측 (노이즈 무시)
  - → Dreamer가 이 방법을 선택한 이유!

논문 연결

Dreamer v3:  h_t, z_t = RSSM(h_{t-1}, z_{t-1}, a_{t-1}, image_t)
                        ↓
             GRU(h)가 시간 의존성 담당 + MLP(z)가 상태 전이 담당

이 실습:    z_{t+1} = MLP(z_t, a_t)
                        ↓
             시간 의존성 없이 단일 스텝만 예측

차이: RSSM은 과거 기억(h)을 유지하여 부분 관측 문제를 해결.
     이 실습은 현재 관측만으로 충분한 CartPole이라 MLP로도 가능.

실행 & 예상 출력

$ python vision_step3_latent_world_model.py

현재 z: torch.Size([6318, 32]), 다음 z: torch.Size([6318, 32])
Latent World Model 파라미터: 41,632개

학습 시작 (100 에폭)...
  에폭  20 | Loss: 0.001234
  에폭  40 | Loss: 0.000567
  ...
  에폭 100 | Loss: 0.000234

잠재 공간 1-step 예측 R²: 0.9779

---

Step 4: 시각화 (vision_step4_visualize.py)

목적

VAE의 압축 품질, World Model의 예측 정확도, 그리고 다중 스텝 상상의 오차 누적을 시각적으로 분석합니다.

Output 차트 4종

차트 1 — VAE Reconstruction (압축 품질)

상단: 원본 이미지 6장 (다양한 시점)
하단: VAE로 압축(32차원) 후 복원한 이미지
의미: 5,000 픽셀 → 32 숫자 → 5,000 픽셀 복원이 얼마나 정확한가
결과: 카트와 막대의 위치/형태가 보존됨, 세부 텍스처만 흐려짐
논문연결: Dreamer의 decoder 품질 평가와 동일한 방법

차트 2 — Training Curves (학습 곡선)

좌: VAE 학습 — Reconstruction Loss (↓) + KL Divergence
우: Latent WM 학습 — MSE Loss (log scale) + R²=0.9779
의미: 두 모델 모두 수렴하여 안정적으로 학습됨
논문연결: Dreamer 논문의 Figure 2 (training curves)와 유사한 형태

차트 3 — 1-Step Prediction (단일 스텝 예측)

행 1: 현재 프레임 (실제)
행 2: 실제 다음 프레임 (ground truth)
행 3: 모델 예측 다음 프레임 (Encoder → WM → Decoder)
의미: 1스텝 예측은 거의 완벽 (R²=0.9779)
논문연결: Dreamer의 1-step reconstruction accuracy

차트 4 — Multi-Step Imagination (다중 스텝 상상)

행 1: 실제 프레임 (t+0 ~ t+7)
행 2: 모델 상상 프레임 (8스텝 연속 예측)
행 3: 차이 히트맵 (빨간색 = 오차 큰 영역)

핵심 관찰:
  t+0~t+2: 거의 완벽 (차이 히트맵이 거의 검정)
  t+3~t+5: 약간의 흐려짐 시작
  t+6~t+7: 오차 누적으로 막대 위치 벗어남

의미: 이것이 Dreamer가 imagination horizon을 15스텝으로
     제한하는 근본적인 이유!
논문연결: Dreamer v3 Section 4.2 "Imagination Horizon"

실행

$ python vision_step4_visualize.py

그래프 1: VAE 재구성 비교... → v1_vae_reconstruction.png
그래프 2: 학습 곡선...       → v2_training_curves.png
그래프 3: 1-step 예측...     → v3_one_step_prediction.png
그래프 4: Multi-step 상상... → v4_multistep_imagination.png

---

5. 실험 결과 종합

지표	수치	의미
수집 프레임	6,318개	300 에피소드 분량
이미지 크기	50×100 grayscale	5,000 픽셀/프레임
VAE 잠재 차원	32	156배 압축
VAE 파라미터	529,633개	인코더+디코더
VAE Recon Loss	4.5523	이미지 복원 오차
VAE KL Loss	4.9320	잠재 분포 정규화
Latent WM 파라미터	41,632개	잠재 공간 전이 모델
Latent WM R²	0.9779	97.8% 정확도
1-step 예측	거의 완벽	차이 히트맵 미미
8-step 상상	오차 누적 관찰	horizon 제한 필요성 확인

2-1 vs 2-2 비교

항목	2-1 (수치)	2-2 (비전)
입력	4차원 벡터	5,000 픽셀
모델 크기	35K params	571K params (16배)
데이터 크기	1.1MB	242MB (220배)
학습 시간	~2분	~12분 (6배)
예측 정확도	R²=0.99+	R²=0.978
핵심 기여	World Model 개념	이미지→잠재공간 파이프라인

---

6. 전체 데이터 흐름도 (End-to-End)

Step 1                Step 2                 Step 3                Step 4
카메라 촬영            VAE 학습               Latent WM 학습        시각화

CartPole           ┌──────────────┐      ┌──────────────┐
  환경  ───────→   │  Encoder     │  →   │  MLP         │
(render)          │  50×100→z(32)│      │  z+a→z'(32)  │
  ↓                │  CNN 3-layer │      │  3-layer     │
  이미지            └──────────────┘      └──────────────┘
  50×100              ↕ 학습                 ↕ 학습
  gray             ┌──────────────┐
                   │  Decoder     │
                   │  z(32)→50×100│
                   │  DeConv 3-lay│
                   └──────────────┘

상상 시 데이터 흐름 (Step 4에서 확인):
이미지 → Encoder → z_0 → WM → z_1 → WM → z_2 → ... → z_n → Decoder → 상상 이미지
                          ↑a_0      ↑a_1           ↑a_{n-1}

---

7. 한계점과 다음 단계

이 실습의 한계

1. 단일 스텝 전이만 학습: 시간 의존성 없음 → RSSM(GRU)이 필요
2. Planning 미포함: VAE + WM만 학습, 실제 행동 선택 없음 → Actor-Critic 필요
3. 오차 누적: 8스텝 이후 상상 품질 급락 → 짧은 horizon + 재계획 필요
4. 이산 행동만: 왼쪽/오른쪽 2개 → 연속 행동 공간으로 확장 필요

이 실습에서 배운 것 → Dreamer로의 다리

배운 것                          Dreamer에서의 활용
────────────────────────────────────────────────────
이미지 → 잠재벡터 압축   →    RSSM의 Encoder (CNN)
잠재 공간에서 전이 예측   →    RSSM의 Dynamics Model
디코더로 이미지 복원      →    RSSM의 Decoder
오차 누적 문제 확인       →    Imagination Horizon 설정
KL 발산으로 정규화        →    Dreamer의 KL balancing

학습 경로

2-1 (입문)              2-2 (중급) ← 현재        2-3 (심화)
수치 상태 + MLP    →    이미지 + VAE + MLP   →    MuJoCo + 연속행동
Random Shooting    →    잠재 예측만           →    CEM + Actor-Critic
이산 행동(2개)     →    이산 행동(2개)        →    연속 행동(다차원)
물리 학습 검증     →    시각 압축 검증        →    로봇 제어 검증