World Models: Minecraft

Introduction

I built a 2.1M-parameter latent world model for Minecraft to explore the core idea behind World Action Models (WAMs), something very relevant to the field of robotics. Using video and player actions, the model learns to predict how the game changes and recursively generates the next future.

Through several model iterations and experiments, I developed a system that responds meaningfully to player inputs, and uncovered the challenges to maintaining visual quality over long rollouts.

What better way to learn about a topic than in the context of one of my favourite video games growing up as a kid.

Background

A world model learns how an environment changes over time. Given the current observation, $o_t$, and an action, $a_t$, it predicts the next observation:

$$\hat{o}_{t+1} = f_\theta(o_t, a_t)$$

And so for Minecraft, the current observation is the current frame we are looking at. The action includes controls such as movement and camera movement. For this experiment, I limited the action space to a nine-dimensional vector that contains movement and camera control:

$$\mathbf{a}_t = \begin{bmatrix} W_t & A_t & S_t & D_t & J_t & \mathrm{Sprint}_t & \mathrm{Sneak}_t & \Delta x_t & \Delta y_t \end{bmatrix} \in [0,1]^7 \times \mathbb{R}^2$$

The first seven values represent keyboard controls, while $\Delta x_t$ and $\Delta y_t$ represent the horizontal and vertical camera movement. Actions such as attacking (typically, left click), using items, and interacting with GUIs like crafting and inventory were excluded for simplicity.

Predicting the next frame in pixel space can be expensive, so I used an encoder to compress each observation into a smaller latent representation:

$$z_t = E(o_t)$$

I then employed a dynamics model to use recent latent states and the current action to predict the next latent state:

$$\hat{z}_{t+1} = F_\theta(z_{t-1}, z_t, a_t)$$

It is important to use frames observed prior to time, $t$, so that the model can infer motion between scenes to determine the future state.

Finally, a decoder converts the predicted latent back into an image in pixel space:

$$\hat{o}_{t+1} = D(\hat{z}_{t+1})$$

Predicting one frame is useful, but a world model often feeds its predictions back into itself to predict further states:

$$\hat{z}_{t+2} = F_\theta(z_t, \hat{z}_{t+1}, a_{t+1})$$

Or more generally:

$$\hat{z}_{t+h+1} = F_\theta(\hat{z}_{t+h-1}, \hat{z}_{t+h}, a_{t+h}), \qquad h=1,\ldots,H$$

The resulting frames form an imagined trajectory through the game. This recursive process is also a huge challenge of world modelling. Each predicted state contains some error, and that error then becomes part of the input of the following prediction. Therefore these errors quickly compound for longer rollouts.

Data Pipeline and Infrastructure

I built a data pipeline around OpenAI's Video Pretraining project which contains public recording of Minecraft gameplay footage.

I used roughly 100GiB of footage, totalling roughly 700 synchronized video/action episodes. My preprocessing involved reducing the footage to 64x64 at 10Hz, filtering data based on my newly defined action space, and separating it across training, validation, and test sets.

The processed videos and actions were then saved to an S3 bucket for training.

Model Development

Flat vs. Spatial Reconstruction

My first autoencoder compressed each frame into a flat vector of 256 values. It preserved broad colours, but block edges and small images were already blurry even before passing the latent through any model. That means the dynamics model that I was building was being asked to make a prediction with a representation that discarded a lot of the detail I cared about.

I replaced the flat vector with a spatial latent arranged as a 16x16 grid, with 16 learned values at each position. I also trained the autoencoder to preserve both pixel values and the boundaries between neighbouring pixels. When tested on frames it had not seen during training, this reduced the reconstruction MSE from 0.00149 to 0.00018 which is an 88% reduction. The qualitative improvement is evident as well.

V1: Deterministic World Model

With a better visual representation in place, I trained V1: a small deterministic world model that receives the previous latent, current latent, and current action, then predicts the next spatial latent. The dynamics model is a convolutional residual network that operates directly on the 16x16 latent grid. The model contains a 1.84M-parameter dynamics model and a 253K-parameter spatial autoencoder. During training, I measured the MSE between the predicted and real next state in both latent and decoded pixel space. V1 does not use random sampling, meaning the same inputs always produce the same prediction.

Because consecutive video frames contain much of the same scene, V1 does not predict the entire next latent from scratch. Instead, it predicts how the current latent should change:

