Games with no game code_

The dream and the actor

• A closed loop. On the chip neither network sees the real game. The actor only sees frames the dream drew, and the dream only sees buttons the actor pressed.
• One core. Both run on core 1 of the RP2350, one after the other, every frame. Core 0 handles USB and the display.
• Learning on the chip. A third network can replace the pretrained actor: a fresh stream-x agent that learns from scratch while it plays inside the dream, on the board.
• Your turn. Any button press hands control from the actor to you, so you play the dream yourself.

The games

Each game is a few hundred lines of numpy (src/tinyworldformer/games/). It renders to a small grid of token ids, and one frame always holds the whole game state. Hidden state is drawn into the frame: the ball's heading is part of its token, and Doom-lite has a minimap.

"Exact frames" counts held-out transitions where every cell of the predicted frame matches the real game, leaving out frames after a random event. A random event can only be sampled, not predicted. Doom-lite's 3D view has 768 cells, so one shaded wall pixel off counts the whole frame as wrong. Its game state is almost always right; see checking the dream.

How the world model works

Frames become tokens

• The frame is cut into square patches: 2×2 cells in Breakout, 4×4 in Doom-lite, Flappy and Pong, 6×6 in Pac-Man. Each game ends up with 36 to 80 tokens.
• A patch's token is the sum of one learned vector per cell. Each position inside the patch has its own table, so the order of the cells is kept.
• Every token also gets a position vector and the action's vector.

One block, looped

• One block: a single transformer block (attention plus a ReLU MLP) with one set of weights, applied 1 to 4 times. The input is added back before every pass.
• Every pass is trained: the output after each pass is trained against the next frame, so at run time the loop count trades speed for accuracy.
• On the Pico: two loops.

The animation shows one held-out Doom-lite step. After one pass the model draws an enemy that is not there; each further pass removes part of it: 92, 35, 12, then 0 wrong cells. The bars on its right are the share of held-out frames each loop count gets exactly right.

Below, one dreamed Doom-lite step decoded after 1, 2, 3 and 4 loops. Red boxes are cells that differ from the real rules.

Exact held-out frames by loop count:

Decoding

The model outputs a score for every cell and every token. Turning the scores into a frame takes three rules, shown above on Pong:

1. Background (walls, bricks, pellets, floor): each cell takes its most likely token on its own.
2. Singletons (the ball, the player, each ghost): one choice over all (cell, token) pairs, so there is always exactly one. For a random event the choice is sampled among candidates within 2.2 logits of the best, so the dream invents a new serve or respawn instead of averaging all of them.
3. Stamps: Pac-Man and the ghosts are 2×2 sprites. 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 drawn from it instead of being predicted. This took Pac-Man from 14% to 65% exact frames with no retraining.

Rule 2 replaced per-cell decoding early on. Decoded cell by cell, the first Breakout model drew two balls at once; one global choice took exact frames from 91.6% to 99.7%.

Checking the dream

Frame accuracy on recorded data is not enough. An agent drives the dream into states the data never had, and any mistake there is something it can exploit. So the check runs while an agent plays:

1. Read the game state back off the dreamed frame (games/<game>.py: read).
2. Apply the real rules to that state and the agent's action.
3. Compare the result with the next dreamed frame, cell by cell, and compare the rewards.

5,000 dreamed steps per game, at 2 loops, with the shipped agent (training/evalwm.py --agent):

Training

1. Recording. Agents play the real game, and every step becomes one example: frame, action, next frame, reward. That gives 1–3 million examples per game. Some episodes start from unusual states (bricks with holes, random positions) so the data covers more of the game.

2. Training. The world model learns to predict the next frame: a cross-entropy loss per cell, 30k steps, about 45 minutes per game on an RTX 4060 Ti. The first 3–4k steps run in float, then int8 rounding switches on (QAT, see quantization). The animation shows one held-out Pong step predicted at six points of a short training run.

3. Deployment. The int8 weights are written into Rust source with every float stored as its exact bits. cargo test checks that the Rust engine reproduces the Python integer spec bit for bit, then the firmware runs the dream on the Pico.

4. Closing the loop. An agent trained in the dream reaches states the recorded data never had. Recording real data with that agent and fine-tuning the world model on everything fixes the dream there.

Each fix below cost a training round to find:

• Hedging. Weighting the loss by "this cell changed" biased the model toward predicting change, so it drew two balls and broke two bricks. The fix was to weight by where things can happen in the input frame, plus the global choice for singletons.
• Coverage. A Breakout agent in the dream got 154 points with 26 bricks left: the dream paid for corner hits without removing the brick. Doom-lite agents snipe from 5 cells away, which scripted play never did. The fix was collecting data with the learned agent.
• Exploits. 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; Breakout needed four rounds.
• Copying. 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.
• Capacity. d=32 could not absorb the wider data for Breakout; d=48 could. Pac-Man's 48×48 frame needed d=64.

The agents: stream-x

Stream AC(λ) (Elsayed, Vasan & Mahmood, 2024) learns from one transition at a time:

• No buffer: no replay buffer, no batches, no target network.
• Updates: eligibility traces carry credit back in time, and ObGD shrinks any step that would overshoot the TD target.
• Networks: a policy and a value network, each 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). Examples: the ball's row, column and heading in Breakout; in Pac-Man, for each direction, the shortest-path distance to the nearest pellet and how close each ghost is.

Below, 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 crashes show.

Breakout

Doom-lite

Flappy Bird

Pac-Man lite

Pong

Greedy return, real game, mean of 10 episodes (snapshots of the same run):

