[CS25 4강] Decision Transformers: Reinforcement Learning Via Sequence Modeling

Decision Transformer를 설명하기 앞서 우리가 왜 이에 대해 얘기하는지 동기를 부여해보겠다. Transformer가 인공지능의 많은 분야에 큰 영향을 미쳤고, NLP, Vision, 단백질 폴딩 문제 등을 해결하였다. 그래서 이 모든 발전이 의사결정이나, 인공 지능을 위한 통합 모델에 가까워지고 있는 것 처럼 보이지만, 인공지능은 단지 perception을 가지는 것 뿐 아니라 perception knowledge가 의사결정을 위해 사용하는 방법에 대해 생각해봐야 한다. Transformer는 RL 모델과 달리 확장에 용이하며 매우 안정적인 training dynamics를 가지고 있다. 따라서 우리는 더 큰 모델과 리소스를 사용할 수 있으며 더 풍부한 분포를 학습할 수 있다. 즉, training 전반적인 절차에 안정성을 가져오고 확장 가능하다는 것은 중요하며 이것이 중요한 이유에 대해 생각해야한다. sequential decision making은 매우 광범위한 영역이며 해당 강연에서 특히 집중할 내용이다.

 

Sequential Decision Making, 즉 RL(강화학습)의 background에 대해 설명하도록 하겠다. RL은 agent 현재 state를 받고, action을 취함으로써 environment와 상호작용하며 action을 환경에 취하면 환경은 다음 state와 해당 action이 얼마나 좋은지에 대한 reward를 주고 이 전체 feedback loop는 intelligent agent의 목표를 계속 수행할 것이다. 그리고 시행착오(try and error)를 통해 다양한 행동을 시도해보고 어떤 action이 어떤 reward로 이어질지 확인하여 state를 action에 매핑하는 policy를 학습한다. 이 policy는 시간 범위 내에서 에이전트의 누적 보상을 최대화하므로 일련의 action을 취하고, aciton sequence에 대해 누적된 보상에 따라 policy가 얼마나 좋은지 판단한다. 이 강연은 offline RL에 대해 특별히 초점을 맞춘다. Offline RL은 환경과 직접 상호작용하며 policy를 학습하는 online RL과 달리 log data를 이용해 agent를 훈련하는 것이다. 이렇게 되면 try and error component가 없으므로 framework를 확장하는 것이 용이하고 이는 매우 흥미로운 점이다.

 

RL policy를 학습하는데 있어, LLM을 학습하는 것보다 더 작은 모델을 사용하지만(NLP와 다르게 task마다 모델을 따로 학습하므로) training이 매우 unstable하다. 현재 RL 알고리즘은 대부분 dynamic programming기반으로 구축되고 있어 매우 불안정한 내부 루프 최적화 문제를 포함하고 있다. 위 그림과 같은 매우 불안정한 reward curve를 보이는 것은 매우 일반적이고 따라서 복잡한 환경에서 에이전트가 성능을 안정적으로 달성할 수 있도록 알고리즘을 개선해야 한다. 

 

Decision Transformer는 매우 심플한 목적을 가지고 있다. state, action, reward 즉 trajectory를 받고 causal transformer에서 만드는 action을 output한다. 즉 이전에 오는 모든 토큰(state, action, reward of preceding timesteps)을 condition으로 받고 action을 out하는 gpt이다. 그리고 여기서 reward는 일반적인 reward가 아닌 return-to-go(앞으로 받을 누적 reward합)이다. 따라서 R_0은 전체 sum of reward for the entire trajectory 이다. 그리고 R_1은 0 step에 받은 reward에 전체 reward 합을 뺀 값이 된다.

 

 

그리고 하나의 토큰이나 패치(nlp, vision)을 받는 일반적인 transformer와 달리 state, action, reward 이렇게 3개의 토큰을 받는다. 그리고 loss function은 continous action에 대해서는 action sequences의 mse이고, discrete action에 대해서는 action sequences의 cross entropy를 사용한다. 이 objective는 매우 안정적이고 regularize가 쉽고 dynamic programming이 필요 없다는 점이 매우 주목할만한 점이다.

 

DT framework는 RL이 probabilistic inference로 알려져 있다는 것으로 연결지어 인스턴스화할 수 있다. 다음 state를 결정하는데 사용되는 state와 action이 있고, optimality를 인코딩하기 위해 일반적으로 auxiliary variables(O)가 있는 그래피컬 모델문제로 RL을 fomulate할 수 있다. 이러한 auxiliary 변수는 implicitly 보상 개념을 인코딩하고, 이 optimality가 true라는 조건이 있음을 암시적으로 의미한다. RL은 최적의 행동을 얻을 수 있도록 상태에서 동작을 매핑하는 정책을 학습하는 작업이다. 따라서 DT에서 Optimality variable은 encoded returns-to-go가 된다. 이러한 optimality variable이 DT에서 encoder returns-to-go가 되는 것은 rollout을 하는 동안 expert returns과 같이 충분히 높은 값을 제공하면 본질적으로 이것이 최적성의 mathematical quantification이라는 조건이 충족되기를 희망하여 DE를 rollout하면 해당 조건이 충족될 수 있다고 희망하는 것이다.

