Автор решил воспроизвести LeWorldModel — небольшую JEPA-архитектуру, которая учит динамику мира по пикселям и действиям. Но вместо Push-T он взял Super Mario Bros. Модель назвали LeMario.
Архитектура работает так: vision encoder сжимает каждый кадр в 192-мерный latent. Action encoder превращает последовательность нажатий кнопок (Left, Right, Up, Down, A, B) в вектор той же размерности. Дальше идёт causal predictor из шести transformer-блоков. Действия попадают в блоки через Adaptive LayerNorm Zero (AdaLN-Zero). Вместо простого конката они задают shift, scale и gate для каждой ветки — attention и MLP. Это позволяет модели, например, при прыжке усилить признаки вертикального движения. На выходе predictor выдаёт три предсказанных будущих latent’а. Их сравнивают с реальными через MSE loss. Чтобы модель не схлопнула все представления в ноль, используют SIGReg — он следит, чтобы реальные latent’ы оставались разнообразными.
LeMario обучили на 737 134 кадрах из 280 эпизодов в 32 уровнях. На тесте модель била baseline «ничего не меняется»: на пяти шагах вперёд ошибка была на 45% меньше. Перемешивание действий ухудшало результат на 47.5%. Модель действительно научилась предсказывать короткие последовательности, зависящие от нажатий.
Потом автор подключил Cross-Entropy Method (CEM) для планирования. Модель перебирала в воображении сотни последовательностей действий и выбирала те, чей финальный latent ближе к latent’у целевого кадра. Без призов, без reward engineering. Звучало круто. Но Марио еле двигался. При цели через 30 пикселей он прошёл четыре.
Автор начал разбираться по частям. Заморозил JEPA и обучил маленький probe (MLP: 192 → 128 → GELU → 128 → 2), чтобы восстановить координаты Марио из latent’а. Горизонтальная позиция восстанавливалась почти идеально (R² = 0.997), вертикальная — плохо (R² = 0.188). Когда автор временно заменил расстояние в latent’е на оценку probe, Марио дошёл до цели почти впритык — от x=40 до x=71 при цели x=72. Позже, с локальным перепланированием, добрался до x=176 при цели x=177.
Проблема оказалась в том, что две далёкие точки в игре могут выглядеть почти одинаково из-за скроллящейся камеры. Encoder считал их близкими, хотя по геймплею между ними пропасть. Разбивка маршрута на промежуточные цели помогла — Марио дошёл до x=314. Но стоило цели потребовать прыжка, планировщик снова валился. Вертикаль в latent’е была слаба.
Автор выделил три главные проблемы. Первая: predictive state не равен control state. Encoder учится предсказывать картинку (фазы врагов, анимацию, таймер), а не то, что контролирует прогресс. Вторая: CEM ищет слабые места модели. Он не замечает, что та ошибается. Третья: условия в Super Mario Bros. кардинально отличаются от Push-T. Там фиксированная камера, плавное движение, тысячи экспертных эпизодов. Здесь — скролл, импульс, ямы, враги, смерть и всего один проход по датасету.
В итоге LeMario не научился проходить игру. Но он показал короткую динамику, зависящую от действий. Автор понял, что копировать архитектуру недостаточно — окружение, данные и метрики не менее важны.