next-state/open-dreamer 딥다이브
— "게임을 통째로 상상해 내는" 프런티어 월드 모델을 JAX로 처음부터 학습하는 레시피

1프로젝트 한 줄 요약

월드 모델은 한 문장으로 "게임 엔진을 신경망으로 통째로 배운 것"이다. 보통 게임은 코드로 짜인 물리·렌더링 엔진이 다음 화면을 계산한다. 월드 모델은 그 엔진을 버리고, 수많은 플레이 영상만 보고 "이 화면에서 이 키를 누르면 다음 화면은 이렇게 된다"를 통째로 학습한다. 그래서 학습이 끝나면, 사람이 방향키를 누를 때마다 모델이 다음 프레임을 그려서 내놓는다. 코드로 짠 마인크래프트가 아니라, 마인크래프트를 흉내 내도록 배운 꿈을 플레이하는 셈이다.

Open Dreamer는 이 아이디어의 최신판인 딥마인드 Dreamer 4(2025) 논문을 JAX/Flax로 처음부터 재현한 오픈 구현이다. Next State 팀(Diego Martì Monsò, Francesco Sacco, Edward Hu)이 Reactor의 후원으로 만들었고, 블로그 제목 그대로 「How to train a frontier-level world model(프런티어급 월드 모델 학습법)」이다. 저장소에는 마인크래프트 세계를 브라우저에서 실시간으로 플레이하는 데모까지 연결돼 있다.

중요한 구분이 하나 있다. 이 저장소는 학습(training) 파이프라인이다 — 토크나이저를 훈련하고, 데이터를 잠재(latent)로 바꾸고, 다이내믹스 모델을 훈련하고, 품질(FVD)을 재는 코드다. 실제로 브라우저에서 실시간으로 돌리는 추론(inference) 코드는 별도 저장소 reactor-team/open-dreamer에 있다. 즉 여기는 "요리법과 주방", 저기는 "손님에게 내는 완성 요리"다.

2왜 주목받는가 — 논문이 아니라 '레시피'를 공개했다

월드 모델 논문은 많다. Dreamer 4도, Genie도, GameNGen도 있다. 그런데 Open Dreamer가 트렌딩에 오른 건 겨냥점이 다르기 때문이다. 논문은 "무엇을 했는지"만 알려주지, "실제로 학습시키다 몇 주씩 막히는 지점"은 안 알려준다. Open Dreamer는 그 막히는 지점과 해법을 통째로 공개했다. 네 가지가 겹친다.

① "거인의 어깨 위로 순간이동" — 프런티어 연구의 진입장벽을 낮춘다

저자들의 표현이 압권이다: "과학에서 가장 오래 걸리는 건 거인의 어깨 위로 기어올라 가는 것인데, 가끔은 어깨 위로 순간이동해서 바로 최전선에서 일하기 시작하는 게 좋다." Dreamer 4 재현에 필요한 아키텍처·학습 코드·하이퍼파라미터·실패 사례를 다 공개해, 월드 모델 연구를 개인·소규모 팀도 시작할 수 있게 만들었다.

② 단일 GPU(CoinRun)로 검증 → 대형(마인크래프트)으로 스케일업

똑똑한 개발 전략이 그대로 담겼다. 처음부터 마인크래프트로 덤비지 않고, 단일 GPU에서 도는 간단한 게임 CoinRun(2D 플랫포머)으로 파이프라인 전 과정을 먼저 완성했다. 디버깅이 쉽고 반복이 빠르다. 코어 로직이 다 돌아가는 걸 확인한 뒤 마인크래프트로 규모만 키웠다. "작게 만들어 검증하고 키운다"는 ML 엔지니어링의 정석을 코드로 보여 준다.

③ 실시간으로 '플레이되는' 결과물

많은 생성 모델이 오프라인으로 영상 몇 개 뽑고 끝난다. Open Dreamer는 Reactor 런타임 위에서 마인크래프트 세계를 브라우저에서 실시간 플레이한다. "Game ⟷ Dream" 토글로 진짜 게임과 모델의 꿈을 프레임 단위로 넘겨 가며 비교할 수 있다. 실시간이 되려면 few-step 확산(4스텝) + KV 캐시 같은 추론 최적화가 필수인데, 그 코드가 generation.py에 그대로 있다(§4).

④ 논문에 없는 '엔지니어링 로그'를 남겼다

