Модели мира или мультиагентное обучение с подкреплением в рамках бюджета

Всем привет! Меня зовут Владимир, и в этом я году оканчиваю магистратуру “Программирование и анализ данных” в Высшей школе экономики в Петербурге, а также работаю в JetBrains Research. В этом посте я хочу рассказать про нашу статью Scalable Multi-Agent Model-Based Reinforcement Learning, которую приняли на большую конференцию AAMAS’22 (A*). Над статьёй я работал вместе с Алексеем Шпильманом, заведующим Центром анализа данных и машинного обучения в Вышке. 

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

Sample efficiency

Возможно, вы слышали про статью от OpenAI. В ней искусственный интеллект учился играть в прятки друг с другом, и в конечном итоге полностью ломал игру, потому что создавал на карте разные конструкции, чтобы превзойти оппонента. Возможно вы также слышали, что всё это обучалось на большом сервере и очень долго. Но насколько долго? Речь идёт о миллиардах взаимодействий со средой, чего невозможно достичь, имея под рукой локальный компьютер или просто сервер поменьше. 

Тут мы сталкиваемся с одной из главных проблем обучения с подкреплением в целом, и в особенности мультиагентного обучения с подкреплением — sample efficiency. Если кратко, то для того, чтобы агенту понять, что ему требуется делать на карте, нужно для начала что-то на этой карте сделать. Оценивая результаты своих действий, а если быть более точным – награду, полученную от среды, агент понимает, что он сделал правильно, но в большинстве случаев, конечно, не понимает. Поэтому процесс повторяется, пока наконец агент не сходится к какому-то локальному максимуму, из которого уже нельзя понять, как ещё улучшить награду. 

А теперь представьте, что помимо этого агента в среде что-то делают другие агенты, которые также не понимают, что происходит. Всё это выливается в огромное количество времени, которое агенты тратят на сбор опыта и взаимодействие со средой, пока наконец обучающий алгоритм не выйдет на что-то приемлемое. Это и есть проблема sample efficiency – огромное количество ресурсов, выделяемое на взаимодействие со средой, в то время как само свойство sample efficiency предполагает обучение агентов на как можно меньшем количестве данных. Штука полезная, как в средах с одним агентом, так и с несколькими. Разница в том, что добиться хороших результатов для сред с одним агентом можно гораздо легче, а применить методы из таких сред в мультиагентных аналогах ожидаемо сложнее. В своей статье я предлагаю первый по-настоящему масштабируемый метод для мультиагентных сред, который способен на порядки уменьшить количество данных, нужных агентам для обучения.

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

Model-based обучение с подкреплением

Сама идея model-based подходов проста и датируется ещё 1991 годом. Но, как и многие другие концепции из машинного обучения, расцветать такие подходы начали только с развитием глубокого обучения. Например, значимым результатом было появление AlphaGo в 2015 году. Однако тут важно понимать, что в AlphaGo та самая модель, которая и является основой таких методов, заранее известна, так как симулировать игру го не является сложной задачей: поставил фишку на стол, поменял текущую конфигурацию доски. 

Куда более сложной задачей является обучение модели в среде, которую мы не можем полностью просимулировать из любого доступного состояния. В качестве такой среды в 2019 году был выбран набор игр Atari, которые в 2014 DeepMind смогли решить с помощью обучения с подкреплением без модели. Для этого потребовались миллионы кадров, скормленных обучаемому алгоритму. Ребятам из Google в 2019 для решения этой задачи понадобилось только 100 тысяч кадров, что эквивалентно нескольким часам игры для человека. Такая значительная разница в количестве, конечно, не без потери качества итоговой политики, всё равно сильно подогрела интерес к model-based подходам. 

Давайте рассмотрим их поподробнее на примере статьи 2019 года:

