ai theory

Reinforcement Learning Basics (18) — Decision Transformer Reads Trajectories Like Sentences

Junyoung Park · 2024-07-12 · 6 min

Introduction

Offline RL learns a Q-function or policy from a fixed dataset. Even there, the Bellman equation and policy gradient remain central.

Decision Transformer changes the point of view.

Treat a trajectory as a sequence of returns, states, and actions, and solve action selection as a sequence-modeling problem.

Just as a language model predicts the next token from preceding tokens, Decision Transformer predicts the next action from a desired return and the experience observed so far.

It does not eliminate the reinforcement-learning problem. It learns from the same trajectory data through a different formulation. This article examines the order of its input tokens and how return-to-go is used during both training and inference.

Turning a Trajectory into Tokens

Consider a trajectory from an offline dataset:

τ=(s0,a0,r0,s1,a1,r1,).\tau = (s_0,a_0,r_0,s_1,a_1,r_1,\ldots).

Decision Transformer groups the return-to-go, state, and action at every time step.

(G0,s0,a0,G1,s1,a1,)(G_0,s_0,a_0,G_1,s_1,a_1,\ldots)

Return-to-go is the sum of discounted rewards remaining from the current time step.

Gt=k=tTγktrkG_t = \sum_{k=t}^{T} \gamma^{k-t}r_k

Finite-horizon benchmarks often set γ=1\gamma=1 and use the plain sum of the remaining rewards.

Gt=rt+rt+1++rTG_t = r_t+r_{t+1}+\cdots+r_T
A trajectory becomes a sentence of return-to-go, state, and action tokens. A causal Transformer predicts the next action from the preceding tokens.

What Does Return-to-Go Ask For?

If the model receives only state sts_t, behavior may be ambiguous when the dataset contains several actions for the same state. Adding return-to-go GtG_t distinguishes which actions occurred when a particular level of remaining performance was desired.

Suppose two trajectories begin from the same point.

τ1:G0=3,a0=LEFT\tau_1: G_0=3, \quad a_0=\text{LEFT} τ2:G0=10,a0=RIGHT\tau_2: G_0=10, \quad a_0=\text{RIGHT}

With only the state, the two actions are mixed together. Once the return condition is included, the model can learn to distinguish RIGHT as the action taken in the high-return trajectory.

Return-to-go does not mean that the model knows the future in advance. During training, the episode has already ended, so we can calculate it. During inference, the user supplies a desired target return as a condition.

Training Is an Action-Prediction Problem

When a causal Transformer predicts the action at time tt, it cannot see future tokens.

a^t=fθ(G0,s0,a0,,Gt,st)\hat a_t = f_\theta \left( G_0,s_0,a_0, \ldots, G_t,s_t \right)

For discrete actions, we can use cross-entropy.

L(θ)=tlogπθ(atGt,st,a<t)L(\theta) = - \sum_t \log \pi_\theta \left( a_t \mid G_{\leq t},s_{\leq t},a_{<t} \right)

For continuous actions, we can use MSE on the action prediction or the negative log-likelihood of a probability distribution.

There is no explicit Bellman backup in this training process. There is no target of the form

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

nor does the model perform new environment rollouts to calculate a policy gradient. It learns with a supervised objective over an offline sequence dataset.

During Inference, the Remaining Return Is Updated at Every Step

After training, we first choose a desired return G0targetG_0^{\text{target}}.

  1. Provide the target return G0G_0 and current state s0s_0.
  2. The model predicts action a0a_0.
  3. The environment returns reward r0r_0 and next state s1s_1.
  4. Reduce the remaining target by the received reward.
  5. Predict the next action with the new context (G1,s1)(G_1,s_1).

In the simple case where γ=1\gamma=1,

Gt+1=Gtrt.G_{t+1} = G_t-r_t.

If the target is 1010 and the first reward is 22, then

G1=102=8.G_1=10-2=8.

The next context effectively says, “From the current state, I want to earn 88 more.”

The target return is updated to the amount remaining at every step. The model does not guarantee that target; it predicts actions that appeared with the corresponding condition in the dataset.

Does a Larger Requested Return Produce Better Behavior?

Not necessarily.

If the highest return in the dataset is 1010 but we provide 100100 during inference, the model receives a condition it never observed during training. The number 100100 alone does not create a new strategy.

Gtargettraining supportG^{\text{target}} \notin \text{training support}

In that case, action prediction may become unstable. In general, we should choose conditions within a return range that the dataset supports and evaluate several candidate values.

Return conditioning is an interface for requesting a desired level of performance, not a guarantee that a planner will achieve it.

Why Decision Transformer Is Appealing

It Can Use Long Histories

In partially observable problems, where the current observation does not fully reveal the state, earlier states and actions can serve as context.

One Model Can Represent Diverse Behaviors

By conditioning separately on low- and high-return trajectories, a single model can express behavior at several performance levels.

It Can Reuse Sequence-Modeling Tools

Transformer architectures, masking strategies, and large-scale supervised-training techniques apply directly.

The Difficulties of RL Have Not Disappeared

Decision Transformer inherits the limitations of its dataset.

  • If the dataset contains no good trajectory, learning that behavior is difficult.
  • A gap may exist between the requested return and what is actually attainable.
  • Long sequences require more computation and context length.
  • The model does not explicitly search through environment dynamics.
  • Out-of-distribution generalization remains a problem in new states.

Nor does sequence modeling make every offline RL problem easier. A value-based method that exploits local rewards and Markov structure may be more data-efficient.

How Is It Different from Behavior Cloning?

Ordinary behavior cloning predicts an action from a state.

π(atst)\pi(a_t\mid s_t)

Decision Transformer also conditions on return and history.

π(atGt,st,a<t)\pi \left( a_t \mid G_{\leq t},s_{\leq t},a_{<t} \right)

This lets it distinguish behavior by desired return in a dataset containing both successes and failures. Still, because it ultimately predicts dataset actions through supervised learning, it remains closely related to imitation learning.

What to Remember

Decision Transformer converts offline trajectories into token sequences of return-to-go, state, and action, then predicts the next action conditioned on a desired return.

  1. Return-to-go is the sum of rewards remaining from the current point to the end of the episode.
  2. During training, GtG_t is calculated from trajectories that have already finished.
  3. During inference, a target return is supplied and reduced by each received reward.
  4. The model learns with an action-prediction loss, without an explicit Bellman backup or on-policy RL loop.
  5. Supplying a high return absent from the dataset does not automatically create a new ability.

The next article connects RLHF, RLAIF, and DPO as methods for learning language-model behavior from human or AI preferences.

References