Энтузиаст воссоздал JEPA-архитектуру Лекуна и обучил её на Super Mario Bros, предсказывает игру, но не проходит

Независимый разработчик, автор блога benjamin-bai.com, решил воспроизвести LeWorldModel, небольшую JEPA-архитектуру (Joint-Embedding Predictive Architecture, предиктивную архитектуру на основе совместных эмбеддингов), которая учится динамике игрового мира по изображениям и действиям. Оригинальная работа о LeWorldModel применяла такую модель для планирования без сигнала вознаграждения (reward-free planning) в тестовой среде Push-T. Автору хотелось глубже разобраться в JEPA-архитектуре Яна Лекуна, а поскольку он любит видеоигры, он написал всю архитектуру с нуля и обучил её на Super Mario Bros, проект получил имя LeMario.

На первых проверках модель выглядела рабочей: она обобщалась на отложенные, не задействованные в обучении эпизоды, учитывала нажатые кнопки и предсказывала состояние на пять шагов вперёд лучше сильных базовых сравнений, а без какого-либо вознаграждения подводила Марио к соседним кадрам-целям с точностью до двух-пяти пикселей. Но стоило отодвинуть цель дальше вглубь уровня, как персонаж переставал надёжно перепрыгивать первое серьёзное препятствие или добираться до удалённой цели-кадра. Именно разбору этого провала посвящён весь дальнейший пост, который автор оформил как честный технический постмортем (разбор причин неудачи), а не только рассказ об успехе.

Архитектура строится на двух энкодерах и предикторе. Визуальный энкодер сжимает каждый кадр Super Mario Bros (изображение 224×224) в вектор из 192 чисел, по сути, приватное описание кадра моделью. Каждый обучающий пример включает четыре таких кадра; между соседними кадрами проходит пять кадров эмулятора, и в каждом фиксируется состояние шести кнопок (влево, вправо, вверх, вниз, A, B). Энкодер действий сжимает такую последовательность в ещё один вектор из 192 чисел. Кадровые латенты и латенты действий поступают в каузальный предиктор, шесть блоков трансформера, где каждый кадр учитывает предыдущие.

Действия внедряются через Adaptive LayerNorm Zero (AdaLN-Zero): отдельно для блока внимания и полносвязного слоя предиктор генерирует сдвиг (смещает признаки кадра), масштаб (усиливает или ослабляет конкретные признаки) и вентиль (определяет, насколько сильно обновление повлияет на состояние), итого шесть управляющих значений на блок. Например, действие «прыжок» могло бы усиливать признаки, отвечающие за вертикальное движение, и ослаблять менее значимые для прыжка признаки. Приставка «Zero» означает, что все эти веса стартуют с нуля, и предиктор постепенно, по ходу обучения, учится, какие вентили открывать под какое действие. После шести блоков трансформера небольшая проекционная головка предсказывает три будущих латентных вектора, которые сравниваются с латентами трёх реальных следующих кадров. Чтобы модель не жульничала, схлопывая все латенты в одинаковые и делая предсказание тривиально «идеальным» (представление коллапсирует), в функцию потерь добавили регуляризатор SIGReg, который заставляет латенты реальных кадров оставаться разнообразными и информативными.

LeMario обучили на 737 134 кадрах из 280 эпизодов на 32 уровнях Super Mario Bros, всего за одну эпоху. Чтобы доказать, что модель выучила именно динамику игры, а не просто минимизировала ошибку (соседние кадры часто почти не отличаются, и предсказание «ничего не изменилось» само по себе уже сильный базовый вариант), автор сравнил LeMario на отложенных эпизодах с двумя базовыми линиями: предсказанием «без изменений» и прогоном с перемешанными, случайными действиями. Перемешивание действий увеличивало ошибку предсказания на один шаг на 20,2%. На горизонте в пять шагов LeMario обгоняла базовую линию «без изменений» на 45,5%, тогда как случайные действия давали результат хуже на 47,5%, чем дальше горизонт предсказания, тем сильнее сказывались именно нажатые кнопки. Вывод автора: модель действительно выучила краткосрочную динамику Super Mario Bros, обусловленную действиями игрока.

