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

Кроме того, можно провести параллели между моделями генерации изображений и мировыми моделями, указывающие на то, что прогресс, достигнутый в моделях генерации изображений, может быть повторен для пользы мировых моделей. С тех пор, как использование трансформеров в фреймворках обработки естественного языка приобрело популярность, появились фреймворки DALL-E и VQGAN. Эти фреймворки реализовали дискретные автокодировщики для преобразования изображений в дискретные токены и смогли построить высокомощные и эффективные модели генерации текста и изображений, используя возможности последовательного моделирования автoregressивных трансформеров. В то же время диффузионные модели получили популярность, и сегодня диффузионные модели установили себя как доминирующий парадигма для генерации высококачественных изображений. Благодаря возможностям, предлагаемым диффузионными моделями и обучением с подкреплением, предпринимаются попытки объединить эти два подхода, с целью воспользоваться гибкостью диффузионных моделей как моделей траекторий, моделей вознаграждения, планировщиков и как политики для аугментации данных в офлайн-обучении с подкреплением.

Мировые модели предлагают перспективный метод для обучения агентов обучения с подкреплением безопасно и эффективно. Традиционно эти модели используют последовательности дискретных латентных переменных для симуляции динамики окружения. Однако это сжатие может упустить из виду важные визуальные детали, необходимые для обучения с подкреплением. В то же время диффузионные модели приобрели популярность для генерации изображений, бросая вызов традиционным методам, использующим дискретные латентные переменные. Вдохновленные этим сдвигом, в этой статье мы поговорим о DIAMOND (DIffusion As a Model Of eNvironment Dreams), агенте обучения с подкреплением, обученном внутри диффузионной мировой модели. Мы исследуем необходимые проектные решения, чтобы сделать диффузию подходящей для моделирования мира, и показываем, что улучшенные визуальные детали приводят к лучшей производительности агента. DIAMOND устанавливает новый эталон на конкурентном тесте Atari 100k, достигая среднего человеческого нормализованного балла 1,46, самого высокого для агентов, обученных полностью внутри мировой модели.

DIAMOND: Диффузия как модель окружающей среды

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

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

Фреймворк DIAMOND (DIffusion As a Model Of eNvironment Dreams) является агентом обучения с подкреплением, обученным внутри диффузионной мировой модели. Фреймворк DIAMOND делает тщательные проектные решения, чтобы обеспечить, что его диффузионная мировая модель остается эффективной и стабильной на протяжении длительных временных горизонтов. Фреймворк предоставляет качественный анализ, чтобы продемонстрировать важность этих проектных решений. DIAMOND устанавливает новый эталон с средним человеческим нормализованным баллом 1,46 на хорошо известном бенчмарке Atari 100k, самом высоком для агентов, обученных полностью внутри мировой модели. Работа в пространстве изображений позволяет диффузионной мировой модели DIAMOND без проблем заменить окружающую среду, предлагая более глубокие знания о поведении мировой модели и агента. Заметно, что улучшенная производительность в определенных играх обусловлена лучшим моделированием критических визуальных деталей. Фреймворк DIAMOND моделирует окружающую среду как стандартный частично наблюдаемый марковский процесс принятия решений с набором состояний, набором дискретных действий и набором наблюдений изображений. Функции перехода описывают динамику окружения, а функция вознаграждения сопоставляет переходы со скалярными вознаграждениями.

DIAMOND: Методология и архитектура

В своей основе диффузионные модели являются классом генеративных моделей, которые генерируют образец, обратив процесс шумоподавления, и черпают вдохновение из термодинамики неравновесных систем. Фреймворк DIAMOND учитывает диффузионный процесс, индексированный непрерывной временной переменной с соответствующими маргинальными и граничными условиями с тривиальной неструктурированной априорной распределением. Кроме того, чтобы получить генеративную модель, которая сопоставляет шум с данными, фреймворк DIAMOND должен обратить процесс, а процесс обратного хода также является диффузионным процессом, работающим в обратном направлении во времени. Кроме того, в любой данной точке времени не тривиально оценить функцию балла, поскольку фреймворк DIAMOND не имеет доступа к истинной функции балла, и модель преодолевает это препятствие, реализуя цель совпадения баллов, подход, который позволяет фреймворку обучать модель балла без знания лежащей в основе функции балла. Модель диффузии на основе балла предоставляет безусловную генеративную модель. Однако условная генеративная модель динамики окружения требуется для служения в качестве мировой модели, и для служения этой цели фреймворк DIAMOND рассматривает общий случай подхода частично наблюдаемого марковского процесса принятия решений, в котором фреймворк может использовать прошлые наблюдения и действия, чтобы аппроксимировать неизвестное марковское состояние. Как продемонстрировано на рисунке 1, фреймворк DIAMOND использует эту историю, чтобы обусловить диффузионную модель, оценить и сгенерировать следующее наблюдение напрямую. Хотя фреймворк DIAMOND может использовать любой решатель СДУ или ОДУ в теории, существует компромисс между количеством функциональных вычислений и качеством выборки, что существенно влияет на стоимость вывода диффузионных моделей.

