LeMario: Training eines JEPA-Weltmodells mit Super Mario Bros — Technischer Leitfaden und Nachbetrachtung

Benjamin Bais LeMario-Projekt reproduziert LeWorldModel – eine kleine Joint-Embedding Predictive Architecture (JEPA) – von Grund auf und trainiert sie auf Super Mario Bros. Die detaillierte technische Beschreibung umfasst die vollständige Architektur, das Training und eine ehrliche Nachlese zu dem, was funktioniert hat und was nicht.

Architekturübersicht

Das Modell verarbeitet vier aufeinanderfolgende Mario-Bilder. Ein Vision-Encoder komprimiert jedes Bild in einen 192-dimensionalen latenten Vektor: z_t = E_θ(x_t), z_t ∈ R^192. Aktionen (Tastenzustände für Links, Rechts, Hoch, Runter, A, B über fünf Emulator-Frames) werden separat über einen Aktions-Encoder in einen weiteren 192-dimensionalen Vektor kodiert.

Die Bild- und Aktions-Latenten werden in einen kausalen Prädiktor mit sechs Transformer-Blöcken eingespeist. Aktionen werden über Adaptive LayerNorm Zero (AdaLN-Zero) injiziert, das pro Zweig (Attention und MLP) Verschiebungs-, Skalierungs- und Gate-Steuerungen erzeugt. Das Gate steuert, wie stark Aktualisierungen den vorhergesagten Zustand beeinflussen. Die Gewichte beginnen bei Null, sodass der Prädiktor allmählich lernt.

Der Prädiktor gibt drei vorhergesagte zukünftige Latenten aus: ẑ₁, ẑ₂, ẑ₃. Diese werden mit den Latenten der drei tatsächlichen nächsten Bilder mittels MSE-Verlust verglichen: L_pred = MSE([ẑ₁,ẑ₂,ẑ₃], [z₁,z₂,z₃]). Um einen Repräsentationskollaps zu verhindern, wird SIGReg-Regularisierung hinzugefügt: L = L_pred + 0,1 * L_SIGReg.

Was funktioniert hat

• Das Modell generalisierte auf nicht gesehene Episoden.
• Es nutzte Aktionsinformationen effektiv.
• Es sagte zukünftige fünf Schritte besser voraus als starke Baselines.
• Rohe belohnungsfreie Planung konnte Mario innerhalb von 2–5 Pixeln an nahegelegene Bildziele bewegen.

Was scheiterte

• Wenn Ziele weiter im Level platziert wurden, konnte Mario nicht zuverlässig über das erste große Hindernis springen.
• Er konnte nicht zu einem einzelnen entfernten Zielbild navigieren.
• Fazit: Das Spiel vorherzusagen bedeutet nicht, Fortschritt zu lernen. Das Modell lernte Pixel-Dynamiken, aber keine langfristige Planung oder zielgerichtetes Verhalten.

Wichtige Erkenntnisse

Bai merkt an, dass die meisten Lehren im Nachhinein offensichtlich erscheinen. Die Lücke zwischen kurzfristiger Vorhersage und langfristiger Planung ist entscheidend – und wird durch den standardmäßigen Vorhersageverlust nicht erfasst. Der Beitrag beschreibt Experimente, die diese Einschränkung schrittweise aufdeckten.