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

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.

re:

Thank you for your advice. We have conducted further experiments on some extra tasks in Table 3 (see appendix C.1 [703-729] in our revised version). Notably, in Maze2D, Doctor outperforms some other baselines by a clear margin, which emphasize 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.

re:

We have conducted additional experiments on Adroit tasks regarding the evaluation of offline-to-online in Table 4 (see appendix C.2 756-770). The results show that Doctor achieves clear improvements on the Adroit tasks, demonstrating its fine-tuning capability in more challenging environments.

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 uses 5x maximum returns in the dataset)

re:

There is no guarantee that having better return alignment helps the final performance. 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.

We further visualize the effect of value alignment in Figure 5 (see appendix C.3, 777-790), showing 0x to 3x the maximum returns in the dataset. Doctor achieves much better alignment across a range of target returns compared to DT and MTM.

Return alignment cannot be solved by just having high target returns. Doctor can achieve better alignment to some extent outside the dataset, but its alignment ability cannot extend too far beyond the maximum return seen during training. Figure 5 shows that the performance of Doctor saturates when the target return is set too high such as 3x maximum return.

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.

re: Yes, better return alignment contributes to more effective online fine-tuning according to the results in table 4, where doctor achieved clear improvements on Adroit.

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?

re:

Motivation: In contrast to prior work, which focuses on improving final performance, our motivation is to extend this to a range of desired returns. 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. From this perspective, the improvement in final performance is merely a byproduct of this alignment.

Regrading the effort of searching for the target return used in inference, the answer is yes. Previous works cannot generate a trajectory that achieves the desired return, as shown in our experiments in Figure 1 and 3, thus they need to search for the target return. In contrast, Doctor can generate a trajectory that achieves the desired return within the dataset and beyond to some extent, thus reducing the effort of searching for the target return.

Q4. Does Doctor improves stitching capability?

re:

Yes, as illustrated above.

Q5. What target return is used for evaluation?

re:

Yes, our target return is set based on the n times maximum return from the dataset as used by DT. Specifically, it is calculated as: target_return = min_return + (max_return-min_return)

We hope these additional experiments and explanations address your concerns. Thank you for your valuable suggestions.

We appreciate your suggestion to further clarify the motivation. We will incorporate this in the next revision.

Thank you again for your valuable feedback and for increasing the score.

Questions

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.

re:

While early works like DT [1] and TT [2] use causal transformers, recent studies such as MTM [3] have shown that bidirectional transformers trained with a combination of random masking and autoregressive masking enable the model to learn a more generalized sequence representation and achieve better performance. Our experiments in Figure 3 also demonstrate that the BiT-based model MTM outperforms the CT-based model DT. Furthermore, our work based on BiT further improves the performance.

[1] Chen, Lili, 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 (2021): 15084-15097.

[2] Janner, Michael, Qiyang Li, and Sergey Levine. "Offline reinforcement learning as one big sequence modeling problem." Advances in neural information processing systems 34 (2021): 1273-1286.

[3] Wu, Philipp, Arjun Majumdar, Kevin Stone, Yixin Lin, Igor Mordatch, Pieter Abbeel, and Aravind Rajeswaran. "Masked trajectory models for prediction, representation, and control." In International Conference on Machine Learning, pp. 37607-37623. PMLR, 2023.

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?

re:

Although Doctor samples N target returns for alignment, we process them as a batch and feed them into the model, resulting in relatively low computational overhead compared to other methods.

Regarding the baseline methods, they are not suitable for the sampling strategy used by Doctor. This is because these methods typically lack a Q-function. Without a Q-function, even if target returns were sampled, these methods would have no mechanism to determine the optimal actions for the target return alignment.

6. What is the time complexity of inference when compared to other baselines?

We have included a detailed comparison of inference speed and computational overhead with other models in the appendix. The results demonstrate that Doctor maintains efficient during inference compared to other approaches.

We hope these additions and clarifications address your concerns.

Weaknesses:

1. The paper lacks a comparison with state-of-the-art methods.

re:

Thank you for pointing out. We have reproduced QT using the officially released code and incorporated the results of RADT and QT into Table 3 (see appendix C.1 [703-729] in our revised version).

2. The performance improvements in D4RL are not significantly better than these baselines. Furthermore, during online fine->tuning, the performance gains are also less pronounced compared to the ODT baseline.

re:

Firstly, 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.

To further showcase the effectiveness of our method, we have expanded our evaluation to include additional benchmark tasks, specifically Adroit and Maze2D. Notably, Doctor performs competitively across these tasks, demonstrating its robustness and generalizability.

Regarding the online fine-tuning performance, since the initial scores are different for each method, the relative improvement may not be directly comparable. We compare the final performance, and Doctor achieves the best overall performance across tasks. Additionally, when the initial performance is similar, for example on Walker2D Medium in Table 2, Doctor achieves better performance during online fine-tuning than ODT.

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?

re:

Interpolating between underrepresented returns and stitching are two different issues. Interpolating involves learning a robust policy that performs well when the data is sparse, while stitching focuses on combining segments to form better sequences. In our work, we aim to address both issues by leveraging the double-check mechanism, which allows the model to interpolate and extrapolate from the dataset.

Regarding the undesirable outcomes due to stitching, this is precisely where the double-check mechanism comes into play. While supervised learning ensures that the generated actions remain within the distribution of the dataset, TD learning builds upon this foundation to improve extrapolation, thereby reducing the risk of generating unrealistic outcomes.

2. In line 167, 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?

re:

We need to ensure consistency between the learning of the target return in the transformer and the action-value in the Q-function. We can use either discounted or undiscounted returns for both the transformer and the Q-function. We chose to use discounted returns.