Сама идея формулируется следующим образом: давайте вместо того, чтобы постоянно взаимодействовать со средой для тестирования новых действий во время обучения агента, будем собирать данные из среды, а именно в какие состояния агент переходит, совершая действия в текущем состоянии. По этим данным можно запустить отдельный процесс обучения – по состоянию и действию агента предсказывать следующее состояние. Всё это, разумеется, обучается при помощи нейронной сети. Такая нейронная сеть в литературе называется “Модель мира”, так как по своей сути является симуляцией реальной среды. Как вы могли догадаться, такую модель мира можно использовать для обучения агента без взаимодействия со средой вообще. Отсюда и главное преимущество model-based подходов – эффективность в использовании опыта или sample-efficiency. 

На данный момент самым сильным model-based подходом, в простонародье state-of-the-art, является DreamerV2 от ребят из Google Brain, выпущенный в 2021 году. Как можно догадаться по суффиксу V2, есть и первая версия алгоритма. Работают они по одной схеме, при этом вторая версия, разумеется, лучше. Поэтому в дальнейшем я буду называть алгоритмом Dreamer именно вторую версию. 

Dreamer

Итак, мы определились с семейством подходов, которые хорошо борются с sample efficiency — это model-based алгоритмы. Мы также нашли среди них лучший – Dreamer. Следующим логичным шагом является применение этого метода в мультиагентных средах. Но перед этим реклама спонсора надо разобраться, чем же Dreamer лучше других существующих подходов и почему у model-based алгоритмов так редко обновляется state-of-the-art. Для этого посмотрим на следующую картинку:

Что же на ней изображено? Это то, как непосредственно обучается политика агента. От самой среды тут только первое состояние (картинка внизу схемы), всё остальное додумывает модель. Отсюда и название – Dreamer. Эта картинка переводится в латентное пространство (проще говоря, то, как модель мира видит текущую картину в среде). Это новое латентное состояние представляется в двух переменных – h и z. Не вдаваясь в технические детали, можно считать, что h – это скрытое состояние модели, отражающее всю траекторию, которую уже прошёл агент, и реализованное с помощью рекуррентной нейронной сети, а z – отражение текущего состояния, обучающееся с помощью VAE. Из этого набора можно предсказать награду r и Value-функцию V, а также само действие агента a. Глядя на это действие, модель может предсказать следующее состояние, уже не опираясь на картинку из среды. Таким образом генерируются траектории для обучения, за счёт чего и достигается sample efficiency.

Внимательный читатель может спросить, а как же обучается сама модель? Давайте рассмотрим следующее изображение:

Здесь показано, как модель мира обучается предсказывать следующее состояние. Если не смотреть на всю схему сразу, а остановиться только на первом шаге (благо все 3 шага являются по сути идентичными), то видно, что картинка снизу преобразуется в латентное состояние z, которое в свою очередь декодируется обратно в ту же самую картинку. Кодируем — декодируем. Да, это VAE. Углубляться мы в них не будем, единственное, что стоит отметить, так это специфичный вид самого z – 32 штуки категориальных распределений по 32 класса каждое. Откуда авторы его достали — непонятно. Но оно работает, причём гораздо лучше того же нормального распределения. 

С лёгкой частью разобрались. Теперь аккуратно смотрим налево и видим ещё один z, но уже с шапочкой. Это то же самое распределение, что и у z без шапочки, только теперь модели надо предсказать его без картинки. Зачем? Очень просто: если бы модели каждый раз требовалась картинка для предсказания самого состояния, как в случае с обычным VAE, то смысла в такой модели нет. Поэтому z с шапочкой является самой главной переменной, которую обучается предсказывать модель – следующее состояние модели, предсказанное только из внутренних переменных. Наконец, как и в предыдущей схеме, агент может совершить действие и перевести скрытое состояние модели на следующий шаг. И так далее пока не кончится реальная траектория.

Из этого описания работы Dreamer можно понять как минимум одно – это сложный с точки зрения реализации алгоритм, в котором легко наделать много багов, а учитывая специфику области, не всегда понятно, есть ли вообще в реализации баг или оно просто не работает. И это только в средах с одним агентом. В мультиагентных средах всё работает по той же схеме – только теперь на вход подаётся не одна картинка, а n – по одной от каждого агента. Действий теперь разумеется тоже n. И хотя идея и лежит на поверхности, именно из-за всех этих сложностей, никто до этого всерьёз не применял model-based подходы в мультиагентных средах. 

