Reinforcement Learning Gradients as Vitamin for Online Finetuning Decision Transformers

Thanks for appreciating our work. Here are our responses:

Q1. Results are noisy, and a more careful approach to statistical significance is needed.

Thanks for pointing this out! We understand the importance of reliable evaluation and statistical significance. For this we reported standard deviations in our learning curves in all main results (Fig. 2,3,5) and most ablations (Fig. 6-7, 9, 16-17, 19). We provided Tab. 1 so that readers can focus on the final performance and its gain over pretraining.

To make our reported results more reliable, according to the reviewer’s advice, we evaluated all main results (MuJoCo, adroit, maze, antmaze) using the rliable library, as suggested by the paper mentioned in the review. Results are provided in Fig. 3 of the global pdf, which shows: our method indeed outperforms all baselines on adroit, MuJoCo, and umaze/medium antmaze. Also, we will update Tab. 1-3, 10-11 by highlighting scores >=90% of the best reward with bold-face following MAHALO [1].

Q2. Number of random seeds.

Thanks for pointing this out. We use 3 seeds for all experiments (a few ablations without standard deviation use only 1 seed). Note, even without ablations this amounts to a large number of experiments: 1) we use 46 different datasets for our main results [4 for maze, 24 for MuJoCo ({random, medium-replay, medium}x{hopper, halfcheetah, walker2d, ant}x{normal, delayed reward}), 12 for adroit ({pen, hammer, door, relocate}x{expert, cloned, human}), and 6 for antmaze]; and 2) we test 6 methods. This amounts to 800+ runs (46 datasets, 6 methods, 3 seeds), each of which requires 6-8h on average (IQL and TD3+BC need less time, but PDT and other methods on complex environments such as adroit require more).

Q3. Try adding forgetting mitigation methods onto TD3 baseline.

Great suggestion! The paper mentioned by the reviewer is pretty recent. It discusses 4 types of forgetting mitigation methods: Behavioral Cloning (BC), Kick Starting (KS), Episodic Memory (EM), and Elastic Weight Consolidation (EWC). Our TD3 baseline and the MLP TD3+BC baseline already use EM as a mitigation method (quote, “simply keeping the examples from the pre-trained task in the replay buffer when training on the new task”).

Note, since TD3 is deterministic without log-likelihood, adapting EWC is out of scope because it requires the Fisher information matrix, which in turn needs the gradient of the log-likelihood. This is also true for BC and KS, which are auxiliary losses of KL-divergences between policies on states drawn from different samplers. However, since BC/KS are essentially soft constraints that prevent the current policy from being too different from the pretrained policy, we can implement BC/KS on continuous action space for deterministic policies by substituting KL divergence with Euclidean distance between actions predicted by the current and pretrained policies. Here, we add a $0.05\cdot\|\|a_{\text{new}}-a_{\text{old}}\|\|_2^2$ penalty term to the actor loss, which is a gradual shift from BC (sampled from offline dataset) to KS (sampled from online rollouts).

To see whether other forgetting mitigation methods work, we also implemented jump-start RL [2]. Concretely, the pretrained policy is used for the first $n$ steps in an episode, before ODT takes over. We set $n=100$ (100% max episode length for pen, 50% max episode length for other adroit environments) initially, and apply an exponential decay rate of $0.99$ for every episode.

We evaluate both methods on adroit environments with the expert dataset (where the forgetting issue is serious) and the cloned dataset (where the agent needs to improve from low-return policies). Results are summarized in Fig. 1. We find that 1) jump-start RL struggles (possibly because it does not directly prevent out-of-distribution updates), and 2) while BC/KS effectively improves both TD3 and TD3+ODT on expert datasets, it also hinders online policy improvement on most cloned datasets.

Q4. Try other baselines such as TD3+BC+transformer and TD3+RvS.

We would like to point out that our work is focused on analyzing and improving ODT. Thus the main baseline for this work is ODT instead of TD3+BC. That being said, we agree that TD3+BC+transformer and TD3+RvS are interesting baselines to test. The results are shown in Fig. 2 in the global pdf. We test the methods on adroit environments with the cloned dataset. The results show that TD3+RvS slightly improves upon TD3+BC, while TD3+BC+transformer improves more. We speculate that an MLP is not expressive enough to model the policy change over different RTGs. Note, both suggestions still fall short of our method.

