실습 #6: MuJoCo 심화 — RSSM-Lite + Actor-Critic으로 Dreamer 재현

시리즈: World Model 실습 | 난이도: ★★★ 고급 | 소요시간: 실행 40분 + 분석 40분

---

1. 목적 (Why)

이 실습에서 하려는 것

"단순 MLP를 GRU 기반 RSSM으로, CEM Planning을 Actor-Critic으로 교체하여 Dreamer v3의 핵심 구조를 직접 구현한다."

실습 #5에서 Pendulum(1관절)을 다뤘다면, 이 실습에서는:

• 고차원 환경: HalfCheetah-v5 (17D 상태, 6D 행동, 6관절 보행)
• 시간 의존 모델: GRU로 시퀀스 정보를 유지하는 RSSM-Lite
• 확률적 잠재 변수: Prior/Posterior 분리 + KL Divergence
• Actor-Critic Planning: CEM 대신 학습된 정책으로 즉시 추론

논문과의 연결

Dreamer v3 (Hafner, 2023)     →  이 실습 (RSSM-Lite)
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
Eq.1: h_t = GRU(h_{t-1})      →  GRU(256) 시간 상태
Eq.2: z_t ~ Cat(logits)       →  z_t ~ N(μ, σ) 가우시안 (간소화)
Eq.3: x_t = decoder(h_t, z_t) →  state = decoder(h, z)
Eq.4: KL(posterior ‖ prior)    →  가우시안 KL loss 적용
Eq.6: V_λ(s_t) TD(λ) 리턴     →  compute_lambda_returns() 구현
Eq.7: Actor 목표               →  REINFORCE + advantage
Eq.8: Critic 목표              →  V(s) ≈ V_λ(s) MSE

선행 학습

실습	내용	이 실습과의 관계
#1 코드 해부	Dreamer v3 구조 이해	RSSM, Actor-Critic 설계 근거
#3 CartPole	MLP World Model	MLP → RSSM 진화 비교
#4 Vision	VAE 잠재 공간	Reparameterization trick 재활용
#5 Embodied AI	CEM Planning	CEM → Actor-Critic 진화 비교

---

2. 환경: HalfCheetah-v5

왜 HalfCheetah인가

