Double Check My Desired Return: Transformer with Value Validation for Offline RL

Authors: Yue Pei, Hongming Zhang, Chao Gao, Martin Müller, Mengxiao Zhu, Hao Sheng, Haogang Zhu
Primary area: reinforcement learning
Keywords: Offline Reinforcement Learning, Transformer

Abstract:
Recently, there has been increasing interest in applying Transformers to offline reinforcement learning (RL). Existing methods typically frame offline RL as a sequence modeling problem and learn actions via Supervised learning (RvS). However, RvS-trained Transformers struggle to align actual returns with desired target returns, especially when dealing with underrepresented returns in the dataset (interpolation) or missed higher returns that could be achieved by stitching sub-optimal trajectories (extrapolation). In this work, we propose a novel method that Double Checks the Transformer with value validation for Offline RL (Doctor). Doctor integrates the strengths of supervised learning (SL) and temporal difference (TD) learning by jointly optimizing the action prediction and value function. SL stabilizes the prediction of actions conditioned on target returns, while TD learning adds stitching capability to the Transformer. During inference, we introduce a double-check mechanism. We sample actions around desired target returns and validate them with value functions. This mechanism ensures better alignment between the predicted action and the desired target return and is beneficial for further online exploration and fine-tuning. We evaluate Doctor on the D4RL benchmark in both offline and offline-to-online settings, demonstrating that Doctor does much better in return alignment, either within the dataset or beyond the dataset. Furthermore, Doctor performs on par with or outperforms existing RvS-based and TD-based offline RL methods on the final performance.

Decision: reject
Reviewer scores: 6, 5, 3, 6

Reviews and discussion:
<thread>
creation date:
1729326892253

thread_type:
official_review

<thread_object>
creator: reviewer_pdqo

<content>
summary:
RvS-trained Transformers lacks 1) the capability to interpolate between underrepresented returns in the dataset, and 2) to stitch information from multiple sub-optimal trajectories into a better one, due to its supervised nature.
The paper introduces Doctor, which aligns the return of the policy with the target return, by selecting the action that has the nearest Q function (learned via expectile regression, similar to IQL) to the target returns.
Doctor shows much better return alignment across various target returns, and pars with / outperforms other RvS methods in terms of performance.

strengths:
1. Doctor shows much better return alignment compared to previous RvS methods.

* Experimental results in Figure 1, Figure 3 shows that Doctor shows much better alignment.
* This alignment can extrapolate for some datasets (i.e. being better than the dataset).

2. The algorithm of Doctor is simple and intuitive.

* The alignment happens in the inference time, and the algorithm is straight-forward (choosing the action with the closest returns)
* Combines SL and IQL-like loss, which does not need to query out-of-distribution actions.

weaknesses:
1. The experimental results is insuffcient on showing the stitching capability of Doctor.
* While stitching is one of the main reason of integrating TD-learning to SL, the performance improvement compared to MTM (which is a pure RvS method) is incremental (708.6 vs 719.7), according to Table 1.
* Generally, D4RL MuJoCo Gym environments is not a good benchmark to evaluate stitching capabilities. Experimental results in D4RL Maze2D tasks or AntMaze tasks, which are designed for stitching, will be helpful on evaluating the stitching capabilities of Doctor.

2. The experimental results is insufficient on showing that "better alignment in returns help offline-to-online fine-tuning".
* Experiment results in Table 2 shows that the performance increase during online fine-tuning is larger for ODT. It looks counterintuitive, considering superior alignment capability of Doctor.
* Again, D4RL MuJoCo Gym environments might be not a good choice to evaluate offline-to-online fine-tuning. For example, Cal-QL[1] uses AntMaze, Kitchen, Adroit tasks for their experiments; experimental results on these environments might better highlight the superiority of Doctor.

[1] Mitsuhiko Nakamoto et al., "Cal-QL: Calibrated Offline RL Pre-Training for Efficient Online Fine-Tuning", NeurIPS 2023.

questions:
Q1. Does having better return alignment helps the final performance?
* i.e., can the problem of return alignment be solved by just having high target returns? (e.g. DT[2] uses 5x maximum returns in the dataset)
* It will be insightful if there are experimental results for ablating the inference-time alignment (with properly tuned target returns, e.g. 1x or 5x maximum return of dataset).

Q2. Does having better return alignment helps online fine-tuning?
* It will be exciting if Doctor is able to show clear improvements in online fine-tuning experiments on AntMaze or Kitchen or Adroit tasks, where offline-to-online methods (e.g. Cal-QL) have showed significant improvements.

Q3. Other than Q1 (final performance) and Q2 (online fine-tuning), can you share your ideas on why having better return alignment is beneficial, if there is any?
* e.g., does it help reducing the effort of searching for the target return used in inference?

