Симуляция DOOM через нейросеть

Десятки, а то и сотни трехмерщиков и разработчиков работают над играми. Впрочем, примерный образ пайплайна всем и так известен. Сегодня мы не рассказываем про нейронку, которая запилит Uncharted 5 или Dark Souls 4 — она интересный кейс среди нейронок последних годов. Обычно, когда мы говорим про генеративные ИИ, мы представляем себе Идеограмм, Stable Diffuison или SORA.

Но вот разработчики с Google создалидвижок, который генерирует игровой процесс уже существующей игры. И пока что это старющий Doom из 90s. Как примерно работает подобная нейронка — рассказываем ниже.

Как утверждают сами разработчики — первым делом они обучили RL-агента. Чтобы собрать данные для обучения этой модели, с целью навязывания учителем, исследователи Google сначал обучают отдельную модель взаимодействию с окружающей средой. Набор данных, который выдает RL-агент после обучения становится датасетом для генеративной ИИ. Кстати, она разработана на базе открытой Stable Diffusion.

Как работает RL-агент?

Он принимает решения о том, какие действия предпринять в каждый момент времени, основываясь на текущем состоянии среды.

Среда в свою очередь откликается на действия агента, изменяя своё состояние и предоставляя агенту новую информацию и вознаграждение.

Основной задачей агента является разработка стратегии (или политики), которая максимизирует совокупное вознаграждение в долгосрочной перспективе. Циклически нейронка улучшает свои действия. ИИ оперирует состояниями (state), которые представляют информацию о текущем положении дел в среде.

Агент осуществляет действия (action), которые изменяют состояние среды. В-третьих, за каждое действие агент получает вознаграждение (reward), которое может быть положительным или отрицательным, в зависимости от его влияния на достижение цели. Наконец, существует функция перехода (transition function), которая описывает, как действия агента изменяют состояние среды.

Короче говоря, типичный RL-агент действует по принципу: получаем пряники и учимся получать больше пряников за правильные действия. Но «ученик» Google — это источник будущих данных для генеративной нейронки, а значит он должен формировать достаточно информации для симуляции реального/человеческого игрового процесса. Никакой максимизации вознаграждения.

Впрочем, принцип обучения строится на установки связей между actions-frames. Так как классический DOOM — это не 144 FPS игра, то проблем с обучением-связками фактически у ресерчеров не возникает. Агент тренируется, получает ответ от окружения и в соответствии учится взаимодействовать с окружением, используя разные игровые навыки.

Выглядит это все примерно так:

S — это пространство скрытых (латентных) состояний среды. В случае DOOM, это, например, данные, хранящиеся в оперативной памяти программы.

O — это наблюдаемые проекции этих скрытых состояний, визуальные данные, которые игрок видит на экране в виде отрисованных пикселей.

V — это функция, которая преобразует скрытые состояния в видимые, то есть логика рендеринга игры (процесс отрисовки графики).

A — это набор действий, которые может совершать игрок, например, нажатия клавиш или движения мышью.

p — это функция, которая описывает вероятность перехода из одного состояния в другое, исходя из выбранных действий и текущего состояния игры.

Пример для DOOM: когда игрок нажимает на клавишу, происходит некий процесс в динамической памяти программы (S), который преобразуется в изменение на экране (O), например, игрок поворачивает голову. Логика рендеринга игры (V) отвечает за то, чтобы правильно показать это изменение на экране. При этом действия игрока (A) влияют на дальнейшее поведение программы, и функция вероятности перехода (p) описывает, как игра реагирует на те или иные команды.

Модель предсказывает, как будет выглядеть следующее наблюдение (например, следующий кадр на экране), исходя из предыдущих действий и наблюдений.

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

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

Изначально Stable Diffuion обучалась на текстовых запросах, но теперь она адаптируется для генерации кадров (или изображений) на основе последовательности действий и наблюдений агента.

Процесс начинается с того, что модель учится на данных агента, которые включают в себя действия (например, нажатия клавиш) и наблюдения (например, предыдущие кадры).

Модель убирает текстовую составляющую, которая была в оригинальной версии, и вместо этого использует последовательность действий агента.

Каждое действие (например, определённое нажатие клавиши) преобразуется в числовое представление — векторное вложение (embedding). Этот процесс заменяет текстовое управление на управление последовательностью действий.

Для того чтобы учесть наблюдения (предыдущие кадры), модель сначала кодирует их в «латентное пространство» с помощью автокодировщика (auto-encoder), сжимая данные.

Затем эти закодированные кадры добавляются к «зашумленным» версиям текущих кадров. Ресерчеры также затестили использование так называемого «cross-attention» для обработки прошлых наблюдений, но этот подход не дал значительного улучшения.

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

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

Когда модель обучается на реальных данных с использованием метода teacher-forcing (когда ей показывают реальные кадры во время обучения), она работает хорошо.

Но когда модель сама начинает предсказывать следующие кадры на основе предыдущих предсказаний (так называемая авторегрессия), качество предсказаний быстро ухудшается.

Это происходит потому, что ошибки накапливаются, и с каждым следующим шагом предсказания становятся всё менее точными. Чтобы предотвратить это ухудшение, во время обучения модель специально «искажает» кадры, добавляя гауссов шум (случайные искажения), а также передаёт уровень этого шума в модель.

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

Модель Stable Diffusion v1.4 использует автокодировщик, который сжимает изображения, переводя их в так называемое латентное пространство (особое представление данных в уменьшенной форме). Этот автокодировщик сжимает участки размером 8×8 пикселей в 4 канала.

