Title: Elastic Decision Transformer
Venue: NeurIPS 2023
Accepted: yes
Keywords: Offline Reinforcement Learning

Trajectory Stitching

Decision Transformer
Abstract: This paper introduces Elastic Decision Transformer (EDT), a significant advancement over the existing Decision Transformer (DT) and its variants. Although DT purports to generate an optimal trajectory, empirical evidence suggests it struggles with trajectory stitching, a process involving the generation of an optimal or near-optimal trajectory from the best parts of a set of sub-optimal trajectories. The proposed EDT differentiates itself by facilitating trajectory stitching during action inference at test time, achieved by adjusting the history length maintained in DT. Further, the EDT optimizes the trajectory by retaining a longer history when the previous trajectory is optimal and a shorter one when it is sub-optimal, enabling it to "stitch" with a more optimal trajectory. Extensive experimentation demonstrates EDT's ability to bridge the performance gap between DT-based and Q Learning-based approaches. In particular, the EDT outperforms Q Learning-based methods in a multi-task regime on the D4RL locomotion benchmark and Atari games.

Elastic Decision Transformer

Abstract

This paper introduces Elastic Decision Transformer (EDT), a significant advancement over the existing Decision Transformer (DT) and its variants. Although DT purports to generate an optimal trajectory, empirical evidence suggests it struggles with trajectory stitching, a process involving the generation of an optimal or near-optimal trajectory from the best parts of a set of sub-optimal trajectories. The proposed EDT differentiates itself by facilitating trajectory stitching during action inference at test time, achieved by adjusting the history length maintained in DT. Further, the EDT optimizes the trajectory by retaining a longer history when the previous trajectory is optimal and a shorter one when it is sub-optimal, enabling it to "stitch" with a more optimal trajectory. Extensive experimentation demonstrates EDT's ability to bridge the performance gap between DT-based and Q Learning-based approaches. In particular, the EDT outperforms Q Learning-based methods in a multi-task regime on the D4RL locomotion benchmark and Atari games. Videos are available at: https://kristery.github.io/edt/ .

1 Introduction

Reinforcement Learning (RL) trains agents to interact with an environment and learn from rewards. It has demonstrated impressive results across diverse applications such as game playing [23, 44], robotics [41, 40, 54], and recommendation systems [12, 1]. A notable area of RL is Offline RL [31], which employs pre-collected data for agent training and proves more efficient when real-time interactions are costly or limited. Recently, the conditional policy approach has shown large potentials in Offline RL, where the agent learns a policy based on the observed state and a goal. This approach enhances performance and circumvents stability issues related to long-term credit assignment. Moreover, the successful Transformer architecture [49], widely

61

45

50

89

Hopper

43

30

61

74

Walker2d

DT

QDT

TS+BC

Ours

used in applications like natural language processing [52, 13, 8] and computer vision [14, 34], has been adapted for RL as the Decision Transformer (DT) [11].

DT utilizes a Transformer architecture to model and reproduce sequences from demonstrations, integrating a goal-conditioned policy to convert Offline RL into a supervised learning task. Despite its competitive performance in Offline RL tasks, the DT falls short in achieving trajectory stitching, a desirable property in Offline RL that refers to creating an optimal trajectory by combining parts of sub-optimal trajectories [19, 9, 59]. This limitation stems from the DT's inability to generate

∗

superior sequences, thus curbing its potential to learn optimal policies from sub-optimal trajectories (Figure 1).

We introduce the Elastic Decision Transformer (EDT), which takes a variable length of the traversed trajectory as input. Stitching trajectories, or integrating the current path with a more advantageous future path, poses a challenge for sequence generation-based approaches in offline RL. Stitching a better trajectory appears to contradict one of the core objectives of sequence generation that a sequence generation model is required to reliably reproduce trajectories found within the training dataset. We suggest that in order to 'refresh' the prediction model, it should disregard 'negative' or 'unsuccessful' past experiences. This involves dismissing past failures and instead considering a shorter history for input. This allows the sequence generation model to select an action that yields a more favorable outcome. This strategy might initially seem contradictory to the general principle that decisions should be based on as much information as possible. However, our proposed approach aligns with this concept. With a shorter history, the prediction model tends to output with a higher variance, typically considered a weakness in prediction scenarios. Yet, this increased variance offers the sequence prediction model an opportunity to explore and identify improved trajectories. Conversely, when the current trajectory is already optimal, the model should consider the longest possible history for input to enhance stability and consistency. Consequently, a relationship emerges between the quality of the path taken and the length of history used for prediction. This correlation serves as the motivation behind our proposal to employ a variable length of historical data as input.

In practice, we train an approximate value maximizer using expectile regression to estimate the highest achievable value given a certain history length. We then search for the history length associated with the highest value and use it for action inference.

Evidence from our studies indicates that EDT's variable-length input sequence facilitates more effective decision-making and, in turn, superior sequence generation compared to DT and its variants. Furthermore, it is computationally efficient, adding minimal overhead during training. Notably, EDT surpasses state-of-the-art methods, as demonstrated in the D4RL benchmark [15] and Atari games [7, 17]. Our analysis also suggests that EDT can significantly enhance the performance of DT, establishing it as a promising avenue for future exploration.

Our Contributions:

We introduce the Elastic Decision Transformer, a novel approach to Offline Reinforcement Learning that effectively addresses the challenge of trajectory stitching, a known limitation in Decision Transformer.

By estimating the optimal history length based on changes in the maximal value function, the EDT enhances decision-making and sequence generation over traditional DT and other Offline RL algorithms.

Our experimental evaluation highlights EDT's superior performance in a multi-task learning regime, positioning it as a promising approach for future Offline Reinforcement Learning research and applications.

2 Preliminaries

In this study, we consider a decision-making agent that operates within the framework of Markov Decision Processes (MDPs) [42]. At every time step t , the agent receives an observation of the world o t , chooses an action a t , and receives a scalar reward r t . Our goal is to learn a single optimal policy distribution P ∗ θ ( a t | o ≤ t , a <t , r <t ) with parameters θ that maximizes the agent's total future return R t = ∑ k>t r k on all the environments we consider.

2.1 Offline Reinforcement Learning

Offline RL, also known as batch RL, is a type of RL where an agent learns to make decisions by analyzing a fixed dataset of previously collected experiences, rather than interacting with an environment in real-time. In other words, the agent learns from a batch of offline data rather than actively exploring and collecting new data online.

Offline RL has gained significant attention in recent years due to its potential to leverage large amounts of pre-existing data and to solve RL problems in scenarios where online exploration is

+

t+1

R

R

a

causaltransformer decoder

Atari

+

t

工

R

k>t

a

observation

return

action

impractical or costly. Examples of such scenarios include medical treatment optimization [33], finance [46], and recommendation systems [55].

Despite its potential benefits, offline RL faces several challenges, such as distributional shift, which occurs when the offline data distribution differs significantly from the online data distribution, and the risk of overfitting to the fixed dataset. A number of recent research efforts have addressed these challenges, including methods for importance weighting [35, 37], regularization [53, 28], and model-based learning [25, 29], among others.

2.2 Decision Transformer

