[ICLR 2022] 3편: Reinforcement Learning as a sequence modeling problem

‘ICLR(The International Conference on Learning Representations, 표현 학습 국제 학회)’은 2013년에 최초로 개최되어 올해로 10주년을 맞은 딥러닝 분야 세계 최고 권위 학회입니다. 인공지능, 데이터 사이언스, 비전, 음성인식 그리고 로보틱스에 이르기까지 딥러닝을 활용하는 모든 분야에 걸친 최첨단의 연구와 논문을 나누는 공유의 장이기도 합니다. 지난 4월 온라인으로 진행된 ‘ICLR 2022’에서 소개된 논문을 LG AI연구원이 소개합니다. - 1편: Efficient Transformer – Applied AI Research Lab 김정희 님 - 2편: Time Series Anomaly Detection – Applied AI Research Lab 심예슬 님 - 3편: Reinforcement Learning as a sequence modeling problem – Applied AI Research Lab 윤든솔 님

강화학습은 연속적 의사 결정 문제를 해결하는 방법으로, 에이전트(agent)가 환경(Environment)에서 현재 state(상태)를 보고 최적의 action(행동)을 수행하여 cumulative reward(보상합)를 최대화할 수 있도록 하는 에이전트의 행동 policy(정책)를 찾는 방법론입니다. 에이전트는 환경과 상호작용하며 다양한 행동을 수행해 reward(보상)라는 행동에 대한 feedback을 받으며 최적의 정책을 찾아갑니다. 아래의 Figure 1은 일반적인 강화학습을 도식화한 것입니다.

(Figure 1) 일반적인 강화학습 구조

최근 강화학습은 딥러닝과 결합하여 심층 강화학습이라는 이름으로 다양한 문제에서 좋은 성능을 보여주고 있습니다. 그러나 강화학습이 매우 유망한 기술임에도 비전이나 자연어 처리만큼 활발하게 적용되고 있지는 않습니다. 그 이유 중 하나는 시뮬레이터의 부재입니다. 심층 강화학습에서 에이전트를 학습시키기 위해서는 많은 양의 데이터가 필요하지만, 현실 세계에서 다양한 행동을 직접 실행해 보기에는 위험성이 높아 데이터를 얻기 어렵습니다. 따라서 데이터를 수월하게 모을 수 있는 시뮬레이터가 필요하지만, 풀고자 하는 현실 세계의 문제들에는 대부분 시뮬레이터가 존재하지 않습니다.

이러한 문제를 해결하기 위해 최근 오프라인 강화학습(Offline RL)이 각광받고 있습니다. 오프라인 강화학습은 기존 강화학습과 다르게 환경과의 상호작용 없이 주어진 데이터를 통해 에이전트를 바로 학습시키는 방법론입니다. 현실 세계에서 에이전트가 직접 행동해가며 데이터를 얻어내기는 힘들지만, 사람들이 만들어낸 데이터는 많습니다. 이러한 데이터를 활용하여 에이전트를 바로 학습시킬 수 있다면 이를 실제로 적용하기가 훨씬 수월할 것입니다. 아래의 Figure 2는 오프라인 강화학습을 도식화한 것입니다.

(Figure 2) 오프라인 강화학습 구조

작년 6월 오프라인 강화학습 분야에서 Decision Transformer(DT)라는 새로운 패러다임이 제안되었습니다. DT에서는 강화학습 문제를 자연어 처리에서 널리 활용되는 sequence modeling 문제로 바꾸어 풀어냈습니다. 특히 DT는 transformer 구조를 강화학습에 성공적으로 적용하고, 강화학습 에이전트를 supervised learning으로 학습시켰다는 점에서 강화학습 분야의 관심을 끌었습니다. 이번 글에서는 최근 강화학습을 supervised learning으로 학습시키는 방법론을 유행시킨 DT와 이를 확장한 모델인 Generalized DT를 소개합니다.

Decision Transformer: Reinforcement Learning via Sequence Modeling

