Meet Open Dreamer: A JAX/Flax Replica of the Dreamer 4 World Mannequin Pipeline, With the Full Coaching Recipe Printed

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

What truly shipped

Two repositories had been launched. next-state/open-dreamer holds the coaching pipeline: a causal video tokenizer, an action-conditioned latent dynamics mannequin, rollout technology, and FVD scoring. reactor-team/open-dreamer holds a minimal native rollout harness that generates frames from an MP4 and an identical motion file.

A 3rd artifact is the browser demo hosted on the Reactor runtime. It streams a generated Minecraft world in actual time and exposes a Sport ⟷ Dream toggle that arms the stream from the actual sport to the world mannequin body by body.

The acknowledged goal was to breed the Dreamer 4 analysis. The analysis workforce intentionally prevented strategies outdoors that analysis paper to maintain the search area slender. They began on CoinRun, a procedurally generated 2D platformer trainable on a single GPU, then scaled the working pipeline to Minecraft/VPT-style gameplay video.

Structure: one spine, two fashions

Each the tokenizer and the dynamics mannequin use the identical block-causal transformer spine. That spine alternates two consideration sorts. House layers propagate data among the many components of a single body. Causal time layers propagate data between frames.

The tokenizer is a transformer-based Masked Autoencoder somewhat than a VAE. The workforce reviews roughly 100× compression and notes the design wants no KL or adversarial loss. Masking, they argue, makes the latent area extra diffusible.

The dynamics mannequin performs next-frame prediction and is skilled with diffusion forcing, stream matching, and shortcut fashions. It additionally predicts the following motion. Reasonably than alternating between a separate transition module and coverage, the rollout is folded into per-timestep blocks of (earlier motion, state, coverage). Spatial consideration runs inside every block; causal temporal consideration connects blocks throughout time.

Critically, world-model tokens can’t learn the agent token. Job and coverage data can due to this fact affect future states solely by the following motion.

The coaching recipe, as configured

The shipped Minecraft configs make the recipe concrete.

The dynamics mannequin is 1.6B parameters: 30 block-causal layers, d_model 1920, 30 consideration heads, and three KV heads for grouped-query consideration. Each fourth layer is a time-attention layer. Every timestep carries 32 realized register tokens, and packing_factor: 2 packs neighbouring tokenizer latents into every dynamics spatial token. Time consideration makes use of a 192-step sliding window.

Coaching runs for 200,000 steps with Muon, a WSD schedule, and a peak studying fee of 3e-4. Shortcut/bootstrap samples swap on at step 100,000 at a 0.25 batch fraction. EMA decay is 0.999.

The tokenizer config emits 512 latent tokens per body at a bottleneck width of 16. Uncooked 360×640 frames are padded to 368×640 so each dimensions divide into 16×16 patches. Encoder depth is 12 at d_model 1536; decoder depth is 8 at d_model 1024. MAE masking chance tops out at 0.9, and LPIPS is utilized at weight 0.2 on half the timesteps.

VPT actions are parsed into 27 binary motion channels plus 121 categorical mouse lessons, with no steady channels.

Computemaxxing and the reminiscence wall

The analysis workforce reviews 57–58% mannequin FLOPs utilization, towards a acknowledged 60% benchmark for wholesome transformer coaching. The reasoning is a roofline argument. On a B200, the crossover between bandwidth-bound and compute-bound sits at 292 FLOP/byte. Feeding 256 frames per GPU pushes the workload previous that ridge level.

Sharding went the opposite approach from expectations. At 1.6B parameters the mannequin state — parameters, gradients, optimizer state, and EMA — occupied roughly 24 GiB, which inserts on a B200. Activations had been the actual value. The analysis workforce tried knowledge parallelism, FSDP, tensor parallelism, and sequence parallelism, then settled on plain knowledge parallelism plus activation checkpointing.

Dataloading was solved by pre-tokenizing your entire dataset into .arrayrecord recordsdata, then utilizing Grain with a GPU-side prefetch buffer. ffmpeg decoding was not quick sufficient to maintain the GPUs fed.

The steadiness part is the actual payload

The analysis workforce is express that stability consumed the most important share of their time. Their key commentary: most stability issues happen regardless of the loss taking place. MSE improves easily whereas technology high quality degrades.

Six fixes are documented. Muon changed LaProp, which spiked randomly and more and more usually, throughout two runs of roughly 400 B200 hours every. EMA weights are handled as necessary for diffusion inference. Combined precision is boundary-sensitive: parameters keep float32, BF16 covers most matmul activations and a spotlight inputs, and float32 is stored for normalization and the dynamics stream output head.

On loss weighting, they use x-prediction with a v-space loss, which reduces to a weighting time period much like Dreamer 4’s however with a squared denominator. They report a small however noticeable enchancment. Minibatch barycentric optimum transport between noise and latent sequences made rollout technology extra steady. μ-parametrization was examined and judged pointless, partly as a result of Muon holds hyperparameters steadier throughout mannequin sizes.

One additional consequence from the CoinRun section: an iso-FLOPs sweep put compute-optimal scaling at roughly $N\propto C^{0.56}$ and $D\propto C^{0.44}$.

What isn’t within the field

The repository doesn’t embody the behaviour-cloning or RL coaching loop; a full Dreamer 4 BC/RL agent loop is listed as an open roadmap merchandise. The CoinRun coverage work described within the publish was not used for Minecraft and was not launched.

The publish additionally doesn’t 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 coaching code and a Minecraft demo.
• The dynamics mannequin is 1.6B parameters, 30 layers, d_model 1920, skilled 200K steps with Muon.
• Reported engineering numbers: 57–58% MFU on B200, 256 frames per GPU, ~24 GiB mannequin state.
• Stability, not throughput, was the bottleneck; loss curves hid most generation-quality regressions.

Take a look at the weblog publish and demo, the coaching repo, the inference repo, and Reactor on X. All credit score for this analysis goes to the researchers of this mission.