world models and using them to dream drive
Aug 25, 2026
You can find the full code here: https://github.com/pandey-ps/world-models/tree/master/drive (it is well commented to explain about the environment of the model).
on a higher level, world models are made up of three parts: an encoder, a state space model, and a decoder. In this post I build these from scratch using the variational autoencoder.
an autoregressive model predicts the next observation from the past as:
\[x_{t+1} \sim p_\theta(x_{t+1} \mid x_{\leq t})\]but for a world model, it also depends on an action the agent takes, as:
\[x_{t+1} \sim p_\theta(x_{t+1} \mid x_{\leq t}, a_t)\]world model acting can be classified in two ways on how it can update its latent:
- open loop (dream):
- closed loop (peek):
open loop needs no observation but lets small errors compound; closed loop re-grounds on every step. for this self driving world model, it is a complete open loop i.e. dream world model.
In this setup:
| Component | Architecture | Details |
|---|---|---|
| encoder | 4-layer Conv2d VAE | 3 -> 32 -> 64 -> 128 -> 256, stride 2, ReLU, Flatten |
| latent space | 32-dim Gaussian | z_dim=32, fc_mu/fc_logvar from 4096 -> 32 |
| state space | MDN-RNN (LSTM + 5 Gaussian mixtures) | hidden_dim=128, num_gaussians=5, LSTM |
| decoder | ConvTranspose | 4x4 -> 8x8 -> 16x16 -> 32x32 -> 64x64 -> sigmoid |
| training | Adam (lr=1e-3), gradient clipping (max_norm=5.0) |
VAE: 30 epochs; RNN: 30 epochs; batch=32 |
The input frame is compressed by the VAE into the latent representation z. The MDN-RNN then uses latent z, the current action, and its hidden memory of previous frames to predict the next state. It gives outcome as a probability distribution, from which we sample the next z. The VAE decoder turns z back into a frame, and the frames are then used to generate the video.
