CoolFace
Modelpublic

flydexo/world-models-carracing-v3

sourceHugging Facemitupdated 3mo agoView on Hugging Face
0likes3downloads
Model Card

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

MetricThis modelPaper (Ha & Schmidhuber)
Best-agent reward (avg / 100 rollouts)915.9906 Β± 21

[image]

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

Files

FileModuleArchitecture
vae.ptVAutoEncoder β€” 4Γ— stride-2 conv encoder [32β†’64β†’128β†’256], mirror deconv decoder, 32-d latent, Ξ²-VAE with free-bits floor (Ξ» = 0.5/dim)
rnn.ptMRNN β€” LSTM (hidden 256) over [z; a] (35-d) + MDN head, 5 Gaussians Γ— 32 dims
controller.ptClinear [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.

Links