Основываясь на этих знаниях, давайте теперь рассмотрим практическую реализацию фреймворка DIAMOND диффузионной мировой модели, включая коэффициенты дрейфа и диффузии, соответствующие определенному выбору подхода диффузии. Вместо того, чтобы выбрать DDPM, естественно подходящий кандидат на эту задачу, фреймворк DIAMOND строится на основе формулировки EDM и учитывает ядро возмущения с действительной функцией диффузионного времени, называемой графиком шума. Фреймворк выбирает предусловители, чтобы сохранить входную и выходную дисперсию для любого уровня голоса. Обучение сети смешивает сигнал и шум адаптивно, в зависимости от уровня деградации, и когда шум низок, и цель становится разницей между чистым и возмущенным сигналом, т. е. добавленным гауссовским шумом. Интуитивно это предотвращает обучающую цель от того, чтобы стать тривиальной в режиме низкого шума. На практике эта цель имеет высокую дисперсию на экстремумах графика шума, поэтому модель выбирает уровень шума из логарифмически нормального распределения, выбранного эмпирически, чтобы сконцентрировать обучение вокруг средних областей графика шума. Фреймворк DIAMOND использует стандартный компонент U-Net 2D для векторного поля и сохраняет буфер прошлых наблюдений и действий, которые фреймворк использует для обусловливания себя. Затем фреймворк DIAMOND конкатенирует эти прошлые наблюдения с следующим шумным наблюдением и входными действиями через адаптивные слои нормализации групп в остаточных блоках U-Net.

DIAMOND: Эксперименты и результаты

Для всесторонней оценки фреймворк DIAMOND выбирает бенчмарк Atari 100k. Бенчмарк Atari 100k состоит из 26 игр, предназначенных для проверки широкого спектра возможностей агента. В каждой игре агент ограничен 100k действиями в окружающей среде, что примерно эквивалентно 2 часам игры человека, чтобы выучить игру до оценки. Для сравнения, неограниченные агенты Atari обычно обучаются в течение 50 миллионов шагов, представляя 500-кратное увеличение опыта. Мы обучили DIAMOND с нуля, используя 5 случайных семян для каждой игры. Каждый запуск обучения требовал около 12 ГБ видеопамяти и занял примерно 2,9 дня на одном Nvidia RTX 4090, что составляет 1,03 года GPU в общей сложности. Следующая таблица предоставляет балл для всех игр, среднее значение и среднее межквартильное значение нормализованных человеческих баллов.

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

Результаты показывают, что DIAMOND работает исключительно хорошо на бенчмарке, превосходя человеческих игроков в 11 играх и достигая сверхчеловеческого среднего нормализованного балла 1,46, устанавливающего новый рекорд для агентов, обученных полностью внутри мировой модели. Кроме того, среднее межквартильное значение DIAMOND сопоставимо с STORM и превышает все другие базовые значения. DIAMOND отличается в окружающих средах, где захват небольших деталей имеет решающее значение, таких как Asterix, Breakout и RoadRunner. Кроме того, как обсуждалось ранее, фреймворк DIAMOND имеет гибкость реализации любой диффузионной модели в своем конвейере, хотя он выбирает подход EDM, было бы естественным выбором выбрать модель DDPM, поскольку она уже реализуется во многих приложениях генерации изображений. Чтобы сравнить подход EDM с реализацией DDPM, фреймворк DIAMOND обучает обе вариации с той же сетевой архитектурой на том же общем статическом наборе данных с более чем 100k кадрами, собранными с помощью экспертной политики. Количество шагов денойзинга напрямую связано со стоимостью вывода мировой модели, и поэтому меньшее количество шагов уменьшит стоимость обучения агента на воображаемых траекториях. Чтобы гарантировать, что наша мировая модель остается вычислительной, сопоставимой с другими базовыми значениями, такими как IRIS, которая требует 16 функциональных вычислений на каждый временной шаг, мы стремимся использовать не более десятков шагов денойзинга, предпочтительно меньше. Однако установка количества шагов денойзинга слишком низко может ухудшить визуальное качество, что приведет к ошибкам, накапливающимся. Чтобы оценить стабильность разных вариантов диффузии, мы отображаем воображаемые траектории, сгенерированные автoregressивно до t = 1000 временных шагов в следующем рисунке, используя разное количество шагов денойзинга n ≤ 10.

