Présentation du projet

Le projet LeMario vise à entraîner un modèle de monde JEPA (Joint-Embedding Predictive Architecture) sur le jeu Super Mario Bros. L'objectif est de reproduire le modèle LeWorldModel, qui a été utilisé pour la planification sans récompense dans le jeu Push-T.

Architecture du modèle

Le modèle est composé de plusieurs éléments : un encodeur de vision, un encodeur d'actions, un prédicteur causal et une tête de projection. L'encodeur de vision compressé chaque trame en une représentation latente de 192 nombres, tandis que l'encodeur d'actions compressé chaque séquence d'actions en une représentation latente de 192 nombres.

frames:  [batch, 4, 3, 224, 224]
actions: [batch, 4, 5, 6]  # Left, Right, Up, Down, A, B

Le prédicteur causal utilise six blocs de transformateurs pour prédire la représentation latente de la prochaine trame en fonction des trames et des actions précédentes. Les actions sont injectées dans les blocs de transformateurs via Adaptive LayerNorm Zero (AdaLN-Zero), qui les convertit en trois types de contrôles : décalage, mise à l'échelle et contrôle de la porte.

Résultats et analyse

Le modèle a été entraîné sur 737 134 trames issues de 280 épisodes répartis sur 32 niveaux de Mario. Les résultats montrent que le modèle a appris les dynamiques du jeu à court horizon conditionnelles aux actions du joueur. Cependant, lorsque le modèle a été utilisé pour contrôler Mario, les résultats ont été décevants, avec Mario à peine capable de bouger.

Une analyse plus approfondie a montré que le modèle avait oublié la position de Mario, avec une erreur moyenne absolue (MAE) de 9,30 pixels pour la position horizontale et de 21,62 pixels pour la position verticale. Cela suggère que le modèle n'a pas appris à représenter correctement la position de Mario dans son espace latent.