PR402 - 13 : Decision Transformer: Reinforcement Learning via Sequence Modeling

 오늘 리뷰할 논문은 Decision Transformer: Reinforcement Learning via Sequence Modeling(Chen et al., 2021)입니다. 이 자체로도 꽤 유명한 논문이기도 하고, 공저자 중에 유명한 사람으로는 피터 아빌 교수가 있습니다. RL problem을 Sequence Modeling으로 살짝 틀어서 Transformer를 통해 풀어낼 수 있다는 내용인데요, 앞으로 RL이 나아갈 방향성에 대한 새로운 패러다임을 제안하고 있습니다.

Link - https://arxiv.org/abs/2106.01345

Intro:

RL problem을 Sequence modeling problem으로 치환하여 simplicity 및 GPT나 BERT등 Transformer 기반 모델들에의 scalability를 높일 수 있다는 내용으로 출발합니다. 기존의 RL problem들이 당면한 long term credit assignment 문제(action을 취한 시점과 reward를 받는 시점의 시차로 인한 학습의 어려움)라던지, deadly triad(모델이 잘 수렴하지 않는 원인 - bootstrapping, off-policy learning, function approximation)들을 피해감으로써 보다 안정적인 학습을 도모할 수 있다고 이야기하고 있습니다. 

기본적으로는 offline RL을 기반으로 하고, 모델 아키텍쳐는 Transformer 기반입니다. 알고리즘을 보면서 이해하는 것이 훨씬 빠를 것 같아 바로 넘어가도록 하겠습니다.

Contents:

기본적으로는 reward, state, action, timestep을 받아서 action을 ouput으로 직접 뱉어내는 구조를 가지고 있습니다. 기존의 RL에서는 대부분 state를 받아서 action을 출력하는 등의 방식으로 function(모델)이 사용되었는데요, 여기서 제안하고자 하는 방법은 state뿐만 아니라 아예 RL system에서 사용되는 정보 전체를 모델의 input으로 집어넣음으로써 RL problem 자체를 sequence modeling problem으로 간주할 수 있다는 아이디어를 바탕으로 하고 있습니다. 

아이디어 자체는 근사하지만, 이 방법이 실제로 성능을 잘 뽑아내기 어려웠던 것은 모든 정보들(r, s, a, t)의 dimension이 다 다를 뿐더러 이를 concatenate해서 집어넣어주는 방법 등을 취해도 뉴럴넷이 이를 이해하기 어렵다는 것입니다. 

결과적으로 위를 가능하게 한 것은 Attention 덕분인데요, query를 통해 각 정보들을 소화하기 쉬운 형태로 만들어주면서 뉴럴넷이 RL system 자체를 이해할 수 있도록 만들어 줄 수 있다는 것이 이 논문에서 이야기하고자 하는 바라고 생각합니다. (개인적인 의견입니다.)

모델을 좀 더 뜯어보면, Decision Transformer 자체는 r, s, a에 t를 기반으로 positional embedding을 넣어주고, transformer를 통과시켜 action을 뽑아내는 구조로 이루어져 있습니다. 학습 과정에서 특이한 점은 reward를 넣어줄 때 해당 timestep에서 구해진 action이 아니라 최종 endstate까지 '얻을' return의 총합(returns-to-go)이라는 점입니다. 저는 여기에 꽤 많은 고민이 녹아들어있다고 생각합니다. 학습 과정을 살펴보면 behavioral cloning과 같은 방식으로 이루어집니다. 데이터셋에서 취해진 action을 복제하는 방향으로 transformer가 학습되는데요, 그렇기 때문에 '좋은' 방향으로 모델을 학습시키는 것이 아니라 '복제'하는 방향으로 모델을 학습시킵니다. 기존의 RL에서는 Temporal Difference 등과 같은 방법으로 reward를 현재가치로 할인시켜 future reward에 대한 학습을 시킬 수 있었습니다. 그러나 그러한 방법을 사용한다면 기존의 RL에 Transformer라는 모델을 사용한 것밖에 되지 않습니다. (이 이유도 굉장히 크다고 생각합니다...ㅋㅋ) 따라서 이를 조금 다른 방식으로 해결해야 했는데, 이를 위해 returns-to-go라는 개념을 만듦으로써 future reward에 대한 지식을 모델이 학습할 수 있도록 하였습니다. 따라서 prediction 단계에서 target return을 설정해줌으로써 원래 모델이 학습했던 trajectory중에서 자동으로 target return에 도달할 수 있는 action에 가깝게 output을 뱉어내도록 유도하여 cloning에서 벗어나 바람직한 action을 취할 수 있도록 한 것입니다. 

