SlowGuess/NeurIPS_2023_OCR huggingface dataset:
4.3 Multi-Task Offline Reinforcement Learning
This section aims to evaluate the multi-task learning ability of our model across diverse tasks, focusing on locomotion and Atari tasks. Locomotion tasks utilize vectorized observations, while Atari tasks depend on image observations. To emphasize the role of trajectory stitching, we restrict our datasets to medium-replay datasets for the four locomotion tasks and datasets derived from DQN Replay [2] for the Atari tasks. Our evaluations span 20 different Atari tasks, with further environment setup details available in the Appendix.
Locomotion. In the locomotion multi-task experiment, we maintain the same model architecture as in the single-task setting. By confining the dataset to medium-replay datasets from four tasks, we increase task complexity and necessitate the offline RL approach to learn and execute these tasks concurrently, while effectively utilizing trajectories generated by random policies. As depicted in Table 2, our proposed EDT successfully accomplishes all four tasks simultaneously without much performance compromise.
Atari. For Atari, we adopt a CNN image encoder used in DrQ-v2 [59] to process stacks of four 84x84 image observations. To ensure fair comparisons, all methods employ the same architecture for the image encoder. Following [30], we incorporate random cropping and rotation for image augmentation. Additional experiment details are delegated to the Appendix for brevity. Performance on each Atari game is measured by human normalized scores (HNS) [36], defined as (score - scorerandom)/(scorehuman - scorerandom), to ensure a consistent scale across each game.
Our experimental results align with those of [30], highlighting that Q-learning-based offline RL approaches encounter difficulties in learning a multi-task policy on Atari games. Despite IQL achieving the highest score in Table1, it demonstrates relative inadequacy in simultaneous multi-task learning as indicated in Table 2 and Figure 5. We leave the raw scores of the 20 Atari games in Appendix B.
Figure 7: The figures illustrate the history length distributions across datasets and training epochs. For each distribution, we collect ten trajectories and derive a 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 of the current trajectory. We illustrated this idea with a toy example.
To validate the benefits of this dynamic history length approach, we tested the EDT using both fixed and variable history lengths. The results, summarized in Table 3, show that the variable history length outperforms the fixed ones, particularly on the "medium-replay" datasets.
These findings suggest that the EDT effectively chooses a history length that yields a higher estimated return. While a shorter history length aids in trajectory stitching, it's also crucial to retain a longer history length to ensure the continuity of optimal trajectories. Therefore, the dynamic adjustment of history length in the EDT is key to its superior performance.
4.5 Ablation Study on Expectile Level  $\alpha$
A key component of EDT is the approximation of the maximal value using expectile learning, as our method depends on accurately estimating these maximal values to choose the optimal history length. Consequently, examining the change in performance relative to the expectile level,  $\alpha$ , provides insight into the necessity of correct history length selection for performance enhancement.
The results, as displayed in Figure 6, suggest that when expectile regression is able to accurately approximate the maximizer, specifically at higher expectile levels, we observe both a higher average performance and lower standard deviation. This suggests that accurate selection of history length not only stabilizes performance but also enhances scores. Conversely, as the expectile level approaches 0.5, the expectile regression's objective shifts towards a mean square error, leading to an estimated value that is more of a mean value than a maximal one. This change makes it a less effective indicator for optimal history length. As a result, we can see a deterioration in EDT's score as the expectile level drops too low, and an increase in standard deviation, indicating inconsistency in the selection of an effective history length.
4.6 Analysis of Optimal History Length Distribution
In our analysis, we examine the history length distributions across various datasets, as depicted in Figure 7. Our findings reveal that the medium-replay dataset, which amalgamates trajectories from multiple policies, yields a distribution closely approximating a uniform distribution. Conversely, the medium dataset, acquired through a singular stochastic policy, exhibits a history length distribution characterized by an increased density at lower history lengths. This observation can be attributed to the prevalence of analogous trajectories within the medium dataset, leading to more frequent occurrences of trajectory stitching than the "medium-replay" dataset. However, it is important to acknowledge that the gains derived from this type of trajectory stitching remain limited, as the trajectories stem from identical distributions. Although performance improvement is observed, as presented in Table 1, it is significantly less pronounced in comparison to the medium-replay dataset.
Contrary to initial expectations, trajectory stitching does not occur as frequently within the medium-replay dataset as within the medium dataset. In fact, the distinct policies within the medium dataset contribute to the reduced instances of trajectory stitching, as their respective state distributions differ from one another. The diversity within the dataset results in a limited number of mutual  $s_t$
instances illustrated in Figure 3. Nevertheless, the proposed EDT method derives substantial benefits from trajectory stitching in this context. The EDT effectively avoids being misled by sub-optimal trajectories within the dataset, demonstrating its capacity to make better decisions regarding history lengths and actions that optimize the current return.
Offline Reinforcement Learning. Offline RL has been a promising topics for researchers since sampling from environments during training is usually costly and dangerous in real-world applications and offline reinforcement learning is able to learn a better policy without directly collecting 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 as a planning algorithm and demonstrates exceptional flexibility across various applications, such as long-horizon dynamics prediction, imitation learning, goal-conditioned reinforcement learning, and offline reinforcement learning.
Recently, there has been a growing interest in incorporating diffusion models into offline RL methods. This alternative approach to decision-making stems from the success of generative modeling, which offers the potential to address offline RL problems more effectively. For instance, [18] reinterprets Implicit Q-learning as an actor-critic method, using samples from a diffusion parameterized behavior policy to improve performance. Similarly, other diffusion-based methods [21, 32, 50, 5, 10, 4] utilize diffusion-based generative models to represent policies or model dynamics, achieving competitive or superior performance across various tasks.
Trajectory Stitching. A variety of methods have been proposed to tackle the trajectory stitching problem in offline RL. The Q-learning Decision Transformer (QDT) [59] stands out as it relabels the ground-truth return-to-go with estimated values, a technique expected to foster trajectory recombination. Taking a different approach, [19] utilizes a model-based data augmentation strategy, stitching together parts of historical demonstrations to create superior trajectories. Similarly, the Best Action Trajectory Stitching (BATS) [9] algorithm forms a tabular Markov Decision Process over logged data, adding new transitions using short planned trajectories. BATS not only aids in identifying advantageous trajectories but also provides theoretical bounds on the value function. These efforts highlight the breadth of strategies employed to improve offline RL through innovative trajectory stitching techniques.
Conclusion. In this paper, we introduced the Elastic Decision Transformer, a significant enhancement to the Decision Transformer that addresses its limitations in offline reinforcement learning. EDT's innovation lies in its ability to determine the optimal history length, promoting trajectory stitching. We proposed a method for estimating this optimal history length by learning an approximate value optimizer through expectile regression.
Our experiments affirmed EDT's superior performance compared to DT and other leading offline RL algorithms, notably in multi-task scenarios. EDT's implementation is computationally efficient and straightforward to incorporate with other DT variants. It outshines existing methods on the D4RL benchmark and Atari games, underscoring its potential to propel offline RL forward.
In summary, EDT offers a promising solution for trajectory stitching, enabling the creation of better sequences from sub-optimal trajectories. This capability can considerably enhance DT variants, leading to improved performance across diverse applications. We are committed to releasing our code.
Limitations. A potential direction for future improvement involves enhancing the speed at which EDT estimates the optimal history. This could make the method suitable for real-time applications that have strict time constraints. While this adaptation is an exciting avenue for future research, it falls outside the primary scope of this paper.
This project is supported by a grant from JST ACT-X, grant number JPMJAX22AC.
[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 Learning, pages 104-114. PMLR, 2020.
[3] Dennis J Aigner, Takeshi Amemiya, and Dale J Poirier. On the estimation of production frontiers: maximum likelihood estimation of the parameters of a discontinuous density function. International economic review, pages 377-396, 1976.
[4] Anurag Ajay, Yilun Du, Abhi Gupta, Joshua Tenenbaum, Tommi Jaakkola, and Pulkit Agrawal. Is conditional generative modeling all you need for decision-making? arXiv preprint arXiv:2211.15657, 2022.
[5] Arthur Argenson and Gabriel Dulac-Arnold. Model-based offline planning. arXiv preprint arXiv:2008.05556, 2020.
[6] Adrià Puigdomènech Badia, Bilal Piot, Steven Kapturowski, Pablo Sprechmann, Alex Vititskyi, 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: An evaluation platform for general agents. Journal of Artificial Intelligence Research, 47:253-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.12026, 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.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, 34:15084-15097, 2021.
[12] Xinshi Chen, Shuang Li, Hui Li, Shaohua Jiang, Yuan Qi, and Le Song. Generative adversarial user model for reinforcement learning based recommendation system. In International Conference on Machine Learning, pages 1052-1061. PMLR, 2019.
[13] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805, 2018.
[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 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.11929, 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.
[52] Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumont, Clement Delangue, Anthony Moi, Pierrick Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, et 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 arXiv:1911.11361, 2019.
[54] Yueh-Hua Wu, Jiashun Wang, and Xiaolong Wang. Learning generalizable dexterous manipulation from human grasp affordance. In Conference on Robot Learning, pages 618-629. PMLR, 2023.
[55] Teng Xiao and Donglin Wang. A general offline reinforcement learning framework for interactive recommendation. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pages 4512-4520, 2021.
[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. Hyperdecision transformer for efficient online policy adaptation. arXiv preprint arXiv:2304.08487, 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 rl. arXiv preprint arXiv:2209.03993, 2022.
[60] Denis Yarats, Rob Fergus, Alessandro Lazaric, and Lerrel Pinto. Mastering visual continuous control: Improved data-augmented reinforcement learning. arXiv preprint arXiv:2107.09645, 2021.
A Additional ablation study
A.1 Step Size Ablation
In Algorithm 1, we introduced the step size parameter  $\delta$ , acting as a balance between search granularity and inference speed. The effects of varying  $\delta$  are depicted in Figure 8. To compute the history length search space, we commence with the maximum history length. For instance, if the maximum history length is  $T = 20$  and the step size  $\delta = 8$ , the history length search space becomes 20, 12, 4.
Figure 8 shows that a narrowed search space leads to a decline in return performance and an increase in standard deviation, corroborating results from Table 3. EDT with a restricted history search space behaves more like EDT with a fixed history length. We found that the EDT works best when we chose a search step of 4.
Figure 8: Ablation study on the step size  $\delta$ . The greater the step size the smaller the length search space is.
This could be because it is tough to make good estimates without being able to interact with the environment directly, as mentioned in previous studies [26, 28]. By increasing the step size, we were able to make our estimates more reliable and less affected by noise and other problems.
A.2 History Length Distribution
We further show more history length distributions for the locomotion tasks here.
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 a histogram of history lengths. The distribution is computed with kernel distribution estimate for better visualization.
Figure 10: The figures illustrate the history length distributions across datasets and training epochs on the medium-replay dataset. For each distribution, we collect ten trajectories and derive a 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 on longer history lengths. The distribution of history lengths in the medium and medium-replay datasets mostly supports this hypothesis.
In the medium-replay dataset, we consistently see a concentration of shorter history lengths. This was expected, given that the nature of the medium-replay dataset is likely to produce shorter history lengths. The distribution often converges towards higher densities for shorter lengths, which aligns with our expectations.
The pattern within the medium dataset, however, is less consistent. This could be attributed to several factors that can either elongate or truncate the history length. Despite these fluctuations, we still
observe a slight inclination towards longer history lengths. This meets our initial assumption but also demonstrates the complexity within the medium dataset.
A.3 Comparison with Diffusion-Based Methods
B.1 Loss Computation Details
In this section, we detail the computation of the tokenized return loss  $\mathcal{L}_{\mathrm{return}}$  and the observation loss  $\mathcal{L}_{\mathrm{observation}}$  in Eq 1. For  $\mathcal{L}_{\mathrm{return}}$ , we first tokenize the predicted return and the target return into  $n$  bins bounded by the MAX and MIN scores obtained from D4RL [15]. We then compute the cross-entropy loss between the predicted value and the target value:
$$
\mathcal {L} _ {\text {r e t u r n}} = C E \left(\hat {R} ^ {t}, R ^ {\prime t}\right), \tag {5}
$$
where  $R^{\prime t}$  suggests the predicted tokenized return. For  $\mathcal{L}_{\mathrm{observation}}$ , we predict the next observation by minimizing the mean square error.
B.2 Learning on Locomotion Tasks
In Section 4, we discussed the application of our approach to a multi-task learning scenario. Specifically, we consolidated medium-replay datasets from Ant, HalfCheetah, Hopper, and Walker2d. To standardize returns across these varied environments, we used scaling statistics (maximum return, minimum return, and return scale) from the official Decision Transformer repository (https://github.com/kzl/decision-transformer). As detailed in Section 3.3, we further segmented the return into 60 discrete buckets and, following the approach of [20, 30], sampled from the top 85th percentile logits within the discrete return distribution.
However, it's important to highlight that our return maximizer,  $\tilde{R}$ , estimates the scaled value directly rather than the discretized one. For the sequence  $\langle \dots, \mathbf{o}, R, a, \dots \rangle$ , we supplement each token embedding with a learned positional embedding. Given the differing state and action dimensions across tasks, we employ an additional state and action embedding layer. This layer transforms the raw state and action representation into a consistent state and action embedding size, while keeping the remainder of the architecture unchanged across tasks.
Approximate  $\tilde{R}$  Estimation. Section 3.4 outlines how we use Bayes' Rule to estimate expert returns. A conventional  $\tilde{R}$  prediction typically involves an autoregressive process, given the necessity of sampling from the discrete  $\hat{R}_{\mathrm{expert}}$  distribution at each time step during the search for the optimal history length. To simplify this process, we follow the approximation strategy used in [30]'s implementation. We mask all return values except the first one, thus making  $\tilde{R}$  solely dependent on  $\mathbf{o}$ ,  $a$ , and the initial return value.
History Length Search Heuristic. In Algorithm 1, we illustrated that a larger  $\delta$  value allows us to infer actions more rapidly. Adopting this method needs to balance inference speed with search accuracy. Based on the concept introduced in Sec. 3.2, the optimal history length at the current state  $s_{t + 1}$  might be close to that of the previous state  $s_t$ . Therefore, given the optimal length  $l_{t}$  at time step  $t$ , we search within the range  $\{l_t - \Delta ,l_t - \Delta +1\dots ,l_t + \Delta \}$  for the optimal length  $l_{t + 1}$  at the next step.
B.3 Multi-Task Learning on Atari Games
The process for action inference in multi-task learning on Atari games closely aligns with that described in Sec. B.2, including the method for approximating  $\tilde{R}$ . However, there are several noteworthy distinctions, which we elaborate upon in this section.
Given that all games utilize grayscale frames of the same dimensions (84x84), we do not need to implement an additional state embedding layer as we did in the Locomotion scenario. Instead, we introduce a shared image encoder for all games, with further details outlined in Sec. B.5.
It's important to note the distinction between the action spaces of Atari games and Locomotion tasks. The action space of Atari games is discrete, so our sampling approach mirrors how we sample from the return distribution: we select from the top 85th percentile of action logits. Inspired by [30], we
discretize the reward signal into  $\{-1,0, + 1\}$ , while the return is split into 120 buckets ranging from  $\{-20,\dots,100\}$ .
Given the complexity and potential pitfalls of learning on 20 Atari games from scratch, we use a subset of GPT-2 to initialize the transformer decoder. This step improves both the convergence rate and the quality of the learned policy. Our dataset consists of 2 training runs from [2], with each run featuring rollouts from 50 checkpoints. Lastly, we enhance the dataset with image augmentations, including random cropping and rotation.
B.4 Atari Games
Due to time and computational constraints, we randomly selected 20 tasks from the 41 tasks in the study by [30]. Details about the game types and number of action spaces can be found in Table 4. We also provide the raw scores of the Atari experiments in Table 5. Given that we transform the original observations to grayscale and rescale them to  $84 \times 84$  images, we illustrate these transformed observations in Figure 11.
Table 4: Atari Games: Name, Category, and Action Space
Game	Category	Action Space
Assault	Shooter	7
Asterix	Platform	18
BankHeist	Strategy	18
BeamRider	Shooter	9
Breakout	Arcade	4
Centipede	Shooter	18
ChopperCommand	Shooter	18
Enduro	Racing	9
FishingDerby	Sports	18
Freeway	Racing	3
Frostbite	Platform	18
Gravitar	Shooter	18
NameThisGame	Shooter	8
Phoenix	Shooter	8
Qbert	Puzzle	6
Seaquest	Shooter	18
TimePilot	Shooter	10
VideoPinball	Pinball	3
WizardofWor	Shooter	10
Zaxxon	Shooter	18