[ 트렌드] 실습 #1: Dreamer v3 코드 해부 — 꿈속에서 배우는 AI의 설계도를 읽다
관
관리자
Lv.1
03-02 09:49
·
조회 20
·
추천 0
실습 #1: Dreamer v3 코드 해부 — 꿈속에서 배우는 AI의 설계도를 읽다
시리즈: World Model 실습 | 난이도: 이론 분석 (코드 실행 없음) | 소요시간: 읽기 30분
1. 목적 (Why)
이 실습에서 하려는 것
"Dreamer v3의 공식 코드를 열어보고, 논문의 수식이 실제로 어떻게 구현되었는지 파일 단위로 해부한다."
논문만 읽으면 추상적인 개념(RSSM, Actor-Critic, Imagination)이 코드에서 어떤 파일, 어떤 클래스, 어떤 함수로 구현되는지 알 수 없습니다. 이 실습에서는:
- Dreamer v3의 코드 구조가 논문의 어떤 개념에 대응하는가?
- 각 파일이 뇌의 어떤 부위에 비유할 수 있는가?
- v1 → v2 → v3으로 무엇이 바뀌었고 왜 바뀌었는가?
논문 정보
- 논문: "Mastering Diverse Domains through World Models" (Hafner et al., 2023)
- 공식 코드: github.com/danijar/dreamerv3
- 관련 시리즈: #6 (Dreamer v1~v3 진화), #7 (핵심 용어 해설)
2. Dreamer v3 전체 아키텍처
논문의 Figure 1을 코드로 매핑
Dreamer v3 — 전체 구조
┌─────────────────────────────────────────────────────────┐
│ │
│ ┌──────────┐ ┌──────────┐ ┌──────────┐ │
│ │ Encoder │ ──→ │ RSSM │ ──→ │ Decoder │ │
│ │ (시각피질)│ │(해마+피질)│ │ (재구성) │ │
│ └──────────┘ └────┬─────┘ └──────────┘ │
│ ↑ │ │
│ 관찰(이미지) ┌───┴───┐ │
│ │ 잠재상태 │ │
│ │ (h, z) │ │
│ └───┬───┘ │
│ │ │
│ ┌───────────┴───────────┐ │
│ │ Imagination │ │
│ │ (RSSM 안에서 상상) │ │
│ └───────────┬───────────┘ │
│ │ │
│ ┌──────────┴──────────┐ │
│ │ Actor │ Critic │ │
│ │ (행동선택) │ (가치평가)│ │
│ │ (기저핵) │ (전두엽) │ │
│ └──────────┴──────────┘ │
│ │
│ ┌──────────┐ │
│ │ Replay │ 경험 저장소 (에피소드 기억) │
│ │ Buffer │ │
│ └──────────┘ │
└─────────────────────────────────────────────────────────┘3. 코드 파일 ↔ 논문 개념 ↔ 뇌 부위 매핑
파일별 상세 해부
| 파일 | 논문 개념 | 뇌 비유 | 핵심 역할 |
|---|---|---|---|
agent.py |
Algorithm 1 (전체 루프) | 전전두엽 (PFC) | 전체 학습 루프 조율 |
nets.py (RSSM) |
Eq. 1~3 (상태 전이) | 해마 + 감각피질 | 세계 모델의 핵심 |
nets.py (Encoder) |
Eq. 4 (관찰 인코딩) | 시각 피질 (V1~V4) | 이미지 → 잠재 벡터 |
nets.py (Decoder) |
Eq. 5 (이미지 재구성) | 연상 피질 | 잠재 벡터 → 이미지 |
behaviors.py |
Eq. 6~8 (Actor-Critic) | 기저핵 + 전두엽 | 행동 정책 + 가치 평가 |
replay.py |
Experience Replay | 에피소드 기억 | 경험 저장/샘플링 |
3-1. agent.py — 전체 조율자 (전전두엽)
역할
논문의 Algorithm 1을 구현합니다. 환경과 상호작용하며 데이터를 모으고, World Model을 학습하고, Actor-Critic을 업데이트하는 전체 루프를 관리합니다.
핵심 흐름
# agent.py의 핵심 구조 (의사 코드)
class Agent:
def train(self, data):
# 1단계: World Model 학습 (현실을 배움)
world_model_loss = self.world_model.train(data)
# 2단계: 상상 (World Model 안에서 미래 생성)
imagined_trajectories = self.world_model.imagine(
policy=self.actor,
horizon=15 # 15스텝 미래까지 상상
)
# 3단계: Actor-Critic 학습 (상상 속에서 정책 개선)
actor_loss = self.actor.train(imagined_trajectories)
critic_loss = self.critic.train(imagined_trajectories)논문 연결
-
Algorithm 1, Line 5:
world_model.train(data)→ 세계 모델 파라미터 업데이트 -
Algorithm 1, Line 7:
world_model.imagine()→ 학습된 모델 안에서 rollout -
Algorithm 1, Line 8-9:
actor/critic.train()→ 상상된 궤적으로 정책 학습
뇌 비유
전전두엽(PFC)처럼 "지금은 경험을 모을 때인지, 학습할 때인지, 행동할 때인지"를 판단하고 조율합니다.
3-2. nets.py — RSSM (해마 + 감각피질)
역할
Dreamer v3의 핵심 중의 핵심. 세상의 물리법칙을 배우는 World Model입니다.
RSSM 구조
RSSM = Recurrent State-Space Model
상태를 2개로 분리:
┌─────────────────────────────────────────────┐
│ h (deterministic, 확정적) │
│ - GRU로 업데이트 │
│ - "지금까지 확실히 알고 있는 것" │
│ - 비유: "내가 지금 고속도로 2차선에 있다" │
│ │
│ z (stochastic, 확률적) │
│ - 확률 분포에서 샘플링 │
│ - "예측은 하지만 불확실한 것" │
│ - 비유: "앞차가 멈출 수도, 안 멈출 수도" │
└─────────────────────────────────────────────┘핵심 수식 → 코드
# 논문 Eq. 1: 확정적 상태 전이
h_t = GRU(h_{t-1}, concat(z_{t-1}, a_{t-1}))
# → nets.py: self.gru(concat(z, a), h)
# 논문 Eq. 2: 사전 분포 (관찰 없이 예측)
z_t ~ p(z_t | h_t)
# → nets.py: prior = self.prior_net(h)
# 상상할 때 사용 (이미지 없이 예측)
# 논문 Eq. 3: 사후 분포 (관찰 반영)
z_t ~ q(z_t | h_t, o_t)
# → nets.py: posterior = self.posterior_net(concat(h, encoded_obs))
# 학습할 때 사용 (이미지 보고 수정)v1 → v3 핵심 변화
v1 (2019): z = 연속 가우시안 (30차원)
→ 문제: 모드 붕괴, 불안정
v2 (2021): z = 이산 카테고리컬 (32 classes × 32 categories)
→ 해결: 더 풍부한 표현, 안정적 학습
v3 (2023): z = 이산 카테고리컬 + Symlog + KL Balancing
→ 해결: 도메인 무관 안정 학습
→ 단일 하이퍼파라미터(size)만 조절논문 연결
- Section 3.1 "World Model": RSSM의 전체 구조 정의
- Eq. 1~3: h, z의 전이 및 분포
- Appendix B: GRU의 구체적 구현 (hidden size = model_size)
뇌 비유
- h (GRU) = 해마의 시간 기억: 순차적 경험을 기억
- z (stochastic) = 감각 피질의 불확실한 해석: 같은 장면도 다르게 해석 가능
3-3. nets.py — Encoder / Decoder (시각피질 / 연상피질)
Encoder: 이미지 → 잠재 벡터
64×64×3 RGB 이미지
↓
CNN 4층 (Conv → ReLU → ...)
↓
Flatten → Linear
↓
잠재 벡터 (RSSM의 posterior에 전달)Decoder: 잠재 벡터 → 이미지 재구성
(h, z) 결합
↓
Linear → Reshape
↓
TransposedCNN 4층
↓
64×64×3 재구성 이미지논문 연결
-
Eq. 4:
q(z_t | h_t, encoder(o_t))— 인코더 출력이 사후 분포에 입력 -
Eq. 5:
p(o_t | h_t, z_t)— 디코더가 관찰을 재구성 - Section 3.1: "The encoder and decoder use CNN architectures"
뇌 비유
- Encoder = 시각 피질 (V1→V2→V4→IT): 이미지에서 점점 추상적 특징 추출
- Decoder = 연상 피질: 기억에서 이미지를 "떠올리는" 과정
3-4. behaviors.py — Actor-Critic (기저핵 + 전두엽)
역할
RSSM이 "세상이 어떻게 돌아가는지" 배웠다면, Actor-Critic은 "그 세상에서 어떻게 행동할지"를 배웁니다.
핵심 구조
Actor (행동 선택):
입력: 잠재 상태 (h, z)
출력: 행동 확률 분포
학습: 상상된 궤적에서 보상 최대화
(h, z) → MLP → action_distribution
→ 비유: 기저핵 — 습관적/자동적 행동 선택
Critic (가치 평가):
입력: 잠재 상태 (h, z)
출력: 예상 보상 (V(s))
학습: TD(λ) 방식으로 가치 추정
(h, z) → MLP → value_estimate
→ 비유: 전두엽 — "이 상황은 좋은가, 나쁜가?"핵심: "상상 속에서" 학습
실제 환경과 상호작용하지 않고 학습하는 과정:
1. RSSM 안에서 15스텝 미래를 상상
z_0 → z_1 → z_2 → ... → z_15
2. 각 상상 상태에서 보상을 예측
r_1, r_2, ..., r_15
3. Actor: 이 보상을 최대화하는 행동을 학습
4. Critic: 각 상태의 가치를 정확하게 추정
→ 이것이 Dreamer의 "Dream Learning"
→ 실제 환경 상호작용 없이 정책 개선 가능
→ 데이터 효율성이 높은 이유논문 연결
- Section 3.2 "Actor Critic": 전체 구조 정의
-
Eq. 6:
V_λ(s_t)— TD(λ) 리턴 계산 -
Eq. 7: Actor 목표 —
max_θ E[Σ V_λ(s_t)] -
Eq. 8: Critic 목표 —
min_ψ E[(V_ψ(s_t) - sg(V_λ(s_t)))²] - v3 혁신: Symlog 예측으로 보상 규모에 무관한 안정 학습
뇌 비유
- Actor = 기저핵: 반복 학습으로 자동화된 행동 패턴
- Critic = 안와전두피질: 상황의 가치 판단 ("이거 해도 될까?")
3-5. replay.py — 경험 저장 (에피소드 기억)
역할
환경과 상호작용한 경험을 저장하고, 학습 시 미니배치로 샘플링합니다.
저장 구조:
Episode = [(obs_0, action_0, reward_0, done_0),
(obs_1, action_1, reward_1, done_1),
...]
샘플링:
랜덤 에피소드에서 연속 50스텝 시퀀스를 추출
→ RSSM은 시퀀스 데이터로 학습하기 때문논문 연결
- Algorithm 1, Line 3-4: 환경 상호작용 → Replay Buffer 저장
- Algorithm 1, Line 5: Buffer에서 배치 샘플링 → 학습
뇌 비유
해마의 에피소드 기억 — 경험을 저장하고, 수면 중 "리플레이"하여 학습 (Memory Replay)
4. v3 핵심 혁신 3가지
혁신 1: Symlog 예측
문제: 게임마다 보상 범위가 다름
Atari: 0~1000점, Minecraft: -1~1점
기존: 보상 클리핑 (정보 손실) 또는 별도 정규화 (도메인별 조정 필요)
v3 해결: symlog(x) = sign(x) · ln(|x| + 1)
→ 큰 값은 압축, 작은 값은 유지
→ 도메인 무관 안정 학습혁신 2: 이산 카테고리컬 표현
v1: z = 가우시안(30차원) → 연속적, 모드 붕괴 위험
v3: z = 32 카테고리 × 32 클래스 = 1024 가능한 상태
→ "몬스터 상태 = 공격/대기/이동 중 택1" 식의 범주형 표현
→ straight-through gradient로 미분 가능혁신 3: 단일 하이퍼파라미터
v1/v2: 도메인마다 하이퍼파라미터 튜닝 필요
v3: "size" 하나만 조절 (XS=8M, S=22M, M=57M, L=200M, XL=400M)
→ 나머지는 모두 size에 비례하여 자동 결정
→ 150+ 태스크에서 단일 설정으로 동작5. 실습 방법: 직접 코드 읽기
환경 설정 (코드를 읽기만 할 경우)
git clone https://github.com/danijar/dreamerv3
cd dreamerv3추천 읽기 순서
1단계: 전체 구조 파악
agent.py → train() 메서드 읽기
→ "아, 이 순서로 학습하는구나"
2단계: 핵심 모델 이해
nets.py → RSSM 클래스
→ GRU, prior, posterior 흐름 추적
3단계: 상상 과정 이해
nets.py → imagine() 메서드
→ "실제 환경 없이 rollout하는구나"
4단계: 행동 학습
behaviors.py → ActorCritic 클래스
→ "상상된 궤적으로 정책을 배우는구나"핵심 확인 포인트
| 확인할 것 | 어디서 | 예상 결과 |
|---|---|---|
| GRU hidden size | nets.py RSSM.init | model_size (기본 512) |
| z의 shape | nets.py RSSM.prior | (batch, 32, 32) = 32 categories × 32 classes |
| imagination horizon | agent.py imagine | 15 스텝 |
| replay sequence length | replay.py | 50 스텝 |
| symlog 함수 | jaxutils.py | sign(x)·ln(|x|+1) |
6. 이 실습과 다음 실습의 관계
실습 #1 (이론) 실습 #2 (설계)
Dreamer v3 코드 해부 → CartPole에 적용할
"설계도를 읽음" "단순화 설계"
이 파일에서 배운 것 다음에서 단순화
─────────────────────────────────────────────
RSSM (GRU+Stochastic) → MLP (단일 스텝)
이산 카테고리컬 z → 4차원 수치 상태
Actor-Critic → Random Shooting
CNN Encoder/Decoder → 없음 (수치 직접 사용)
15스텝 Imagination → 10스텝 Imagination7. 핵심 요약
| 논문 개념 | 코드 위치 | 뇌 비유 | 한 줄 설명 |
|---|---|---|---|
| World Model | nets.py (RSSM) | 해마 | 세상의 물리법칙 학습 |
| Encoder | nets.py (Encoder) | 시각피질 | 이미지 → 잠재 벡터 |
| Decoder | nets.py (Decoder) | 연상피질 | 잠재 벡터 → 이미지 복원 |
| Actor | behaviors.py | 기저핵 | 행동 선택 정책 |
| Critic | behaviors.py | 전두엽 | 상태 가치 평가 |
| Replay | replay.py | 에피소드 기억 | 경험 저장/재생 |
| Agent | agent.py | 전전두엽 | 전체 루프 조율 |
◀
실습 #4 체계적 정리: Vision World Model — VAE·잠재공간·Dreamer의 눈
실습 #2: CartPole World Model 설계 — Dreamer를 극단적으로 단순화하기
▶
💬 0
로그인 후 댓글 작성
첫 댓글을 남겨보세요!