[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://t.me/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).