Decision Transformer: Reinforcement Learning via Sequence Modeling
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Michael Laskin, Pieter Abbeel, Aravind Srinivas, Igor Mordatch
Статья: https://arxiv.org/abs/2106.01345
Code: https://github.com/kzl/decision-transformer
Разбор от Янника: https://www.youtube.com/watch?v=-buULmf7dec

Работа от примерно той же группы авторов, что и в предыдущем посте про “Pretrained Transformers as Universal Computation Engines” (https://t.me/gonzo_ML/711). Тоже по сути про универсальность трансформеров для решения разных задач, теперь для RL.

Здесь авторы подошли к задаче offline reinforcement learning и моделируют её как задачу conditional sequence modeling, используя авторегрессионный трансформер, аналогичный GPT.

Offline RL по сути несколько специфическая область, когда агент не имеет возможности итеративно взаимодействовать со средой, и вынужден довольствоваться как-то полученным (возможно каким-то совершенно другим неизвестным агентом) накопленным и сохранённым опытом в виде массива наблюдений.

Трансформеры уже применялись в RL (и мы даже писали про них, например здесь: https://t.me/gonzo_ML/148), но обычно они встраиваются в традиционные RL алгоритмы, например, обучая policy через TD learning (в упомянутой выше работе использовалась on-policy адаптация метода Maximum a Posteriori Policy Optimization, называющаяся V-MPO).

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

На вход GPT-like модели подаётся траектория, закодированная последовательностью троек (R, s, a), где a — это действие, s — состояние среды, а R в данном случае не полученный при этой комбинации (s,a) reward, а return-to-go (сумма будущих вознаграждений), и это сделано для того, чтобы моделировать будущие действия не на основании прошлых вознаграждений, а на основании будущего желаемого возврата.

Для получения следующего действия на вход прилетает вся прошлая траектория + текущие значения желаемого возврата R_t и текущее состояние среды s_t, и для позиции этого s_t предсказывается действие a_t. То есть для K шагов получается 3K токенов (я бы даже сказал 3K-1).

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

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

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

В общем всё довольно прямолинейно. Обучают на датасете траекторий, сгенерированных какими-то полиси (в каждом эксперименте своими, есть варианты).

Оцениваются на задачах дискретного и непрерывного управления (соответственно несколько игр из бенчмарка Атари и бенчмарк D4RL в OpenAI Gym), сравнивают с model-free методами TD learning (Conservative Q-Learning, CQL, который на тот момент вроде как был SoTA, и с некоторыми другими), а также с imitation learning.

Для Атари это трансформер размерности 128 с 6 слоями и 8 головами, для Gym всего трёхслойный с одной головой.

На Атари метод сравним с CQL и бьёт более старые алгоритмы. На задачах непрерывного управления из D4RL почти на всех задачах Decision Transformer побил традиционные RL методы.

Авторы пытаются анализировать, не занимается ли Decision Transformer behaviour cloning’ом и показывают, что он более эффективен, чем просто imitation learning на подмножестве датасета.

Отдельно изучают, важна ли глубина контекста, и да, на контексте длины 1 трансформер работает существенно хуже.