이게 진짜 차별점이다. 블로그의 절반이 Stability(안정성) 이야기다 — "MSE 손실은 매끄럽게 내려가는데 생성 품질은 나빠진다"는, 며칠~몇 주씩 사람을 괴롭히는 문제들. Muon vs LaProp 옵티마이저, EMA 필수, 혼합정밀도 경계, x-예측 + v-공간 손실, 최적수송(OT)까지 — "우리가 겪은 실패와 고친 법"을 순서대로 문서화했다. 논문 재현을 해 본 사람이면 이 로그의 가치를 안다.

3기술 스택 전체 지도

Open Dreamer의 스택은 머리끝부터 발끝까지 JAX다. 저자들이 JAX를 고른 이유는 명확하다 — "XLA 컴파일러가 효율적이고 확장 가능한 SPMD(단일 프로그램·다중 데이터) 프로그램에 탁월하기 때문". 즉 코드 한 벌로 여러 GPU에 자동 분산시키기 좋다는 것. 패키지 이름부터 tiny-dreamer-4, 설명은 "Dreamer 4 implementation in pure Jax"다.

스택에서 눈여겨볼 선택 세 가지. 첫째, Flax의 새 API nnx를 쓴다(구 linen 아님) — 파이썬 객체처럼 상태를 다루는 최신 방식이라 KV 캐시 같은 가변 상태를 자연스럽게 담는다. 둘째, optax를 깃 최신 버전에서 직접 당겨온다(tool.uv.sources) — 실험적 Muon 옵티마이저가 안정판에 없기 때문. 셋째, 데이터를 미리 토큰화해 ArrayRecord로 저장하고 Grain으로 읽는다 — ffmpeg로 영상을 매번 디코딩하면 GPU를 못 먹여서다(§6의 dataloading).

핵심 의존성 (pyproject.toml 기준)

4아키텍처 심화 분석

전체 그림은 의외로 단순하다. 모델이 딱 둘이고, 백본(뼈대)은 하나다. ①토크나이저가 영상 프레임을 작은 잠재 토큰으로 압축하고(≈100배), ②다이내믹스 모델이 그 잠재 공간에서 "다음 프레임"을 확산(diffusion)으로 그린다. 둘 다 같은 블록-인과 트랜스포머(block-causal transformer)를 뼈대로 쓴다.

백본 — 공간(Space) 어텐션과 시간(Time) 어텐션의 교대

두 모델이 공유하는 뼈대의 핵심은 어텐션을 두 종류로 쪼갠 것이다. 영상은 "여러 장의 그림이 시간순으로 이어진 것"이라, 정보가 흘러야 할 방향이 둘이다.

• Space Layer(공간 층): 한 프레임 안의 토큰들끼리 정보 교환. "이 화면에서 나무 옆엔 흙이 있다" 같은 공간 관계.
• Time Layer(시간 층, 인과적): 서로 다른 프레임 사이 정보 교환. 단 과거만 본다(causal) — 미래를 훔쳐보면 안 되니까. time_every=4, 즉 4개 층마다 1개가 시간 층이다.

① 토크나이저 — "VAE를 버리고" 마스킹 오토인코더로

영상 프레임을 그대로 다이내믹스에 넣으면 너무 크다. 그래서 먼저 잠재 공간으로 ≈100배 압축한다. 보통 이럴 때 VAE를 쓰는데, 저자들의 표현이 솔직하다: "VAE를 학습해 본 적 있다면 그게 얼마나 끔찍한 경험인지 안다." 그래서 대신 트랜스포머 기반 마스킹 오토인코더(MAE)를 택했다.

• 마스킹이 잠재를 '확산 친화적'으로 만든다: 인코더 패치의 일부를 가리고(mae_p_max=0.9) 복원시키면, 잠재 공간이 매끄러워져 뒤 단계 확산이 쉬워진다.
• KL 손실도, 적대적(GAN) 손실도 없다: VAE의 골칫거리(KL 균형, GAN 불안정)를 통째로 피한다. 손실은 MSE 재구성 + LPIPS(지각 손실)뿐.
• 빠르고 확장 가능: 인코더 512 latent × 16차원, depth 12 / 디코더 depth 8.

② 다이내믹스 — '다음 프레임 그리기'를 확산으로

월드 모델의 심장이다. 근본 원리는 "다음 프레임 예측" — 각 프레임을 이전 프레임들에 조건부인 이미지로 생성한다. 이때 그냥 예측이 아니라 diffusion forcing + flow matching + shortcut models라는 세 겹의 기법으로 훈련한다. 그리고 결정적으로, 한 타임스텝을 블록으로 접는다:

