Offline Reinforcement Learning as One Big Sequence Modeling Problem

1 Introduction . The standard treatment of reinforcement learning relies on decomposing a long-horizon problem into smaller , more local subproblems . In model-free algorithms , this takes the form of the principle of optimality ( Bellman , 1957 ) , a recursion that leads naturally to the class of dynamic programming methods like Q-learning . In model-based algorithms , this decomposition takes the form of single-step predictive models , which reduce the problem of predicting high-dimensional , policy-dependent state trajectories to that of estimating a comparatively simpler , policy-agnostic transition distribution . However , we can also view reinforcement learning as analogous to a sequence generation problem , with the goal being to produce a sequence of actions that , when enacted in an environment , will yield a sequence of high rewards . In this paper , we consider the logical extreme of this analogy : does the toolbox of contemporary sequence modeling itself provide a viable reinforcement learning algorithm ? We investigate this question by treating trajectories as unstructured sequences of states , actions , and rewards . We model the distribution of these trajectories using a Transformer architecture ( Vaswani et al. , 2017 ) , the current tool of choice for capturing long-horizon dependencies . In place of the trajectory optimizers common in model-based control , we use beam search ( Reddy , 1977 ) , a heuristic decoding scheme ubiquitous in natural language processing , as a planning algorithm . Posing reinforcement learning , and more broadly data-driven control , as a sequence modeling problem handles many of the considerations that typically require distinct solutions : actor-critic algorithms require separate actors and critics , model-based algorithms require predictive dynamics models , and offline RL methods often require estimation of the behavior policy ( Fujimoto et al. , 2019 ) . These Code is available at trajectory-transformer.github.io 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . components estimate different densities or distributions , such as that over actions in the case of actors and behavior policies , or that over states in the case of dynamics models . Even value functions can be viewed as performing inference in a graphical model with auxiliary optimality variables , amounting to estimation of the distribution over future rewards ( Levine , 2018 ) . All of these problems can be unified under a single sequence model , which treats states , actions , and rewards as simply a stream of data . The advantage of this perspective is that high-capacity sequence model architectures can be brought to bear on the problem , resulting in a more streamlined approach that could benefit from the same scalability underlying large-scale unsupervised learning results ( Brown et al. , 2020 ) . We refer to our model as a Trajectory Transformer ( Figure 1 ) and evaluate it in the offline regime so as to be able to make use of large amounts of prior interaction data . The Trajectory Transformer is a substantially more reliable long-horizon predictor than conventional dynamics models , even in Markovian environments for which the standard model parameterization is in principle sufficient . When decoded with a modified beam search procedure that biases trajectory samples according to their cumulative reward , the Trajectory Transformer attains results on offline RL benchmarks that are competitive with the best prior methods designed specifically for that setting . Additionally , we describe how variations of the same decoding procedure yield a model-based imitation learning method , a goal-reaching method , and , when combined with dynamic programming , a state-of-the-art planner for sparse-reward , long-horizon tasks . Our results suggest that the algorithms and architectural motifs that have been widely applicable in unsupervised learning carry similar benefits in RL . 2 Related Work . Recent advances in sequence modeling with deep networks have led to rapid improvement in the effectiveness of such models , from LSTMs and sequence-to-sequence models ( Hochreiter & Schmidhuber , 1997 ; Sutskever et al. , 2014 ) to Transformer architectures with self-attention ( Vaswani et al. , 2017 ) . In light of this , it is tempting to consider how such sequence models can lead to improved performance in RL , which is also concerned with sequential processes ( Sutton , 1988 ) . Indeed , a number of prior works have studied applying sequence models of various types to represent components in standard RL algorithms , such as policies , value functions , and models ( Bakker , 2002 ; Heess et al. , 2015a ; Chiappa et al. , 2017 ; Parisotto et al. , 2020 ; Parisotto & Salakhutdinov , 2021 ; Kumar et al. , 2020b ) . While such works demonstrate the importance of such models for representing memory ( Oh et al. , 2016 ) , they still rely on standard RL algorithmic advances to improve performance . The goal in our work is different : we aim to replace as much of the RL pipeline as possible with sequence modeling , so as to produce a simpler method whose effectiveness is determined by the representational capacity of the sequence model rather than algorithmic sophistication . Estimation of probability distributions and densities arises in many places in learning-based control . This is most obvious in model-based RL , where it is used to train predictive models that can then be used for planning or policy learning ( Sutton , 1990 ; Silver et al. , 2008 ; Fairbank , 2008 ; Deisenroth & Rasmussen , 2011 ; Lampe & Riedmiller , 2014 ; Heess et al. , 2015b ; Janner et al. , 2020 ; Amos et al. , 2021 ) . However , it also figures heavily in offline RL , where it is used to estimate conditional distributions over actions that serve to constrain the learned policy to avoid out-of- distribution behavior that is not supported under the dataset ( Fujimoto et al. , 2019 ; Kumar et al. , 2019a ; Ghasemipour et al. , 2021 ) ; imitation learning , where it is used to fit an expert ’ s actions to obtain a policy ( Ross & Bagnell , 2010 ; Ross et al. , 2011 ) ; and other areas such as hierarchical RL ( Peng et al. , 2017 ; Co-Reyes et al. , 2018 ; Jiang et al. , 2019 ) . In our method , we train a single high-capacity sequence model to represent the joint distribution over sequences of states , actions , and rewards . This serves as both a predictive model and a behavior policy ( for imitation ) or behavior constraint ( for offline RL ) . Our approach to RL is most closely related to prior model-based methods that plan with a learned model ( Chua et al. , 2018 ; Wang & Ba , 2020 ) . However , while these prior methods typically require additional machinery to work well , such as ensembles in the online setting ( Kurutach et al. , 2018 ; Buckman et al. , 2018 ; Malik et al. , 2019 ) or conservatism mechanisms in the offline setting ( Yu et al. , 2020 ; Kidambi et al. , 2020 ; Argenson & Dulac-Arnold , 2021 ) , our method does not require explicit handling of these components . Modeling the states and actions jointly already provides a bias toward generating in-distribution actions , which avoids the need for explicit pessimism ( Fujimoto et al. , 2019 ; Kumar et al. , 2019a ; Ghasemipour et al. , 2021 ; Nair et al. , 2020 ; Jin et al. , 2021 ; Yin et al. , 2021 ; Dadashi et al. , 2021 ) . Our method also differs from most prior model-based algorithms in the dynamics model architecture used , with fully-connected networks parameterizing diagonalcovariance Gaussian distributions being a common choice ( Chua et al. , 2018 ) , though recent work has highlighted the effectiveness of autoregressive state prediction ( Zhang et al. , 2021 ) like that used by the Trajectory Transformer . In the context of recently proposed offline RL algorithms , our method can be interpreted as a combination of model-based RL and policy constraints ( Kumar et al. , 2019a ; Wu et al. , 2019 ) , though our approach does not require introducing such constraints explicitly . In the context of model-free RL , our method also resembles recently proposed work on goal relabeling ( Andrychowicz et al. , 2017 ; Rauber et al. , 2019 ; Ghosh et al. , 2021 ; Paster et al. , 2021 ) and reward conditioning ( Schmidhuber , 2019 ; Srivastava et al. , 2019 ; Kumar et al. , 2019b ) to reinterpret all past experience as useful demonstrations with proper contextualization . Concurrently with our work , Chen et al . ( 2021 ) also proposed an RL approach centered around sequence prediction , focusing on reward conditioning as opposed to the beam-search-based planning used by the Trajectory Transformer . Their work further supports the possibility that a high-capacity sequence model can be applied to reinforcement learning problems without the need for the components usually associated with RL algorithms . 3 Reinforcement Learning and Control as Sequence Modeling . In this section , we describe the training procedure for our sequence model and discuss how it can be used for control . We refer to the model as a Trajectory Transformer for brevity , but emphasize that at the implementation level , both our model and search strategy are nearly identical to those common in natural language processing . As a result , modeling considerations are concerned less with architecture design and more with how to represent trajectory data – potentially consisting of continuous states and actions – for processing by a discrete-token architecture ( Radford et al. , 2018 ) . 3.1 Trajectory Transformer . At the core of our approach is the treatment of trajectory data as an unstructured sequence for modeling by a Transformer architecture . A trajectory τ consists of T states , actions , and scalar rewards : τ = ( s1 , a1 , r1 , s2 , a2 , r2 , . . . , sT , aT , rT ) . In the event of continuous states and actions , we discretize each dimension independently . Assuming N -dimensional states andM -dimensional actions , this turns τ into sequence of length T ( N+M+1 ) : τ = ( . . . , s1t , s 2 t , . . . , s N t , a 1 t , a 2 t , . . . , a M t , rt , . . . ) t = 1 , . . . , T. Subscripts on all tokens denote timestep and superscripts on states and actions denote dimension ( i.e. , sit is the i th dimension of the state at time t ) . While this choice may seem inefficient , it allows us to model the distribution over trajectories with more expressivity without simplifying assumptions such as Gaussian transitions . Algorithm 1 Beam search 1 : Require Input sequence x , vocabulary V , sequence length T , beam width B 2 : Initialize Y0 = { ( ) } 3 : for t = 1 , . . . , T do 4 : Ct ← { yt−1 ◦ y | yt−1 ∈ Yt−1 and y ∈ V } // candidate single-token extensions 5 : Yt ← argmax Y⊆Ct , |Y |=B logPθ ( Y | x ) // B most likely sequences from candidates 6 : end for 7 : Return argmax y∈YT logPθ ( y | x ) We investigate two simple discretization approaches : 1 . Uniform : All tokens for a given dimension correspond to a fixed width of the original continuous space . Assuming a per-dimension vocabulary size of V , the tokens for state dimension i cover uniformly-spaced intervals of width ( max si −min si ) /V . 2 . Quantile : All tokens for a given dimension account for an equal amount of probability mass under the empirical data distribution ; each token accounts for 1 out of every V data points in the training set . Uniform discretization has the advantage that it retains information about Euclidean distance in the original continuous space , which may be more reflective of the structure of a problem than the training data distribution . However , outliers in the data may have outsize effects on the discretization size , leaving many tokens corresponding to zero training points . The quantile discretization scheme ensures that all tokens are represented in the data . We compare the two empirically in Section 4.2 . Our model is a Transformer decoder mirroring the GPT architecture ( Radford et al. , 2018 ) . We use a smaller architecture than those typically used in large-scale language modeling , consisting of four layers and four self-attention heads . ( A full architectural description is provided in Appendix ? ? . ) Training is performed with the standard teacher-forcing procedure ( Williams & Zipser , 1989 ) used to train sequence models . Denoting the parameters of the Trajectory Transformer as θ and induced conditional probabilities as Pθ , the objective maximized during training is : L ( τ ) = T∑ t=1 ( N∑ i=1 logPθ ( sit | s < it , τ < t ) + M∑ j=1 logPθ ( ajt | a < j t , st , τ < t ) +logPθ ( rt | at , st , τ < t ) ) , in which we use τ < t to denote a trajectory from timesteps 0 through t− 1 , s < it to denote dimensions 0 through i− 1 of the state at timestep t , and similarly for a < jt . We use the Adam optimizer ( Kingma & Ba , 2015 ) with a learning rate of 2.5× 10−4 to train parameters θ .

Summary: The authors propose a sequence modeling approach for reinforcement learning (RL) problems. This proposed approaches enables to tackle diverse range of problem settings and unifies many of the standard components RL algorithms (e.g. as policies, models, and value functions) under a single sequence model. The authors present empirical results that illustrate the effectiveness of their model on imitation learning, goal-reaching, and offline reinforcement learning problems.