Дальше автор попробовал не просто предсказывать, а планировать. По текущему кадру и кадру-цели энкодер строит два латента, а метод перекрёстной энтропии (Cross-Entropy Method, CEM) сэмплирует сотни последовательностей действий, прогоняет каждую через LeMario, оценивает, насколько предсказанный финальный латент близок к латенту цели, оставляет лучших кандидатов и повторяет процесс заново. Первый содержательный тест с целью подальше сразу показал проблему: старт на x=40, цель-кадр на x=72, «сырое» планирование JEPA+CEM почти не сдвинуло Марио, он остановился на x=44. Автору предстояло понять, что именно сломалось: предиктор, сам CEM или энкодер, который мог вовсе не улавливать положение Марио.

Чтобы разобраться, автор заморозил визуальный энкодер и обучил отдельный небольшой зонд, извлекающий из 192-числового латента координаты Марио в системе эмулятора. Результат резко разошёлся по осям: горизонтальная позиция восстанавливалась почти идеально, средняя ошибка 9,3 пикселя, R² = 0,997, а вертикальная гораздо хуже: средняя ошибка 21,62 пикселя, R² = 0,188. Латентное представление действительно содержало полезную информацию о положении персонажа, но качество этой информации сильно различалось по осям.

Когда автор временно стал ранжировать кандидатов CEM не по «сырому» латентному расстоянию, а по горизонтальной позиции, которую предсказывал зонд, сама модель по-прежнему воображала будущее, а CEM по-прежнему выбирал действия, менялась только оценка кандидатов, результат резко улучшился: для цели x=72 Марио дошёл от x=40 до x=71, а с локальным перепланированием, до x=176 при цели x=177. Это был первый по-настоящему рабочий прогон в проекте: он показал, что нужная информация в латенте есть, а выученная динамика способна её туда сдвигать. Возможно, дело было в том, что самая первая цель просто была слишком похожа на стартовый кадр и не создавала достаточного «давления» в латентном пространстве для планирования.

Автор убрал зонд и повторил «сырое» латентное планирование с целями примерно на середине уровня, в мирах 1-1, 2-1 и 3-1. На уровне 3-1 Марио продвинулся заметно дальше, с x=40 до x≈290, 307, и погибал уже у первого серьёзного препятствия, то есть планирование перестало быть полностью бесполезным. Но тут же выявилась новая проблема: реальное расстояние до цели оставалось огромным, 1 442 пикселя игрового мира, тогда как энкодер оценивал латентное расстояние финальной сцены до цели всего в 0,164 (CEM предсказывал 0,153, то есть сам предиктор не «галлюцинировал», ошибался именно энкодер). Причина, прокручивающаяся камера Super Mario Bros: два визуально похожих кадра могут относиться к совершенно разным точкам уровня. Модель добросовестно делала то, о чём её просили, искала путь между двумя эмбеддингами, но застревала в месте, которое просто похоже на цель по картинке, а не является ею на самом деле.

Тогда автор разбил маршрут человека-игрока на более мелкие промежуточные кадры-цели, чтобы соседние цели визуально отличались сильнее. Это помогло с движением: в одном из прогонов «сырое» планирование без зонда довело Марио до x=314, лучший результат в проекте без вспомогательного зонда. Но надёжнее оценка прогресса не стала: до второй промежуточной цели Марио сначала перелетел (x=283), затем скорректировался назад и остановился на x=239, в пяти пикселях от эталона, однако тест всё равно засчитал попытку как провал, поскольку итоговый латент не совпал с латентом цели достаточно точно (вероятно, из-за HUD, игровых индикаторов поверх кадра, или других мелких деталей изображения). Следующая цель требовала прыжка, и планировщик снова не справился, это совпало с более ранним выводом зонда о слабом представлении вертикальной позиции. Итог серии: более мелкие цели помогли с горизонтом планирования, но не сделали латентное расстояние надёжной мерой прогресса.

