NewsMacroOpen Dreamer: Reactor Releases Open-Source JAX/Flax Reproduction of Dreamer 4 World Model Pipeline with Full Training Recipe

Open Dreamer: Reactor Releases Open-Source JAX/Flax Reproduction of Dreamer 4 World Model Pipeline with Full Training Recipe

Author: MarkTechPost·

Key Takeaways

  • Open Dreamer provides an open-source implementation of the Dreamer 4 pipeline with a causal video tokenizer, action-conditioned latent dynamics model, rollout utilities, and FVD evaluation code.
  • The Minecraft dynamics model uses 1.6 billion parameters across 30 block-causal layers and is configured for 200,000 training steps with the Muon optimizer.
  • Reactor reports 57% to 58% model FLOPs utilization on NVIDIA B200 GPUs, with activations identified as the main memory cost rather than model state.
  • The team documented multiple stability measures, including EMA use, mixed-precision boundaries, optimizer changes, loss weighting, and minibatch optimal transport.
  • The release does not include behavior-cloning or reinforcement-learning training loops, and the project has not published FVD scores.
Open Dreamer: Reactor Releases Open-Source JAX/Flax Reproduction of Dreamer 4 World Model Pipeline with Full Training Recipe

A research group operating under the name Reactor has released Open Dreamer, an open-source implementation of the Dreamer 4 world-model pipeline built in JAX and Flax NNX. World models learn to predict how an environment evolves from observation data, enabling agents to plan, learn, and generate video without interacting with a real system. The Dreamer series, developed by Danijar Hafner, has been one of the most studied lines of work in model-based reinforcement learning, with Dreamer 4 extending the approach toward large-scale video world models. The project aims to faithfully reproduce the Dreamer 4 research methodology while making the full training pipeline — including stability fixes and compute configurations — publicly accessible, a level of detail that large-scale world-model releases have typically omitted.

Released Artifacts

Two code repositories were published. The first, next-state/open-dreamer, contains the complete training pipeline: a causal video tokenizer, an action-conditioned latent dynamics model, rollout generation utilities, and FVD (Fréchet Video Distance) scoring. The second, reactor-team/open-dreamer, provides a minimal local rollout harness capable of generating video frames from an input MP4 file paired with a corresponding action file.

A third deliverable is a browser-based demo running on the Reactor runtime. It streams a generated Minecraft environment in real time and includes a Game ⟷ Dream toggle that transitions the video stream between the actual game and the world model output on a frame-by-frame basis.

The team stated that their objective was to reproduce the Dreamer 4 research deliberately, avoiding techniques outside the original paper to maintain a narrow search space. Development began on CoinRun, a procedurally generated 2D platformer that can be trained on a single GPU, before the working pipeline was scaled to Minecraft/VPT-style gameplay video. Minecraft has served as a benchmark domain for learning agents from gameplay demonstrations since OpenAI's Video PreTraining (VPT) project, which trained an inverse dynamics model on contractor-labeled YouTube footage to produce large-scale action-labeled gameplay data.

Additional details are available on the project blog, the arXiv paper, and Reactor's announcement on X.

Architecture: One Backbone, Two Models

Both the tokenizer and the dynamics model share the same block-causal transformer backbone, which alternates between two types of attention. Space layers propagate information across elements within a single frame, while causal time layers propagate information across frames.

The tokenizer is designed as a transformer-based Masked Autoencoder (MAE) rather than a traditional variational autoencoder (VAE). The team reports approximately 100× compression and notes that this design eliminates the need for KL divergence or adversarial losses. They argue that the masking approach renders the latent space more amenable to diffusion-based generation.

The dynamics model performs next-frame prediction using diffusion forcing, flow matching, and shortcut models. It also predicts the next action. Instead of alternating between a separate transition module and a policy module, the rollout is integrated into per-timestep blocks structured as (previous action, state, policy). Spatial attention operates within each block, and causal temporal attention links blocks across time.

A critical design constraint: world-model tokens cannot read the agent token. Consequently, task and policy information can influence future states only through the predicted next action.

Training Configuration

The shipped Minecraft configuration files specify the training recipe in detail.

