Offline Reinforcement Learning as One Big Sequence Modeling Problem
Authors: Michael Janner
Qiyang Li
Sergey Levine
Venue: NeurIPS 2021 Spotlight (2021)
Keywords: Reinforcement learning
transformers

Abstract: Reinforcement learning (RL) is typically viewed as the problem of estimating single-step policies (for model-free RL) or single-step models (for model-based RL), leveraging the Markov property to factorize the problem in time. However, we can also view RL as a sequence modeling problem: predict a sequence of actions that leads to a sequence of high rewards. Viewed in this way, it is tempting to consider whether powerful, high-capacity sequence prediction models that work well in other supervised learning domains, such as natural-language processing, can also provide simple and effective solutions to the RL problem. To this end, we explore how RL can be reframed as "one big sequence modeling" problem, using state-of-the-art Transformer architectures to model distributions over sequences of states, actions, and rewards. Addressing RL as a sequence modeling problem significantly simplifies a range of design decisions: we no longer require separate behavior policy constraints, as is common in prior work on offline model-free RL, and we no longer require ensembles or other epistemic uncertainty estimators, as is common in prior work on model-based RL. All of these roles are filled by the same Transformer sequence model. In our experiments, we demonstrate the flexibility of this approach across imitation learning, goal-conditioned RL, and offline RL.

Offline Reinforcement Learning as One Big Sequence Modeling Problem
Michael Janner
Qiyang Li
Sergey Levine
University of California at Berkeley
(janner, geli)\berkeley. edu svlevinedeecs berkeley wtw.eeeeeeesbbbbeee.........
edu
Abstract
Reinforcement learning (RL) is typically concerned with estimating stationary policies single-step models, leveraging the Markov property to factorize prob- lems time, However, we can also view RL as a generic sequence modeling problem, with the goal being produce sequence of actions that leads to a sequence of high rewards. Viewed in this way, it is tempting to consider whether high-capacity sequence prediction models that work well in other domains, such as natural- lanural-language processing, can also provide effective solutions to the RL problem. To this end, explore how RL be tackled with the tools of sequence modeling, using a Transformer architecture to model distributions over trajectories and repurposing beam search planning algorithm. Framing RL sequence modeling problem simplifies a range design decisions, allowing us to dispense with many of the components common in offline RL algorithms. We demonstrate the flexibility of this approach across long-horizon dynamics prediction, imitation learning, goal-conditioned RL, and offline RL. Further, show that this approach can be combined with existing model-free algorithms to yield a state-of-the-art planner in sparse-reward, long-horizon tasks.
1 Introduction
The standard treatment of reinforcement learning relies on decomposing long-horizon problem into smaller. more local subproblems. model-free algorithms, this takes the form of the principle of optimality (Bellman, 1957), 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, volicy-dependent state trajectories to that of estimating a comparatively simpler, policy agliss-agnostic transition distribution.
However, we can also view reinforcement learning as analogous to sequence generation problem, with the goal being to produce sequence of actions that, when enacted in 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 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 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). heuristic decoding scheme ubiquitous in natural language processing, as planning algorithm.
Posing reinforcement learning, and more broadly data-driven control, as 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 RI methods often require estimation of the behavior policy (Fujimoto et al., 2019). These
Code available at trajectory crajectory-trrasffrmee trrassfrrraaf.rrr.oo.....oo .iitrrggttt..... io
35th Conference on Neural Information Processing Systems (NeurIPS 2021).

---

3
2
3
M
St
t
t
a
a
aj
at
Tt
1
Trajectory Transformer
2 b
3 St
3 a
St
o
i
a
a
O
a
Figure 1 (Architecture) The Trajectory Transformer trains on sequences of (autoregressively dis- cretized) states, actions, and rewards. Planning with the Trajectory Transformer mirrors the sampling procedure used to generate sequences from language model.
components estimate different densities 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 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 single sequence model, which treats states, actions, and rewards as simply stream of data. The advantage of this perspective is that high-capacity sequence model architectures can be brought bear the problem, resulting in same scalability underlying large-scale unsupervised learning results (Brown et al., 2020)-
streamlined approach that could benefit from the
We refer to our model Trajectory Transformer (Figure 1) and evaluate it in the offline regime so as to able to make use of large amounts of prior interaction data. The Trajectory Transformer is substantially more reliable onng-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 competitive with the best prior methods designed specifically for that setting. Additionally, describe how variations of the same decoding procedure yield a model-based imitation learning method, goal-reaching method, and, when combined with dynamic programming, a state-of-the-an 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-sequenee models (Hochreiter Schmidhuber, 1997; Sutskever et al., 2014) to Transformer architectures with self-attention (Vaswani 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, number of prior works have studied applying sequence models of various types to represent components in standard RL algorithms, such policies, value functions, and models (Bakker, 2002; Heess 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 work is different: 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.
aim replace much of the RL pipeline possible with
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., 201 Janner et al., 2020; Amos al., 2021). However, it also figures heavily in offline RI., where 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 is used fit an expert's actions to obtain policy (Ross & Bagnell, 2010; Ross et al., 201 1); and other areas such as hierarchical RL (Peng et al., 2017; Co-Reyes al., 2018; Jiang et al., 2019). 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 predictive model and a behavior policy (for imitation) or behavior constraint (for offline RL).
Our approach RL is most closely related to prior model--ased methods that plan with a learned model (Chua et al., 2018; Wang & Ba, 2020). However, while these prior methods typically require additional machinery work well, such as ensembles the online setting (Kurutach al., 2018; Buckman et al., 2018; Malik et al., 2019) or conservatism mechanisms in the offline setting (Yu 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 diagonal- covariance Gaussian distributions being a common choice (Chua et al., 2018), though recent work has highlighted the effectiveness of autoregressive state prediction (Zhang al., 2021) like that used by the Trajectory Transformer. In the context of recently proposed offline RL algorithms,, our method can be interpreted
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, 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 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 ceam-search--ee baarch planning used by the Trajectory Transformer. Their work further supports the possibility that a high-capacity se- quence 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, 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 the implementation level, both our model and search strategy are nearly identical to those common in natural language processing. As 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 discrete-token architecture (Radford et al., 2018). 8).
3.1 Trajectory Transformer
At the core of our approach the treatment of trajectory data as an unstructured sequence for modeling by a Transformer architecture. A trajectory t consists of T states, actions, and scalar rewards:
s1 a r1, 82, a2, ................................ ST, ar, ar,,rr) rT).
In the event of continuous states and actions, we discretize each dimension independently. Assuming N N-dimensional states and M -ddimennsinnl actions, this turns T into sequence of length T(N+ + +M+1):
(
. a. oz
M re....)
t=
Subscripts on all tokens denote timestep and superscripts on states and actions denote dimension (i.e., of is the 7h dimension of the state at time 1). 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.
3