The Decision Transformer architecture, introduced by [11], approaches the offline RL problem as a type of sequence modeling. Unlike many traditional RL methods that estimate value functions or compute policy gradients, DT predicts future actions based on a sequence of past states, actions, and rewards. The input to DT includes a sequence of past states, actions, and rewards, and the output is the next action to be taken. DT uses a Transformer architecture [49], which is composed of stacked self-attention layers with residual connections. The Transformer architecture has been shown to effectively process long input sequences and produce accurate outputs.

Despite the success of being applied to offline RL tasks, it has a limitation in its ability to perform "stitching." Stitching refers to the ability to combine parts of sub-optimal trajectories to produce an optimal trajectory. This approach can lead to a situation where the agent follows a sub-optimal trajectory that provides an immediate reward, even if a different trajectory leads to a higher cumulative reward over time. This limitation of DT is a significant challenge in many offline RL applications, and addressing it would greatly enhance the effectiveness of DT in solving real-world problems.

3 Elastic Decision Transformer

In this section, we present Elastic Decision Transformer (EDT), a model that automatically utilizes a shorter history to predict the next action when the traversed trajectory underperforms compared to those in the training dataset. The mechanism allows the model to switch to a better trajectory by forgetting 'unsuccessful' past experiences, thus opening up more possibilities for future trajectories. We further propose a method to estimate the maximum achieveable return using the truncated history, allowing EDT to determine the optimal history length and corresponding actions.

3.1 Reinforcement Learning as Sequence Modeling

In this paper, we adopt an approach to offline reinforcement learning that is based on a sequence modeling problem. Specifically, we model the probability of the next token in the sequence (denoted as τ ) based on all the tokens that come before it. The sequences we model can be represented as:

where t is a time step and ˆ R is the return for the remaining sequence. The sequence we consider here is similar to the one used in [30] whereas we do not include reward as part of the sequence and we predict an additional quantity ˜ R that enables us to estimate an optimal input length, which we will cover in the following paragraphs. Figure 2 presents an overview of our model architecture. It should

be noted that we also change the way to predict future observation from standard DT [11], where the next observation is usually directly predicted from a t through the causal transformer decoder.

3.2 Motivation

We propose a shift in the traditional approach to trajectory stitching. Instead of focusing on training phases, we aim to achieve this stitching during the action inference stage. This concept is illustrated in Figure 3 using a simplified example. In this scenario, we consider a dataset, D , comprising only two trajectories: D = ( s a t -1 , s t , s a t +1 ) , ( s b t -1 , s t , s b t +1 ) . A sequence model trained with this dataset is likely to predict the next states in a manner consistent with their original trajectories.

In the given example, if the model starts at state s b t -1 , it will choose to retain the history ( s t ) upon reaching state s t , as ˜ R ( s t ) > ˜ R ( s t -1 , s t ) . Conversely, if the model initiates from state s a t -1 , it will preserve the history ( s a t -1 , s t ) when decision-making at s t , as ˜ R ( s a t -1 , s t ) ≥ ˜ R ( s t ) . From the above example, we understand that the optimal history length depends on the quality of the current trajectory we've traversed, and it can be a specific length anywhere between a preset maximal length and a single unit.

To estimate the optimal history length in a general scenario, we propose solving the following optimization problem:

where τ T denotes the history length T . More precisely, τ T takes the form:

3.3 Training objective for Maximum In-Support Return

In the EDT, we adhere to the same training procedure as used in the DT. The key distinction lies in the training objective - we aim to estimate the maximum achievable return for a given history length in EDT. To approximate the maximum operator in max τ T ∈ D ˆ R t ( τ T ) , we employ expectile regression [38, 3], a technique often used in applied statistics and econometrics. This method has previously been incorporated into offline reinforcement learning; for instance, IQL [26] used expectile regression to estimate the Q-learning objective implicitly. Here, we leverage it to enhance our estimation of the maximum expected return for a trajectory, even within limited data contexts. The α ∈ (0 , 1) expectile of a random variable X is the solution to an asymmetric least squares problem, as follows:

where L α 2 ( u ) = | α -1 ( u < 0) | u 2 .

Through expectile regression, we can approximate max τ T ∈ D ˆ R t ( τ T ) :

We estimate ˜ R t by applying an empirical loss of Equation 2 with a sufficiently large α (we use α = 0 . 99 in all experiments). The only difference in training EDT compared to other DT variants is the use of Equation 2, making the training time comparably shorter. We summarize our objective as:

at

causaltransformerdecoder

↑

ot

R;(T1)

R1(T2)

Bayes' Rule

T2

ot-1,Rt-1,at-1,o

RW

R;(T2)

causaltransformerdecoder

at-1

ot

TT

t-1,Rt-1

at-1,0

Rk(TT)

Optimal length w

Algorithm 1 EDT optimal history length search

1: Input: a query sequence τ = ⟨ ..., o t -1 , R t -1 , a t -1 , o t ⟩ and EDT model θ EDT

2: for w = 1 , 1 + δ, 1 + 2 δ, ..., T do

3: Obtain ˜ R t ( τ w ) with truncated sequence τ w = ⟨ o t -w +1 , R t -w +1 , a t -w +1 , ..., o t ⟩ and θ EDT

4: end for

5: Compute P ( R t ) with θ EDT for τ w that has the highest ˜ R t and then estimate P ( R t | expert t , ... ) with Equation 4.

where L observation and L action are computed with a mean square error, L return is a cross-entropy loss, and L max is an empirical estimate of Equation 2. We set c r = 0 . 001 to balance scale differences between mean square error and cross-entropy loss. In tasks with discrete action spaces like Atari, we optimize the action space as the tokenized return objective L return using cross-entropy with weight 10 c r . Our training method extends the work of [30] by estimating the maximum expected return value for a trajectory using Equation 2. This estimation aids in comparing expected returns of different trajectories over various history lengths. Our proposed method is not only easy to optimize, but can also be conveniently integrated with other DT variants. As such, it marks a significant advance in developing efficient offline reinforcement learning approaches for complex decision-making tasks.

3.4 Action Inference During Test time

During action inference phase in test time, we first (1) estimate the maximum achievable return ˜ R i for each history length i . Subsequently, (2) we predict the action by using the truncated traversed trajectory as input. The trajectory is truncated to the history length that corresponds to the highest value of ˜ R i . These steps are elaborated in Figure 4.

To identify the history length i that corresponds to the highest ˜ R t i , we employ a search strategy as detailed in Algorithm 1. As exhaustively searching through all possible lengths from 1 to T may result in slow action inference, we introduce a step size δ to accelerate the process. This step size not only enhances inference speed by a factor of δ , but also empirically improves the quality of the learned policy. An ablation study on the impact of the step size δ is provided in Appendix A. For all experiments, we set δ = 2 to eliminate the need for parameter tuning.

To sample from expert return distribution P ( R t , ... | expert t ) , we adopt an approach similar to [30] by applying Bayes' rule P ( R t , ... | expert t ) ∝ P ( expert t | R t , ... ) P ( R t , ... ) and approximate the distribu-

Table 1: Baseline comparisons on D4RL [15] tasks. Mean of 5 random training initialization seeds, 100 evaluations each. Our result is highlighted. The results of QDT, TS+BC, and S4RL are adopted from their reported scores. Following [26], we emphasize in bold scores within 5 percent of the maximum per task ( ≥ 0 . 95 · max ) .

