Introducing Open Dreamer: JAX/Flax reproduction of the Dreamer 4 world model pipeline, complete training recipe now available

A small group of AI researchers (Reactor) has released Open Dreamer, an open implementation of the Dreamer 4 world model pipeline written in JAX and Flax NNX.

what was actually shipped

Two repositories have been released. next-state/open-dreamer holds the training pipeline, including the causal video tokenizer, action-conditional latent dynamics model, rollout generation, and FVD scoring. actor-team/open-dreamer maintains a minimal local rollout harness that generates frames from MP4 and matching action files.

The third artifact is a browser demo hosted on the Reactor runtime. Publishes a game ⟷ Dream Toggle that streams the generated Minecraft world in real time and passes the stream frame by frame from the real game to the world model.

The stated aim was to replicate the Dreamer 4 study. The research team deliberately avoided methods outside of the research paper to keep the scope of the study narrow. They started with CoinRun, a procedurally generated 2D platformer that can be trained on a single GPU, and then expanded their work pipeline to Minecraft/VPT style gameplay videos.

Architecture: One backbone, two models

Both tokenizers and dynamics models use the same block causal transformation backbone. Its backbone alternates between two attention types. Spatial layers convey information between elements of a single frame. The causal temporal layer propagates information between frames.

The tokenizer is a transformer-based masked autoencoder rather than a VAE. The team reports approximately 100x compression and notes that the design requires no KL or adversarial losses. They argue that masking makes the latent space more diffuse.

The dynamics model performs next frame prediction and is trained using diffusion forcing, flow matching, and shortcut models. It also predicts your next action. Rather than switching between separate migration modules and policies, rollouts are collapsed into timestep-by-timestep blocks (previous actions, states, and policies). Spatial attention is performed within each block. Causal temporal attention connects blocks across time.

Importantly, world model tokens cannot read agent tokens. Therefore, task and policy information can only affect future states through the following actions:

Set training recipes

The included Minecraft configuration fleshes out the recipes.

Dynamics model has 1.6B parameters: 30 blocks causal layer, d_model 1920, 30 attention heads for grouped query attention, and 3 KV heads. Every fourth layer is a time attention layer. Each timestep contains 32 learned register tokens. packing_factor: 2 We pack adjacent tokenizer potentials into each dynamics spatial token. Time attention uses a 192-step sliding window.

Training is performed for 200,000 steps using Muon, a WSD schedule, and a peak learning rate of 3e-4. The shortcut/bootstrap sample is turned on at step 100,000 in 0.25 batch fractions. EMA decay is 0.999.

The tokenizer configuration issues 512 latent tokens per frame with a bottleneck width of 16. The raw 360×640 frame is padded to 368×640, so both dimensions are split into 16×16 patches. The encoder depth is 12. d_model 1536. Decoder depth is 8 d_model 1024. The MAE masking probability is highest at 0.9, and LPIPS is applied with a weight of 0.2 at half the timestep.

VPT actions are parsed into 27 binary action channels and 121 categorical mouse classes. There are no continuous channels.

Computemaxing and memory walls

The researchers report that their model’s FLOP utilization is 57-58%, compared to a healthy Transformer training benchmark of 60%. That argument is an attic argument. For B200, the crossover between bandwidth limit and compute limit is 292 FLOPs/byte. Serving 256 frames per GPU will push your workload past that ridge point.

Sharding went in the opposite direction than expected. With 1.6B parameters, the model states (parameters, gradients, optimizer states, EMA) occupy approximately 24 GiB, which fits in B200. Activation had real costs. After experimenting with data parallelism, FSDP, tensor parallelism, and sequence parallelism, the research team settled on simple data parallelism and activation checkpointing.

Data loading was solved by pre-tokenizing the entire dataset. .arrayrecord Save the file and use Grain in the prefetch buffer on the GPU side. ffmpeg decoding was not fast enough to feed the GPU.

Stability section is the actual payload

The research team revealed that they spent the most time on stability. Their key observation: Most stability problems occur in spite of Loss is reduced. MSE improves steadily, but the generation quality decreases.

Six fixes are documented. Muon replaced LaProp, which spiked randomly and frequently in two runs of B200 hours of about 400 hours each. EMA weights are treated as essential for diffusion inference. Mixed precision is bounds dependent. Parameters remain float32, BF16 covers most matmul activation and attention inputs, and float32 is kept for normalization and dynamics flow output heads.

Loss weighting uses x predictions with v spatial losses. This will be a weighting term similar to Dreamer 4, but with a squared denominator. They report small but noticeable improvements. Rollout generation is now more stable with mini-batch centroid-optimal transport between noise and latent sequences. The μ parameterization was tested and determined not to be necessary. This is also because Muon keeps hyperparameters more stable across model sizes.

Further results from the CoinRun phase: iso-FLOPs sweep nearly achieves optimal scaling for computation $N\propto C^{0.56}$ and $D\propto C^{0.44}$.

what’s not in the box

The repository does not contain behavior cloning or RL training loops. The complete Dreamer 4 BC/RL agent loop is listed as an open roadmap item. The CoinRun policy efforts described in the post were not used in Minecraft and were never released.

The FVD score is also not published in this post. scripts/eval_fvd.py comes with an I3D-based harness configured for 4 context frames and a 240 frame horizon.

Important points

• Open Dreamer recreates the Dreamer 4 pipeline in JAX/Flax NNX using training code and a Minecraft demo.
• The dynamics model has 1.6B parameters, 30 layers, d_model In 1920, he trained 200,000 steps with Muon.
• Reported engineering numbers: 57-58% MFU on B200, 256 frames per GPU, up to 24 GiB model state.
• Stability, not throughput, was the bottleneck. The loss curves concealed most generational quality regressions.

---

Check out our blog posts and demos, training repository, inference repository, and Reactor on X. All credit for this study goes to the researchers of this project.