Pendulum (실습 #5)                HalfCheetah (이 실습)
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
관절: 1개                         관절: 6개
상태: 3D (cosθ, sinθ, θ̇)         상태: 17D (관절각, 각속도, 중심속도 등)
행동: 1D 토크 [-2, 2]            행동: 6D 토크 [-1, 1]
물리: 단순 회전                   물리: 다관절 보행 역학
보상: 각도 기반                   보상: 전진 속도 - 제어 비용

HalfCheetah는 Dreamer 논문 시리즈의 표준 벤치마크 중 하나이며, MuJoCo 물리 시뮬레이터의 다관절 보행은 MLP로는 시간 의존성을 포착하기 어려워 GRU 기반 모델이 필요한 자연스러운 환경이다.

환경 구조

┌─────────── HalfCheetah-v5 ───────────┐
│                                       │
│  State (17D):                        │
│    - rootx position                   │
│    - rootz position                   │
│    - rooty angle                      │
│    - 6 joint angles (bthigh~ffoot)    │
│    - rootx velocity                   │
│    - rootz velocity                   │
│    - rooty angular velocity           │
│    - 6 joint angular velocities       │
│                                       │
│  Action (6D): 6 joint torques [-1,1] │
│  Reward: forward_speed - ctrl_cost    │
│  Episode: 500 steps (truncation)      │
└───────────────────────────────────────┘

---

3. 파이프라인 개요

Step 1: 데이터 수집              Step 2: RSSM-Lite 학습
  HalfCheetah-v5                   GRU + Prior/Posterior
  200 ep × 500 steps               시퀀스 기반 학습
  → 100K 전환                      KL + Recon + Reward loss
  → cheetah_data.pkl               → rssm_lite.pt (R²=0.9844)

Step 3: Actor-Critic 학습        Step 4: 시각화
  상상 속 rollout 15스텝            6종 차트
  TD(λ) 리턴 계산                  RSSM 학습곡선, 상상궤적
  Actor + Critic 동시 학습         AC 학습곡선, 성능비교
  → actor_critic.pt                시리즈 진화, Dreamer 달성도

---

4. Step 1: 데이터 수집

수집 설정

항목	설정
환경	HalfCheetah-v5
정책	Random (env.action_space.sample())
에피소드	200 × 최대 500 steps
수집량	100,000 전환
저장	cheetah_data.pkl

결과

총 전환: 100,000
평균 에피소드 보상: -143.3 ± 55.7
상태 shape: (100000, 17)
행동 shape: (100000, 6)

Random 정책의 평균 보상이 음수인 이유: 전진 보상을 얻으려면 조율된 다관절 운동이 필요한데, 랜덤 행동은 에너지만 소모(ctrl_cost)하고 전진하지 못한다.

---

5. Step 2: RSSM-Lite World Model

MLP → RSSM 진화의 핵심

실습 #5 (MLP):
  s_{t+1} = f(s_t, a_t)    ← 한 스텝만 본다

이 실습 (RSSM-Lite):
  h_t = GRU(h_{t-1}, [s_t, a_t])     ← 과거 전체를 기억
  z_t ~ N(μ(h_t), σ(h_t))             ← 불확실성 표현
  s_{t+1} = decoder(h_t, z_t)         ← 시간 + 확률 결합
  r_t = reward_head(h_t, z_t)         ← 보상 동시 예측

모델 아키텍처

┌────────────────── RSSM-Lite ──────────────────┐
│                                                │
│  입력: [s_t(17D), a_t(6D)] = 23D               │
│         ↓                                      │
│  GRU(23 → 256)     ← 시간 의존 상태 h          │
│         ↓                                      │
│  ┌─── Prior ───┐   ┌── Posterior ──┐           │
│  │ h → μ,σ     │   │ [h,obs] → μ,σ│           │
│  │ (관찰 없이) │   │ (관찰 반영)   │           │
│  └──────┬──────┘   └──────┬────────┘           │
│         │    KL Loss ←───→│                    │
│         ↓                 ↓                    │
│       z (32D) ← Reparameterization trick       │
│         ↓                                      │
│  ┌── Decoder ─────┐  ┌── Reward Head ──┐      │
│  │ [h,z] → s_{t+1}│  │ [h,z] → r_t    │      │
│  │ (288→256→256→17)│  │ (288→128→1)    │      │
│  └────────────────┘  └─────────────────┘      │
│                                                │
│  Loss = recon + 0.5×reward + 0.1×KL           │
│  Parameters: 565,906                           │
└────────────────────────────────────────────────┘

Prior vs Posterior — 왜 둘 다 필요한가

네트워크	입력	용도	비유
Prior	h만	상상 모드 (관찰 없이 예측)	눈 감고 예측
Posterior	h + 관찰	학습 모드 (관찰로 수정)	눈 뜨고 수정

KL Divergence가 하는 일: Prior를 Posterior에 가깝게 밀어준다. 결과적으로 상상 모드에서도 관찰이 있을 때만큼 정확한 예측이 가능해진다.

학습 결과

시퀀스: 5,000개 (길이 32)
에폭: 100
최종 Loss: train=0.01907, val=0.02102
상태 예측 R²: 0.9844

R²=0.9844는 17차원 상태 + 다관절 물리에서 매우 좋은 결과다. 이는 GRU의 시간 의존 상태 유지가 HalfCheetah의 관성/운동량 동역학을 효과적으로 포착하고 있음을 의미한다.

---

6. Step 3: Actor-Critic

CEM → Actor-Critic 진화의 핵심

CEM (실습 #5):
  매 스텝마다:
    64개 후보 × 10 horizon × 3 iterations = 1,920 forward
    → 가장 좋은 행동 1개 선택

Actor-Critic (이 실습):
  학습 후:
    1회 forward pass → 즉시 행동 출력
    → 실시간 제어 가능

Actor 네트워크

입력: features = [h_t, z_t] (256 + 32 = 288D)
       ↓
  MLP(288 → 256 → 256)  with SiLU
       ↓
  ┌── μ head ──┐  ┌── log_σ head ──┐
  │ → tanh(μ)  │  │ → clamp(-5, 0) │
  │ [-1, 1]    │  │ → exp(log_σ)   │
  └────────────┘  └────────────────┘
       ↓
  action ~ N(μ, σ) → tanh → [-1, 1]

Critic 네트워크

입력: features = [h_t, z_t] (288D)
       ↓
  MLP(288 → 256 → 256 → 1)  with SiLU
       ↓
  V(s): 상태 가치 (스칼라)

TD(λ) 리턴 — Dreamer v3 Eq.6

V_λ(s_t) = r_t + γ × [(1-λ) × V(s_{t+1}) + λ × V_λ(s_{t+1})]

γ=0.99 (할인율), λ=0.95 (TD와 MC 사이 보간)

• λ=0 → 순수 TD(1-step): V(st) = r_t + γV(s{t+1})
• λ=1 → 순수 Monte Carlo: V(st) = r_t + γr{t+1} + γ²r_{t+2} + ...
• λ=0.95 → Dreamer의 선택: 두 방법의 장점 결합

학습 방식 — 상상 속 학습

┌──── Imagination Rollout (15 steps) ────┐
│                                         │
│  1. 데이터에서 시작 상태 샘플 (B=32)     │
│  2. RSSM-Lite로 features 추출           │
│  3. Actor가 features → action 출력      │
│  4. RSSM-Lite가 action → 다음 상태 상상 │
│  5. 15스텝 반복 → 궤적 수집             │
│                                         │
│  Critic 학습: V(s) → TD(λ) 리턴 근사   │
│  Actor 학습: advantage × log_prob 최대화│
│                                         │
│  ※ 실제 환경 상호작용 ZERO             │
│  → 순수 상상 속에서 정책 개선            │
└─────────────────────────────────────────┘

학습 결과

학습 반복: 200
Actor 파라미터: 140,294
Critic 파라미터: 140,033

성능:
  Random:       -117.3 ± 44.1
  Actor-Critic: -139.6 ± 2.5

분산 비교:
  Random 분산:  44.1 (매우 불안정)
  Actor 분산:    2.5 (극도로 안정)

성능 분석

절대 수치에서 Actor가 Random을 이기지 못한 이유와 이 결과의 의미:

1. 데이터 한계: Random 정책으로만 수집한 100K 데이터로는 HalfCheetah의 전진 동역학을 충분히 학습할 수 없다. Dreamer v3 논문에서는 환경과 상호작용하며 지속적으로 데이터를 수집한다.

2. 학습량: 200 iteration은 매우 짧다. 실제 Dreamer v3는 수백만 스텝 학습한다.

3. 극적으로 낮은 분산 (2.5 vs 44.1): Actor는 "일관된 전략"을 학습했다. Random은 운 좋으면 높고 나쁘면 매우 낮지만, Actor는 항상 안정적인 행동을 출력한다. 이는 정책 학습이 작동하고 있다는 강력한 증거다.

4. Dreamer 전체 루프와의 차이: 실제 Dreamer는 "환경 상호작용 → 데이터 수집 → 모델 학습 → 상상 속 정책 학습" 을 반복하는 순환 구조다. 이 실습에서는 한 번의 데이터로 모델+정책을 학습했으므로 탐색의 피드백 루프가 없다.

---

7. 시각화 (6종 차트)

차트	내용	핵심 관찰
1. RSSM 학습 곡선	Recon/KL/Reward/Val Loss	모든 loss 안정적 수렴, KL이 초기에 빠르게 감소
2. 상상 궤적	실제 vs 예측 (30스텝)	단기(~10스텝) 정확, 장기에서 발산
3. AC 학습 곡선	Actor/Critic Loss 동시 추적	Critic이 먼저 수렴, Actor는 진동하며 개선
4. 성능 비교	Random vs Actor 박스플롯	Actor의 극도로 낮은 분산이 핵심
5. 시리즈 진화	2-1 → 3-2 전체 비교표	MLP→RSSM, RS→CEM→AC 진화 경로
6. Dreamer 달성도	15개 항목 중 9개 구현 (60%)	핵심 구조는 구현, 고급 기법은 미구현

---

8. 전체 시리즈 비교

│ 실습 │ 환경        │ 모델         │ Planning        │ 복잡도  │
├──────┼─────────────┼──────────────┼─────────────────┼─────────┤
│ #3   │ CartPole    │ MLP (35K)    │ Random Shooting │ ★☆☆   │
│ #4   │ CartPole    │ VAE+MLP      │ Random Shooting │ ★★☆   │
│ #5   │ Pendulum    │ MLP+SiLU     │ CEM             │ ★★☆   │
│ #6   │ HalfCheetah │ RSSM-Lite    │ Actor-Critic    │ ★★★   │
│      │             │ (566K)       │ (280K)          │         │

모델 진화 비교

MLP (실습 #3~5)              RSSM-Lite (실습 #6)
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
시간 처리: 없음 (1-step)     GRU로 과거 전체 요약
확률성: 없음 (결정적)        z ~ N(μ, σ) 불확실성 표현
Prior/Post: 없음             분리 → 상상 모드 가능
KL Loss: 없음               Prior ≈ Posterior 정렬
Planning: 탐색 기반          학습된 정책으로 즉시 추론
추론 비용: N×H forward       1 forward (실시간)

Planning 방법 진화

│ 방법            │ 실습  │ 추론 비용        │ 학습 필요 │
├─────────────────┼───────┼──────────────────┼───────────┤
│ Random Shooting │ #3,#4 │ N×H forward      │ ✗         │
│ CEM             │ #5    │ N×H×I forward    │ ✗         │
│ Actor-Critic    │ #6    │ 1 forward        │ ✓         │

---

9. Dreamer v3 달성도

구현 완료 (9/15)

• ✅ GRU 기반 시간 상태 (h_t)
• ✅ 잠재 변수 (z_t) 확률적 모델링
• ✅ Prior/Posterior 분리
• ✅ KL Divergence 손실
• ✅ 보상 예측 헤드
• ✅ Actor-Critic 학습
• ✅ TD(λ) 리턴 계산
• ✅ Imagination rollout
• ✅ SiLU 활성화 함수

미구현 (6/15)

• ⬜ 이산 카테고리컬 표현 (32×32 one-hot → 가우시안으로 대체)
• ⬜ Symlog 예측 (보상/가치 안정화 기법)
• ⬜ LayerNorm 전면 적용
• ⬜ CNN Encoder/Decoder (이미지 입력)
• ⬜ Continue 예측 (에피소드 종료 예측)
• ⬜ 단일 하이퍼파라미터 (task-agnostic)

달성도: 60% — 핵심 구조는 모두 구현

미구현 항목들은 성능 최적화를 위한 "고급 기법"이며, World Model의 근본 원리(시간 모델링 + 확률적 표현 + 상상 속 학습)는 이 실습에서 모두 체험했다.

---

10. 배운 점과 한계

핵심 인사이트

1. GRU의 힘: 단순 MLP로는 불가능한 시간 의존 동역학을 GRU가 자연스럽게 포착한다. HalfCheetah의 관성, 운동량 같은 물리 현상은 과거 상태의 축적이 필수적이다.

2. Prior/Posterior 분리의 우아함: 학습 시에는 관찰을 사용하고(Posterior), 상상 시에는 관찰 없이(Prior) 예측하는 이중 구조는 "모델을 학습하는 것"과 "모델을 사용하는 것"을 깔끔하게 분리한다.

3. 상상 속 학습의 효율성: Actor-Critic이 환경과 한 번도 상호작용하지 않고 World Model 안에서만 학습한다는 것은 놀라운 효율이다. 이것이 Dreamer가 sample-efficient한 핵심 이유다.

4. 분산의 의미: 성능의 절대치보다 분산이 더 중요할 수 있다. 실제 로봇에서는 "가끔 잘하는 정책"보다 "항상 안정적인 정책"이 훨씬 가치있다.

한계와 다음 단계

• 데이터 다양성: Random 정책 데이터만으로는 전진 동역학을 학습하기 어렵다. 다양한 수준의 정책 데이터가 필요하다.
• 학습-실행 루프: 실제 Dreamer는 환경 상호작용과 모델 학습을 반복한다. 이 순환 없이는 탐색이 제한된다.
• 이미지 입력: 현재는 상태 벡터를 직접 사용하지만, 실제 로봇은 카메라 입력을 사용한다. CNN Encoder 추가가 다음 과제다.

---

실행 방법

cd 실습코드/3-2_MuJoCo_심화

# Step 1: 데이터 수집 (~3분)
python step1_collect_mujoco.py

# Step 2: RSSM-Lite 학습 (~5분)
python step2_train_rssm_lite.py

# Step 3: Actor-Critic 학습 (~3분)
python step3_actor_critic.py

# Step 4: 시각화
python step4_visualize.py

필요 패키지

pip install gymnasium[mujoco] torch numpy matplotlib