TinyWorldFormer executa um modelo de mundo em loop e um ator de RL em um Raspberry Pi Pico 2 W

tinyworldformer é um projeto open-source de SheepSeb que executa duas pequenas redes neurais em um Raspberry Pi Pico 2 W: um world model que gera frames de jogo e um ator pré-treinado que joga esses frames. Nenhuma das redes vê o jogo real no chip — o ator só lê frames produzidos pelo world model, e o world model só recebe ações do ator.

Arquitetura

• World model: transformer em loop com quantização int8. 26k–131k parâmetros, 2–8M multiplicações-soma por frame. A entrada é o frame atual + ação; a saída é o próximo frame + recompensa.
• Ator: MLP com duas camadas ocultas, float32. 9k–16k parâmetros, 9k–16k multiplicações-soma por frame. A entrada é o frame; a saída é a ação.
• Ambos rodam no core 1 do RP2350; o core 0 cuida do USB e do display.
• Um ator pré-treinado pode ser trocado por um agente de RL stream-x novo que aprende do zero na própria placa. Os botões passam o controle para um jogador humano.

Jogos e resultados de benchmark

Cada jogo tem algumas centenas de linhas de NumPy em src/tinyworldformer/games/, renderizando para uma pequena grade de IDs de token. Um frame contém o estado inteiro do jogo, então o estado oculto é desenhado no frame (a direção da bola faz parte do token dela; Doom-lite tem um minimapa).

• Breakout — frame 20×16, 26k parâmetros, 4,3M MACs, 99,99% de frames exatos.
• Doom-lite — visão raycast 32×32 de um labirinto 8×8, 55k parâmetros, 4,2M MACs, 83,6% de frames exatos (um único pixel sombreado de parede já conta o frame inteiro como errado).
• Flappy Bird — frame 24×24, 46k parâmetros, 2,0M MACs, 100% de frames exatos.
• Pong — frame 20×32, 46k parâmetros, 2,3M MACs, 99,45% de frames exatos.
• Pac-Man lite — frame 48×48, labirinto 23×21, 249 pastilhas, 131k parâmetros, 8,2M MACs, 86,4% de frames exatos.

"Frames exatos" é a proporção de transições de validação em que cada célula do frame previsto corresponde ao jogo real (frames logo após eventos aleatórios são excluídos).

Como o world model funciona

Os frames são divididos em patches quadrados (2×2 células no Breakout, 4×4 no Doom-lite, Flappy Bird e Pong, 6×6 no Pac-Man), gerando 36–80 tokens por jogo. Cada embedding de patch soma um vetor aprendido por célula com uma tabela separada por posição dentro do patch, mais vetores de posição e ação. O modelo é um único bloco transformer (atenção + MLP ReLU) com um conjunto de pesos aplicado várias vezes — um transformer em loop — o que mantém a memória baixa o suficiente para um microcontrolador.

Executando os jogos reais

uv run python -m tinyworldformer.games.play --game pacman

Substitua por breakout, doomlite, flappy ou pong. O repositório também inclui tinyworldformer-rs, documentação, testes e uma demo no navegador.

Para desenvolvedores que trabalham com inferência em edge ou pequenos world models, a lição prática é o trade-off memória vs. computação: reutilizar os pesos de um único bloco transformer permite que um modelo de 26k parâmetros rode em alguns milhões de MACs por frame, permanecendo dentro dos limites do Pico 2 W.