Decision Transformer: модель для управления роботами
Введение
Трансформер — одна из наиболее успешных моделей в машинном обучении. Именно он стоит за ChatGPT и другими большими языковыми моделями. Робототехники решили не оставаться в стороне и начали использовать трансформер для управления роботами. Одна из первых известных работ по этой теме — Decision Transformer. О ней мы и поговорим в этой статье 😊
Чтобы познакомиться с работой, для начала нам нужно вспомнить, что такое обучение с подкреплением (Reinforcement Learning; RL), в частности оффлайн RL, а также архитектура трансформера.
Коротко про обучение с подкреплением
📌Глубокое обучение с подкреплением (RL) — метод для создания агентов, принимающих решения.
Эти агенты стремятся научиться оптимальному поведению (политике / стратегии) путём взаимодействия с окружающей средой за счёт метода проб и ошибок, а также получения вознаграждений в виде обратной связи.
Поскольку RL основано на гипотезе вознаграждения, все цели могут быть описаны как максимизация ожидаемого накопленного вознаграждения (expected cumulative reward). То есть цель агента — максимизировать своё накопленное вознаграждение, называемое возвратом (return).
Агенты в RL, как и в обучении с учителем, обучаются на некоторых данных. Вопрос заключается в том, как они его собирают. В онлайн-RL агент сам собирает данные за счёт взаимодействия с окружающей средой. Затем он оперативно использует их для обучения (обновления своей политики).
Такой подход к RL предполагает, что агент либо обучается непосредственно в реальном мире, либо у нас есть симулятор. Если же симулятора нет — его необходимо создать. А этот процесс может быть сложным (реальный мир тяжело воспроизвести в симуляторе), дорогостоящим и небезопасным: агент использует недостатки симулятора, если они есть и если они помогают ему максимизировать возврат.
С другой стороны, при оффлайн-обучении с подкреплением агент использует только данные, собранные другими агентами или с помощью демонстраций человека. Он не взаимодействует с окружающей средой.
Процесс выглядит следующим образом:
- Создаётся набор данных при помощи одной или нескольких политик и / или демонстрации, собранные человеком.
- Выполняется оффлайн-RL на этом наборе данных для обучения политики.
У этого метода есть недостаток: проблема контрафактических запросов. Она возникает, если агент решает сделать что-то, для чего у нас нет данных. Например, он хочет повернуть направо на перекрёстке, но для этого у нас нет данных траектории.
На таком поверхностном уровне оффлайн-RL может показаться аналогичной обучению с демонстрацией (имитационному обучению), где мы обучаем модель воспроизводить траектории из собранного, как правило, человеком датасета. Однако это не так — оффлайн-RL не просто воспроизводит траектории, а пытается склеить их для достижения максимального возврата (см. рисунок ниже). Достигается это путём обучения на временных разностях (об этом мы подробнее рассказывали в прошлой статье).
Коротко про архитектуру трансформера
О трансформере можно написать целый пост (или даже два 🙂), поэтому сейчас мы вспомним только основную идею. Трансформер состоит из множества слоёв (само)внимания, соединённых остаточными связями (residual connections). Слой внимания принимает на вход \( n \) эмбеддингов и возвращает тоже
\( \text{выход}i = \sum{j=1, j\neq i}^n a_j v_j, \quad a_j = \text{softmax}\left(\sum q_i \cdot k_j\right) \)
Визуализация работы слоя (само)внимания показана на иллюстрации ниже:
В каузальном трансформере для расчёта \( i- \)й выхода слоя используются только эмбеддинги до \( i-1 \); иными словами, эмбеддингу не разрешается смотреть вперед (в будущее). Визуализация работы каузального трансформера на рисунке ниже:
Для детального погружения в принципы работы трансформера рекомендуем заглянуть в блогпост Елены Войты.
Decision Transformer
📌Decision Transformer (DT) предлагает использовать архитектуры моделирования последовательностей (трансформера) вместо обучения политики с помощью методов RL.
DT использует желаемое значение возврата, прошлые состояния и действия, чтобы сгенерировать будущее действие, которое в идеале должно привести к нужному возврату. Можно сказать, что DT — генеративный подход к моделированию траекторий, тогда как классический RL — дискриминативный.
Репрезентация траекторий
В RL траектории представляют собой последовательности состояний \( s \), действий \( a \) и вознаграждений \( r \):
\( \tau=(s_0,a_0,r_0,\dots,s_T, a_T, r_T) \)
где \( T \) обозначает длину траектории. Используя такое определение, мы вряд ли сможем заставить трансформер сделать что-то полезное, потому что прошлые вознаграждения нам мало говорят о будущих. Мы хотим, чтобы модель генерировала действия на основе желаемых будущих возвратов, а не прошлых вознаграждений. Поэтому мы представляем траекторию как последовательность будущих возвратов \( \hat R \), состояний \( s \) и действий \( a \):
\( \tau=(\hat R_0, s_0,a_0, \dots,\hat R_T, s_T, a_T) \)
Во время инференса мы можем задать желаемый возврат (например, 1 для успеха или 0 для неудачи), а также начальное состояние среды в качестве условия для начала генерации. После выполнения сгенерированного действия для текущего состояния мы уменьшаем целевой возврат на полученное вознаграждение (\( \hat R_{t+1} = \hat R_t — r_t \)) и повторяем процесс до завершения эпизода.
Архитектура
DT представляет собой каузальный декодерный трансформер, подобно GPT моделям. На вход модели подаются последние \( K \) кортежей \( (\hat R, s,a) \), которые представляют собой \( 3K \) токенов.
Чтобы получить эмбеддинги токенов, мы обучаем линейный слой для каждого типа модальности. Он проецирует исходные данные в размерность эмбеддинга, после чего применяется нормализация слоя. Для сред с визуальными данными состояние передаётся в свёрточный энкодер вместо линейного слоя.
Кроме того, для каждого временного шага обучается собственный эмбеддинг, который добавляется к каждому токену. В этом отличие от стандартного позиционного эмбеддинга, используемого в трансформерах, так как один временной шаг соответствует трём токенам (см. Рис. 6). Затем токены обрабатываются моделью GPT, которая предсказывает токены будущих действий с помощью авторегрессионного моделирования.
Обучение
В качестве данных для обучения выступают оффлайн-траектории (собранные человеком или другими стратегиями / политиками). Мы выбираем мини-батчи с длиной последовательности 3К из этого набора данных. Для входного токена \( s_t \) модель обучается предсказывать действие \( a_t \) либо с помощью кросс-энтропийной функции потерь для дискретных действий (например, повернуть налево или направо), либо через среднеквадратичную ошибку для непрерывных действий (например, положение устройства захвата робота). Значения функции потерь для каждого временного шага усредняются.
Авторы не обнаружили, чтобы предсказание состояний или целевых возвратов улучшало производительность. Алгоритм обучения и оценки модели для непрерывных действий показан на рисунке 6:
Сравнение с другими подходами
Авторы сравнили DT с двумя другими подходами на трёх бенчмарках — оффлайн-RL (TD learning) и обучение из демонстраций (behavior cloninig):
Обычно DT не показывает выдающихся результатов. Задачи, где DT действительно преобладает над другими методами, — это задачи долгосрочного планирования (Key-to-Door).
Эта задача основана на табличной среде и состоит из трёх последовательных фаз:
- На первой фазе агент находится в комнате с ключом.
- На второй фазе агент перемещается в пустую комнату.
- На третьей фазе агент оказывается в комнате с дверью.
Агент получает бинарное вознаграждение за достижение двери на третьей фазе, но только при условии, что он подобрал ключ на первой фазе. Эта задача сложна с точки зрения распределения ответственности, так как вознаграждение должно быть корректно перенесено от начала эпизода к его завершению с пропуском действий промежуточной фазы.
С другой стороны, такие задачи самые сложные в робототехнике, и DT — первый шаг к их успешному решению 😇
Заключение
В этой статье мы познакомились Decision Transformer — моделью для управления роботами на основе существующего датасета его траекторий.
Модель представляет собой альтенативный подход для генерации траекторий: вместо методов обучения с подкреплением она использует условную генерацию последовательностей.
DT особенно выделяется в задачах долгосрочного планирования. Если вас заинтеровала работа, и вы хотели бы поиграться, то у Hugging Face есть отличный туториал об этом 🙂
Ресурсы
