TinyWorldFormer uses compact neural networks and a looped transformer to generate playable games and train an AI agent directly on a low-cost microcontroller.
A developer has demonstrated how a small microcontroller can generate game environments using AI and learn to play them without depending on a desktop computer or cloud service. Called TinyWorldFormer, the project runs two neural networks on a Raspberry Pi Pico 2 W: a world model that predicts what happens next in a game and an actor network that selects actions based on the generated frames.
The world model employs the looped transformer structure, which involves using the same transformer for a number of times instead of using multiple layers that have been separately trained. The world model gets the current game state frame and an action as input data, and from there, it outputs the next state and its reward. The role of the actor is to analyze the frame output by the world model and select the next action. There are five games supported by this system. They are Breakout, Doom-lite, Flappy Bird, Pong, and Pac-Man lite.
The Raspberry Pi Pico 2 W employs the RP2350 MCU. Which includes Arm Cortex-M33 CPU cores and 520KB of SRAM memory. In order to accommodate the models within the given memory constraints, the world model employs 8-bit quantisation of integers, whereas the actor works with 32-bit floating-point numbers. The software is programmed in the no_std Rust language using Embassy. One processor core handles game simulation and AI inference, while the other manages USB communication and the optional display.
An interesting feature is that the agent can continue learning while running on the microcontroller. The project uses Stream AC(λ), a reinforcement-learning algorithm that updates the actor using individual interactions rather than a large replay buffer. Players can also control games using physical buttons while the agent learns. According to the developer, a Flappy Bird agent achieved an average score of more than 300 points after 250,000 learning steps.
The project is released as open-source software, with its Python implementation, Rust engine, training scripts and firmware available for modification. Developers can also use its model-export tools to convert trained PyTorch models into embedded Rust code. The repository includes instructions for training, testing, exporting models and flashing the Pico 2 W, making it a useful starting point for experimenting with compact AI models on microcontrollers.
For more information, click here.
