Генерируем игровые миры из видео с TinyWorlds

Представьте, что можно создать интерактивный 2D-мир, просто показав нейросети видео с геймплеем. Без ручного кодирования физики, без прописывания логики персонажей. Звучит как научная фантастика? Еще недавно так и было, но проект Genie от Google DeepMind показал, что это возможно. А сегодня, благодаря репозиторию TinyWorlds, мы можем заглянуть под капот этой технологии и даже запустить ее у себя.

TinyWorlds — это минималистичная реализация архитектуры Genie, созданная, чтобы разработчики и исследователи могли понять, как работают так называемые "мировые модели" (world models). Давайте разберемся, что это такое и почему это может быть интересно каждому из нас.

Что еще за "мировые модели"?

Если говорить просто, мировая модель — это нейросеть, которая учится понимать "законы физики" и правила конкретного мира. Посмотрев множество видео, она строит внутреннее представление о том, как объекты взаимодействуют, как движется персонаж, и что произойдет в следующий момент времени.

По сути, это попытка сжать всю сложность мира в компактную модель, способную предсказывать его будущее состояние. Это не просто генерация красивых картинок, а создание симуляции, с которой можно взаимодействовать. Зачем это нужно?

• Для обучения роботов: Модель может дать роботу понимание физического мира.
• Для симуляций: Можно обучать другие нейросети внутри такого "карманного мира", что гораздо быстрее и дешевле, чем в реальности.
• Для игр и творчества: А это самое интересное для нас. Представьте себе процедурную генерацию не просто уровней, а целых интерактивных реальностей со своей уникальной логикой.

Главный трюк: учимся без учителя

Ключевая проблема при обучении таких моделей — данные. Обычно нейросети для игр требуют гигантские датасеты с пометками вида "кадр -> действие игрока -> следующий кадр". Собирать их долго и дорого.

И вот здесь кроется магия Genie и TinyWorlds. Эта архитектура способна обучаться на видео без каких-либо данных о действиях игрока. Модель сама "додумывает", какое скрытое действие произошло между двумя кадрами. Это открывает доступ к безграничному источнику данных — миллионам геймплейных роликов с YouTube.

Как устроен TinyWorlds: разбираем по косточкам

Проект состоит из трех ключевых модулей, которые работают в связке. Архитектура довольно изящна и напоминает современные языковые модели, только вместо текста здесь — "слова" из мира видеоигр.

1. Видео-токенизатор: превращаем кадры в "слова"

Первый шаг — превратить видео в последовательность дискретных токенов, понятных для трансформера. Этот модуль работает как продвинутый видеокодек:

• Он берет кадр, делит его на патчи (небольшие квадраты).
• С помощью сверточных сетей и специального пространственно-временного трансформера (Space-Time Transformer) он сжимает каждый патч в компактный вектор.
• Затем эти векторы квантуются, то есть превращаются в дискретные токены из заранее определенного словаря.

На выходе мы получаем "текст", который описывает каждый кадр. Это позволяет модели работать не с бесконечным пространством пикселей, а с конечным набором "слов", что значительно упрощает задачу.

2. Токенизатор действий: вычисляем скрытые ходы

Это сердце всей системы и самая остроумная часть. Этот модуль смотрит на два последовательных кадра и пытается угадать, какое действие привело от первого ко второму.

• Энкодер получает два кадра и на выходе выдает один токен действия.
• Декодер пытается восстановить второй кадр, используя первый кадр и этот самый токен действия.

В процессе обучения модель учится кодировать в этот токен самую важную информацию об изменении между кадрами. Например, если персонаж прыгнул, токен будет содержать информацию о прыжке. Так модель и строит свой словарь действий: "прыжок", "движение влево", "атака" и т.д., даже не зная этих слов.

3. Модель динамики: предсказываем будущее

Это уже классический трансформер, который работает как модель для генерации текста. Он получает на вход последовательность токенов кадров и токенов действий, а его задача — предсказать токены следующего кадра.

Обучается он по принципу, похожему на BERT: часть токенов в последовательности маскируется, и модель должна их угадать. Во время работы (инференса) происходит следующее:

1. Игрок "нажимает кнопку", то есть выбирает один из выученных токенов действия.
2. Модель динамики, получив историю кадров и это действие, генерирует токены нового кадра.
3. Видео-декодер превращает эти токены обратно в изображение, которое видит игрок.

И этот цикл повторяется, создавая интерактивный геймплей в реальном времени.

Пробуем сами: от "Зельды" до "Соника"

Автор репозитория позаботился о том, чтобы мы могли не только изучить код, но и поиграть с готовыми моделями. Процесс запуска довольно прост:

1. Клонируем репозиторий и устанавливаем зависимости:

   git clone https://github.com/AlmondGod/tinyworlds.git
   cd tinyworlds
   pip install -r requirements.txt
   

2. Скачиваем готовые датасеты (есть наборы по играм PicoDoom, Pong, Zelda, Pole Position и Sonic):

   python scripts/download_assets.py datasets --pattern "zelda_frames.h5"
   

3. Или скачиваем уже обученную модель (например, для Sonic) и запускаем инференс:

   # Скачиваем чекпоинты
   python scripts/download_assets.py models --suite-name sonic
   # Запускаем
   python scripts/run_inference.py --config configs/inference.yaml -- use_latest_checkpoints=true dataset=SONIC
   

Код хорошо документирован и использует современные возможности PyTorch, такие как torch.compile и DistributedDataParallel для ускорения обучения.

Кому и зачем это нужно?

TinyWorlds — это не готовый движок для вашей следующей игры. Это, в первую очередь, образовательный и исследовательский инструмент.

• Для студентов и исследователей ИИ: Это прекрасная возможность разобраться в архитектуре мировых моделей на живом, работающем примере. Код чистый и не перегружен лишними деталями.
• Для геймдев-разработчиков: Это взгляд в будущее. Технологии, подобные этой, однажды могут кардинально изменить подходы к созданию контента, тестированию и даже дизайну игровых механик.
• Для энтузиастов: Это просто очень круто! Возможность своими руками запустить модель, которая учится правилам игры по видео, вдохновляет и расширяет представление о возможностях нейросетей.

Автор активно призывает к контрибьюции, оставив в коде и README множество пометок TODO с идеями для улучшений: от внедрения Mixture of Experts до добавления новых датасетов. Так что если вы ищете интересный open-source проект для участия — это отличный кандидат.

В общем, TinyWorlds — яркий пример того, как прорывная идея из исследовательской лаборатории DeepMind становится доступной всему сообществу. Обязательно загляните на GitHub, поставьте звездочку и попробуйте создать свой собственный маленький мир.