---

Algorithm Beam search
1: Require Input sequence x, vocabulary V. sequence length T, beam width B
2: Initialize Yo {{)
3: for 1,...,T6 do
4: C - V1-109 Y E Ye-1 and y E v} // candidate single-token extensions
5: Yi argmax log PY
// B most likely sequences from candidates
Ygc, [Y
6: end for
7: Return argmax log Po(y |x
yεУr
We investigate two simple discretization approaches:
1. Uniform: All tokens for given dimension correspond fixed width of the original
continuous space. Assuming a per-dimension vocabulary size of V. the tokens for state dimension cover uniformly-spacee intervals of width (maxs' min s')
2. Quantile: All tokens for given dimension account for an equal amount of probability mass under the empirical data distribution; each token accounts for 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 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 the data. We compare the two empirically in Section 4.2.
Our model is Transformer decoder mirroring the GPT architecture (Radford et al., 2018). We use 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 0 and induced conditional probabilities as Po. the objective maximized during training is:
[(T) = Σ Σ log (s S TE)) +
Σ log Po (a a St, T<t) + log (o a St. T<()):
in which use Ter denote trajectory from timesteps through 1, s.¹ denote dimensions 0 through i 1 of the state timestep t, and similarly for a We use the Adam optimizer (Kingma & Ba, 2015) with learning rate of 2.5 x 10 4 to train parameters t.
3.2 Planning with Beam Search
We now describe how sequence generation with the Trajectory Transformer can be repurposed for control, focusing on three settings: imitation learning, goal-conditioned reinforcement learning, and offline reinforcement learning. These settings listed in increasing amount of required modification on top of the sequence model decoding approach routinely used natural language processing.
The core algorithm providing the foundation of our planning techniques, beam search, is described in Algorithm 1 for generic sequences. Following the presentation in Meister al. (2020), we have overloaded log Pa x) to define the likelihood of set of sequences in addition to that of a single sequence: log (Y x) = Eyer log Po x). We use () to denote the empty sequence and to represent concatenation.
Imitation learning. When the goal is to reproduce the distribution of trajectories in the training data, optimize directly for the probability of trajectory т. This situation matches the goal of sequence modeling exactly and as such we may use Algorithm 1 without modification by setting the conditioning input x to the current state s; (and optionally previous history Tex).
The result of this procedure tokenized trajectory т. beginning from current state s,, that has high probability under the data distribution. If the first action a, in the sequence is enacted and beam

---

search is repeated, we have receding horizon-controolee controller. This approach resembles a long- aongoorizon model-based variant of behavior cloning, in which entire trajectories are optimized match those of reference behavior instead of only immediate state-conditioned actions. If we set the predicted sequence length to be the action dimension, our approach corresponds exactly to the simplest form of behavior cloning with an autoregressive policy.
Goal-conditioned reinforcement learning. Transformer architectures feature a "causal" attention mask ensure that predictions only depend on previous tokens in sequence. In the context of natural language, this design corresponds to generating sentences in the linear order in which they are spoken as opposed to an ordering reflecting their hierarchical syntactic structure (see, however, Gu et al. 2019 for a discussion of non-left-to-right sentence generation with autoregressive models). In the context of trajectory prediction, this choice instead reflects physical causality, disallowing future events to affect the past. However, the conditional probabilities of the past given the future are still well-defined, allowing us condition samples not only on the preceding states, actions, and rewards that have already been observed, but also any future context that we wish to occur. If the future context is state at the end of a trajectory, we decode trajectories with probabilities of the form:
(s 5 ST)
We can use this directly as goal-reaching method by conditioning desired final state ST. If we always condition sequences on final goal state, we may leave the lower-diagonal attention mask intact and simply permute the input trajectory to {sT, S2, 57-1) By prepending the goal state to the beginning of a sequence, we ensure that all other predictions may attend to it without modifying the standard attention implementation. This procedure for conditioning resembles prior methods that use supervised learning to train coal-conditioned policies (Ghosh et al., 2021) I) and is also related to relabeling techniques in model-free RL (Andrychowicz et al., 2017). In our framework, is identical to the standard subroutine in sequence modeling: inferring the most likely sequence given available evidence.
Offline reinforcement learning. The beam search method described in Algorithm 1 optimizes sequences for their probability under the data distribution. By replacing the pog--obbbbilities of transitions with the predicted reward signal, we can use the same Trajectory Transformer and behavior. Appealing to the control as inference graphical
search strategy for model (Levine, 2018), we are in effect replacing a transition's 's log-probabilig in beam search with its log-probability of optimality.
Using beam-search as a -arradd-ixiimzznngg procedure has the risk of leading to myopic behavior. To address this issue, we augment each transition in the training trajectories with reward-to-go: R₁ R 7 and include it as an additional quantity, discretized identically to the others, to be predicted after immediate rewards r.- During planning, we then have access to value estimates from our model to add to cumulative rewards. While acting greedily with respect to such Monte Carlo value estimates is known to suffer from poor sample complexity and convergence to suboptimal behavior when online data collection is not allowed, we only use this reward-to- -togooo-oooo-ooooo estimate as a heuristic guide beam search, and hence our method does not require the estimated values to be as accurate as in methods that rely solely on the value estimates to select actions.
In offline RL, reward-to-go estimates are functions of the behavior policy that collected the training data and do not, in general, correspond to the values achieved by the Trajectory Transformer-derivedd policy. Of course, it is much simpler to learn the value function of the behavior policy than that of the optimal policy, since we can simply use Monte Carlo estimates without relying on Bellman updates. A value function for improved policy would provide better search heuristic, though requires invoking the tools of dynamic programming. In Section 4.2 we show that the simple reward-to- to-go estimates are sufficient for planning with the Trajectory Transformer in many environments, but that improved value functions are useful in the most challenging settings, such as sparse-reward tasks.
Because the Trajectory Transformer predicts reward and reward-to-ge only every N + M + 1 tokens, we sample all intermediate tokens according to model log-probabilities, as in the imitation learning and goal-reaching settings. More specifically, we sample full transitions (St, ar, rt, R₂) using beam search, treat these transitions as our vocabulary, and filter sequences of transitions by those with the highest cumulative reward plus reward-to-go estimate.
5

