Skip to content
Bud's Wiki

Open Dreamer — Open-source Dreamer 4 world-model implementation in JAX/Flax

TL;DR

Section titled “TL;DR”

Open Dreamer is a from-scratch JAX/Flax NNX reimplementation of the Dreamer 4 world-model pipeline (Hafner, Yan, Lillicrap 2025; arXiv 2509.24527), trained on Minecraft/VPT-style gameplay and playable in real time via an in-browser demo hosted on Reactor. The repo bundles a causal video tokenizer, an action-conditioned latent dynamics model, ArrayRecord-based data pipelines, Orbax checkpointing, and an FVD evaluator, structured as a six-step training/eval workflow. A companion inference repository (reactor-team/open-dreamer) hosts the runtime for local rollouts on user-supplied MP4 + action sequences. The stated roadmap gap is the full Dreamer 4 behavior-cloning / RL agent training loop — currently unimplemented.

Key claims

Section titled “Key claims”
  • The pipeline factorizes Dreamer 4 into two independently trained stages: a causal video tokenizer on raw MP4 clips, then an action-conditioned latent dynamics model on tokenized episodes [README §Workflow].
  • The tokenizer can be trained on shorter windows than tokenization time — e.g. 16-frame training windows on fixed 256-frame raw records, later tokenized as full 256-frame latent episodes [README §Train the tokenizer].
  • Tokenization is a separate offline pass that emits latent shards plus a metadata/latent_stats.npz (mean/std/num_samples/num_videos); those statistics must be copied into the latent-dataset config before dynamics training [README §Tokenize the dataset].
  • The default backing dataset format is Minecraft/VPT-style shard-*.array_record files: each raw record is a pickled Python dict containing MP4 bytes, a video shape, actions, and optional source metadata; tokenized records are written as msgpack ArrayRecords with tokenizer latents and actions [README §Dataset].
  • The training stack targets a CUDA-12-compatible JAX build with uv lock-file dependency management on Python 3.11 [README §Requirements].
  • Evaluation supports both mode=generate (save MP4s only) and mode=evaluate (compute FVD from previously generated videos) via scripts/eval_fvd.py [README §Evaluate].
  • The Roadmap explicitly lists the full Dreamer-4 BC/RL agent training loop as unimplemented; the current release covers only the world model + tokenizer + FVD eval [README §Roadmap].
  • The live in-browser demo hands the video stream between a real Minecraft game and the world model frame by frame (“Game ⟷ Dream toggle”), with the real-time model running in the cloud on Reactor rather than locally [README §Try it now].
  • The project cites Jasmine (p-doom/jasmine) — a scalable JAX-based world-modeling codebase — as a direct prior for the JAX/Flax approach [README §References].

Method

Section titled “Method”

Open Dreamer follows the Dreamer 4 recipe as a two-stage latent world model: (1) a causal video tokenizer compresses raw video frames into a latent sequence with a temporal causality constraint so future frames cannot leak into past latents, and (2) an action-conditioned latent dynamics model is trained autoregressively on those latents together with the shifted action sequence. Both models are Flax NNX modules under dreamer/models.py, with training loops in dreamer/training.py, rollout / denoising utilities in dreamer/generation.py, and Orbax checkpoint bundles in dreamer/checkpointing.py. Data flows through Grain + ArrayRecord pipelines under dreamer/data/ — MP4 bytes are decoded on the fly for tokenizer training, and tokenizer latents are re-serialized as msgpack ArrayRecords for dynamics training. FVD evaluation lives in dreamer/fvd/ and is driven by scripts/eval_fvd.py with a fvd_chunk_size parameter controlling rollout batching.

Configuration is Hydra/OmegaConf under configs/, with the six-step workflow — prepare raw MP4 shards → train tokenizer → tokenize dataset → copy latent stats → train dynamics → generate + FVD — driven entirely by editing YAML files (tokenizer.yaml, tokenize.yaml, dynamics.yaml, eval_fvd.yaml) plus a raw-video dataset config (dataset/minecraft_vpt.yaml) and a latent-video dataset config (dataset/minecraft_vpt_latent.yaml). The live demo is a Next.js app in site/ and is not part of the training pipeline.

Results

Section titled “Results”

No quantitative results are reported in the README beyond pointing at the live demo — FVD numbers, training compute, or comparisons against the original Dreamer 4 numbers are not included at filing time. The referenced upstream paper is arXiv:2509.24527 (Hafner, Yan, Lillicrap 2025), and the referenced prior JAX codebase is p-doom/jasmine.

Why it’s interesting

Section titled “Why it’s interesting”

Open Dreamer is the first filed open-source Dreamer 4 reimplementation the wiki has seen, and it lands in the middle of an unusually populated bracket. It sits alongside AlayaWorld: Interactive Long-Horizon World Modeling (Alaya Lab, ~13B LTX-2.3-based interactive video WFM), HY-World 1.5 (WorldPlay): A Systematic Framework for Interactive World Modeling with Real-Time Latency and Geometric Consistency (Tencent HY-World 1.5), Advancing Open-source World Models (LingBot-World) (LingBot-World v1) and LingBot-World 2.0 / LingBot-World-Infinity — Infinite Worlds with Versatile Interactions (LingBot-World 2.0) as an open-weights + open-code alternative to closed flagships Project Genie: Experimenting with infinite, interactive worlds (Project Genie / Genie 3), Runway GWM-1, and Odyssey-2 in the real-time interactive category. What differentiates it is the substrate: rather than adapting a large pretrained video DiT (LTX-2.3, Wan, Cosmos), Open Dreamer trains a causal video tokenizer + action-conditioned latent dynamics stack from scratch on Minecraft/VPT data — the classical Dreamer factorization ported to modern JAX rather than a video-diffusion-first architecture. This is the Ha–Schmidhuber–Hafner latent-world-model lineage the World Foundation Models concept page catalogs but rarely sees a modern open reproduction of.

The JAX/Flax NNX + ArrayRecord + Grain stack also makes this a concrete referenceable training-infra artifact for the Distributed training parallelism concept — a rare filed instance of a full-stack world-model training pipeline in JAX (previous JAX WM entries have been at the framework/library level: facebookresearch/vjepa2 — official PyTorch code and checkpoints for V-JEPA 2 and V-JEPA 2.1 is PyTorch-only, x-jepa — Explorations into JEPA approaches (lucidrains) is PyTorch, SpectraX — True MPMD Pipeline Parallelism for JAX is JAX-infra without a WM). Complementary rather than competitive with x-jepa — Explorations into JEPA approaches (lucidrains) on the “portable minimal reference for a world-model lineage” role — x-jepa bundles the LeCun/JEPA latent-predictive pole, Open Dreamer occupies the Hafner/Dreamer generative-latent pole.

See also

Section titled “See also”