Q5. Where does the problem of value learning come from?

Good question! Generally, we found that there is a tradeoff between more information and training stability: while longer context length allows the decision transformer to utilize more information, it also causes increased input complexity, i.e., noise from the environment leads to reduced stability and slower convergence in training. Note, this differs from LLM works, where context length improves generalization thanks to extremely high expressivity, huge amounts of data, and less noisy evaluation. This is especially the case with bootstrapping involved, as the fitting target depends on its own output. In fact, Parisotto et al. [4] suggest that if a recurrent agent can first learn a Markovian policy and value function, then the following training could be stabilized.

Q6. Some limitations are left out, e.g., lack of image-based testbed and increased computational cost.

Thanks for pointing these out! Those are important questions, and we will discuss those in a revised limitation section. Regarding computational cost, we discussed the impact of the RL gradient on training in Appendix H: we found that our method is only marginally (20%) slower than ODT. For the MuJoCo experiment, ODT needs ~6 hours to train on our machine.

Thanks for appreciating our work. Below are responses to questions:

Q1. TD3 gradient for policy update is not clear.

Thanks for pointing this out! There is a typo in Eq. (2) which should read as follows:

$\min\_{\mu^{\text{DT}}}\mathbb{E}\_{\tau\sim D}\Bigg[\frac{1}{T\_{\text{train}}}\sum\_{t=1}^{T\_{\text{train}}}\left[-\alpha Q\_{\phi\_1}(s\_t,\mu^{\text{DT}}(s\_{0:t},a\_{0:t-1},\text{RTG}\_{0:t},\text{RTG}=\text{RTG}\_{\text{real}})) + \|\|\mu^{\text{DT}}(s\_{0:t},a\_{0:t-1},\text{RTG}\_{0:t},\text{RTG}=\text{RTG}\_{\text{real}})-a\_t\|\|\_2^2\right]\Bigg],$

i.e., the $a_t$ in Q should be the action generated by the decision transformer.

Q2. Some notations are confusing.

We introduced RTG in line 82-83 (“RTG, the target total return”), RTG_eval in line 88 (“a desired return RTG_eval is specified”) and line 96, and RTG_real in line 100 (the real return-to-go). We understand that they are not introduced concurrently, which makes it a bit confusing. We will modify this in the revised version.

Q3. How does DT prioritize trajectories with higher return as mentioned in point (2) in line 182?

DT prioritizes trajectories because the actions are learned to be generated conditioned on the Return-To-Go (RTG), and we use high RTG as the condition in evaluation. Consider as an example training with two sets of trajectories extending from the same state: one set with ground truth RTG being 0 and the other set with ground truth RTG being 1. If inference asks DT to generate trajectories conditioned on RTG=1, then the generated trajectory will be more similar to the latter set of training trajectories.

Q4. The limitation that the use of transformers requires more computational cost is not mentioned.

Thanks for pointing this out. We are happy to include a discussion in the limitation section in the revised version. Meanwhile, we would like to mention two points:

1. Our work aims to analyze and improve ODT. We thus compare computational cost primarily to ODT instead of MLP-based solutions;

2. In Appendix H we stated that our proposed solution is only marginally (20%) slower than ODT: while the use of RL gradients slows down training, the actor training overhead is negligible (since it only contains an MLP critic inference to get the Q-value), and the critic only takes about 20% time to train. We found that on our machine, for the MuJoCo experiment, ODT requires about 6 hours to train.

Thanks for appreciating our work. Below are responses to questions:

Q1. test ODT baselines that tackle online finetuning with alternative methods, e.g., better exploration policy and gradually increasing target RTG.

Great suggestions! Currently, ODT explores due to the stochasticity of its policy. We enhance ODT in two ways and provide results in the global pdf: 1) ODT+gradually increasing target RTG (which is similar to curriculum learning in RL); and 2) a cheating baseline ODT with oracle exploration, i.e., jump-start RL [2], using as the guide policy an expert policy trained via IQL on expert data. Concretely, the expert policy is used for the first $n$ steps in an episode, before ODT takes over. We set $n=100$ (100% max episode length for adroit pen, 50% max episode length for other adroit environments) initially, and apply an exponential decay rate of $0.99$ for every episode.