---

Feedforwart Transformer Reference
Figure 2 (Prediction visualization) A qualitative comparison of length- 100 trajectories generated by the Trajectory Transformer and a feedforward Gaussian dynamics model from PETS, otate-of-the-aar planning algorithm Chua et al. (2018). Both models were trained on trajectories collected by single policy, for which true trajectory is shown for reference. Compounding errors in the single-step model lead to physically implausible predictions, whereas the Transformer generated trajectory is visually indistinguishable from those produced by the policy acting in the actual environment. The paths of the feet and head are traced through space for depiction of the movement between rendered frames.
-ransforrer--eeeerate
We have taken sequence-modeling route to what could be described as fairly simple-looking model- based planning algorithm, in that sample candidate action sequences, evaluate their effects using predictive model, and select the reward-maximizing trajectory. This conclusion is in part due to the close relation between sequence modeling and trajectory optimization. There is one dissimilarity, however, that is worth highlighting: by modeling actions jointly with states and sampling them using the same procedure, we can prevent the model from being queried on out-of-distribution actions. The alternative, of treating action sequences unconstrained optimization variables that do not depend on state (Nagabandi et al., 2018), can more readily lead to model exploitation, as the problem of maximizing reward under learned model closely resembles that of finding adversarial examples for a classifier (Goodfellow et al., 2014).
4 Experiments
Our experimental evaluation focuses (1) the accuracy of the Trajectory Transformer long- horizon predictor compared to standard dynamics model parameterizations and (2) the utility of sequence modeling tools namely beam search control algorithm in the context of offline reinforcement learning, imitation learning, and goal-reaching.
4.1
Model Analysis
We begin by evaluating the Trajectory Transformer long-horizon coiii-conditioned predictive model. The usual strategy for predicting trajectories given policy is rollout with single-step model, with actions supplied by the policy. Our protocol differs from the standard approach not only in that the model is not Markovian, but also that does not require access to a policy to make predictions the outputs of the policy are modeled alongside the states encountered by that policy. Here, we focus only on the quality of the model's predictions; we use actions predicted by the model for imitation learning method in the next subsection.
Trajectory predictions. Figure 2 depicts a visualization of predicted 100-timestep trajectories from our model after having trained on dataset collected by trained humanoid policy. Though model-based methods have been applied to the humanoid task, prior works tend to keep the horizon

---

Humanoid
Partially-Observed Humanoid
100
70
log likelihood
80
log likelihood
60
50
40
0
30
10
timestep
timestep
Transformer
Markovian Transformer
Feedforward
Discrete oracle
Figure 3 (Compounding model errors) We compare the accuracy of the Trajectory Transformer (with uniform discretization) to that of the probabilistic feedforward model ensemble (Chua et al., 2018) over the course of a planning horizon in the humanoid environment, corresponding to the trajectories visualized in Figure 2. The Trajectory Transformer has substantially better enror com- pounding with respect prediction horizon than the feedforward model. The discrete oracle the maximum log likelihood attainable given the discretization size; see Appendix ?? for discussion.
intentionally short to prevent the accumulation of model errors (Janner et al., 2019; Amos et al., 2021). The reference model is the probabilistic ensemble implementation of PETS (Chua al., 2018); we tuned the number models within the ensemble, the number of layers, and layer sizes, but were unable to produce a model that predicted accurate sequences for more than few dozen steps. In contrast, we see that the Trajectory Transformer's 's long-horizon predictions are substantially more accurate, remaining visually indistinguishable from the ground-truth trajectories even after 100 predicted steps. To our knowledge, no prior model-based RL. algorithm has demonstrated predicted rollouts of such accuracy and length tasks of comparable dimensionality.
Error accumulation. A quantitative account of the same finding is provided in Figure 3, which evaluate the model's accumulated error versus prediction horizon. Standard predictive models tend to have excellent single-step errors but poor long-horizon accuracy, so instead of evaluating test-set single-step likelihood, we sample 1000 trajectories from fixed starting point estimate the per-timestep state marginal predicted by each model. We then report the likelihood of the states visited by the reference policy on held-out of trajectories under these predicted marginals. To evaluate the likelihood under our discretized model, we treat each bin as uniform distribution over its specified range; by construction, the model assigns zero probability outside of this range.
To better isolate the source the Transformer's improved accuracy over standard single-step models, also evaluate Markovian variant of our architecture. This ablation has truncated context window that prevents it from attending to more than one timestep in the past. This model performs similarly to the trajectory Transformer on fully-observed environments. suggesting that architecture differences and increased expressivity from the autoregressive state discretization play large role in the trajectory Transformer's long-horizon accuracy. We construct partially-observe the same humanoid environment, in which each dimension of every state is masked out with 50% probability (Figure right), and find that, as expected, the long-horizon conditioning plays larger role in the model's accuracy in this setting.
version of
Attention patterns. We visualize the attention maps during model predictions in Figure We find two primary attention patterns. The first is a discovered Markovian strategy. in which state prediction attends overwhelmingly the previous transition. The second is qualitatively striated, with the model attending to specific dimensions in multiple prior states for each state prediction. Simultaneously, the action predictions attend to prior actions more than they do prior states. The action dependencies contrast with the usual formulation of behavior cloning, in which actions are function of only past states, but is reminiscent of the action filtering technique used in some planning algorithm to produce smoother action sequences (Nagabandi al., 2019).
7

---

