Car Racing World Model

A PyTorch implementation of the World Models architecture (Ha & Schmidhuber, 2018) applied to the CarRacing-v3 environment from Gymnasium.

The agent learns entirely inside a learned "dream" — it never trains the controller in the real environment. Instead, a VAE compresses visual observations into a compact latent space, an MDN-RNN predicts future latent states, and a lightweight Controller is evolved with CMA-ES using imagined rollouts.

---

Architecture

Observation (96×96 RGB)
        │
        ▼
  ┌───────────┐
  │    VAE    │  →  latent vector z (32-dim)
  └───────────┘
        │
        ▼
  ┌───────────┐
  │  MDN-RNN  │  →  predicts next z + reward
  └───────────┘
        │
        ▼
  ┌────────────┐
  │ Controller │  →  action (steering, gas, brake)
  └────────────┘
        │
     evolved via CMA-ES on dream rollouts

Components

File	Description
VAE.py	Convolutional VAE — encodes 64×64 frames into a 32-dim latent space
MDRNN.py	MDN-RNN — LSTM with a Mixture Density Network head; predicts future latent states and rewards
controller.py	Linear controller — maps (z, h) → action
collect_data.py	Collects random rollouts from the real environment and saves them as .npz files
train_vae.py	Trains the VAE on collected frames
train_mdrnn.py	Encodes the dataset with the trained VAE, then trains the MDN-RNN
cma_es.py	Optimises the controller weights using CMA-ES over batched dream rollouts
dream.py	Runs imagination rollouts (single and batched) entirely in latent space

---

Setup

pip install torch gymnasium[box2d] numpy tqdm cma

Tested with Python 3.10+. On Apple Silicon, MPS acceleration is used automatically.

---

Usage

1. Collect training data

python collect_data.py

Saves random rollout episodes to ./data/ as compressed .npz files (frames, actions, rewards).

2. Train the VAE

python train_vae.py

Trains the visual model on collected frames. Saves checkpoint to checkpoint/.

3. Train the MDN-RNN

python train_mdrnn.py

Loads the trained VAE, encodes the dataset into latent sequences, and trains the memory model.

4. Evolve the Controller

python cma_es.py

Uses CMA-ES to optimise controller weights entirely through dream rollouts in latent space.

---

Model Details

VAE — 4-layer convolutional encoder/decoder, latent dim = 32, trained with MSE reconstruction loss + β-weighted KL divergence.

MDN-RNN — Single-layer LSTM (hidden dim = 256) feeding into a Mixture Density Network (5 Gaussian mixtures) over the latent space. Also predicts per-step reward via a separate linear head.

Controller — Single linear layer mapping the concatenated latent vector z (32-dim) and LSTM hidden state h (256-dim) directly to a 3-dim continuous action (steering, gas, brake).

---

References

• Ha, D., & Schmidhuber, J. (2018). World Models. arXiv:1803.10122
• Gymnasium CarRacing-v3