Давайте теперь посмотрим на первый такой метод, который мы предложили в своей статье – MAMBA. Рассмотрим вкратце, зная, что за всем этим стоит Dreamer, модифицированный для мультиагентных сред. Все технические подробности вы можете прочитать непосредственно в нашей статье.

MAMBA 

Тут после Dreamer уже не должно быть ничего страшного. Идём глазами слева направо и понимаем, как обучается модель с политикой агента. 

Вот у нас были латентные состояния z и действия агентов a c предыдущего шага. Эти вектора мы передаём Communication Block, который по сути является обычным стэком Attention. Тот в свою очередь выдаёт вектора e для каждого агента, которым он может обновить свою модель мира. Для этого опять же используется рекуррентная нейронная сеть, переводящая скрытое состояние h на текущий шаг, а уже из этого состояния и наблюдения агента o (считай, картинка), предсказывается текущее латентное состояние z у каждого агента, из которого, как и в случае с Dreamer, можно предсказать всё, что нужно агенту.

На этой схеме показано, какой информацией агенты делятся друг с другом непосредственно в среде. Всё теми же латентными состояниями z и действиями a, которые всё также передаются Communication Block, с помощью которого каждый агент обновляет свою собственную модель мира. И так на каждом шаге в среде. Тут стоит отметить, что как z, так и a – это дискретные переменные, которые можно очень эффективно передавать другим агентам. Вообще, такой способ обмена сообщениями имеет несколько интересных свойств, которые мы описываемы в статье. Вкратце можно сказать, что такие сообщения очень напоминают слова – ведь слова это тоже дискретные переменные, которые неплохо справляются с описанием предметов.

Среды

Собственно, а на чём всё это тестировать? Вопрос на самом деле всегда остаётся открытым и в большинстве случаев упирается в бюджет. Понятно, что решать Starcraft II или Dota доступно только крупным компаниям типа OpenAI или DeepMind. Поэтому исследователи выкручиваются как могут, постепенно наращивая сложность сред, которые могут решаться условно на 1 gpu. 

Первым из таких важных наращений был Starcraft Multi-Agent Challenge (SMAC), этакий Старкрафт на минималках. Для ощущения масштаба приведу следующую картинку:

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

Стоит помнить, что 15 агентов – это уже очень много для текущих алгоритмов обучения с подкреплением, а в SMAC можно разогнаться и до 25. Среда вышла в 2018 году, она до сих пор является одним из главных бенчмарков в мире мультиагентных алгоритмов обучения с подкреплением, во многом благодаря своей повышенной сложности по сравнению с конкурентами, коих не очень то много. Мы отклонятся от традиций не стали и тоже тестируем MAMBA в SMAC, тем более что так гораздо удобнее сравниваться с предшественниками, которые тоже показывали свою мощь на этом бенчмарке.

Но прогресс не стоит на месте, за последние четыре года постепенно выходили новые мультиагентные среды, и  некоторые из них достойны внимания. Одной из таких сред является Flatland, симулирующий трафик 2d паровозиков по железнодорожным путям. Достаточно посмотреть на картинку, чтобы всё понять:

Среда сложна для обучения, по ней ежегодно проводятся соревнования, про одно из которых мы писали год назад, когда наша команда из ВШЭ — Санкт-Петербург заняла первое место. Наверное, главным преимуществом этой среды является гибкое регулирование сложности – добавить паровозиков или расширять карту не представляет никаких проблем, ведь сама карта генерируется процедурно. Что интересно, большинство уже ставших классическими подходов в MARL там попросту не работают без дополнительных ухищрений, что позволяет с ними не сравниваться позволяет сравнить MAMBA только с победившим на соревновании решением.

Эксперименты

В рубрике эксперименты графики говорят сами за себя.

Flatland (Метрика: Процент доехавших поездов):

SMAC (Метрика: Процент побед):

Всех победили, всё понятно.

Заключение

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

Другие работы наших магистрантов и аспирантов:

• Градиенты в нейронных сетях для поиска аномалий в данных

• Личное или социальное? Как добиться кооперации в мультиагентной среде