기존 강화학습에서는 에이전트의 정책을 학습시키기 위해 action-value function(행동-가치 함수)을 활용합니다. Action-value function Q(s,a)란, 주어진 state s에서 action a를 했을 때 얻을 수 있는 보상합에 대한 기댓값입니다. 즉 해당 행동이 현재 상태에서 얼마큼 값어치가 있는지를 표현했다고 볼 수 있습니다. 그리고 action-value function으로부터 상태 s에서 어떠한 행동을 해야 하는지를 알려주는 정책 policy π(s)를 얻어낼 수 있습니다.

수십 개의 neural network layers를 쌓는 비전이나 자연어 처리와 달리, 일반적으로 심층 강화학습에서는 두, 세 가지 정도로 적은 neural network layers를 쌓아 action-value function 및 policy를 표현합니다. 그 이유는 action-value function을 안정적으로 학습시키기가 매우 어렵기 때문입니다. Action-value function의 학습을 불안정하게 만드는 주된 이유 중 하나는 action-value function을 학습시킬 때 활용되는 Temporal Difference(TD) learning입니다. TD learning은 아래의 식과 같습니다.

TD learning에서는 현재 상태 s와 a에 대한 action-value function를 학습시킬 때 다음 상태 s’에 대한 action-value function을 사용하여 target을 만듭니다. 이때 현재 학습 중인 불안정한 action-value function가 target에 활용되기 때문에 그 결과가 불안정한 학습으로 이어집니다. 이에 TD learning을 활용하여 action-value function을 학습시킨 기존의 심층 강화학습에서는 layer를 깊게 쌓게 되면 학습이 더욱더 불안정해질 수 있어 복잡한 구조의 layer 활용을 기피해왔습니다.

What is DT?

본 연구에서는 강화학습 문제를 sequence modeling 문제로 보고 supervised learning을 활용하여 policy를 학습시키는 새로운 방법론을 제안합니다. Decision Transformer(DT)^[1]는 TD learning을 활용하지 않아 transformer와 같은 복잡한 구조의 layer를 활용하면서 안정적인 policy를 학습시킬 수 있었습니다. Figure 3은 Decision Transformer의 작동 과정입니다.

(Figure 3) 본 논문에서 제안한 Decision Transformer 작동 과정

기존 심층 강화학습에서는 action-value function을 학습하여 각 행동이 얼마나 가치 있는지를 측정하고, 이를 바탕으로 주어진 상황에서 어떠한 행동을 해야 하는지를 결정하는 policy를 학습했지만, DT에서는 action-value function이 존재하지 않습니다. 대신, 해당 상태에서 원하는 가치를 얻고 싶을 때 어떠한 행동을 해야 하는지를 바로 학습합니다. 그리고 이를 학습시키기 위해 autoregressive model인 transformer를 활용합니다

 

How to train?

DT를 학습시키는 방법은 자연어 처리에서 seq2seq 모델을 학습시키는 방법과 매우 유사합니다. State, action, reward의 sequence로 구성된 trajectory $\left({\mathrm{s}}_1,a_1,r_1,\dots ,s_T\mathit{,}{a}_T\mathit{,}{r}_T\right)$가 주어졌을 때, 각 시점 t에서 에피소드 끝까지의 cumulative reward인 return-to-go $\overbrace{{\mathrm{R}}_{\mathrm{t}}}=\sum _{t\text{'}=t}^T{r}_{t\text{'}}$를 구합니다. 그리고 trajectory에 return-to-go를 추가하여 $\widehat{{\mathrm{R}}_1},s_1,a_1,\dots ,\widehat{{\mathrm{R}}_T},s_T,a_T$와 같이 재구성한 후 K의 길이를 갖는 subsequence를 활용하여 학습시킵니다. 즉 transformer의 입력으로 $\left(\widehat{{\mathrm{R}}_{t-K-1}},s_{t-K-1},a_{t-K-1},\dots ,\widehat{{\mathrm{R}}_t},s_t\right)$가 주어지면 $a_t$를 예측하도록 학습됩니다. 이때 행동이 discrete라면 cross entropy loss를, continuous라면 mean square error loss를 활용합니다. 아래는 DT를 학습시키는 코드입니다.

일반적인 강화학습에서는 return-to-go를 출력으로 하는 action-value function Q(s,a)로부터 policy를 얻습니다. 하지만 DT에서는 conditional policy를 활용하여 return-to-go가 policy의 추가적인 입력으로 들어가게 됩니다. 따라서 policy는 현재 s에서 return-to-go를 얻기 위해서는 a라는 행동을 해야 한다는 것을 바로 학습합니다

 

