TinyWorldFormer Runs a Looped World Model and RL Actor on a Raspberry Pi Pico 2 W

tinyworldformer is an open-source project by SheepSeb that runs two small neural networks on a Raspberry Pi Pico 2 W: a world model that generates game frames and a pretrained actor that plays them. Neither network ever sees the real game on the chip — the actor only reads frames produced by the world model, and the world model only receives actions from the actor.

Architecture

• World model: looped transformer with int8 quantization. 26k–131k parameters, 2–8M multiply-adds per frame. Input is current frame + action; output is next frame + reward.
• Actor: MLP with two hidden layers, float32. 9k–16k parameters, 9k–16k multiply-adds per frame. Input is frame; output is action.
• Both run on core 1 of the RP2350; core 0 handles USB and the display.
• A pretrained actor can be swapped for a fresh stream-x RL agent that learns from scratch on the board. Button presses hand control to a human player.

Games and benchmark results

Each game is a few hundred lines of NumPy in src/tinyworldformer/games/, rendering to a small grid of token IDs. One frame holds the entire game state, so hidden state is drawn into the frame (ball heading is part of its token; Doom-lite has a minimap).

• Breakout — 20×16 frame, 26k params, 4.3M MACs, 99.99% exact frames.
• Doom-lite — 32×32 raycast view of an 8×8 maze, 55k params, 4.2M MACs, 83.6% exact frames (a single shaded wall pixel counts the whole frame wrong).
• Flappy Bird — 24×24 frame, 46k params, 2.0M MACs, 100% exact frames.
• Pong — 20×32 frame, 46k params, 2.3M MACs, 99.45% exact frames.
• Pac-Man lite — 48×48 frame, 23×21 maze, 249 pellets, 131k params, 8.2M MACs, 86.4% exact frames.

"Exact frames" is the share of held-out transitions where every cell of the predicted frame matches the real game (frames right after random events are excluded).

How the world model works

Frames are split into square patches (2×2 cells in Breakout, 4×4 in Doom-lite, Flappy Bird and Pong, 6×6 in Pac-Man), giving 36–80 tokens per game. Each patch embedding sums one learned vector per cell with a separate table per position inside the patch, plus position and action vectors. The model is a single transformer block (attention + ReLU MLP) with one set of weights applied multiple times — a looped transformer — which keeps memory low enough for a microcontroller.

Running the real games

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

Substitute breakout, doomlite, flappy, or pong. The repo also includes tinyworldformer-rs, docs, tests, and a browser demo.

For developers working on edge inference or small world models, the practical takeaway is the memory-vs-compute trade: reusing one transformer block's weights lets a 26k-parameter model run in a few million MACs per frame while staying under the Pico 2 W's limits.