들어가며...

Offline RL은 고정 Dataset에서 Q-function이나 Policy를 학습했다. 이때도 Bellman equation과 Policy Gradient가 중심이었다.

Decision Transformer는 관점을 바꾼다.

Trajectory를 Return, State, Action으로 이루어진 Sequence로 보고, 다음 Action을 예측하는 Sequence modeling 문제로 풀자.

언어 Model이 앞의 Token을 보고 다음 Token을 예측하듯, Decision Transformer는 원하는 Return과 지금까지의 경험을 보고 다음 Action을 예측한다.

강화학습 문제를 없앤 것이 아니라 같은 Trajectory Data를 다른 방식으로 학습하는 것이다. 이번 글에서는 입력 Token의 순서와 Return-to-go가 실제 Training과 Inference에서 어떻게 사용되는지 살펴본다.

Trajectory를 Token으로 바꾸기

Offline Dataset의 한 Trajectory를,

$$\tau = (s_0,a_0,r_0,s_1,a_1,r_1,\ldots)$$

라고 하자.

Decision Transformer는 각 시간의 Return-to-go, State, Action을 묶는다.

$$(G_0,s_0,a_0,G_1,s_1,a_1,\ldots)$$

여기서 Return-to-go는 현재 시점부터 남은 Discounted reward의 합이다.

$$G_t = \sum_{k=t}^{T} \gamma^{k-t}r_k$$

Finite-horizon benchmark에서는 $\gamma=1$로 두어 남은 Reward의 단순 합을 사용하는 경우도 많다.

$$G_t = r_t+r_{t+1}+\cdots+r_T$$

Return-to-go는 무엇을 요청하는 값일까

State $s_t$만 입력하면 같은 State에서 Dataset에 여러 행동이 있을 때 어떤 행동을 따라야 할지 모호할 수 있다. Return-to-go $G_t$를 함께 넣으면 “이 정도의 남은 성과를 원할 때 어떤 Action이 이어졌는가”를 구분할 수 있다.

예를 들어 같은 출발점의 두 Trajectory가 있다고 하자.

$$\tau_1: G_0=3, \quad a_0=\text{LEFT}$$ $$\tau_2: G_0=10, \quad a_0=\text{RIGHT}$$

State만 보면 두 Action이 섞이지만 Return condition을 함께 보면 높은 Return의 Trajectory에서 사용된 RIGHT를 구분해 학습할 수 있다.

Return-to-go는 미래를 미리 안다는 뜻이 아니다. Training data에서는 Episode가 이미 끝났으므로 계산할 수 있고, Inference에서는 사용자가 원하는 목표 Return을 조건으로 넣는다.

Training은 Action 예측 문제다

Causal Transformer는 시간 $t$의 Action을 예측할 때 미래 Token을 볼 수 없다.

$$\hat a_t = f_\theta \left( G_0,s_0,a_0, \ldots, G_t,s_t \right)$$

Discrete action이라면 Cross-entropy를 사용할 수 있다.

$$L(\theta) = - \sum_t \log \pi_\theta \left( a_t \mid G_{\leq t},s_{\leq t},a_{<t} \right)$$

Continuous action이라면 Action prediction의 MSE나 확률 분포의 Negative log-likelihood를 사용할 수 있다.

이 Training 과정에는 명시적인 Bellman backup이 없다.

$$r+\gamma\max Q(s',a')$$

도 없고, 환경을 새로 Rollout해 Policy Gradient를 계산하지도 않는다. Offline Sequence Dataset에 대한 Supervised objective로 학습한다.

Inference에서는 Return을 줄여가며 사용한다

학습이 끝난 뒤에는 원하는 Return $G_0^{\text{target}}$을 먼저 정한다.

1. 목표 Return $G_0$과 현재 State $s_0$을 입력한다.
2. Model이 Action $a_0$를 예측한다.
3. 환경에서 Reward $r_0$와 Next state $s_1$을 받는다.
4. 받은 Reward만큼 남은 목표를 줄인다.
5. 새 Context $(G_1,s_1)$로 다음 Action을 예측한다.