How to inference?

학습 후 inference 또한 seq2seq 모델과 유사합니다. 아래의 코드와 같이 초기 상태 s에서 원하는 return-to-go 값을 넣어 준 후 action a를 뽑아냅니다. Action a를 환경에 적용하고 환경으로부터 나온 다음 state s’과 reward r를 받아 s’에서의 return-to-go를 다시 구해준 뒤 DT에 넣어주어 다음 상태 s’에서의 action a’를 뽑아냅니다.

Result

논문에서 제안한 모델인 DT는 각각 Figure 4와 5에서 볼 수 있듯 Offline RL에서 널리 사용되는 D4RL benchmark와 Atari benchmark에서 높은 성능을 보여주었습니다.

(Figure 4) DT의 D4RL benchmark에서의 성능

(Figure 5) DT의 Atari benchmark에서의 성능

본 연구는 offline RL 문제를 sequence modeling 문제로 바라보는 새로운 패러다임을 제안하였습니다. 이를 통해 기존 오프라인 강화학습의 알고리즘들과 달리 supervised loss를 활용하여 모델을 학습하였고, 강력하지만 복잡한 구조를 갖는 transformer를 성공적으로 적용하였습니다.

GENERALIZED DECISION TRANSFORMER FOR OFFLINE HINDSIGHT INFORMATION MATCHING

해당 연구는 기존 DT를 일반화시켜 확장한 연구로 2022 ICLR에서 발표되었습니다. Generalized Decision Transformer(GDT)^[2]는 policy에 future trajectory information을 condition하는 DT뿐만 아니라 기존의 다양한 연구들이 모두 hindsight information matching (HIM) 문제로 묶일 수 있다는 새로운 framework를 제안합니다.

State marginal matching

어떠한 정책 π를 rollout 했을 때 상태 s에 방문할 확률인 state marginal distribution은 아래와 같이 정의될 수 있습니다.

${\mathrm{p}}_t^{\pi }\left(s\right)={\iint }_{s_{0:t},a_{0:t-1}}^{}\prod _t^{}p\left(s_t\overline{)s_{t-1},a_{t-1}}\right)\pi \left(a_t{\overline{)s}}_t\right)$
${\mathrm{p}}^{\pi }\left(s\right)=\left(1-\gamma \right)\sum _{t\text{'}}{\gamma }^{t\text{'}}{\mathrm{p}}_{t\text{'}}^{\pi }\left(s_{t\text{'}}=s\right)$

State marginal matching (SMM)은 policy의 state marginal distribution ${\mathrm{p}}^{\pi }\left(s\right)$과 target의 state marginal distribution ${\mathrm{p}}^{*}\left(s\right)$ 사이를 최소화하는 정책 $\mathrm{\pi }$를 찾는 objective이며 아래와 같이 정의될 수 있습니다.

D는 KL divergence와 같은 distribution divergence measure이며, SMM objective는 주로 target을 expert demonstration으로 두어 expert를 모방하는 정책을 찾는 imitation learning 쪽에서 많이 활용되는 objective입니다.

Hindsight Information matching

저자는 hindsight experience replay (HER)^[3], temporal difference model (TDM)^[4], Learning from play (LfP)^[5], return-conditioned RL^[1]과 같이 future trajectory information을 match 하는 trajectory를 생성할 수 있는 policy를 학습시키는 다양한 연구들이 hindsight information matching (HIM)을 풀고 있다고 합니다.

Hindsight information matching을 정의하기 전에 추가로 정의할 내용이 몇 가지 있습니다.