Breakout: bricks minus lost balls. Doom-lite: kills minus hits in 500 steps. Flappy Bird: pipes in 3,000 steps. Pac-Man: pellets minus lives lost. Pong: points won minus lost in 3,000 steps.

• Flappy Bird reaches the maximum by 250k steps.
• Breakout is still improving at 500k; the 1M-step agents clear the board (56.0).
• Pong's agent learns not to miss, then plateaus: beating the CPU needs angled shots off the paddle's ends, which it rarely finds.

Trained only in the dream. These agents never saw the real game. Scores are measured in the real game (greedy, 50 episodes):

Without shaping, Doom-lite agents camp in the start corner and fire: the enemy always comes to them. --idle-penalty 0.1 costs 0.1 for each step spent in a cell visited more than 6 times in the last 20 steps. The position is read off the minimap, so the penalty works inside the dream too. With it the agent patrols and gets twice the kills.

Learning on the Pico. tinyworldformer-rs/src/learner.rs is the same algorithm in no_std Rust, with backpropagation written by hand. It matches 200 PyTorch updates to within 3.5e-6. On the host, a fresh agent inside the Rust int8 Flappy dream averages 6 pipes per episode after 25k steps, 121 after 100k, and 300+ after 250k.

Quantization and the Pico

• Weights: every matrix in the world model is int8, with one scale per output row (scale = max|w| / 127).
• Activations: int8, with one scale per tensor and per loop index (the loops see different ranges).
• Matmuls: int8 × int8 sums into int32 (two multiply-adds per SMLAD instruction on the RP2350), then one float multiply by the two scales.
• Softmax: a 1024-entry exp table, with probabilities stored as uint8.
• Float32: norms, residuals and embeddings. The RP2350 has an FPU.
• QAT: training applies the same rounding (straight-through estimator), so the int8 model is the model that was trained.
• Bit-exact: rustexport/intref.py is the int8 model in numpy with every float operation in a fixed order. The Rust engine reproduces it bit for bit, and cargo test checks golden vectors for every exported env.

• Measuring: flash and RAM are measured from the firmware ELF (size -A), built with the Display Pack driver. They include the code, USB, the display driver and the on-device learner.
• The Pico 2 W: 4 MB of flash and 520 KB of SRAM (512 KiB of it in the main memory map).
• Room to shrink: the float32 embedding tables are the largest part. They could be int8 too.

Reusing the export. src/tinyworldformer/rustexport/ imports only numpy and torch, and a test keeps it that way, so other projects can use it:

1. Build a model from rustexport.qat (QLinear, LoopedBlock).
2. Train it with ActQuant.enabled = True after a float warm-up.
3. rustexport.convert.looped_model turns it into the integer spec.
4. int_model_arrays and mlp_arrays with codegen.write_module write the Rust source; rustexport.Golden writes the vectors a Rust test checks.

On the chip

The firmware (tinyworldformer-rs/firmware) is no_std Rust on embassy. Nothing is allocated at run time: the weights are static arrays in flash, and every buffer is a fixed-size static in RAM.

Core 1 runs the game. One pass of its loop (dream_loop in firmware/src/main.rs) is one frame:

1. The actor picks a button. It reads features off the current frame (agent::observe) and runs its MLP. If on-device learning is on, the learning agent picks instead. If you are pressing buttons, your press is used.
2. The dream draws the next frame. The int8 transformer runs 2 loops over the frame and the button (Engine::forward). Decoding turns its scores into the next frame: background argmax per cell, each singleton placed once (sampled for random events, with an xorshift RNG), then the sprite stamps. The reward is the argmax of the reward head, and lives and episode ends are counters on top.
3. The learner updates (if learning is on), from this one transition, then forgets it.
4. The frame goes to core 0 through a two-slot channel. If core 0 is still busy with the previous frame, this one is dropped: core 1 never waits for the screen.

When you play, core 1 also slows to a playable speed (66–150 ms per frame, depending on the game). The agent runs as fast as the chip allows.

Core 0 talks to the outside world (embassy async tasks):

• USB serial: each frame goes out as coloured ANSI text with a status line (model time, actor time, fps, env and agent names). Keys come back in.
• Display Pack (optional, --features display): a 320×240 ST7789 over SPI at up to 62.5 MHz. Only cells that changed since the last frame are redrawn.
• Buttons: polled every 20 ms.

The cores share only the frame channel and a few atomics: who is playing, the pressed key, learning on/off, the loop count.

Where the memory goes:

The memory map gives 4 MB of flash and 512 KiB of RAM. Every game fits: Breakout uses 361 KB of flash and 299 KB of RAM; Pac-Man, the largest, 600 KB and 460 KB.

The inner loop. Most of the time goes into int8 dot products (dot_i8 in src/lib.rs):

• Loads: one 32-bit load fetches 4 weights.
• Unpacking: SXTB16 unpacks bytes 0/2 and 1/3 into 16-bit halves.
• Multiply-adds: SMLAD does two multiply-adds in one instruction.
• Attention: its uint8 probabilities times int8 values use UXTB16 instead of SXTB16.
• Float work: norms, the exp table lookup and the final scale per output stay in float32 on the FPU.
• Speed: a frame at 2 loops is 2–4M multiply-adds. On a desktop CPU the same Rust code dreams a frame in 0.7 ms (Flappy Bird) to 1.5 ms (Breakout).

Learning at boot. The board starts with a fresh agent learning on the device. Press l over USB, or hold B for 1 s on the Display Pack, to switch to the pretrained actor, and again to switch back.