3. In lines 185–186, how does the action-value provide the model with the ability to stitch sub-optimal trajectories ?

re:

In our method, the action-value function is trained via implicit Q-learning [1], an in-sample multi-step dynamic programming method that excels at stitching sub-optimal trajectories. We further enhance the stitching capability by incorporating our double-check mechanism, thereby reducing the risk of generating unrealistic outcomes.

[1] Kostrikov, Ilya, Ashvin Nair, and Sergey Levine. "Offline reinforcement learning with implicit q-learning." arXiv preprint arXiv:2110.06169 (2021).

Weaknesses:

1. I am not sure about the novelty as self-supervised learning and value function learning are two common techniques.

re:

We acknowledge that self-supervised learning and value function learning are common techniques. Since all transformers are based on self-supervised learning and all value-based RL methods use value function learning, the novelty lies in the specific problem we are solving and how we solve it.

Problem: 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.

Method: 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.

We hope this clarifies the novelty of our work.

2. Ablation studies about the design choices of the proposed framework are needed.

re:

We explored the impact of different target return sampling sizes N in Figure 4, as N increases, Doctor achieves better alignment with the given target return. We also included the results for different transformer architectures in Figure 3, demonstrating that the BiT-based transformer (MTM) is a better choice than causal transformers (DT).

3. The performance over the baselines on different tasks seems inconsistent, as shown in Table 1.

It is unrealistic to expect Doctor to outperform baselines on every individual task. This is consistent with observations in prior work, where different methods excel in different scenarios due to varying task characteristics. Furthermore, we would like to emphasize that our primary goal is to align actual returns with desired target returns. The improvement in final performance is merely a byproduct of this alignment.

We hope these additions and clarifications address your concerns.

Weaknesses:

1. Model Choice (Bidirectional vs. Causal Transformer): The use of a bidirectional transformer, rather than a causal transformer, is >not fully justified.

re:

While early works like DT [1] and TT [2] use causal transformers, recent studies such as MTM [3] have shown that bidirectional transformers trained with a combination of random masking and autoregressive masking enable the model to learn a more generalized sequence representation and achieve better performance. Our experiments in Figure 3 also demonstrate that the BiT-based model MTM outperforms the CT-based model DT. Furthermore, our work based on BiT further improves the performance.

[1] Chen, Lili, 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 (2021): 15084-15097.

[2] Janner, Michael, Qiyang Li, and Sergey Levine. "Offline reinforcement learning as one big sequence modeling problem." Advances in neural information processing systems 34 (2021): 1273-1286.

[3] Wu, Philipp, Arjun Majumdar, Kevin Stone, Yixin Lin, Igor Mordatch, Pieter Abbeel, and Aravind Rajeswaran. "Masked trajectory models for prediction, representation, and control." In International Conference on Machine Learning, pp. 37607-37623. PMLR, 2023.

2.Limitations Discussion: The paper lacks a detailed discussion of the limitations

re:

One potential limitation of our approach is that we additionally learn action-value heads, which introduces extra computational overhead during inference. We summarize the inference time and computational overhead (see appendix D.3 [848-856] in our revised version).

Another limitation is that since the goal of our work is to align actual returns with desired target returns, the improvement in final performance is only a byproduct of this alignment. The model may not always outperform other baselines that focus on policy improvement for every individual task, as shown in Table 1.

3.Inference Complexity: The inference complexity of the proposed method may be significantly higher than that of DT.

re:

During inference, Doctor processes one batch at a time and utilizes a unified architecture to generate both action predictions and value estimates simultaneously, enabling fully parallel computation. To address the concern, we have included a detailed comparison of inference speed and computational overhead with other models in the appendix (see appendix D.3 [848-856] in our revised version). The results demonstrate that Doctor maintains efficient during inference compared to other approaches.

4.Performance and Overhead: While the performance is competitive, the gains are marginal compared to the baselines.

re:

The computational overhead of our model is indeed higher compaired with MTM due to the inclusion of value function training. However, we have found this increase to be within an acceptable range.

We have also conducted further experiments on some extra tasks (see appendix C.1 [703-729] in our revised version). For instance, in tasks like Maze2D, which emphasize the model's ability to effectively "stitch" together different parts of the sequence, Doctor outperforms some other baselines by a clear margin. Thus, we believe this additional cost is justified by the better alignment and on par or superior performance in a variety of tasks.

5.Evaluation Scope: The evaluation is conducted only on the D4RL benchmark, which may introduce a bias in the results.

re:

We have expanded our evaluation to include additional benchmark tasks, specifically Adroit, Maze2D. These tasks cover a broader range of environments, each emphasizing different aspects of the model’s ability. We found that Doctor's performance is competitive across these tasks, demonstrating its robustness and generalizability. We have included a detailed discussion of these results in the revised manuscript.

6.Reward Modeling Assumption

re:

We would like to clarify that our method does not rely on the 'Reward Modeling Assumption'. The better alignment and more robust performance of Doctor come from the double-check mechanism. Previous works such as DT assume that the return-conditioned learned model can generate a trajectory that achieves the desired return. We demonstrate that this is not the case, as previous works only care about improvement in the final performance and ignore the alignment assumption. In our work, we introduce a Q-function to double-check the alignment with the target returns, serving as a twofold verification (see appendix A in our revised version for more explanation.

Questions:

1.Have you test your method in some robotics benchmarks?

re:

Yes, Gym locomotion tasks are considered robotics benchmarks. To further validate our method, we have conducted additional experiments on maze2d and adroit tasks (see appendix C.1).

2.Can the inference process be accelerated?

re:

Since Doctor processes one batch at a time thus the computational overhead is relatively low as shown in Table. Please refer to the respond to the weakness.