LeMario: Entrenando un modelo mundial JEPA en Super Mario Bros — Guía técnica y análisis posterior

El proyecto LeMario de Benjamin Bai reproduce LeWorldModel — una pequeña Arquitectura Predictiva de Incrustación Conjunta (JEPA) — desde cero y lo entrena en Super Mario Bros. El recorrido técnico detallado cubre la arquitectura completa, la configuración de entrenamiento y una autopsia franca de lo que funcionó y lo que no.

Resumen de la Arquitectura

El modelo procesa cuatro fotogramas consecutivos de Mario. Un codificador visual comprime cada fotograma en un vector latente de 192 dimensiones: z_t = E_θ(x_t), z_t ∈ R^192. Las acciones (estados de botones para Izquierda, Derecha, Arriba, Abajo, A, B a lo largo de cinco fotogramas del emulador) se codifican por separado en otro vector de 192 dimensiones mediante un codificador de acciones.

Los latentes de fotogramas y acciones alimentan un predictor causal con seis bloques transformer. Las acciones se inyectan mediante Adaptive LayerNorm Zero (AdaLN-Zero), que produce controles de desplazamiento, escala y compuerta por rama (atención y MLP). La compuerta controla cuán fuertemente las actualizaciones afectan el estado predicho. Los pesos comienzan en cero para que el predictor aprenda gradualmente.

El predictor genera tres latentes futuros predichos: ẑ₁, ẑ₂, ẑ₃. Estos se comparan con los latentes de los tres fotogramas reales siguientes utilizando pérdida MSE: L_pred = MSE([ẑ₁,ẑ₂,ẑ₃], [z₁,z₂,z₃]). Para evitar el colapso de representación, se añade regularización SIGReg: L = L_pred + 0.1 * L_SIGReg.

Lo que Funcionó

• El modelo generalizó a episodios no vistos.
• Utilizó la información de acción de manera efectiva.
• Predijo futuros de cinco pasos mejor que líneas base sólidas.
• La planificación cruda sin recompensa podía mover a Mario dentro de 2-5 píxeles de objetivos de imagen cercanos.

Lo que Falló

• Cuando los objetivos se colocaban más lejos en el nivel, Mario no podía saltar de manera confiable el primer obstáculo importante.
• No podía navegar hacia una sola imagen de objetivo distante.
• Conclusión: predecir el juego ≠ aprender a progresar. El modelo aprendió dinámicas a nivel de píxel, pero no planificación a largo plazo ni comportamiento orientado a objetivos.

Lecciones Clave

Bai señala que la mayoría de las lecciones parecen obvias en retrospectiva. La brecha entre la predicción a corto plazo y la planificación a largo plazo es crítica — y no se captura con la pérdida de predicción estándar. El artículo detalla experimentos que gradualmente expusieron esta limitación.