Tinyworldformer: Running Looped Transformers on Raspberry Pi Pico 2 W

Overview

Tinyworldformer demonstrates that complex world models and reinforcement learning agents can run on resource-constrained hardware by using looped transformers and int8 quantization. The system enables a Raspberry Pi Pico 2 W to generate game frames and learn to play them in real-time, effectively allowing the microcontroller to "dream" and act within its own simulated environments.

Architecture: Looped Transformers and Actors

The system relies on two distinct neural networks running on the RP2350 chip:

  • The World Model: A looped transformer that predicts the next frame and reward based on the current frame and action. By reusing a single transformer block multiple times with the same weights, the model achieves significant memory savings, making it viable for microcontrollers. It uses 26k–131k parameters depending on the game complexity.
  • The Actor: A Multi-Layer Perceptron (MLP) that observes the generated frames and selects actions. It is lightweight, utilizing only 9k–16k parameters.

Both networks operate on core 1 of the RP2350, while core 0 manages USB communication and display output. The system avoids using the "real" game logic during inference, relying entirely on the world model's predictions.

Game Implementation and Performance

The project supports five games: Breakout, Doom-lite, Flappy Bird, Pong, and Pac-Man. Each game is implemented in a few hundred lines of NumPy code, rendering to a small grid of token IDs. The world model's accuracy varies by game, with high fidelity achieved in simpler titles like Flappy Bird (100% exact frame prediction) and complex ones like Pac-Man (86.4% at 2 loops).

To ensure the model remains stable, the project employs specific decoding rules, such as placing "singleton" objects (like the ball or player) to prevent the model from hallucinating multiple instances of the same object. The world model is validated by comparing its predictions against the actual game rules during agent training.

Quantization and Deployment

To fit within the Pico 2 W's 512 KB of SRAM and 4 MB of flash, the project uses Quantization-Aware Training (QAT). Every matrix in the world model is converted to int8 format. The firmware is written in no_std Rust, ensuring deterministic performance without dynamic memory allocation.

Performance metrics for the largest game, Pac-Man, show that the system consumes approximately 600 KB of flash and 460 KB of SRAM, fitting comfortably within the hardware limits. The use of the RP2350's FPU allows for efficient handling of norms and residuals, while integer-based dot products are optimized using specific assembly instructions (SMLAD).

Insights on AI on Microcontrollers

Community discussion highlights the potential for this approach in future AI applications. Commenters noted that such projects—focused on minimal, locally run models—represent a significant direction for the future of embedded AI.

"I honestly believe that these neat projects (smallest possible, locally run LLMs) are the real future of AI on the long run." — skrellm

Others pointed out that while the hardware constraints are impressive, the underlying logic could be integrated into traditional game development. By using these models for procedural generation or adaptive NPC behavior, developers could potentially solve common limitations found in standard procedural generation techniques.

Sources

Related

  • Project
    SheepSeb/tinyworldformer

    A framework for running tiny, looped transformer world models and reinforcement learning agents on a Raspberry Pi Pico 2 W, allowing the device to simulate and learn to play games entirely on-chip.

  • Project
    TencentARC/WorldCrafter

    WorldCrafter is a research‑grade video generation model that creates camera‑controllable, consistent videos from a single image or a text prompt. It stores an implicit 3‑D‑aware memory so the scene stays coherent across viewpoints. The repo provides a full‑size Base model and a distilled Fast model, scripts for image‑to‑video and text‑to‑video inference, automatic prompt generation via Qwen‑VL, and an open‑source interactive web demo. Setup uses Python 3.11 with PyTorch 2.10 (CUDA 12.8) via uv or conda, and model weights are hosted on Hugging Face.

  • Dispatch
    WorldTUI: How a 360M Model Was Trained to Pretend to Be a Terminal

    WorldTUI is a 360M-parameter model fine-tuned from Qwen2.5-Coder-0.5B to simulate a terminal by predicting screen updates from keypresses, with no shell or filesystem behind it, running locally in a browser tab at 208 MB.

  • Dispatch
    Hugging Face Transformers timm Integration

    Hugging Face has introduced the TimmWrapper, allowing any model from the PyTorch Image Models (timm) library to be used seamlessly within the transformers ecosystem for inference, quantization, and fine-tuning.

  • Project
    nicklashansen/dreamer4

    A PyTorch implementation of the Dreamer 4 world model designed for multi-task continuous control, featuring a causal tokenizer and an interactive dynamics model.