Tiny world model "dreams" games on a Raspberry Pi Pico 2 W

What it's about

A developer going by SheepSeb has published a project on GitHub called "tinyworldformer". It runs a so-called world model on a Raspberry Pi Pico 2 W. A world model is a neural network that learns how an environment works. It takes the current frame and an action, then predicts the next frame along with the reward. So no game engine computes the game. The network "dreams" it. A second, even smaller network plays this dreamed game: it looks at the generated frame and picks the next key. The Pico 2 W is a microcontroller that costs about $5. According to SlashGear, it has the RP2350 chip with a Cortex-M33 (up to 150 MHz), 520 kB of SRAM and 4 MB of flash.

The two networks

The real game never runs on the chip. The actor only sees frames generated by the world model, and the world model only receives actions from the actor. Both run on core 1 of the RP2350, while core 0 handles USB and the display. Any key press hands control over to you, so you can play the dream yourself.

Looped transformer and decoding

The author was interested in "looped transformers". A single transformer block (attention plus a ReLU MLP) is applied several times with the same weights, instead of stacking different layers. That saves memory. The frame is split into square patches (2×2 to 6×6 cells), which yields 36 to 80 tokens. The block runs 1 to 4 times, and each pass is trained against the target frame. At run time, this lets you trade speed for accuracy. The Pico uses two passes.

According to the author, turning the model's outputs into an image was almost as important as the model itself. There are three rules. Background cells each take their most likely token individually. Single objects (ball, player, ghost) are placed exactly once, using a global choice across all cell-token pairs. For random events (serve, respawn), the decoder rolls the dice among nearly equally good candidates. Pac-Man and the ghosts are 2×2 sprites where only one cell carries state and the rest is copied. The global choice raised the exact-frame rate of the Breakout model from 91.6% to 99.7%. Before that, the model had drawn two balls.

The five games and their accuracy

"Exact" means every cell of the predicted frame matches the real game. In Doom-lite, even a single wrong shading pixel counts as an error. The author says the number makes the game look worse than it plays. Extra loops help most there: Doom-lite gets 59.8% with 1 loop, 83.6% with 2 and 85.9% with 3.

An agent pushes the dream into states that were never in the data. So the author also checked the model in play. He read the game state from the dreamed frame, applied the real rules and compared the results. Over 5,000 steps per game, Flappy Bird had no errors, and Doom-lite and Pong had one wrong state each. Pac-Man, however, had 742 wrong states and 38 wrong rewards.

Training and agents

Agents first play the real game, and 1–3 million steps per game are recorded. The world model is trained for 30k steps. The first 3–4k run in float, then training switches to int8 rounding (quantization-aware training). That way, the int8 model is exactly the model that was trained. The weights are exported as Rust source code, and tests check for bit-exact agreement with the Python reference.

Agents exploited weaknesses in the dream. One Breakout agent scored 154 points with 26 bricks still standing, because bricks in the corner didn't disappear. Agents trained only in the dream reached 210 points there (the real maximum is 56) but only 28 in the real game. So the author repeatedly collected fresh data with the agent and retrained. In Doom-lite, noise during training kept the enemy from vanishing from the frame.

The agents use Stream AC(λ). It learns from one transition at a time, with no replay buffer, batches or target network. The board starts with a freshly learning agent, and a key press switches to the pretrained one. In the real game, an agent reaches the following scores after 500k steps: 45.6 points in Breakout (max 56), 332 in Flappy Bird (max) and 217.4 in Pac-Man (max 249).

Quantization and memory

Matrices are int8 with one scale per output row, and activations are int8 too. The RP2350 instruction SMLAD performs two multiply-adds at once. Norms, residuals and embeddings stay in float32, because the chip has an FPU. The firmware is no_std Rust on embassy with no dynamic memory allocation. Breakout needs 361 KB of flash and 299 KB of RAM, and Pac-Man about 600 KB and 460 KB. The author notes that the float32 embeddings take up most of that and could be int8 as well.

What you need to try it

• A Raspberry Pi Pico 2 W (about $5), optionally a Display Pack (320×240 ST7789) and a USB cable. Without a display, the output runs as colored ANSI text over USB serial.
• A Rust toolchain. Flash with ./flash.sh --env doomlite (hold BOOTSEL while plugging in).
• No hardware? Run ./flash.sh --env pong --host in the terminal, or use the WebAssembly version in the browser (local web server).
• Retraining needs Python with uv and PyTorch. The models are small, and the source doesn't require a GPU.

Our take

The project shows nicely that world models in the style of GameNGen or World Models (Ha & Schmidhuber) can shrink to tiny scales if you pick games small enough. The honestly documented failure modes (the agent's exploits in the dream) are interesting. On the critical side, the author hasn't yet measured the frame rate on the Pico. The games are heavily simplified, and Pac-Man, with 742 faulty states, shows clear weaknesses. The other sources (console builds, Zotac CM5 board) don't relate to this story. SlashGear only supplies the Pico 2's chip specs.