Results are summarized in Fig. 1 of the global pdf. Curriculum RTG does not work, probably because the task is too hard and cannot be improved by random exploration without gradient guidance. Also, even with oracle exploration, ODT is not guaranteed to succeed: it fails on the hammer environment where TD3+ODT succeeds, probably because of insufficient expert-level data and an inability to improve with random exploration.

Q2. TD3 slows down training.

Our method is only 20% slower than ODT without TD3 gradients. In Appendix H, we discussed the impact of the introduced RL gradient on training: the critic only takes about 20% of the time to train, while the actor overhead is negligible (since it only contains an MLP critic inference to get the Q-value). We found that on our machine, for the MuJoCo experiment, ODT requires about 6 hours to train. We will add part of the discussion in Appendix H to the limitation section of the revised paper.

Q3. Are the hyperparameters tuned on the random datasets?

No, as suggested in Tab. 9 in the appendix, we use the same hyperparameters for every variant (e.g., medium / medium-replay / random, expert / human / cloned) on the same environment. Note, our parameters on the random datasets are aligned with the ODT medium environment and not tuned further.

Thanks for appreciating our work. Below are responses to questions:

Q1. Overclaim of “boosting performance”, and the need to refine claims with properties where TD3+ODT should outperform ODT.

Thanks for pointing this out. While TD3+ODT is generally better than ODT, we are happy to rephrase our claim to reflect that there are some cases where ODT performs better than TD3+ODT. To be specific, we expect ODT to struggle (either fail completely or fail to further improve) with medium-to-low quality offline data. This is supported by our result on adroit environments with cloned/human dataset, antmaze, and MuJoCo (especially medium-replay and random where there are many trajectories with low ground truth RTGs). In contrast, if offline data has good quality (e.g., adroit expert), or the reward signal is sparse (e.g., delayed reward in Appendix G.1) so that RL struggles, then we expect ODT to work as good or better.

Q2. [minor] Figure in Section 3.1 could be improved for better understanding.

Thanks a lot for the advice! We will modify the figure accordingly in the revised version.

Q3. Ablations on DDPG+ODT.

Great suggestion! Unfortunately, due to the rebuttal time limit, we can only provide ODT+DDPG results on some environments for now. We will provide results of ODT+DDPG on all additional environments in the revised version. The current results are shown in Fig. 2 of the global pdf. We find DDPG+ODT does not work. We speculate that this is due to the rougher landscape of the estimated Q-value without smoothing, delayed update, and double Q-learning by TD3.

Q4. Impact of context length on the approach.

We conducted ablations on both the context length at training time ($T_{\text{train}}$) and evaluation / rollout time ($T_{\text{eval}}$). The former is shown in Fig. 18 of Appendix G.3, and the latter is shown in Fig. 4 (b). Generally, we found that there is a tradeoff between more information and training stability: while longer context length allows the decision transformer to utilize more information, it also causes increased input complexity, i.e., noise from the environment leads to reduced stability and slower convergence in training. Note, this differs from LLM works, where context length improves generalization thanks to extremely high expressivity, huge amounts of data, and less noisy evaluation. This is especially the case when bootstrapping is involved, as the fitting target depends on its own output. In fact, in prior work, Parisotto et al. [4] suggest in their Sec. 3.1 that if a recurrent agent can first learn a Markovian policy and value function at the start of its training, then the training process could be stabilized.

Q5. Number of random seeds.

We use 3 seeds for all experiments (a few ablations without standard deviation use only 1 seed). Note, even without ablations this amounts to a large number of experiments: 1) we use 46 different datasets for our main results [4 for maze, 24 for mujoco ({random, medium-replay, medium}x{hopper, halfcheetah, walker2d, ant}x{normal, delayed reward}), 12 for adroit ({pen, hammer, door, relocate}x{expert, cloned, human}), and 6 for antmaze]; and 2) we test 6 methods. This amounts to 800+ runs (46 datasets, 6 methods, 3 seeds), each of which requires 6-8h on average (IQL and TD3+BC need less time, but PDT and other methods on complex environments such as adroit require more).