Decision Transformer: Reinforcement Learning via Sequence Modeling
결정 트랜스포머: 시퀀스 모델링을 통한 강화학습
Lili Chen, Kevin Lu, Aravind Rajeswaran, et al. (2021)
강화학습 문제를 시퀀스 모델링 문제로 재구성하여, Transformer가 (리턴, 상태, 행동)의 시퀀스를 조건부로 생성함으로써 기존 RL 알고리즘 없이도 오프라인 RL 과제를 수행할 수 있음을 보여주었다.
배경
기존 오프라인 강화학습(Offline RL)은 사전 수집된 데이터셋에서 최적 정책을 학습해야 하는데, 분포 이동(distribution shift) 문제로 인해 과대추정된 Q-값이 불안정한 학습을 초래하는 근본적 어려움이 있었다. CQL, IQL 등의 보수적 알고리즘이 제안되었지만 복잡한 제약 조건이 필요했다. 한편 GPT 시리즈의 성공은 시퀀스 예측의 놀라운 잠재력을 보여주었다.
핵심 아이디어
Decision Transformer는 강화학습을 '원하는 리턴을 달성하는 행동 시퀀스를 생성하는 조건부 시퀀스 모델링'으로 재정의한다. 에이전트의 궤적을 (R̂₁, s₁, a₁, R̂₂, s₂, a₂, ...) 형태의 시퀀스로 표현하되, R̂ₜ는 해당 시점 이후의 목표 리턴(return-to-go)이다. GPT 아키텍처의 인과적(causal) Transformer가 이 시퀀스에서 다음 행동을 자기회귀적으로 예측한다. 테스트 시 높은 목표 리턴을 조건으로 주면 모델이 해당 리턴을 달성하는 행동을 생성한다. 벨만 방정식, 정책 그래디언트, 행동 복제 같은 기존 RL 프레임워크를 전혀 사용하지 않는다.
방법론
오프라인 데이터셋의 궤적을 (return-to-go, 상태, 행동) 트리플렛 시퀀스로 변환한다. 각 모달리티에 대해 선형 임베딩 레이어를 적용하고, 타임스텝 임베딩을 더한 후 GPT-2 스타일의 Transformer에 입력한다. 최근 K 타임스텝의 컨텍스트를 유지하며, 행동 예측에 대한 교차 엔트로피(이산 행동) 또는 MSE(연속 행동) 손실만으로 학습한다.
주요 결과
Atari(이산 행동)와 D4RL(연속 행동) 벤치마크에서 CQL 등 기존 오프라인 RL 알고리즘과 동등하거나 우수한 성능을 달성했다. 특히 보상이 희소한(sparse reward) 환경에서 기존 TD 학습 기반 방법보다 크게 우수했으며, 이는 장기 신용 할당(long-term credit assignment)에서 시퀀스 모델링의 이점을 보여준다.
임팩트
Decision Transformer는 RL과 시퀀스 모델링의 경계를 허물어, 파운데이션 모델 시대에 RL 문제를 접근하는 새로운 패러다임을 제시했다. 이후 Trajectory Transformer, GATO, RT-2 등 시퀀스 모델링 기반 의사결정 연구의 폭발적 성장을 이끌었으며, LLM의 인컨텍스트 학습 능력을 의사결정에 활용하는 연구 방향을 열었다.