Open Dreamer: JAX/Flax 기반 Dreamer 4 파이프라인의 오픈소스 구현, 전체 학습 레시피 포함

연구자 Francesco Sacco, Diego Martí, Edward Hu는 Reactor의 후원을 받아 Dreamer 4의 오픈 구현체인 Open Dreamer를 JAX/Flax NNX로 공개했습니다. 학습용(토크나이저, 역학 모델, 롤아웃, FVD 점수 산출)과 로컬 추론용 두 개의 저장소가 공개되었습니다. 브라우저 데모는 생성된 Minecraft 세계를 스트리밍하며, 실제 게임과 월드 모델 간 전환을 허용합니다. 아키텍처는 토크나이저와 역학 모델 모두에 블록-인과 트랜스포머 백본을 사용합니다. 토크나이저는 트랜스포머 기반의 마스크 오토인코더로, KL 손실이나 적대적 손실 없이 약 100배 압축을 달성합니다. 역학 모델(파라미터 16억 개)은 확산 강제(diffusion forcing), 흐름 정합(flow matching), 숏컷 모델을 사용하여 다음 프레임 예측을 수행하고, 다음 행동도 예측합니다. 학습은 200,000 스텝 동안 Muon 옵티마이저, WSD 스케줄, 피크 학습률 3e-4, 이엠에이(EMA) 감쇠 0.999로 진행됩니다. 팀은 B200 GPU에서 모델 FLOPs 활용률 57-58%를 달성했으며, 모델 상태는 약 24GiB에 맞춰집니다. 활성화(activation)가 주요 메모리 비용이었으며, 데이터 병렬화와 활성화 체크포인팅으로 해결했습니다. 안정성 문제가 가장 큰 과제였으며, 여섯 가지 해결책으로는 Muon 옵티마이저, 필수 EMA 가중치, 신중한 혼합 정밀도, v-공간 손실 가중치, 미니배치 무게중심 최적 수송(barycentric optimal transport), μ-파라미터화의 중요성 축소 등이 포함됩니다. 저장소에는 행동 복제나 강화학습 학습 루프가 포함되어 있지 않으며, FVD 점수는 공개되지 않았습니다.