Автор сводит все сбои к трём причинам. Во-первых, предсказательное состояние, не то же самое, что состояние для управления: энкодер обучен представлять всё, что помогает предсказывать будущие кадры (положение камеры, фазу врага, анимацию, таймер), а не то, что нужно для контроля прогресса к цели. Во-вторых, CEM ищет именно слабые места модели: он честно находит действия, которые, по мнению LeMario, ведут к цели, но не способен распознать, что сама модель ошибается, тем самым он обнажает её слабости: путаницу визуально похожих, но разных локаций и слабое представление вертикального движения. В-третьих, Super Mario Bros изменил условия, на которых строился успех исходного метода Push-T: там использовались близкие цели из экспертных траекторий, фиксированная камера и плавное движение, при которых визуальное сходство кадров совпадало с реальным прогрессом, а модель обучалась десять эпох на 20 000 экспертных эпизодов. LeMario же обучалась одну эпоху на 280 эпизодах по 32 уровням: близкие цели превратились в цели за полуровня, фиксированная камера, в прокручивающуюся, а плавное движение, в инерцию, прыжки, ямы, врагов, анимацию и смерть. Автор скопировал архитектуру, но не сразу учёл, что вместе с ней изменились и условия, от которых зависела работоспособность самого метода планирования.

В итоге автор признаёт: несмотря на то что LeMario так и не научилась полноценно проходить Super Mario Bros, ей всё же удалось выучить достоверную краткосрочную динамику игры, обусловленную действиями игрока. Весь путь, включая ошибки и тупики, задокументирован как подробный технический разбор, а не только как история успеха.

Ключевые факты

LeMario, независимый проект по мотивам JEPA-архитектуры Яна Лекуна (конкретно статьи о LeWorldModel, которая обучалась планированию в среде Push-T): визуальный и «действийный» энкодеры сжимают кадры 224×224 и последовательности из шести кнопок в векторы из 192 чисел, шесть блоков трансформера-предиктора получают действия через AdaLN-Zero (сдвиг, масштаб и вентиль на блок), а регуляризатор SIGReg не даёт латентам схлопнуться в одинаковые.
Модель обучили на 737 134 кадрах (280 эпизодов, 32 уровня, всего одна эпоха): перемешивание действий увеличивало ошибку предсказания на 20,2%, а на горизонте в пять шагов LeMario обгоняла базовую линию «без изменений» на 45,5%, подтверждение, что модель выучила динамику, зависящую именно от действий игрока.
Зонд поверх замороженного энкодера показал резкую асимметрию: горизонтальная позиция Марио восстанавливается почти идеально (ошибка 9,3 пикселя, R² = 0,997), вертикальная, гораздо хуже (ошибка 21,62 пикселя, R² = 0,188).
«Сырое» латентное планирование через CEM почти не сдвигало Марио к дальней цели (x=40→44 вместо цели x=72), но с оценкой по зонду результат вырос до x=71, 176; на целях в полуровня выяснилось, что реальная дистанция была 1 442 пикселя, а латентное расстояние, всего 0,164, потому что прокручивающаяся камера делает далёкие места визуально похожими на цель.
Автор называет три причины провала: предсказательное представление не равно представлению для управления; CEM находит и эксплуатирует слабости модели вместо честного пути к цели; ключевые условия успеха исходного метода Push-T (фиксированная камера, близкие цели, 10 эпох на 20 000 эпизодов) не перенеслись на Super Mario Bros (прокрутка, инерция, всего одна эпоха на 280 эпизодах).

Почему это важно

