PEARL, VariBAD와 같은 inference based meta-RL의 structured (graphical) representation 버전
World-model의 multi-task/meta RL 버전
기존의 방법들이 implicit latent task varible을 inference하는 구조였다면, AdaRL은 이를 task간의 domain-shared latent state $s_t$와 domain-specific change factor $\theta_k$ 로 나누어 explicit한 inference를 함
이를 가능하도록 학습하기 위해 두 time step동안의 state dimension, action, reward 그리고 이들에 영향을 주는 domain specific parameters $\theta_k$사이의 상호관계 graph를 아래와 같이 가정
여기서 $c$는 그래프에서 각 요소 사이의 edge에 대한 mask parameter이며 이로 인해 최소한의 필요한 representation만 남음
이러한 latent variable를 inference하는 encoder와 각각의 edge들을 end-to-end 학습하기위해 VAE (MiSS-VAE; Multi-model Structured Sequential Variational Auto-Encoder)구조를 제안하여 사용
이렇게 구한 explicit한 compact latent variable만 있으면 policy를 학습하기에 충분하다는것을 appendix에서 증명
알고리즘의 검증은 Cartpole과 Atari Pong에서 진행
explicit한 task representation덕분에 다양한 task의 variation에서 in/out of distribution 모두에서 기존 알고리즘들보다 비슷하거나 효율적인것을 확인
개인적인 생각
explicit하게 task representation을 나누어 학습하는것이 좋다. graph를 사용하여 이를 학습가능하게 하는부분이 기발하다. 다만 리뷰어도 언급한 부분으로 이걸 저자가 interpretable이라고 부르기엔 이에대한 검증이 부족해 보인다.
지금까지 많은 Meta-RL 논문들이 reward 기반 multiple task에 좀더 집중을 하여 검증을 해왔기에, state dynamics(혹은 transition)이 달라지는 multiple task에 대한 접근은 상대적으로 실험이 부족한감이 없지않아 있었다. 반면 이 연구에선 기존 연구들과 달리 이러한 state dynamics의 변화에 따른 adaptation을 reward보다 더 중점적으로 다룬단 차별점이 있다. 다만 이는 저자의 말에 따르면 reward가 달라지는 task에선 20~50 step의 적은 sample만으론 adaptation이 어려워서가 이유이기도하다.
같은 맥락에서 리뷰어도 이야기를 한 부분인데, state dynamics에 잘 adaptation하는 알고리즘이라면 그 장점을 더 잘 보여줄 수 있게 locomotion과 같은 task에서도 검증을 했다면 어떨까 싶다.
실제로 이렇게 structured latent estimation을 할 경우 단순히 baseline대비 성능이 좋다는 내용만 본문에 써둔게 아쉽다. 실제 추론된 $\theta_k$의 검증은 appendix에 있는데, 환경의 물리적 변화와의 연관성이 충분히 보이고 있어서 제안하는 structured self-supervised approach의 유효성을 어느정도 입증하고 있다.
Cartpole에서 실제 물리적 변화에 따른 추론된 latent variable의 경향성
policy optimization이 필요없는 meta-adaptation을 강조하고 있는데, 이는 VariBAD나 PEARL과 같은 inference based meta-RL들에서 이미 보여준 부분이긴하다. 리뷰어도 언급한 부분인데 이 inference에 대한 설명이 본문에 안보여 Appendix를 읽기 전엔 이해가 다소 어렵다.
domain index $k$를 input으로 사용하는 부분이 multi-task RL의 색깔이 있기때문에, meta-RL관점에서 이 알고리즘이 generalization이 잘 될지도 궁금하다.
알고리즘이 좀 복잡하다. 상당히 많은 notation이 본 paper와 appendix에 섞여 있어서 어렵긴 하나, 이론적 증명이 탄탄한것이 큰 장점이어서 높은 rating을 받은듯하다.
저자들이 causal reinforcement learning을 연구한 사람들이라 structural causal model의 개념을 도입한것 같다.
질량과 마찰력에 대한 OOD 조건에서 실험결과 기존 model-based RL 및 model-free RL 알고리즘 (MBPO & SAC)보다 더 나은 task generalization 및 sample efficeincy를 보여줌
또한 학습조건의 질량과 마찰력이 한쪽으로 쏠린 unbalanced distribution (90:5:5)에서도 다른 알고리즘 대비 더 나은 generalization 및 sample efficient한 성능을 보여줌
개인적인 생각
model-based RL이 real world에서 쓰이기위해 꼭 필요하다고 생각했던, 그리고 causal inferece의 인기를 생각하면 조만간 나올거라 생각했던 counterfactual performance를 다룬 연구라 반가웠다.
다만 다른 몇몇 casual inference연구의 결과에서도 그랬지만, 이 연구도 모델에 counterfactual property에 대한 loss가 따로 들어가지 않는 경우라 생성된 counterfactual data에 대한 신뢰성이 사실상 없다. (21.11.09 코멘트: 그래서인가 오늘 공개된 openrewiew에서 5353점을 받았다.)
Structured casual graph를 쓰는만큼 앞으로는 더욱 inductive bias를 잘 활용할 수 있는 방향으로, 그래서 학습된 world model이 실제 환경의 dynamics에 대한 근본적인 이해를 바탕으로 만든 causal inference가 접목되지 않을까 싶다.
이 논문은 multi-task RL과 그 연장선으로 볼 수 있는데, meta-learning에도 causal inference를 접목하는 연구가 general intelligence를 최종 목표로하는 분야의 특성상 곧 나올것 같다.
기존 memory based meta-RL의 대표 알고리즘인 RL2에서 RNN을 Transformer로 대체한 버전의 알고리즘
Transformer의 구조가 왜 meta-learning에 부합하는지 신경과학적으로 해석함
한 transition의 embedding은 감각에서 들어오는 신경과학에서의 working memory로 보고, 이 들에 attention mechanism을 적용한 것을 신경과학에서의 reinstatement mechanism와 같다고 하며 Transformer 각 step의 output이 episodic memory에 해당한다고 해석
즉, 이러한 working memory들이 episodic memory로 합쳐지고 무엇보다 이게 task 분포를 proxy한다고 가설
RNN보다 Transformer의 sequential representation 능력이 좋기 때문에, MetaWorld에서 RL2의 상위호완 성능을 보여줌
새 step이 들어올때마다 Transformer의 input을 queue처럼 사용하기때문에, 사실상 RNN과 같이 쓸 수 있어 RL2, VariBAD와 같은 online adaptation이 가능한 알고리즘
Episodic adaptation인 MAML이나 PEARL에 비하면 매우 빠른것이 장점
OOD에서는 PEARL이나 MAML보다 TrMRL 및 RL2가 높은 성능을 보여주어 보다 효과적인 representation을 생성함
Working memory가 각 Task에 대해 잘 분리 된다는것을 latent visualization으로 보여줌
개인적인 생각
RL2에서 RNN을 Transformer로 바꾼것이 사실 이 연구의 알고리즘적인 contribution이다. 하지만 RNN을 Transformer로 대체하는건 사실 다른 분야에선 더 이상 contribution이 되지 못한다. (21.11.09 코멘트: 오늘 공개된 openreview 점수가 역시나 5533이다) 그래서 그런지 Transformer의 평범한 구조들에 대한 meta-learning측면의 신경과학적인 해석에 공을 상당히 들였다.
Reward signal만으로 capacity가 큰 Transformer까지 end-to-end로 학습하는구조라 학습이 매우 어렵다. 특히 episodic memory중에서 1개로만 policy loss가 back prop 들어가는 구조라 충분히 학습이 되려나 싶었는데 나름 성능이 좋다고해서 신기하다. 하지만 역시나 학습이 불안정하다는걸 강조하며 Ad hoc으로 network initializaton이 쓰인다.
알고리즘과 환경의 특성상 "Metaworld가 아닌 Mujoco에선 PEARL이 더 performance가 좋을것 같은데 figure가 안보이네"라고 생각하며 읽었는데, 역시나 Appendix로 빼둔것이었고 PEARL이 압도적으로 잘된다. 논문을 쓰는 전략적인 측면에선 나름 알고리즘을 돋보일 수 있는 실험들 위주로 잘 배치한 것 같긴하지만 리뷰어들의 의견이 궁금하다. (21.11.09 코멘트: 오늘 공개된 openreview 점수가 역시나 5533이다)
이 논문에선 episodic memory가 task를 proxy한다고 가설을 세웠다. 하지만 정작 task를 구분하는 loss가 없어서 어찌될지 궁금한데, episodic memory의 latent space를 visualized한 결과를 안보여줬다. Reward기반의 task들이라 linear한 구분이 당연히 잘 될 수 밖에 없는 working memory의 latent space만 visualization해서 task구분이 되는듯하다고 써둔것이 아쉽다.
OOD에 대한 해석에서 memory based 알고리즘이 optimization based나 context based 알고리즘보다 더 representation이 효과적이라고 했으나, 개인적으론 memory based 알고리즘 역시 (혹은, 특히나) generalization에 대한 loss가 딱히 없기 때문에 잘못된 해석인것 같다. 실제로 MAML, PEARL보다 높다 뿐이지 해당 halfcheetah에서의 score를 보면 제대로 동작한다고 보긴 어려워, RL2 및 TrMRL에 의해 형성되는 task latent space의 명확함이 오히려 떨어져서 생기는 현상이라고 생각된다.
Author : Amir Feder, Katherine A. Keith, Emaad Manzoor, Reid Pryzant, Dhanya Sridhar, Zach Wood-Doughty, Jacob Eisenstein, Justin Grimmer, Roi Reichart, Margaret E. Roberts, Brandon M. Stewart, Victor Veitch, Diyi Yang Paper Link :https://arxiv.org/abs/2109.00725
A Transformer-Based Variational Autoencoder for Sentence Generation (IJCNN 2019) Paper: https://ieeexplore.ieee.org/document/8852155 - Domain: Natural language generation
- RNN기반 VAE의 long-term input에 대한 posterior collapse와 같은 문제를 해결하고자함 - Encoder는 기존 transformer와 구조는 같고 출력이 Gaussian distribution - Test조건에선 decoder와 샘플링 파트만 사용한다고 가정하여 cross-attention 레이어는 제거하고 마지막에 autoregressive LSTM generator를 추가
Transformer VAE: A Hierarchical Model for Structure-aware and Interpretable Music Representation Learning (ICASSP 2020) Paper: https://ieeexplore.ieee.org/document/9054554 Related blog: https://mip-frontiers.eu/2020/08/20/transformer-vae.html - Domain: Music generation
- Structure awareness와 interpretability를 모두 만족하는 음악생성모델을 만들고자 함 - Encoder는 기존 transformer와 구조는 같고 네트워크 출력으로 입력 토큰 개수 만큼의 Gaussian distribution - Decoder역시 기존의 transformer구조와 같이 autoregressive한 구조이며, cross-atention을 encoder에서 샘플된 입력 길이 만큼의 latent varible과 수행
T-CVAE: Transformer-Based Conditioned Variational Autoencoder for Story Completion (IJCAI 2019) Paper:https://www.ijcai.org/proceedings/2019/727 Code:https://github.com/sodawater/T-CVAE - Domain: Story completion
- 문단의 맥락을 파악하여 빈 문장을 적절하게 생성하고자 함 - Encoder와 decoder의 입력이 같은 distrubution set이어서 layer를 공유함 - Encoder는 [나머지 문장; 생성중인 문장]을 self-attention하여 latent variable을 추론하고, decoder는 생성중인 문장을 [encoder 각 레이어 output; decoder의 output]에 cross-attention - generator는 deocoder의 output을 condition으로하면서 샘플된 z를사용하여 빈 문장을 autoregressive하게 reconstruction
Variational Transformer Networks for Layout Generation (CVPR 2021) Paper: https://arxiv.org/abs/2104.02416 Blog: https://ai.googleblog.com/2021/06/using-variational-transformer-networks.html - Domain: Layout design - 그럴듯한 레이아웃을 가능한 다양하게 생성하고자 함 - self-attention레이어가 주어진 레이아웃 내부의 element사이의 관계를 파악하며 representation distribution에 압축 - VAE에 Transformer의 attention mechanism을 추가하여 생성 레이아웃의 context이해와 다양성을 동시에 추구 - BERT의 auxiliary token과 같은 역할로서 encoder output을 bottleneck으로 사용 후 샘플된 z를 autoregressive decoder의 첫번째 입력으로 넣어줌 - Decoder는 cross-attention이 빠져 상대적으로 덜 expressive한 구조이지만 그 덕분에 posterior collapse는 줄어듬
Attention-based generative models for de novo molecular design (Chemical Science 2021) Paper:https://pubs.rsc.org/en/content/articlelanding/2021/SC/D1SC01050F Code:https://github.com/oriondollar/TransVAE - Domain: Molecular design - 분자 구조 생성에 있어서 RNN, RNN+Attention, TransformerVAE의 장단점을 비교하고자 함 - Encoder는 기존 transformer와 구조는 같고 출력부분에 Convolution bottleneck을 추가해 Gaussian distribution을 추론 - Decoder역시 기존의 transformer구조와 같이 autoregressive한 구조이며, cross-atention을 encoder에서 샘플된 입력 길이 만큼의 latent varible과 수행 - RNN + Attention 이 TransformerVAE보다 OOD에 더 잘 대응하는 반면, TransformerVAE가 더 interpretable 함
Transformer-based Conditional Variational Autoencoder for Controllable Story Generation (Arxiv 2021) Paper: https://arxiv.org/abs/2101.00828 Code: https://github.com/fangleai/TransformerCVAE - Domain: Story generation
- Encoder는 unmasked/bi-directional self-attention 구조이며 출력부분에 attention-average block과 bottleneck을 추가해 Gaussian distribution을 추론 - Decoder는 GPT-2의 autoregressive한 구조이며, encoder에서 샘플된 latent variable을 1. input, 2. pseudo self-attention, 3. softmax에 전달하는 구조를 각각 시도 - Pre-trained model을 사용하여 posterior collapse 줄이려 시도
Finetuning Pretrained Transformers into Variational Autoencoders (Workshop on EMNLP 2021) Paper:https://arxiv.org/abs/2108.02446 Code:https://github.com/seongminp/transformers-into-vaes - Domain: Language modeling
- Posterior collapse를 보다 효과적으로 줄이고자 함 - Pretrained-T5 모델을 VAE에 사용하되 decoder는 self-attention없이 encoder에서 샘플된 z와 cross-attention만을 진행하여 decoder의 expressive power를 제한 - KL loss를 0으로 두고 학습한 후 full VAE를 학습하는 두 단계로 진행
Sequence 데이터에 대한 latent space를 형성하고자 하면서도 더 나은 context의 이해를 위해 attention을 사용하려는 시도에 대한 논문들을 훑어보았다.
대부분 크게 다르지 않았지만, decoder 부분 self-attention의 expressive power가 오히려 sequential VAE의 posterior collapse를 악화시킬수도 있기에 위 연구들 중에서는 이를 피하기위한 고민을 한 논문들과 context representation이 실제로 얼마나 잘 학습되었는지를 보여주는 논문들이 다른 연구들보다 상대적으로 의미있는것 같다.
학습은 카이스트 황보재민 교수님께서 만드신 RaiSim시뮬레이터에 Unitree사의 A1로봇을 넣고 환경변수(Friction, payload, center of mass, motor strength, terrain height 등)를 다양하게 바꿔주고 지면조건도 다르게 하며 meta training을 수행.
Meta training 환경변수와 Meta Test환경변수의 distribution이 겹치되 test가 더 넓게 설정
아키텍처는 두개의 phase로 진행.
phase 1: 시뮬레이션 상의 환경변수 $e_{t}$를 인코딩한 latent task variable $z_{t}$를 conditional policy에 주는 multi-task meta RL을 먼저 학습.
이 결과 fine-tuning없이 바로 실제 A1로봇에 policy와 Adaptation Module을 deploy했을때, 처음 보는 환경(지형, payload 등)에서도 실패 없이 잘 걷는것을 확인
Context-based adaptation이 잘 되는지 확인하고자 latent variable $z_{t}$을 분석한 결과, 오일이 뿌려진 미끄러운 지형에 돌입하면서 빠르게 latent variable이 바뀌는것을 확인할 수 있고 미끄러움이 지속됨에 따라 latent variable도 이를 잘 캡쳐하는것을 확인함.
Unexpected payload가 주어졌을때도 latent variable에 잘 반영이 됨
개인적인 생각
1.빠른 adaptation, 2.높은 'structured' representation power, 3. causal inference로의 확장 가능성, 4.tuning이 필요없는 adaptation의 장점으로인해 Context-based Meta RL이 Real world RL을 위한 solution중 하나가 될거라 생각해왔는데, 버클리와 FacebookAI에서 이렇게 직접 보여주니 좋다.
Real world 문제이니 만큼 이외에도 많은 detail이 들어갔지만, 그럼에도 이정도의 복잡도를 가진 로봇이 deploy후 realworld에서 튜닝없이 바로 동작하는건 고무적이다.
Author : Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Michael Laskin, Pieter Abbeel, Aravind Srinivas, Igor Mordatch Paper Link :https://arxiv.org/abs/2106.01345
바로 이전 포스팅인 Trajectory Transformer의 발표 딱 하루전에 먼저 발표된 논문. (기본 개념이 거의 같은 관계로 이 논문은 정리보다 요약으로 대신하며, 많은부분 두 논문을 비교하고자 함)
Trajectory Transformer가 GPT아키텍처로 RL의 transitions들을 sequence modeling을 하여, 이전 trajectory가 input으로 들어갈때 미래의 trajectory $\tau$를 output으로 하는 모델이라면
Decision Transformer는 input은 마찬가지로 trajectory이지만 output은 action $a$를 output으로 하는 모델이라는 차이가 있음.
Trajectory Transformer와 마찬가지로 offline RL에서의 활용을 주로 삼고 있으나, Trajectory Transformer는 transformer 아키텍처를 planning을 위한 model로 쓰기때문에 model-based offline RL인 반면
Decision Transformer는 transformer 아키텍처가 바로 action을 내뱉기 때문에 model-free offline RL.
Decision Transformer가 추구하는 Offline RL의 예시를 그림으로 잘 설명해주는데, 아래와같이 그래프 navigation task에 대한 랜덤한 trajectory 데이터들이 주어져있을 때, 이 경로들을 stitching하여 새로운 최단 경로를 찾아내는것.
여기서 주의해야할 점이 있는데 바로 마지막 generation path에서 노란색 선 옆에 쓰인 return값이 -3이라는것으로, 데이터셋에 있는 return을 그대로 stitching하는 수준이지
이걸 마치 "두 스텝 이전이니 -2로 return을 생성해야 하는군" 처럼 모델이 이해한다던가 하는 데이터셋 이상의 trajectory를 생성하는건 아니라는 것. (DL유투버 Yannic Kilcher가 -2여야 할것같다고 영상에서 잘못 설명하기도)
이렇다 보니, 적절한 action을 생성하기 위해 의도적으로 return을 매번 감소시켜 입력해줘야하는 등 조금은 억지스러운 Transformer 아키텍처의 활용이 필요.
이런 점에선 Beam Search 알고리즘을 써서 Conservation과 Reward Maximization을 동시에 하는 Trajectory Transformer가 더 자연스러운 sequence model의 활용이지 않을까 싶음.
반면 Trajectory Transformer보다 실험을 매우 다양하게 했다는 장점이 있음
Atari를 통해 high-dim과 long-term credit assignment의 성능을 확인
Qbert에선 잘 안됐지만 다른 환경에선 offline RL SOTA인 CQL과 유사
D4RL 벤치마크를 통해 continuous control에서의 성능을 확인
대부분의 데이터셋과 환경에서 CQL보다 나은 성능 확인BC는 데이터셋에서 상위 % 성능의 데이터를 빼서 return의 조절없이 Decision Transformer를 학습한 경우
이 실험은 Trajectory Transformer에서도 했는데다 두 논문에서의 대조군인 CQL성능이 거의 비슷해서 두 논문의 성능을 간접적으로 비교 가능함.
두 논문에서 차이가 가장 나는 부분은 replay buffer의 데이터인 "Medium-Replay(Decision Transformer) = mixed(Trajectory Transformer)"인데, Decision Transformer는 대체로 데이터셋의 최고성능보다 크게 뛰어난 성능을 못내는반면 Trajectory Transformer는 그 이상의 성능을 내는 모습을 보여줌.
실제로 Decision Transformer에서 return을 조절해가며 얼마나 잘 복원되는지를 분석한 결과, 아래와 같이 데이터의 분포를 있는 그대로 잘 학습했으나 데이터셋 최고성능 이상의 return을 조건으로 주었을땐 대부분의 환경에서 데이터셋의 성능을 upper limitation으로 갖는것을 확인함.
또한 과거의 context가 얼마나 중요한지를 확인함.
Trajectory Transformer에선 Humanoid 환경에서 비슷한 실험을 했는데, long-term credit assignment가 상대적으로 안중요한 환경이어서인지 context의 길이가 크게 중요하지 않았음.
Long-term credit assignment가 특히나 중요한 Key-to-Door 환경에서 실험을 해준 부분이 좋았는데, 아래와 같이 과거 context에 따라서 예측되는 reward가 크게 달라지는것을 볼 수 있으며, transformer 아키텍처의 장점인 attension map으로 이를 해석적으로 확인함.
또한 sequence modeling 방식의 접근은 reward의 density가 상대적으로 덜 중요함에 따라 기존 bellman backup 및 policy gradient방식과 대비하여 sparse reward에서 큰 장점을 가지는것을 보여줌.
마지막으로 왜 pessimism이 필요 없는지에 대해선 Trajectory Transformer와 같이 value function을 approximate하는 과정에서 오는 오차가 없기 때문이라는 공통된 의견을 냄
개인적인 생각
위에서 중간중간 언급했다시피, 이렇게 데이터셋의 action을 생성하는 접근으론 stitching은 되겠지만 사실상 imitation learning에서 크게 벗어나지 않는 접근이 되어 데이터셋의 퀄리티에 영향을 받을수 밖에 없는것 같다.
때문에 transformer 아키텍처를 RL에 활용하는데엔, Trajectory Transformer 처럼, pre-trained language모델의 데이터에 대한 뛰어난 representation power에 집중하여 world 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로 구성됨
같은 방법으로 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 방법을 사용함
$\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 할수 있음.
이때 마지막 $\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로 새롭게 접근하는 방법들이 많이 나올 수 있을것 같다.
