LeWorldModel 1 entraîne un modèle de monde à sélectionner un starter Pokémon

Contexte et objectif

L’auteur a choisi Pokémon Red (1996) comme environnement test pour un modèle de monde basé sur l’architecture JEPA de Yann LeCun. Le jeu comporte 1 026 créatures, mais le scénario réduit à la sélection d’un starter dans le laboratoire du Professeur Oak, qui ne nécessite que douze pressions du bouton A. L’objectif initial était de faire planifier au modèle la séquence d’actions menant à ce résultat, sans recourir à une fonction de récompense explicite.

Architecture du modèle et apprentissage

LeWorldModel 1, implémenté sur une RTX 3080 Ti, utilise un encodeur d’images qui transforme chaque capture d’écran en un vecteur d’embedding de 192 dimensions. Le processus d’apprentissage se déroule entièrement dans cet espace latent : l’observation actuelle o_t est encodée en z_t, l’action a_t (ex. A) est concaténée, puis un prédicteur estime l’embedding futur \hat z_{t+1}. La perte principale est l’erreur quadratique moyenne entre \hat z_{t+1} et l’embedding réel z_{t+1} obtenu en ré‑encodant la capture suivante.

loss = MSE(\hat z_{t+1}, z_{t+1})

Pour éviter que l’espace d’embedding se contracte (phénomène de « collapse »), l’auteur applique le régulateur SIGReg, qui contraint la variance des dimensions d’embedding à rester élevée, garantissant ainsi une représentation discriminante.

Problèmes rencontrés et solutions

La première tentative a échoué : le modèle a appris à errer dans le laboratoire ou à annuler systématiquement avec le bouton B, ne produisant jamais la séquence de douze A. L’analyse a révélé que le plan généré était trop long pour être couvert par la fenêtre de prédiction du modèle, entraînant une perte de gradient faible et un apprentissage instable. La solution a consisté à introduire un rollout fine‑tuning : après chaque prédiction, le modèle exécute réellement l’action dans l’émulateur, récupère la nouvelle observation et ré‑entraîne le prédicteur sur cette trajectoire. Cette boucle de rétroaction a permis de stabiliser la perte et d’apprendre une dynamique plus fiable.

Évaluation et limites

Après le second entraînement, le modèle a réussi à planifier la séquence de douze A et à sélectionner le starter Squirtle. La réussite repose sur la capacité du modèle à généraliser le déplacement « gauche » ou « droite » au-delà du laboratoire, démontrant une forme d’intuition spatiale dans l’espace latent. Cependant, la portée reste limitée : le test ne couvre qu’une action triviale, et aucune mesure de performance n’est fournie pour des tâches plus complexes (batailles, exploration du monde). De plus, l’absence de récompense explicite rend difficile l’évaluation quantitative de la politique générée. Enfin, le besoin de fine‑tuning par rollout indique que le modèle seul, entraîné uniquement sur la perte d’embedding, ne suffit pas à garantir la convergence vers des plans utiles dans des environnements plus riches.