Présentation

Le projet TinyWorldFormer combine deux réseaux de neurones très compacts pour générer et jouer à des jeux simples directement sur un microcontrôleur Raspberry Pi Pico 2 W (RP2350). Le premier réseau, appelé world model, prédit l’état suivant du jeu et la récompense à partir du cadre actuel et de l’action proposée. Le second, l'actor, reçoit le cadre synthétisé par le world model et choisit l’action suivante. Les deux réseaux s’exécutent simultanément sur le cœur 1 du RP2350, tandis que le cœur 0 gère l’USB et l’affichage.

Architecture du modèle

Le cœur du système repose sur un unique bloc de transformeur composé d’une couche d’attention suivie d’un MLP à activation ReLU. Ce bloc est réutilisé plusieurs fois (de 1 à 4 boucles) avec les mêmes poids, une technique appelée looped transformer. La réutilisation des poids réduit fortement l’empreinte mémoire, ce qui est crucial sur un dispositif disposant de seulement 264 KB de RAM. Chaque passage du bloc ajoute l’entrée originale avant le calcul, ce qui permet d’ajuster le compromis vitesse‑précision : plus de boucles augmentent la fidélité de la prédiction mais allongent le temps de calcul. Sur le Pico, deux boucles sont retenues comme compromis optimal pour maintenir un taux de rafraîchissement acceptable.

Les cadres sont découpés en patchs carrés (2×2 cellules pour Breakout, 4×4 pour Doom‑lite, Flappy Bird et Pong, 6×6 pour Pac‑Man), ce qui produit entre 36 et 80 tokens selon le jeu. Chaque token est encodé comme la somme de vecteurs appris pour chaque cellule du patch, avec une table distincte selon la position du cellule afin de préserver l’ordre. Un vecteur de position global et un vecteur représentant l’action sont ajoutés à chaque token avant l’étape d’attention.

Implémentation sur le Pico

Les deux réseaux sont quantifiés pour tenir dans la mémoire du RP2350 et sont exécutés en natif grâce à l’unité de calcul en virgule flottante (FPU) du microcontrôleur. Le code source, écrit en Python/Numpy, est compilé et flashé sur la carte à l’aide de la commande :

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

Cette ligne lance l’émulation du jeu Pac‑Man dans le terminal, démontrant que le même code peut être exécuté en mode simulation ou directement sur le matériel. Le modèle pré‑entraîné de l’actor peut être remplacé par un agent d’apprentissage en ligne (stream‑x) qui ajuste ses poids pendant que le board « rêve » du jeu, offrant ainsi une boucle d’apprentissage entièrement embarquée.

Performances et limites

Le critère d’évaluation principal est le taux de Exact frames, c’est‑à‑dire la proportion de transitions retenues où chaque cellule du cadre prédit correspond exactement à l’état réel. Les jeux contenant des éléments très détaillés, comme Doom‑lite (768 cellules), affichent des scores faibles car une seule cellule erronée invalide le cadre complet, même si l’état global reste correct. En pratique, le modèle atteint 92 % d’erreurs après une boucle, 35 % après deux, 12 % après trois, et 0 % après quatre passes, illustrant l’effet de la boucle d’attention.

Les contraintes principales proviennent de la capacité de calcul du RP2350 : augmenter le nombre de boucles ou la taille des tokens entraîne une chute du taux de rafraîchissement, rendant le jeu moins réactif. De plus, la quantification nécessaire pour réduire la taille du modèle introduit une perte de précision qui se manifeste surtout sur les jeux à haute résolution de grille. Malgré ces limites, le projet montre qu’il est possible de faire cohabiter un world model et un agent d’action sur un microcontrôleur à faible coût, ouvrant la voie à des expériences d’apprentissage par renforcement entièrement embarquées.