world-models-carracing-v3

🏎️ 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×3 frame into a 32-d latent z
• 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

left: what the car sees · middle: the frame round-tripped through V · right: the next frame as M predicts it, one step ahead.

Files

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