Q4. Does Doctor improves stitching capability?
* Experimental results in D4RL Maze2D tasks or AntMaze tasks, which are specifically designed for stitching, would be informative to assess the stitching capability of Doctor.

Q5. What target return is used for evaluation?
* Can you share how did you find it? (e.g. is it the maximum return from the dataset?)


[2] Lili Chen et al., Decision Transformer: Reinforcement Learning via Sequence Modeling, NeurIPS 2021.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730187605083

thread_type:
official_review

<thread_object>
creator: reviewer_p2gx

<content>
summary:
This paper introduces Double Checks the Transformer with Value Validation for Offline Reinforcement Learning (Doctor), a framework designed to address alignment issues between actual returns and desired target returns, especially in cases involving underrepresented or higher returns. Specifically, this approach integrates additional temporal-difference (TD) learning into the Transformer architecture, sampling actions around desired target returns and validating them through value functions. Experiments conducted in both offline and offline-to-online settings demonstrate improved alignment with target returns, showcasing the effectiveness of the proposed approach.

strengths:
1. The paper is well-written and easy-to-read.
2. The motivation is clearly presented, focusing on addressing alignment issues within sequence modeling methods.

weaknesses:
1. The paper lacks a comparison with state-of-the-art methods, such as RADT[1], which also addresses target return alignment, and QT[2], which enhances stitching capability by incorporating the TD method.
2. The performance improvements in D4RL are not significantly better than these baselines, diminishing the perceived effectiveness of the proposed framework. Although this work integrates the TD learning component into the MTM network, the improvement over MTM is marginal. Furthermore, during online fine-tuning, the performance gains are also less pronounced compared to the ODT baseline.

[1] Tsunehiko Tanaka, et al. Return-Aligned Decision Transformer. 2024

[2] Shengchao Hu, et al. Q-value Regularized Transformer for Offline Reinforcement Learning. 2024 ICML.

questions:
1. Does interpolating between underrepresented returns and stitching information address the same issue? While stitching sub-optimal trajectories can yield better performance, it could also lead to undesirable outcomes not observed in the dataset, representing underrepresented returns. Thus, should the primary focus be on improving the stitching capability?
2. In line 167, $R_t$ is defined as the discounted return, while Decision Transformer (DT) treats it as the undiscounted sum of rewards. What is the rationale behind using the discounted return in this context?
3. In lines 185–186, how does the action-value $q_t$ provide the model with the ability to stitch sub-optimal trajectories?
4. Why was a bi-directional Transformer chosen, given that most DT-based approaches utilize a causal Transformer to predict actions in an auto-regressive manner? This causal structure is typically important during inference to predict future actions based on historical sequences.
5. Is the comparison of results in Table 1 equitable? The Doctor method requires sampling N target returns for alignment. Do the baseline methods also allow for such extensive sampling of target returns to construct their final results?
6. What is the time complexity of inference when compared to other baselines?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730475543514

thread_type:
official_review

<thread_object>
creator: reviewer_lbgr

<content>
summary:
This paper proposes to combine supervised learning and temporal difference learning for offline reinforcement learning. First, the model is pretrained by randomly masking a subset of a trajectory and predicting it. In temporal difference learning, the model is trained to predict the action value. Experiments show promising results.

strengths:
1. It looks reasonable to combine supervised learning and temporal difference learning for offline reinforcement learning. 
2. The paper is easy to follow.
3. Experiments show promising results.

weaknesses:
1. I am not sure about the novelty as self-supervised learning and value function learning are two common techniques.
2. Ablation studies about the design choices of the proposed framework are needed.
3. The performance over the baselines on different tasks seems inconsistent, as shown in Table 1.

questions:
Please refer to the weakness.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731079593615

thread_type:
official_review

<thread_object>
creator: reviewer_r5yx

<content>
summary:
The paper introduces a novel method called "Doctor" that integrates supervised learning (SL) and temporal difference (TD) learning to improve the alignment of predicted actions with desired target returns in offline reinforcement learning (RL). The method uses a bidirectional transformer and a double-check mechanism during inference to validate actions with value functions, ensuring better alignment and improved performance. The work is evaluated on the D4RL benchmark, demonstrating competitive performance compared to existing RvS-based and TD-based methods.

strengths:
1. The core idea of the paper is interesting and well-motivated, addressing the challenge of aligning predicted actions with desired returns in offline RL.
2. The paper is well-written and easy to follow, making it accessible to readers with a background in RL and transformers.
3. The experimental results on the D4RL benchmark show that the proposed method, Doctor, outperforms several baselines in some tasks, indicating its potential effectiveness.

