LeMario: Training a JEPA World Model on Super Mario Bros
AI learns Super Mario Bros and predicts the future.
1The LeMario Project
LeMario is a project that aims to learn world dynamics in Super Mario Bros using the Joint-Embedding Predictive Architecture (JEPA) model. This model is designed to predict the future from pixels and actions.
2JEPA Architecture
JEPA is an AI architecture proposed by LeCun. This architecture aims to predict future states by learning the relationship between images and actions. LeMario applied this architecture to Super Mario Bros.
3Model Training
LeMario was trained on 737,134 frames and 280 episodes. The model passed various tests to learn the game's dynamics and succeeded in short-term predictions.
4Prediction Success
LeMario outperformed baseline methods in predicting one-step and five-step futures. The importance of actions increased significantly in five-step predictions.
5Model Limitations
Although LeMario could predict the game's dynamics, it failed to learn strategies to progress Mario. It particularly struggled with reaching distant goals.
6Technical Details
The model compressed images and actions into latent vectors. These vectors were processed through transformer blocks to predict future states.
7Adaptive LayerNorm Zero
AdaLN-Zero is a technique that integrates actions into transformer blocks. This technique controls the effect of actions on latent features and improves predictions.
8Loss Function
The model used the MSE loss function to minimize prediction errors. Additionally, SIGReg loss was added to maintain the diversity of latent vectors.
9Analysis of Results
LeMario succeeded in short-term predictions but failed to develop long-term strategies. This shows that the model learned the game's dynamics but couldn't progress.
10Future Steps
Future versions of LeMario could develop new techniques to learn long-term strategies. This could enable the model to fully master the game.
11Project Significance
LeMario is an important project that demonstrates AI's potential to learn video games. Such models could be used for more complex tasks in the future.
12Conclusion
LeMario successfully learned the world dynamics in Super Mario Bros but failed to progress. This project highlights the limits and potential of AI.