README.md

---
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](assets/demo.mp4)
- [twitter post](https://x.com/sachin_singh092/status/2066544447224156402?s=20)
- [Space](https://huggingface.co/spaces/build-small-hackathon/neural-boy)


## 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).

![Speed vs denoising steps](assets/bench/bench_speed.png)

| 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.

![Dream-vs-real fidelity vs horizon](assets/bench/bench_fidelity.png)

| 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},
}
```