[정리] Trajectory Transformer: Reinforcement Learning as One Big Sequence Modeling Problem (Michael Janner, NeurIPS 2021 spotlight)

Author : Michael Janner, Qiyang Li, Sergey Levine
Paper Link (Arxiv버전) : https://arxiv.org/abs/2106.02039

Paper Link (NeurIPS버전) : https://trajectory-transformer.github.io/trajectory-transformer-neurips-2021.pdf

Website : https://trajectory-transformer.github.io/

Code: https://github.com/JannerM/trajectory-transformer

 

Trajectory Transformer

Single-Step Model

Long-horizon predictions of the Trajectory Transformer compared to those of a feedforward single-step dynamics model.

 

 

0. Abstract

• 기존 RL은 factorized single-step의 Markov property를 중요시 여기는 프레임 워크임
• 하지만 RL을 "one big sequence modeling"문제로 본다면, NLP에서 이미 상당한 성능을 보인 Transformer와 같은 high-capacity sequence prediction 모델을 그대로 가져다 쓸 수 있음.
• 이렇게 할 경우 기존의 offline model-free RL에서 필요로 했던 constraint나 uncertainty estimator가 필요 없어져 model-based RL과 같은 장점이 생김.
• RL에 대한 이러한 재접근이 다양한 문제(long-horizon dynamics prediction, imitation learning, goal-conditioned RL, offline RL)에서 효과를 보인다는것을 검증함.

1. Introduction

• 기존 RL은 long-horizon문제를 Markv property에 따라 single-step subproblem으로 나누어 최적화하는것을 기본 원리로 하며 이는 model-free RL의 Q-learning, model-based RL의 single-step prediction으로 다루어짐.
• Actor-critic, model-based, offline RL은 서로 다른 분포를 학습하지만 state, action, reward를 단순한 stream of data로 본다면 모두 single sequence model로 다룰 수 있으며, high-capacity를 가진 sequence model 아키텍처를 그대로 적용하여 GPT-3와 같은 scalability의 장점을 얻을 수 있을것.
• "RL을 sequence generation 문제로 본다면 기존의 unsupervised sequence model을 적용하여 새로운 RL알고리즘을 만들 수 있을까?" 에대한 분석을 하고자 Transformer 아키텍처를 사용하여 state, action, reward의 sequence를 다루고 planning 알고리즘으론 beam search를 사용하려 함.
• 이러한 접근을 Trajectory Transformer 라고 이름 짓고, offline RL, model-baed imitation learning에 적용해보고자함

2. Recent Work

• LSTM, seq2seq, 그리고 Transformer 아키텍처까지 빠르게 발전한 sequence 모델을 RL에 적용한 경우가 다양하나 모두 RL의 업데이트 알고리즘은 그대로 사용한 반면, 여기선 가능한 많은 RL파이프라인(특히 알고리즘)을 sequence model의 representation capacity로 대체하고자 함.
• RL에서는 predictive 모델 (for model-based RL) / behaviror policy (for imitation learning) / behavior constraint (for offline RL) 등의 다양한 분포들을 학습하는 접근이 있으나, single high-capacity sequence model로 하여금 state / action / reward sequence들을 interchangeable하게 다루는 joint distribution을 학습하도록 하면 이러한 분포들의 역할을 모두 수행 할 수 있을것.
• Trajectory Transformer는 학습된 모델로 planning을 하는 model-based RL과 가장 유사하지만 ensemble이 필요없으며 특히나 offline setting에서 conservatism이나 pessimism 구조가 명시적으론 필요없다는 점이 다른데, 이는 state와 action을 jointly모델링 하는데서 in-distribution action을 생성하도록 함축적인 bias가 들어가기 때문.
• 이 논문이 발표되기 하루전에 피터아빌 팀에서 Decision Transformer를 발표했으며, RL의 알고리즘 없이도 high-capacity sequence model만으로 RL문제를 다루고 그 가능성을 입증한다는 점에서 본 논문과 본질적으로 추구하는 컨셉이 같음.

 

3. Reinforcement Learning and Control as Sequence Modeling

• 본 논문에서 제안하는 접근을 Trajectory Transformer라고 새로 명명 했으나 구현 측면에선 NLP에서의 sequence 및 탐색 모델과 거의 동일하므로, 아키텍처 부분보다 trajectory 데이터를 어떻게 다루었는지에 더 집중하고자 함.

3.1 Trajectory Transformers

• 제안하는 접근방식의 핵심은 trajectory를 Transformer 아키텍처를 위한 비정형화된 sequence로 다루는것으로,  trajectory $\tau$는 아래와 같이 $N$차원 state, $M$차원 action, 스칼라 reward로 구성됨

