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://hottg.com/gonzo_ML/711). Тоже по сути про универсальность трансформеров для решения разных задач, теперь для RL.
Здесь авторы подошли к задаче offline reinforcement learning и моделируют её как задачу conditional sequence modeling, используя авторегрессионный трансформер, аналогичный GPT.
Offline RL по сути несколько специфическая область, когда агент не имеет возможности итеративно взаимодействовать со средой, и вынужден довольствоваться как-то полученным (возможно каким-то совершенно другим неизвестным агентом) накопленным и сохранённым опытом в виде массива наблюдений.
Трансформеры уже применялись в RL (и мы даже писали про них, например здесь: https://hottg.com/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 трансформер работает существенно хуже.
>>Click here to continue<<
