LeMario: Treinando um Modelo Mundial JEPA em Super Mario Bros — Guia Técnico e Análise Pós-Morte

O projeto LeMario de Benjamin Bai reproduz o LeWorldModel — uma pequena Arquitetura Preditiva de Embeddings Conjuntos (JEPA) — do zero e o treina no Super Mario Bros. A análise técnica detalhada cobre toda a arquitetura, configuração de treinamento e uma avaliação franca do que funcionou e do que não funcionou.

Visão Geral da Arquitetura

O modelo processa quatro quadros consecutivos de Mario. Um codificador visual comprime cada quadro em um vetor latente de 192 dimensões: z_t = E_θ(x_t), z_t ∈ R^192. As ações (estados dos botões para Esquerda, Direita, Cima, Baixo, A, B ao longo de cinco quadros do emulador) são codificadas separadamente em outro vetor de 192 dimensões por meio de um codificador de ações.

Os latentes de quadro e ação alimentam um preditor causal com seis blocos transformer. As ações são injetadas via Adaptive LayerNorm Zero (AdaLN-Zero), que produz controles de deslocamento, escala e porta por ramo (atenção e MLP). A porta controla a intensidade com que as atualizações afetam o estado previsto. Os pesos começam em zero, para que o preditor aprenda gradualmente.

O preditor gera três latentes futuros previstos: ẑ₁, ẑ₂, ẑ₃. Eles são comparados com os latentes dos três próximos quadros reais usando perda MSE: L_pred = MSE([ẑ₁,ẑ₂,ẑ₃], [z₁,z₂,z₃]). Para evitar colapso de representação, a regularização SIGReg é adicionada: L = L_pred + 0.1 * L_SIGReg.

O que Funcionou

• O modelo generalizou para episódios não vistos.
• Usou informações de ação de forma eficaz.
• Previu futuros de cinco passos melhor que linhas de base fortes.
• O planejamento bruto sem recompensa conseguiu mover Mario dentro de 2 a 5 pixels de metas de imagem próximas.

O que Falhou

• Quando as metas foram colocadas mais adiante no nível, Mario não conseguiu pular o primeiro obstáculo grande de forma confiável.
• Não conseguiu navegar em direção a uma única meta de imagem distante.
• Conclusão: prever o jogo ≠ aprender a progredir. O modelo aprendeu dinâmicas em nível de pixel, mas não planejamento de longo prazo ou comportamento direcionado a objetivos.

Lições Principais

Bai observa que a maioria das lições parece óbvia em retrospecto. A lacuna entre previsão de curto prazo e planejamento de longo prazo é crítica — e não é capturada pela perda de previsão padrão. O artigo detalha experimentos que gradualmente expuseram essa limitação.