🏎️ World Models — CarRacing-v3
A faithful reproduction of Ha & Schmidhuber's World Models
on CarRacing-v3. The agent factorises into three parts — Vision, Memory,
Controller — trained in that order:
- V — a β-VAE that compresses each
64×64×3frame into a 32-d latentz - M — an MDN-RNN (LSTM-256, 5-mixture density head) that predicts the next latent
p(z′ | z, a, h) - C — a single 867-parameter linear layer mapping
[z; h] → action, evolved with CMA-ES
Only the controller ever touches the reward; V and M are trained once, self-supervised, then frozen.
🎯 Result
| Metric | This model | Paper (Ha & Schmidhuber) |
|---|---|---|
| Best-agent reward (avg / 100 rollouts) | 915.9 | 906 ± 21 |
left: what the car sees · middle: the frame round-tripped through V · right: the next frame as M predicts it, one step ahead.
Files
| File | Module | Architecture |
|---|---|---|
vae.pt | V | AutoEncoder — 4× stride-2 conv encoder [32→64→128→256], mirror deconv decoder, 32-d latent, β-VAE with free-bits floor (λ = 0.5/dim) |
rnn.pt | M | RNN — LSTM (hidden 256) over [z; a] (35-d) + MDN head, 5 Gaussians × 32 dims |
controller.pt | C | linear [z(32); h(256)] → a(3), 867 params, CMA-ES (popsize 64, avg 16, σ 0.3) |
model.py | — | the module definitions |
config.yaml | — | hyperparameters for instantiation |
Usage
import torch
from omegaconf import OmegaConf
from huggingface_hub import hf_hub_download
from model import AutoEncoder, RNN # model.py from this repo
repo = "flydexo/world-models-carracing-v3"
cfg = OmegaConf.load(hf_hub_download(repo, "config.yaml"))
vae = AutoEncoder(cfg)
vae.load_state_dict(torch.load(hf_hub_download(repo, "vae.pt"), map_location="cpu"))
rnn = RNN(cfg)
rnn.load_state_dict(torch.load(hf_hub_download(repo, "rnn.pt"), map_location="cpu"))
# Controller: a plain linear [z; h] -> action
ctrl = torch.nn.Linear(cfg.controller.state_dim + cfg.controller.hidden_dim,
cfg.controller.action_dim)
ctrl.load_state_dict(torch.load(hf_hub_download(repo, "controller.pt"), map_location="cpu"))
Rollout loop: encode obs → z, concat [z; h] → controller → action, step env, feed
[z; a] through the RNN to advance the hidden state h.
Reproduction notes
The gap between a naïve implementation (~600) and the paper (~906) came down to a few details:
- VAE — sum-reduced reconstruction paired with a free-bits KL floor (λ = 0.5/dim), KL scaled consistently against the recon term. No posterior collapse — all 32 latents stay alive.
- MDN-RNN — trained on
z ~ N(μ, σ)sampled every batch (not the mean μ); softmax temperature applied only at sampling, never inside the training loss; correct mixture sampling. - Controller — input is
[z; h](latent plus the RNN hidden state). - CMA-ES — population 64, 16 rollouts averaged per candidate, σ = 0.3.
Links
- 📄 Paper: World Models (Ha & Schmidhuber, 2018)
- 🤗 Collection: World Models
- 📊 Live training dashboards (Trackio): VAE sweep · RNN / controller