Glonce

tinyworldformer runs a looped-transformer game world model on a Raspberry Pi Pico 2 W

modelsThe Glonce editorsHacker News🔥69Glonce rating69Impact40Novelty60Relevance95Surprise60Players10Credibility70
tinyworldformer runs a looped-transformer game world model on a Raspberry Pi Pico 2 W

tinyworldformer is a GitHub project that runs two small neural networks on a Raspberry Pi Pico 2 W at the same time. The first is a world model: given the current frame and an action, it predicts the next frame and the reward. The second is an actor: it looks at the generated frame and picks the next action. On the chip, neither network ever touches the real game. The actor sees only frames the world model made up, and the world model gets its actions only from the actor. Both run on core 1 of the RP2350, while core 0 handles USB and the display.

The author's motivation, as written in the README, was an interest 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 the aim was to see whether a looped world model could be small enough for a microcontroller. The Pico 2 W was picked because it is cheap and has an FPU. The games are small custom ones, each a few hundred lines of numpy, so every generated frame can be checked against the real rules: Breakout, Pong, Flappy Bird, Pac-Man and a reduced Doom called Doom-lite. Later the author added a streaming reinforcement-learning agent so the board can also learn to play inside its own dream. The pretrained actor can be swapped for a fresh agent that learns from scratch on the board, and any button press hands control to the human.

