LeMario: Training a JEPA World Model on Super Mario Bros
2 hours ago
- #Reinforcement Learning
- #JEPA
- #World Models
- The author trained a Joint-Embedding Predictive Architecture (JEPA) model, named LeMario, on Super Mario Bros to learn world dynamics from pixels and actions.
- The model included a vision encoder, action encoder, and causal predictor with Adaptive LayerNorm Zero (AdaLN-Zero) for action conditioning, and used SIGReg to prevent representation collapse.
- LeMario outperformed baselines in short-horizon prediction on held-out episodes, showing it learned action-conditioned dynamics.
- However, when using reward-free planning with the Cross-Entropy Method (CEM) to reach distant image goals, Mario struggled with obstacles like jumps and navigation.
- Probing revealed the latent space encoded horizontal position well but vertical position weakly, explaining planning failures.
- Issues included latent distance not aligning with controllable progress, CEM exploiting model weaknesses, and differences from the original Push-T setup (e.g., scrolling camera, fewer episodes).
- Despite failures, the project provided insights into world models, emphasizing the importance of dataset, environment assumptions, and systematic testing.