Meet Open Dreamer: A JAX/Flax Reproduction of the Dreamer 4 World Model Pipeline, With the Full Training Recipe Published

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.

Two repositories were released. next-state/open-dreamer holds the training pipeline: a causal video tokenizer, an action-conditioned latent dynamics model, rollout generation, and FVD scoring. reactor-team/open-dreamer holds a minimal local rollout harness that generates frames from an MP4 and a matching action file.

A third artifact is the browser demo hosted on the Reactor runtime. It streams a generated Minecraft world in real time and exposes a Game ⟷ Dream toggle that hands the stream from the real game to the world model frame by frame.

The stated objective was to reproduce the Dreamer 4 research . The research team deliberately avoided methods outside that research paper to keep the search space narrow. They started on CoinRun, a procedurally generated 2D platformer trainable on a single GPU, then scaled the working pipeline to Minecraft/VPT-style gameplay video.

Both the tokenizer and the dynamics model use the same block-causal transformer backbone. That backbone alternates two attention types. Space layers propagate information among the elements of a single frame. Causal time layers propagate information between frames.

The tokenizer is a transformer-based Masked Autoencoder rather than a VAE. The team reports roughly 100× compression and notes the design needs no KL or adversarial loss. Masking, they argue, makes the latent space more diffusible.

The dynamics model performs next-frame prediction and is trained with diffusion forcing, flow matching, and shortcut models. It also predicts the next action. Rather than alternating between a separate transition module and policy, the rollout is folded into per-timestep blocks of (previous action, state, policy). Spatial attention runs inside each block; causal temporal attention connects blocks across time.

Critically, world-model tokens cannot read the agent token. Task and policy information can therefore influence future states only through the next action.

The shipped Minecraft configs make the recipe concrete.

The dynamics model is 1.6B parameters: 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 packing_factor: 2 packs neighbouring 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/bootstrap samples switch on at step 100,000 at a 0.25 batch fraction. EMA decay is 0.999.