## Article: Open Dreamer — Reproducing Dreamer 4 in JAX/Flax NNX
A small group of AI researchers (Reactor) have released **Open Dreamer**, an open implementation of the Dreamer 4 world-model pipeline written in JAX and Flax NNX. The work provides training code, a local rollout harness, and a browser-based Minecraft demo that streams generated video in real time. The stated goal was to reproduce the Dreamer 4 research under a narrow, controlled search space—starting with CoinRun and scaling to Minecraft/VPT-style gameplay.
### What Actually Shipped
Two main repositories were released:
– `next-state/open-dreamer`: Contains the full training pipeline, including a causal video tokenizer, action-conditioned latent dynamics model, rollout generation, and FVD scoring.
– `reactor-team/open-dreamer`: A minimal local rollout harness that generates frames from an MP4 and a matching action file.
Additionally, a browser demo hosted on the Reactor runtime streams a generated Minecraft world in real time and includes a Game ⟷ Dream toggle that can hand over control frame by frame.
### Architecture: One Backbone, Two Models
Both the tokenizer and the dynamics model share the same **block-causal transformer backbone**, which alternates:
– **Space layers** that propagate information within a single frame
– **Causal time layers** that propagate information across frames
The **tokenizer** is a transformer-based Masked Autoencoder (MAE) rather than a VAE, reportedly achieving about **100× compression** without KL or adversarial loss. The design relies on masking to make the latent space more diffusible.
The **dynamics model** performs next-frame prediction using diffusion forcing, flow matching, and shortcut models. It also predicts the next action. The rollout is folded into per-timestep blocks containing (previous action, state, policy). Spatial attention operates within each block, while causal temporal attention connects blocks across time. Notably, world-model tokens cannot read the agent token, so task and policy information can only influence future states through the next action.
### The Training Recipe, as Configured
The shipped Minecraft configurations make the research concrete:
– The dynamics model is **1.6B parameters**, with 30 block-causal layers, `d_model=1920`, 30 attention heads, and 3 KV heads for grouped-query attention. Every fourth layer is a time-attention layer.
– Each timestep carries **32 learned register tokens**, and a `packing_factor: 2` packs neighboring tokenizer latents into each dynamics spatial token.
– Time attention uses a **192-step sliding window**.
– Training runs for **200,000 steps** with **Muon**, a **WSD schedule**, and a peak learning rate of **3e-4**.
– Shortcut and bootstrap samples switch on at step **100,000** at a **0.25 batch fraction**.
– Exponential moving average decay is **0.999**.
The tokenizer config emits **512 latent tokens per frame** with a bottleneck width of **16**. Raw **360×640** frames are padded to **368×640** so both dimensions divide evenly into **16×16** patches. Encoder depth is **12** at `d_model=1536`; decoder depth is **8** at `d_model=1024`. MAE masking probability tops out at **0.9**, and LPIPS is applied at weight **0.2** on half the timesteps.
VPT actions are parsed into **27 binary action channels** plus **121 categorical mouse classes**, with no continuous channels.
### Computemaxxing and the Memory Wall
The team reports **57–58% model FLOPs utilization (MFU)** on a B200, against a stated healthy benchmark of **60%**. This is analyzed as a roofline problem: on B200, the crossover between bandwidth-bound and compute-bound sits at **292 FLOP/byte**, and pushing **256 frames per GPU** moves the workload past that ridge point.
Sharding behaved unexpectedly. Although the model state (parameters, gradients, optimizer state, and EMA) occupies roughly **24 GiB**—fitting within a single B200—activations were the real memory cost. The team tried data parallelism, FSDP, tensor parallelism, and sequence parallelism, then settled on **plain data parallelism plus activation checkpointing**.
Dataloading was solved by **pre-tokenizing** the entire dataset into `.arrayrecord` files and loading them with **Grain** and a **GPU-side prefetch buffer**. In their experience, **ffmpeg decoding was not fast enough** to keep the GPUs fed.
### The Stability Section Is the Real Payload
The research team emphasizes that **stability consumed the largest share of their effort**. A key observation is that **most stability problems occur despite the loss going down**—that is, MSE can improve smoothly while generation quality degrades.
Six fixes are documented:
– **Muon** replaced **LaProp**, which spiked randomly and increasingly often.
– **EMA weights** are treated as mandatory for diffusion inference.
– Mixed precision keeps **parameters in float32**, **BF16** for most matmul and attention inputs, and **float32** for normalization, sinusoidal helpers, and the dynamics flow output head.
– **x-prediction with v-space loss** (L_v = ||x – x̂||² / (1-τ)²) reports a small but noticeable improvement.
– **Minibatch barycentric optimal transport** made rollout generation more stable.
– **μ-parametrization** was tested but found unnecessary, partly because Muon already stabilizes hyperparameters across model sizes.
Additional results from the CoinRun phase show an **isofLOPs sweep** where compute-optimal scaling is roughly **N ∝ C^0.56** and **D ∝ C^0.44**.
### What Is Not in the Box
– The repository does **not** include the behavior cloning or RL training loop; a full Dreamer 4 BC/RL agent loop is listed as an open roadmap item.
– The CoinRun policy work described in the post was **not used for Minecraft** and was **not released**.
– The post does **not publish FVD scores**, although `scripts/eval_fvd.py` ships with an I3D-based harness configured for 4 context frames and a 240-frame horizon.
### Key Takeaways
– Open Dreamer reproduces the Dreamer 4 pipeline in JAX/Flax NNX, with training code and a Minecraft demo.
– The dynamics model is **1.6B parameters**, trained for **200K steps** with **Muon**.
– Reported engineering numbers include **57–58% MFU** on B200, **256 frames per GPU**, and approximately **24 GiB** of model state.
– **Stability**, not throughput, was the main bottleneck; loss curves often masked generation-quality regressions.
Check out the blog post and demo, the training repo, the inference repo, and Reactor on X. All credit for this research goes to the researchers behind the project.