How the model works. Each game renders to a small grid of token ids, and one frame must hold the whole game state, so hidden state is drawn into the frame (the ball's heading is part of its token; Doom-lite has a minimap). 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, giving 36 to 80 tokens per game. 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. The Pico uses two loops to keep a decent frame rate. In one held-out Doom-lite step, the model after one pass draws an enemy that is not there, and further passes leave 92, 35, 12, then 0 wrong cells.

Decoding scores into a frame turned out to matter almost as much as the model. It ended up as three rules. Background cells (walls, bricks, pellets, floor) each take their most likely token independently. 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 model samples a new serve or respawn instead of averaging them. Pac-Man and the ghosts are 2×2 sprites with only one state-carrying cell each, and the other three cells are stamped from it. The second rule was added because the first Breakout model drew two balls at once; the one global choice took exact frames from 91.6% to 99.7%. "Exact frames" means the share of held-out transitions where every cell of the predicted frame matches the real game, leaving out frames right after a random event. Doom-lite's 3D view has 768 cells, so one wrong shaded wall pixel fails the whole frame; the author says its number looks worse than it plays and the game state is almost always right.

Training and the dream check. Agents first play the real game and every step is recorded as an example (frame, action, next frame, reward), 1 to 3 million per game, with some episodes starting from unusual states. The world model trains for 30k steps with a per-cell cross-entropy loss; the first 3 to 4k steps run in float, then int8 rounding (quantisation-aware training) switches on. Accuracy on recorded data was not enough, because an agent drives the dream into states the data never had. So the main check runs while an agent plays: the game state is read off the dreamed frame, the real rules are applied to it and the agent's action, and the result is compared with the next dreamed frame and reward, over 5,000 dreamed steps per game at two loops with the shipped agent.

The agents found exploits, and the README lists the fixes. Weighting the loss by "this cell changed" biased the model toward predicting change, so it drew two balls and broke two bricks; weighting by where things can happen in the input frame, plus the global singleton choice, fixed it. A Breakout agent in the dream got 154 points with 26 bricks left, because the world model did not remove bricks from the corner when they were hit. Doom-lite agents learned to snipe from 5 cells away. Agents trained only in the dream scored 210 there (the real maximum is 56) and 28 in the real game, which the author says is what World Models (Ha & Schmidhuber) predicts; each round of collecting real games with the agent and fine-tuning the world model closed the exploit that agent had found. The Doom-lite dream sometimes lost the enemy sprite for good, because the model copied the previous view instead of redrawing it from the minimap; noise augmentation (train --noise 0.3, as in GameNGen) fixed it. Wider models were needed for some games: d=32 could not absorb the wider Breakout data and d=48 could, and Pac-Man's 48×48 frame needed d=64.

The agents use Stream AC(λ) (Elsayed, Vasan & Mahmood, 2024), which learns from one transition at a time with no replay buffer, no batches and no target network. The policy and value networks each have two hidden layers of 64 units (32 on the Pico) with LayerNorm. Without shaping, Doom-lite agents camp in the start corner and fire; an idle penalty of 0.1 per step spent in a cell visited more than 6 times in the last 20 steps makes the agent patrol and get twice the kills. A hand-written no_std Rust learner matches 200 PyTorch updates to within 3.5e-6, and on a host machine a fresh agent in the Rust int8 Flappy Bird dream averages 6 pipes per episode after 25k steps, 121 after 100k and 300+ after 250k.

On the chip. Every world-model matrix is int8 with one scale per output row; activations are int8 with one scale per tensor and per loop index; matmuls sum int8 × int8 into int32; softmax uses a 1024-entry exp table; norms, residuals and embeddings stay float32. Training applies the same rounding, so the int8 model is the model that was trained, and the Rust engine reproduces the numpy integer reference bit for bit, checked by cargo test with golden vectors. The float32 embedding tables are the largest part of the footprint and could be int8 too. The firmware is no_std Rust on embassy, with nothing allocated at run time: weights are static arrays in flash and buffers are fixed-size statics in RAM. Each pass of core 1's loop is one frame: the actor picks a button, the int8 transformer runs two loops and decoding builds the next frame, and if on-device learning is on the learner updates from that single transition and forgets it. Frames go to core 0 through a two-slot channel; if core 0 is still busy, the frame is dropped so core 1 never waits for the screen. Over USB serial each frame goes out as coloured ANSI text with a status line, and keys come back in. An optional Display Pack is a 320×240 ST7789 display over SPI at up to 62.5 MHz. When a human plays, core 1 slows to 66 to 150 ms per frame depending on the game; the agent runs as fast as the chip allows.

Key facts

  • Two small networks run on a Raspberry Pi Pico 2 W: a looped-transformer world model that generates the next frame and reward, and an actor that plays inside the generated game; neither touches the real game on the chip.
  • The world model is one transformer block with one set of weights applied 1 to 4 times; the Pico uses two loops. In one Doom-lite example, wrong cells fall from 92 to 35 to 12 to 0 across four loops.
  • A decoding rule that places each singleton (such as the ball) with one global choice raised Breakout exact frames from 91.6% to 99.7%.
  • Agents exploited the dream: one Breakout agent scored 154 points with 26 bricks left, and dream-only agents scored 210 in the dream (real maximum 56) but 28 in the real game; repeated fine-tuning rounds closed the exploits.
  • Weights are int8 and the no_std Rust engine matches the numpy integer reference bit for bit; a human can take over with any button press, at 66 to 150 ms per frame.

Why it matters

A world model that generates a game's next frame and reward is usually a heavy piece of software. This project shows a version small enough to run alongside its own player on a cheap microcontroller. The README says the idea came from looped transformers, which reuse one block's weights and so save memory. The write-up is also candid about what goes wrong when an agent trains inside a learned model: it finds and exploits the model's mistakes, which the README ties to the prediction in World Models (Ha & Schmidhuber).

Who it affects

The audience is hobbyists and researchers interested in world models, looped transformers, quantised inference and embedded Rust. The README also says the rustexport package imports only numpy and torch, so other projects can reuse the int8 export path: build a model from rustexport.qat, train it with activation quantisation after a float warm-up, convert it, and write the Rust source.

How to use it

The real games can be played in a terminal with uv run python -m tinyworldformer.games.play --game pacman (or breakout, doomlite, flappy, pong). On the board, the dream is streamed over USB serial as coloured ANSI text with a status line, and key presses come back in. The pretrained actor can be swapped for a fresh Stream AC(λ) agent that learns from scratch on the board, and any button press gives control to the human. An optional Display Pack (320×240 ST7789) is supported through a display feature flag in the firmware. The README's Pico 2 W hardware note: 4 MB of flash and 520 KB of SRAM, 512 KiB of it in the main memory map.

How solid is it

This is a single author's README, so the numbers are self-reported. Some checks are strong by construction: the Rust engine is tested bit for bit against the numpy integer reference with golden vectors, and the Rust learner matches 200 PyTorch updates to within 3.5e-6. The accuracy metric is held-out exact frames, and the author adds an agent-driven check over 5,000 dreamed steps per game. The source gives no per-game exact-frame accuracy table, parameter counts or measured flash and RAM figures for the firmware in the text. It also gives no frame rate for the agent-driven run on the Pico.

Risks and caveats

The games are small custom numpy games (Doom-lite, not Doom), not the commercial titles. Exact-frame accuracy for Doom-lite understates quality by the author's own account, because one wrong pixel fails a frame, so that number is hard to read at face value. Frames right after a random event are excluded from the metric. The 92, 35, 12, 0 sequence is a single held-out example, and the 210 / 56 / 28 scores come from earlier dream-only agents, not from a final Pico run. Agents in the dream keep finding model errors, so each fix closes only the exploit found so far. The README's numbers for fresh-agent Flappy Bird learning (6, 121, 300+ pipes) are from a host machine, not the board.