Un minuscule WorldFormer exécute un modèle du monde en boucle et un acteur RL sur un Raspberry Pi Pico 2 W

tinyworldformer est un projet open source de SheepSeb qui exécute deux petits réseaux de neurones sur un Raspberry Pi Pico 2 W : un modèle du monde qui génère des images de jeu et un acteur pré-entraîné qui y joue. Aucun des deux réseaux ne voit jamais le jeu réel sur la puce — l'acteur ne lit que les images produites par le modèle du monde, et le modèle du monde ne reçoit que les actions de l'acteur.

Architecture

• Modèle du monde : transformer en boucle avec quantification int8. 26k–131k paramètres, 2–8M multiplications-additions par image. L'entrée est l'image courante + l'action ; la sortie est l'image suivante + la récompense.
• Acteur : MLP avec deux couches cachées, float32. 9k–16k paramètres, 9k–16k multiplications-additions par image. L'entrée est l'image ; la sortie est l'action.
• Les deux tournent sur le cœur 1 du RP2350 ; le cœur 0 gère l'USB et l'écran.
• Un acteur pré-entraîné peut être remplacé par un agent RL stream-x vierge qui apprend de zéro sur la carte. Des pressions sur les boutons passent le contrôle à un joueur humain.

Jeux et résultats de benchmark

Chaque jeu fait quelques centaines de lignes de NumPy dans src/tinyworldformer/games/, rendant une petite grille d'ID de tokens. Une image contient tout l'état du jeu, donc l'état caché est dessiné dans l'image (la direction de la balle fait partie de son token ; Doom-lite a une mini-carte).

• Breakout — image 20×16, 26k paramètres, 4,3M MACs, 99,99 % d'images exactes.
• Doom-lite — vue raycast 32×32 d'un labyrinthe 8×8, 55k paramètres, 4,2M MACs, 83,6 % d'images exactes (un seul pixel de mur ombré compte toute l'image comme fausse).
• Flappy Bird — image 24×24, 46k paramètres, 2,0M MACs, 100 % d'images exactes.
• Pong — image 20×32, 46k paramètres, 2,3M MACs, 99,45 % d'images exactes.
• Pac-Man lite — image 48×48, labyrinthe 23×21, 249 pastilles, 131k paramètres, 8,2M MACs, 86,4 % d'images exactes.

« Images exactes » est la proportion de transitions de validation où chaque cellule de l'image prédite correspond au jeu réel (les images juste après des événements aléatoires sont exclues).

Comment fonctionne le modèle du monde

Les images sont découpées en patchs carrés (2×2 cellules dans Breakout, 4×4 dans Doom-lite, Flappy Bird et Pong, 6×6 dans Pac-Man), donnant 36 à 80 tokens par jeu. Chaque plongement de patch additionne un vecteur appris par cellule avec une table séparée par position dans le patch, plus des vecteurs de position et d'action. Le modèle est un seul bloc transformer (attention + MLP ReLU) avec un seul jeu de poids appliqué plusieurs fois — un transformer en boucle — ce qui maintient la mémoire suffisamment basse pour un microcontrôleur.

Lancer les jeux réels

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

Remplacez par breakout, doomlite, flappy ou pong. Le dépôt inclut aussi tinyworldformer-rs, de la documentation, des tests et une démo navigateur.

Pour les développeurs travaillant sur l'inférence en périphérie ou les petits modèles du monde, le point pratique à retenir est le compromis mémoire-calcul : réutiliser les poids d'un seul bloc transformer permet à un modèle de 26k paramètres de tourner en quelques millions de MACs par image tout en restant dans les limites du Pico 2 W.