81
a
ag.
8:455
S
a
a
Figure 4 (Attention patterns) We observe two distinct types of attention masks during trajectory prediction. In the first, both states and actions dependent primarily on the immediately preceding transition, corresponding to a model that has learned the Markov property. The second strategy has striated appearance, with state dimensions depending most strongly on the same dimension of multiple previous timesteps. Surprisingly, actions depend more on past actions than they do on past states, reminiscent of the action smoothing used in some trajectory optimization algorithms (Nagabandi et al., 2019). The above masks are produced by first- and third-layer attention head during sequence prediction on the hopper benchmark; reward dimensions are omitted for this visualization.
4.2 Reinforcement Learning and Control
Offiine reinforcement learning. We evaluate the Trajectory Transformer number of envi- ronments from the D4RI, offline benchmark suite (Fu et al., 2020), including the locomotion and AntMaze domains. This evaluation the most difficult of our control settings, reward-maximizing behavior is the most qualitatively dissimilar from the types of behavior that are normally associated with unsupervised modeling namely, imitative behavior. Results for the locomotion environ- ments are shown in Table 1. We compare against five other methods spanning other approaches to data-driven control: (1) schavior-regularine
actor-critic (BRAC: Wu al. 2019) and conservative seree-offthheeeeee Q-learning (CQL; Kumar et al. 2020a) represent the current state-of of ttt-hherrt art in model-fri frre offline RI; model-based offline planning (MBOP: Argenson & Dulac-Arnold 2021) is the best-perforning prior offline trajectory optimization technique; decision transformer (DT; Chen et al. (2021)) is mencurrently-develogee sequence-modeling approach that uses conditionii instead of planning; and behavior-cloning (BC) provides the performance of pure imitative method.
The Trajectory Transformer performs on par with or better than all prior methods (Table 1). The two discretization variants of the Trajectory Transformer, uniform and quantile, perform similarly on all environments except for HalfCheetah-Me Maeihhh-eiimm--- Medium-Expert, where the large range of the velocities prevents the uniform discretization scheme from recovering the precise actuation required for enacting the expert policy. As result, the quantile discretization approach achieves a return of more than twice that of the uniform discretization.
Combining with Q-functions. Though Monte Carlo value estimates are sufficient for many stan- dard offline RL benchmarks, in sparse-reward and long-horizon settings they become too uninforma- tive to guide the beam beam-search-based planning procedure. In these problems, the value estimate from the Transformer can be replaced with Q-function trained via dynamic programming. We explore this combination by using the Q-function from the implicit Q-learning algorithm (IQL; Kostrikov et al. 2021) on the AntMaze navigation tasks (Fu et al., 2020), for which there is only sparse reward upon reaching the goal state. These tasks evaluate temporal compositionality because they require stitching together multiple zero-reward trajectories in the dataset to reach a designated goal.
AntMaze results are provided in Table 2. Q- Q-guided Trajectory Transformer planning outperforms all prior methods on all maze sizes and dataset compositions. particular, outperforms the IQL method from which we obtain the Q-function, underscoring that planning with Q Q-function
More attention visualizations be found at trajectory-traasforre gitthb.cco io/attention

---

Dataset
Environment
BC
MBOP
BRAC
CQL
DT
TT
(uniform)
TT (quantile)
Med-Ean Med-Expert
HalfCheetah
59.9
105.9
41.9
91.6
86.8
40.8 #2.3
95.0 #0.2
Med-Expert
Hopper
79.6
55.1
0.9
105.4
107.6
106.0
10.28
110.0 12.7
Med-Expert
Walker2d
36.6
70.2
81.6
108.8
108.1
91.0 12.8
101.9 16.8
Medium
HalfCheetah
44.6
46.3
44.0
42.6
#0.31
46.9 #0.4
Medium
Hopper
63.9
48.8
31.3
58.5
67.6
67.4 #2.9
61.1 #3.6
Medium
Walker2d
77.3
41.0
81.1
72.5
74.0
81.3 22.1
79.0 12.8
Med-Replay
HalfCheetah
4.3
42.3
17.7
45.5
36.6
44.1 1.09
41.9 1:255
Med-Replay
Hopper
27.6
12.4
0.6
95.0
82.7
99.4 #3.2
91.5 #3.6
Med-Replay
Walker2d
36.9
9.7
0.9
77.2
66.6
79.4 #3.3
82.6 6699
Average
47.7
47.8
36.9
77.6
74.7
72.6
78.9
Table (Offline reinforcement learning) The Trajectory Transformer (TT) performs on par with or better than the best prior offline reinforcement learning algorithms on D4RL locomotion (v2) tasks. Results for TT variants correspond to the mean and standard eπror over 15 random seeds (5 independently trained Transformers and 3 trajectories per Transformer). We detail the sources of the performance for other methods in Appendix ??.
80
Average iimmalzedd nutm
60
40
20
BC
MBOP
BRAC
CQL
Trajectary
Transfrmme
Transformer
Behavior Cloning
Trojectory Optimisation
Temporal Difference
Sequence Modeling
Figure 5 (Offfine averages) A plot showing the average per-algorithm performance Table 1, with bars colored according crude algorithm categorization. In this plot, "Trajectory Transformer" refers to the quantile discreization variant.
search heuristic can be less susceptible errors in the Q-function than policy extraction. However, because the Q-guided planning procedure still benefits from the temporal compositionality of both dynamic programming and planning, outperforms return-conde couuu-ccndiiioning approaches, such as the Decision Transformer, that suffer due to the lack of complete demonstrations in the AntMaze datasets.
Imitation and goal-reaching. We additionally plan with the Trajectory Transformer using standard likelihood- likelihood-maximiziii maximizingiiggg beam search, as opposed to the return- mauum-maximizing version used for offline RL. We find that after training the model on datasets collected by expert policies (Fu et al., 2020), using beam search as a receding-horize horizon controller achieves an average normalized return of 104% and 109% in the Hopper and Walker2d environments, respectively, using the same evaluation protocol of 15 runs described as the offline RL results. While this result is perhaps unsurprising, as behavior cloning with standard feedforward architectures is already able to reproduce the behavior of the expert policies, it demonstrates that a decoding algorithm used for language modeling can be effectively repurposed for control.
Finally, we evaluate the goal-reaching variant of beam-search, which conditions on a future desired state alongside previously encountered states. We use continuous variant of the classic four rooms environment as testbed (Sutton al., 1999). Our training data consists of trajectories collected by pretrained goal-reaching agent, with start and goal states sampled uniformly at random across the state space. Figure depicts routes taken by the planner. Anti-causal conditioning on a future state allows for beam search to be used as a goal-reaching method. No reward shaping, or rewards of any sort, are required; the planning method relies entirely on goal relabeling. An extension of this experiment to procedurally generatte maps is described in Appendix ??.

