GitHub - SheepSeb/tinyworldformer

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.

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.

The games 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).

How the world model works Frames split into tokens 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. Looped Transformer 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.

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): 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 world model samples a new serve or respawn instead of averaging them. 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%. Checking the world model 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.

Some episodes start from unusual states, like bricks with holes or random positions, so the data covers more of the game. Then the world model learns to predict the next frame with a per-cell cross-entropy loss: 30k steps, with the first 3–4k steps run in float, then int8 rounding switches on (QAT). The int8 weights are written into Rust source with every float stored as its exact bits. We check them with cargo test to reproduces the Python integer spec bit for bit, and then the firmware runs on the Pico. An agent trained in the world model reaches states the recorded data never had, so record real games with that agent and fine-tune the world model on everything.

The agent in different games stated exploiting them so those are some fixes to it: 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 didn't remove bricks from the corner when it 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, as World Models (Ha & Schmidhuber) predicts. Each collect-and-fine-tune round closed the exploit that agent had found. The Doom-lite dream sometimes lost the enemy sprite and never drew it again: the model copied the previous view instead of redrawing it from the minimap. Noise augmentation fixed it (train --noise 0.3, as in GameNGen): input views get random pixels and sometimes lose the sprite. d=32 couldn't absorb the wider Breakout data and d=48 could. Pac-Man's 48×48 frame needed d=64. The agents: stream-x The agents use Stream AC(λ) (Elsayed, Vasan & Mahmood, 2024), which learns from one transition at a time: no replay buffer, no batches, no target network. Eligibility traces carry credit back in time, and ObGD shrinks any step that would overshoot the TD target.

The policy and value networks each have two hidden layers of 64 units (32 on the Pico) with LayerNorm. The input is a feature vector read off the frame (training/streamx.py: observe_*, mirrored in Rust in tinyworldformer-rs/src/agent.rs). For example, the ball's row, column and heading in Breakout, or in Pac-Man, for each direction, the shortest-path distance to the nearest pellet and how close each ghost is. Below is one training run per game (real game, 500k steps) at six stages: untrained, the end, and the snapshots where the score has covered 25, 50, 75 and 95% of the way. Each panel plays greedily from the same start and skips a few steps per frame.

"try" counts the episodes played so far, so you can see the crashes.