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://news.1rj.ru/str/gonzo_ML/711). Тоже по сути про универсальность трансформеров для решения разных задач, теперь для RL.
Здесь авторы подошли к задаче offline reinforcement learning и моделируют её как задачу conditional sequence modeling, используя авторегрессионный трансформер, аналогичный GPT.
Offline RL по сути несколько специфическая область, когда агент не имеет возможности итеративно взаимодействовать со средой, и вынужден довольствоваться как-то полученным (возможно каким-то совершенно другим неизвестным агентом) накопленным и сохранённым опытом в виде массива наблюдений.
Трансформеры уже применялись в RL (и мы даже писали про них, например здесь: https://news.1rj.ru/str/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 трансформер работает существенно хуже.
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://news.1rj.ru/str/gonzo_ML/711). Тоже по сути про универсальность трансформеров для решения разных задач, теперь для RL.
Здесь авторы подошли к задаче offline reinforcement learning и моделируют её как задачу conditional sequence modeling, используя авторегрессионный трансформер, аналогичный GPT.
Offline RL по сути несколько специфическая область, когда агент не имеет возможности итеративно взаимодействовать со средой, и вынужден довольствоваться как-то полученным (возможно каким-то совершенно другим неизвестным агентом) накопленным и сохранённым опытом в виде массива наблюдений.
Трансформеры уже применялись в RL (и мы даже писали про них, например здесь: https://news.1rj.ru/str/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 трансформер работает существенно хуже.
GitHub
GitHub - kzl/decision-transformer: Official codebase for Decision Transformer: Reinforcement Learning via Sequence Modeling.
Official codebase for Decision Transformer: Reinforcement Learning via Sequence Modeling. - kzl/decision-transformer
👍1
И в целом метод неплохо выучивает полиси, где надо делать long-term credit assignment, даже если обучается на рандомных полиси. В ситуации с разреженными и отложенными наградами также работает достойно.
Такие вот дела. Были когда-то прикольные работы, где GPT играла в Го (https://arxiv.org/abs/2007.03500) или шахматы (https://arxiv.org/abs/2008.04057), а здесь тот же универсальный трансформер по типу GPT решает более-менее любые RL задачи. Интересно, что будет, когда обучат какую-нибудь гигантскую модель на пиллиарды параметров.
Такие вот дела. Были когда-то прикольные работы, где GPT играла в Го (https://arxiv.org/abs/2007.03500) или шахматы (https://arxiv.org/abs/2008.04057), а здесь тот же универсальный трансформер по типу GPT решает более-менее любые RL задачи. Интересно, что будет, когда обучат какую-нибудь гигантскую модель на пиллиарды параметров.
arXiv.org
The Go Transformer: Natural Language Modeling for Game Play
This work applies natural language modeling to generate plausible strategic moves in the ancient game of Go. We train the Generative Pretrained Transformer (GPT-2) to mimic the style of Go...
Online RL, Off-policy RL и Offline RL. Отсюда: https://arxiv.org/abs/2005.01643
[Trajectory Transformer] Offline Reinforcement Learning as One Big Sequence Modeling Problem
Michael Janner, Qiyang Li, Sergey Levine
Статья: https://arxiv.org/abs/2106.02039
Code: https://github.com/JannerM/trajectory-transformer
Ещё одна работа про замену компонентов reinforcement learning (RL) алгоритмов на стандартные решения для sequence learning, а именно трансформерные блоки декодера для авторегрессионной генерации (всё как в GPT). В предыдущей работе про Decision Transformer (DT) (https://news.1rj.ru/str/gonzo_ML/719) использовали reward conditioning, а в этой, где решение называют Trajectory Transformer (TT), делают планирование на основе beam search. Получается как бы довольно простой идейно model-based алгоритм.
В обоих решениях это довольно заметный отход от работы с MDP (Markov Decision Process), потому что марковское свойство больше не действует, генерация происходит с учётом всей предыдущей траектории, что может быть полезно.
Оцениваются также на задачах Offline RL, чтобы можно было использовать большие объёмы предыдущих накопленных интеракций.
Траектория длины T задаётся как последовательность токенов, описывающая T состояний (размерности N), действий (размерности M) и скалярных вознаграждений, получается итоговая длина T*(N+M+1).
Для работы с непрерывными состояниями и действиями делают дискретизацию, пробуют два варианта — равномерную (фиксированный шаг разбиения пространства) и квантильную (каждому токену достаётся одинаковая масса вероятности).
Архитектура — стандартный декодер по типу GPT, 4 слоя, 4 головы внимания, 128 размерность. Beam width 256, горизонт планирования 15, контекст размера 5 (5 предыдущих кортежей (s,a,r,R)) плюс ещё несколько тонких настроек. Обучаются с обычным teacher forcing, лосс по факту стандартный авторегрессионный, только из-за многомерности состояний и действий, каждое следующее измерение авторегрессионно зависит от предыдущих.
Проверяют в трёх типах экспериментов:
1) Imitation learning, когда надо выучить распределение траекторий в обучающей выборке и оптимизируется вероятность конкретной траектории. Тут всё как в обычном sequence modeling.
2) Goal-conditioned RL, когда задаётся некий будущий контекст, но это в целом тоже сводится к предыдущей задаче, когда в начало траектории можно добавить целевое состояния. А-ля prompting такой.
3) Offline RL, где (лог)вероятности переходов заменяются на сигнал подкрепления и задача модели — максимизировать reward, а чтобы уменьшить риск близорукости beam search в такой постановке, в траектории для каждого временного отсчёта добавляется reward-to-go (как в общем и в предыдущем DT).
В экспериментах показывают, что модель хорошо генерит траектории и на примере предсказания траектории гуманоида на 100 шагов вперёд видно, что траектория стабильна и близка к референсу, чего не скажешь о традиционном методе, который делает rollout полиси шаг за шагом и в котором ошибки быстро накапливаются.
В работе сделали вариант Markovian Transformer с глубиной контекста всего 1 и он оказался примерно такого же качества на fully-observed средах (что интересно), но хуже на partially-observed.
При визуализации внимания выделяются два паттерна: 1) модель выучила марковскую стратегию (когда смотрит преимущественно на предыдущий шаг), и 2) модель смотрит на конкретные измерения множества предыдущих состояний или действий.
В задачах Offline RL выбрали несколько сред из набора в D4RL, показали, что Trajectory Transformer (TT) лучше других методов типа BRAC или CQL (которые вроде как SoTA), а также лучше недавнего Decision Transformer (DT).
Michael Janner, Qiyang Li, Sergey Levine
Статья: https://arxiv.org/abs/2106.02039
Code: https://github.com/JannerM/trajectory-transformer
Ещё одна работа про замену компонентов reinforcement learning (RL) алгоритмов на стандартные решения для sequence learning, а именно трансформерные блоки декодера для авторегрессионной генерации (всё как в GPT). В предыдущей работе про Decision Transformer (DT) (https://news.1rj.ru/str/gonzo_ML/719) использовали reward conditioning, а в этой, где решение называют Trajectory Transformer (TT), делают планирование на основе beam search. Получается как бы довольно простой идейно model-based алгоритм.
В обоих решениях это довольно заметный отход от работы с MDP (Markov Decision Process), потому что марковское свойство больше не действует, генерация происходит с учётом всей предыдущей траектории, что может быть полезно.
Оцениваются также на задачах Offline RL, чтобы можно было использовать большие объёмы предыдущих накопленных интеракций.
Траектория длины T задаётся как последовательность токенов, описывающая T состояний (размерности N), действий (размерности M) и скалярных вознаграждений, получается итоговая длина T*(N+M+1).
Для работы с непрерывными состояниями и действиями делают дискретизацию, пробуют два варианта — равномерную (фиксированный шаг разбиения пространства) и квантильную (каждому токену достаётся одинаковая масса вероятности).
Архитектура — стандартный декодер по типу GPT, 4 слоя, 4 головы внимания, 128 размерность. Beam width 256, горизонт планирования 15, контекст размера 5 (5 предыдущих кортежей (s,a,r,R)) плюс ещё несколько тонких настроек. Обучаются с обычным teacher forcing, лосс по факту стандартный авторегрессионный, только из-за многомерности состояний и действий, каждое следующее измерение авторегрессионно зависит от предыдущих.
Проверяют в трёх типах экспериментов:
1) Imitation learning, когда надо выучить распределение траекторий в обучающей выборке и оптимизируется вероятность конкретной траектории. Тут всё как в обычном sequence modeling.
2) Goal-conditioned RL, когда задаётся некий будущий контекст, но это в целом тоже сводится к предыдущей задаче, когда в начало траектории можно добавить целевое состояния. А-ля prompting такой.
3) Offline RL, где (лог)вероятности переходов заменяются на сигнал подкрепления и задача модели — максимизировать reward, а чтобы уменьшить риск близорукости beam search в такой постановке, в траектории для каждого временного отсчёта добавляется reward-to-go (как в общем и в предыдущем DT).
В экспериментах показывают, что модель хорошо генерит траектории и на примере предсказания траектории гуманоида на 100 шагов вперёд видно, что траектория стабильна и близка к референсу, чего не скажешь о традиционном методе, который делает rollout полиси шаг за шагом и в котором ошибки быстро накапливаются.
В работе сделали вариант Markovian Transformer с глубиной контекста всего 1 и он оказался примерно такого же качества на fully-observed средах (что интересно), но хуже на partially-observed.
При визуализации внимания выделяются два паттерна: 1) модель выучила марковскую стратегию (когда смотрит преимущественно на предыдущий шаг), и 2) модель смотрит на конкретные измерения множества предыдущих состояний или действий.
В задачах Offline RL выбрали несколько сред из набора в D4RL, показали, что Trajectory Transformer (TT) лучше других методов типа BRAC или CQL (которые вроде как SoTA), а также лучше недавнего Decision Transformer (DT).
GitHub
GitHub - jannerm/trajectory-transformer: Code for the paper "Offline Reinforcement Learning as One Big Sequence Modeling Problem"
Code for the paper "Offline Reinforcement Learning as One Big Sequence Modeling Problem" - GitHub - jannerm/trajectory-transformer: Code for the paper "Offline Reinforcem...
Со сравнением с DT непросто. Во-первых, в работе есть два варианта TT, с равномерной и квантильной дискретизацией. Первый хуже DT, второй лучше. Во-вторых, непонятно, откуда цифры для DT, сами имплементировали или взяли из оригинальной статьи. По цифрам получается, что взяли из статьи, но при этом дропнули один из экспериментов, а может с ним DT обогнал бы? В-третьих, с какой именно архитектурой сравнивались. Здесь в работе был 4-слойный 4-головый трансформер, а в работе про DT, если они взяли цифры из неё, был использован 3-слойный 1-головый, что как бы странно сравнивать, ибо при такой небольшой разнице в цифрах почти наверняка можно дотюнить DT до более высокого качества. В общем меня лично превосходство TT над DT не убеждает и по мне такое сравнение не особо честное.
Также в работе попробовали комбинацию с Q-функцией из другого метода, чтобы заменить ею Monte Carlo оценку, которая на длинных горизонтах не очень информативна. Q-guided TT оказался лучше альтернатив.
В отдельном эксперименте попробовали goal conditioning, когда будущее целевое состояние также передаётся в модель (они это называют anti-causal conditioning) и beam search по факту используется как метод достижения цели. Работает.
Такие дела. Интересное в целом направление про замену различных классических частей RL на универсальную архитектуру и универсальную задачу (авторегрессионное моделирование). Попахивает очередным подкреплением пресловутого The Bitter Lesson Саттона (http://www.incompleteideas.net/IncIdeas/BitterLesson.html).
Также в работе попробовали комбинацию с Q-функцией из другого метода, чтобы заменить ею Monte Carlo оценку, которая на длинных горизонтах не очень информативна. Q-guided TT оказался лучше альтернатив.
В отдельном эксперименте попробовали goal conditioning, когда будущее целевое состояние также передаётся в модель (они это называют anti-causal conditioning) и beam search по факту используется как метод достижения цели. Работает.
Такие дела. Интересное в целом направление про замену различных классических частей RL на универсальную архитектуру и универсальную задачу (авторегрессионное моделирование). Попахивает очередным подкреплением пресловутого The Bitter Lesson Саттона (http://www.incompleteideas.net/IncIdeas/BitterLesson.html).
🔥1