- Information statistics $\mathrm{I}\left({\tau }_t\right)$는 partial trajectory ${\mathrm{\tau }}_{\mathrm{t}}=\left(s_t,a_t,s_{t+1},\dots ,s_T,a_T\right)$의 statistical property를 capture 하는 함수입니다. 예를 들어 mean과 variance를 구하는 함수가 될 수도 있습니다.
- Feature function $\mathrm{\varphi }\left(·,·\right)$은 state s와 action a를 입력으로 받아 어떠한 feature를 출력하는 function입니다. 예를 들어 feature function $\mathrm{\varphi }$은 reward function r(s,a) 가 될 수도 있고 state의 sub-dimension이 뽑아내는 함수가 될 수도 있습니다.
- State와 action을 feature function으로 표현한 partial trajectory ${\mathrm{\tau }}_{\mathrm{t}}^{\mathrm{\varphi }}=\left(\mathrm{\varphi }\left({\mathrm{s}}_{\mathrm{t}},{\mathrm{a}}_{\mathrm{t}}\right),\dots ,\mathrm{\varphi }\left({\mathrm{s}}_T,{\mathrm{a}}_T\right)\right)$가 주어졌을 때, 이에 대한 information statistics는 ${\mathrm{I}}^{\mathrm{\varphi }}\left({\mathrm{\tau }}_{\mathrm{t}}\right)$로 표현됩니다.

Hindsight information matching은 z를 추가적인 입력으로 받는 conditional policy $\mathrm{\pi }\left(\mathrm{a}\overline{)\mathrm{s},\mathrm{z}}\right)$가 어떠한 information statics value z를 만족하면서 rollout 될 수 있도록 학습시키는 문제이며 아래와 같이 정의됩니다.

$\underset{\pi }{\mathrm{m}\mathrm{ax}}{E}_{z~p\left(z\right),r~p_z^{\pi }\left(r\right)}\left[-D\left({\mathrm{I}}^{\mathrm{\varphi }}\left(\mathrm{\tau }\right),z\right)\right]$

모방하고 싶은 trajectory가 주어졌을 때, 해당 trajectory의 s와 trajectory로부터 계산되는 ${\mathrm{I}}^{\mathrm{\varphi }}\left(\tau \right)$를 z로 두어 conditional policy의 입력으로 학습시키면 policy는 우리가 모방하고 싶은 trajectory를 따라가는 policy가 될 수 있습니다. 이때 feature function $\mathrm{\varphi }$를 reward function으로 두고 $\mathrm{\gamma }$-discounted summation을 aggregate function으로 활용하면 information statistics는 ${\mathrm{I}}^{\mathrm{\varphi }}\left({\tau }_t\right)=\sum _{\mathrm{t}}{\mathrm{\gamma }}^{\mathrm{t}}{\mathrm{r}}_{\mathrm{t}}$로 정의되면서 DT가 됩니다. 그리고 Feature function $\mathrm{\varphi }$이 state의 position을 출력하는 함수이고 aggregate function이 단순히 trajectory ${\mathrm{\tau }}_t^{\varphi }$의 마지막을 뽑아내는 함수라면 information statistics ${\mathrm{I}}^{\mathrm{\varphi }}\left({\mathrm{\tau }}_t\right)={\mathrm{\varphi }}_{\mathrm{T}}$가 되고 HER과 같은 goal conditioned RL이 됩니다. 이처럼 feature function과 information statistics에 따라 다양한 variant들이 나올 수 있으며, 이러한 framework를 Generalized DT라 말합니다. 아래 Figure 6의 테이블은 기존의 다양한 연구들이 어떻게 GDT framework로 묶일 수 있는지 보여줍니다.

(Figure 6) 이전의 다양한 연구들과 GDT framework의 연결

(Figure 7) 본 논문에서 제안한 GDT의 작동 과정

저자는 policy에 target trajectory의 information statistics가 주어졌을 때, target trajectory를 생성하는 문제인 offline multi-task state-marginal matching (SMM)과 offline multi-task imitation learning (IL)을 효과적으로 풀 수 있는 GDT 기반의 새로운 variants인 Categorical-DT와 Bidirectional-DT를 제안합니다. 두 variants는 기존의 DT의 구조를 아주 조금만 바꾸면 얻을 수 있는 모델입니다. 제안한 모델들을 소개하기에 앞서, 해당 논문이 풀고자 하는 문제인 offline multi-task SMM과 IL에 대해 간략하게 설명해 드리겠습니다.

