LeMario: Обучение мировой модели JEPA на Super Mario Bros — Техническое руководство и постмортем

Проект LeMario Бенджамина Бая воспроизводит LeWorldModel — небольшую Joint-Embedding Predictive Architecture (JEPA) — с нуля и обучает её на Super Mario Bros. Подробное техническое описание охватывает полную архитектуру, настройку обучения и откровенный разбор того, что сработало, а что нет.

Обзор архитектуры

Модель обрабатывает четыре последовательных кадра Mario. Кодировщик изображений сжимает каждый кадр в 192-мерный латентный вектор: z_t = E_θ(x_t), z_t ∈ R^192. Действия (состояния кнопок Left, Right, Up, Down, A, B за пять кадров эмулятора) кодируются отдельно в другой 192-мерный вектор через кодировщик действий.

Латентные представления кадров и действий поступают в каузальный предиктор с шестью блоками трансформера. Действия внедряются с помощью Adaptive LayerNorm Zero (AdaLN-Zero), который создаёт сдвиг, масштаб и управление стробированием для каждой ветви (внимание и MLP). Стробирование контролирует, насколько сильно обновления влияют на предсказанное состояние. Веса начинаются с нуля, чтобы предиктор учился постепенно.

Предиктор выдаёт три предсказанных будущих латента: ẑ₁, ẑ₂, ẑ₃. Они сравниваются с латентами трёх реальных следующих кадров с помощью MSE-потерь: L_pred = MSE([ẑ₁,ẑ₂,ẑ₃], [z₁,z₂,z₃]). Для предотвращения коллапса представлений добавляется регуляризация SIGReg: L = L_pred + 0.1 * L_SIGReg.

Что сработало

• Модель обобщалась на новые эпизоды.
• Эффективно использовала информацию о действиях.
• Предсказывала будущее на пять шагов лучше сильных бейзлайнов.
• Планирование без вознаграждения могло перемещать Марио в пределах 2-5 пикселей от близких целевых изображений.

Что не сработало

• Когда цели располагались дальше по уровню, Марио не мог надёжно перепрыгивать первое крупное препятствие.
• Не удавалось двигаться к единственному удалённому целевому изображению.
• Вывод: предсказывать игру ≠ научиться достигать прогресса. Модель выучила динамику на уровне пикселей, но не долгосрочное планирование или целенаправленное поведение.

Ключевые уроки

Бай отмечает, что большинство уроков кажутся очевидными задним числом. Разрыв между краткосрочным предсказанием и долгосрочным планированием критичен — и не улавливается стандартной функцией потерь при предсказании. В статье описаны эксперименты, которые постепенно выявили это ограничение.