Running Looped World Models on Raspberry Pi Pico 2 W
Deploying neural networks to microcontrollers for real-time game simulation and reinforcement learning.
1Architecting Tiny World Models
The project utilizes two small neural networks running on the RP2350 microcontroller to simulate and play games. The world model predicts the next frame and reward based on current inputs, while the actor network selects actions. By using looped transformers, the system reuses weight sets across layers to minimize memory usage on hardware.
2Hardware Constraints and Optimization
The Raspberry Pi Pico 2 W was selected for its cost-effectiveness and integrated FPU. To ensure performance, the system employs int8 quantization for weights and activations. Matrix multiplications utilize the SMLAD instruction for dual multiply-adds, while norms and residuals remain in float32 to maintain precision during the inference loop on the chip.
3Game State and Tokenization
Games are rendered into small grids of token IDs, where each frame encapsulates the entire state, including hidden variables like ball heading or minimaps. Frames are divided into square patches, ranging from 2x2 to 6x6 cells, providing 36 to 80 tokens per game. Embeddings combine learned vectors with positional and action-specific data.
4Looping for Accuracy Trade-offs
The model applies a single transformer block with shared weights one to four times per frame. Increasing the loop count improves prediction accuracy but reduces the frame rate. On the Pico, two loops are typically used to balance speed and visual fidelity, effectively removing artifacts like phantom enemies in the Doom-lite environment.
5Stream-X Reinforcement Learning
Agents use the Stream AC(λ) algorithm, which learns from individual transitions without replay buffers or target networks. Eligibility traces facilitate temporal credit assignment, while ObGD prevents overshooting TD targets. The policy and value networks are constrained to 32 hidden units on the Pico to fit within the 512 KB SRAM limit.
6Training and Quantization Strategy
Models undergo Quantization Aware Training (QAT) with a straight-through estimator to ensure the int8 model matches training performance. The process involves 30k training steps, switching to int8 rounding after initial float warm-up. Rust firmware reproduces the Python integer specifications bit-for-bit, verified by cargo tests against golden vectors for every exported environment.
7Core Management on RP2350
The RP2350's dual-core architecture is leveraged to isolate tasks. Core 1 handles the game loop, including the world model and agent inference, while Core 0 manages USB communication and display drivers. This separation ensures that frame rendering and user input polling remain responsive despite the computational load of the neural networks.
8Real-time Learning on Device
The system supports on-device learning, allowing the agent to fine-tune its policy within its own dreamed environment. Users can toggle between a pretrained actor and a fresh agent that learns from scratch. This capability demonstrates the feasibility of continuous reinforcement learning on resource-constrained hardware without external compute or large memory buffers.
9Memory and Flash Footprint
Firmware is written in no_std Rust to avoid runtime allocations, using static arrays for weights. Breakout requires 361 KB of flash and 299 KB of SRAM, while Pac-Man uses 600 KB and 460 KB. These figures include the display driver, USB stack, and the on-device learner, fitting comfortably within the 4 MB flash.
10WebAssembly Portability
The same int8 engine compiles to WebAssembly, enabling the model to run in browsers without external dependencies. This implementation shares the core logic with the Pico firmware, allowing for consistent behavior across platforms. Users can run simulations at 1-5 ms per frame, facilitating rapid testing and visualization of the learned agent behaviors.
11Debugging and Verification
The system includes robust verification tools, such as comparing dreamed frames against real game rules. By reading the game state back from the dreamed frame, the learner can identify and correct discrepancies. This feedback loop is critical for preventing agents from exploiting inaccuracies in the world model's simulated environment during training.
12Future of Embedded AI
This project highlights the potential for running complex, generative world models on low-power microcontrollers. By combining weight-sharing architectures with efficient int8 quantization and no_std Rust, developers can deploy sophisticated agents that learn and adapt in real-time. This approach opens new possibilities for autonomous, edge-based interactive systems and embedded game environments.
Original article: Hacker News: Front Page — Running Looped World Models on Raspberry Pi Pico 2 W