Elastic Decision Transformer

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.

Strengths

- This paper is well presented and easy to follow. - The idea to dynamically adjust history to achieve trajectory stitching is novel and well motivated. - Experiment results show that the proposed method seems robust enough to benefit from optimal and sub-optimal trajectories. - 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. - Chosen baselines might not cover all recent progresses in offline RL, especially those specifically designed to address trajectory stitching. - Comparison to recent methods based on Diffusion model, including Diffuser[Janner et al., 2022] and Decision Diffuser[Ajay et al., 2023], is missed. - Though the paper claims EDT to be generally applicable for all DT variants, it is only evaluated against the original DT. - Typo, L234 should be Table 2 instead of Figure 2. References: - Janner et al., Planning with Diffusion for Flexible Behavior Synthesis, ICML 2022. - Ajay et al., Is Conditional Generative Modeling all you need for Decision Making?, ICLR 2023.

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]? - Can we demonstrate how the EDT idea could improve the performance of other DT variants? References: - Janner et al., Planning with Diffusion for Flexible Behavior Synthesis, ICML 2022. - Ajay et al., Is Conditional Generative Modeling all you need for Decision Making?, ICLR 2023.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

4 excellent

Contribution

2 fair

Limitations

See above.

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.

Strengths

** Strengths** 1. The concept of Elastic Decision Transformer (EDT) offers a novel approach to tackle the issue of trajectory stitching. 2. 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. 3. The empirical results demonstrating EDT's superior performance in various benchmarks provide robust support for the proposed model.

Weaknesses

** Weaknesses ** 1. 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]. 2. 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. 3. 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. [1] Xu, Mengdi et al. “Prompting Decision Transformer for Few-Shot Policy Generalization.” ArXiv abs/2206.13499 (2022): n. pag. [2] Xu, Mengdi et al. “Hyper-Decision Transformer for Efficient Online Policy Adaptation.” ArXiv abs/2304.08487 (2023): n. pag. [3] Lee, Kuang-Huei et al. “Multi-Game Decision Transformers.” ArXiv abs/2205.15241 (2022): n. pag.

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. 2. 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. 3. How does the EDT handle scenarios where the distinction between optimal and sub-optimal trajectories is not clear? 4. How would EDT perform with extremely long trajectories where the computation cost may become prohibitive? 5. Could you expand on how the computation overhead during training is minimized?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

The major limitation of this paper is value maximizer training speed, as illustrated in the paper. I don't find any obvious broader societal impacts.

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.

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. This 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. The modifications to DT are simple, well-justified, and effective. The paper is overall clearly written and easy to follow.

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. **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. **Minor points.** Line 94: DT uses rewards-to-go, or return, not rewards. Line 205: Extra quote mark at the end of the line. Structure 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. It 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. **Missing references** The paper considered the multi-task regime, but some related works are missing: Reed, Scott, et al. "A generalist agent." arXiv preprint arXiv:2205.06175 (2022). Xu, Mengdi, et al. "Prompting decision transformer for few-shot policy generalization." international conference on machine learning, 2022.

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? 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? 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?)

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

3 good

Contribution

4 excellent

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).

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.

Strengths

* The concept of adaptively adjusting the history length input to the DT model is intriguing. * The introduced solution addresses the trajectory stitching problem in the original DT model. * The authors utilize expectile return regression to select the value of each candidate history length * 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. * 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. * 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. * 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. * The inclusion of an algorithm table to detail the entire process would have significantly aided reader understanding. * There is an apparent lack of effort to make the paper self-contained.

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? * Why isn't there a comparison with Multi-game DT[29], especially for multi-task experiments? * 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. * The calculation process for P(R^t) in Equation (4) is unclear. Could the authors provide more details on this? * 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. * 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. * Should "Figure 2" on line 234 be "Table 2"?

Rating

4: Borderline reject: Technically solid paper where reasons to reject, e.g., limited evaluation, outweigh reasons to accept, e.g., good evaluation. Please use sparingly.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

2 fair

Presentation

1 poor

Contribution

2 fair

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. * 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. * 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.

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.

Strengths

- The proposed method seems to work better than DT and is competitive with offline RL methods. - The experimental results are quite thorough with multiple domains and relevant comparisons. - EDT also seems especially strong in the multi-task setting, where Q-learning-based offline RL methods typically struggle. - 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. - 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. - In the introduction, the messaging could be a bit clearer by mentioning that this observation applies to DT / sequence modeling approaches.

Questions

- Are the plots in Fig. 7 illustrating the best history length according to the value estimator? - 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? - 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. - What does "-ONE" indicate? Nits: - Could error bars be provided for the Atari results? - Line 212: casual -> causal ================= POST-REBUTTAL RESPONSE =================== Thank you for the detailed response. I still am in favor of acceptance.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

4 excellent

Contribution

3 good

Limitations

Yes

Response to authors

Thank you for the detailed response. I still am in favor of acceptance.

Thanks for your response

I thank the authors for their thorough responses. My questions and concerns are addressed and I'd like to keep the score. The additional experiment results confirm that the search process during inference does not consume much more computation time. However, it would still be great if the Transformer can learn to determine the optimal history length (for example, learning which states to mask to generate a better trajectory). I believe this is an interesting future direction and is not necessary for this paper. So EDT indeed trains on trajectories of different lengths. Reading the relevant sections again, I did find one sentence talking about this in Line 164-165, “This estimation aids in comparing expected returns of different trajectories over various history lengths.” It would be great to make this clearer, possibly by describing the training algorithm using pseudocode.