---

Dataset
Environment
BC
CQL
IQL
DT
TT (+Q)
Umaze
AntMaze
54.6
74.0
87.5
59.2
100.0 #0.0
Medium-Play
AntMaze
0.0
61.2
71.2
0.0
93.3 16.4
Modium-Diverns
AntMaze
0.0
53.7
70.0
0.0
100.0
Large-Play
AntMaze
0.0
15.8
39.6
0.0
66.7 112.2
Large-Diverse
AntMaze
0.0
14.9
47.5
0.0
60.0 #12.7
Average
10.9
44.9
63.2
11.8
84.0
Table (Combining with Q-functions) Performance on the sparse-reward AntMaze (v0) navigation task. Using Q-function search heuristic with the Trajectory Transformer (TT (+-)) outperforms policy extraction from the Q-function (IQL) and return- connn--oniiionnngggg approaches like the Decision Transformer (DT). We report means and standard error over 15 random seeds for TT (+Q);; baseline results are taken from Kostrikov et al. (2021).
Figure 6 (Goal-reaching) Trajectories collected by TTO with anti-causal goal-state conditioning in continuous variant of the four rooms environment. Trajectories visualized through all encountered states, with color becoming more saturated as time progresses. Note that these curves depict real trajectories collected by the controller and sampled sequences. The starting state depicted by and the goal state by Best viewed in color.
passing
5 Discussion and Limitations
We have presented a sequence modeling view on reinforcement learning that enables to derive single algorithm for diverse range of problem settings, unifying many the standard components of reinforcement learning algorithms (such as policies, models, and value functions) under single sequence model. The algorithm involves training sequence model jointly on states, actions, and rewards and sampling from it using minimally modified beam search. Despite drawing from the tools of large-scale language modeling instead of those normally associated with control, we find that this approach is effective imitation learning, goal-reaching, and offline reinforcement learning.
However, prediction with Transformers is currently slower and more resource- intennsive than predic- tion with the types of single-step models often used in model-based control, requiring up to multiple seconds for action selection when the context window grows too large. This precludes real-time wal-time control with standard Transformers for most dynamical systems. While the beam-search-basee planner conceptually an instance model-predictive control,
as such could be applicable wherever model- based RI. is, in practice the slow planning also makes online RI. experiments un- wieldy. CComputationally-fffiicc
Transformer architectures (Tay et al., 2021) could potentially cut runtimes down substantially.) Further, we have chosen to discretize continuous data to fit a standard architecture instead of modifying the architecture to handle continuous inputs. While we found this design to be much more effective than conventional continuous dynamics models, it does in principle impose an upper bound on prediction precision.
This paper is an investigation of minimal type of algorithm that can be applied to RL problems. While one of the interesting implications of our results that problems can be reframed as supervised learning tasks with an appropriate choice of model, the most practical instantiation of this idea may come from combinations with dynamic programming techniques, as suggested by the effectiveness of the Trajectory Transformer with Q-guided planning.
10

---

Code References
We used the following open-source libraries for this work: NumPy (Harris et al., 2020), PyTorch (Paszke et al., 2019), and minGPT (Karpathy, 2020).
Acknowledgements
We thank Ethan Perez and Max Kleiman-Weiner for helpful discussions and Ben Eysenbach for feedback on an early draft. M.J. thanks Karthik Narasimhan for early inspiration about parallels between language modeling and model-based reinforcement learning. This work was partially supported by computational resource donations from Microsoft. M.J. is supported by fellowships from the National Science Foundation and the Open Philanthropy Project.
References
Amos, B., Stanton, S., Yarats, D., and Wilson, A. G. On the model-based stochastic value gradient for continuous reinforcement learning. In Conference Learning for Dynamics and Control, 2021.
Andrychowicz, M., Wolski, F., Ray, A., Schneider, J., Fong, R., Welinder, P., McGrew, B., Tobin, J., Abbeel, P., and Zaremba, W. Hindsight experience replay. In Advances in Neural Information Processing Systems, 2017.
Argenson, A. and Dulac-Amold, G. Model- based offline planning. In International Conference on Learning Representations, 2021.
Bakker, B. Reinforcement learning with long short-term memory. In Neural Information Processing Systems, 2002.
Bellman, R. Dynamic Programming. Dover Publications, 1957.
Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., Agarwal, S., Herbert- Voss, A., Krueger, G., Henighan, T., Child, R., Ramesh, A., Ziegler, D., Wu, J., Winter, C., Hesse, C.,, Chen, M., Sigler, E., Litwin, M., Gray, S., Chess, B., Clark, J., Berner, C., McCandlish, S., Radford, A., Sutskever, and Amodei, D. Language models are few-shot learners. In Advances in Neural Information Processing Systems, 2020.
Buckman, J., Hafner, D.. Tucker, G., Brevdo, E., and Lee, H. Sample-efficient reinforcement learning with stochastic ensemble value expansion. Advances in Neural Information Processing Systems, 2018.
Chen, L., Lu, K., Rajeswaran, A., Lee, K., Grover, A., Laskin, M., Abbeel, P., Srinivas, A., and Mordatch, I. Decision Transformer: Reinforcement learning via sequence modeling. arXiv preprint arXiv:2106.01345. 2021.
Chiappa, S., Racaniere, S., Wierstra, D., and Mohamed, S. Recurrent environment simulators. In International Conference Learning Representations, 2017.
Chua, K., Calandra, R., McAllister, R., and Levine, S. Deep reinforcement learning handful of trials using probabilistic dynamics models. In Advances in Neural Information Processing Systems. 2018.
Co-Reyes, J., Liu, Y., Gupta, A., Eysenbach, B., Abbeel, P., and Levine, S. Self-consistent trajectory autoencoder: Hierarchical reinforcement learning with trajectory embeddings. In International Conference on Machine Learning, 2018.
Dadashi, R., Rezacifar, S., Vieillard, N., Hussenot, L., Pietquin, O., and Geist, M. Offline reinforce- ment learning with pseudometric learning. In International Conference on Machine Learning, 2021.
Deisenroth, M. and Rasmussen, C. PILCO: A model-based and data-efficient approach policy search. In International Conference Machine Learning, 2.11.

