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 proposes a sequence-model approach for offline RL, which does not need to solve cumbersome dynamic programming subproblems. It utilizes a transformer architecture to effectively model state-action sequence which is conditional to the return-to-go. The proposed method shows competitive or superior performance compared to the previous state-of-the-art in various offline RL datasets.