$$\hat z_{t+1}=z_t+\Delta_\theta(z_{t-1},z_t,a_t).$$

This allows the model to focus on what changed between frames, such as camera movement or the player moving forward, while carrying the rest of the current state into the next prediction. After learning one-step prediction, I trained V1 over ten recursive steps so that it also learned to operate on its own previous outputs.

I evaluated the final model on 5,000 sequences it had not seen during training. At the model's 10 Hz frame rate, twenty recursive steps represent approximately two seconds of imagined gameplay. After twenty recursive steps, V1 achieved a pixel MSE of 0.0196, compared with 0.0321 for a baseline that simply repeated the last real frame: a 39.1% reduction. When the player actions were shuffled, its error increased by 40.8%. This shows that V1 was doing more than copying the previous frame: its predictions meaningfully depended on the supplied actions.

The visual results show both its success and its limitation. Forward movement, camera turns, and idle inputs produce visibly different futures, but the frames become progressively smoother as prediction errors accumulate. V1 learned how the controls affect the world, but it struggled to preserve the world's visual detail over longer rollouts.

V1 Ablations

I tested the three most obvious explanations with controlled ablations. Increasing the number of unique training windows from 10K to 100K reduced twenty-step pixel MSE from 0.0232 to 0.0209 and made the model more action-dependent. Training on recursively generated context for ten steps produced a better fidelity/control balance than either five or twenty steps. Increasing dynamics capacity from 255K to 1.84M parameters reduced twenty-step MSE again, from 0.0206 to 0.0196. These are consistent improvements, but the filmstrips still look broadly alike: none of the runs produces a sharp long-horizon simulation.

V2: The Diffusion Hypothesis

V1 always predicts a single next state. When the training data contains several plausible futures, minimizing MSE can encourage the model to average between them, producing a blurry prediction. My hypothesis was that diffusion could avoid this averaging by learning to generate a plausible next state instead of predicting one fixed answer.

For V2, I kept the same spatial autoencoder but replaced the deterministic dynamics model with a 19.3M-parameter diffusion model. The diffusion model uses a multiscale U-Net with cross-attention to condition each denoising step on the action history. Rather than predicting the next latent directly, it begins with random noise and gradually turns that noise into a predicted next state. This also means that the same inputs can produce different predictions.

V2 was trained on 409,910 clean sequences. To test whether it actually used the controls, I replaced the correct action with an action from a different sequence while keeping the rest of the input unchanged. This made its prediction loss 13.6% worse, showing that the supplied action influenced its prediction. However, the recursive rollouts revealed a larger problem. Each generated frame introduced small changes that were not necessarily consistent with the previous scene. When those frames were fed back into the model, the changes accumulated and the geometry quickly began to drift.

V2 sometimes produced more texture than V1 in an individual frame, but it was less stable over time. Trees, terrain, and camera geometry changed in ways that were not explained by the player's actions. V1 became blurry while preserving the broad scene; V2 produced a more detailed image that was less consistent with the world it had generated one step earlier. Diffusion changed the appearance of the error, but it did not solve the underlying problem of maintaining a coherent world over a recursive rollout.

Closing Thoughts

Despite its limitations, V1 successfully captured Minecraft's short-term dynamics: in the first one or two predicted frames, movement is visible in the direction implied by the player's action. V2 then tested whether MSE averaging was the main cause of blur. The diffusion model occasionally produced more texture, but its world drifted more quickly. This showed that generating a plausible individual frame was not enough; maintaining a consistent scene across recursive predictions remained the harder problem.

The gap to larger Minecraft world models is still substantial. V1 contains 2.1M parameters and was trained on a laptop, compared with the 500M-parameter Oasis model and the substantially larger MineWorld and Matrix-Game experiments. My ablations do not prove that scale alone solves the problem, but they suggest that the project was not one small adjustment away from a sharp long-term simulation. A future model should use balanced movement data and recorded position and camera angles, then be evaluated on movement direction, magnitude, and idle stability, not pixel MSE alone.

The broader lesson I would carry into robotics is that visual representation and action-conditioned dynamics must develop together. The flat latent discarded important spatial detail before the dynamics model could use it, while the spatial latent preserved a much stronger representation of the scene. Projects such as EgoScale suggest a promising path forward: use large collections of egocentric human demonstrations to learn about objects, geometry, and how scenes change, then use robot action data to connect that visual knowledge to actions the robot can execute.

You can check out the code here.