Quantized world model runs five full games inside a Raspberry Pi Pico 2 W

A project published on GitHub by a developer identified as SheepSeb shows just how far an AI model can shrink without losing practical usefulness. tinyworldformer generates and plays five classic games, Breakout, Pong, Flappy Bird, a simplified version of Doom (called Doom-lite), and another of Pac-Man, entirely inside a Raspberry Pi Pico 2 W. No internet connection, no GPU, no remote server.

The Pico 2 W is Raspberry Pi Foundation's low-cost microcontroller based on the RP2350 chip: it has a built-in FPU, 4 MB of flash, and 520 KB of SRAM (512 KiB in the main memory map). Running a generic language model on it would be unfeasible, but the project doesn't attempt that. It trains a "world model" small enough to predict, frame by frame, how a specific game behaves.

Two networks, each with its own job

The system splits the task between two separate networks, which never see the actual game:

• The world model: a looped transformer quantized to int8, with 26,000 to 131,000 parameters depending on the game. It takes the current frame plus the chosen action and predicts the next frame and the reward, at a cost of 2 to 8 million multiply-adds per frame.
• The actor: a two-hidden-layer MLP, in float32, with 9,000 to 16,000 parameters. It only looks at the frame the world model "dreamed up" and decides which button to press.

Both networks run on core 1 of the RP2350, which is dual-core, while core 0 handles USB and the display. The pre-trained actor can be swapped out for an agent that learns from scratch directly on the board.

Why a looped transformer

The architecture choice comes from the author's specific interest in reusing weights across passes, instead of stacking different layers, which drastically cuts memory usage.

I got interested in looped transformers, where one transformer block is applied several times with the same weights instead of stacking different layers.

SheepSeb, author of the tinyworldformer project

In practice, the same attention-plus-MLP block runs 1 to 4 times over the same frame, adding the input back in before each pass. On the Pico, to keep a playable frame rate, the project uses two passes. Each extra pass removes errors from the generated frame: in one Doom-lite example, the model draws an enemy wrong on the first pass, and each following pass erases part of it, until the error reaches zero on the fourth.

Frame decoding follows three rules: background cells (walls, blocks, dots) independently receive the most likely token; unique elements (the ball, the player, each ghost) are chosen only once among all possible cell-and-token pairs, guaranteeing that there are never two balls at the same time; and larger sprites, like Pac-Man and the ghosts, have only one cell with real state, with the rest stamped from it.

How accurate is the generated game

The repository measures the fraction of frames, among test data never seen during training, in which all predicted cells exactly match the real game:

The Doom-lite number looks worse than it is: its first-person view has 768 cells, and a single wrong wall pixel already drops the entire frame from the count, even if the game state remains correct. Beyond accuracy on recorded data, the project runs an agent inside its own dream and compares the state read back from the generated frame with what the real rules would produce: across 5,000 simulated steps per game, most errors are concentrated in a few cells per frame, not in broken game states.

Fitting into kilobytes of flash

Every weight in the world model is int8, with one scale per output row; the activations are also int8, with one scale per tensor and per loop index. The multiplications accumulate int8 by int8 into int32 using the RP2350's SMLAD instruction, two per cycle. The int8 weights alone range from about 20 KB, in Breakout, to a little over 70 KB in the larger model, Pac-Man lite, according to the repository's memory table.

Adding up firmware, display drivers, execution buffers, and the agent, the full package fits comfortably on the hardware: Breakout uses 361 KB of flash and 299 KB of SRAM; Pac-Man, the heaviest game, uses 600 KB of flash and 460 KB of SRAM, within the Pico 2 W's limit of 4 MB of flash and 512 KB of SRAM. The int8 weights themselves are written as static arrays in Rust, with each float converted to its exact bits, and a test (cargo test) ensures that the Rust calculation reproduces the Python code bit for bit.

Learning inside its own dream

The default agent uses Stream AC(λ), a reinforcement learning algorithm by Elsayed, Vasan, and Mahmood (2024) that updates on every transition, with no replay buffer, no batches, and no target network. It's light enough to run on-device: pressing l over USB, or holding the B button for 1 second on the optional Display Pack, toggles between the pre-trained actor and an agent learning from scratch directly on the board.

The repository reinforces a well-known point from the world models literature (Ha & Schmidhuber's work is cited directly): an agent trained only inside the dream learns to exploit flaws in the generative model and racks up scores well above the game's real ceiling, while performance collapses in the actual game. Each round of collecting real data and readjusting the world model closes one of those gaps.

What this changes for those coding outside the cloud

For those developing embedded systems, IoT, or just wanting to test inference without relying on a rented GPU, tinyworldformer is a complete, reproducible example of an end-to-end pipeline: training with quantization-aware training, exporting to no_std Rust on top of the Embassy framework, and firmware that runs on a chip costing a few dollars. The same int8 engine, with no external imports and no wasm-bindgen, also compiles to WebAssembly and runs on a static page in the browser.

Before flashing anything, you can run the real versions of the games directly in the terminal:

uv run python -m tinyworldformer.games.play --game pacman

And flashing a specific model and agent onto the board is a single command:

./flash.sh --env doomlite --agent doom_agent_roam_p0.3

One data point remains open, which the author himself admits not having measured yet: the real frame rate on the physical hardware. On desktop, the same Rust code generates a frame in 0.7 ms (Flappy Bird) to 1.5 ms (Breakout); on the Pico, when a human plays, the speed is deliberately reduced to something playable, between 66 and 150 ms per frame, but the performance of the agent alone, at the chip's limit, still has no published number.