neural-drive-model / README.md
codelion's picture
drop MLX mentions (published separately)
6bfaf08 verified
|
Raw
History Blame Contribute Delete
2.06 kB
---
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