SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Yueh-Hua Wu $^{12}$  Xiaolong Wang $^{1*}$  Masashi Hamaya $^{2*}$
$^{1}$ UC San Diego  $^{2}$ OMRON SINIC X
yuwo88@ucsd.edu
This paper introduces Elastic Decision Transformer (EDT), a significant advancement over the existing Decision Transformer (DT) and its variants. Although DT purports to generate an optimal trajectory, empirical evidence suggests it struggles with trajectory stitching, a process involving the generation of an optimal or near-optimal trajectory from the best parts of a set of sub-optimal trajectories. The proposed EDT differentiates itself by facilitating trajectory stitching during action inference at test time, achieved by adjusting the history length maintained in DT. Further, the EDT optimizes the trajectory by retaining a longer history when the previous trajectory is optimal and a shorter one when it is sub-optimal, enabling it to "stitch" with a more optimal trajectory. Extensive experimentation demonstrates EDT's ability to bridge the performance gap between DT-based and Q Learning-based approaches. In particular, the EDT outperforms Q Learning-based methods in a multi-task regime on the D4RL locomotion benchmark and Atari games. Videos are available at: https://kristery.github.io/edt/.
Reinforcement Learning (RL) trains agents to interact with an environment and learn from rewards. It has demonstrated impressive results across diverse applications such as game playing [23, 44], robotics [41, 40, 54], and recommendation systems [12, 1]. A notable area of RL is Offline RL [31], which employs pre-collected data for agent training and proves more efficient when real-time interactions are costly or limited. Recently, the conditional policy approach has shown large potentials in Offline RL, where the agent learns a policy based on the observed state and a goal. This approach enhances performance and circumvents stability issues related to long-term credit assignment. Moreover, the successful Transformer architecture [49], widely
Figure 1: Normalized return with medium-replay datasets. The dotted gray lines indicate normalized return with medium datasets. By achieving trajectory stitching, our method benefits from worse trajectories and learns a better policy.
used in applications like natural language processing [52, 13, 8] and computer vision [14, 34], has been adapted for RL as the Decision Transformer (DT) [11].
DT utilizes a Transformer architecture to model and reproduce sequences from demonstrations, integrating a goal-conditioned policy to convert Offline RL into a supervised learning task. Despite its competitive performance in Offline RL tasks, the DT falls short in achieving trajectory stitching, a desirable property in Offline RL that refers to creating an optimal trajectory by combining parts of sub-optimal trajectories [19, 9, 59]. This limitation stems from the DT's inability to generate
* equal advising
superior sequences, thus curbing its potential to learn optimal policies from sub-optimal trajectories (Figure 1).
We introduce the Elastic Decision Transformer (EDT), which takes a variable length of the traversed trajectory as input. Stitching trajectories, or integrating the current path with a more advantageous future path, poses a challenge for sequence generation-based approaches in offline RL. Stitching a better trajectory appears to contradict one of the core objectives of sequence generation that a sequence generation model is required to reliably reproduce trajectories found within the training dataset. We suggest that in order to 'refresh' the prediction model, it should disregard 'negative' or 'unsuccessful' past experiences. This involves dismissing past failures and instead considering a shorter history for input. This allows the sequence generation model to select an action that yields a more favorable outcome. This strategy might initially seem contradictory to the general principle that decisions should be based on as much information as possible. However, our proposed approach aligns with this concept. With a shorter history, the prediction model tends to output with a higher variance, typically considered a weakness in prediction scenarios. Yet, this increased variance offers the sequence prediction model an opportunity to explore and identify improved trajectories. Conversely, when the current trajectory is already optimal, the model should consider the longest possible history for input to enhance stability and consistency. Consequently, a relationship emerges between the quality of the path taken and the length of history used for prediction. This correlation serves as the motivation behind our proposal to employ a variable length of historical data as input.
In practice, we train an approximate value maximizer using expectile regression to estimate the highest achievable value given a certain history length. We then search for the history length associated with the highest value and use it for action inference.
Evidence from our studies indicates that EDT's variable-length input sequence facilitates more effective decision-making and, in turn, superior sequence generation compared to DT and its variants. Furthermore, it is computationally efficient, adding minimal overhead during training. Notably, EDT surpasses state-of-the-art methods, as demonstrated in the D4RL benchmark [15] and Atari games [7, 17]. Our analysis also suggests that EDT can significantly enhance the performance of DT, establishing it as a promising avenue for future exploration.
Our Contributions:
- We introduce the Elastic Decision Transformer, a novel approach to Offline Reinforcement Learning that effectively addresses the challenge of trajectory stitching, a known limitation in Decision Transformer.
- By estimating the optimal history length based on changes in the maximal value function, the EDT enhances decision-making and sequence generation over traditional DT and other Offline RL algorithms.
- Our experimental evaluation highlights EDT's superior performance in a multi-task learning regime, positioning it as a promising approach for future Offline Reinforcement Learning research and applications.
In this study, we consider a decision-making agent that operates within the framework of Markov Decision Processes (MDPs) [42]. At every time step  $t$ , the agent receives an observation of the world  $o_t$ , chooses an action  $a_t$ , and receives a scalar reward  $r_t$ . Our goal is to learn a single optimal policy distribution  $P_\theta^*(a^t | o^{\leq t}, a^{<t}, r^{<t})$  with parameters  $\theta$  that maximizes the agent's total future return  $R_t = \sum_{k > t} r^k$  on all the environments we consider.
2.1 Offline Reinforcement Learning
Offline RL, also known as batch RL, is a type of RL where an agent learns to make decisions by analyzing a fixed dataset of previously collected experiences, rather than interacting with an environment in real-time. In other words, the agent learns from a batch of offline data rather than actively exploring and collecting new data online.
Offline RL has gained significant attention in recent years due to its potential to leverage large amounts of pre-existing data and to solve RL problems in scenarios where online exploration is
Figure 2: An overview of our Elastic Decision Transformer architecture.  $\tilde{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 to the fixed dataset. A number of recent research efforts have addressed these challenges, including methods for importance weighting [35, 37], regularization [53, 28], and model-based learning [25, 29], among others.
2.2 Decision Transformer
The Decision Transformer architecture, introduced by [11], approaches the offline RL problem as a type of sequence modeling. Unlike many traditional RL methods that estimate value functions or compute policy gradients, DT predicts future actions based on a sequence of past states, actions, and rewards. The input to DT includes a sequence of past states, actions, and rewards, and the output is the next action to be taken. DT uses a Transformer architecture [49], which is composed of stacked self-attention layers with residual connections. The Transformer architecture has been shown to effectively process long input sequences and produce accurate outputs.
Despite the success of being applied to offline RL tasks, it has a limitation in its ability to perform "stitching." Stitching refers to the ability to combine parts of sub-optimal trajectories to produce an optimal trajectory. This approach can lead to a situation where the agent follows a sub-optimal trajectory that provides an immediate reward, even if a different trajectory leads to a higher cumulative reward over time. This limitation of DT is a significant challenge in many offline RL applications, and addressing it would greatly enhance the effectiveness of DT in solving real-world problems.
3 Elastic Decision Transformer
In this section, we present Elastic Decision Transformer (EDT), a model that automatically utilizes a shorter history to predict the next action when the traversed trajectory underperforms compared to those in the training dataset. The mechanism allows the model to switch to a better trajectory by forgetting 'unsuccessful' past experiences, thus opening up more possibilities for future trajectories. We further propose a method to estimate the maximum achievable return using the truncated history, allowing EDT to determine the optimal history length and corresponding actions.
3.1 Reinforcement Learning as Sequence Modeling
In this paper, we adopt an approach to offline reinforcement learning that is based on a sequence modeling problem. Specifically, we model the probability of the next token in the sequence (denoted as  $\tau$ ) based on all the tokens that come before it. The sequences we model can be represented as:
$$
\tau = \langle \dots , o ^ {t}, \hat {R} ^ {t}, a ^ {t}, \dots \rangle ,
$$
where  $t$  is a time step and  $\hat{R}$  is the return for the remaining sequence. The sequence we consider here is similar to the one used in [30] whereas we do not include reward as part of the sequence and we predict an additional quantity  $\tilde{R}$  that enables us to estimate an optimal input length, which we will cover in the following paragraphs. Figure 2 presents an overview of our model architecture. It should
be noted that we also change the way to predict future observation from standard DT [11], where the next observation is usually directly predicted from  $a^t$  through the causal transformer decoder.
3.2 Motivation
We propose a shift in the traditional approach to trajectory stitching. Instead of focusing on training phases, we aim to achieve this stitching during the action inference stage. This concept is illustrated in Figure 3 using a simplified example. In this scenario, we consider a dataset,  $D$ , comprising only two trajectories:  $D = (s_{t-1}^{a}, s_{t}, s_{t+1}^{a}), (s_{t-1}^{b}, s_{t}, s_{t+1}^{b})$ . A sequence model trained with this dataset is likely to predict the next states in a manner consistent with their original trajectories.
Figure 3: A Toy example to illustrate the motivation of EDT. The figure shows an offline RL dataset that contains only two trajectories  $(s_{t - 1}^{a}, s_{t}, s_{t + 1}^{a}), (s_{t - 1}^{b}, s_{t}, s_{t + 1}^{b})$ .
To overcome this, we propose a method that enables trajectory stitching, where the model starts from  $s_{t-1}^b$  and concludes at  $s_{t+1}^a$ . This is achieved by adaptively adjusting the history length. We introduce a maximal value estimator,  $\tilde{R}$ , which calculates the maximum value
among all potential outcomes within the dataset. This allows us to determine the optimal history length that maximizes  $\tilde{R}$ .
In the given example, if the model starts at state  $s_{t-1}^b$ , it will choose to retain the history  $(s_t)$  upon reaching state  $s_t$ , as  $\tilde{R}(s_t) > \tilde{R}(s_{t-1}, s_t)$ . Conversely, if the model initiates from state  $s_{t-1}^a$ , it will preserve the history  $(s_{t-1}^a, s_t)$  when decision-making at  $s_t$ , as  $\tilde{R}(s_{t-1}^a, s_t) \geq \tilde{R}(s_t)$ . From the above example, we understand that the optimal history length depends on the quality of the current trajectory we've traversed, and it can be a specific length anywhere between a preset maximal length and a single unit.
To estimate the optimal history length in a general scenario, we propose solving the following optimization problem:
$$
\underset {T} {\arg \max } \underset {\tau_ {T} \in D} {\max } \hat {R} ^ {t} \left(\tau_ {T}\right), \tag {1}
$$
where  $\tau_{T}$  denotes the history length  $T$ . More precisely,  $\tau_{T}$  takes the form:
$$
\tau_ {T} = \langle o ^ {t - T + 1}, \hat {R} ^ {t - T + 1}, a ^ {t - T + 1}, \dots , o ^ {t - 1}, \hat {R} ^ {t - 1}, a ^ {t - 1}, o ^ {t}, \hat {R} ^ {t}, a ^ {t} \rangle .
$$
3.3 Training objective for Maximum In-Support Return
In the EDT, we adhere to the same training procedure as used in the DT. The key distinction lies in the training objective - we aim to estimate the maximum achievable return for a given history length in EDT. To approximate the maximum operator in  $\max_{\tau_T \in D} \hat{R}^t(\tau_T)$ , we employ expectile regression [38, 3], a technique often used in applied statistics and econometrics. This method has previously been incorporated into offline reinforcement learning; for instance, IQL [26] used expectile regression to estimate the Q-learning objective implicitly. Here, we leverage it to enhance our estimation of the maximum expected return for a trajectory, even within limited data contexts. The  $\alpha \in (0,1)$  expectile of a random variable  $X$  is the solution to an asymmetric least squares problem, as follows:
$$
\underset {m _ {\alpha}} {\arg \min } \mathbb {E} _ {x \in X} \left[ L _ {2} ^ {\alpha} (x - m _ {\alpha}) \right],
$$
where  $L_2^\alpha (u) = |\alpha -\mathbb{1}(u < 0)|u^2$
Through expectile regression, we can approximate  $\max_{\tau_T\in D}\hat{R}^t (\tau_T)$
$$
\tilde {R} _ {T} ^ {t} = \max  _ {\tau_ {T} \in D} \hat {R} ^ {t} \left(\tau_ {T}\right) \approx \underset {\tilde {R} ^ {t} \left(\tau_ {T}\right)} {\arg \min } \mathbb {E} _ {\tau_ {T} \in D} \left[ L _ {2} ^ {\alpha} \left(\tilde {R} ^ {t} \left(\tau_ {T}\right) - \hat {R} ^ {t}\right) \right]. \tag {2}
$$
We estimate  $\tilde{R}^t$  by applying an empirical loss of Equation 2 with a sufficiently large  $\alpha$  (we use  $\alpha = 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:
$$
\mathcal {L} _ {\mathrm {E D T}} = c _ {r} \mathcal {L} _ {\text {r e t u r n}} + \mathcal {L} _ {\text {o b s e r v a t i o n}} + \mathcal {L} _ {\text {a c t i o n}} + \mathcal {L} _ {\max }, \tag {3}
$$
Figure 4: The figure illustrates the action inference procedure within the proposed Elastic Decision Transformer. Initially, we estimate the value maximizer,  $\tilde{R}_i$ , for each length  $i$  within the search space, as delineated by the green rectangle. Subsequently, we identify the maximal value from all  $\tilde{R}_i$ , which provides the optimal history length  $w$ . Utilizing this optimal history length, we estimate the expert value at time step  $t$ , denoted as  $\tilde{R}_{w,e}^t$ , 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_i^t$  during the estimation process for  $\tilde{R}_i$  and we present the inference here for clarity.
Algorithm 1 EDT optimal history length search
1: Input: a query sequence  $\tau = \langle \dots, o^{t-1}, R^{t-1}, a^{t-1}, o^t \rangle$  and EDT model  $\theta_{EDT}$
2: for  $w = 1, 1 + \delta, 1 + 2\delta, \dots, T$  do
3: Obtain  $\tilde{R}^t (\tau^w)$  with truncated sequence  $\tau^w = \langle o^{t - w + 1},R^{t - w + 1},a^{t - w + 1},\dots,o^t\rangle$  and  $\theta_{EDT}$
4: end for
5: Compute  $P(R^t)$  with  $\theta_{EDT}$  for  $\tau^w$  that has the highest  $\tilde{R}^t$  and then estimate  $P(R^t | \text{expert}^t, \ldots)$  with Equation 4.
where  $\mathcal{L}_{\mathrm{observation}}$  and  $\mathcal{L}_{\mathrm{action}}$  are computed with a mean square error,  $\mathcal{L}_{\mathrm{return}}$  is a cross-entropy loss, and  $\mathcal{L}_{\mathrm{max}}$  is an empirical estimate of Equation 2. We set  $c_{r} = 0.001$  to balance scale differences between mean square error and cross-entropy loss. In tasks with discrete action spaces like Atari, we optimize the action space as the tokenized return objective  $\mathcal{L}_{\mathrm{return}}$  using cross-entropy with weight  $10c_{r}$ . Our training method extends the work of [30] by estimating the maximum expected return value for a trajectory using Equation 2. This estimation aids in comparing expected returns of different trajectories over various history lengths. Our proposed method is not only easy to optimize, but can also be conveniently integrated with other DT variants. As such, it marks a significant advance in developing efficient offline reinforcement learning approaches for complex decision-making tasks.
3.4 Action Inference During Test time
During action inference phase in test time, we first (1) estimate the maximum achievable return  $\tilde{R}_i$  for each history length  $i$ . Subsequently, (2) we predict the action by using the truncated traversed trajectory as input. The trajectory is truncated to the history length that corresponds to the highest value of  $\tilde{R}_i$ . These steps are elaborated in Figure 4.
To identify the history length  $i$  that corresponds to the highest  $\tilde{R}_i^t$ , we employ a search strategy as detailed in Algorithm 1. As exhaustively searching through all possible lengths from 1 to  $T$  may result in slow action inference, we introduce a step size  $\delta$  to accelerate the process. This step size not only enhances inference speed by a factor of  $\delta$ , but also empirically improves the quality of the learned policy. An ablation study on the impact of the step size  $\delta$  is provided in Appendix A. For all experiments, we set  $\delta = 2$  to eliminate the need for parameter tuning.
To sample from expert return distribution  $P(R^{t}, \ldots | \text{expert}^{t})$ , we adopt an approach similar to [30] by applying Bayes' rule  $P(R^{t}, \ldots | \text{expert}^{t}) \propto P(\text{expert}^{t} | R^{t}, \ldots) P(R^{t}, \ldots)$  and approximate the distribu
Dataset	DT	QDT	TS+BC	S4RL	IQL	EDT (Ours)
hopper-medium	60.7 ± 4.5	57.2 ± 5.6	64.3 ± 4.2	78.9	63.8 ± 9.1	63.5 ± 5.8
hopper-medium-replay	61.9 ± 13.7	45.8 ± 35.5	50.2 ± 17.2	35.4	92.1 ± 10.4	89.0 ± 8.3
walker-medium	71.9 ± 3.9	67.5 ± 2.0	78.8 ± 1.2	93.6	79.8 ± 3.0	72.8 ± 6.2
walker-medium-replay	43.3 ± 14.3	30.3 ± 16.2	61.5 ± 5.6	30.3	73.6 ± 6.3	74.8 ± 4.9
halfcheetah-medium	42.5 ± 0.4	42.3 ± 2.5	43.2 ± 0.3	48.8	47.3 ± 0.2	42.5 ± 0.9
halfcheetah-medium-replay	34.9 ± 1.6	30.0 ± 11.1	39.8 ± 0.6	51.4	44.1 ± 1.1	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
ant-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 is highlighted. The results of QDT, TS+BC, and S4RL are adopted from their reported scores. Following [26], we emphasize in bold scores within 5 percent of the maximum per task ( $\geq 0.95 \cdot \max$ ).
tion of expert-level return with inverse temperature  $\kappa$  [24, 48, 47, 43]:
$$
P \left(R ^ {t} \mid \text {e x p e r t} ^ {t}, \dots\right) \propto \exp (\kappa R ^ {t}) P \left(R ^ {t}\right). \tag {4}
$$
While it may initially appear feasible to directly use the predicted  $\tilde{R}$  as the expert return, it's important to note that this remains a conservative maximum operation. Empirically, we have found that Eq. 4 encourages the pursuit of higher returns, which consequently enhances the quality of the actions taken.
Our experiments are designed to address several key questions, each corresponding to a specific section of our study:
- Does EDT significantly outperform DT and its variants? (Sec. 4.2, 4.3)
- Is the EDT effective in a multi-task learning regime, such as Locomotion and Atari games? (Sec. 4.3)
- Does a dynamic history length approach surpass a fixed length one? (Sec. 4.4)
- How does the expectile level  $\alpha$  impact the model's performance? (Sec. 4.5)
- How does the quality of datasets affect the predicted history lengths? (Sec. 4.6)
We also provide an additional ablation study in Appendix A due to space constraints.
In the subsequent section, we draw comparisons with two methods based on the Decision Transformer: the original Decision Transformer (DT) [11] and the Q-learning Decision Transformer (QDT) [59]. Additionally, we include a behavior cloning-based method (TS+BC) [19], as well as two offline Q-learning methods, namely S4RL [45] and IQL [26], in our comparisons.
It is important to note that QDT and TS+BC are specifically designed to achieve trajectory stitching. QDT accomplishes this by substituting collected return values with estimates derived from Conservative Q-Learning [28]. Conversely, TS+BC employs a model-based data augmentation strategy to bring about the stitching of trajectories."
4.2 Single-Task Offline Reinforcement Learning
For locomotion tasks, we train offline RL models on D4RL's "medium" and "medium-replay" datasets. The "medium" dataset comes from a policy reaching about a third of expert performance. The 'medium-replay', sourced from this policy's replay buffer, poses a greater challenge for sequence modeling approaches such as DT.
We conclude our locomotion results in Table 1. Since the proposed model estimates the return of the current sequence, reward information is not required during test time.
Our observations indicate that the proposed EDT consistently outperforms the baseline DT and its variants on the majority of the datasets, with a notable performance advantage on the "medium
Task	DT-1	IQL-1	EDT-1 (Ours)
hopper	51.2	59.8	76.9
walker	29.8	52.6	74.1
halfcheetah	30.5	40.4	36.8
ant	79.8	82.3	88.6
sum	191.3	235.1	276.4
Figure 5: The average HNS comparison on 20 Atari games. The results are evaluated with three trials.
Table 2: Evaluation results on multi-task regime. Mean of 5 random training initialization seeds, 100 evaluations each. The training dataset is a mixture of medium-replay datasets from the four locomotion tasks. Our main result is highlighted.
Dataset	EDT (w=20)	EDT (w=5)	EDT (Ours)
hopper-m	63.6 ± 5.2	57.8 ± 7.0	63.5 ± 5.8
hopper-mr	67.6 ± 27.7	62.6 ± 26.9	89.0 ± 8.3
walker-m	65.6 ± 11.7	62.6 ± 14.3	72.8 ± 6.2
walker-mr	64.5 ± 12.9	44.3 ± 8.7	74.8 ± 4.9
halfcheetah-m	42.0 ± 0.4	42.3 ± 0.8	42.5 ± 0.9
halfcheetah-mr	33.5 ± 8.0	36.4 ± 7.4	37.8 ± 1.5
ant-m	90.8 ± 16.0	95.6 ± 8.0	97.9 ± 8.0
ant-mr	80.0 ± 17.0	82.3 ± 15.9	92.0 ± 4.1
sum	507.6	483.9	570.3
Table 3: Mean of 5 random training initialization 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.
Figure 6: Ablation study on expectile level  $\alpha$ . Expectile objective reduces to Mean Square Error when  $\alpha = 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, a task that DT and its variants cannot accomplish. Although EDT doesn't fully outperform IQL in the single-task learning, it does bridge the gap between Q-learning-based methods and DT by performing trajectory stitching with the estimated maximum return.