[RL] Decision Transformer에 대한 이해

이 글은 논문 Decision Transformer: Reinforcement Learning via Sequence Modeling에 대한 이해를 돕기 위한 포스팅이다. 논문 외 추가적인 내용에 관한 참고 문서들은 페이지 맨 아래에 적어놓았다.

Abstract

이 논문은 기존의 강화학습과 달리 sequential modeling을 통한 강화학습을 이용하여 여러 문제들을 풀 수 있다는 것을 보여준다. Decision Transformer는 causally masked Transformer를 이용함으로써 Reinforcement Learning (RL)의 문제들을 sequential modeling으로 바꾼다.

위 이미지는 Decision Transformer 구조를 도식화한 것으로, Return-to-go $\hat{R}$, states $s$, actions $a$를 입력으로 취해 다음 time-step의 action을 예측한다. 예측된 action는 autoregressive하게 다음 step의 입력으로 취해진다. 세부적인 이야기는 뒤에서 다루겠다.

Causal Language Models (CLM) [1]

Masked Language Model (MLM)은 문장 내 특정 단어을 masking 하고 masked 단어를 예측하도록 학습한다. CLM도 비슷하지만 좌측의 단어들만을 고려하여 masked 단어를 예측하도록 한다. 결과적으로, MLM은 모델을 bidirectional 하게 만들고 CLM은 unidirectional 하게 만든다. Decision Transformer가 CLM을 채택한 이유는 autoregressive 모델을 가능하게 하기 위함이다.

 

Introduction & Preliminaries

Sequential decision making problems을 가지고 전통적인 RL을 대체할만큼의 이점들이 있는가???

Sequence modeling objective를 사용하여 수집된 경험에 대해 Transformer 모델을 훈련함으로써 long term credict assignment를 위한 bootstrapping의 필요를 피할 수 있다. 그러므로 RL를 불안정하게 만드는 "deadly triad" 중 하나를 회피하며 discounting factor $\gamma$ 도 또한 사용하지 않아도 된다. 추가적으로 Transformer는 긴 sequences를 모델링할 수 있고 self-attention을 거쳐 직접적으로 credit assignment를 수행할 수 있다. 이는 보상 전파가 느리고 "distractor" signal을 받기 쉬운 Bellman backups과는 대조적이다. 또한 behaviors의 넓은 분포를 잘 모델링할 수 있어 generalization과 transfer를 하기에 더 좋다.

위의 가설들을 파헤치기 위해 offline RL을 고려한다. Offline RL task는 error propagation과 value overestimation으로 인해 challenging하지만, sequence modeling objective를 가지고 훈련할 때는 자연스러운 작업이다. States, actions, returns의 sequences에 대해 autoregressive model을 훈련함으로써 policy sampling을 autoregressive generative modeling으로 축소한다.

Credit assignment problem [5]

수많이 행했던 각 행동이 결과에 무슨 영향을 끼쳤는지 모르는 문제다. 예를 들면 100번의 행동 중 30번째 행동이 좋은 보상을 얻는데 방해가 됐는지 도움이 되었는지 모른다.

Bootstrapping

같은 종류의 추정 값에 대해서 업데이트를 할 때, 한 개 혹은 그 이상의 추정 값을 사용하는 방법이다. 예를 들면, SARSA에서 Q-value를 업데이트하기 위한 수식은 다음과 같다.

$$Q(s,a)\leftarrow Q(s,a)+\alpha(R_{t+1}+\gamma Q(s',a')-Q(s,a))$$

여기서 $R_{t+1}+\gamma Q(s',a')$는 $Q(s,a)$의 실제 값에 대한 추정 치 TD target이다. 반면, Monte Carlo의 업데이트 수식은 다음과 같고 따라서 bootstrapping 방식이 아니다.

$$Q(s,a)\leftarrow Q(s,a)+\alpha (G_t-Q(s,a))$$

Deadly triad [5]

RL을 불안정하게 만드는 다음 세 가지 요소다.

1) Function approximation : Q-value를 function[e.g) ML]을 통해 추정하여 업데이트를 하는 경우로, 한 state에 대해 value가 업데이트될 때 다른 state에 대한 value가 부적절하게 바뀔 수 있는 문제가 있다.

2) Bootstrapping : 위의 설명 참고

3) Off-policy learning : target policy보다는 다른 행동 policy에 의해 생산된 transitions 또는 episodes의 분포에 대해 훈련하는 방식이다.

Bellman backup [6]

Bellman expectation equation에서 미래의 value로 현재의 value를 계산하는 것

Offline RL [4]

전통적은 active RL은 environment와 상호작용하며 학습을 한다. 그러나 의학 진단 프로세스나 dialog 모델 등에서 environment와 상호작용하며 policy를 학습하는 것은 불가능에 가깝다. 그래서 offline RL은 policy $\pi_\beta$ (정답이 되는 행동 등)을 가지고 suboptimal 데이터 집합 $\mathcal{D}$를 수집하고 이로부터 policy를 학습한다. $\mathcal{D}$는 훈련동안 바뀌지 않으며 policy는 완전히 학습이 끝나면 그냥 바로 배포된다. 정리하자면, exploration 없이 고정된 데이터만 가지고 exploitation을 수행하는 것이다. 이 논문에서는 다음과 같이 정의한다.

$MDP = \{S,\mathcal{A},P,R\}$

Trajectory $\mathcal{T}=(s_0,a_0,r_0,s_1,a_1,r_1,...,s_T,a_T,r_T)$

Sequence of states, actions, and rewards

