Elastic Decision Transformer
Authors: Yueh-Hua Wu
Xiaolong Wang
Masashi Hamaya
Venue: NeurIPS 2023 poster (2023)
Keywords: Offline Reinforcement Learning
Trajectory Stitching
Decision Transformer
TL;DR: We propose the Elastic Decision Transformer, an advanced offline Reinforcement Learning algorithm that effectively achieves trajectory stitching and optimizes performance through dynamic adjustment of history length.

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
Yueh-Hua Wu 12
Xiaolong Wang
Masashi Hamaya?*
UC San Diego *OMRON SINIC X
yuw088@ucsd. edu
Abstract
This paper introduces Elastic Decision Transformer (EDT), significant advance- ment over the existing Decision Transformer (DT) and its variants, Although DT purports to generate an optimal trajectory. empirical evidence suggests strug- gles with trajectory stitching, a process involving the generation of an optimal or near-optimal trajectory from the best parts of a set 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 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 D'T-hased and Q Learning- based approaches. In particular, the EDT outperforms Q Learning- based methods a multi-task regime on the D4RL locomotion benchmark and Atari games. Videos are available at: https: htpps://krisee /ttiistrrte...r gi thub. io/edt/.
1 Introduction
Reinforcement Learning (RL) trains agents to interact with an environment and learn from re- wards. It has demonstrated impressive results across diverse applications such as game play- ing [23, 44], robotics [41, 40, 54], and recom- mendation systems [12, 1]. A notable area of RL. Offline RL [31], which employs pre-collected data for agent training and proves more efficient when real-time interactions are costly lim- ited. Recently, the conditional policy approach has shown large potentials in Offline RI,, where the agent learns policy based on the observed state and goal. This approach enhances perfor- mance and circumvents stability issues related long-term credit assignment. Moreover, the successful Transformer architecture [49], widely used applications like natural language processing [52, 13, 8] and computer vision [14, 34], has been adapted for RL as the Decision Transformer (DT) [11].
Hopper
Walker2d
89
74
61
6I
50
30
Figure 1: Normalized return with medium- replay datasets. The dotted gray lines indicate nor- malized return with medium datasets. By achiev- ing trajectory stitching, our method benefits from worse trajectories and learns better policy.
DT utilizes Transformer architecture model and reproduce sequences from demonstrations, integrating a goal-conditioned policy to convert Offline RI. into a supervised learning task. Despite its competitive performance in Offline RI 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
"cqual advising
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

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 more advantageous future path, poses a challenge for sequence generation- better trajectory appears contradict one of the core objectives of sequence generation that sequence generation model is required to reliably reproduce trajectories found within the training dataset. We suggest that in order 'refresh' the prediction model, it should disregard 'negative' or "unsuccessful' past experiences. This involves dismissing past failures and instead considering 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 based on as much information as possible. However, our proposed approach aligns with this concept. With shorter history, the prediction model tends to output with higher variance, typically considered a weakness 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, relationship emerges between the quality of the path taken and the length of history used for prediction. This correlation serves the motivation behind our proposal to employ a variable length of historical data input.
approaches in offline RL. Stitching
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 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-oftthe-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 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, 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 traditional DT and other Offline RL algorithms.
Our experimental evaluation highlights EDT's s superior performance in multi-task learning regime, positioning a promising approach for future Offline Reinforcement Learning research and applications.
2 Preliminaries
In this study, we consider decision-making agent that operates within the framework of Markov Decision Processes (MDPs) [42]. At every time step f, the agent receives an observation of the world or, chooses an action at, and receives a scalar reward rt- Our goal is to learn a single optimal policy distribution Pt (at)oStt "ost, rrt) with parameters 0 that maximizes the agent's total future return Ekse nd on all the environments we consider.
R.
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 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 RI. 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

---