---

Fairbank, M. Reinforcement learning by value gradients. arXiv preprint arXiv:0803.3539 2008.
Fu, J., Kumar, A., Nachum, O., Tucker, G., and Levine, S. D4RL: Datasets for deep data-driven reinforcement learning. arXiv preprint arXiv:2004.072199 2020.
Fujimoto, S., Meger, D., and Precup. D. Off-policy deep reinforcement learning without exploration. In International Conference on Machine Learning, 2019.
Ghasemipour, S. K. S., Schuurmans, D., and Gu, S. S. EMaQ: Expected-max Q-learing operator for simple yet effective offline and online RL. 2021.
Ghosh, D., Gupta, A., Reddy, A., Fu, Devin, C. M., Eysenbach. B., and Levine, S. Learning reach goals via iterated supervised learning, In International Conference on Learning Representations, 2021.
Goodfellow, I. J., Shlens, J., and Szegedy, C. Explaining and harnessing adversarial examples. arXiv preprint arXiv:1412.6553 2014.
Gu, J., Liu, Q. and Cho, K. Insertion-based Decoding with Automatically Inferred Generation Order. Transactions of the Association for Computational Linguistics, 2019.
Harris, C. R., Millman, K. J., van der Walt, S. J., Gommers, R., Virtanen, P., Cournapeau, D., Wieser, E., Taylor, J., Berg, S., Smith, J., Kern, R., Picus, M., Hoyer, S., van Kerkwijk, M. H., Brett, M., Haldane, A., del Río, J. F., Wiebe, M., Peterson, P.. Gérand-Marchant, P., Sheppard, K., Reddy, T., Weckesser, W. Abbasi, H.. Gohlke, C., and Oliphant, T. E. Array programming with NumPy. 2020.
Nature, 585(7825)3357-366
Heess, N., Hunt, J. J., Lillicrap, T., and Silver, D. Memory-based control with recurrent neural networks. In Neural Information Processing Systems Deep Reinforcement Learning Workshop. 2015a.
Heess, N., Wayne, G., Silver, D., Lillicrap. T., Tassa, Y., and Erez, T. Learning continuous control policies by stochastic value gradients. In Advances in Neural Information Processing Systems, 2015b.
Hochreiter, S. and Schmidhuber, J. Long short-term memory. Neural computation, 903):1335-1780 1997.
Janner, M., Fu, J., Zhang, M., and Levine, S. When to trust your model: Model-based policy optimization. In Advances in Neural Information Processing Systems, 2019.
Janner, M., Mordatch, I., and Levine, S. 7-models: Generative temporal difference learning for
infinite-hone -hriitee-orzzonn prediction. In Advances in Neural Information Processing Systems, 2020.
Jiang, Y., Gu, S., Murphy, K., and Finn, C. Language as abstraction for hierarchical deep reinforcement learning. In Advances in Neural Information Processing Systems, 2019.
Jin, Y., Yang. Z., and Wang. Z. Is pessimism provably efficient for offline RL? In International Conference on Machine Learning, 2021.
Karpathy, A. minGPT: A minimal pytorch re-implementation the openai gpt training, 2020. URL https://gitiee hot..///aaaaaaaatt//nnn......
Kidambi, R., Rajeswaran, A., Netrapalli, P., and Joachims, T. MOReL: Model-based offline rein-
forcement learning. In Advances in Neural Information Processing Systems, 2020.
Kingma, D. P. and Ba. Adam: A method for stochastic optimization. In International Conference
on Learning Representations, 2015.
Kostrikov, I., Nair, A., and Levine. S. Offline reinforcement learning with implicit q-learning. arXiv preprint 1rXiv:2110.061 2021.
Kumar, A., Fu, J., Tucker, G., and Levine, Stabilizing off-policy Q-learning via bootstrapping error reduction. In Advances in Neural Information Processing Systems, 2019a.
12

---

Kumar, A.. Peng. X. B., and Levine, S. Reward-conditione policies. arXiv preprint arXiv:1912.13465 2019b.
Kumar, A., Zhou, A., Tucker, G., and Levine, S. Conservative Q-learning for offline reinforcement learning. Advances in Neural Information Processing Systems, 2020a.
Kumar, S., Parker, J., and Naderian, P. Adaptive transformers in RL. arXiv preprini arXivv:000400371,, 2020b.
Kurutach, T., Clavera, I., Duan, Y., Tamar, A., and Abbeel, P. Model-ensemble trust-region policy optimization. In International Conference Learning Representations, 2018.
Lampe, T. and Riedmiller, M. Approximate model-assisted neural fitted Q-iteration. International Joint Conference Neural Networks, 2014.
Levine, S. Reinforcement learning and control as probabilistic inference: Tutorial and review, arXiv preprint arXiv:1805.00900 2018.
Malik, A., Kuleshov, V., Song, J., Nemer, D., Seymour, H., and Ermon, S. Calibrated model-baser deep reinforcement learning. In International Conference on Machine Learning, 2019.
Meister, C., Cotterell, R., and Vieira, T. If beam search the answer, what was the question? In Empirical Methods in Natural Language Processing. 2020.
Nagabandi, A., Kahn, G., S. Fearing, R., and Levine, S. Neural network dynamics for model-based deep reinforcement learning with model-free fine-tuning, In International Conference on Robotics and Automation, 2018.
Nagabandi, A., Konoglie, K., Levine, S., and Kumar, V. Deep Dynamics Models for Learning Dexterous Manipulation. In Conference Robot Learning. 2019.
Nair, A., Dalal, M., Gupta, A., and Levine, S. Accelerating online reinforcement learning with offline datasets. arXiv preprint arXiv:2006.098: 093599 2020.
Oh, J., Chockalingam, V., Lee, H., et al. Control of memory, active perception, and action in Minecraft. In International Conference Machine Learning, 2016.
Parisotto, E. and Salakhutdinov, Efficient transformers in reinforcement learning using actor- learner distillation. In International Conference Learning Representations, 2021.
Parisotto, E., Song, F., Rae, Pascanu, R., Gulcehre, C.. Jayakumar, S., Jaderberg. M., Kauf- man, R. L.,, Clark, A., Noury, S., et al. Stabilizing transformers for reinforcement learning. In International Conference on Machine Learning, 2020.
Paster, K., Mcllraith, S. A., and Ba, J. Planning from pixels using inverse dynamics models. In International Conference on Learning Representations, 2021.
Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z., Gimelshein,
N., Antiga, L., Desmaison, A., Kopf. A., Yang, E., De Z., Raison, M., Tejani, A.. Chilamkurthy, S., Steiner, B., Fang, L., Bai, J., and Chintala, S. Pytorch: An imperative style, nigh-performance deep learning library. In Advances in Neural Information Processing Systems. 2019.
Peng, X. B., Berseth, G., Yin, K., and Van De Panne, M. Deepl.oco: Dynamic locomotion skills using hierarchical deep reinforcement learning. ACM Transactions on Graphics, 2017.
Radford, A., Narasimhan, K., Salimans, T., and Sutskever, I. Improving language understanding by generative mre-training. 2018.
Rauber, P., Ummadisingu, A., Mutz, F., and Schmidhuber, J. Hindsight policy gradients. In International Conference Learning Representations, 2019,
Reddy, R. Speech understanding systems: Summary of results of the five-year research effort at Carnegie Mellon University, 1977.
13

