Running Looped Transformer World Models on Raspberry Pi Pico 2 W
Deploying lightweight neural networks for real-time game simulation and agent learning on microcontrollers.
1Microcontroller-Based World Models
The project implements a looped transformer architecture on the RP2350-powered Raspberry Pi Pico 2 W. By reusing a single transformer block across multiple passes, the model minimizes memory footprint while maintaining predictive accuracy. This approach enables complex game state simulation directly on low-power hardware without requiring external cloud processing or large memory buffers.
2Dual-Network Architecture
The system utilizes two distinct neural networks: a world model that predicts future frames and rewards, and an actor network that selects actions. Both networks operate exclusively on the generated dream state, allowing the agent to learn and play within its own internal simulation. This decoupling ensures the agent remains functional even without real-time game input.
3Weight Reusing Transformers
Unlike standard deep stacks, this model applies the same transformer weights multiple times. This loop-based design trades computational speed for significant memory savings, fitting the entire model into the Pico's 520 KB SRAM. The loop count is adjustable, allowing users to balance frame rate against prediction precision based on specific game requirements and performance targets.
4Quantization and Performance
The model employs int8 quantization for weights and activations to maximize throughput on the RP2350. Using SMLAD instructions, the system performs two multiply-add operations per cycle, significantly accelerating dot product calculations. Float32 precision is reserved for normalization and embedding layers, leveraging the chip's built-in FPU to maintain stability during complex inference tasks.
5Stream-X Agent Integration
Agents utilize the Stream AC(λ) algorithm, which learns from individual transitions without relying on replay buffers or batch processing. Eligibility traces allow the agent to credit past actions effectively, while ObGD prevents overshooting target values. This lightweight learning mechanism enables the agent to adapt and improve its strategy directly on the microcontroller during runtime.
6Game State Representation
Games are rendered into small grids of token IDs, where each frame encapsulates the entire state, including hidden variables like ball heading or minimap data. The input is partitioned into square patches, with embeddings summing learned vectors for each cell position. This compact representation ensures that even complex games fit within the 4 MB flash.
7Verification and Training
Models undergo Quantization-Aware Training (QAT) using a straight-through estimator to ensure the int8 model matches the performance of its float32 counterpart. Rust-based golden tests verify that the firmware implementation reproduces Python-based integer specifications bit-for-bit. This rigorous validation process guarantees that the deployed model behaves exactly as intended within the constrained embedded environment.
8Core Utilization Strategy
The RP2350's dual-core architecture is leveraged to ensure smooth operation. Core 1 handles the game logic, world model inference, and agent decision-making, while Core 0 manages USB communication and display output. This separation prevents I/O bottlenecks from interrupting the real-time simulation, maintaining consistent frame timing even during intensive learning cycles on the board.
9Handling Complex Dynamics
To prevent agents from exploiting simulation errors, the system employs idle penalties for repetitive movement patterns. By reading the game state directly from the dreamed frame, the learner can compare predicted outcomes against real rules. This feedback loop forces the agent to patrol effectively, significantly improving performance in environments like the Doom-lite 3D view.
10Firmware and Memory Constraints
The firmware is written in no_std Rust, avoiding dynamic memory allocation to ensure reliability. Weights are stored as static arrays in flash memory, while all buffers are pre-allocated at compile time. This design allows the entire system, including the learner and display driver, to operate within the strict 512 KB SRAM limit of the Pico.
11WebAssembly Portability
The same no_std Rust engine compiles to WebAssembly, enabling the world model to run in a browser. This allows users to reproduce the exact episodes generated on the hardware, frame for frame, at high speeds. The WASM module includes the full learning loop, providing a consistent development environment across both native hardware and web platforms.
12Future of Embedded AI
This project demonstrates that sophisticated world models and reinforcement learning agents can thrive on commodity microcontrollers. By combining looped architectures with efficient quantization and static memory management, developers can deploy autonomous agents in resource-constrained settings. This methodology provides a blueprint for future edge-AI applications requiring real-time simulation and adaptive behavior on minimal hardware.
Original article: Hacker News — Running Looped Transformer World Models on Raspberry Pi Pico 2 W