왜 이렇게 블록으로 접나? 원래 롤아웃은 "월드가 상태를 만들고 → 정책이 액션을 뽑고"를 번갈아 한다. 두 모듈을 오가면 느리고, 표현·파라미터도 못 나눈다. 블록으로 접으면 모든 타임스텝을 병렬로 학습하고, 월드와 에이전트가 같은 트랜스포머를 공유한다. register 토큰(n_register=32)은 "Vision Transformers Need Registers" 논문에서 온, 계산용 공용 메모장이다.

추론 — τ-ladder로 4스텝 만에, KV 캐시로 실시간

generation.py가 실시간의 정체다. 학습은 최대 노이즈 해상도 k_max=256으로 하지만, 추론은 τ-사다리(τ-ladder)를 단 4스텝만 밟는다(num_steps=4, d=1/4). 각 스텝은 "깨끗한 latent 예측(x̂) → 현재와 섞기"를 반복한다.

포인트 둘. 첫째, 문맥 프레임을 완전히 깨끗하게 두지 않고 살짝 노이즈(τ_ctx=0.9)를 준다 — 학습·추론 분포를 맞추기 위한 diffusion forcing의 트릭. 둘째, KV 캐시를 '링 버퍼'로 구현해(고정 크기 배열에 순환 기록) 슬라이딩 윈도우 어텐션을 O(1)로 돌린다. 이 둘이 없으면 매 프레임 전체를 다시 계산해야 해서 실시간이 불가능하다.

5디렉토리 구조 해부

구조는 깔끔한 연구 코드베이스의 전형이다. dreamer/에 모델·학습·데이터·평가 로직이 다 있고, scripts/가 실제 실행 진입점, configs/가 Hydra YAML, site/가 Next.js 데모다. 워크플로 순서(토크나이저 학습 → 토큰화 → 다이내믹스 학습 → 평가)가 scripts/의 파일명에 그대로 드러난다.

6학습 포인트

이 저장소의 최고 가치는 대형 생성 모델을 실제로 안정적으로·빠르게 학습시키는 노하우다. 블로그 절반이 "손실은 내려가는데 품질이 나빠지는" 안정성 문제와 그 해법이다. 주제(월드 모델)를 걷어내도, 아래는 어떤 대형 확산·시퀀스 모델에도 통하는 교훈이다.

① Muon 옵티마이저 — LaProp의 '손실 스파이크'를 잡다

Dreamer 시리즈는 원래 LaProp을 썼는데, 학습이 길어질수록 손실이 무작위로 튀는(spike) 현상이 잦아졌다. 저자들은 2D 행렬 가중치에 Muon(+ 임베딩·norm·bias엔 AdamW)을 적용해 훨씬 안정적이고 낮은 손실을 얻었다. Muon은 요즘 LLM 학습에서 급부상한 옵티마이저로, Newton-Schulz 반복(muon_ns_steps=5)으로 그래디언트를 직교화한다. "옵티마이저 하나가 안정성을 바꾼다"는 실전 사례.

② 확산 추론엔 EMA가 '필수'

확산 모델에서 온라인(학습 중) 가중치로 추론하면 결과가 확 나빠진다. 파라미터의 지수이동평균(EMA, θ_ema ← β·θ_ema + (1-β)·θ)로 추론해야 한다(ema_decay=0.999). 학습 코드가 온라인·EMA 두 벌의 가중치를 항상 함께 관리하는 이유다. "왜 내 확산 모델은 학습 손실은 좋은데 샘플이 이상하지?"의 흔한 답.

③ 혼합정밀도 — '어디를 float32로 둘지'가 생사를 가른다

최신 가속기는 BF16에 최적화됐지만 일부 연산은 고정밀이 필요하다. 규칙: 파라미터는 float32 보관, 대부분의 행렬곱·어텐션 입력은 BF16, 그러나 정규화(norm)·사인파 임베딩·어텐션 가중치·다이내믹스 flow 출력 헤드는 float32. 이 경계를 잘못 그으면 학습·생성이 둘 다 불안정해진다. config에 dtype=bfloat16, param_dtype=float32가 분리된 이유.

④ x-예측 + v-공간 손실 — 작지만 확실한 품질 개선