tion of expert-level return with inverse temperature κ [24, 48, 47, 43]:

While it may initially appear feasible to directly use the predicted ˜ R as the expert return, it's important to note that this remains a conservative maximum operation. Empirically, we have found that Eq. 4 encourages the pursuit of higher returns, which consequently enhances the quality of the actions taken.

4 Experiments

Our experiments are designed to address several key questions, each corresponding to a specific section of our study:

Does EDT significantly outperform DT and its variants? (Sec. 4.2, 4.3)

Is the EDT effective in a multi-task learning regime, such as Locomotion and Atari games? (Sec. 4.3)

Does a dynamic history length approach surpass a fixed length one? (Sec. 4.4)

How does the expectile level α impact the model's performance? (Sec. 4.5)

How does the quality of datasets affect the predicted history lengths? (Sec. 4.6)

We also provide an additional ablation study in Appendix A due to space constraints.

4.1 Baseline Methods

In the subsequent section, we draw comparisons with two methods based on the Decision Transformer: the original Decision Transformer (DT) [11] and the Q-learning Decision Transformer (QDT) [59]. Additionally, we include a behavior cloning-based method (TS+BC) [19], as well as two offline Q-learning methods, namely S4RL [45] and IQL [26], in our comparisons.

It is important to note that QDT and TS+BC are specifically designed to achieve trajectory stitching. QDT accomplishes this by substituting collected return values with estimates derived from Conservative Q-Learning [28]. Conversely, TS+BC employs a model-based data augmentation strategy to bring about the stitching of trajectories."

4.2 Single-Task Offline Reinforcement Learning

For locomotion tasks, we train offline RL models on D4RL's 'medium' and 'medium-replay' datasets. The 'medium' dataset comes from a policy reaching about a third of expert performance. The ''medium-replay', sourced from this policy's replay buffer, poses a greater challenge for sequence modeling approaches such as DT.

We conclude our locomotion results in Table 1. Since the proposed model estimates the return of the current sequence, reward information is not required during test time.

Our observations indicate that the proposed EDT consistently outperforms the baseline DT and its variants on the majority of the datasets, with a notable performance advantage on the "medium-













































replay" datasets. These findings provide strong evidence that our approach is highly effective in stitching together sub-optimal trajectories with high return proportion, a task that DT and its variants cannot accomplish. Although EDT doesn't fully outperform IQL in the single-task learning, it does bridge the gap between Q-learning-based methods and DT by performing trajectory stitching with the estimated maximum return.

4.3 Multi-Task Offline Reinforcement Learning

This section aims to evaluate the multi-task learning ability of our model across diverse tasks, focusing on locomotion and Atari tasks. Locomotion tasks utilize vectorized observations, while Atari tasks depend on image observations. To emphasize the role of trajectory stitching, we restrict our datasets to medium-replay datasets for the four locomotion tasks and datasets derived from DQN Replay [2] for the Atari tasks. Our evaluations span 20 different Atari tasks, with further environment setup details available in the Appendix.

Locomotion. In the locomotion multi-task experiment, we maintain the same model architecture as in the single-task setting. By confining the dataset to medium-replay datasets from four tasks, we increase task complexity and necessitate the offline RL approach to learn and execute these tasks concurrently, while effectively utilizing trajectories generated by random policies. As depicted in Table 2, our proposed EDT successfully accomplishes all four tasks simultaneously without much performance compromise.

Atari. For Atari, we adopt a CNN image encoder used in DrQ-v2 [59] to process stacks of four 84 x 84 image observations. To ensure fair comparisons, all methods employ the same architecture for the image encoder. Following [30], we incorporate random cropping and rotation for image augmentation. Additional experiment details are delegated to the Appendix for brevity. Performance on each Atari game is measured by human normalized scores (HNS) [36], defined as ( score -scorerandom ) / ( scorehuman -scorerandom ) , to ensure a consistent scale across each game.

Our experimental results align with those of [30], highlighting that Q-learning-based offline RL approaches encounter difficulties in learning a multi-task policy on Atari games. Despite IQL achieving the highest score in Table1, it demonstrates relative inadequacy in simultaneous multi-task learning as indicated in Table 2 and Figure 5. We leave the raw scores of the 20 Atari games in Appendix B.



















































































4.4 Dynamic History Length vs. Fixed History Length

In Sec. 3.2, we proposed the concept of EDT, which adjusts history length based on the quality of the current trajectory. We illustrated this idea with a toy example.

To validate the benefits of this dynamic history length approach, we tested the EDT using both fixed and variable history lengths. The results, summarized in Table 3, show that the variable history length outperforms the fixed ones, particularly on the "medium-replay" datasets.

These findings suggest that the EDT effectively chooses a history length that yields a higher estimated return. While a shorter history length aids in trajectory stitching, it's also crucial to retain a longer history length to ensure the continuity of optimal trajectories. Therefore, the dynamic adjustment of history length in the EDT is key to its superior performance.

4.5 Ablation Study on Expectile Level α

A key component of EDT is the approximation of the maximal value using expectile learning, as our method depends on accurately estimating these maximal values to choose the optimal history length. Consequently, examining the change in performance relative to the expectile level, α , provides insight into the necessity of correct history length selection for performance enhancement.

The results, as displayed in Figure 6, suggest that when expectile regression is able to accurately approximate the maximizer, specifically at higher expectile levels, we observe both a higher average performance and lower standard deviation. This suggests that accurate selection of history length not only stabilizes performance but also enhances scores. Conversely, as the expectile level approaches 0 . 5 , the expectile regression's objective shifts towards a mean square error, leading to an estimated value that is more of a mean value than a maximal one. This change makes it a less effective indicator for optimal history length. As a result, we can see a deterioration in EDT's score as the expectile level drops too low, and an increase in standard deviation, indicating inconsistency in the selection of an effective history length.

4.6 Analysis of Optimal History Length Distribution

In our analysis, we examine the history length distributions across various datasets, as depicted in Figure 7. Our findings reveal that the medium-replay dataset, which amalgamates trajectories from multiple policies, yields a distribution closely approximating a uniform distribution. Conversely, the medium dataset, acquired through a singular stochastic policy, exhibits a history length distribution characterized by an increased density at lower history lengths. This observation can be attributed to the prevalence of analogous trajectories within the medium dataset, leading to more frequent occurrences of trajectory stitching than the 'medium-replay' dataset. However, it is important to acknowledge that the gains derived from this type of trajectory stitching remain limited, as the trajectories stem from identical distributions. Although performance improvement is observed, as presented in Table 1, it is significantly less pronounced in comparison to the medium-replay dataset.

Contrary to initial expectations, trajectory stitching does not occur as frequently within the mediumreplay dataset as within the medium dataset. In fact, the distinct policies within the medium dataset contribute to the reduced instances of trajectory stitching, as their respective state distributions differ from one another. The diversity within the dataset results in a limited number of mutual s t

instances illustrated in Figure 3. Nevertheless, the proposed EDT method derives substantial benefits from trajectory stitching in this context. The EDT effectively avoids being misled by sub-optimal trajectories within the dataset, demonstrating its capacity to make better decisions regarding history lengths and actions that optimize the current return.

5 Related Work