SMM과 IL은 train 중 보지 못한 unseen target trajectory의 information statistics ${\mathrm{I}}^{\mathrm{\varphi }}\left(\mathrm{\tau }\right)$가 주어졌을 때, 학습된 policy가 target trajectory를 얼마나 잘 생성하는지(따라가는지)를 측정하는 distribution matching 문제입니다. 이때 저자는 feature function이 reward function이나 sub-dimension을 뽑아내는 function과 같이 특정한 function으로 주어졌을 경우 SMM 문제로 정의하였고, feature function이 identity function ($\mathrm{\varphi }$=s)일 경우, 즉 target trajectory의 ${\mathrm{I}}^{\mathrm{\varphi }}\left(\mathrm{\tau }\right)$가 단순히 full state trajectory로 주어졌을 때, 이를 IL 문제로 정의하였습니다.

Categorial DT and Bidirectional-DT

최근 distributional RL^[6]은 기존의 일반적인 RL보다 다양한 분야에서 좋은 성능을 보였습니다. 일반적인 RL에서는 Q function을 학습시킬 때 target인 expected discounted reward sum을 scalar 값으로 두어 예측하도록 학습시킵니다. 이와 다르게 distribution RL에서는 target의 expected return 값이 아닌 distribution 자체를 예측하도록 학습시킵니다.

CDT는 scalar 값이 아닌 distribution 자체를 학습에 활용한다는 아이디어를 DT에 적용한 variant입니다. DT에서 information statistics ${\mathrm{I}}^{\mathrm{\varphi }}\left({\mathrm{\tau }}_{\mathrm{t}}\right)$는 $\sum _t{\gamma }^{t}r\left(s_t,a_t\right)$였다면, CDT에서는 histogram$\left({\mathrm{r}}_{\mathrm{t}},\gamma \right)$로 정의됩니다. CDT는 Distributional RL과 같이 continuous value를 갖는 reward를 discretize하여 categorical distribution으로 표현한 뒤 이를 information statistics로 두어 모델을 학습시킵니다.

BDT는 DT에 한 층의 transformer를 더 쌓은 구조입니다. 새로 추가된 transformer는 순서가 뒤집힌 state sequence를 입력으로 받아 z를 출력하며, 이는 transformer policy의 information statistics로 활용됩니다. 즉 이전 모델들은 미리 정의된 feature function과 aggregate function에 의해 information statistics가 계산되었지만, BDT에서는 이를 학습 가능한 function으로 대체시켰습니다.

Result

저자는 D4RL benchmark에서 offline multi-task SMM/IL을 위한 metric을 정의하며 새로운 benchmark tasks을 제안하였고, CDT와 BDT를 해당 tasks에서 평가하였습니다. Target의 state marginal distribution과 학습된 policy의 state marginal distribution을 정확히 알 수 있다면 density를 통해 거리를 비교하는 것이 가장 정확하지만, 우리는 단순히 샘플에만 접근할 수 있습니다. 따라서 저자는 샘플 및 discretization을 활용하여 두 continuous distribution을 categorical distribution으로 근사시킨 후 거리를 구합니다.

아래의 테이블은 offline multi-task SMM에서의 CDT의 성능을 보여줍니다. Figure 8의 table은 feature function이 state의 sub-dimension인 x-velocity를 가져오는 function일 때이며 Figure 9의 table은 feature function이 reward function일 때의 성능을 나타냅니다. 두 가지 케이스 모두에서 CDT가 다른 모델들보다 target distribution을 잘 매칭할 수 있다는 것을 확인할 수 있습니다.

(Figure 8) feature function이 x-velocity일 때 CDT의 성능

(Figure 9) feature function이 reward일 때 CDT의 성능

추가적으로 저자는 CDT의 generalization 성능을 확인 해보기위해 bi-modal distribution matching과 meta RL과 IL에서 많이 사용되는 cheetah-velocity matching에서도 실험을 진행합니다. Bi-modal distribution을 갖는 trajectory를 만들기 위하여 저자는 앞으로 달리는 치타의 trajectory와 뒤로 백덤블링을 하는 치타의 trajectory를 반으로 잘라 이어 붙여 bi-modal trajectory를 만들었습니다. 아래 Figure 10의 그래프는 target distribution, CDT의 distribution, DT의 distribution을 보여줍니다. DT는 bi-modal distribution을 거의 matching하지 못하는 반면 CDT는 target과 유사하게 matching하는 모습을 볼 수 있습니다. 또한 Figure 11의 테이블을 보면 CDT는 복잡한 multi-task 문제를 풀기 위하여 고안된 Meta-BC와 유사한 성능을 보여준다는 것을 알 수 있습니다.

