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 launched Open Dreamer, an open implementation of the Dreamer 4 world-model pipeline written in JAX and Flax NNX.
What really 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 era, 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 third artifact is the browser demo hosted on the Reactor runtime. It streams a generated Minecraft world in actual time and exposes a Game ⟷ Dream toggle that palms the stream from the actual sport to the world mannequin body by body.
The acknowledged goal was to breed the Dreamer 4 research. The analysis crew intentionally prevented strategies exterior that analysis paper to maintain the search house slim. 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.
Architecture: one spine, two fashions
Both the tokenizer and the dynamics mannequin use the identical block-causal transformer spine. That spine alternates two consideration varieties. Space layers propagate data amongst the components of a single body. Causal time layers propagate data between frames.
The tokenizer is a transformer-based Masked Autoencoder fairly than a VAE. The crew studies roughly 100× compression and notes the design wants no KL or adversarial loss. Masking, they argue, makes the latent house extra diffusible.
The dynamics mannequin performs next-frame prediction and is educated with diffusion forcing, movement matching, and shortcut fashions. It additionally predicts the subsequent motion. Rather 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 not learn the agent token. Task and coverage data can subsequently affect future states solely via the subsequent 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. Every fourth layer is a time-attention layer. Each 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.
Training 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. Raw 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 crew studies 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 different 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 price. The analysis crew tried knowledge parallelism, FSDP, tensor parallelism, and sequence parallelism, then settled on plain knowledge parallelism plus activation checkpointing.
Dataloading was solved by pre-tokenizing the complete dataset into .arrayrecord information, then utilizing Grain with a GPU-side prefetch buffer. ffmpeg decoding was not quick sufficient to maintain the GPUs fed.
The stability part is the actual payload
The analysis crew is specific that stability consumed the largest share of their time. Their key remark: most stability issues happen regardless of the loss taking place. MSE improves easily whereas era 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 obligatory for diffusion inference. Mixed 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 movement 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 era extra secure. μ-parametrization was examined and judged pointless, partly as a result of Muon holds hyperparameters steadier throughout mannequin sizes.
One additional consequence from the CoinRun part: an iso-FLOPs sweep put compute-optimal scaling at roughly and .
What is just not in the field
The repository doesn’t embrace 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 in the submit was not used for Minecraft and was not launched.
The submit 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_model1920, educated 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.
Check out the blog post and demo, the training repo, the inference repo, and Reactor on X. All credit score for this analysis goes to the researchers of this challenge.
The submit Meet Open Dreamer: A JAX/Flax Reproduction of the Dreamer 4 World Model Pipeline, With the Full Training Recipe Published appeared first on MarkTechPost.
