Add World Models CarRacing-v3 (V+M+C), 915.9 best-agent reward

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+controller_triptych.gif filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,89 @@
+---
+license: mit
+library_name: pytorch
+tags:
+  - reinforcement-learning
+  - world-models
+  - vae
+  - mdn-rnn
+  - cma-es
+  - car-racing
+  - gymnasium
+pipeline_tag: reinforcement-learning
+---
+
+# 🏎️ World Models — CarRacing-v3
+
+A faithful reproduction of Ha & Schmidhuber's [*World Models*](https://arxiv.org/abs/1803.10122)
+on `CarRacing-v3`. The agent factorises into three parts — **V**ision, **M**emory,
+**C**ontroller — 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** |
+
+![Controller-driven episode](controller_triptych.gif)
+
+*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.
+
+## Links
+
+- 📄 Paper: [World Models](https://arxiv.org/abs/1803.10122) (Ha & Schmidhuber, 2018)
+- 🤗 Collection: [World Models](https://huggingface.co/collections/flydexo/world-models-6a493823e48400161b1cd828)
+- 📊 Live training dashboards (Trackio): [VAE sweep](https://huggingface.co/spaces/flydexo/ha_schmidhuber-vae) · [RNN / controller](https://huggingface.co/spaces/flydexo/ha_schmidhuber)

config.yaml

@@ -0,0 +1,12 @@
+# Hyperparameters needed to instantiate model.py (AutoEncoder / RNN / Controller).
+# Load with OmegaConf and pass to the constructors — see the README for a snippet.
+rnn:
+  hidden_size: 256
+  num_mix: 5
+  z_dim: 32
+  action_dim: 3
+  temp: 0.2          # softmax temperature — applied only at sampling, never in training
+controller:
+  state_dim: 32      # z
+  hidden_dim: 256    # h (LSTM hidden state)
+  action_dim: 3

controller.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e11cd312ef37b54802d2ea5ab20ac5817627eec6ea3ce65760c393d57c963d73
+size 5401

model.py

@@ -0,0 +1,164 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.distributions.normal import Normal
+from torch.nn.modules.rnn import LSTM
+
+
+class PrintShape(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        print("hey", x.shape)
+        return x
+
+
+class Dense(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.mu = nn.Linear(1024, 32)
+        self.log_sigma = nn.Linear(1024, 32)
+
+    def forward(self, x):
+        x = x.flatten(start_dim=1)
+        mu = self.mu(x)
+        log_sigma = self.log_sigma(x)
+        z = mu + torch.exp(log_sigma) * torch.randn_like(mu)
+        return z, mu, log_sigma
+
+
+class Fit(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        return x.unsqueeze(-1).unsqueeze(-1)
+
+
+class AutoEncoder(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.cfg = cfg
+        self.conv = nn.Sequential(
+            nn.Conv2d(3, 32, 4, 2),
+            nn.ReLU(),
+            nn.Conv2d(32, 64, 4, 2),
+            nn.ReLU(),
+            nn.Conv2d(64, 128, 4, 2),
+            nn.ReLU(),
+            nn.Conv2d(128, 256, 4, 2),
+            nn.ReLU(),
+        )
+        self.dense = Dense()
+        self.decoder = nn.Sequential(
+            nn.Linear(32, 1024),
+            Fit(),
+            nn.ConvTranspose2d(1024, 128, 5, 2),
+            nn.ReLU(),
+            nn.ConvTranspose2d(128, 64, 5, 2),
+            nn.ReLU(),
+            nn.ConvTranspose2d(64, 32, 6, 2),
+            nn.ReLU(),
+            nn.ConvTranspose2d(32, 3, 6, 2),
+            nn.Sigmoid(),
+        )
+
+    def encode(self, x):
+        # returns (z, mu, log_sigma)
+        return self.dense(self.conv(x))
+
+    @staticmethod
+    @torch.compile
+    def kl_divergence(mu, log_sigma):
+        # KL(N(mu, sigma^2) || N(0,1)) summed over latent dims, averaged over the batch.
+        # log_sigma is log-STD (Dense samples with std = exp(log_sigma)), so variance is
+        # exp(2*log_sigma). This is 2x the textbook KL -- the global 0.5 is dropped to match
+        # the sum-reduced MSE recon (also 2x a unit-variance Gaussian NLL), keeping the recon:KL
+        # scale (and thus beta / the free-bits floor) consistent.
+        var = torch.exp(2 * log_sigma)
+        return 0.5 * (mu.pow(2) + var - 2 * log_sigma - 1).sum(-1).mean()
+
+    def forward(self, x):
+        # x.shape = B * C * H * W
+        z, mu, log_sigma = self.encode(x)
+        # (z,mu,log_sigma).shape = B * 32
+        x_recon = self.decoder(z)
+        kl = self.kl_divergence(mu, log_sigma)
+        return x_recon, kl
+
+
+class MDN(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        h = cfg.rnn.hidden_size
+        self.gaussians = cfg.rnn.num_mix
+        self.z_dim = cfg.rnn.z_dim
+        self.temp = cfg.rnn.temp
+        # self.layer = nn.Sequential(nn.Linear(h, h), nn.ReLU())
+        self.probs_layer = nn.Linear(h, self.gaussians)
+        self.means = nn.Linear(h, self.gaussians * self.z_dim)
+        self.stds = nn.Linear(h, self.gaussians * self.z_dim)
+
+    def forward(self, x):
+        # fix #5: return distribution params for NLL loss, not a sampled point
+        # x.shape = (B, 256)
+        # x = self.layer(x)
+        temp = self.temp if not (self.training) else 1
+        pi = F.softmax(self.probs_layer(x) / temp, dim=-1)  # (B, 5)
+        mu = self.means(x).view(-1, self.gaussians, self.z_dim)  # (B, 5, 32)
+        sigma = torch.exp(self.stds(x)).view(
+            -1, self.gaussians, self.z_dim
+        )  # (B, 5, 32)
+        return pi, mu, sigma
+
+    def sample(self, pi, mu, sigma):
+        # fix #6: correct mixture sampling — pick one component, then sample from it
+        # pi: (B, 5), mu/sigma: (B, 5, 32)
+        k = torch.multinomial(pi, num_samples=1).squeeze(
+            -1
+        )  # (B,) — hard component draw
+        B = mu.shape[0]
+        mu_k = mu[torch.arange(B), k]  # (B, 32)
+        sigma_k = sigma[torch.arange(B), k]  # (B, 32) — temperature scales uncertainty
+        return Normal(mu_k, sigma_k).sample()  # (B, 32)
+
+    @staticmethod
+    def loss(pi, mu, sigma, target, mask=None):
+        # Works for any prefix shape: (B, 32) or (B, T, 32)
+        log_pi = torch.log(pi + 1e-8)  # (..., K)
+        log_prob = Normal(mu, sigma).log_prob(target.unsqueeze(-2))  # (..., K, 32)
+        nll = -torch.logsumexp(log_pi + log_prob.sum(-1), dim=-1)  # (...)
+        return nll[mask].mean() if mask is not None else nll.mean()
+
+
+class RNN(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.lstm = LSTM(
+            cfg.rnn.z_dim + cfg.rnn.action_dim, cfg.rnn.hidden_size, batch_first=True
+        )
+        self.mdn = MDN(cfg)
+
+    def forward(self, z, a, hidden=None):
+        # z: (B, T, z_dim) or (T, z_dim) for single episode
+        # a: (B, T, action_dim) or (T, action_dim)
+        x = torch.cat([z, a], dim=-1)  # (B, T, 35) or (T, 35)
+        if x.dim() == 2:
+            x, squeeze = x.unsqueeze(0), True
+        else:
+            squeeze = False
+        output, hidden = self.lstm(x, hidden)  # (B, T, 256)
+        B, T, H = output.shape
+        pi, mu, sigma = self.mdn(output.reshape(B * T, H))
+        pi = pi.view(B, T, self.mdn.gaussians)
+        mu = mu.view(B, T, self.mdn.gaussians, self.mdn.z_dim)
+        sigma = sigma.view(B, T, self.mdn.gaussians, self.mdn.z_dim)
+        if squeeze:
+            pi, mu, sigma, output = (
+                pi.squeeze(0),
+                mu.squeeze(0),
+                sigma.squeeze(0),
+                output.squeeze(0),
+            )
+        return pi, mu, sigma, hidden, output

rnn.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a563113bfc6dc4dafbfb1b8b61dae411e0d232a9829194692bef4e4ebec70674
+size 1538453

vae.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a7a76f3e62e6d0ac258ef4a82cd9c62e183964f87fdc129b540b3370934e8038
+size 17404153