(Figure 10) target, CDT, DT의 uni-modal과 bi-modal trajectory에서의 state의 sub-dimension의 distribution

(Figure 11) 비교 모델들과 CDT의 uni-modal과 bi-modal trajectory에서의 성능

Cheetah-velocity matching 문제는 다양한 속도로 달리는 치타 로봇의 trajectories를 학습 중에 보지 못한 속도로 달리게 했을 때 원하는 속도에 맞추어 얼만큼 잘 달리는지를 측정하는 문제입니다. 저자는 [0.0, 3.0] 사이에서 0.1 간격으로 존재하는 trajectories 중 0.5, 1.5, 2.5 속도의 trajectories를 test set으로, 나머지를 train set으로 활용하였습니다. 아래 Figure 12의 표와 같이 CDT는 기존 방법론에 비해 unknown target trajectory에 대해 좋은 generalization 성능을 보여줍니다.

(Figure 12) 비교 모델들과 CDT의 cheetah-velocity matching에서의 성능

아래 Figure 13의 테이블은 offline multi-task IL에서의 BDT의 성능을 보여줍니다. SMM과 달리 IL에서는 학습된 policy가 따라가고 싶은 target trajectory의 full state만이 주어집니다. 저자는 BDT와 DT의 또다른 variant들을 비교합니다. DT-AE는 auto-encoding을 활용하고, DT-CPC는 contrastive loss를 활용하여 state에 대한 feature function $\mathrm{\varphi }\left(\mathrm{s}\right)$을 unsupervised learning으로 학습시킨 variant이고, DT-E2E는 $\mathrm{\varphi }\left(\mathrm{s}\right)$를 단순히 DT의 supervised loss를 활용하여 학습시킨 variant입니다.

저자는 DT-AE와 DT-CPC를 학습시킬 때 세 가지 전략을 사용합니다. 첫 번째는 단순히 feature function을 unsupervised loss만으로 학습시킨 모델이고, 두 번째는 DT의 supervised loss와 함께 동시에 학습시킨 모델(joint)입니다. 그리고 세 번째는 unsupervised loss로 pretrain한 후 weight를 고정시킨 상태에서 DT의 supervised loss를 학습시킨 모델(frozen)입니다. 추가로 feature function의 출력으로 나온 state sequence를 summation으로 aggregate하여 단일 벡터로 만들어주고, 이를 DT의 z로 넣어주었습니다. BDT의 경우 transformer 구조를 활용하기 때문에 summation과 같은 aggregate function을 필요로 하지 않습니다. BDT는 다른 모든 learnable $\mathrm{\varphi }\left(\mathrm{s}\right)$를 활용하는 variants에 비해 IL에서 훨씬 좋은 distribution matching 성능을 보여줍니다.

(Figure 13) offline multi-task IL 에서의 BDT의 성능

Conclusion

지금까지 transformer를 활용하여 강화학습 문제를 sequence modeling 문제로 접근한 연구들에 대해 살펴보았습니다. 최근 위의 연구들을 더욱더 확장하여 large-scale generalist model을 만드는 연구가 활발하게 진행되고 있습니다. 기존의 DT나 GDT는 state와 action space를 공유하는 환경에서만 multi task가 가능하다는 한계를 지녔습니다. Multi-Game Decision Transformers^[7]에서는 state와 action space가 다른 46개의 Atari 게임에서 단일 모델로 학습을 성공시켰으며, A Generalist Agent^[8]에서는 Atari 게임뿐만 아니라 robot control, dialogue 및 image captioning까지 가능한 단일 모델을 제안하였습니다. 이처럼 Vision, NLP, 그리고 강화학습 간의 경계가 무너지며, 다양한 tasks를 수행할 수 있는 단일 거대 모델에 대한 연구가 시작되고 있습니다. 이와 같은 최근의 연구 방향성을 보면 사람과 같은 AI가 생각보다 빠른 시일 내에 나올 수도 있을 것 같습니다.