weaknesses:
1. Model Choice (Bidirectional vs. Causal Transformer): The use of a bidirectional transformer, rather than a causal transformer, is not fully justified. Since inference is typically performed in a causal manner, it would be beneficial to provide more insight into why a bidirectional transformer was chosen and how it impacts the model's performance during inference.
2. Limitations Discussion: The paper lacks a detailed discussion of the limitations of the proposed work. It would be valuable to acknowledge potential drawbacks or scenarios where the method might not perform as well, setting realistic expectations for future research.
3. Inference Complexity: The inference complexity of the proposed method may be significantly higher than that of DT. A detailed comparison of the inference speed and computational overhead relative to other models would strengthen the paper.
4. Performance and Overhead: While the performance is competitive, the gains are marginal compared to the baselines. Given the higher computational cost, a more in-depth analysis of the trade-offs between performance and computational resources is necessary.
5. Evaluation Scope: The evaluation is conducted only on the D4RL benchmark, which may introduce a bias in the results. To ensure the robustness and generalizability of the proposed method, it would be valuable to conduct additional experiments on other offline RL benchmarks.
6. Reward Modeling Assumption: if I understand correctly, the underlying assumption is that value prediction is more general/robust than reward modeling for aligning the ground-truth reward during inference, which needs further theoretical justification. An analysis or explanation supporting this claim would add depth to the paper.

questions:
1. Have you test your method in some robotics benchmarks?
2. Can the inference process be accelerated?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732642326160

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
**Global Response**

We sincerely thank all the reviewers for their valuable feedback on our work.

To address common questions and share additional experimental results, we have prepared a global response covering the following key aspects:

**1: Additional experimental results**

To further demonstrate the effectiveness of our method, we have expanded our evaluation to include additional benchmark tasks. Specifically, we have added three Maze2D tasks and three Adroit tasks in the offline setting as well as three Adroit tasks in the offline-to-online setting.

We have also compared our method with additional baselines, including the state-of-the-art method [1]. Furthermore, we provide additional visualizations of the alignment experiment [776-790], along with evaluations of inference time and training overhead [847-856].

All new results are included in the revised paper.

[1] Shengchao Hu, et al. Q-value Regularized Transformer for Offline Reinforcement Learning. 2024 ICML.

**2: Why Doctor acheives on par performance in some tasks?**

We would like to emphasize that our primary goal is to align actual returns with desired target returns, as demonstrated in Figure 3. The improvement in final performance is merely a byproduct of this alignment and not a direct optimization target. From this perspective, it is not surprising that Doctor does not significantly outperform other baselines that focus on policy improvement.

We would like to emphasize our motivation and novelty as follows:

**Motivation**: Previous works focus solely on improving final performance. In contrast, our work aims to extend this to a range of desired returns, presenting a new and more challenging task. Achieving this allows us to extract various levels of performance policies, which is particularly valuable for scenarios like game AI, where NPCs with diverse skill levels are essential for creating balanced and engaging gameplay.

**Novelty**: The core novelty of our work lies in the double-check mechanism, which combines self-supervised learning and value function learning. This enables the model to align predicted values with target returns more accurately while also enhancing its ability to fine-tune through further online exploration. Our method is the first to combine self-supervised learning and value function learning in this manner, and we believe this is a significant contribution to the field.

**3: Revised paper**


We have incorporated the reviewers' suggestions into the revised paper and uploaded it, with all major updates highlighted in blue for your convenience. The key revisions are as follows:

1) We have provided additional justification of value alignment (in appendix A).

2) Our experiments have been extended, comparing with additional baselines across more benchmark tasks (in appendix C).

3) We have included new visualizations for the alignment experiment, along with evaluations of inference time and training overhead (in appendix D.3).

4) More related papers have been cited.

Thank you again for your feedback. We are looking forward to hearing your thoughts.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This work proposes a method called Double Checks the Transformer with value validation for offline RL. This method integrates supervised learning (SL) and temporal difference (TD) learning to improve the alignment of predicted actions with desired target returns in offline reinforcement learning (RL). The work is evaluated on the D4RL benchmark, demonstrating competitive performance compared to existing RvS-based and TD-based methods.

After the rebuttal and discussion, this paper receives review scores of 3, 5, 6, 6. Reviewers raised several concerns:
1) lack of comparisons with recent methods
2) The performance on D4RL is not significantly better than other methods. 
3) The motivation for aligning with the target return is not adequately justified in the current version of the article.
4) a lack of ablation study
5) novelty is limited.

There are discussions between the reviewers and the authors. The reviewers who gave negative comments both replies and weren't fully satisfied with the responses from the authors. AC has checked the submission, the reviews, the rebuttal, and discussion. AC sided with the negative reviewers and believed this work needs further effort to revise and improve the method performance. Thus, a rejection is recommended.