The dynamics model comprises 1.6 billion parameters across 30 block-causal layers, with d_model of 1920, 30 attention heads, and 3 KV heads using grouped-query attention. Every fourth layer functions as a time-attention layer. Each timestep carries 32 learned register tokens, and a packing factor of 2 consolidates neighboring tokenizer latents into each dynamics spatial token. Time attention operates over a 192-step sliding window.

Training spans 200,000 steps using the Muon optimizer with a WSD (warmup-stable-decay) schedule and a peak learning rate of 3e-4. Shortcut and bootstrap sampling activates at step 100,000 with a 0.25 batch fraction. Exponential moving average (EMA) decay is set at 0.999.

The tokenizer configuration produces 512 latent tokens per frame at a bottleneck width of 16. Raw 360×640 frames are padded to 368×640 so both spatial dimensions divide evenly into 16×16 patches. Encoder depth is 12 with d_model 1536; decoder depth is 8 with d_model 1024. MAE masking probability reaches 0.9 at maximum, and LPIPS loss is applied at a weight of 0.2 on half of the timesteps.

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

Compute Performance and Memory Strategy

The team reports 57–58% model FLOPs utilization (MFU), compared to a commonly cited 60% benchmark for healthy transformer training. Their analysis uses a roofline argument: on an NVIDIA B200, the crossover between bandwidth-bound and compute-bound operation occurs at 292 FLOP/byte. Processing 256 frames per GPU pushes the workload beyond that ridge point into the compute-bound regime.

Sharding decisions defied initial expectations. At 1.6 billion parameters, the full model state — parameters, gradients, optimizer state, and EMA — occupied approximately 24 GiB, which fits within a single B200. Activations, rather than model state, proved to be the primary memory cost. The team experimented with data parallelism, fully sharded data parallelism (FSDP), tensor parallelism, and sequence parallelism before settling on plain data parallelism combined with activation checkpointing.

For dataloading, the team pre-tokenized the entire dataset into .arrayrecord files and used Grain with a GPU-side prefetch buffer. Standard ffmpeg-based decoding was insufficient to keep the GPUs saturated.

Stability Engineering

The research team explicitly states that stability issues consumed the largest share of their development time. Their central observation: most stability problems arise even as the loss continues to decrease. Mean squared error (MSE) improves smoothly while generation quality simultaneously degrades — a phenomenon that makes conventional loss-based monitoring unreliable for diffusion-based world models.

Six specific fixes are documented:

  1. Optimizer change: Muon replaced LaProp, which exhibited random and increasingly frequent spikes across two separate training runs of approximately 400 B200 hours each.

  2. EMA as mandatory: EMA weights are treated as essential for diffusion inference, not optional.

  3. Mixed precision boundaries: Parameters remain in float32, BF16 covers most matmul activations and attention inputs, and float32 is retained for normalization layers and the dynamics flow output head.

  4. Loss weighting: The team uses x-prediction with a v-space loss, which reduces to a weighting term similar to Dreamer 4's formulation but with a squared denominator. They report a small but noticeable improvement.

  5. Optimal transport: Minibatch barycentric optimal transport applied between noise and latent sequences improved rollout generation stability.

  6. μ-parametrization: Tested and deemed unnecessary, partly because Muon maintains hyperparameter stability more effectively across model sizes.

One additional finding from the CoinRun development phase: an iso-FLOPs sweep estimated compute-optimal scaling at approximately N ∝ C^0.56 and D ∝ C^0.44.

What Is Not Included

The repository does not contain the behaviour-cloning (BC) or reinforcement learning (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 project documentation was not used for the Minecraft implementation and was not released.

The project does not publish FVD scores, although the repository ships scripts/eval_fvd.py, an I3D-based evaluation harness configured for 4 context frames and a 240-frame horizon.

Summary of Key Engineering Figures

  • Dynamics model: 1.6B parameters, 30 layers, d_model 1920, trained for 200,000 steps with Muon
  • Hardware efficiency: 57–58% MFU on NVIDIA B200 GPUs, 256 frames per GPU, approximately 24 GiB model state
  • Primary challenge: Stability, not throughput — loss curves concealed most generation-quality regressions