Offline Reinforcement Learning. Offline RL has been a promising topics for researchers since sampling from environments during training is usually costly and dangerous in real-world applications and offline reinforcement learning is able to learn a better policy without directly collecting stateaction pairs. Several previous works have utilized constrained or regularized dynamic programming to mitigate deviations from the behavior policy [39, 51, 27, 16, 56].

Decision Transformer and its variants [11, 30, 59, 58, 57] have been a promising direction for solving offline RL from the perspective of sequence modeling. Trajectory Transformer (TT) [22] models distributions over trajectories using transformer architecture. The approach also incorporates beam search as a planning algorithm and demonstrates exceptional flexibility across various applications, such as long-horizon dynamics prediction, imitation learning, goal-conditioned reinforcement learning, and offline reinforcement learning.

Recently, there has been a growing interest in incorporating diffusion models into offline RL methods. This alternative approach to decision-making stems from the success of generative modeling, which offers the potential to address offline RL problems more effectively. For instance, [18] reinterprets Implicit Q-learning as an actor-critic method, using samples from a diffusion parameterized behavior policy to improve performance. Similarly, other diffusion-based methods [21, 32, 50, 5, 10, 4] utilize diffusion-based generative models to represent policies or model dynamics, achieving competitive or superior performance across various tasks.

Trajectory Stitching. A variety of methods have been proposed to tackle the trajectory stitching problem in offline RL. The Q-learning Decision Transformer (QDT) [59] stands out as it relabels the ground-truth return-to-go with estimated values, a technique expected to foster trajectory recombination. Taking a different approach, [19] utilizes a model-based data augmentation strategy, stitching together parts of historical demonstrations to create superior trajectories. Similarly, the Best Action Trajectory Stitching (BATS) [9] algorithm forms a tabular Markov Decision Process over logged data, adding new transitions using short planned trajectories. BATS not only aids in identifying advantageous trajectories but also provides theoretical bounds on the value function. These efforts highlight the breadth of strategies employed to improve offline RL through innovative trajectory stitching techniques.

6 Discussion

Conclusion. In this paper, we introduced the Elastic Decision Transformer, a significant enhancement to the Decision Transformer that addresses its limitations in offline reinforcement learning. EDT's innovation lies in its ability to determine the optimal history length, promoting trajectory stitching. We proposed a method for estimating this optimal history length by learning an approximate value optimizer through expectile regression.

Our experiments affirmed EDT's superior performance compared to DT and other leading offline RL algorithms, notably in multi-task scenarios. EDT's implementation is computationally efficient and straightforward to incorporate with other DT variants. It outshines existing methods on the D4RL benchmark and Atari games, underscoring its potential to propel offline RL forward.

In summary, EDT offers a promising solution for trajectory stitching, enabling the creation of better sequences from sub-optimal trajectories. This capability can considerably enhance DT variants, leading to improved performance across diverse applications. We are committed to releasing our code.

Limitations. A potential direction for future improvement involves enhancing the speed at which EDT estimates the optimal history. This could make the method suitable for real-time applications that have strict time constraints. While this adaptation is an exciting avenue for future research, it falls outside the primary scope of this paper.

References

MMehdi Afsar, Trafford Crump, and Behrouz Far. Reinforcement learning based recommender systems: A survey. ACM Computing Surveys , 55(7):1-38, 2022.

Rishabh Agarwal, Dale Schuurmans, and Mohammad Norouzi. An optimistic perspective on offline reinforcement learning. In International Conference on Machine Learning , pages 104-114. PMLR, 2020.

Dennis J Aigner, Takeshi Amemiya, and Dale J Poirier. On the estimation of production frontiers: maximum likelihood estimation of the parameters of a discontinuous density function. International economic review , pages 377-396, 1976.

Anurag Ajay, Yilun Du, Abhi Gupta, Joshua Tenenbaum, Tommi Jaakkola, and Pulkit Agrawal. Is conditional generative modeling all you need for decision-making? arXiv preprint arXiv:2211.15657 , 2022.

Arthur Argenson and Gabriel Dulac-Arnold. Model-based offline planning. arXiv preprint arXiv:2008.05556 , 2020.

Adrià Puigdomènech Badia, Bilal Piot, Steven Kapturowski, Pablo Sprechmann, Alex Vitvitskyi, Zhaohan Daniel Guo, and Charles Blundell. Agent57: Outperforming the atari human benchmark. In International conference on machine learning , pages 507-517. PMLR, 2020.

Marc G Bellemare, Yavar Naddaf, Joel Veness, and Michael Bowling. The arcade learning environment: An evaluation platform for general agents. Journal of Artificial Intelligence Research , 47:253-279, 2013.

Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems , 33:1877-1901, 2020.

Ian Char, Viraj Mehta, Adam Villaflor, John M Dolan, and Jeff Schneider. Bats: Best action trajectory stitching. arXiv preprint arXiv:2204.12026 , 2022.

Huayu Chen, Cheng Lu, Chengyang Ying, Hang Su, and Jun Zhu. Offline reinforcement learning via high-fidelity generative behavior modeling. arXiv preprint arXiv:2209.14548 , 2022.

Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems , 34:15084-15097, 2021.

Xinshi Chen, Shuang Li, Hui Li, Shaohua Jiang, Yuan Qi, and Le Song. Generative adversarial user model for reinforcement learning based recommendation system. In International Conference on Machine Learning , pages 1052-1061. PMLR, 2019.

Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805 , 2018.

Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.11929 , 2020.

Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning. arXiv preprint arXiv:2004.07219 , 2020.

Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. In International conference on machine learning , pages 2052-2062. PMLR, 2019.

Caglar Gulcehre, Ziyu Wang, Alexander Novikov, Thomas Paine, Sergio Gómez, Konrad Zolna, Rishabh Agarwal, Josh S Merel, Daniel J Mankowitz, Cosmin Paduraru, et al. Rl unplugged: A suite of benchmarks for offline reinforcement learning. Advances in Neural Information Processing Systems , 33:7248-7259, 2020.

Philippe Hansen-Estruch, Ilya Kostrikov, Michael Janner, Jakub Grudzien Kuba, and Sergey Levine. Idql: Implicit q-learning as an actor-critic method with diffusion policies. arXiv preprint arXiv:2304.10573 , 2023.

Charles A Hepburn and Giovanni Montana. Model-based trajectory stitching for improved offline reinforcement learning. arXiv preprint arXiv:2211.11603 , 2022.

Ari Holtzman, Jan Buys, Li Du, Maxwell Forbes, and Yejin Choi. The curious case of neural text degeneration. arXiv preprint arXiv:1904.09751 , 2019.

Michael Janner, Yilun Du, Joshua B Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. arXiv preprint arXiv:2205.09991 , 2022.

Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as one big sequence modeling problem. Advances in neural information processing systems , 34:12731286, 2021.

Lukasz Kaiser, Mohammad Babaeizadeh, Piotr Milos, Blazej Osinski, Roy H Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, et al. Modelbased reinforcement learning for atari. arXiv preprint arXiv:1903.00374 , 2019.

Hilbert J Kappen, Vicenç Gómez, and Manfred Opper. Optimal control as a graphical model inference problem. Machine learning , 87:159-182, 2012.

Rahul Kidambi, Aravind Rajeswaran, Praneeth Netrapalli, and Thorsten Joachims. Morel: Model-based offline reinforcement learning. Advances in neural information processing systems , 33:21810-21823, 2020.