$\gamma=1$인 단순한 경우,

$$G_{t+1} = G_t-r_t$$

이다.

목표가 $10$이고 첫 Reward가 $2$라면,

$$G_1=10-2=8$$

이 된다. 다음에는 “현재 State에서 앞으로 $8$을 더 얻고 싶다”는 Context를 준다.

원하는 Return을 크게 쓰면 더 잘할까

반드시 그렇지 않다.

Dataset의 최고 Return이 $10$인데 Inference에서 $100$을 넣으면, Model은 Training 중 보지 못한 Condition을 받는다. $100$이라는 숫자만으로 새로운 전략이 생기지 않는다.

$$G^{\text{target}} \notin \text{training support}$$

이면 Action prediction이 불안정할 수 있다. 일반적으로 Dataset에서 충분히 지지되는 Return 범위 안에서 Condition을 고르고 여러 값을 평가해야 한다.

Return conditioning은 원하는 성과를 요청하는 Interface이지, Planner가 목표를 반드시 달성한다는 보장이 아니다.

Decision Transformer가 매력적인 이유

긴 History를 사용할 수 있다

현재 관측만으로 State를 완전히 알 수 없는 Partially observable problem에서 앞의 State와 Action History를 Context로 사용할 수 있다.

다양한 Behavior를 한 Model에 담을 수 있다

낮은 Return과 높은 Return의 Trajectory를 조건으로 구분해 하나의 Model이 여러 수준의 행동을 표현할 수 있다.

Sequence modeling 도구를 사용할 수 있다

Transformer architecture, Masking, Large-scale supervised training 기법을 그대로 활용할 수 있다.

하지만 RL의 어려움이 사라진 것은 아니다

Decision Transformer도 Dataset의 한계를 그대로 가진다.

• 좋은 Trajectory가 Dataset에 없으면 그 행동을 배우기 어렵다.
• 원하는 Return과 실제 가능한 Return 사이에 Gap이 생길 수 있다.
• 긴 Sequence는 계산량과 Context length를 요구한다.
• Environment dynamics를 명시적으로 Search하지 않는다.
• 새로운 State에서 OOD generalization 문제가 남는다.

모든 Offline RL 문제가 Sequence modeling으로 더 쉬워지는 것도 아니다. Local reward와 Markov structure를 잘 이용하는 Value-based method가 더 Data-efficient할 수 있다.

Behavior Cloning과 무엇이 다른가

일반 Behavior Cloning은 State에서 Action을 예측한다.

$$\pi(a_t\mid s_t)$$

Decision Transformer는 Return과 History도 조건으로 사용한다.

$$\pi \left( a_t \mid G_{\leq t},s_{\leq t},a_{<t} \right)$$

이렇게 성공과 실패가 섞인 Dataset에서 원하는 Return에 따라 행동을 구분할 수 있다. 그래도 Dataset Action을 예측하는 Supervised learning이라는 점에서는 Imitation Learning과 가깝다.

이번 글에서 기억할 것

Decision Transformer는 Offline Trajectory를 Return-to-go, State, Action의 Token sequence로 바꾸고, 원하는 Return을 조건으로 다음 Action을 예측한다.

1. Return-to-go는 현재부터 Episode 끝까지 남은 Reward의 합이다.
2. Training에서는 이미 끝난 Trajectory로 $G_t$를 계산한다.
3. Inference에서는 목표 Return을 넣고 받은 Reward만큼 줄여간다.
4. 명시적인 Bellman backup이나 On-policy RL loop 없이 Action prediction loss로 학습한다.
5. Dataset에 없는 높은 Return을 입력한다고 새로운 능력이 자동으로 생기지는 않는다.

다음 글에서는 언어 Model의 답변을 사람 또는 AI의 Preference로 학습하는 RLHF, RLAIF, DPO를 하나의 흐름으로 연결한다.

참고 자료

• Chen et al., Decision Transformer: Reinforcement Learning via Sequence Modeling
• Levine et al., Offline Reinforcement Learning: Tutorial, Review, and Perspectives
• Sutton & Barto, Reinforcement Learning: An Introduction, 2nd Edition