Мы наблюдаем, что использование DDPM (а) в этом режиме приводит к серьезным ошибкам, накапливающимся, что вызывает быстрое выход мировой модели из распределения. В отличие от этого, диффузионная мировая модель на основе EDM (б) остается намного более стабильной на протяжении длительных временных горизонтов, даже с одним шагом денойзинга. Воображаемые траектории с диффузионными мировыми моделями на основе DDPM (слева) и EDM (справа) показаны. Начальное наблюдение в момент t = 0 одинаково для обоих, и каждая строка соответствует уменьшающемуся количеству шагов денойзинга n. Мы наблюдаем, что генерация на основе DDPM страдает от ошибок, накапливающихся, и меньшее количество шагов денойзинга приводит к более быстрому накоплению ошибок. В отличие от этого, диффузионная мировая модель DIAMOND остается намного более стабильной, даже для n = 1. Оптимальное одношаговое предсказание является ожиданием над возможными реконструкциями для данного шумного входа, которое может быть вне распределения, если постериорное распределение является многомодальным. Хотя некоторые игры, такие как Breakout, имеют детерминированные переходы, которые можно точно смоделировать с помощью одного шага денойзинга, другие игры демонстрируют частичную наблюдаемость, что приводит к многомодальным распределениям наблюдений. В этих случаях итеративный решатель необходим для руководства процедурой выборки к конкретному режиму, как проиллюстрировано в игре Boxing в следующем рисунке. Следовательно, фреймворк DIAMOND устанавливает n = 3 во всех наших экспериментах.

Вышеуказанный рисунок сравнивает одношаговую (верхний ряд) и многошаговую (нижний ряд) выборку в Boxing. Движения черного игрока непредсказуемы, что приводит к тому, что одношаговая денойзинга интерполирует между возможными результатами, в результате чего получается размытое предсказание. В отличие от этого, многошаговая выборка производит четкое изображение, направляя генерацию к конкретному режиму. Интересно, что, поскольку политика контролирует белого игрока, его действия известны мировой модели, что исключает двусмысленность. Следовательно, и одношаговая, и многошаговая выборка правильно предсказывают положение белого игрока.

На вышеуказанном рисунке траектории, воображаемые DIAMOND, обычно демонстрируют более высокое визуальное качество и являются более верными истинной окружающей среде по сравнению с теми, которые воображаются IRIS. Траектории, сгенерированные IRIS, содержат визуальные несоответствия между кадрами (выделены белыми коробками), такие как враги, отображаемые как вознаграждения и наоборот. Хотя эти несоответствия могут повлиять только на несколько пикселей, они могут существенно повлиять на обучение с подкреплением. Например, агент обычно стремится целиться в вознаграждения и избегать врагов, поэтому эти небольшие визуальные несоответствия могут сделать более сложным обучение оптимальной политики. Рисунок показывает последовательные кадры, воображаемые IRIS (слева) и DIAMOND (справа). Белые коробки выделяют несоответствия между кадрами, которые появляются только в траекториях, сгенерированных IRIS. В Asterix (верхний ряд) враг (оранжевый) становится вознаграждением (красным) во втором кадре, затем возвращается к врагу в третьем, и снова к вознаграждению в четвертом. В Breakout (средний ряд) кирпичи и счет являются несоответствующими между кадрами. В Road Runner (нижний ряд) вознаграждения (маленькие синие точки на дороге) отображаются несоответствующим образом между кадрами. Эти несоответствия не появляются в DIAMOND. В Breakout счет надежно обновляется на +7, когда красный кирпич разрушается.

Вывод

В этой статье мы говорили о DIAMOND, агенте обучения с подкреплением, обученном внутри диффузионной мировой модели. Фреймворк DIAMOND делает тщательные проектные решения, чтобы обеспечить, что его диффузионная мировая модель остается эффективной и стабильной на протяжении длительных временных горизонтов. Фреймворк предоставляет качественный анализ, чтобы продемонстрировать важность этих проектных решений. DIAMOND устанавливает новый эталон с средним человеческим нормализованным баллом 1,46 на хорошо известном бенчмарке Atari 100k, самом высоком для агентов, обученных полностью внутри мировой модели. Работа в пространстве изображений позволяет диффузионной мировой модели DIAMOND без проблем заменить окружающую среду, предлагая более глубокие знания о поведении мировой модели и агента. Заметно, что улучшенная производительность в определенных играх обусловлена лучшим моделированием критических визуальных деталей. Фреймворк DIAMOND моделирует окружающую среду как стандартный частично наблюдаемый марковский процесс принятия решений с набором состояний, набором дискретных действий и набором наблюдений изображений. Функции перехода описывают динамику окружения, а функция вознаграждения сопоставляет переходы со скалярными вознаграждениями. Фреймворк DIAMOND моделирует окружающую среду как стандартный POMDP или частично наблюдаемый марковский процесс принятия решений с набором состояний, набором дискретных действий и набором наблюдений изображений. Функции перехода описывают динамику окружения, а функция вознаграждения сопоставляет переходы со скалярными вознаграждениями.