README.md

metadata

title: Neural Boy
emoji: 🎮
colorFrom: pink
colorTo: indigo
sdk: gradio
sdk_version: 6.18.0
python_version: '3.10'
app_file: app.py
pinned: false
license: mit
short_description: Atari games dreamed frame-by-frame by a diffusion model
tags:
  - track:wood
  - achievement:offgrid
  - achievement:offbrand
  - achievement:sharing
  - achievement:fieldnotes

Neural Boy

Classic atari games with no ROM, no emulator, no game code. Every frame is generated live by a small (~4M parameter) diffusion world model conditioned on the previous frames and the button you're holding. The weights are the cartridge. Nobody programmed the rules; the model learned them by watching gameplay.

The idea came from Neural Computers (https://arxiv.org/abs/2604.06425) and I decided to try it out. Since it emulated the entire operating system, I was looking for a smaller way to use it. So I decided to try diffusion models that are trained onto atari games and use them to emulated the games without any ROMs or emulators. So we have a 4M parameter model with 64*64 frames and we use 4 frames as an episode to train the model. For each input, we denoise a frame using an episode and an action (eg move left) on that episode and then feed it back in autoregressively.

The weights of the games (the cartridges) have the frozen weight of the games which give 14 FS at 10 denoising steps.

Demo

Watch it play:

• demo video
• twitter post
• Space

Controls

• arrow keys — move.
• A / B — fire. For games without a standalone fire action (e.g. KungFuMaster) the button resolves to a directional punch.
• Restart: Restart the current gameplay.
• Drag a cartridge into the slot (or click it) to load a game.

Benchmarks

The whole "console" is one tiny diffusion world model, small enough to be the cartridge, fast enough to play.

Denoiser parameters	4.41 M (4,406,851)
Checkpoint on disk	54.3 MB
Frame resolution	64×64, 4-frame context
Peak inference VRAM	71 MB

Speed

Per-frame latency / throughput on an NVIDIA H100 at 64×64. Denoising steps trade sharpness for speed; the app defaults to 10 (~14 FPS, comfortably playable).

denoising steps	latency / frame	FPS
5	35.0 ms	28.6
10	69.8 ms	14.3
20	141.5 ms	7.1

Fidelity

Open-loop "dreaming": seed the model with 4 real frames, then let it predict forward on its own output under the real action sequence, and compare each dreamed frame to the emulator's ground truth (averaged over 8 rollouts, 10 steps). One-step prediction is near-perfect (PSNR 36 dB / SSIM 0.98); quality then degrades gracefully as autoregressive error accumulates, exactly what you expect from a world model with no access to the real game.

frames ahead	PSNR (dB)	SSIM
1	36.3	0.981
8	25.1	0.922
16	24.7	0.904
32	23.8	0.894
64	23.0	0.877

Deploying

This Space needs a dedicated GPU (it runs a continuous, autoregressive diffusion loop with one generated frame per step). ZeroGPU is not a good fit without a batched-rollout rewrite. It works on a CPU device too, albeit with high latency.

The Space is self-contained: src/ (model code) and assets/games/* (the visible cartridges) must live alongside app.py in the Space repo.

The .pt files are large, so they are tracked with Git LFS.

To run locally, do python3 app.py.

Citation

@inproceedings{alonso2024diffusionworldmodelingvisual,
      title={Diffusion for World Modeling: Visual Details Matter in Atari},
      author={Eloi Alonso and Adam Jelley and Vincent Micheli and Anssi Kanervisto and Amos Storkey and Tim Pearce and François Fleuret},
      booktitle={Thirty-eighth Conference on Neural Information Processing Systems}}
      year={2024},
      url={https://arxiv.org/abs/2405.12399},
}