XLand-100B: первый в мире большой датасет для контекстного обучения с подкреплением
Хабр, привет! Меня зовут Александр Никулин, я аспирант МФТИ и один из исследователей научной группы «Адаптивные агенты» в Институте AIRI.
Как можно понять из названия, наша группа заинтересована в создании адаптивных агентов, способных обобщаться на новые задачи после обучения. Направление это относительно новое и в литературе именуется как контекстное обучение с подкреплением (далее in-context RL). И мы активно двигаем его вперед! Совсем недавно выпустили две статьи, обе приняты на ICML 2024, а ещё среду на JAX со множеством задач для мета-обучения. Мы обязательно расскажем о них чуть позже (подписывайтесь!), а в этой статье хочется затронуть наш недавний препринт. В нем мы представили и выложили в open-source огромный (по меркам RL) и пока единственный датасет для in-context RL. На сбор траекторий для 40к задач и 130B транзиций потребовалось 50000 GPU-часов. Эту работу мы проделали совместно с коллегами из лаборатории T-Bank AI Research.
Датасетом уже можно пользоваться, так что рассказываем и надеемся на будущий акцепт статьи! Ну, а начнем чуть издалека, расскажу, что такое in-context learning, как он появился в RL и почему нам понадобился собственный датасет.
Контекстное и мета-обучение
На самом деле, сказав, что область новая, я чуть чуть схитрил! Конечно же, идея обучения мета‑алгоритмов, то есть алгоритмов, которые сами по себе представляют алгоритмы обучения, появилась относительно давно. В отличие от мультизадачного обучения, мы хотим получить модель, не просто хорошо работающую на всех тренировочных задачах одновременно, а способную максимально эффективно дообучаться на новых за минимальное время. Мета‑обучение существует и в более традиционной области обучения с учителем и в обучении с подкреплением. Для интересующихся: о краткой истории мета‑обучения и основных методах можно узнать подробнее в блогпосте Lilian Weng из Open AI или в недавнем большом обзоре.
Схематическая разница междумультизадачным и мета-обучением. Источник: meta-world.github.io
Однако под новизной я имел в виду способ получения мета‑моделей, набирающий большую популярность в последнее время. Думаю, всем, кто игрался с большими языковыми моделями, например ChatGPT, он знаком, даже если вы не знаете формального названия.
Пример контекстного обучения. Вначале в промпт подаются примеры решения задачи, а дальше модель справляется сама.
Наверное, вы слышали, а может даже пробовали сами, что ChatGPT можно научить новой задаче, просто подав в контекст промпт с некоторым количеством примеров и решений для нее. Для широкого спектра задач это сработает. И это удивительно, ведь никто не учил модель специально решать вашу задачу, она сама разобралась на основе контекста и данных ей примеров как правильно отвечать на новых данных. Данный феномен называется контекстное обучение (in‑context learning или ICL) и по сути своей, является и примером мета‑обучения тоже!
Феномен контекстного обучения был впервые во всей красе продемонстрирован с выходом GPT-3, после чего некоторое время научное сообщество думало, что это свойство исключительно огромных языковых моделей на основе архитектуры трансформер. Однако после было обнаружено, что на него способны и совсем маленькие трансформеры, не трансформеры вовсе, и, в целом, явление это не уникальное для языка, а возникает и в других доменах, например в картиночных генеративных моделях. Долгое время основным ограничением был размер контекста из‑за квадратичности механизма внимания и, как следствие, дороговизны длинных промптов. С появлением миллионных контекстов стало понятно, что контекстное обучение часто гораздо проще и эффективнее более привычного дообучения под конкретную задачу.
Контекстное обучение с подкреплением
Концептуально, контекстное обучение очень хорошо ложится на RL (см. контекстуальный марковский процесс принятия решений). Учитывая, как иногда анекдотически трудно бывает с нуля обучить агента даже на нескольких задачах с помощью обучения с подкреплением, возможность получить агента, способного к быстрой и эффективной адаптации на новых задачах, крайне привлекательна. В конце концов, разве не в этом наша конечная цель — получить аналог базовой (foundation) модели для RL?
Возможный потенциал такого подхода показал Ada от DeepMind, вышедший в прошлом году. Оказалось, что при должном размахе обучения, сложности и разнообразия сред можно получить агента, способного адаптироваться к новым нетривиальным задачам со скоростью равной человеческой, а то и быстрее. Очень советую глянуть видео с презентацией статьи!
Несмотря на успех, Ada обучался с помощью более‑менее традиционного алгоритма из RL, наследуя тем самым все нестабильности обучения и сложности с масштабированием. Можно ли проще?
На самом деле, да. Ведь, если задуматься, в обучении с подкреплением у нас есть траектории взаимодействия агента со средой — суть дискретные последовательности. Да, они состоят из трех модальностей (состояние, действие, награда), но глобально мало чем отличаются от языка с точки зрения трансформера. Можно подавать их в трансформер и предсказывать следующее действие.
Однако так агент не научится максимизировать награду, только имитировать политику в датасете. Можно обучать, обуславливаясь на итоговую суммарную награду в текущей траектории, а во время тестирования промптить наибольшей возможной. Получится Decision Transformer (DT), наверное, самый популярный сейчас метод на основе трансформера в RL.
Визуализация подхода DT. Обратите внимание, что контекст начинается с будущей суммарной наградой. Источник: Lili Chen et al. / arXiv, 2021
Но для контекстного обучения DT подходит плохо, и на самом деле таких способностей у него замечено не было. Во‑первых, разброс суммарной награды может значительно различаться от задачи к задаче, или мы в целом можем заранее его не знать. А во‑вторых, мы не учим трансформер максимизировать награду явно.
Представим, что у нас есть полные истории обучения каких‑то базовых RLных алгоритмов. Награда в них, очевидно, растет с течением обучения. Будем обучать на этом трансформер, подавая траектории как есть, то есть, не меняя порядка и не ломая упорядоченность по награде. Как думаете, что произойдет? А произойдет следующее.
При маленьком размере контекста у нас получится обычный behavioral cloning (BC), так как трансформер будет просто пытаться предсказать действие политики в датасете на основе краткой истории. Однако если мы попробуем расширить контекст достаточно, чтобы вобрать в себя историю улучшения базового алгоритма, например четыре или больше эпизодов, то внезапно трансформер после обучения приобретет in‑context learning свойства! А именно: при тестировании на новой задаче он на основе своего собственного контекста начнет максимизировать награду и решать ранее не виденные задачи (но из схожего с тренировочным распределения). Более того, начнет решать лучше, чем эксперт в данных.
Почему так происходит? У трансформера не остается выбора. Для того, чтобы минимизировать функцию потерь ему придется «дистиллировать» базовый алгоритм, собравший данные. Ключевое здесь, что именно алгоритм улучшения, а не какую‑то одну конкретную усредненную политику. Таким образом, при наличии подходящих данных, у нас получился крайне простой и стабильный метод для обучения мета‑алгоритма обучения с подкреплением, наследующий все открытия для эффективного обучения трансформеров из NLP области. Первым так обучать предложила статья от DeepMind, назвав это Algorithm Distillation (AD). С нее фактически и началась in‑context RL область.
Визуализация подхода AD. В отличие от DT, в контекст подаются состояния, действия и сырые награды в естественном порядке. Источник: Michael Laskin et al. / arXiv, 2022
Любопытно: дальнейший теоретический анализ показал, что обученный таким способом трансформер может реализовывать целые оптимальные RL алгоритмы внутри себя. Причем алгоритмы разнообразные, от простых Thompson Sampling, Lin‑UCB, до многих вариаций TD‑learning, лежащего в основе современных DQN‑подобных алгоритмов. Реализуют ли они их на самом деле, и какие конкретно — вопрос открытый. Теоретический анализ лишь доказывает, что они на это способны. Если вам интересна тема mechanistic interpretability, то вы можете попытать удачу и попробовать декодировать выученные трансформером алгоритмы. Например, недавно так обнаружили, что трансформеры неявно выучивают скрытые модели мира и планируют действия. Так что тема крайне увлекательная!
Проблема контекстного обучения в RL
Только заинтересовавшись in‑context RL, мы быстро поняли что ей не хватает сложного и интересного бенчмарка, а также больших датасетов. Все существующие работы проверяли свои методы на крайне простых средах и были вынуждены каждый раз пересобирать датасеты с нуля, т.к. не существовало стандартизации. У нас, как исследователей, это быстро начало вызывать фрустрацию, ибо в каждой статье среды и датасеты были похожи внешне, но обязательно находились незаметные различия (о которых авторы статей забывали упомянуть), в итоге влияющие на результаты. Мы решили что должны это исправить, и, кажется, у нас получилось!
Преждем чем тратить время и ресурсы на сбор нового датасета, мы, конечно же, посмотрели на все существующие, вдруг их можно приспособить для наших нужд? Оказалось, что из всех нам известных датасетов (особенно крупных) ни один не подходит под требования in‑context RL: не содержит полные истории обучения базовых алгоритмов и не обладает высоким разнообразием задач. Если мы хотим изучать контекстное обучение со всех сторон, например scaling laws по количеству и сложности уникальных задач, а не просто токенов, как в NLP, нам нужен по настоящему разнообразный и большой датасет, содержащий тысячи задач.
Сравнение нашего датасета с существующими. Ни один не позволяет изучать контекстное обучение.
Было решено, что нужно собирать самим. Но где взять среду с множеством уникальных и нетривиальных задач, достаточно быструю, чтобы осилить обучения тысяч RLных агентов с нуля? По счастливой случайности, недавно мы выпустили такую среду сами! Назвали ее XLand‑MiniGrid, т.к. она вдохновлена и берет лучшее от XLand от DeepMind и от широко известной в комьюнити минималистичной MiniGrid. Она представляет собой grid‑world с разреженной наградой, некоторым набором целей и правил, способных динамически изменять окружение. Задачи могут генерироваться процедурно в виде бинарных деревьев.
Визуализация окружения в XLand-MiniGrid. Показан краткий пример действия одного правила, преобразующего два конкретных объекта в новый.
Пример одной задачи в XLand-MiniGrid. В корне дерева находится цель, которую нужно достигнуть. Любая другая вершина задает правила, с помощью которых можно получать новые необходимые для цели объекты из уже существующих.
Представьте, что каждый раз, когда вы заходите в Minecraft, дерево крафта рандомизируется, и вас просят добыть какой‑то сложный предмет. Вам предстоит увлекательная задача с помощью экспериментов и хорошей памяти открыть это дерево крафта самостоятельно. Более того, решив одну такую задачу, вы не сможете генерализоваться на новую (ведь новое дерево крафта скрыто от вас), вам придется начинать все сначала, то есть быть адаптивным — выучить сам процесс обучения и эффективного исследования среды.
Благодаря тому, что среда с нуля написана на JAX, она поддерживает GPU и TPU, а, значит, мы можем обучаться на десятках тысяч параллельных сред одновременно, полностью избегая CPU-боттлнека, присущего большинству пайплайнов обучения в RL.
Сбор датасета
Для сбора датасета мы выбрали medium-1m и trivial-1m — наборы задач из заранее сгенерированных в XLand‑MiniGrid. Один самый простой для датасета поменьше, чтобы быстрее итерироваться, и один более сложный, для основного датасета. Мы заранее выбрали суммарно 40к задач различной сложности из этих наборов. В качестве базового алгоритма взяли наиболее известный и стабильный PPO, добавив GRU в качестве модуля памяти. На каждой задаче обучали на множестве сред и GPU параллельно. Во время обучения собирали все логи с первых 32 сред, чтобы не раздувать размер датасета слишком сильно.
К сожалению, первая попытка сбора провалилась. Оказалось, что большая часть задач слишком сложная для базового PPO+RNN даже если обучать на миллиардах транзиций. Награда слишком разреженная и шансы случайно решить минимальны. Пробовали разные способы улучшить исследование среды агентом, но без особого успеха.
В итоге сошлись к сбору в две фазы. Изначально правила и цель скрыты от агента. Мы слегка схитрили и во время первой фазы сделали их открытыми и подавали агенту как дополнительный вход (goal‑conditioning), обучая сразу на множестве различных задач. Такой агент после обучения мог неплохо решать новые задачи в zero‑shot режиме на основе информации о задаче, поданной на вход. Далее, для какой‑то конкретной задачи мы дообучали агента, полностью маскируя всю информацию о правилах и целях. По сути, после предобучения и маскирования мы получали политику, усредненную по всем задачам, что значительно улучшило способности исследования среды у агента и ускорило сходимость на сложных задачах.
Награда в течение обучения усредненная по множеству задач. Как видно, после предобучения агент достигает гораздо лучшего результата в среднем.
Сбор данных занял в сумме две недели. На сбор траекторий для 40к задач и 130B транзиций потребовалось 50 000 GPU‑часов. После сбора мы отфильтровали некоторое количество задач, если агент так и не смог решить их, или во время сбора на них произошел сбой по той или иной причине. Данные хранятся в формате hdf5 с включенной компрессией. Мы немного поигрались с настройками компрессии и размером кэша, сжав датасет в 15 раз и всего в два раза замедлив сэмплирование. Итоговые статистики можно увидеть в таблице ниже:
Валидация датасета
Наконец, после сбора надо было удостовериться, что получившийся датасет отвечает всем нужным нам требованиям. Мы смотрели на две вещи.
Во‑первых, мы построили графики обучающих историй, увидев, что траектории естественным образом упорядочены по возрастанию суммарной награды за эпизод, как и должно быть.
Во‑вторых, мы воспроизвели и обучили Algorithm Distillation (AD) и Decision‑Pretrained Transformer (DPT), два наиболее популярных метода для in‑context RL, чтобы продемонстрировать, что обучаясь на наших данных агент на самом деле проявляет способности к контекстному обучению. Выяснили интересное.
Оказалось, что DPT не умеет из коробки работать со средами где не выполняется марковское свойство, что ранее упоминалось лишь вскользь в оригинальной работе. Касательно AD: после обучения он смог решать простые задачи и показывать признаки контекстного обучения на чуть более сложных, однако далеко от идеала. Вероятно, его нужно учить дольше, а может, нужны новые, более эффективные идеи?
Визуализация тестовой награды AD. На простых задачах контекстное обучение быстро улучшает награду (около 50 эпизодов). Однако на более сложных оно либо пока что не наблюдается, либо быстро выходит на плато далеко от оптимальной награды.
Итог
Как видите, работы предстоит еще много! На данный момент датасет уникален в своей нише, как по размеру, так и по своим свойствам. Мы планируем проверить еще много гипотез на его основе и, я уверен, значительно поднять state‑of‑the‑art в in‑context RL области. Надеемся, он окажется полезным и вам, ведь его можно использовать и для многих других областей тоже, например offline RL.
В любом случае, я надеюсь, что помимо самого датасета, неформальное и краткое введение в эту новую область поведало вам что‑то новое и интересное. Оставайтесь на связи и подписывайтесь на каналы AIRI и нашей команды. Обязательно ставьте звездочки датасету и среде на GitHub. Если вас заинтересовала тема или вы в целом интересуетесь RL и научной работой, хотите попробовать себя в этом, то мы всегда рады с вами пообщаться. Пишите мне или руководителю группы Владу Куренкову.
