Decision Transformer: Reinforcement Learning via Sequence Modeling

1Our code is available at : https : //sites.google.com/berkeley.edu/decision-transformer 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . 1 Introduction . Recent work has shown transformers [ 1 ] can model large-scale distributions of semantic concepts , including capable zero-shot generalization in language [ 2 ] and impressive out-of-distribution image generation [ 3 ] . This stands in sharp contrast to much work in reinforcement learning ( RL ) , which learns a single policy to model a particular narrow behavior distribution . Given the diversity of applications and impact of transformer models , we seek to examine their application to sequential decision making problems . In particular , instead of using transformers as an architectural choice for traditional RL algorithms [ 4 , 5 ] , we seek to study if trajectory modeling ( analogous to language modeling ) can serve as a replacement for conventional RL algorithms . We consider the following shift in paradigm : instead of training a policy through conventional RL algorithms like temporal difference ( TD ) learning [ 6 ] , the dominant paradigm in RL , we will train transformer models on collected experience using a sequence modeling objective . This will allow us to bypass the need for bootstrapping to propagate returns – thereby avoiding one of the “ deadly triad ” [ 6 ] known to destabilize RL . It also avoids the need for discounting future rewards , as typically done in TD-learning , which can induce undesirable short-sighted behaviors . Additionally , we can make use of existing transformer frameworks widely used in language and vision that are easy to scale , utilizing a large body of work studying stable training of transformer models ; this approach removes the need for specialized RL frameworks by appealing only to commonplace supervised learning systems . Given their demonstrated ability to model long sequences and wide data distributions , transformers also have other advantages . Transformers can perform credit assignment directly via self-attention , in contrast to Bellman backups which slowly propagate rewards and are prone to “ distractor ” signals [ 7 ] . This can enable transformers to still work effectively in the presence of sparse or distracting rewards . Furthermore , a transformer modeling approach can model a wide distribution of behaviors , enabling better generalization and transfer . While “ upside-down ” reinforcement learning ( UDRL ) [ 8 , 9 , 10 ] also uses a supervised loss conditioned on a target return , our work is motivated by sequence modeling rather than supervised learning and seeks to benefit from modeling long sequences of behaviors . See Section 6 for more discussions about related works . We explore our hypothesis by considering offline RL , where we will task agents with learning policies from suboptimal data – producing maximally effective behavior from fixed , limited experience . This task is traditionally challenging due to error propagation and value overestimation [ 11 ] . However , it is a natural task when training with a sequence modeling objective . By training an autoregressive model on sequences of states , actions , and returns , we reduce policy sampling to autoregressive generative modeling . We can specify the expertise of the policy – which “ skill ” to query – by manually setting the return tokens , acting as a prompt for generation . Illustrative example . To get an intuition for our proposal , consider the task of finding a shortest path on a directed graph posed as an RL problem . The reward is 0 when at the goal node and−1 otherwise . We train a GPT [ 12 ] model to predict next token in a sequence of returns-to-go ( sum of future rewards ) , states , and actions . Training only on random walk data – with no expert demonstrations – we can at test time generate optimal trajectories by adding a prior to generate highest possible returns ( see more details and empirical results in the Appendix ) and subsequently generate actions conditioned on that . Thus , by combining the tools of sequence modeling with hindsight return information , we achieve policy improvement without the need for dynamic programming . Motivated by this observation , we propose Decision Transformer , where we use the GPT architecture to autoregressively model trajectories ( shown in Figure 1 ) . We study whether sequence modeling can perform policy optimization by evaluating Decision Transformer on offline RL benchmarks in Atari [ 13 ] , OpenAI Gym [ 14 ] , and Key-to-Door [ 15 ] environments . We show that – without using dynamic programming – Decision Transformer performs comparably on these benchmarks to state-of-the-art model-free offline RL algorithms [ 16 , 17 ] . Furthermore , in tasks where long-term credit assignment is required , Decision Transformer capably outperforms RL algorithms . With this work , we hope to bridge vast recent progress in transformer models with RL problems . 2 Preliminaries . 2.1 Offline reinforcement learning . We consider learning in a Markov decision process ( MDP ) described by the tuple ( S , A , P , R ) . The MDP tuple consists of states s ∈ S , actions a ∈ A , transition dynamics P ( s′|s , a ) , and a reward function r = R ( s , a ) . We use st , at , and rt = R ( st , at ) to denote the state , action , and reward at timestep t , respectively . The goal in reinforcement learning is to learn a policy which maximizes the expected return E [ ∑T t=1 rt ] in an MDP . In offline reinforcement learning , instead of obtaining data via environment interactions , we only have access to some fixed limited dataset consisting of trajectories from the environment . This setting is harder as it removes the ability for agents to explore the environment and collect additional feedback . 2.2 Transformers . Transformers were proposed by Vaswani et al . [ 1 ] as an architecture to efficiently model sequences . They consist of stacked self-attention layers with residual connections . Each self-attention layer receives n embeddings { xi } ni=1 corresponding to unique input tokens , and outputs n embeddings { zi } ni=1 , preserving the input dimensions . The i-th token is mapped via linear transformations to a key ki , query qi , and value vi . The i-th output of the self-attention layer is given by weighting the values vj by the normalized dot product between the query qi and other keys kj : zi = n∑ j=1 softmax ( { 〈qi , kj′〉 } nj′=1 ) j · vj . ( 1 ) This allows the layer to assign “ credit ” by implicitly forming state-return associations via similarity of the query and key vectors ( maximizing the dot product ) . In this work , we use the GPT architecture [ 12 ] , which modifies the transformer architecture with a causal self-attention mask to enable autoregressive generation , replacing the summation/softmax over the n tokens with only the previous tokens in the sequence ( j ∈ [ 1 , i ] ) . We defer the other architecture details to the original papers . 3 Method . In this section , we present Decision Transformer , which models trajectories autoregressively with minimal modification to the transformer architecture , as summarized in Figure 1 and Algorithm 1 . Trajectory representation . The key desiderata in our choice of trajectory representation are ( a ) it should enable transformers to learn meaningful patterns and ( b ) we should be able to conditionally generate actions at test time . It is nontrivial to model rewards since we would like the model to generate actions based on future desired returns , rather than past rewards . As a result , instead of modeling the rewards directly , we model the returns-to-go R̂t = ∑T t′=t rt′ . This leads to the following trajectory representation which is amenable to autoregressive training and generation : τ = ( R̂1 , s1 , a1 , R̂2 , s2 , a2 , . . . , R̂T , sT , aT ) . ( 2 ) Architecture . We feed the last K timesteps into Decision Transformer , for a total of 3K tokens ( one for each modality : return-to-go , state , or action ) . To obtain token embeddings , we learn a linear layer for each modality , which projects raw inputs to the embedding dimension , followed by layer normalization [ 18 ] . For environments with visual inputs , the state is fed into a convolutional encoder instead of a linear layer . Additionally , an embedding for each timestep is learned and added to each token – note this is different than the standard positional embedding used by transformers , as one timestep corresponds to three tokens . The tokens are then processed by a GPT [ 12 ] model , which predicts future action tokens via autoregressive modeling . Training . We sample minibatches of sequence length K from the dataset . The prediction head corresponding to the input token st is trained to predict at – either with cross-entropy loss for discrete actions or mean-squared error for continuous actions – and the losses for each timestep are averaged . We did not find predicting the states or returns-to-go to be necessary for good performance , although it is possible ( as shown in Section 5.3 ) and would be an interesting study for future work . Evaluation . During evaluation rollouts , we specify a target return based on our desired performance ( e.g. , specify maximum possible return to generate expert behavior ) as well as the environment starting state , to initialize generation . After executing the generated action , we decrement the target return by the achieved reward and obtain the next state . We repeat this process of generating actions and applying them to obtain the next return-to-go and state until episode termination . Algorithm 1 Decision Transformer Pseudocode ( for continuous actions ) # R , s , a , t : returns -to-go , states , actions , or timesteps # K : context length ( length of each input to DecisionTransformer ) # transformer : transformer with causal masking ( GPT ) # embed_s , embed_a , embed_R : linear embedding layers # embed_t : learned episode positional embedding # pred_a : linear action prediction layer # main model def DecisionTransformer ( R , s , a , t ) : # compute embeddings for tokens pos_embedding = embed_t ( t ) # per -timestep ( note : not per -token ) s_embedding = embed_s ( s ) + pos_embedding a_embedding = embed_a ( a ) + pos_embedding R_embedding = embed_R ( R ) + pos_embedding # interleave tokens as ( R_1 , s_1 , a_1 , ... , R_K , s_K ) input_embeds = stack ( R_embedding , s_embedding , a_embedding ) # use transformer to get hidden states hidden_states = transformer ( input_embeds=input_embeds ) # select hidden states for action prediction tokens a_hidden = unstack ( hidden_states ) . actions # predict action return pred_a ( a_hidden ) # training loop for ( R , s , a , t ) in dataloader : # dims : ( batch_size , K , dim ) a_preds = DecisionTransformer ( R , s , a , t ) loss = mean ( ( a_preds - a ) * * 2 ) # L2 loss for continuous actions optimizer.zero_grad ( ) ; loss.backward ( ) ; optimizer.step ( ) # evaluation loop target_return = 1 # for instance , expert -level return R , s , a , t , done = [ target_return ] , [ env.reset ( ) ] , [ ] , [ 1 ] , False while not done : # autoregressive generation/sampling # sample next action action = DecisionTransformer ( R , s , a , t ) [ -1 ] # for cts actions new_s , r , done , _ = env.step ( action ) # append new tokens to sequence R = R + [ R [ -1 ] - r ] # decrement returns -to -go with reward s , a , t = s + [ new_s ] , a + [ action ] , t + [ len ( R ) ] R , s , a , t = R [ -K : ] , ... # only keep context length of K

Summary: This paper uses the Transformer model as an auto-regressive model to predict the actions (decisions) givens the state inputs, past actions, and more importantly, the return-to-go. Training such a policy becomes a supervised-learning problem, similar to behavior cloning. The paper applies this method to offline RL problems and demonstrates that policy trained in this supervised manner achieves competitive results compared to SOTA model-free offline RL algorithms.