Artificial Intelligence1 min read

A tiny world model running games on a Raspberry Pi Pico 2 W

By · Published by Everything Blog

In short

The architecture divides labor across the RP2350 processor, dedicating one core to USB and display management while the other runs the generative and decision-making models. To operate within strict hardware limits, the system tokenizes visual frames into patch grids written in NumPy that encode vital state details like headings or minimaps. Users can let an autonomous model play, take manual control of the generated environment via a button press, or train a new reinforcement learning agent directly on the hardware.

Key points

  • Dual neural networks can simulate game mechanics and play them simultaneously without usi…: Dual neural networks can simulate game mechanics and play them simultaneously without using the original underlying game code.
  • Hardware efficiency is achieved by executing between 2 million and 8.2 million multiply-a…: Hardware efficiency is achieved by executing between 2 million and 8.2 million multiply-accumulate operations per frame using looped architectures.
  • The platform supports dynamic interactions, allowing human players to interrupt the simul…: The platform supports dynamic interactions, allowing human players to interrupt the simulation or training fresh reinforcement learning agents directly on the device.

Sources checked by our analysis pipeline: 1

Using 2 small neural networks to generate and play the game at the same time on a Raspberry Pi Pico 2 W. The first one creates the game: given the current frame and the action, it predicts the next frame and the reward. The second one plays that game: it uses the generated frame and picks an action.

Every frame above was sampled by the world model network, and the actions were taken by the actor network.

I got interested 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 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.

Contents: the…

Artificial Intelligence