The return of a trajectory at time-step $t$ : $R_t=\sum^T_{t'=t}r_{t'}$

The sum of future rewards from that timestep

강화 학습의 목표는 MDP에서 return의 기댓값 $\mathbb E[\sum^T_{t=1}r_{t'}]$을 최대화하는 policy를 학습하는 것이다. Offline RL에서는 환경과 상호작용을 거쳐 데이터를 얻는 대신 임의의 policies의 trajectory rollouts로 구성된 데이터셋(고정됨)에만 접근한다. Offline RL의 이러한 설정 때문에 환경을 탐험하고 추가적인 피드백을 수집하는 agent의 능력이 제거되어 더 어렵다.

Exploration [2,3]

기본 RL은 학습할 수 있는 사전 생성 데이터가 없기에 모든 이용가능한 가능성을 시도하고 전반적인 reward를 최대화하기 위한 path (or trajectory)를 찾아야 한다. 따라서 agent는 여러 states에서 수많은 actions을 시도해야 한다. 이러한 탐색을 exploration이라고 한다.

Exploitation [2,3]

보상을 최대화하기 위해 agent가 학습 정보를 이용하는 것

Value overestimation

잘못된 action에 대한 Q-value (or 다른 value)가 과평가되어지는 문제

Transformer

Transformer는 sequential 데이터를 효율적으로 모델링하는 아키텍쳐로, residual connections를 가진 stacked self-attention으로 구성된다. 각 self-attention layer는 unique input tokens에 대응하는 n개의 embeddings $\{x_i\}^n_{i=1}$를 받고 input dimensions이 보존된 n개의 embeddings $\{z_i\}^n_{i=1}$을 출력한다. 자세한 설명은 [NLP] Transformer 모델 분석 (Self-Attention)을 참고하길 바란다.

앞서 언급했듯 self-attention layer는 query, key vectors의 유사도를 통해 state-return associations을 함축적으로 형성함으로써 credit을 할당하도록 한다. 입력으로 time-step에 대한 시퀀스가 들어가고 그를 이용해 가중치를 계산한다는 느낌을 생각하면 자연스럽게 받아들여진다.

 

Method

Decision Transformer는 Transformer 아키텍쳐를 최소한 수정하여 자기회귀적으로(autoregressively) trajectories를 모델링할 수 있도록 한다.

Trajectory representation

입력에 대해 Transformers가 의미 있는 patterns을 학습할 수 있어야 하고, 조건부적으로 test time에서 actions을 생성할 수 있어야 한다. 따라서 autoregressive training과 generation을 처리할 수 있는 trajectory representation은 다음과 같다.

$$\mathcal T=(\hat R_1,s_1,a_1,\hat R_2,s_2,a_2,...,\hat R_T,s_T,a_T)$$

과거 보상들보다 미래의 desired returns에 기반된 actions을 생성하도록 모델링하길 원하기 때문에 rewards을 바로 모델에 feeding하는 대신 returns-to-go $\hat R_t=\sum^T_{t'=t}r_{t'}$을 가지고 모델에 feeding한다. 추가적으로, 현재 state에 대해 생성된 action을 실행한 후 에피소드가 종료될때까지 달성된 보상만큼 target return을 줄인다. (Decrement returns-to-go with reward) > $R\leftarrow R + [R[-1]-r]$

Architecture

위 그림은 논문 내용을 바탕으로 Decision Transformer의 아키텍쳐를 구체화한 것이다. 모델의 입력으로 마지막 $K$ time-steps을 feed한다. 각 time-step은 return-to-go, state, action, time-step index로 구성된다. 해당 입력에 대한 token embeddings을 얻기 위해 raw inputs을 linear layer를 가지고 embedding dimension으로 project하여 각 modality에 대해 학습한다. 그 뒤 layer normalization을 적용한다. Visual input을 가진 환경에 대한 state는 linear layer 대신 convolutional encoder로 feed된다. 추가적으로 각 time-step을 위한 embedding을 학습하고 앞서 언급한 각 tokens에 더해진다. 여기서 time-step에 대한 embedding은 기존 transformer에서의 positional embedding이 아닌 세 개의 tokens에 대응되는 하나의 time-step에 대한 것이다. 그리고나서 tokens은 GPT 모델에 의해 처리되며 autoregressive modeling을 거쳐 미래의 action token을 예측한다.

Training

먼저, Offline trajectories 데이터셋 $\mathcal D$로부터 sequence 길이 $K$만큼 미니배치를 샘플링한다. Input token $s_t$에 대응하는 prediction head는 $a_t$를 예측하도록 훈련된다. 손실 함수는 discrete actions이라면 cross-entropy loss를, continous actions이라면 mean-squared error를 사용한다. 또한 각 time-step에 대한 losses를 평균하여 최종 loss로 계산한다. Performance를 향상시키기 위해 states 또는 returns-to-go를 예측하게 하지 않았지만 구조 상 쉽게 추가 가능하다.

References

[1] Understanding Masked Language Models (MLM) and Causal Language Models (CLM) in NLP

[2] RL — Reinforcement Learning Terms

[3] The Complete Reinforcement Learning Dictionary

[4] S. Levine, A. Kumar, G. Tucker, and J. Fu, "Offline Reinforcement Learning: Tutorial, Review, and Perspectives on Open Problems", Arxiv, 2020

[5] H. v. Hasselt, Y. Doron, F. Strub, M. Hessel, N. Sonnerat, and J. Modayil, "Deep Reinforcement Learning and the Deadly Triad", Arxiv, 2018

[6] Bellman Optimality