$\tau=\left \{s_{t}^{0}, s_{t}^{1}, \cdots,s_{t}^{N-1},a_{t}^{0}, a_{t}^{1}, \cdots,a_{t}^{M-1},r_{t}\right \}_{t=0}^{T-1}$

• 이때 $i$는 $i$번째 차원을 의미하며, $t$는 시간으로 0부터 시간 $T$까지의 trajectory.
• Continuous space일 경우 discretize를 해줌.
• NLP와 마찬가지로 tokenization을 해주는데, state의 각 차원을 offset을 주어 토큰화 하여 다음과 같음.

$\bar{s}_{t}^{i}=\left \lfloor V\frac{s_{t}^{i}-l^{i}}{r^{i}-l^{i}} \right \rfloor + Vi$

• 여기서 $\left \lfloor \cdot  \right \rfloor$는 floor function이며, $V$는 차원 토큰의 크기.
• 같은 방법으로 action 토큰 $\bar{a}_{t}^{j}$도 차원별로 offset $V\times \left ( N + j \right )$을 주어 tonkenize 하며 discretized reward 토큰 $\bar{a}_{t}^{j}$역시 offset $V\times \left ( N + M \right )$로 tokenize함.
• Gaussian transition과 같은 단순화 가정 없이 이렇게 개별로 tokenize함으로서 trajectory에대한 분포를 더욱 expressive하게 만듬 
• Trajectory Transformer는 GPT 아키텍처의 Transformer decoder를 차용했으나, 상대적으로 작은 크기의 4개 레이어와 6개 self-attention head로 구성됨
• Trajectory Transformer의 파라메터를 $\theta$라고 할때 학습 objective는 다음과 같고 이때 정답을 input으로 주는 teacher-forcing 방법을 사용함

$\mathcal{L}(\bar{\tau})=\sum_{t=0}^{T-1}\left ( \sum_{i=0}^{N-1}\log P_{\theta}\left (\bar{s}_{t}^{i} \mid 
 \bar{s}_{t}^{<i}, \bar{\tau}_{<t}\right ) + \sum_{j=0}^{M-1}\log P_{\theta}\left ( \bar{a}_{t}^{j} \mid 
 \bar{a}_{t}^{<j}, \bar{s}_{t}, \bar{\tau}_{<t} \right ) +\log P_{\theta}\left ( \bar{r}_{t} \mid \bar{a}_{t}, \bar{s}_{t}, \bar{\tau}_{<t} \right ) \right )$

• $\bar{\tau}_{<t}$는 시간 t까지의 tokenized trajectory이나 self-attention의 quadratic complexity때문에 512개의 토큰으로 제한을 두기로 하고 총 $\frac{512}{N+M+1}$의 horizon을 사용함.

 

3.2 Transformer Trajectory Optimization

• 제안한 Trajectory Transformer를 control문제에 어떻게 적용하지 다뤄보고자 하며 기본 NLP알고리즘에서 추가되는 변형이 많아지는 순서로 다음 세가지 세팅이 있음
  1. Imitation learning
  2. Goal-conditioned RL
  3. Offline RL
• 이와같은 Trajectory Transformer의 control 문제에서의 변형들을 아울러 Transformer Trajectory Optimization (TTO)라고 정의.

3.2.1 Imitation learning

• Imitation learning과 같이 task의 목적이 학습데이터의 distribution을 복원하는 경우는 sequence modeling의 기존 목적과 일치하므로, 변형없이 beam search 알고리즘과 같으 탐색알고리즘을 적용하며 사용가능함
• 그 결과 현재의 state $s_{t}$로 시작하는 tokenized trajectory $\bar{\tau}$를 생성하며, 이때의 action $\bar{a}_{t}$ 는 reference 행동을 모방하는것으로서 model-based behavior cloning 역할을 함

3.2.2 Goal conditioned RL

• Transformer 아키텍처는 "causal" attention mask기반의 이전 토큰에만 의존한 다음 토큰 예측을 그 특징으로 하며, 이는 물리적으로 미래가 과거에 영향을 미치지 않는 "physical causality"와도 일맥상통함.
• 하지만 self-attention 아키텍처에 기반하고 있기 때문에 과거 뿐만아니라 미래 역시 다룰 수 있어, 아래와 같이 마지막 state가 주어졌을때의 conditional probability역시 decode 할수 있음.

$P\left ( \bar{s}_{t}^{i} \mid \bar{s}_{t}^{<i}, \bar{\tau}_{<t},\bar{s}_{T-1}  \right )$