flow matching은 보통 속도장을 예측(v-prediction)하지만, 저자들은 깨끗한 데이터를 직접 예측(x-prediction)하되 v-공간 손실 가중을 쓴다: L_v = ||x - x̂||² / (1-τ)². 월드 모델은 문맥 프레임이 전부 x-공간(실제 픽셀)이라, "다음 프레임도 x-공간으로 예측"하는 게 자연스러운 귀납편향이다. Dreamer 4보다 분모에 제곱을 더한 미세 변형으로 소폭 개선을 얻었다.

⑤ 최적수송(OT) — 흐름의 '길'을 짧게

"최적수송과 flow matching은 천생연분"이다. OT는 노이즈와 데이터를 이동 거리가 최소가 되게 짝지어, 흐름의 경로를 짧고 덜 모호하게 만든다. Open Dreamer는 미니배치 barycentric OT(OTT-JAX의 Sinkhorn, max_iter=50)로 노이즈-데이터 시퀀스를 통째로 커플링해, 롤아웃 생성을 더 안정화했다.

⑥ 스케일링 법칙 — "compute의 제곱근"이라는 자연의 선물

CoinRun 토크나이저로 iso-FLOPs 스케일링을 재 보니, compute-optimal 파라미터 수와 데이터 수가 compute의 제곱근으로 붙었다: N ∝ C^0.56, D ∝ C^0.44. Karpathy의 인용처럼 "복잡한 함수일 수도 있었는데, 자연이 정확히 직선으로 결정했다(뭐지??)". dreamer/scaling.py의 ScalingContext가 토큰 예산·최대 스텝을 이 법칙으로 유도한다.

⑦ MFU 58% 뽑기 — 루프라인·샤딩·활성화 체크포인팅

대형 학습의 하드웨어 효율 챕터다. B200 GPU에서 MFU(모델 FLOPs 이용률) 57~58%를 달성했는데(60%면 매우 건강), 방법은: (a) GPU당 256프레임을 먹여 루프라인의 ridge point(292 FLOP/byte)를 넘겨 compute-bound로 만들고, (b) 샤딩은 DP/FSDP/TP/SP를 다 시도해 본 뒤 단순 데이터 병렬(DP)이 최선(모델 상태 24GiB가 이미 들어맞아 FSDP 통신 이득이 없음), (c) 긴 영상의 활성화가 OOM을 내니 활성화 체크포인팅으로 일부만 저장·나머지는 역전파 때 재계산.

⑧ 데이터로딩 — GPU를 굶기지 않기

마지막 병목은 "GPU를 계속 배부르게 하기". 영상은 무겁고 ffmpeg 디코딩이 느려 GPU를 못 먹인다. 해법: 데이터셋 전체를 미리 토큰화해 ArrayRecord로 저장하고, Grain으로 읽으며 GPU 측 프리페치 버퍼로 데이터를 미리 올려 둔다. "연산 최적화보다 데이터 파이프라인이 병목"인 흔한 현실.

7하드웨어 / 시스템 요구사항

요구사항은 "어디까지 재현하느냐"에 따라 극과 극이다. 파이프라인을 배우고 CoinRun으로 돌려 보는 건 단일 GPU로 충분하지만, 블로그의 마인크래프트 1.6B 모델을 재현하려면 데이터센터급 GPU가 필요하다. 공통 전제는 CUDA 12 호환 JAX 환경이다.

8직접 해볼 수 있는 실습 과제

9관련 기술 심화 학습 로드맵

10핵심 키워드 사전

11참고 링크

• GitHub 저장소(학습): github.com/next-state/open-dreamer (pure JAX · 라이선스 잠정 독점 · 분석 시점 2026-07 초기 커밋)
• GitHub 저장소(추론): github.com/reactor-team/open-dreamer (로컬 실시간 롤아웃)
• TrendShift: trendshift.io/repositories/92173 (2026-07-25 Daily #18)
• 웹사이트·블로그·데모: next-state.github.io/open-dreamer (「How to train a frontier-level world model」 · Game⟷Dream 실시간 데모)
• 핵심 소스: dreamer/models.py(백본·두 모델) · dreamer/generation.py(τ-ladder 추론) · dreamer/training.py(학습 스텝) · configs/tokenizer.yaml·dynamics.yaml(하이퍼파라미터)
• 배경 논문: Dreamer 4 arXiv 2509.24527 · Shortcut Models arXiv 2410.12557 · ViT Need Registers arXiv 2309.16588
• 관련 코드베이스: p-doom/jasmine (JAX 월드 모델링 코드베이스) · optax · Grain