마지막 prediction 단계에서는 최근의 K개 정보들을 사용하여 context(history)를 가져갈 수 있도록 구성하고 있습니다.

최대한 간결하게 Decision Transformer 요약을 준비해봤는데, 잘 전달이 되었을지 모르겠네요. 그럼 이제 성능을 보러 가실까요?

Experiments:

Atari, OpenAI Gym, Key-To-Door problem에서의 성능을 비교하고 있습니다. Atari나 OpenAI Gym에서는 사실 기존의 알고리즘을 크게 뛰어넘고 있지는 못하지만, Key-To-Door problem에서는 압도하고 있네요. 후술하겠지만, long term credit assignment 문제는 상대적으로 잘 해결하고 있다는 사실을 알 수 있습니다.

논문에서는 5. Discussion 부분에서 추가적으로 몇 가지 논의들을 이어나가고 있습니다. 몇 가지 추려서 적어보겠습니다.

Q . DT는 Behavioral Cloning을 하고 있는게 아닌가? 

제가 앞서 설명드렸듯이, training 과정이 dataset에서 취한 action을 복제하는 방향으로 이루어지기 때문에 returns-to-go에 대한 이해가 없다면 behavioral cloning과 다를 바가 없다고 생각할 수 있습니다. BC보다 낫다는 것을 증명하기 위해 '%BC'라는 방법을 사용하는데요, dataset에서 좋은 모습을 보인 x%의 데이터들만 사용한 BC의 성능들과 비교하여 DT가 BC보다 나은 무언가를 가지고 있다는 것을 간접적으로 증명하고 있습니다.

Q . 굳이 K만큼의 history를 같이 넣어야 하나? 현재 상태만 가지고 어떤 action을 취할지 결정하면 불충분한가?

K=1로 설정한 값과 비교하였을때 성능 차이가 굉장히 많이 나는 점을 강조하여 답하고 있습니다.

Q . Long term credit assignment 문제를 잘 해결하고 있는건가?

Credit assignment가 굉장히 중요한 이슈인 key-to-door problem에서 우수한 성능을 보이는 것을 바탕으로 해당 문제에 강력한 모습을 보이는 것을 간접적으로 증명하고 있습니다. key-to-door problem은 agent가 door에 도착하여야 하는데, key를 갖고 도착하여야 reward가 제공되는 경우를 말합니다. 즉, agent가 했던 과거의 'key를 가지러 가는 action'을 통해 door에 도착하였을 때의 reward가 같이 제공되기 때문에 long term credit assignment 성능을 판가름하는데 좋은 지표가 될 수 있습니다.

이외에도 discussion 관련 몇 내용이 더 있지만, 이쯤에서 줄이도록 하겠습니다.

Conclusion:

물론 RL problem 자체를 sequential problem으로 치환한 것은 참신하지만, 과연 성능이 앞으로도 뒷받침될까에 대한 부분은 의문이 생깁니다. 성능 부문에서의 장점으로는 credit assignment issue 부문에서의 강점인데, 사실 이마저도 data가 많아지게 되면 다른 알고리즘에 밀리는 모습을 보여주고 있습니다. 결국 이 논문은 performance보다는 novelty나 new paradigm 측면에서 강점을 갖고 있는데, 해당 방법론이 메인스트림이 되기 위해서는 결국 performance가 뒷받침되어야 합니다. 사실 기존의 RL 모델들도 최근에는 한 계단 올라선, 그런 혁신적인 성능을 보여주고 있지는 못하는 것 같아서 어쩌면 가능할지도 모르겠습니다. 

더 발전시킬 만한 부분은 아마도 reward setting 부분일 것 같습니다. returns-to-go라는 개념은 충분히 매력적이지만 첫 시도에서 제안된 방법이기 때문에 분명히 건드릴 부분이 존재할 거라고 생각합니다. 그렇지만 계속 생각해볼수록 기존의 RL으로 흘러들어가게 되어서 쉽지 않을지도 모르겠네요... 아무튼! 오늘은 이만 마치도록 하겠습니다.

읽어주셔서 감사합니다!

+틀린 부분에 대한 코멘트는 언제나 환영합니다 :)