• 이때 마지막 $\bar{s}_{T-1}$는 일어나길 바라는 상태로서 지속적으로 입력 sequence에 줄 경우,   trajectory transformer를 goal-reaching 방식으로 사용가능함.
• 실제 구현에선 미래의 goal state token을 sequence의 제일 앞에 항상 붙이는 방법으로, 기존 Transformer의 causal attention mask를 그대로 사용가능함.
• 이러한 접근은 기존에 supervised learning을 goal-conditioned policy에 적용하던 접근 혹은 relabeling 접근과 유사함.

3.2.3 Offline RL

• Beam search (BS) 는 sequence planning의 한 방법으로 가장 높은 log-probability를 가지는 sequence들을 선택하는 알고리즘 

• BS 알고리즘에서 각 토큰의 log-probability대신 reward를 사용하면 reward-maximizing planning 알고리즘으로 사용가능
• 하지만 single-step reward-maxing으론 근시적인 planning만 가능하므로, 대신 아래의 reward-to-go (discounted return) 값을 전처리하여 학습 trajectory의 각 transition step에 reward 토큰 다음에 오는 새로운 토큰으로 추가

$R_{t}=\sum_{t'=t}^{T-1}\gamma^{t'-t}r_{t'}$

• 이로서 planning과정에서 reward-to-go의 예측값, 즉 value를 추정하지만 BS알고리즘에서만 휴리스틱 가이드로서 사용하므로 offline setting에서의 value estimation 문제와 달리 정확할 필요가 없어진다는 장점이 생김.
• 구체적으론, transition을 하나의 단어로 취급하여 likelihood-maximzing BS알고리즘을 사용하여 가장 확률이 높은 transition을 먼저 샘플링 한 뒤, 이 중에서 reward와 value가 높은 transition을 필터링함으로써 passimism없이도 기존의 offline RL과 같이 conservative한 planning효과를 내는것이 가능해짐.

 

4. Experiments

• 실험 파트에서는 다음 두가지에 집중하여 검증을 진행함
  1. long-horizon 예측에 대한 기존 single-step 예측 모델 대비 Trajectory Transformer 모델의 정밀성
  2. Offline RL / imitation learning / goal-reaching 세 가지 문제에서 BS알고리즘을 sequence 모델링 툴로 사용한 것의 제어 알고리즘으로서의 효용성

4.1 Model Analysis

4.1.1 Trajectory prediction

humanoid policyd 데이터셋을 reference로 하여 학습된 모델로 100 step을 예측한 결과 비교

• Single-step만 예측하는 Markovian 조건 하에서 ensemble dynamics 모델을 학습하고 policy의 action에 따른 transition을 생성하는 probabilistic ensembles with trajectory sampling (PETS)알고리즘의 경우 feedforward로 few dozon step이 넘어가면 오차가 매우 커짐 (Figure 1의 세번째 줄) 
• 반면 Trajectory Transformer는 reference (첫번째 줄) 와 차이없는 모습을 100 step 넘게 안정적으로 예측하는것을 보여주었으며 (두번째 줄) 이는 model-based RL방식에서 SOTA임.

4.1.2 Error accumulation

• 위 실험에서 보여준 누적 오차를 정량적으로 평가하기위해 두 모델에서 각각 1000 trajectory를 샘플링하여 per-timestep state marginal을 구하고 reference state의 likeihood를 비교한 결과, Trajectory Transformer가 월등한 대비성능을 보여줌
• 또한 Trajectory Tranformer로 하여금 과거의 1개 step만 참고하도록 변경한 Markovian Transformer의 경우도 비슷한 예측 성능을 보여줌으로서, Transformer 아키텍처 그 자체와 autoregressive state discretization으로 인해 향상된 expressivity가 long-horizon accuracy에 큰 역할을 한다는것을 확인.
• state 차원을 랜덤하게 절반 masking한 partially-observed 케이스에서 original Trajectory Transformer가 abblation대비 높은 성능을 보여주어 long-horizon conditioning의 accuracy에대한 역할을 확인할 수 있음.

4.1.3 Attention patterns

• Trajectory Transformer는 두가지 attention pattern을 보여줌
  1. 직전 transition에 집중하는것으로 Markovian strategy를 학습한 패턴
  2. State는 동일한 dimension에 집중하며 action은 과거 state보다 과거 action에 집중하는 패턴
• 두번째 패턴의 경우 action이 과거 state에만 의존하는 behavior cloning과 반대되는 결과이며, 몇 trajectory optimization알고리즘들에서 쓰인 action filtering 기법과 닮음.

 

4.2 Reinforcement Learning and Control

4.2.1 Offline RL

• D4RL offline RL 벤치마크에서 reward-maximizing TTO알고리즘을 검증하며 이를 다음 4가지 접근방법과 비교함.
  1. Conservative Q-learning (CQL; model-free SOTA)
  2. Model-based offline policy optimization (MOPO; model-based SOTA)
  3. Model-based offline planning (MBOP; single-step dynamics 모델 사용)
  4. Behavior cloning (BC)
• 그리고 3가지 MuJoco 환경에서 다음 3가지 데이터셋으로 학습을 진행함.
  1. medium: 중간정도의 점수를 가지는 policy로 만들어낸 데이터
  2. med-expert: medium 데이터와 최고점수를 가지는 policy로 만들어낸 데이터의 혼합
  3. mixed: medium을 학습시키기 위해 사용된 replay buffer 데이터 
• 실험 결과 모든 데이터넷에서 대체로 기존 알고리즘들과 동등하거나 보다 뛰어난성능을 보여주었음.
• 단, HalfCheetah의 med-expert데이터에서는 expert데이터의 성능이 상당히 좋아지는 바람에 discretization이 세세하게 되지 못하여 성능이 낮은것으로 추정.

4.2.2 Imitation and goal-reaching

• Behavior cloning 성능을 확인하고자 likelihood-maximizing TTO 알고리즘을 hopper와 walker2d환경에 적용한 결과, 각각 behavior policy대비 104%와 109%의 return을 보여줌.

• Goal-reaching 성능을 확인하고자 likelihood-maximizing TTO 알고리즘을 goal-stated conditioned sequence를 주어 four rooms환경에 적용한 결과, reward shaping이나 어떤 reward 정보 없이 goal relabeling만으로도 plannindg이 잘 되는것을 확인함.

 

5. Discussion

• Large-scale language 모델에서 사용되는 Transformer를 기존 강화학습의 분리된 구조 대신 하나의 통합된 sequence 모델로서 beam search알고리즘과 함께 사용한 결과 imitation learning, goal-reaching, offline RL에서 효용성을 보여줌.
• 구조가 간편하고 유연해지는 대신 기존의 single-step 예측모델에 비하면 느리다는 단점이 있으나, 현재의 Transformer의 경량화 흐름을 비추어보면 문제없을것으로 예상.
• Continuous space 데이터의 경우 disctretize해주어 사용하기때문에 accuracy가 떨어진다는 단점이 있지만 이는 다양한 discretization방법으로 어느정도 해결가능할것으로 예상.
• Transformer를 사용하여 RL문제를 SL문제로 다루는것이 처음은 아니지만, 이러한 접근이 Markov property를 학습하고 offline RL에서 기존 방법과 비교할만한 결과를 보여준다는것을 입증한것에서 의의가 있음.

 

 

개인적인 의견

• 기존의 Model-based RL에서도 말만 RL이지 policy부분은 전혀 RL이 아닌 경우가 있어서 사실상 Model-based RL 접근으로 볼 수 있을것 같다. 특히 Model-based RL에서 가장 중요한게 model부분인데 기존의 single-step model은 markov property 조건때문에 real-world에서 적용이 어려워보이는 느낌이 강했었다. 이러한 문제를 Transformer를 사용하면 쉽게 해결할 수 있어져 decision making AI가 real-world로 나아갈수 있는 방향이 크게 늘어날것 같은 생각이 든다.
• Model-based RL의 대표적인 model인 world-model은 RNN 기반인데, NLP에서 Transformer가 RNN을 대체한 만큼 world-model에서도 Transformer가 할 수 있는게 많지 않을까 싶다.
• 의료 도메인에서 특히나 치료/추천형 AI 연구를 하는 입장에서 AI의 판단근거를 항상 고려해야 하기에, 개인적으로 Transformer아키텍쳐는 attention map이 늘 매력적으로 다가왔다. 물론 attention의 XAI 측면은 자연어 생성 분야에서 여전히 말이 많긴하지만, Decision making으로 나아가고자 하는 이 논문이 추구하는 방향성에서 충분히 의미있는 정보가 될것은 분명한듯하다.
• Goal-conditioned reaching역시 재밌는 부분이다. policy의 경우 Goal-conditional policy, Contexual policy 등 추가적인 정보를 policy에 주어 더 복잡한 task를 푸는 연구가 많다. 이러한 task들을 conditional-sequence로 새롭게 접근하는 방법들이 많이 나올 수 있을것 같다.