codelion/neural-drive-model

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.

▶ 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