자세하게 얘기하면 optimality variable이라고 하는 것은 문제를 푸는 goal이다. 예를 들어서 binary라고 하면 1이 optimality variable이 된다. 따라서 test time에 encoder return-to-go를 주면 true 즉 optimal action을 generate하게 되는 것을 희망하는 것이다. (하지만 실제로는 DT에서 이 return to go가 잘 동작하지 않는 것으로 알고 있다.) 물론 우리의 learning은 완벽하지 않으므로 return to go를 주는 것은 gurantee를 보장하지 않는다. 하지만 transformer가 returns-to-go를 some notion of optimality를 interprete하도록 학습하는 것이다. binary case 같은 경우는 가능하겠지만, 우리는 continous variable을 배우기 때문에 이가 optimal인지 non-optimal인지 알기 쉽지 않다.

 

그리고 DT가 주목을 받은 이유는 간단한 모델 구조와 직관적인 프레임워크에 비해 실험에 대한 성능이 성공적으로 보이기 때문이다. model-free offlline RL 모델에 버금가는 성능을 보였다.

 

그리고 test time에 return conditioning를 따라가도록 evaluation이 잘 되었다고 얘기하며 (하지만 실제로는 이 return to go가 유효하진 않다고 한다.) return conditioning을 통해 multi-task가 가능하다고 한다. 그리고 Seaquest task 결과를 보면 모델이 extrapolate를 할 수 있다고 하는데, 즉 training 데이터셋에 있는 성능을 능가할 수 있는 policy를 모델이 찾을 수 있다고 주장하는 것이다. 각 다른 파트의 trajectory를 stitching하여 extrapolation이 가능해 offline dataset으로 training하는 경우 가질 수 있는 sub-optimal 문제를 완화할 수 있다고 하며 이는 CQL같은 baseline은 할 수 없는 unique한 것이라고 얘기한다.

그리고, dynamics가 stochastic인 경우, 즉 같은 state에서 같은 action을 했을 때 같은 next state가 보장되지 않는 경우는 모델이 corrleation을 잡기 어려워 성능이 저하될 수 있다.  

그리고 LSTM과 같은 sequence model은 large sequence를 처리할 수 있다는 점에서, markovian에 대한 장점이 작용할 것으로 보이지만 sparse reward환경 같은 경우는 context를 잘 인지하는 transformer가 더 도움이 된다.

 

 

Key-to-Door 환경은 1단계에서 수집한 열쇠를 이용해 문을 열어야 하는 환경이다. 1단계에서 열쇠를 잡았는지, 3단계에서 문을 열었는지 여부에 따라 0또는 1의 보상을 받게 되고 과거에 일어난 이벤트가 미래에 영향을 미치는 credit assignment 문제에 적합한 환경이다. critic은 일반적으로 long-term 누적 보상합에 있어 policy가 좋은 action을 취했는지 평가한다. 만약 DT가 critic의 역할을 잘 할 수 있도록 훈련되었다면 action을 취하는 것 대신, 즉 action이 output이 되는 대신, reward를 out한다면 우리는 같은 causal transformer를 사용해 이전 transition을 보고 reward를 선택하려고 시도할 수 있다. 그리고 실험을 통해 key-to-door의 3가지 phrase마다 reward probability가 매우 많이 변하는 흥미로운 패턴을 발견할 수 있었다. 1단계에서 열쇠를 주면 보상 확률은 거의 비슷하게 시작하지만 에이전트가 열쇠를 집지 않을 것으로 보이면, 보상이 줄어들기 시작하고, 이후 단계에서 열쇠가 없으면 문을 열고 나가 보상을 받을 확률이 거의 없어지기 때문에 에피소드 내에서 보상확률이 거의 0에 가깝게 유지되는 것을 볼 수 있다. 이는 DT가 단순히 좋은 actor일 뿐 아니라 long-term period와 assignment에서 매우 인상적인 critic이 될 수도 있음을 시사한다. 하지만 action을 out하도록 모델을 만든 것은 reward를 out하는 구조로 바꾼다고 해서 performance improvement가 일어나지는 않았기 때문이라고 한다.

 

 

요약하자면, DT는 RL을 Sequence modeling 문제로 해결한 첫번째 시도였으며 이전 접근법들에 비해 많은 장점들이 있다. 디자인이 심플하고, 확장에 용이하여 RL scaling에 큰 도움이 될 것이다. 

 

DT는 RL 에이전트의 Scaling과 Generalization을 구축하는 방법을 재고하는 첫 번째 작업이다. 이를 기반으로 Multimodality, Multi-task, Multi-agent에 활용가능할 것으로 보이고 이는 future work로 남겨두겠다.

 

추가적으로 한 학생분이 질문으로 RL literature에서 Transformer가 더 많이 등장할 것 같냐고 질문하였는데, 발표자분은 더 많은 work가 쏟아져 나올 것 같다고 생각한다고 했다. 하지만 transformer를 활용한 작업은 매우 많지만 중요하게 해결해야할 작업만 완료된 것일뿐이며 exploration같은 non-trivial한 문제를 해결해야 한다고 한다. 발표자의 생각으로는 transformer의 loss function에 exploration을 위한 텀을 두면 성능이 크게 저하되는지 확실하게 밝혀진 바가 없지만 이는 매우 까다로우며, 따라서 아직까지는 transformer의 loss function은 cross entropy loss를 사용하고 있다. 우리가 이 병목현상을 넘어야 하며 넘을 수 없는 것이 RL의 미래라고 믿고 싶지 않다고 한다.