A tiny world model running games on a Raspberry Pi Pico 2 W A developer built a world model that runs entirely on a Raspberry Pi Pico 2 W, pairing a looped transformer (26k–131k int8 parameters, 2–8M multiply-adds per frame) that predicts the next frame and reward with a two-hidden-layer float32 MLP actor (9k–16k parameters) that picks actions, with both networks running on core 1 of the RP2350 while core 0 handles USB and the display. Across five games the world model reproduced held-out frames exactly in 83.6% to 100% of transitions, including 99.99% for Breakout (26k parameters, 4.3M MACs) and 100% for Flappy Bird (46k parameters, 2.0M MACs), with frames immediately after random events excluded. The board can also swap the pretrained actor for a stream-x agent that learns from scratch on the chip inside its own generated dream, and any button press hands control to the player. Using 2 small neural networks to generate and play the game at the same time on a Raspberry Pi Pico 2 W. The first one creates the game: given the current frame and the action, it predicts the next frame and the reward. The second one plays that game: it uses the generated frame and picks an action. Every frame above was sampled by the world model network, and the actions were taken by the actor network. I got interested in looped transformers, where one transformer block is applied several times with the same weights instead of stacking different layers. Reusing weights is cheap on memory, so I wanted to see if a looped world model could be small enough to run on a microcontroller. I picked the Pico 2 W because it is cheap and has an FPU, and small games because I could write the real versions directly into numpy and check every generated frame against them. Later I added a streaming RL agent so the board could also learn to play inside its own dream. Contents: the two networks the-world-model-and-the-actor · the games the-games · the world model how-the-world-model-works · the agents the-agents-stream-x · quantization quantization-and-the-pico · on the chip on-the-chip · flashing flashing-a-game-and-an-agent · browser in-the-browser · reproduce reproduce · code layout code-layout | | network 1: the world model | network 2: the actor | |---|---|---| | job | frame + action → next frame + reward | frame → action | | architecture | looped transformer, q int8 | MLP with two hidden layers, float32 | | size | 26k–131k parameters | 9k–16k parameters | | cost per frame | 2–8M multiply-adds | 9k–16k multiply-adds | | trained | on recorded real games how how-the-world-model-works | with stream-x, in the real game or inside the dream how the-agents-stream-x | The size depends on the game, some games require more parameters to learn them to be playable, we tried to get to keep the number of the parameters small enough. On the chip neither network ever uses the real game: the actor only sees frames the world model, and the world model only gets actions from the actor. Both run on core 1 of the RP2350, while core 0 handles USB and the display. The pretrained actor can be swapped for a fresh stream-x agent that learns from scratch on the board. And any button press hands control to you, so you can play the dream yourself. Each game is a few hundred lines of numpy in src/tinyworldformer/games/ . It renders to a small grid of token ids, and one frame always holds the whole game state, so hidden state has to be drawn into the frame: the ball's heading is part of its token, and Doom-lite has a minimap. | game | frame | rules | random event | world model | exact frames, 2 loops | |---|---|---|---|---|---| | Breakout | 20×16 | paddle, ball, 56 bricks, 5 lives | the serve after a miss | 26k params, 4.3M MACs | 99.99% | | Doom-lite | 32×32 | raycast first-person view of an 8×8 maze; an enemy walks shortest paths to you; you shoot | the enemy's respawn | 55k params, 4.2M MACs | 83.6% see below | | Flappy Bird | 24×24 | gravity, flaps, scrolling pipes | the next gap height | 46k params, 2.0M MACs | 100% | | Pong | 20×32 | you against a CPU paddle that reacts late, first to 5 | the serve after a point | 46k params, 2.3M MACs | 99.45% | | Pac-Man lite | 48×48 | 23×21 maze, 249 pellets, two ghosts on shortest paths, 3 lives | the ghosts' respawn | 131k params, 8.2M MACs | 86.4% | "Exact frames" is the share of held-out transitions where every cell of the predicted frame matches the real game. Frames right after a random event are left out, since those can only be sampled, not predicted. Doom-lite's 3D view has 768 cells and a single shaded wall pixel off counts the whole frame as wrong, so its number looks worse than it plays; the game state is almost always right see checking the dream checking-the-dream . You can play the real games in a terminal: uv run python -m tinyworldformer.games.play --game pacman or breakout , doomlite , flappy , pong . The frame is cut into square patches: 2×2 cells in Breakout, 4×4 in Doom-lite, Flappy Bird and Pong, 6×6 in Pac-Man. That gives 36 to 80 tokens per game. A patch's embedding is the sum of one learned vector per cell, with a separate table for each position inside the patch so the cell order isn't lost. Every token also gets a position vector and the action's vector. The model is a single transformer block attention plus a ReLU MLP with one set of weights, applied 1 to 4 times, with the input added back before every pass. The output after every pass is trained against the next frame, so at run time the loop count trades speed for accuracy. On the Pico to keep a decent frame rate we use uses two loops. The animation shows one held-out Doom-lite step. After one pass the model draws an enemy that isn't there, and each further pass removes part of it: 92, 35, 12, then 0 wrong cells. The bars on the right are the share of held-out frames each loop count gets exactly right. Here is a dreamed Doom-lite step decoded after 1, 2, 3 and 4 loops, with red boxes on cells that differ from the real rules: Exact held-out frames by loop count: | loops | Breakout | Doom-lite | Flappy Bird | Pong | Pac-Man | |---|---|---|---|---|---| | 1 | 99.68% | 59.8% | 99.99% | 98.47% | 79.6% | | 2 | 99.99% | 83.6% | 100% | 99.45% | 86.4% | | 3 | 100% | 85.9% | 100% | 99.60% | 85.7% | | 4 | 99.96% | 84.7% | 100% | 99.53% | 84.4% | The model outputs a score for every cell and every token. Turning those scores into a frame turned out to matter almost as much as the model, and it ended up as three rules shown above on Pong : 1. Background cells walls, bricks, pellets, floor each take their most likely token independently. 2. Singletons the ball, the player, each ghost are placed with one choice over all cell, token pairs, so there is always exactly one of each. After a random event the choice is sampled among candidates within 2.2 logits of the best, so the world model samples a new serve or respawn instead of averaging them. 3. Pac-Man and the ghosts are 2×2 sprites, but only one cell of each carries state the mouth holds the heading; the ghost's lower-left cell holds whether a pellet is underneath . The other three cells are stamped from it instead of predicted. We implemented Rule 2 because the first Breakout model drew two balls at once; one global choice took exact frames from 91.6% to 99.7%. Accuracy on recorded data wasn't enough. An agent drives the dream into states the data never had, and any mistake there is something it can exploit. So the main check runs while an agent plays: read the game state back off the dreamed frame games/