Ilya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit q-learning. arXiv preprint arXiv:2110.06169 , 2021.

Aviral Kumar, Justin Fu, Matthew Soh, George Tucker, and Sergey Levine. Stabilizing offpolicy q-learning via bootstrapping error reduction. Advances in Neural Information Processing Systems , 32, 2019.

Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. Advances in Neural Information Processing Systems , 33:11791191, 2020.

Byung-Jun Lee, Jongmin Lee, and Kee-Eung Kim. Representation balancing offline modelbased reinforcement learning. In International Conference on Learning Representations , 2021.

Kuang-Huei Lee, Ofir Nachum, Mengjiao Yang, Lisa Lee, Daniel Freeman, Winnie Xu, Sergio Guadarrama, Ian Fischer, Eric Jang, Henryk Michalewski, et al. Multi-game decision transformers. arXiv preprint arXiv:2205.15241 , 2022.

Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint arXiv:2005.01643 , 2020.

Zhixuan Liang, Yao Mu, Mingyu Ding, Fei Ni, Masayoshi Tomizuka, and Ping Luo. Adaptdiffuser: Diffusion models as adaptive self-evolving planners. arXiv preprint arXiv:2302.01877 , 2023.

Ran Liu, Joseph L Greenstein, James C Fackler, Jules Bergmann, Melania M Bembea, and Raimond L Winslow. Offline reinforcement learning with uncertainty for treatment strategies in sepsis. arXiv preprint arXiv:2107.04491 , 2021.

Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. In Proceedings of the IEEE/CVF international conference on computer vision , pages 10012-10022, 2021.

Eric Mitchell, Rafael Rafailov, Xue Bin Peng, Sergey Levine, and Chelsea Finn. Offline metareinforcement learning with advantage weighting. In International Conference on Machine Learning , pages 7780-7791. PMLR, 2021.

Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Marc G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. nature , 518(7540):529-533, 2015.

Ashvin Nair, Abhishek Gupta, Murtaza Dalal, and Sergey Levine. Awac: Accelerating online reinforcement learning with offline datasets. arXiv preprint arXiv:2006.09359 , 2020.

Whitney K Newey and James L Powell. Asymmetric least squares estimation and testing. Econometrica: Journal of the Econometric Society , pages 819-847, 1987.

Xue Bin Peng, Aviral Kumar, Grace Zhang, and Sergey Levine. Advantage-weighted regression: Simple and scalable off-policy reinforcement learning. arXiv preprint arXiv:1910.00177 , 2019.

Matthias Plappert, Marcin Andrychowicz, Alex Ray, Bob McGrew, Bowen Baker, Glenn Powell, Jonas Schneider, Josh Tobin, Maciek Chociej, Peter Welinder, et al. Multi-goal reinforcement learning: Challenging robotics environments and request for research. arXiv preprint arXiv:1802.09464 , 2018.

Athanasios S Polydoros and Lazaros Nalpantidis. Survey of model-based reinforcement learning: Applications on robotics. Journal of Intelligent & Robotic Systems , 86(2):153-173, 2017.

Martin L Puterman. Markov decision processes. Handbooks in operations research and management science , 2:331-434, 1990.

Ross D Shachter. Probabilistic inference and influence diagrams. Operations research , 36(4):589-604, 1988.

David Silver, Thomas Hubert, Julian Schrittwieser, Ioannis Antonoglou, Matthew Lai, Arthur Guez, Marc Lanctot, Laurent Sifre, Dharshan Kumaran, Thore Graepel, et al. A general reinforcement learning algorithm that masters chess, shogi, and go through self-play. Science , 362(6419):1140-1144, 2018.

Samarth Sinha, Ajay Mandlekar, and Animesh Garg. S4rl: Surprisingly simple self-supervision for offline reinforcement learning in robotics. In Conference on Robot Learning , pages 907-917. PMLR, 2022.

Farzan Soleymani and Eric Paquet. Financial portfolio optimization with online deep reinforcement learning and restricted stacked autoencoder-deepbreath. Expert Systems with Applications , 156:113456, 2020.

Emanuel Todorov. Linearly-solvable markov decision problems. Advances in neural information processing systems , 19, 2006.

Marc Toussaint. Robot trajectory optimization using approximate inference. In Proceedings of the 26th annual international conference on machine learning , pages 1049-1056, 2009.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems , 30, 2017.

Zhendong Wang, Jonathan J Hunt, and Mingyuan Zhou. Diffusion policies as an expressive policy class for offline reinforcement learning. arXiv preprint arXiv:2208.06193 , 2022.

Ziyu Wang, Alexander Novikov, Konrad Zolna, Josh S Merel, Jost Tobias Springenberg, Scott E Reed, Bobak Shahriari, Noah Siegel, Caglar Gulcehre, Nicolas Heess, et al. Critic regularized regression. Advances in Neural Information Processing Systems , 33:7768-7778, 2020.

Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, et al. Transformers: Stateof-the-art natural language processing. In Proceedings of the 2020 conference on empirical methods in natural language processing: system demonstrations , pages 38-45, 2020.

Yifan Wu, George Tucker, and Ofir Nachum. Behavior regularized offline reinforcement learning. arXiv preprint arXiv:1911.11361 , 2019.

Yueh-Hua Wu, Jiashun Wang, and Xiaolong Wang. Learning generalizable dexterous manipulation from human grasp affordance. In Conference on Robot Learning , pages 618-629. PMLR, 2023.

Teng Xiao and Donglin Wang. A general offline reinforcement learning framework for interactive recommendation. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 35, pages 4512-4520, 2021.

Haoran Xu, Li Jiang, Li Jianxiong, and Xianyuan Zhan. A policy-guided imitation approach for offline reinforcement learning. Advances in Neural Information Processing Systems , 35:40854098, 2022.

Mengdi Xu, Yuchen Lu, Yikang Shen, Shun Zhang, Ding Zhao, and Chuang Gan. Hyperdecision transformer for efficient online policy adaptation. arXiv preprint arXiv:2304.08487 , 2023.

Mengdi Xu, Yikang Shen, Shun Zhang, Yuchen Lu, Ding Zhao, Joshua Tenenbaum, and Chuang Gan. Prompting decision transformer for few-shot policy generalization. In international conference on machine learning , pages 24631-24645. PMLR, 2022.

Taku Yamagata, Ahmed Khalil, and Raul Santos-Rodriguez. Q-learning decision transformer: Leveraging dynamic programming for conditional sequence modelling in offline rl. arXiv preprint arXiv:2209.03993 , 2022.

Denis Yarats, Rob Fergus, Alessandro Lazaric, and Lerrel Pinto. Mastering visual continuous control: Improved data-augmented reinforcement learning. arXiv preprint arXiv:2107.09645 , 2021.

A Additional ablation study

A.1 Step Size Ablation

In Algorithm 1, we introduced the step size parameter δ , acting as a balance between search granularity and inference speed. The effects of varying δ are depicted in Figure 8. To compute the history length search space, we commence with the maximum history length. For instance, if the maximum history length is T = 20 and the step size δ = 8 , the history length search space becomes 20 , 12 , 4 .

Figure 8 shows that a narrowed search space leads to a decline in return performance and an increase in standard deviation, corroborating results from Table 3. EDT with a restricted history search space behaves more like EDT with a fixed history length. We found that the EDT works best when we chose a search step of 4 .