Однако при предсказании игровых кадров такая компрессия вызывает артефакты — небольшие искажения, особенно заметные на мелких деталях, например, на индикаторах (HUD) внизу экрана игры. Чтобы улучшить качество изображений, авторы модели решили настроить только декодер автокодировщика, который отвечает за восстановление сжатых данных обратно в изображение.

Для этого они использовали функцию потерь MSE (среднеквадратическую ошибку), которая измеряет разницу между предсказанными и реальными кадрами. Это позволило уменьшить артефакты и улучшить визуальное качество.

Важно отметить, что эта настройка декодера проводится отдельно от настройки основной части модели (U-Net), которая отвечает за генерацию изображений.

Модель Stable Diffusion v1.4 использует автокодировщик, который сжимает изображения, переводя их в так называемое латентное пространство (особое представление данных в уменьшенной форме).

Этот автокодировщик сжимает участки размером 8×8 пикселей в 4 канала. Однако при предсказании игровых кадров такая компрессия вызывает артефакты — небольшие искажения, особенно заметные на мелких деталях, например, на индикаторах (HUD) внизу экрана игры.

Чтобы улучшить качество изображений, авторы модели решили настроить только декодер автокодировщика, который отвечает за восстановление сжатых данных обратно в изображение.

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

Настройка декодера проводится отдельно от настройки основной части модели (U-Net), которая отвечает за генерацию изображений. Авторегрессия (предсказание следующих кадров на основе предыдущих) не затрагивается этой настройкой, так как она работает с латентными представлениями (сжатыми данными), а не с пикселями напрямую.

Обычно для создания качественного изображения требуется множество шагов, но в этом исследовании они обнаружили, что достаточно всего 4 шага, что позволило генерить изображения с частотой 20 кадров в секунду.

При этом каждый шаг удаления шума занимает около 10 миллисекунд, а вместе с автокодировщиком полная задержка на один кадр составляет 50 миллисекунд — 20 кадров в секунду, что достаточно для симуляции игры DOOM в риал тайме.

Когда они пытались использовать всего один шаг удаления шума, качество заметно ухудшалось, поэтому авторы также экспериментировали с методом дистилляции модели. Этот метод помогает модели работать с одним шагом, что повышает частоту кадров до 50 FPS, но качество изображения всё равно ухудшается. Поэтому они предпочитают использовать 4 шага без дистилляции.

Обучается агент через Proximal Policy Optimization (PPO), предложенный Шульманом и его коллегами в 2017 году.

В качестве базовой архитектуры для извлечения признаков используется простая сверточная нейронная сеть (CNN), подобная той, что была предложена Mnih и соавторами в 2015 году.

В ходе обучения агент получает уменьшенные версии изображений игровых кадров и карты с разрешением 160×120. Кроме того, у агента есть доступ к последним 32 действиям, которые он выполнил.

Сеть признаков преобразует каждое изображение в векторное представление размером 512. Далее, актор и критик PPO представлены двуслойной полносвязной сетью (MLP), которая работает на основе комбинации выходов сети признаков изображений и последовательности предыдущих действий.

Обучение агента для игры происходит в среде Vizdoom — тест-полигон для тестирования.

Итого:

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

Важными гиперпараметрами остаются: коэффициент дисконтирования (γ = 0.99) и коэффициент энтропии (0.1), который регулирует случайность действий агента.

В ходе каждой итерации модель обучается в течение 10 эпох с использованием пакетов данных размером 64 и с постоянной скоростью обучения 1e-4. Общее количество шагов взаимодействия с окружающей средой — 10 миллионов.

Что касается симуляционных моделей, они обучаются на основе уже предварительно обученной модели Stable Diffusion версии 1.4.

В процессе обучения все параметры архитектуры U-Net размораживаются, что позволяет модели адаптироваться к новой задаче.

Для этого используется пакет данных размером 128 с постоянной скоростью обучения 2e-5, а оптимизация осуществляется с помощью метода Adafactor без весового затухания и с обрезкой градиентов до значения 1.

Параметризация потерь в диффузионной модели изменена на предсказание скорости (v-prediction). В процессе обучения вероятность удаления контекста из предыдущих кадров составляет 10%, что позволяет применять технику Classifier-Free Guidance (CFG) во время вывода модели.

Обучение происходит на 128 устройствах TPU-v5e с параллелизацией данных. Большинство результатов в работе получены после 700 тысяч обучающих шагов. При этом для улучшения качества обучения используется метод добавления шума, при котором максимальный уровень шума составляет 0.7, а для его учета используется 10 эмбеддинговых корзин.

Для обучения модели декодера латентных признаков используются те же параметры, что и для основного денойзера, за исключением того, что для декодера используется более крупный пакет данных размером 2048.

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

Все изображения, как в процессе обучения, так и в процессе вывода, имеют разрешение 320×240 и дополнены до размера 320×256. Модель обучается с контекстом длиной 64, то есть модель получает свои собственные последние 64 предсказания и последние 64 действия агента в качестве входных данных.

По итогу ребята получают отличный результат, который повторяет игру на 80%, но с некоторыми ньюансами. К сожалению, при длительной игре — нейронка выдает артефакты с накоплением. Впрочем, подход достаточно революционный и кто знает… может, в купе с новой технологией NeRF, мы увидим что-то действительно выдающееся в ближайшие пять лет.