| --- |
| license: other |
| tags: |
| - world-model |
| - diffusion |
| - supertuxkart |
| - webgpu |
| - onnx |
| - pytorch |
| library_name: pytorch |
| --- |
| |
| # Neural Drive — SuperTuxKart World Model |
|
|
| An action-conditioned **world model for SuperTuxKart**: it predicts the next frame from the last |
| 7 frames plus your controls, with **no game engine underneath**. Trained in PyTorch on gameplay |
| (actions extracted automatically), and shipped in two runtimes from the same weights. |
|
|
| <video controls autoplay loop muted src="https://huggingface.co/codelion/neural-drive-model/resolve/main/preview.mp4"></video> |
|
|
| **▶ Live in-browser demo (WebGPU):** https://huggingface.co/spaces/codelion/neural-drive |
|
|
| - **Resolution** 384×192 · **Action space** 6-dim `[accelerate, brake, left, right, nitro, drift]` |
| - **Codec** ConvVAE (z_ch 8, /8 spatial) · **Dynamics** LatentDiT (dim 768 / depth 12 / patch 3, **130.8M**), rectified flow + per-frame diffusion forcing |
| - **Sampling** few-step distilled (**k = 2**) |
| |
| ## Runtimes |
| |
| | runtime | where | notes | |
| |---|---|---| |
| | **PyTorch** | CUDA / CPU | training + reference (fp32) | |
| | **ONNX / WebGPU** | any browser | fp16, runs on the visitor's GPU | |
| |
| ## Files |
| |
| ### PyTorch checkpoints (fp32) |
| | file | what | |
| |---|---| |
| | `dit.pt` | dynamics, 100k steps (k=8 teacher) | |
| | `dit_distilled.pt` | few-step distilled student (k=2) | |
| | `vae.pt` | ConvVAE codec | |
| | `decoder_big.pt` | larger LPIPS decoder | |
| | `config.json` | architecture config | |
|
|
| ### Web / ONNX (fp16 — the in-browser build) |
| | file | what | |
| |---|---| |
| | `web/genframe_fp16.onnx` | single-graph frame generator: prefills the 7-frame context once, then runs the k=2 denoise steps internally (KV-cache in the graph). ~262 MB | |
| | `web/vae.onnx` | VAE decoder (fp32, to keep decode stable) | |
| | `web/seeds.bin` | 18 spawn context states, raw float32 `[18,7,8,24,48]` | |
| | `web/meta.json` | shapes + config for the web runtime | |
|
|
| ## Reference |
|
|
| **Playing SuperTuxKart in Your Head: Learning Interactive World Models from Mobile Gameplay Videos** |
| — https://x.com/latent_node/status/2033769450592485543 |
| |