2

4

6

8

0

20

40

60

80

100

Return

Medium-Replay (Hopper)

This could be because it is tough to make good estimates without being able to interact with the environment directly, as mentioned in previous studies [26, 28]. By increasing the step size, we were able to make our estimates more reliable and less affected by noise and other problems.

A.2 History Length Distribution

We further show more history length distributions for the locomotion tasks here.















































































































































































































































































































































Our initial hypothesis regarding the Ant task was that the medium-replay dataset would primarily consist of shorter history lengths, while the medium dataset would be more focused on longer history lengths. The distribution of history lengths in the medium and medium-replay datasets mostly supports this hypothesis.

In the medium-replay dataset, we consistently see a concentration of shorter history lengths. This was expected, given that the nature of the medium-replay dataset is likely to produce shorter history lengths. The distribution often converges towards higher densities for shorter lengths, which aligns with our expectations.

The pattern within the medium dataset, however, is less consistent. This could be attributed to several factors that can either elongate or truncate the history length. Despite these fluctuations, we still

observe a slight inclination towards longer history lengths. This meets our initial assumption but also demonstrates the complexity within the medium dataset.

A.3 Comparison with Diffusion-Based Methods

B Experiment Details

B.1 Loss Computation Details

In this section, we detail the computation of the tokenized return loss L return and the observation loss L observation in Eq 1. For L return, we first tokenize the predicted return and the target return into n bins bounded by the MAX and MIN scores obtained from D4RL [15]. We then compute the cross-entropy loss between the predicted value and the target value:

where R ′ t suggests the predicted tokenized return. For L observation, we predict the next observation by minimizing the mean square error.

B.2 Learning on Locomotion Tasks

