You can deploy this model and start earning money today!

flydexo/
world-models-carracing-v3

Name: world-models-carracing-v3
Availability: InStock

pytorch

reinforcement-learning

Views

No views yet

🏎️ 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

python

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.

flydexo/world-models-carracing-v3