---

Ross, S. and Bagnell, D. Efficient reductions for imitation learning. In International Conference on Artificial Intelligence and Statistics, 2010.
Ross, S., Gordon, G., and Bagnell, D. A reduction of imitation learning and structured prediction to no-regret online learning. In International Conference on Artificial Intelligence and Statistics, 2011.
Schmidhuber, J. Reinforcement learning upside down: Don't predict rewards-just map them to actions. arXiv preprint (tt.:::99/22227733 2019.
Silver, D., Sutton, R. S., and Müller, M. Sample-based learning and search with permanent and transient memories. In International Conference on Machine Learning, 2008.
Srivastava, R. K., Shyam, P., Mutz, F., Jaškowski, W., and Schmidhuber, J. Training agents using upside-down reinforcement learning. arXiv preprint arXiv: 1912.022/ 2019.
Sutskever, I., Vinyals, O., and Le, Q. V. Sequence to sequence learning with neural networks. In Advances in Neural Information Processing Systems, 2014.
Sutton, R. S. Learning to predict by the methods of temporal differences. Machine Learning, 1988.
Sutton, R. S. Integrated architectures for learning, planning, and reacting based approximating dynamic programming. In International Conference on Machine Learning, 1990.
Sutton, R. S., Precup, D., and Singh, S. Between MDPs and abstraction in reinforcement learning. Artificial Intelligence, 1999.
A framework for temporal
Tay, Y., Dehghani, M., Abnar, S., Shen, Y., Bahri, D., Pham, P., Rao, J., Yang, L., Ruder, S., and Metzler, D. Long range arena: A benchmark for efficient transformers. In International Conference on Learning Representutions, 2021.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L.., Gomez, A. N., Kaiser, E.,, and Polosukhin, I. Attention is all you need. In Advances in Neural Information Processing Systems, 2017.
Wang, T. and Ba, J. Exploring model-based planning with policy networks. In International Conference on Learning Representations, 2020.
Williams, and Zipser, D. A learning algorithm for continually running fully recurrent neural networks. Neural computation, 1989.
Wu, Y., Tucker, G., and Nachum, O. Behavior regularized offline reinforcement learning. arXiv preprint xrXiv:191..1131,, 2019.
Yin, M., Bai, Y., and Wang, Y.-.. Near-optimal offline reinforcement learning via double variance reduction. arXiv preprint arXiv:2102.01748. 2021.
Yu, T., Thomas, G., Yu, L., Ermon, S., Zou, J., Levine, S., Finn, C., and Ma. T. MOPO: Model-based offline policy optimization. In Advances in Neural Information Processing Systems, 2020.
Zhang, M. R., Paine, T., Nachum, O., Paduraru, C., Tucker, G., ziyu wang, and Norouzi, M. Autoregressive dynamics models for offline policy evaluation and optimization. In International Conference on Learning Representations, 2021.
14