In Section 4, we discussed the application of our approach to a multi-task learning scenario. Specifically, we consolidated medium-replay datasets from Ant, HalfCheetah, Hopper, and Walker2d. To standardize returns across these varied environments, we used scaling statistics (maximum return, minimum return, and return scale) from the official Decision Transformer repository ( https://github.com/kzl/decision-transformer ). As detailed in Section 3.3, we further segmented the return into 60 discrete buckets and, following the approach of [20, 30], sampled from the top 85th percentile logits within the discrete return distribution.

However, it's important to highlight that our return maximizer, ˜ R , estimates the scaled value directly rather than the discretized one. For the sequence ⟨ ..., o , R, a, ... ⟩ , we supplement each token embedding with a learned positional embedding. Given the differing state and action dimensions across tasks, we employ an additional state and action embedding layer. This layer transforms the raw state and action representation into a consistent state and action embedding size, while keeping the remainder of the architecture unchanged across tasks.

Approximate ˜ R Estimation. Section 3.4 outlines how we use Bayes' Rule to estimate expert returns. A conventional ˜ R prediction typically involves an autoregressive process, given the necessity of sampling from the discrete ˆ R expert distribution at each time step during the search for the optimal history length. To simplify this process, we follow the approximation strategy used in [30]'s implementation. We mask all return values except the first one, thus making ˜ R solely dependent on o , a , and the initial return value.

History Length Search Heuristic. In Algorithm 1, we illustrated that a larger δ value allows us to infer actions more rapidly. Adopting this method needs to balance inference speed with search accuracy. Based on the concept introduced in Sec. 3.2, the optimal history length at the current state s t +1 might be close to that of the previous state s t . Therefore, given the optimal length l t at time step t , we search within the range { l t -∆ , l t -∆+1 ..., l t +∆ } for the optimal length l t +1 at the next step.

B.3 Multi-Task Learning on Atari Games

The process for action inference in multi-task learning on Atari games closely aligns with that described in Sec. B.2, including the method for approximating ˜ R . However, there are several noteworthy distinctions, which we elaborate upon in this section.

Given that all games utilize grayscale frames of the same dimensions (84x84), we do not need to implement an additional state embedding layer as we did in the Locomotion scenario. Instead, we introduce a shared image encoder for all games, with further details outlined in Sec. B.5.

It's important to note the distinction between the action spaces of Atari games and Locomotion tasks. The action space of Atari games is discrete, so our sampling approach mirrors how we sample from the return distribution: we select from the top 85th percentile of action logits. Inspired by [30], we

discretize the reward signal into {-1 , 0 , +1 } , while the return is split into 120 buckets ranging from {-20 , ..., 100 } .

Given the complexity and potential pitfalls of learning on 20 Atari games from scratch, we use a subset of GPT-2 to initialize the transformer decoder. This step improves both the convergence rate and the quality of the learned policy. Our dataset consists of 2 training runs from [2], with each run featuring rollouts from 50 checkpoints. Lastly, we enhance the dataset with image augmentations, including random cropping and rotation.

B.4 Atari Games

Due to time and computational constraints, we randomly selected 20 tasks from the 41 tasks in the study by [30]. Details about the game types and number of action spaces can be found in Table 4. We also provide the raw scores of the Atari experiments in Table 5. Given that we transform the original observations to grayscale and rescale them to 84x84 images, we illustrate these transformed observations in Figure 11.

B.5 Image Encoder

To ensure a fair comparison, we standardize the image encoder across EDT and the two baseline approaches. We adopt the image encoder from DrQ-v2 [60] for this purpose.

The architecture of the image encoder is as followed:

1 convolution with stride 2, output channels 32, kernel size 3. (ReLU).

3 convolution with stride 1, output channels 32, kernel size 3. (ReLU).

1 fully connected layer and H output dimensions.

Assault

Asterix

BankHeist

BeamRider

Breakout

Centipede

ChopperCommand

Enduro

Freeway

Frostbite

Gravitar

NameThisGame

Phoenix

Seaquest

TimePilot

VideoPinball

WizardOfWor

Zaxxon

口口口

5

15

0-0

国

*电

OTET

IHHNHI

g00000

C Inter-Quartile Mean of the Atari results

Alongside the main paper's results, we present the InterQuartile Mean (IQM) of our findings. Notably, IQL outperforms DT in terms of IQM. This occurs due to the variance in task difficulties, despite our use of HNS to balance scales across tasks. The difficulty disparity among games leads to significant variance. However, despite this, our proposed EDT still significantly surpasses both baselines, demonstrating its robust performance.

In the original Decision Transformer model, estimating the remaining return-to-go is dependent on a preset expected return-to-go value. This requirement presents a challenge when applied to multi-task environments, where return-to-go values can significantly vary across different tasks.

The issue is particularly pronounced in Atari games where high scores often necessitate long trajectories.

These long trajectories may consequently yield negative estimates for the return-to-go, which presents an issue for the model's performance and application.



















Thus, the unique characteristics of different tasks highlight the limitations of a one-size-fits-all approach in the context of return-to-go estimation in the Decision Transformer model.

D Architecture Details

We adopt the causal transformer decoder architecture and process the sequences as follows.

Transformer Blocks: Composed of masked causal attention and a multilayer perceptron (MLP), with layer normalization and residual connections. Activation function used is GELU (Gaussian Error Linear Unit).

· Embedding Layers:

-State Embedding: A linear layer followed by positional (time) embeddings addition.

-Action Embedding: A linear layer followed by positional (time) embeddings addition.

-Return to Go Embedding: A linear layer followed by positional (time) embeddings addition.

-Timestep Embedding: An embedding layer for encoding timesteps.

· Prediction Heads:

-State Prediction: A linear layer taking as input the concatenated action and state embeddings.

-Action Prediction: A linear layer followed by a Tanh activation function.

-Return-to-go Prediction: A linear layer.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper tackles the \"trajectory stitching\" problem of Decision Transformer. It proposes a method called Elastic Decision Transformer (EDT) that dynamically adjusts the history at inference time such that longer history is maintained for optimal trajectories and shorter history is maintained for sub-optimal ones. Empirical evidence shows that EDT can better achieve trajectory stitching than original DT. Furthermore, it is able to outperform methods based on Q-Learning on certain D4RL tasks and Atari games.\n\nSTRENGTHS:\n- This paper is well presented and easy to follow.\n\n- The idea to dynamically adjust history to achieve trajectory stitching is novel and well motivated.\n\n- Experiment results show that the proposed method seems robust enough to benefit from optimal and sub-optimal trajectories.\n\n- Code is submitted. Though I didn't run them.\n\nWEAKNESSES:\n- The applicability of the proposed method seems to be limited since it is only applicable to DT and its variants.\n\n- Chosen baselines might not cover all recent progresses in offline RL, especially those specifically designed to address trajectory stitching.\n\n- Comparison to recent methods based on Diffusion model, including Diffuser[Janner et al., 2022] and Decision Diffuser[Ajay et al., 2023], is missed.\n\n- Though the paper claims EDT to be generally applicable for all DT variants, it is only evaluated against the original DT.\n\n- Typo, L234 should be Table 2 instead of Figure 2.\n\nReferences:\n- Janner et al., Planning with Diffusion for Flexible Behavior Synthesis, ICML 2022.\n- Ajay et al., Is Conditional Generative Modeling all you need for Decision Making?, ICLR 2023.\n\nQUESTIONS:\n- How is the method compared to diffusion-based planning methods for offline RL, such as Diffuser[Janner et al., 2022] and Decision Diffuser[Ajay et al., 2023]?\n\n- Can we demonstrate how the EDT idea could improve the performance of other DT variants?\n\nReferences:\n- Janner et al., Planning with Diffusion for Flexible Behavior Synthesis, ICML 2022.\n- Ajay et al., Is Conditional Generative Modeling all you need for Decision Making?, ICLR 2023.\n\nLIMITATIONS:\nSee above.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes the Elastic Decision Transformer (EDT), an improvement over the existing Decision Transformer (DT). DT has had notable limitations with trajectory stitching, i.e., generating optimal or near-optimal trajectories from sub-optimal ones. EDT addresses this issue by modifying the history length utilized in DT, optimizing the trajectory by maintaining a longer history for optimal trajectories and a shorter one for sub-optimal ones. This adjustment allows the model to \"stitch\" more optimal trajectories.\n\nSTRENGTHS:\n** Strengths**\n\n1. The concept of Elastic Decision Transformer (EDT) offers a novel approach to tackle the issue of trajectory stitching.\n2. The paper successfully translates the abstract theoretical concepts into practical, tested methods. The methodology of using varying history lengths based on trajectory quality is both innovative and intuitive.\n3. The empirical results demonstrating EDT's superior performance in various benchmarks provide robust support for the proposed model.\n\nWEAKNESSES:\n** Weaknesses **\n1. While the paper presents compelling evidence of EDT's superiority over the original DT, it would further enrich the discussion to include comparative analyses with other state-of-the-art DT variants. Examples of such variants include Prompt DT [1], Hyper DT [2], and Multi-Game DT [3].\n\n2. A potential concern arises from the fact that the EDT's performance hinges heavily on the value maximizer. While effective in function, it tends to be less efficient in terms of training speed. Exploring and discussing ways to enhance this aspect could be beneficial.\n\n3. The paper could benefit from a more elaborate explanation of the value maximizer's mechanics, particularly the implementation of expectile regression. Understanding this component better could aid readers in appreciating the model's inner workings.\n\n[1] Xu, Mengdi et al. “Prompting Decision Transformer for Few-Shot Policy Generalization.” ArXiv abs/2206.13499 (2022): n. pag.\n[2] Xu, Mengdi et al. “Hyper-Decision Transformer for Efficient Online Policy Adaptation.” ArXiv abs/2304.08487 (2023): n. pag.\n[3] Lee, Kuang-Huei et al. “Multi-Game Decision Transformers.” ArXiv abs/2205.15241 (2022): n. pag.\n\nQUESTIONS:\n1. I recommend the authors include further results for comparison with the state-of-the-art algorithms I've referenced in the weaknesses section, namely Prompt DT, Hyper DT, and Multi-Game DT. Such comparative analysis would be instrumental in providing a clearer understanding of the proposed method's efficacy.\n\n2. Upon careful scrutiny of the experimental results, I noted a lack of evaluations concerning the training speed associated with different step sizes denoted by \"$\\delta$\". It would be worthwhile to explore and discuss how this parameter influences both the learning time and the overall performance of the model.\n\n3. How does the EDT handle scenarios where the distinction between optimal and sub-optimal trajectories is not clear?\n\n4. How would EDT perform with extremely long trajectories where the computation cost may become prohibitive?\n\n5. Could you expand on how the computation overhead during training is minimized?\n\nLIMITATIONS:\nThe major limitation of this paper is value maximizer training speed, as illustrated in the paper.\nI don't find any obvious broader societal impacts.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes Elastic Decision Transformer (EDT), which helps Decision Transformer (DT) generalize offline data by trajectory stitching. Their changes to the DT are: 1) during training, adding an extra term to the objective to estimate the maximum Q values starting from a state (similar to the Implicit Q-Larning algorithm in the literature); 2) during inference, it determines the optimal length of trajectory to put in the prompt to most-likely generate a trajectory with a high return. EDT is evaluated with other state-of-the-art offline RL algorithms (including both classical and DT-based ones) on D4RL and Atari domains.\n\nSTRENGTHS:\nThis work addresses a well-known problem in DT, the ability to generalize offline data by stitching different trajectories. There are earlier works that try to address the same problem (Q-Learning Decition Transformer, for example), but the earlier works do not achieve as good performance as EDT.\n\nThis work brings the performance of DT-based offline RL methods to match the best algorithms in the literature (IQL). The authors did a thorough evaluation of EDT and baseline algorithms, on D4RL and Atari datasets, which are common benchmarks for offline RL.\n\nThe modifications to DT are simple, well-justified, and effective. The paper is overall clearly written and easy to follow.\n\nWEAKNESSES:\n**Contributions.** Does this simply implement the idea behind the IQL algorithm in DT? If this is the case, it should be emphasized somewhere in the paper. I have some questions about this in the question section below.\n\n**Clarity.** Need background on expectile regression. It may be helpful to include a figure like Fig. 1 (left) in the IQL paper (Kostrikov et al., 2021) to make this paper more self-competent.\n\n**Minor points.**\n\nLine 94: DT uses rewards-to-go, or return, not rewards.\n\nLine 205: Extra quote mark at the end of the line.\n\nStructure of Sec. 3: Line 115: \\tilde{R} should be clearly defined here, it’s a “maximal value estimator.” It’s defined on a trajectory segment.\n\nIt would be great to change the color of Fig. 3. The node labels ($S^a_{t-1}$ and $S^b_{t-1}$) are a bit hard to read.\n\n**Missing references**\n\nThe paper considered the multi-task regime, but some related works are missing:\n\nReed, Scott, et al. \"A generalist agent.\" arXiv preprint arXiv:2205.06175 (2022).\n\nXu, Mengdi, et al. \"Prompting decision transformer for few-shot policy generalization.\" international conference on machine learning, 2022.\n\nQUESTIONS:\nMy understanding is that training always uses the same history length ($T$). Does it help to use different history lengths during training as well? It seems that training and inference use histories of different lengths (always length of $T$ for training, length of $\\leq T$ for inference). Would it harm the performance?\n\nSec. 4.4 shows that different history lengths do affect the performance. I wonder if the optimal history length search can be avoided without harming the performance: Given the estimate of the maximum value starting from the state, can the model learn to attend to the right length of history? Or can we make optimal history length search as part of the training? For example, we can find the maximum length during the training and let the model only attend to states within the lengths?\n\nAbout the empirical results and the contribution of the paper: EDT significantly improves the performance of DT, which is similar to the performance of IQL. Can the authors elaborate the implications of this? Should we simply use IQL since it already has the best performance? What are the reasons for using EDT? (more computationally efficient, simpler to implement?)\n\nLIMITATIONS:\nAs the authors have mentioned in the paper, EDT adds a computational overhead in the inference process of DT by choosing the history length (Alg. 1).", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces the Elastic Decision Transformer (EDT), an extension of the Decision Transformer (DT) model that automatically determines the history length input during inference. Specifically, the EDT estimates the value of each candidate history length with the expectile return regression and chooses the length that yields the maximum value. They validate the effectiveness of EDT on RL benchmarks, demonstrating that its performance surpasses that of the DTs and other offline RL algorithms, particularly in the context of multi-task scenarios.\n\nSTRENGTHS:\n* The concept of adaptively adjusting the history length input to the DT model is intriguing.\n* The introduced solution addresses the trajectory stitching problem in the original DT model.\n* The authors utilize expectile return regression to select the value of each candidate history length\n* The authors demonstrate its utility in RL benchmark environments.\n\nWEAKNESSES:\n* The design and implementation of the Elastic Decision Transformer (EDT) model appear ad-hoc, with insufficient discussions on the situations where the model is effective or ineffective.\n    * The selection process for history length, leaning towards an optimistic approach by possibly ignoring potentially inconvenient histories, calls into question the model's adaptability in non-Markovian or stochastic environments where the original DT might be suitable.\n    * Considering the inherent attention mechanism of the Transformer architecture, it remains unclear whether the proposed method of adjusting the input history length is beneficial universally, particularly in cases of sufficient data availability.\n* The presentation of the methods lacks clarity, with missing definitions for variables and loss functions. The necessity to refer to external sources ([29]) for key details detracts from the paper's self-contained nature.\n    * The inclusion of an algorithm table to detail the entire process would have significantly aided reader understanding.\n    * There is an apparent lack of effort to make the paper self-contained.\n\nQUESTIONS:\n* What history length is used during training? Is it a fixed length, or does the EDT model learn \\tilde R for each potential length?\n* Why isn't there a comparison with Multi-game DT[29], especially for multi-task experiments?\n* Could the authors clarify line 162? In particular, what the authors mean by \"we optimize the action space as the return objective L return using cross-entropy.\" Although I understand that the authors are estimating the return distribution with a categorical return distribution, there is no explicit mention of this, which can lead to confusion.\n* The calculation process for P(R^t) in Equation (4) is unclear. Could the authors provide more details on this?\n* What is \"expert return distribution\"? There seems to be no mention of the \"expert.\" Moreover, the definitions of the returns R, \\tilde R, and \\hat R are needed.\n* The definition of the loss is unclear, as it requires a reference to [29]. Also, the additional losses L_return and L_observation, which were not used in the original DT, have been added to the loss, so an explanation is necessary.\n* Should \"Figure 2\" on line 234 be \"Table 2\"?\n\nLIMITATIONS:\n* The paper lacks a comprehensive exploration of the specific situations where the model excels or falls short. It would be valuable to see a deeper investigation into the behavior of the model across various scenarios, especially regarding the impact of history-length selection in different types of environments. In other words, if the Elastic Decision Transformer (EDT) essentially narrows down the scope of application of the original Decision Transformer (DT) by taking advantage of Markov properties, it is unclear what unique benefits it retains when compared with other comparable models like IQL or CQL.\n* The history length selection process appears to be overly optimistic, possibly overlooking histories that may be inconvenient to the EDT model. This approach could limit the model's adaptability in stochastic non-Markovian environments. To address this, the authors could conduct experiments in such environments to validate the model's robustness.\n* This paper isn't entirely self-contained; understanding crucial aspects of the algorithm, particularly those related to returns and losses, necessitates referencing external resources such as [29]. This undermines the paper's self-contained nature and can hinder the reader's understanding.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-10"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes to vary the history length provided to a Decision Transformer to enable it to stitch suboptimal trajectories into better ones. The motivation behind this, the authors argue, is that the model can shift to an optimal trajectory best given only the current state (and not the history), while providing the history can better help the model maintain an optimal trajectory once found. Concretely, this method EDT trains a value function to estimate the highest achievable return given a history length, and searches for the best history length for action inference. In experiments on the D4RL benchmark and Atari games, EDT is shown to outperform DT, which they show struggles with trajectory stitching.\n\nSTRENGTHS:\n- The proposed method seems to work better than DT and is competitive with offline RL methods.\n- The experimental results are quite thorough with multiple domains and relevant comparisons.\n- EDT also seems especially strong in the multi-task setting, where Q-learning-based offline RL methods typically struggle.\n- The solution is intuitive: it essentially varies the receptive field at inference time, because the model would otherwise just follow the training trajectories given the full history.\n\nWEAKNESSES:\n- The main limitation of the method seems to be inference cost, as it requires a search over all possible history lengths at each time-step. While the search space can be reduced by defining a coarser range of possible history lengths, this can also compromise the performance of the method.\n- It would be helpful to study a separate setting which _only_ requires stitching of offline trajectories (maybe a gridworld example?) to see how EDT compares to Q-learning-based offline RL more directly.\n- In the introduction, the messaging could be a bit clearer by mentioning that this observation applies to DT / sequence modeling approaches.\n\nQUESTIONS:\n- Are the plots in Fig. 7 illustrating the best history length according to the value estimator?\n- If so, I’m curious what the performance of DT would look like if we set the history length to the mode of the distributions shown in Fig. 7 (i.e., either 3 or 4). Are the results in Table 3 illustrating this but for history length 5 and 20? Or, are there other differences between DT and EDT?\n- How much data is needed to train the value estimator? What happens if you train the estimator on a subset of the DT data or on a superset of the DT data? This would help get a sense for whether the data requirement of the estimator is greater than or smaller than the DT.\n- What does \"-ONE\" indicate?\n\nNits:\n- Could error bars be provided for the Atari results?\n- Line 212: casual -> causal\n\n\n================= POST-REBUTTAL RESPONSE ===================\n\nThank you for the detailed response. I still am in favor of acceptance.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-19"}