flydexo/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

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