Tiny world model imagines and plays five games on a Raspberry Pi Pico 2 W

Two networks playing inside an imagination

A project named tinyworldformer, published on GitHub by developer SheepSeb and surfaced on the Hacker News front page, runs a complete learned simulation on a Raspberry Pi Pico 2 W. Two small neural networks cooperate: a world model that takes the current frame plus an action and predicts the next frame and its reward, and an actor that reads each generated frame and picks the next action. On the microcontroller neither network ever consults the real game — the actor only sees frames the world model produces, and the world model only receives actions from the actor.

Why a looped transformer

According to the project's README, the starting point was an interest in looped transformers, where one transformer block is reused several times with the same weights instead of stacking many layers. Weight reuse keeps memory usage low, which suggested a world model might fit on a microcontroller. The Pico 2 W was picked because it is inexpensive and includes a floating-point unit, and each game is only a few hundred lines of NumPy, so every predicted frame can be checked against the true rules.

How the two networks are built

The world model is a single transformer block (attention plus a ReLU MLP) applied one to four times, with the input added back before each pass. Every pass is trained to predict the next frame, so the loop count lets the developer spend more compute for better predictions at run time; the firmware uses two loops to keep a decent frame rate. Depending on the game, the model has 26k to 131k parameters, costs 2–8 million multiply-adds per frame, and is quantized to int8. The actor is a float32 MLP with two hidden layers and 9k–16k parameters. Both run on core 1 of the RP2350 chip, while core 0 handles USB and the display. The pretrained actor can be swapped for a fresh stream-x reinforcement-learning agent that learns from scratch on the board, and any button press hands control to a human player.

Five games, measured cell by cell

Frames are token grids that always hold the whole state: a ball's heading is encoded in its token, and the Doom clone carries a minimap. Breakout runs on a 20×16 grid (26k parameters, 99.99% of held-out transitions predicted with every cell correct), Flappy Bird on 24×24 (46k, 100%), Pong on 20×32 (46k, 99.45%), a raycast Doom-lite on 32×32 (55k, 83.6%) and a Pac-Man lite with two ghosts on 48×48 (131k, 86.4%). The README cautions that Doom-lite's number understates quality, since its first-person view has 768 cells and one wrongly shaded pixel marks the entire frame wrong. Extra loops help modestly — three loops lift Doom-lite to 85.9% — before returns flatten.

Decoding mattered as much as learning

Turning per-cell scores into a frame required three hand-written rules. Background cells take their most likely token independently. Singletons — the ball, the player, each ghost — are placed by one global choice over all cell-token pairs, guaranteeing exactly one of each; that rule alone lifted Breakout from 91.6% to 99.7% exact frames after an early model drew two balls. After random events such as serves and respawns, the choice is sampled among near-best candidates so the model invents variety instead of an average. Pac-Man sprites predict state in one designated cell and stamp the other three from it.

Checking the imagination

Offline accuracy was not enough, the author writes, because agents drive the simulation into states the training data never covered and exploit any mistake. The main check therefore runs while an agent plays: the game state is read back from each imagined frame, the real rules are applied, and the result is compared with the next imagined frame. Over 5,000 dreamed steps per game at two loops, Flappy Bird and Pong were essentially flawless, Doom-lite produced one wrong state, and Pac-Man lite accumulated 742 wrong states and 38 wrong rewards — the largest model is the least reliable.

Training and deployment

Training data came from agents playing the real games: one to three million recorded transitions per game, some starting from scrambled states. The world model trains for 30k steps with per-cell cross-entropy, switching on int8 rounding after the first few thousand steps as quantization-aware training. Final weights are written into Rust source as exact bit patterns and verified bit-for-bit with tests before flashing. Because agents trained inside the simulation reach novel states, their games are recorded and folded back in to fine-tune the model. The README also documents a subtle failure: weighting the loss by whether a cell changed biased the model toward predicting change, producing duplicate balls and bricks, fixed by weighting where change is possible plus the singleton rule.

Why it matters

This is a compact demonstration of how small learned dynamics models have become: tens of thousands of parameters capture game physics well enough that a reinforcement-learning agent can train inside the simulation, offline, on a hobbyist microcontroller. It also exposes the two practical problems that come with that. Learned simulators need careful decoding and closed-loop verification, because agents will find and exploit any hallucination. And the recipe here — record, train, redeploy agents inside the imagination, record again — mirrors the feedback loop behind much larger world-model research, at a scale anyone can reproduce on a desk.