Open Dreamer 리뷰: JAX/Flax로 다시 만든 Dreamer 4 월드 모델 파이프라인, 학습 레시피까지 공개
요즘 AI가 게임·로봇 시뮬레이션을 스스로 익히게 만드는 기술이 큰 주목을 받고 있습니다. 그 핵심이 바로 "월드 모델"이라고 부르는 것입니다. 이번에 공개된 Open Dreamer는 Google DeepMind의 최신 월드 모델인 Dreamer 4를 JAX/Flax 기반으로 직접 재현(reproduction)한 오픈소스 프로젝트입니다. 누구나 학습 레시피와 안정화 트릭까지 그대로 따라 해볼 수 있도록 정돈되어 있어, 학계와 산업 현장에서 화제가 되고 있습니다.
Open Dreamer가 무엇인가요?
쉽게 말하면 Open Dreamer는 "AI가 상상 속에서 세상을 돌려보는 훈련장치"입니다. 실제 로봇이나 게임 엔진을 매번 돌리지 않아도, AI 내부에 작은 시뮬레이터(월드 모델)를 만들고 그 위에서 수백만 번의 가상 경험을 쌓으면서 행동 정책을 학습합니다.
원조인 Dreamer 4는 Google DeepMind가 2025년 NeurIPS에서 발표한 최신 모델로, 비디오 자기회귀(self-attention 기반 시퀀스 모델)와 강화학습을 결합해 Minecraft와 같은 복잡한 환경에서도 장기적 계획을 잘 세우는 것으로 유명합니다. Open Dreamer는 이 Dreamer 4의 구조와 학습법을 가능한 한 동일하게 JAX/Flax로 다시 코딩해 누구나 손쉽게 돌려볼 수 있도록 만든 "재현 프로젝트"입니다.
한눈에 보는 Open Dreamer 핵심 특징
- 100% 오픈소스: Dreamer 4와 동일한 구조를 JAX/Flax로 다시 작성
- 전체 학습 레시피 공개: 데이터 구성부터 옵티마이저, 손실함수까지
- 안정화 트릭 포함: 대규모 학습에서 무너지지 않도록 만든 노하우 포함
- 거의 모든 코드가 단일 파일: 1,200줄의 깔끔한 학습 스크립트
어떤 구조로 만들어졌나요? (Architecture)
Open Dreamer의 핵심 구조는 크게 두 개의 신경망으로 나뉩니다.
1) 월드 모델 (World Model)
월드 모델은 "AI의 상상 시뮬레이터"입니다. 현재 화면(이미지)과 과거 행동 로그를 입력으로 받아, 다음 프레임이 어떻게 보일지를 픽셀 단위로 예측합니다. 이때 단순한 점화(prediction)가 아니라 분포 전체를 예측하기 때문에 같은 상황에서도 다양한 미래가 가능합니다.
학습에는 두 가지 최신 기법이 결합되었습니다.
디퓨전 포싱 (Diffusion Forcing) — 자기회귀(다음 토큰 예측)와 노이즈 제거 기반 시퀀스 생성(Diffusion)을 합친 최신 생성 패러다임입니다. "자연스러운 비디오의 시간적 인과관계"와 "다양한 미래 샘플링"을 동시에 잡을 수 있어, 장기 호라이즌 비디오 생성에서 강력한 성능을 보입니다.
플로우 매칭 (Flow Matching) — 연속 정규화 플로우(Continuous Normalizing Flow)를 시뮬레이션 없이 훨씬 가볍게 학습시키는 방법입니다. 학습 난이도가 낮고, 샘플 품질은 비슷하거나 더 좋아서 최근 확산 모델 대안으로 빠르게 자리 잡고 있습니다.
2) 액터-크리틱 (Actor-Critic) 정책 네트워크
월드 모델이 만든 "가상 세계"에서 정책을 학습시킵니다. 정책(액터)은 "지금 무엇을 할지"를 결정하고, 크리틱은 "그 행동이 얼마나 좋은지" 점수를 매깁니다. 정책은 분포형(distribution)으로 표현되어 있어 탐색(exploration)과 활용(exploitation) 사이의 균형을 자연스럽게 잡습니다.
학습 레시피는 어떻게 되나요?
Open Dreamer가 공개한 학습 과정은 크게 세 단계로 나뉩니다.
① 인터리브드 사전학습 (Interleaved Pre-training)
먼저 마스크드 오토인코더(Masked Autoencoder, MAE) 방식으로 비디오 토크나이저를 학습합니다. 영상 프레임의 일부를 가린 뒤, 가려진 부분을 맞추게 하는 자기지도학습(self-supervised learning)입니다. 인간이 퍼즐을 맞추듯, 모델이 "세상의 구조"를 자연스럽게 익히도록 돕습니다.
② 월드 모델 학습
토크나이저 위에서 디퓨전 포싱 + 플로우 매칭을 결합해 다음 프레임을 생성하도록 훈련합니다. 단순히 "다음 한 장"이 아니라, 수십~수백 스텝 앞까지 자연스럽게 이어지는 영상을 만들도록 합니다.
③ 액터-크리틱 정책 학습
고정된 월드 모델 안에서 PPO(Proximal Policy Optimization) 계열의 분포형 정책 그래디언트로 행동을 학습합니다. 정책과 가치 함수를 함께 업데이트하기 때문에 보상이 드문 환경에서도 안정적으로 수렴합니다.
옵티마이저는 무엇을 쓰나요?
뮤온 옵티마이저 (Muon Optimizer) — 모멘텀과 뉴턴-슐츠 반복법(Newton-Schulz iteration)을 결합해, Adam 계열 대비 메모리를 훨씬 적게 쓰면서도 대형 모델 학습에서 더 좋은 수렴 속도를 보이는 차세대 옵티마이저입니다. Open Dreamer는 이 Muon을 모든 네트워크에 일관되게 적용했습니다.
메모리 벽은 어떻게 넘었나요?
월드 모델 학습에서 가장 큰 적은 GPU 메모리입니다. 한 번의 순전파(forward pass)에 비디오 토큰이 수십만 개씩 필요하기 때문입니다. Open Dreamer는 이 문제를 두 가지 방식으로 해결했습니다.
완전 분할 데이터 병렬 처리 (FSDP, Fully Sharded Data Parallel) — 모델 파라미터와 옵티마이저 상태를 여러 GPU에 잘게 쪼개(shard) 보관하는 대규모 분산 학습 기법입니다. GPU 수에 거의 비례해 메모리가 절약되며, 대규모 모델 학습의 표준 방식으로 자리 잡았습니다.
바리센트릭 최적 수송 (Barycentric Optimal Transport) — 두 분포 사이의 "비용"을 계산할 때 샘플링 점들을 분포에 맞춰 정렬해 비교하는 기법입니다. 노이즈와 데이터 매칭이 정밀해져 손실 계산이 안정적이고, 대규모 배치에서도 효율적으로 작동합니다.
결과적으로 Open Dreamer는 16×H100 GPU에서 학습이 가능하도록 설계되어, 일반 연구팀도 손쉽게 시도해볼 수 있게 됐습니다.
안정성을 위한 핵심 트릭
대규모 월드 모델 학습은 학습이 갑자기 무너지거나 발산(divergence)하는 경우가 흔합니다. Open Dreamer는 이를 막기 위해 다음 네 가지를 적용했습니다.
- EMA(지수이동평균) 체크포인트: 매 스텝 평가 가중치를 따로 유지해 평가·배포 시 안정성 확보
- Q-러닝 정규화: 가치 함수의 폭발적인 증가를 막아 학습 후반부 발산 방지
- 보상 스케일링: 보상 신호를 일정한 범위로 정규화해 정책 그래디언트의 노이즈 감소
- 결정적 평가: 평가 시 행동 샘플링을 고정해 공정하고 재현 가능한 측정
성능은 어느 정도인가요?
Open Dreamer는 Dreamer 4와 거의 동일한 품질을 재현했습니다. 평가 지표로는 프레셰 비디오 거리(FVD, Fréchet Video Distance)를 사용하는데, 이는 실제 영상과 생성 영상의 분포가 얼마나 가까운지를 측정하는 표준 지표입니다.
프레셰 비디오 거리 (FVD) — 실제 비디오와 AI가 만든 비디오의 통계적 거리를 측정하는 지표입니다. 값이 낮을수록 생성 품질이 좋다는 의미이며, 비디오 생성 모델의 표준 벤치마크로 널리 쓰입니다.
Minecraft·CSGO 등 다양한 환경에서 Dreamer 4의 수치를 거의 그대로 따라잡았고, 단일 코드베이스로 재현되었기 때문에 향후 후속 연구의 강력한 베이스라인이 될 전망입니다.
아직 박스에 들어 있지 않은 것들
아무리 잘 구성된 재현 프로젝트라도 원본과 다른 부분이 있습니다. Open Dreamer에서 빠진 항목은 다음과 같습니다.
- 원본의 일부 정규화 트릭(예: weight decay, learning rate warmup 세부값)
- 고해상도(128×128 이상) 학습 결과의 부분적 재현
- 전체 평가 환경의 100% 동일 구성(일부 환경은 라이선스 이슈로 대체)
그럼에도 "거의 동일한 품질을 단일 파일로 재현"했다는 점만으로 충분한 가치가 있다는 평가가 많습니다.
자주 묻는 질문 (FAQ)
- Q1. 월드 모델이 정확히 무엇인가요?
- AI가 "세상을 어떻게 이해하는지" 학습한 내부 시뮬레이터입니다. 실제 환경을 매번 돌리지 않아도, 이 시뮬레이터 안에서 가상 경험을 쌓으며 정책을 학습할 수 있어 로봇·게임·자율주행 등 다양한 분야에서 활용됩니다.
- Q2. Dreamer 4와 다른 점은 무엇인가요?
- Dreamer 4는 Google DeepMind의 독점 구현입니다. Open Dreamer는 이를 JAX/Flax라는 오픈소스 프레임워크로 다시 작성해 누구나 학습 레시피까지 그대로 따라 해볼 수 있도록 만든 재현 프로젝트입니다.
- Q3. 제 개인 GPU로도 돌릴 수 있나요?
- 풀 학습은 16×H100급이 필요하지만, 소형 환경에서 추론(inference)과 미세조정(fine-tuning)은 단일 고성능 GPU로도 시도해볼 수 있도록 코드가 모듈화되어 있습니다.
- Q3. 왜 JAX/Flax를 선택했나요?
- JAX는 Google이 만든 고성능 수치 계산 라이브러리로, TPU와 GPU에서 동일한 코드로 빠르게 동작합니다. Flax는 JAX 위의 신경망 라이브러리(NNX 모듈 포함)로, 함수형 프로그래밍과 깔끔한 모듈 구성이 가능합니다.
핵심 정리 (Key Takeaways)
Open Dreamer는 월드 모델 연구의 "진짜 기준점"을 오픈 생태계로 끌어내린 의미 있는 프로젝트입니다. 이제 연구팀은 독점 코드 없이도 최신 월드 모델 학습을 재현하고, 자신만의 변형을 시도해볼 수 있게 됐습니다. AI가 "상상"을 통해 배우는 시대의 서막이라고 할 수 있습니다.