История наглядно и с цифрами показывает разрыв, который часто прячется за красивыми демонстрациями моделей мира (world models): система, которая отлично предсказывает будущие кадры игры, не обязательно способна планировать действия для достижения цели. LeMario почти идеально предсказывала физику Super Mario Bros на несколько шагов вперёд и неплохо кодировала горизонтальную позицию персонажа в латенте, но при реальном планировании через CEM спотыкалась уже на первом серьёзном препятствии. Ещё показательнее конкретный механизм провала: прокручивающаяся камера заставляла энкодер путать визуально похожие, но далёкие друг от друга участки уровня, из-за чего латентное расстояние само по себе оказалось ненадёжным сигналом прогресса. Это конкретный, воспроизводимый и количественно измеренный пример более общей проблемы, с которой сталкивается любой, кто пытается строить агентов и планировщиков поверх обучаемых моделей мира, от исследований JEPA Яна Лекуна до игровых и робототехнических систем.

Кому это важно

В первую очередь, исследователям и инженерам, которые строят агентов и планировщиков поверх обучаемых моделей мира и архитектур самообучения по представлениям, включая семейство JEPA. Полезно и всем, кто использует латентное планирование (CEM, MPC и похожие методы поиска в «воображении» модели) в играх и робототехнике. Отдельная аудитория, те, кто учится через самостоятельное воспроизведение чужих статей и архитектур с нуля: пост показывает, как выглядит честный технический постмортем с конкретными числами вместо общих слов об успехе.

Как это применить

Автор фактически формулирует чек-лист для похожих проектов. Проверяйте представление модели отдельным зондом именно на той величине, которая нужна для задачи (в данном случае, координатах персонажа), а не полагайтесь только на низкую ошибку предсказания кадров. Тестируйте планирование не только на близких целях: успех на соседних кадрах может маскировать полную беспомощность на целях подальше, как и произошло с LeMario (x=40→44 вместо x=72). Отдельно проверяйте, не путает ли энкодер визуально похожие, но разные по факту состояния, особенно если среда содержит повторяющиеся виды вроде прокручивающейся камеры. И прежде чем переносить метод планирования из чужой статьи на новую задачу, сверяйте скрытые условия, при которых он вообще работал: масштаб обучающих данных (у Push-T, 10 эпох на 20 000 экспертных эпизодов против одной эпохи на 280 эпизодах у LeMario), поведение камеры, близость целей к старту.

Можно ли доверять

Это личный проект и запись в блоге на benjamin-bai.com, а не рецензируемая научная статья; обсуждение на Hacker News на момент публикации собрало умеренные 69 баллов и 9 комментариев. При этом доверие держится не на статусе источника, а на прозрачности метода: автор приводит точные числа почти для каждого эксперимента, проценты ошибки предсказания, точность зонда по каждой оси, пиксельные и латентные расстояния до целей, и честно описывает провалы наравне с успехами, оформляя материал именно как постмортем, а не как рекламный анонс. Такая проверяемая конкретика и готовность показывать отрицательный результат, сильный сигнал достоверности, хотя это по-прежнему отчёт одного автора без независимого воспроизведения результатов третьей стороной.

Риски и подводные камни

Сам автор выделяет три системных риска подхода. Первый, предсказательное представление не равно представлению, удобному для управления: энкодер учится сохранять всё, что помогает предсказывать картинку (положение камеры, анимацию, фазу врага, таймер), а не то, что нужно для контроля прогресса к цели. Второй, CEM ищет именно слабые места модели: он честно находит действия, которые, по мнению модели, ведут к цели, но не умеет распознавать, что сама модель ошиблась, из-за этого малейшая неточность представления превращается в уверенно неверный план. Третий, латентное расстояние хрупко как мера прогресса: из-за прокручивающейся камеры визуально похожие, но далёкие сцены получают маленькое латентное расстояние, и планировщик застревает в ложно «близких» местах; отдельная слабость, вертикальная позиция (прыжки) кодируется заметно хуже горизонтальной (R² 0,188 против 0,997), поэтому цели, требующие прыжка, проваливались чаще. Есть и общий методологический риск: копируя архитектуру из чужой статьи, легко упустить, что перенос её на новую среду меняет и скрытые условия, от которых зависел исходный успех метода.

«Модель научилась предсказывать игру, но это не значит, что она научилась в ней продвигаться.»

— автор проекта LeMario, блог benjamin-bai.com