How to train a frontier-level world model

willemhelmet2 pts0 comments

How to train a frontier-level world model

How to train a frontier-level world model

Sponsored by

Reactor

Authors

Diego Marti Monso*

Francesco Sacco*

Edward Hu

* Equal contribution

Repository

github.com/next-state/open-dreamer

Date

July 21, 2026

DOI

10.5281/zenodo.21475232

One of the most time-consuming parts of doing science is climbing up to the shoulders of giants. While this can be a very instructive experience, sometimes it’s nice to be teleported right on top of them and get to work right away, pushing the boundary of human knowledge.

In the last few months, with support from Reactor, we’ve been working on a frontier-level world model. We are releasing and open-sourcing the model and training code, and sharing all the lessons we learned along the way in this blog post. We want to make frontier research accessible to everyone and hope this will accelerate progress in the field.

You can also play with our model thanks to the Reactor runtime!

A survival guide to training a world model

We do these things not because they are easy, but because we thought they were going to be easy.

Figure 1: The search space.

When starting a new project, you have to define:

The goal.

The starting point.

The search space.

Our main objective was to reproduce the Dreamer4 paper . That gave us the search space: essentially, whichever methods and architecture it used. It also gave us the goal: building a Minecraft world model. As a starting point, we chose a simpler game that can be trained on a single GPU: CoinRun.

We purposefully refrained from using different methods that were not in the original paper as that would have made the search space wider.

The starting point

Figure 2: CoinRun. Run to get the coin.

This is CoinRun: it’s a 2D platform game where each level is procedurally generated. As the name suggests, the goal is to run to get the coin.

With this simple game, it is possible to do every step of the Dreamer4 paper, with some nice added benefits:

Just one GPU (much simpler to debug).

A fast iteration loop (much shorter training time).

Once done, the codebase will have all the core logic ready and working.

The search space

The architecture has two main parts: the tokenizer and the dynamics model. Both use the same model-architecture backbone: the block-causal transformer.

What sets it apart is that there are two different attention-based layers:

Space Layer: information propagates among all the elements that represent a single frame.

Causal Time Layer: information propagates between different frames.

The Tokenizer (Autoencoder)

If you have ever trained a VAE in the past, you know how terrible of an experience that is. Most of the time the wisest thing to do is to give up and take an off-the-shelf one.

In this project we’ve trained a transformer-based Masked Autoencoder with ~100x compression, and we probably could have aimed higher.

This architecture is really cool because:

Masking makes the latent space more diffusible.

There is no need for KL or adversarial loss.

It’s very fast to train.

It’s scalable.

All of the above make it easier to train than a typical VAE.

The Dynamics Model (World Model)

World models fundamentally work by doing “next frame prediction”. Each frame is generated as an image conditioned on the previous ones. The dynamics model is trained with diffusion forcing with flow matching and shortcut models .

Figure 3: Tokenizer architecture.

Figure 4: Diffusion forcing training (left) and inference (right).

Other than predicting just the next frame, we worked on predicting the next action as well.<br>Typically, the rollout alternates between the world and the agent. Given the previous action $a_{t-1}$, the dynamics predicts the next state $s_t$; the policy $\pi_t$ observes that state and samples the next action:

Figure 5: A rollout alternates between predicting the next state and sampling the next action.

Evaluating the transition and policy as separate modules would require repeatedly moving between them. It would also prevent them from sharing the same representation and transformer parameters. Instead, we would like to process all timesteps in parallel during training and reuse the same computation for both the world and the agent.

This architecture folds the sequence into blocks:

$$<br>B_t = (a_{t-1}, s_t, \pi_t)<br>$$

Within each block, spatial attention implements $(a_{t-1} \to s_t \to \pi_t)$.

Figure 6: Folding the rollout into timestep blocks. Spatial attention operates within each block, while causal temporal attention connects blocks across time.

What does $s_t$ contain? The state is the part generated through diffusion—or, more precisely, the model diffuses its latent spatial tokens. The state computation therefore includes:

a noise-level and shortcut step-size token describing how far the model should advance in one denoising step, following Shortcut Models ;

spatial latent tokens representing the visual state;

register tokens, which provide...

model world next state space level

Related Articles