LeMario : Entraîner un modèle du monde JEPA sur Super Mario Bros — Guide technique et rétrospective

Le projet LeMario de Benjamin Bai reproduit LeWorldModel — une petite architecture prédictive conjointe (JEPA) — à partir de zéro et l'entraîne sur Super Mario Bros. La présentation technique détaillée couvre l'architecture complète, la configuration d'entraînement et un post-mortem honnête de ce qui a fonctionné et de ce qui n'a pas fonctionné.

Aperçu de l'Architecture

Le modèle traite quatre images consécutives de Mario. Un encodeur visuel compresse chaque image en un vecteur latent de dimension 192 : z_t = E_θ(x_t), z_t ∈ R^192. Les actions (états des boutons pour Gauche, Droite, Haut, Bas, A, B sur cinq images d'émulateur) sont encodées séparément en un autre vecteur de dimension 192 via un encodeur d'actions.

Les latents d'image et d'action alimentent un prédicteur causal avec six blocs de transformeurs. Les actions sont injectées via Adaptive LayerNorm Zero (AdaLN-Zero), qui produit des contrôles de décalage, d'échelle et de porte par branche (attention et MLP). La porte contrôle la force des mises à jour sur l'état prédit. Les poids commencent à zéro pour que le prédicteur apprenne progressivement.

Le prédicteur produit trois latents futurs prédits : ẑ₁, ẑ₂, ẑ₃. Ceux-ci sont comparés aux latents des trois vraies images suivantes en utilisant la perte MSE : L_pred = MSE([ẑ₁,ẑ₂,ẑ₃], [z₁,z₂,z₃]). Pour éviter l'effondrement de représentation, la régularisation SIGReg est ajoutée : L = L_pred + 0.1 * L_SIGReg.

Ce qui a fonctionné

• Le modèle s'est généralisé à des épisodes non vus.
• A utilisé efficacement les informations d'action.
• A prédit des futurs à cinq pas mieux que de fortes références.
• La planification sans récompense brute pouvait déplacer Mario à 2-5 pixels d'images objectifs proches.

Ce qui a échoué

• Lorsque les objectifs étaient placés plus loin dans le niveau, Mario ne pouvait pas sauter de manière fiable le premier obstacle majeur.
• Impossible de naviguer vers une seule image objectif lointaine.
• Conclusion : prédire le jeu ≠ apprendre à progresser. Le modèle a appris la dynamique au niveau des pixels mais pas la planification à long terme ni le comportement orienté vers un but.

Leçons Clés

Bai note que la plupart des leçons semblent évidentes a posteriori. L'écart entre la prédiction à court terme et la planification à long terme est critique — et n'est pas capturé par la perte de prédiction standard. L'article détaille des expériences qui ont progressivement révélé cette limitation.