D4RL
O
t+1
R
R
causal transformer decoder
Atari
R
a
observation
return
action
Figure 2: An overview of our Elastic Decision Transformer architecture. R is the prediction of the maximum return. We also show the environments we used in the experiments on the right. We adopt four tasks for D4RL [15] and 20 tasks for Atari games.
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 the fixed dataset. A number of recent research efforts have addressed these challenges, including methods for importance weighting [35, 37J, regularization [53, 28], and model-based learning [25, 29], among others.
2.2 Decision Transformer
The Decision Transformer architecture, introduced by [1 [] 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 sequence of past states, actions, and rewards, and the output the next action to be taken. DT uses 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, has a limitation in its ability perform "stitching." Stitching refers to the ability to combine parts of sub-optimal trajectories to produce an optimal trajectory. This approach can lead to situation where the agent follows sub-optimal trajectory that provides an immediate reward, even if a different trajectory leads to higher cumulative reward over time. This limitation of DT is a significant challenge 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), 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 better trajectory by forgetting 'unsuccessful' past experiences, thus opening up more possibilities for future trajectories. We further propose 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 sequence modeling problem. Specifically, we model the probability of the next token in the sequence (denoted as r) based on all the tokens that come before The sequences we model can represented as:
('.., ft', aa....).
where is time step and ft is the return for the remaining sequence. The sequence consider here similar to the one used in [30] whereas we do not include reward as part the sequence and we predict an additional quantity ît that enables us to estimate an optimal input length, which we will cover in the following paragraphs. Figure 2 presents 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 usually directly predicted from a' through the causal transformer decoder.
3.2 Motivation
We propose shift in the traditional approach to trajec- tory stitching. Instead of focusing on training phases, aim to achieve this stitching during the action in- ference stage. This concept is illustrated in Figure 3 using simplified example. In this scenario, con- sider a dataset, D. comprising only two trajectories: D (59_1, St: f6+1), (sf_1, 844)) A sequence model trained with this dataset is likely to predict the next states in manner consistent with their original tra- jectories.
Value
S
S
S
S
b S
Timesteps A Toy example to il- lustrate the motivation of EDT. The figure shows an offiin RL dataset that contains only two trajectories (sh-1, $(11) (si_1, St, s(-1).
To overcome this, we propose method that enables trajectory stitching, where the model starts from si-1 and concludes sh+1. This is achieved by adaptively adjusting the history length. We introduce a maximal value estimator, R, which calculates the maximum value among all potential outcomes within the dataset. This allows to determine the optimal history length that maximizes R.
Figure 3:
In the given example, the model starts at state s it will choose to retain the history (s:) upon reaching state st, as R(s;)) > R(s=-1, Conversely, if the model initiates from state sy it will preserve the history 1. s2) when decision-making at 8t, R( St) > R(s;). From the above example, we understand that the optimal history length depends on the quality of the current trajectory we've ve traversed, and it be specific length anywhere between preset maximal length and single unit.
To estimate the optimal history length in general scenario, we propose solving the following optimization problem:
arg max max (TT)), TrED
T
where Ty denotes the history length T. More precisely, Ty takes the form:
TI = (o' 1 R
a- --++1 d-1, R a'-1 6-1,00 0 R aa).
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 given history length in EDT. To approximate the maximum operator in max,ye R'(T)), 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 E (0, 1) expectile of a random variable X is the solution to an asymmetric least squares problem, as follows:
arg min ExEx [L -ma)), m.
where (u) |a= 1u < 0)
Through expectile regression, we can approximate max (T(r)): i max it' (TT) Rs arg min Ery [L] (ft' (rr) - DD
R((rT)
We estimate RR by applying an empirical loss of Equation with a sufficiently large a (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:
LEDT C, Cretum + Lotservation + Laction + Lmax:

---

R
causal transformer decoder
max
R
R
w.e
R
Bayes' Rule
1
max
Ry(F2)
R
causal transformer decoder
Be(TT)
max
Optimal length w
Figure 4: The figure illustrates the action inference procedure within the proposed Elastic Decision Transformer. Initially, we estimate the value maximizer, R,, for each length within the search space, as delineated by the green rectangle. Subsequently, we identify the maximal value from all R. which provides the optimal history length w. Utilizing this optimal history length, we estimate the expert value at time step denoted as R,,er by Bayes' Rule. Finally, the action prediction is accomplished via the causal transformer decoder, which is indicated by the blue rectangle. In practice, we retain the distribution of R during the estimation process for I2 and present the inference here for clarity.
Algorithm EDT optimal history length search
1: Input: query sequence (.... 0' R a o²) and EDT model EDDT
2: for w = 1,,1 6, + 26,.., do
3:
Obtain Jt' (T")) with truncated sequence T" (ot-wt1, w+1 pt-=+11 a w+- o') and θEDT
4: end for
5: Compute P())) with OEDT for that has the highest it and then estimate P( lexpert' with Equation 4.
where Lobservation and Laction are computed with a mean square error, Lream is a cross-entropy loss, and Lmax an empirical estimate of Equation 2. We set G, 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 the tokenized return objective Zreum using cross-entropy with weight 10c,. 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 significant advance in developing efficient offline reinforcement learning approaches for complex decision vecisinn-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, 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. These steps are elaborated in Figure 4.
To identify the history length i that corresponds to the highest RC we employ a search strategy as detailed in Algorithm 1. As exhaustively searching through all possible lengths from to T may result in slow action inference, we introduce step size ō to accelerate the process. This step size not only enhances inference speed by factor of 8, but also empirically improves the quality of the learned policy. An ablation study on the impact of the step size is provided Appendix A. For all experiments, set to eliminate the need for parameter tuning.
To sample from expert return distribution (..'xer/exppertt) we adopt an approach similar to [30] by applying Bayes' rule .e./eer)reerr) x P(expent' ...) ....) and approximate the distribu-
5

---

Dataset
DI
QDT
TS+BC
S4RL
IQL
EDT (Ours)
hopper-medium
60.7 14.
5.6
64.3 4.2
78.9
63.8 + 9.1
63.5
nopper-medium-rellyy
61.9 61.9133.7
45.8
30.2
35.4
92.1 10.1 10.4
89.0 8.3
malker-medium
3.9
2.0
7..
93.6
3.0
72.8 6.2
walker- medium-replay
43.3 14.
36.2
6..
30.3
73.6
6.3
74.8 4.9
hawfcheetahh-meeeee
42.5 1
42.3
43.2
0.3
48.8
47.3 - 0.2
42.5 0.9
halfcheetah- malfcheetab-medium medium- replay
34.9 = 1.6
11.1
39.8 0.6
51.4
44.11111
37.8 + 1.5
average
52.5
45.5
56.3
56.4
66.7
63.4
ant-medium
92.5 + 5.1
99.9
+ 5.8
97.9 + 8.0
unt-medium-replay
87.9 + 4.9
91.2 + 7.2
92.0 + 4.1
average
90.2
95.5
94.9
Table 1: Baseline comparisons on D4RL [15] tasks. Mean of 5 random training initialization seeds, 100 evaluations each. Our result 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 (2 0.95 max).
tion of expert-level return with inverse temperature [24, 48, 47, 43]:
P(R' expeet', ....) cxp(xR' P(R').
While it may initially appear feasible to directly use the predicted as the expert return, it's important note that this remains conservative maximum operation. Empirically, we have found that Eq. encourages the pursuit of higher returns, which consequently enhances the quality of the actions taken.
4 Experiments
Our experiments are designed 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 Locomotion and Atari games? (Sec. 4.3)
Does dynamic history length approach surpass 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 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.
that QDT and TS+BC are specifically designed achieve trajectory stitching. QDT accomplishes this by substituting collected return values with estimates derived from Conser- vative Q-Learning [28]. Conversely, TS+BC employs model-based data augmentation strategy to bring about the stitching of trajectories."
It important
4.2 Single-Task Offline Reinforcement Learning
For locomotion tasks, train offline RL models D4RL's "medium" and "medium-replas" datasets. The "medium" dataset comes from 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 required during test time.
current sequence, reward information
Our observations indicate that the proposed EDT consistently outperforms the baseline DT and its variants the majority of the datasets, with notable performance advantage on the "medium-

---

Task
DT-1
IQL-1
EDT- I (Ours)
100%
hopper
51.2
59.8
76.9
93.0%
walker
29.8
52.6
74.1
halfcheetah
30.5
40.4
36.8
Mean
67.1%
50%
ant
79.8
82.3
88.6
48.9%
sum
191.3
235.1
276,4
Table 2: Evaluation results on multi-task regime. Mean of random training initial- ization seeds, 100 evaluations each. The train- ing dataset is mixture of mepium-replay datasets from the four locomotion tasks. Our main result highlighted.
EDT-ONE
DT-ONE
(QL-ONE
Figure 5: The average HNS comparison on 20 Atari games. The results are evaluated with three trials.
Dataset
EDT (w=20)
(W=5)
(Ours)
100
hopper-nt
5.2
hopper-mr
67.6 +: 27.7
62.6 4 26.9
89.0
8.3
walker-m
65.6 11.7
62.6
72.8 + 6.2
walker
64.5 12.9
44.3 1.3::..
74.8
4.9
halfcheetah-m
42.0
Return
89.0%
ant m
hulfchectah-mr
90.8 +: 16.0
33.5
36.4 74
95.6 +: 8.0
97.9
37.8
+
8.0
1.5
80.1%
69.5%
73.2%
ant-mr
80.0 + 17.0
82.3
15.9
92.0
4.1
sum
507.6
$70.3
Table 3: Mcan of random training initial- ization seeds, 100 evaluations each. In the Dataset column, "m" indicates medium, "mr" indicates medium-replay dataset. The "w" stands for history length. Our main results are highlighted.
0.99
0.9
0.7
0.5
Figure 6: Ablation study expectile level a. Expectile objective reduces to Mean Square Error when a 0.5. Evaluated on Hopper and medium-replay dataset.
replay" datasets. These findings provide strong evidence that our approach is highly effective in stitching together sub-optimal trajectories with high return proportion, task that DT and its variants cannot accomplish. Although EDT doesn't 't fully outperform IQL in the single-task learning, it does bridge the gap between Q-learning-basee methods and DT by performing trajectory stitching with the estimated maximum return.
4.3 Multi-Task Offfine 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, 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 our proposed EDT successfully accomplishes all four tasks simultaneously without much performance compromise.
Atari. For Atari, we adopt CNN image encoder used in DrQ-v2 [59] to process stacks of four 84x84 image observations. To ensure fair comparisons, all methods employ the same architecture for the image encoder. Following [30], 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 (score score / (score scorezandom). to ensure a consistent scale across each game.
Our experimental results align with those of [30], highlighting that Q-learning-basse 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 and Figure 5. We leave the raw scores of the 20 Atari games in Appendix B.

---

Walker2d (medium-replay)
Walker2d (medium)
Epoch
0.14
Epoch
100
0.12
100
Density
Density
0.10
400
0.08
400
0.04
0.02
0.02
15
20
History length
History length
Figure 7: The figures illustrate the history length distributions across datasets and training epochs. For each distribution, we collect ten trajectories and derive histogram of history lengths. The distribution is computed with kernel distribution estimate for better visualization.
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 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 shorter history length aids in trajectory stitching, it s also erucial retain 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 choose the optimal history length. Consequently, examining the change in performance relative to the expectile level, a, 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 higher expectile levels, observe both higher average performance and lower standard deviation. This suggests that accurate selection of history length not only stabilizes performance but also enhances scores. Conversely, the expectile level approaches 0.5, the expectile regression's 's objective shifts towards a mean square error, leading to an estimated value that more of mean value than a maximal one. This change makes 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 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 distribution closely approximating uniform distribution. Conversely, the medium dataset, acquired through a singular stochastic policy, exhibits history length distribution characterized by an increased density at lower history lengths. This observation can attributed to the prevalence of analogous trajectories within the medium dataset, leading to more frequent occurrences 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-repla dataset.
Contrary to initial expectations, trajectory stitching does not oceur as frequently within the medium- replay dataset 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 limited number of mutual st
8

---

instances illustrated in Figure 3. Nevertheless, the proposed EDT method derives substantial benefits from trajectory stitching 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.
Related Work
Offline Reinforcement Learning. Offline RL has been promising topics for researchers since sampling from environments during training is usually costly and dangerous real-world applications and offline reinforcement learning is able to learn better policy without directly collecting state- action 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 plamning algorithm and demonstrates exceptional flexibility across various applica- tions, such as long-houg hong-horizon dynamics prediction, imitation learning, conlicoond reinforcement learning, and offline reinforcement learning.
Recently, there has been 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 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-bar -aaood generative models to represent policies 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 truth return-to-go with estimated values, technique expected to foster trajectory recombina- tion. 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 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.
Discussion
Conclusion. In this paper, we introduced the Elastic Decision Transformer, 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 mulli-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 RI. forward.
In summary, EDT offers a promising solution for trajectory stitching, enabling the creation of better sequences from sub-optimal trajectories. This capability considerably enhance DT variants, leading to improved performance across diverse applications. We are committed to releasing our code.
Limitations. potential direction for future improvement involves enhancing the speed which EDT estimates the optimal history. This could make the method suitable for real-time applications that have strict time constraints. While this adaptation an exciting avenue for future research, falls outside the primary scope of this paper.

---

Acknowledgement
This project is supported by a grant from JST ACT-X, grant number JPMIAX22AC.
References
[1] M Mehdi Afsar, Trafford Crump, and Behrouz Far. Reinforcement learning based recommender systems: A survey. ACM Computing Surveys, 55(7):1-38, 2022.
[2] Rishabh Agarwal, Dale Schuurmans, and Mohammad Norouzi. An optimistic perspective on offline reinforcement learning. In International Conference on Machine Leamiin, pages 104-114. PMLR, 2020.
[3] Dennis Aigner, Takeshi Amemiya, and Dale Poirier. On the estimation of production frontiers: maximum likelihood estimation the parameters of a discontinuous density function. International economic review, pages 377-396, 1976.
[4] Anurag Ajay, Yilun Du, Abhi Gupta, Joshua Tenenbaum, Tommi Jaakkola, and Pulkit Agrawal. Is conditional generative modeling all you need for mocisionnmmainn? arXiv preprint arXiv:2211.15657 2022.
[5] Arthur Argenson and Gabriel Dulac-Amold. Model-based offline planning. arXiv preprint arXiv:2008.05556, 2020.
[6] Adrià Puigdomènech Badia, Bilal Piot, Steven Kapturowski, Pablo Sprechmann, Alex Vitvit- skyi, Zhaohan Daniel Guo, and Charles Blundell. Agent57: Outperforming the atari human benchmark. In International conference on machine learning, pages 507-517. PMLR, 2020.
[7] Marc G Bellemare, Yavar Naddaf, Joel Veness, and Michael Bowling. The arcade learning environment: evaluation platform for general agents. Journal of Artificial Intelligence Research, 47:253-277 279, 2013.
[8] 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.
[9] Ian Char, Viraj Mehta, Adam Villaflor, John M Dolan, and Jeff Schneider. Bats: Best action trajectory stitching. arXiv preprint arXiv:2204. 0222204..2200 /0026, 2022.
[10] Huayu Chen, Cheng Lu, Chengyang Ying, Hang Su, and Jun Zhu. Offline reinforcement learning via high-fidelity generative behavior modeling. arXiv preprint arXiv:2209.14544 14548, 2022.
[11] 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, 14:10044-15097, 2021.
[12] Xinshi Chen, Shuang Li, Hui Li, Shaohua Jiang, Yuan Qi, and Le Song. Generative adver- sarial user model for reinforcement learning based recommendation system. In International Conference on Machine Learning, pages 1062-106. PMLR, 2019.
[13] Jacob Devlin, Ming-Wai Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint 1rXiv:/810.0000 04005, 2018.
[14] 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 16x 16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.1 1929, 2020.
[15] 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.
10

---

[16] 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.
[17] Caglar Gulcehre, Ziyu Wang. Alexander Novikov, Thomas Paine, Sergio Gómez, Konrad Zolna, Rishabh Agarwal, Josh S Merel, Daniel Mankowitz, Cosmin Paduraru, et al. RI unplugged: A suite of benchmarks for offline reinforcement learning. Advances in Neural Information Processing Systems, 33:7248-7253 2020.
[18] Philippe Hansen-EEstuuch 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.
[19] Charles A Hepburn and Giovanni Montana. Model-based trajectory stitching for improved offline reinforcement learning, arXiv preprint arXiv:2211.11605 2022.
[20] Ari Holtzman, Jan Buys, Li Du, Maxwell Forbes, and Yejin Choi. The curious case of neural text degeneration. arXiv preprint arXiv: (90:::9944000 09751, 2019.
[21] Michael Janner, Yilun Du, Joshua B Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. arXiv preprint (.Xii:2205.09999) 2022.
[22] Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as one big sequence modeling problem. Advances in neural information processing systems, 34:1273- 1286, 2021.
[23] Lukasz Kaiser, Mohammad Babaeizadeh, Piotr Milos, Blazej Osinski, Roy Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, et al. Model- based reinforcement learning for atari. arXiv preprint arXiv:1903.00074, 2019.
[24] Hilbert J Kappen, Vicenç Gómez, and Manfred Opper. Optimal control as a graphical model inference problem. Machine learning, 87:159-182, 2012.
[25) Rahul Kidambi, Aravind Rajeswaran, Praneeth Netrapalli, and Thorsten Joachims. Morel: Model-based offline reinforcement learning. Advances in neural information processing systems, 33:21810-22 123,, 2020.
[26] Ilya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit q-learning arXiv preprint arXiv:2110.06699, 2021.
[27] Aviral Kumar, Justin Fu, Matthew Soh, George Tucker, and Sergey Levine. Stabilizing off- policy q-learning via bootstrapping error reduction. Advances in Neural Information Processing Systems, 32, 2019.
[28] Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative gelarnning for offline reinforcement learning. Advances in Neural Information Processing Systems, 33:1179- 179- 1191, 2020.
[29] Byung-Jun Lee, Jongmin Lee, and Kee-Fung Kim. Representation balancing offline model- hased reinforcement learning. In International Conference on Learning Representations, 2021.
[30] Kuang-Huei Lee, Ofir Nachum, Mengjiao Yang, Lisa Lee, Daniel Freeman, Winnie Xu, Ser- gio Guadarrama, lan Fischer, Eric Jang, Henryk Michalewski, et al. Multi-game decision transformers. arXiv preprint arXiv:2205. 1..05..5200 15241, 2022.
[31] Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint arXiv:2005.01645 2020.
[32] Zhixuan Liang, Yao Mu, Mingyu Ding, Fei Ni, Masayoshi Tomizuka. and Ping Luo. Adaptdif- fuser: Diffusion models as adaptive self-evolving planners. arXiv preprint arXiv:2302.01877, 2023.
[33] 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: 2021.

---

[34] 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.
[35] Eric Mitchell, Rafael Rafailov, Xue Bin Peng, Sergey Levine, and Chelsea Finn. Offline meta- reinforcement learning with advantage weighting. In International Conference on Machine Learning, pages 7780-7791. PMLR. 2021.
[36] Volodymyr Mnih, Koray Kavukeuoglu, 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-333 2015,
[37] Ashvin Nair, Abhishek Gupta, Murtaza Dalal, and Sergey Levine. Awac: Accelerating online reinforcement learning with offline datasets. arXiv preprint arXiv:2006.09355, 2020.
[38] Whitney K Newey and James L Powell. Asymmetric least squares estimation and testing. Econometrica: Journal of the Econometric Society, pages 819-847, 1987.
[39] Xue Bin Peng, Aviral Kumar, Grace Zhang, and Sergey Levine. sreaantgg-weighhee regression: Simple and scalable off-policy reinforcement learning. arXiv preprint arXiv: 1910001177, 2019.
[40] 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.09444, 2018.
Applications on robotics. Journal of Intelligent & Robotic Systems, 86(2):153-173, 2017. [41] Athanasios S Polydoros and Lazaros Nalpantidis. Survey of model-based reinforcement learning:
[42] Martin Puterman. Markov decision processes. Handbooks in operations research and management science, 2:331-434, 1990.
[43] Ross D) Shachter. Probabilistic inference and influence diagrams. Operations research, 36(4):589-604, 1988.
[44] David Silver, Thomas Hubert. Julian Schrittwieser, Ioannis Antonoglou, Matthew Lai, Arthur Guez, Mare Lanctot, Laurent Sifre, Dharshan Kumaran, Thore Graepel, et al. general reinforcement learning algorithm that masters chess, shogi, and go through self-play. Science, 962(6419):1140-11
2018
[45] Samarth Sinha, Ajay Mandlekar, and Animesh Garg. S4rl: Surprisingly simple sulf-supervision for offline reinforcement learning in robotics. In Conference Robot Learning, pages 907-917. PMLR, 2022.
[46] Farzan Soleymani and Eric Paquet. Financial portfolio optimization with online deep rein- forcement learning and restricted stacked autoencoder-ddeppbreen Expert Systems with Applications, 156: 13456. 2020.
[47] Emanuel Todorov. Linearly-solvable markov decision problems. Advances in neural information processing systems, 19, 2006.
[48] Mare Toussaint. Robot trajectory optimization using approximate inference. In Proceedings of the 26th annual international conference machine learning, pages 1049-1056, 2009.
[49] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in newral information processing systems, 30, 2017.
[50] Zhendong Wang, Jonathan J Hunt, and Mingyuan Zhou Diffusson policies an expressive policy class for offline reinforcement learning. arXiv preprint arXivv2208.061193 2022.
[51] 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.
12

---

[52] Thomas Wolf. Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue. Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, al. Transformers: State- of-the-art natural language processing. In Proceedings of the 2020 conference on empirical methods in natural language processing: system demonstrations, pages 38 45, 2020.
[53] Yifan Wu, George Tucker, and Ofir Nachum. Behavior regularized offline reinforcement learning. arXiv preprint
2019.
[54] Yueh-Hua Wu, Jiashun Wang, and Xiaolong Wang. Learning generalizable dexterous manipula- tion from human grasp affordance. In Conference on Robot Learning, pages 618-629. PMLR, 2023.
[55] Teng Xiao and Donglin Wang. general offline reinforcement learning framework for in- teractive recommendation. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pages 4512 4512-4520, 2021.
[56] 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:4085- 4098, 2022.
[57] Mengdi Xu, Yuchen Lu, Yikang Shen, Shun Zhang, Ding Zhao, and Chuang Gan. Hyper- decision transformer for efficient online policy adaptation. arXiv preprint arXiv:2304.08887, 2023.
[58] 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.
[59] Taku Yamagata, Ahmed Khalil, and Raul Santos-Rodriguez. Q-learning decision transformer: Leveraging dynamic programming for conditional sequence modelling in offline arXiv preprint arXiv:2209.03999. 2022.
[60] Denis Yarats, Rob Fergus. Alessandro Lazaric, and Lerrel Pinto.
Mastering visual continuous control: Improved dugmentd reinforcement learning. arXiv preprint arXiv:2107.06644 2021.
13

---

A Additional ablation study
A.1 Step Size Ablation
In Algorithm 1, we introduced the step size pa- rameter 8, acting granularity and inference speed. The effects of varying 8 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 7 = 20 and the step size d 8, the history length search space becomes 20, 12,
balance between search
Medium-Replay (Hopper)
80
Retur
60
40
20
Figure 8 shows that narrowed search space leads to a decline in return performance and an increase in standard deviation, corroborating results from Table 3. EDT with restricted his- tory 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 This could be because it is tough to make good estimates without being able to interact with the environment directly, mentioned in previous studies [26, 28]. By increasing the step size, were able to make our estimates more reliable and less affected by noise and other problems.
Step Size
Figure 8: Ablation study the step size ō. The greater the step size the smaller the length search space is.
A.2
History Length Distribution
We further show more history length distributions for the locomotion tasks here.
WalkerZd (medium)
Hopper (mediium)
Ant (medium)
Halfcheeiah (mediuum)
0.14
1114
0.12
Donsit
8
0.2
641
608
10 28
0.03
15
0.0
15
600
70
15
Hisdory longth
Window longth
History longth
Histtry longth
Figure 9: The figures illustrate the history length distributions across datasets and training epochs on the medium dataset. For each distribution, we collect ten trajectories and derive histogram of history lengths. The distribution is computed with kernel distribution estimate for better visualization.
WalkerZd (medium-roplay)
Hopper (medium-replay)
Ant (medium-replay)
Halfcheetah mmeuuuu-rpppaa/ ay)
6.12
Town
3.175
190
0.10
0115
6:06
Purch
20155
6.15
0075
8.825
0005
0.08
15
0.000
15
0008
13
24
C.00
Hildtoyy longth
Histrry length
Hastboy longth
Hilttoy longth
Figure 10: The figures illustrate the history length distributions across datasets and training epochs on the medium-replay dataset. For each distribution, we collect trajectories and derive histogram of history lengths. The distribution is computed with kernel distribution estimate for better visualization.
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 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 concentration of shorter history lengths. This expected, given that the nature of the medium-repll dataset 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, less consistent. This could be attributed several factors that either clongate truncate the history length. Despite these fluctuations, still
14

---

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, detail the computation of the tokenized return loss Cram and the observation loss Lobservation in Eq 1. For Lreum. 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:
Lretmm CE((f, R²",
(5)
where Rd suggests the predicted tokenized return. For Lobservation: we predict the next observation by minimizing the mean square error.
B.2 Learning Locomotion Tasks
In Section 4, we discussed the application of our approach to a multi-task learning scenario. Specif- ically, 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:/ htt//gitttt github. commkzl/ kkl//eciiioonn httt://aiiiiioottasfffooooeeooe segmented the return into 60 discrete buckets and, following the approach of [20, 30J, sampled from the top 85th percentile logits within the discrete return distribution.
As detailed in Section 3.3, we further
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, ...). supplement each token embed- ding with a learned positional embedding. Given the differing state and action dimensions across tasks, employ 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 ? Estimation. Section 3.4 outlines how we use Bayes' Rule estimate expert returns. A conventional R prediction typically involves an autoregressive process, given the necessity of sam- pling from the discrete Rexpert 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 solely dependent on o, a, and the initial return value.
History Length Search Heuristic. In Algorithm 1, we illustrated that 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 $141 might be close to that of the previous state st. Therefore, given the optimal length 4, at time step i, we search within the range (4e A,4 A 1..,, + A} for the optimal length 4441 at the next step.
B.3 Multi-Task Learning 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 additional state embedding layer as we did in the Locomotion scenario. Instead, 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 discrete, so our sampling approach mirrors how we sample from the return distribution: select from the top 85th percentile of action logits. Inspired by [30],
15

---

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 subset of GPT-2 initialize the transformer decoder. This step improves both the convergence rate and the quality of the learned policy. Our dataset consists of training runs from [2], with each run featuring rollouts from 50 checkpoints. Lastly, enhance the dataset with image augmentations, including random cropping and rotation.
16

---

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 Given that we transform the original observations to grayscale and rescale them to 84x84 images, we illustrate these transformed observations in Figure 11.
Table 4: Atari Games: Name, Category, and Action Space
Game
Category
Action Space
Assault
Shooter
Asterix
Platform
18
BankHeist
Strategy
18
BeamRider
Shooter
Breakout
Arcade
Centipede
Shooter
18
ChopperCommand
Shooter
18
Enduro
Racing
FishingDerby
Sports
18
Freeway
Racing
Frostbite
Platform
Gravitar
Shooter
18
NameThisGame
Shooter
Phoenix
Shooter
Qbert
Puzzle
Seaquest
Shooter
TimePilot
Shooter
10
VideoPinball
Pinball
WizardofWor
Shooter
Zaxxon
Shooter
18
Table 5: Raw Atari scores for the three offline RL approaches. The scores average results
three trials. The Human Avg and Random scores are adopted from [6].
Games
Human Avg
Random
EDT-ONE (80M)
DT-ONE (80M)
IQL-IQM (80M)
Asterix
8503.3
210
12089.6
7731.
994.6
Assault
742.0
222.4
1849.2
1260.0
964.3
BunkHeist
753.1
14.2
5.0
9.3
21.3
BcamRider
16926.5
363.9
6469.9
4864.6
3865.3
Breakout
30.5
1.7
220.1
180.8
70.9
Centipede
12017.0
2090.9
2164.0
2290.4
2940.4
ChopperCommand
7387.8
811
3491.6
2774.8
331.7
Enduro
860.5
0
823.6
557.1
877.8
FishingDerby
-38.7
-91.7
-1.7
-43.3
30.6
Freeway
29.6
0
26.0
13.5
37.8
Frostbite
4334.7
1837.6
1558.1
488.5
Gravitar
3351.4
173
70.5
119.8
-37.1
NameThisGame
8049.0
2292.3
7529.4
5939.7
3954.3
Phoenix
7242.0
761.4
4280.5
3369.6
1893.1
Qbert
13455.0
163.9
11739.8
7978.4
11035.5
Seaquest
42054.0
68.4
4410.9
2628.3
756.7
TimePilot
$229.2
3568.0
2737.7
3036.7
2311.3
VideoPinball
17667.9
0
875.7
476.0
WizardOrWor
4756.5
563.5
222.2
384.8
498.2
Zaxxon
9173.3
32.5
233.4
141.7
-6.5
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 the image encoder is as followed:
I convolution with stride 2, output channels 32, kernel size 3. (ReLU).
17

---

3 convolution with stride 1, output channels 32, kernel size 3. (ReLU).
fully connected layer and output dimensions.
Asteri
Bankéleist
Beam
Breakout
WITH 01
Centipede
ChenperCammend
FichinTeetty
Trosth
Genstrrr
NameThhiGGaee
Phoesix
Ooert
ReuteeR
Sesquest
TimePia at
VideoPinball
WizersWNWor
7ascon
Figure All Atari games used for our experiments. We adopted gray scale and reshape observations to images of size 84x84.
C Inter-Quartile Mean of the Atari results
Alongside the main paper's results, we present the Inter- Quartile Mean (IQM) of our findings. Notably, IQL outperforms DT in terms of IQM. This occurs due to of HNS
the variance in task difficulties, despite to balance scales across tasks. The difficulty disparity among games leads to significant variance. However, de- spite this, our proposed EDT still significantly surpasses both baselines, demonstrating its robust performance.
50%
IQM
44.7%
In the original Decision Transformer model, estimating the remaining return-to-go is dependent pected return-to-go value. This requirement presents challenge when applied to multi-lask environments, where return-to-go values can significantly vary across different tasks.
preset ex-
0%
EDT-ONE
DT-ONE
IQL-ONE
Figure 12: Results of learning from 20 Atari games in terms of IQM.
The issue particularly pronounced Atari games where high scores often necessitate long trajectories. These long trajectories may consequently yield negative estimates for the return-to-go, which presents issue for the model's performance and application.
18

---

Thus, the unique characteristics of different tasks highlight the limitations of a one- cee--izz-fit-lll approach in the context of returm-to-g to-go estimation the Decision Transformer model.
D Architecture Details
We adopt the causal transformer decoder architecture and process the sequences follows.
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).
Transformer Blocks:
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 embed- dings.
Action Prediction: A linear layer followed by Tanh activation function.
Retuun-to-go Prediction: A linear layer.
Table 6: plper-parameters. We list all hyper-parameters for D4RL and Atari games below.
Parameter
Setting (D4RL)
Setting (Atari)
Observation down-sampling
84 84
Frames stacked
4
Frames skip
4
Terminal on loss of life
True
Augmentation (random cropping)
True
Augmentation (random rotation)
True
Discount factor
1
0.9974
Gradient clipping
True
True
Reward clipping
True
True
Optimizer
AdamW
Adam W
Optimizer (learning rate)
0.0001
0.0001
Optimizer (weight decay)
0.0001
0.0001
Maximum history length
20
30
inverse temperature (x)
10
10
Expectile level
0.99
0.99
Mix precision
True
True
Evaluation episodes
100
32
Action loss coefficient (Caction)
I
0.01
Return loss coefficient (Cretu))
0.001
0.001
State loss coefficient (Lstate)
1
Return maximizer loss coefficient (Cmxx)
0.5
0.5
Number of return output bins
60
120
Observation normalization
True
True
Batch size
256
256
Sample from the top 85th percentile logits (return)
True
True
Sample from the top 85th percentile logits (action)
True
Step size (8)
2
2
GPU
3090
V100
19


Reviews:
[{"review_id": "qT3gPsjFCP", "reviewer": "Reviewer_kCB8", "summary": "This 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.", "questions": "- 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.", "limitations": "Yes", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- 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.", "weaknesses": "- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their insightful comments. We address your comments in the following.\n\n---\n\n**Q**: “The main limitation of the method seems to be inference cost”\n\n**A**: Thank you for pointing this out. We have conducted an ablation study to compare the wall clock time with the proposed $\\delta$ during the history length search. The following results suggest that although EDT indeed requires additional time for searching the length, the difference is quite small when compared with sampling from environments like gym. The searching approach can be potentially improved by learning a predictor to accelerate the inference speed. \n\n|               $\\delta$               |    2   |    4   |    6   |    8   |   DT   |\n|:------------------------------------:|:------:|:------:|:------:|:------:|:------:|\n| wall clock time for 1000 steps (sec) | $12.6$ | $12.1$ | $11.8$ | $11.5$ | $11.4$ |\n\n---\n\n**Q**: “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\n**A**: Directly visualizing the trajectory results using a simpler task, such as gridworld, and comparing the results with a Q-learning-based approach is indeed an interesting idea. In a simple gridworld scenario, the Q-learning-based method might choose to go for the grids with the highest estimated return. On the other hand, vanilla DT may directly follow the paths presented in demonstrations. EDT, however, might show a pattern similar to the Q-learning-based approach. An extreme case of this is when EDT adopts a unit history length throughout action inference. This can result in identical results to the Q-learning-based approach, assuming the learned Q-values and returns are closely matched.\nHowever, due to limited time, we were unable to provide a comprehensive visualization at this stage. We plan to work on visualizing the learned policy in this direction for a future version and will include these insights on our project website in the future. \n\n---\n\n\n**Q**: “In the introduction, the messaging could be a bit clearer by mentioning that this observation applies to DT / sequence modeling approaches.”\n\n**A**: Thank you for your suggestions. We will improve the clarity of our manuscript in the introduction section.\n\n---\n\n\n**Q**: Are the plots in Fig. 7 illustrating the best history length according to the value estimator?\n\n**A**: Yes. In Figure 7, we plotted the distribution of the best history length in trajectories.\n\n---\n\n**Q**: 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\n**A**: Table 3 illustrates EDT's performance when the history length is fixed to 5 and 20. Since we observed that some datasets prefer a shorter 'fixed history length' while others favor a longer one, we provide a thorough comparison by presenting the performance for history lengths of 5 and 20. Additionally, in the table below, we show EDT's performance when the history length is set to 3. These results suggest that although a shorter history length may be preferred based on the outcome of history length search, it can also be critical to use a longer history length when necessary. This demonstrates EDT's ability to utilize history length appropriately.\n\n|                   |    hopper-m   |    hopper-mr   |    walker-m   |    walker-mr   | halfcheetah-m | halfcheetah-mr |\n|:-----------------:|:-------------:|:--------------:|:-------------:|:--------------:|:-------------:|:--------------:|\n| EDT (fixed $T=3$) | $60.5\\pm 4.6$ | $57.8\\pm 22.7$ | $68.3\\pm 8.1$ | $57.0\\pm 14.6$ | $42.1\\pm 0.4$ |  $16.6\\pm 4.2$ |\n\nIn the table, ‘m’ suggests ‘medium’ and ‘mr’ suggests ‘medium-replay’.\n\nThe differences between DT and EDT can be summarized as follows:\nDuring training, we additionally optimize the expectile regression loss in EDT. This allows us to estimate the maximum return a history length can possibly achieve.\nDuring action inference, EDT utilizes the estimated values from the approximate return maximizer, obtained through expectile regression, to identify the optimal history length.\n\n\n---\n\n**Q**: 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\n**A**: We conducted an ablation study on the size of demonstrations, where the ratio indicates the proportion of demonstrations used for training. For this experiment, we used the 'hopper-medium-replay' demonstration, and the results are summarized in the table below. The results show that EDT still outperforms the DT baseline with a sufficient amount of demonstrations. However, when the demonstration ratio falls below 0.4, both methods fail to learn meaningful policies. An interesting observation is that although EDT outperforms the baseline DT when the demonstration ratio is greater than or equal to 0.6, the difference becomes less significant when the ratio is less than 0.8. This indicates that having enough demonstrations is one of the key requirements to benefit from trajectory stitching in EDT.\n\n| ratio |   1.0  |   0.8  |   0.6  |  0.4  |  0.2  |\n|:-----:|:------:|:------:|:------:|:-----:|:-----:|\n|  EDT  | $89.0$ | $65.2$ | $19.6$ | $5.3$ | $5.4$ |\n|   DT  | $61.9$ | $59.7$ | $18.2$ | $5.1$ | $5.5$ |\n\n---\n\n**Q**: What does \"-ONE\" indicate?\n\n**A**: The “-ONE” denotes that we employ a single model to learn all tasks in a multi-task scenario, distinguishing it from the setting used in the single-task experiments.\n\n---\n\nPlease do not hesitate to let us know if you have any additional comments.\n"}, {"review_id": "1paPe4dVDA", "reviewer": "Reviewer_GPqH", "summary": "This 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.", "questions": "* 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\"?", "limitations": "* 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. ", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 1, "contribution": 2, "strengths": "* 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.", "weaknesses": "* 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their insightful comments.\n\n---\n\n**Q**: The design of the EDT appears ad-hoc.\n\n**A**: The components proposed in EDT such as using expectile regression to estimate the maximum return and optimal history length search, are not particularly designed for specific cases. Instead, the implementation of EDT are rather simple and straightforward since we would like to purely demonstrate the idea–namely, using a varying history length can be helpful for trajectory stitching. It's important to note that the core contribution of this work is to provide a novel method for trajectory stitching. This is a key factor in why DT-based approaches might fall behind the SOTA in single-task scenarios.\n\n---\n\n**Q**: 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\n**A**: EDT doesn't necessarily resolve the trajectory stitching issue in non-Markovian environments since we consider the Markov Decision Process in this work. We've demonstrated improvements across both locomotion and Atari tasks, in both single-task and multi-task scenarios. This work opens up the possibility for DT-based approaches to achieve comparable performance in single-task scenarios and to significantly outperform other methods in multi-game scenarios (18% improvement in locomotion tasks, 101% improvement in Atari tasks).\n\n---\n    \n**Q**: Considering the 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\n**A**: We utilized the same number of demonstrations as those presented in D4RL, an offline RL benchmark for model evaluations. In our experiments, we demonstrated the superiority of EDT over the DT baseline in both single-task and multi-task scenarios. The improvements were significant: 27% in single-task D4RL, 44% in multi-task D4RL, and 40% in multi-task Atari . The attention mechanism of the Transformer architecture enables DT to predict actions robustly in accordance with training trajectories. However, this may be insufficient for DT to choose a better action that may differ from the decision presented in the demonstrations.\n\n---\n\n**Q**: What history length is used for training?\n\n**A**: During the training process, the EDT model learns the values of $\\tilde{R}$ for each potential length. We use the same data loader as other Decision Transformer models, and we sample sub-trajectories with a context length $T$. For example, if $T=3$, the sampled input will be $[o_{t-2}, R_{t-2}, a_{t-2}, o_{t-1}, R_{t-1}, a_{t-1}, o_{t}, R_{t}, a_{t}]$, and the output will be to predict $[(R_{t-2}, \\tilde R_{t-2}), a_{t-2}, o_{t-1}, (R_{t-1}, \\tilde R_{t-1}), a_{t-1}, o_{t}, (R_{t}, \\tilde R_{t}), a_{t}, o_{t+1}]$. In this context, the model learns to predict both $R_{t-2}$ and $\\tilde R_{t-2}$ based solely on $o_{t-2}$. Due to the inherent nature of Decision Transformer training, the EDT naturally learns the value of $\\tilde R$ for each potential length.\n\n---\n\n**Q**: Didn't compare with Multi-Game DT (MGDT).\n\n**A**: The training script for MGDT has not been released, and EDT without varying history length essentially serves as an approximation to MGDT. The main difference is that their model doesn't estimate the next state and $\\tilde{R}$. Consequently, in our experiments, we compare the proposed EDT to an architecture without varying length, serving as a re-implementation of MGDT. To further illustrate this, we provide the performance of EDT without varying history length, using it as an approximation for the performance of MGDT in our multi-task experiment.\n|               | Hopper | Walker | Halfcheetah |   Ant  | Atari (Mean %) |\n|:-------------:|:------:|:------:|:-----------:|:------:|:------------:|\n| EDT (fixed T) | $56.5$ | $33.7$ |    $31.0$   | $82.5$ |    $75.5$   |\n\n---\n\n**Q**: Explain \"we optimize the action space as the return objective L return using cross-entropy.\"\n\n**A**: Thank you for pointing this out. We will revise our manuscript to clarify that we use tokenized values for the return.\n\n---\n\n**Q**: The calculation process for $P(R^t)$ is unclear.\n\n**A**: To derive $P(R^t|\\text{expert}^t)$, we first apply Bayes’ Rule: $P(R^t|\\text{expert}^t) \\propto P(\\text{expert}^t|R^t)P(R^t)$. Then, we define a binary classifier that classifies whether an instance is from expert demonstrations. This classification is done with a probability proportional to the future return, using an inverse temperature $\\kappa$. Similar approaches have been studied extensively in previous works [1,2,3]. As a result, we approximate $P(\\text{expert}^t|R^t)$ with $\\exp(\\kappa R^t)$, which gives $P(R^t|\\text{expert}^t) \\propto \\exp(\\kappa R^t)P(R^t)$.\n\n---\n\n**Q**: What is \"expert return distribution\"? Define the returns $R$, $\\tilde R$, and $\\hat R$.\n\n**A**: The term 'expert' here suggests 'optimal.' Regarding the return definition, we use $\\hat R$ for normal return, $\\tilde R$ for the maximum return estimator, and $R$ otherwise. We will make changes to the manuscript to enhance clarity.\n\n---\n\n**Q**: The additional losses not used in the DT, have been added to the loss.\n\n**A**: In the paper, we mentioned that the return and observation losses are computed with cross-entropy loss and mean-square error, respectively. We also illustrate the architecture that EDT uses to predict return and the next observation in Figure 2. To improve clarity, we will modify the manuscript to highlight that we employ tokenized return and detail the loss computation. \n\n---\n\n[1] Lee et al., Multi-Game Decision Transformers\n\n[2] Kappen et al., Optimal control as a graphical model inference problem\n\n[3] Toussaint et al., Robot trajectory optimization using approximate inference\n"}, {"review_id": "gwFhzUBbmS", "reviewer": "Reviewer_ak2g", "summary": "This 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.", "questions": "My 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", "limitations": "As 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).", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "This 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", "weaknesses": "**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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their insightful comments. We address your comments in the following.\n\n\n---\n\n**Q**: “Does this simply implement the idea behind the IQL algorithm in DT?”\n\n**A**: Although both IQL and EDT employ expectile regression, they use it for different purposes. In IQL, expectile regression is utilized to implicitly estimate the Q-Learning objective, specifically serving as a Q-value maximizer with respect to actions. In contrast, EDT employs expectile regression to estimate return concerning previous history, composed of observations, previous returns, and actions. This enables us to choose the optimal history length. While the loss itself may be similar, the computation and the ways we apply them are quite distinct.\n\n---\n\n**Q**: “Need background on expectile regression.”.\n\n**A**: Thank you for your suggestions. Due to limited space, we were unable to cover more information and details regarding expectile regression. We will provide more detailed background knowledge about expectile regression in the Appendix of our revised manuscript.\n\n---\n\n**Q**: “Minor points about clarity and missing references”.\n\n**A**: Thank you for pointing this out. We will revise the manuscript according to your suggestions.\n\n---\n\n**Q**: \"My 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\n**A**: During EDT training, the $\\tilde{R}$ is trained with respect to different history lengths and this can be achieved with a standard data loader and DT training. To accomplish this, we sample sub-trajectories with a context length $T$. For example, let’s say $T=3$, and the sampled input will be $[o_{t-2}, R_{t-2}, a_{t-2}, o_{t-}, R_{t-1}, a_{t-1}, o_{t}, R_{t}, a_{t}]$, with the task being to predict $[(R_{t-2}, \\tilde R_{t-2}), a_{t-2}, o_{t-1}, (R_{t-1}, \\tilde R_{t-1}), a_{t-1}, o_{t}, (R_{t}, \\tilde R_{t}), a_{t}, o_{t+1}]$. In this case, the model is trained to predict $R_{t-2}$ and $\\tilde R_{t-2}$ using only $o_{t-2}$, which corresponds to a unit history length. Due to the inherent structure of DT training, EDT naturally learns the value of $\\tilde{R}$ for each potential length.\t\n\n---\n\n**Q**: \"Sec. 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\n**A**: Finding the optimal history length without conducting a history length search is indeed an intriguing research question. In this work, we introduce a straightforward search method to estimate the optimal length. While this approach requires slightly more inference time (a difference that is almost negligible, as shown in the table below), it effectively illustrates the idea that varying the history length enables DT to achieve trajectory stitching and derive better performance. As a potential enhancement, one could directly predict the optimal length by training a classifier on an offline RL dataset, labeled with optimal lengths as predicted by the trained EDT model. This method would allow the classifier to estimate the optimal length with an inference time roughly equivalent to that of the standard DT, simplifying the process.\n\n|               $\\delta$               |    2   |    4   |    6   |    8   |   DT   |\n|:------------------------------------:|:------:|:------:|:------:|:------:|:------:|\n| wall clock time for 1000 steps (sec) | $12.6$ | $12.1$ | $11.8$ | $11.5$ | $11.4$ |\n\n---\n\n**Q**: \"About 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\n**A**: The experimental results show that EDT greatly improves the performance of DT, indicating that trajectory stitching is one of the major issues that hinder DT's competitiveness with other SOTA methods in single-task scenarios. Indeed, IQL is a simple yet effective Q-learning approach, but our experiments demonstrate that EDT yields superior results in multi-task learning. The Multi-Game Decision Transformers [1] also reveals that it significantly outperforms Q-learning-based approaches, such as CQL, in multi-game scenarios without directly addressing trajectory stitching. In conclusion, our experiments suggest that EDT is a more favorable option than IQL for multi-task learning.\n\n---\n\nPlease do not hesitate to let us know if you have any additional comments.\n\n[1] Lee et al., Multi-Game Decision Transformers (2022), https://arxiv.org/abs/2205.15241\n"}, {"review_id": "oXjkF1Mvru", "reviewer": "Reviewer_C3CW", "summary": "This 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. ", "questions": "1. 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?", "limitations": "The major limitation of this paper is value maximizer training speed, as illustrated in the paper.\nI don't find any obvious broader societal impacts.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "** 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.", "weaknesses": "** 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their insightful comments. We address your comments in the following.\n\n---\n\n**Q**: 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.\n**A**: Although EDT employs an additional expectile regression loss compared to DT, the training time is not significantly affected. This is because the expectile regression loss involves only an additional one-layer MLP, and the computation is similar to that of MSE. Below, we compare the wall clock time for training both DT and EDT to provide a more comprehensive result. (3090 GPU, 64 batch size)\n\n|                                            |   DT   |   EDT  |\n|:------------------------------------------:|:------:|:------:|\n| wall clock training time (500 epochs, sec) | $1442$ | $1449$ |\n\n\n---\n\n**Q**: “The paper could benefit from a more elaborate explanation of the value maximizer's mechanics, particularly the implementation of expectile regression.”\n\n**A**: Thank you for the constructive advice. During EDT training, the $\\tilde{R}$ is trained with respect to different history lengths and this can be achieved with a standard data loader and DT training. To accomplish this, we sample sub-trajectories with a context length $T$. For example, let’s say $T=3$, and the sampled input will be $[o_{t-2}, R_{t-2}, a_{t-2}, o_{t-}, R_{t-1}, a_{t-1}, o_{t}, R_{t}, a_{t}]$, with the task being to predict $[(R_{t-2}, \\tilde R_{t-2}), a_{t-2}, o_{t-1}, (R_{t-1}, \\tilde R_{t-1}), a_{t-1}, o_{t}, (R_{t}, \\tilde R_{t}), a_{t}, o_{t+1}]$. In this case, the model is trained to predict $R_{t-2}$ and $\\tilde R_{t-2}$ using only $o_{t-2}$, which corresponds to a unit history length. Due to the inherent structure of DT training, EDT naturally learns the value of $\\tilde{R}$ for each potential length. The implementation of expectile regression for $\\tilde R$ is fundamentally the same as the implementation of MSE loss for $R$, with the exception that we use expectile regression (as defined in Eq. 2). This can be computed with just a few lines of Python code, as shown below:\n```\ndef expectile_loss(diff, expectile=0.99):\n   weight = torch.where(diff > 0, expectile, (1 - expectile))\n   return weight * (diff**2)\n```\nWe will add more details regarding the implementation of expectile regression in our revised manuscript.\n\n---\n\n**Q**: “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\n**A**: Thank you for your suggestions! The studies on Prompt DT and Hyper DT explore potential directions to enable a DT model to adapt to novel tasks. On the other hand, Multi-Game DT, which is similar to EDT, is focused on large-scale multi-task learning with DT. We will include Prompt DT and Hyper DT in our related work section to provide a more comprehensive survey.\n\n---\n\n**Q**: “It would be worthwhile to explore and discuss how parameter $\\delta$ influences both the learning time and the overall performance of the model.”\n\n**A**: Thank you for pointing this out. The parameter $\\delta$ does not influence the learning time as action inference with varying history length is not required during training. This, however, does not imply that the training length is fixed. Instead, the value $\\tilde{R}$ is learned for each history length. We incorporate the parameter $\\delta$ only in the evaluation phase, where the history length search is performed to find the optimal length. In the Table below, we summarize the ablation study on the value of $\\delta$ and the corresponding inference time on a 3090 GPU.\n\n|               $\\delta$               |    2   |    4   |    6   |    8   |   DT   |\n|:------------------------------------:|:------:|:------:|:------:|:------:|:------:|\n| wall clock time for 1000 steps (sec) | $12.6$ | $12.1$ | $11.8$ | $11.5$ | $11.4$ |\n\n---\n\n**Q**: How does the EDT handle scenarios where the distinction between optimal and sub-optimal trajectories is not clear?\n\n**A**: It is actually easier for EDT to combine optimal parts of trajectories when the distinction between optimal and sub-optimal trajectories is not clear. This is because, in such a case, EDT can identify more common or similar parts between these trajectories and then switch to the one with a higher estimated return.\n\n---\n\n**Q**: How would EDT perform with extremely long trajectories where the computation cost may become prohibitive?\n\n**A**: In our implementation, instead of keeping all the previous history, we maintain a maximum history length $T$ to store previous trajectories. In our experiments, we use $T=20$ for D4RL tasks and $T=30$ for Atari experiments. We can alleviate the computation cost by either using a shorter $T$ or a greater $\\delta$ as proposed in the paper, which will reduce the number of queries for each action inference.\n\n---\n\n**Q**: Could you explain how the computation overhead during training is minimized?\n\n**A**: The additional computational overhead in training EDT stems solely from the expectile regression used for maximum return estimation. This loss calculation is efficiently carried out through the expectile regression loss (Eq. 2) using an additional one-layer MLP, as illustrated in Figure 2. As the EDT’s training only involves this additional expectile regression loss, its training time is almost identical to that of the vanilla DT. Therefore, the complexity does not significantly impact the efficiency of our method, enabling us to minimize the computational overhead.\n\n---\n\nPlease do not hesitate to let us know if you have any additional comments.\n"}, {"review_id": "G4DkDAjqMu", "reviewer": "Reviewer_nm2U", "summary": "This 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.", "questions": "- 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.", "limitations": "See above.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "- 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.", "weaknesses": "- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their insightful comments. We address your comments in the following.\n\n---\n\n**Q**: \"The applicability of the proposed method seems to be limited since it is only applicable to DT and its variants.\"\n\n**A**: In this work, we primarily focus on improving DT and its variants. However, the motivation and intuition behind our approach are not confined solely to DT but extend to sequence generation models as well. For models that predict future actions or statistics and struggle with trajectory stitching, the intuition proposed in this work—namely, using a shorter history length for more accurate prediction results—could potentially be applicable.\n\n---\n\n**Q**: “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**A**: \nWe present the performance comparison on D4RL in the Table below, including the reported values from Decision Diffuser [1] to ensure a fair comparison. While diffusion-based models sometimes show comparable or even superior performance in single-task scenarios, it's important to note that our EDT model demonstrates scalability across multi-game scenarios with a large amount of mixed demonstrations—an area where diffusion models are unexplored. Even a DT-based method like Multi-Game Decision Transformers [2], which does not directly address the trajectory stitching issue, substantially outperforms Q-learning-based approaches such as CQL.\n\nIn our multi-task learning experiments, EDT maintains a clear advantage, achieving both superior scalability and performance. We believe these results highlight the unique contributions and potential applications of EDT.\n\n|                           | Diffuser [3] | Decision Diffuser [1] |   DT   |   EDT  |\n|:-------------------------:|:--------:|:-----------------:|:------:|:------:|\n|    hopper-medium-replay   |  $96.8$  |      $100.0$      | $61.9$ | $89.0$ |\n|    walker-medium-replay   |  $61.2$  |       $75.0$      | $43.3$ | $74.8$ |\n| halfcheetah-medium-replay |  $42.2$  |       $39.3$      | $34.9$ | $37.8$ |\n|       hopper-medium       |  $58.5$  |       $79.3$      | $60.7$ | $63.5$ |\n|       walker-medium       |  $79.7$  |       $82.5$      | $71.9$ | $72.8$ |\n| halfcheetah-medium-replay |  $44.2$  |       $49.1$      | $42.5$ | $42.5$ |\n\n**Q**: \"Can we demonstrate how the EDT idea could improve the performance of other DT variants?\"\n\n**A**: An example that illustrates the EDT idea in comparison to other DT variants can be seen in its application to the Online Decision Transformer (ODT) [4]. In ODT, action prediction is modeled as a stochastic policy, optimized by minimizing the negative log-likelihood (NLL) loss. To demonstrate the impact of EDT, we compare two methods: 1) Offline RL using stochastic action output, optimized with NLL loss, and 2) Offline RL incorporating EDT, but using the same stochastic action and NLL loss. Both methods follow standard DT training protocols. The experimental results, summarized in the table below, indicate that incorporating EDT into the architecture and training loss proposed by ODT leads to enhanced performance.\n\n|                        | hopper-medium | hopper-medium-replay |\n|:----------------------:|:-------------:|:--------------------:|\n|    ODT architecture    | $58.9\\pm 5.1$ |    $65.3\\pm 13.2$    |\n| ODT architecture + EDT | $62.3\\pm 3.8$ |     $82.8\\pm 8.1$    |\n\n---\n\nPlease do not hesitate to let us know if you have any additional comments.\n\n[1] Ajay et al., Is Conditional Generative Modeling all you need for Decision-Making? (2023), https://arxiv.org/abs/2211.15657\n\n[2] Lee et al., Multi-Game Decision Transformers (2022), https://arxiv.org/abs/2205.15241\n\n[3] Janner et al., Planning with Diffusion for Flexible Behavior Synthesis (2022), https://arxiv.org/abs/2205.09991\n\n[4] Zheng et al., Online Decision Transformer (2022), https://proceedings.mlr.press/v162/zheng22c.html\n"}]

Author rebuttal:
We thank all reviewers for their thoughtful comments and insights. We will revise our manuscript based on your feedback, and we have responded to each of your individual comments.

Summary of revisions: 
 * Fixing typo: “casual“ -> “causal“, “Figure 2“ -> “Table 2“ on line 234.
 * Adding explanation regarding tokenized return and details about return and observation loss computation.
 * Including notes about the interchangeable use of terms ‘expert’ and ‘optimal’ and the definition of returns ($R$, $\hat R$, $\tilde R$).
 * Adding related work Prompt DT and Hyper DT.
 * Implementation details and background knowledge regarding the expectile regression.
 * Improving clarity of the introduction section.

We would also like to provide a comparison of the training and inference times in terms of wall clock time. This demonstrates that the additional computational cost induced by the expectile regression loss and the optimal history length search is not significant. The experiments were conducted on a 3090 GPU, and the results show that the training time remains almost unchanged. Meanwhile, the additional action inference time constitutes only around 10% of the total inference time since the sampling from the gym environment takes considerably more time than action inference.

Wall clock time comparison (training):
|                                            |   DT   |   EDT  |
|:------------------------------------------:|:------:|:------:|
| wall clock training time (500 epochs, sec) | $1442$ | $1449$ |

Wall clock time vs. $\delta$ (action inference):
|               $\delta$               |    2   |    4   |    6   |    8   |   DT   |
|:------------------------------------:|:------:|:------:|:------:|:------:|:------:|
| wall clock time for 1000 steps (sec) | $12.6$ | $12.1$ | $11.8$ | $11.5$ | $11.4$ |

Again, we thank the reviewers for their constructive feedback. We believe that all comments have been addressed, but are happy to address any further comments from reviewers.

Best,
Author of 134


Meta-review: Decision transformers are a recent development in offline RL that aim to turn offline RL into a sequence modeling problem. In principle, DTs can "stich" the optimal parts of suboptimal trajectories together to create optimal behavior. In practice however, trajectory stiching is not found to empirically take place. This paper proposes a new method for training and deploying DTs (called Elastic Decision Transformer) and shows that it is better suited for stiching suboptimal trajectories together. Thus the proposed approach is a step towards realizing the potential benefits of DTs.

Main strengths: The proposed method is an intuitive modification to decision transformers that improves empirical performance and the paper includes a comphrensive empirical evaluation.

Main Concerns: The proposed method increases the latency of action selection using the trained transformer. The author rebuttal included wall clock time measurements that show the increase is significant but. not very large.

Decision transformers have recently become popular as a means to bring the power of sequence modeling to offline RL. This paper takes a step towards realizing some of the potential of DTs and thus I think it makes a useful contribution to the RL research community.
Decision: Accept (poster)