Reviews:
[{"review_id": "yMVL3dxFZA", "reviewer": "Reviewer_Ldp2", "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.", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "*Originality*: In my view, the usage of contemporary sequence modeling toolbox as a for sequential decision-making tasks is sufficiently novel. It provides new prospective on how to make use of highly successful transformer archtecture in various contexts and get good results. The related work adequately cited?\n\n*Quality*: The submission is technically sound. Claims are supported by the large number of experimental results. The visualisations and qualitative examples are also assisting the reader.\n\n*Clarity*: The manuscript is clear, well organized and easy to follow. \n\n*Significance*: In my view, these are interesting results that the community can learn from and build upon. \n\n### Cons\n\n- The authors admit that predictions with Transformers is slower and more resource intensive. This is expected. However, I think the reader would benefit from knowing how much slower TTO is compared to baselines. Hence I ask the authors to add that information for their experiments (imitation learning, goal-conditioned RL, and offline RL).\n- I'd be curious to see if goal-reaching is possible in more challenging domains compared to 4 rooms. Perhaps in procedurally generated mazes that change at every episode (e.g. MiniGrid).\n\n# Post rebuttal update\n\nI thanks the authors for their response and for running additional experiments using MiniGrid. I've also read the reviews by my fellow reviewers and the authors rebuttal to them (including many additional experiments). Therefore, I have increased my score by 1.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes. Several limitations of the work are described in the Discussion section. There are no ethical concerns, as far as I'm concerned.\n\nSee above for suggestions.\n\n", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "6oBJoUFOj-K", "reviewer": "Reviewer_oBoc", "summary": "This paper attempts to reformulate RL as a sequence modeling problem and hence exploit the recent advances in sequence modeling like Transformers to design an RL agent. Authors show that with little modification to existing sequence to sequence models and beam search techniques, they can perform well in imitation learning, goal-conditioned RL, and offline RL.", "questions": "", "limitations": "", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "While this is an interesting research direction, I have several major concerns about this work.\n\n1. I do not think that throwing away all the dynamic programming techniques and replacing them with Transformers will help us solve RL. While I disagree with this idea, my decision is not based on the fact that I do not like this approach.\n2. I feel like authors do not highlight the limitations of this simple sequence to sequence approach well enough. Limitations mentioned in section 5 are not serious limitations! Would this approach work when there is stochasticity? Non-stationarity? More diverse tasks? This simple approach would suffer a lot in complex environments and this has to be made clear. While it is ok for having these limitations, it is crucial for this paper to have an explicit limitations section so that it does not mislead the community.\n3. Why did the authors decide to go with the proposed discretization approach? It would be nice to see the performance of some baseline representations to understand the significance of this representation.\n4. Line 133 -> The maximum tokens in 512. What are the values of N and M? How many transitions can you keep in the history if your maximum token length is 512? This is very useful information to know.\n5. Again, Figure-1 experiments are very misleading. You are comparing Transformers models with PETS. PETS uses simple feedforward networks. I would like to see the comparison between Transformers and PlaNet (Hafner et al 2019).\n6. Figure 3 - Are these the only two types of attention you have seen in the trained model? How much cherry-picking happened here? Please explain the selection process here.\n7. How is this work different from [2]? It is not clear whether the performance trends that authors show here would be the same if they use proper recurrent models in their baselines as well.\n\n\nMinor comments:\n\n1. Figure 2 and Figure 4 - Please specify the number of runs.\n2. Figure 4 - What is the y axis?\n\nReferences:\n\n[1] Learning Latent Dynamics for Planning from Pixels. Hafner et al 2019\n[2] Reinforcement Learning Upside Down: Don't Predict Rewards -- Just Map Them to Actions. Schmidhuber 2019\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Check the main review.\n", "review_text": "", "extra_scores": {"ethical_concerns": "None.", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5"}, "extra_text": {}}, {"review_id": "4bf_iBy0ZzN", "reviewer": "Reviewer_6knX", "summary": "In a departure from traditional RL approaches, the paper applies transformer-based sequence modeling (with beam search) to a variety of RL tasks. Evaluations show the proposed Trajectory Transformer (TTO) to be competitive with baselines on three domains.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "- Originality\n\nThe proposed approach is novel and well-motivated, given the successes of transformers in sequence modeling.\n\n\n- Quality\n\nThe experimental results indicate that TTO can be usefully applied to at least a few, somewhat peripheral (non-online) tasks that fall within the broad collection of work associated with RL. However, the paper is worded in a way that is easy to interpret too broadly, starting with the catchy title, which could make many readers think of online RL, since that is the core of the field. Some readers may not even realize that online RL is excluded by the final sentence of the abstract, since imitation learning and offline RL are common pretraining steps performed prior to online RL training, and since goal conditioning often appears in the online RL setting. Expectations could be better managed by making the title less all-inclusive. For instance, “**Non-Online** Reinforcement Learning as One Big Sequence Modeling Problem”. (I’m sure the authors could think of something better than I just did.) And line 308 in the conclusion could also be more carefully qualified, as “**certain** reinforcement learning problems”.\n\nAny reader who initially assumes that this work reduces online RL to sequence modeling, freeing them from the tyranny of estimating cumulative reward, will be disappointed by the fine print that follows. As explained in section 3.2, this paper’s imitation learning task is not of the usual kind, and the goal-conditioned task is outside the online setting. Line 86 says “Our model treats states, actions, and rewards interchangeably”, but this is not quite true for the offline RL task, where reward-to-go is used as the beam search heuristic. Line 193 argues that these reward estimates need not be particularly accurate, but why assume that online RL requires more accurate value estimates than beam search does?\n\nI don’t believe that a bold new RL approach such as sequence modeling needs to achieve SOTA or even be applicable to the most common RL problems in order to merit publication. But I do believe the claims should be more explicitly scoped.\n\n\n- Clarity\n\nApart from needing greater clarity in its claims, the paper is generally clear and well-written. I see a few other areas for improvement.\n\nLine 107 says that continuous states and actions were discretized “for processing by a discrete-token architecture”, leaving readers to assume that the transformer architecture itself is being discussed. But transformers have already been successfully applied to online RL in a variety of ways, using both continuous and discrete observations and/or actions. So it would help to explain more directly that the discrete-token architecture chosen for this work was minGPT along with its discrete token embedding and decoding machinery.\n\nThis work replaces some of the traditional RL machinery with beam search, which is a tricky tool with its own practical limitations that deserve much greater attention for the benefit of potential users in the RL community. Algorithm 1 describes beam search at a certain level, but the algorithm box or accompanying text should include basic explanations of the individual symbols and how they are to be used in an RL task. In addition, the critical issues like beam size, branching factor, search depth, and heuristic pruning factors need to be given concrete numbers in the discussion of each task. Many of these details could be relegated to an appendix, but they are as important as any hyperparameter value. \n\nWhat is the interpretation of the vertical axis in Figure 4 and Figure 6?\n\nThe binary footprint of Figure 5 should be compressed. It dramatically slows down my browser every time I try to scroll past it.\n\n\n- Significance\n\nThis is an interesting new approach that may deliver benefits in certain types of RL problems, and may even lead to useful online techniques.\n\n\n- Conclusion\n\nGiven the shortcomings described above, I cannot support publication of the paper in its current form. But I believe the work makes important contributions which deserve a wide audience, so I hope the paper’s claims can be clarified so as not to mislead over-anxious readers.\n\n- POST-RESPONSES UPDATE\n\nThe authors have addressed all of my reservations. With these revisions, the paper seems properly scoped and easier to understand. I have raised my score by 2 points.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "yes", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "12"}, "extra_text": {}}]

Author rebuttal:

Meta-review: The reviews were thorough and there was good interaction with the authors. 

I'd like to suggest that the authors may have come off as overly aggressive in their pushing for additional reviewer interaction, especially when the interaction was already quite good. This aggressiveness may backfire sometime in the future. 

That said, the paper does seem to make a reasonable contribution and the delta between what is promised and the original submission is not overly large. 
Decision: Accept (Spotlight)