Unleashing the Power of Pre-trained Language Models for Offline Reinforcement Learning

We thank the reviewer for the constructive comments and suggestions. We address each of your comments in the following.

Weaknesses:

Q1: In section 5.5 Ablations, while empirical results indicate the superiority of the former, the absence of a deeper analysis of the choice to use MLPs warrants consideration. It is advisable to provide further insight into the theoretical basis and motivations for this decision.

A1: In Section 5.5, we propose that “We find that simple linear projections could not fully exploit the cross-domain knowledge from language pre-training, and thus our design to replace linear projections with MLPs is critical.” As a supplement, we hypothesize that this is because the stronger representation power prevents inductive bias in pre-trained Transformers from being lost during the fine-tuning process. On the other hand, LoRA restricts the model’s learning ability, thus deeper embeddings are required. Actually, this improvement is motivated by our initial explorations: in experiments, we found that w. MLP is much more effective than w.o. MLP on MuJoCo tasks when the Transformer layer is frozen. An overall ablation study on this component further strengthens our statement, shown in Appendix G, Fig 12.

Questions:

Q2: In section 4.2 you mentioned that you used LORA to inject low-rank matrices into attention weights Q, K and V only and freeze all other weights inside the Transformer, given that there are other large MLPs inside it, what is the rationale of only applying LoRA to Q, K and V?

A2: We inject low-rank matrices into Q, K, V for mainly two reasons. The first reason is that Q,K, and V matrices will interact with the embeddings directly in the process of calculating Attention, and since we have changed the embeddings, adaptations of these weights are intuitively required. The second reason is that as is mentioned in the original LoRA paper [1] page 5, line 8, for simplicity and parameter-efficiency, we can only adapt attention weights and get good results. We follow their method.

Q3: In section 5.5, the benchmark tasks you used for comparison change sometimes, I’m curious how you select ablation benchmarks for showing different components in your framework works better?

A3: We have enhanced the ablations during the rebuttal phase and we hope that these full and consistent ablation experiments in Appendix G could show the consistent advantages of our proposed techniques in leveraging the pre-trained language model for offline RL tasks.

Q4: It would be nice to see how different scales of GPT-2 could affect the performance on your benchmarks.

A4: Yes, we agree with you that this question is worth exploring. Due to the limited computing resources and the fact that GPT2-small has demonstrated the effectiveness of this approach, we only compare the effects of GPT2-medium on some tasks. The experimental results show that in most cases, there is no significant difference between GPT2-medium and GPT2-small, while in some tasks (Kitchen Complete (30%) and Kitchen Partial (100%)), GPT2-medium slightly exceeds GPT2-small, shown in Appendix F, Fig 10. We intend to further explore this aspect in our future research projects.

[1] Hu, et al. “LoRA: Low-Rank Adaptation of Large Language Models” arXiv preprint arXiv:2106.09685

Thanks for the response. I think that generally addresses my concerns and I would keep my score.

We thank the reviewer for the constructive comments and suggestions. We address each of your comments in the following.

Weaknesses:

Q1: Since the paper only studied one LLM (GPT-2) and one RL model (DT), I wonder if the same methodology generalizes to other LLMs and RL models. Since LLMs are improving quickly, it's important that the conclusion in the paper also holds for newer and stronger LLMs.

A1: Yes, we agree with that. As explained in Appendix I.1, we have explored stronger language models, such as the application of LLaMA in the DT framework. However, as far as we know, there has been no successful application of LLaMA to DT, and due to the limitation of computing resources, we have not been able to achieve systematic and significant results. Most of the methods mentioned in Related Works “Transformers for decision making” are worthy of applying our methods, which are beyond the goal of this work.

Q2: The baselines used for comparison seem a bit outdated. There have been much stronger baselines in offline RL in recent years, especially diffusion-based methods such as Diffusion-QL and DD. The authors should include these as stronger baselines.

A2: Thank you very much for your constructive opinion. In recent work focused on Decision Transformer [1][2], we find them not using diffusion policy as the baseline, which may be due to the large gap compared with DT. However, as recent strong baselines, we think they are worth comparing. We have tried it out on our benchmarks, shown in Appendix F, Table 12, prominently showing the power of our method under poor-data regime.

Q3: The experiments only include the medium datasets in D4RL and are missing the medium-replay and medium-expert datasets. I expect LaMO to perform well on medium-replay as these datasets are of low-quality and should highlight the advantage of language pretraining.

A3: We have conducted experiments on Medium-Replay and Medium-Expert datasets, shown in Appendix F, Table 8,9, the results solidate the advantages of our method. As mentioned in Appendix B, we explored both Complete and Partial datasets in the Kitchen environment, while we focused on Medium datasets in the MuJoCo and Atari environments. We emphasize again: Medium data sets are mainly used in our work because they meet two properties that we want: 1. low-quality 2. fair and robust for sampling in poor data regime, as we stated in Appendix B, "The utilization of this dataset is aimed at minimizing variations in quality among different trajectories".

Q4: The first paragraph of Section 4.1 is a bit confusing, it made me think that the authors pretrained GPT-2 themselves. It should be stated clearly that the authors used the pretrained GPT-2 from HF and only pretrained its special variants.

A4: Thank you for pointing out that our statements are a bit confusing, and we have rephrased it in our updated paper.

Questions:

Q5: Can we replace the language loss with a regularization loss so that the fine-tuned model is not too far away from the original pre-trained model? For example, we can minimize the discrepancy between the embeddings of the fine-tuned and pre-trained models.

A5: In this work we actually use language loss as a regularization term, which is particularly effective in certain tasks, as we stated in Section 5.5. Minimizing embedding discrepancy is interesting and reasonable, but we think it cannot be used in our model in a straightforward way: We are using two distinct sets of embedding for language and motion control respectively. As a matter of the fact, the method of reducing the discrepancy of two embeddings has been tried by Wiki-RL, where they utilize K-means clustering and similarity loss, which turns out not to be very helpful, as stated in their paper [3] Section 5.6 Table 6.

[1] Wu, et al. “Elastic Decision Transformer” arXiv preprint arXiv:2307.02484

[2] Yamagata, et al. “Q-learning Decision Transformer: Leveraging Dynamic Programming for Conditional Sequence Modeling in Offline RL” arXiv preprint arXiv:2209.03993

[3] Reid, et al. “Can Wikipedia Help Offline Reinforcement Learning?” arXiv preprint arXiv:2201.12122

We thank the reviewer for the constructive comments and suggestions. We address each of your comments in the following.

Weaknesses:

Q1: Figure 1, in Reid et al paper, they show that a pre-trained DT tends to give improved performance, why in your figure DT gets better performance in Kitchen and Atari?

A1: We strongly agree with the point that “pre-trained DT tends to give improved performance”, while we would like to add that techniques proposed in our work are actually required to “unleash” the power of pre-trained DT. In MuJoCo tasks, Wiki-RL gains higher scores in each task compared with DT, which is both examined by [1] and our work. However, their work has not tested Wiki-RL in the Kitchen environment, which is harder than MuJoCo. In the Atari environment, the dataset they adopt (1% DQN-replay dataset) is different from ours (100%, 10% D4RL Medium dataset). After running extensive experiments, we find that Wiki-RL could not prominently leverage the power of pre-trained DT in these two difficult settings although they obtain advantages in MuJoCo tasks. We propose this is due to the fact that fully fine-tuning the pre-trained model could compromise the general knowledge of pre-trained models, and we find that full fine-tuning could not work at all in certain tasks. This phenomenon is presented in the ablation part Appendix G, Fig 13.

Q2: Did you finetune the methods you are comparing to on the benchmarks you are studying?

A2: Yes, we have looked into the hyperparameters configuration of baselines on our benchmarks, which investigate the poor data regime of D4RL. Specifically, we run the baseline on MuJoCo and Kitchen environments based on CORL (https://github.com/tinkoff-ai/CORL) and run the baseline on Atari using d3rlpy (https://github.com/takuseno/d3rlpy). For both DT-based and value-based methods, we tuned parameters on representative tasks from each domain of our benchmark. It is crucial to note that, within the range of settings of hyperparameters explored, the current parameters adopted by our baselines have been identified as not worse than others, and the results are shown in Appendix F, Fig 11.

Q3: The other thing is only some of the tasks in each benchmark are studied, so it is also a bit concerning whether there will still be a big performance gap between proposed method and baseline when other tasks are also tested.

A3: To ensure strong alignment with the existing literature on DT and offline RL research, we selected tasks identical to those used in [1][2][3][4], including MuJoCo domains such as Hopper, Walker2d, Halfcheetah, Reacher, and Atari domains like Breakout, Qbert, and Pong. Additionally, we included a challenging domain, Kitchen, where our approach demonstrated promising results. It is crucial to emphasize that our choice of tasks is not due to any limitations in applicability to other tasks; rather, we specifically focused on tasks that are most commonly encountered in this research direction, ensuring maximum comparability. To solidify this point, we conduct additional tests on tasks Freeway, Asterix in Atari and Ant in MuJoCo, as shown in Appendix F, Table 10,11, which further prove the claims made in our Experiments part.

Q4: Technical novelty of the method is a bit lacking. Essentially compared to DT+Wiki, a different finetuning method is used, and the projection and output layers are made bigger. Neither of these are new techniques.

A4:

We acknowledge that the implementations we utilized, such as LoRA, MLP embedding, and language loss, are neither novel nor overly complex. However, we would like to underscore that the novelty of our work does not solely lie in any individual module we employed. Instead, our innovation stems from the exploration of the scientific question—whether language pre-training can help offline RL, which is inherently novel and insightful. We successfully demonstrated the potential of this direction by effectively combining simple existing methods, highlighting this as the distinctive aspect of our work.

As articulated in Related Works “LLM for decision making” part, many recent studies focus on using language models as high-level planners or are limited to tasks primarily involving textual input. In contrast, our approach investigates the possibility of efficiently leveraging the inductive bias of language models for low-level motion control without expanding model parameters. We provide an affirmative answer to this question. Furthermore, devising and correctly combining these implementations is non-trivial, as evidenced by our ablation experiments. Incorrectly using the pre-trained model resulted in suboptimal performance. Our framework, which combines all these implementations and adeptly adapts the pre-trained model, unleashes the potential of pre-trained language models for offline RL.

We thank the reviewer for the constructive comments and suggestions. We address each of your comments in the following.

Weaknesses:

Major:

Q1: I have a major concern over the validity of the experimental protocol and results presented in the paper from the observation that the performance reported for the baselines differs significantly from their performance as reported in their original papers. In particular，

Q1.1: BC and CQL in the D4RL paper obtain scores significantly divergent from the scores reported in this paper. Can the authors explain this divergence, and if due to different training budgets/settings explain the motivation behind the different settings? E.g. CQL and BC achieve 43.8 and 33.8, respectively, on Franka-complete in D4RL, but are reported to achieve 0 scores in this paper. (Overfitting of CQL is mentioned, but does not clarify the discrepancy)

A1.1:

Yes, the difference arises from the reported metric. Some offline RL papers use the performance of the best several checkpoints, while for more scientific study, we use “the average performance over the last 20K training steps out of a total of 100K training steps with evaluations conducted every 2500 training steps” as stated in Section 5.1. This is because in the context of offline RL, the RL agent could never interact with the environment during training， and thus mitigating overfitting is a crucial aspect of a method's effectiveness. This choice aligns with our perspective on the importance of robustness in the evaluation of offline RL methods.

Therefore, when an algorithm exhibits overfitting on a certain task (shown in Appendix F, Fig 9) after an excessive number of gradient steps, differences in metrics can result in discrepancies between our results and the values reported in the D4RL paper. We are also willing to report the results on the metric of top performance, as shown in Appendix F, Table 13, which shows the consistent advantages of our method and solidates the validity of our claims.

Q1.2: Also BC and CQL are reported to perform better on Kitchen with 1% of the datasets (or other fractions) than on Kitchen with full datasets. Same for Wiki-RL and DT on some specific splits.

A1.2:

Yes, this is reasonable because of the property of the Kitchen dataset: the Partial dataset contains data that cannot complete the task and data that can, as described in Appendix B. Therefore, when the dataset is being downsampled, the number of varieties of data would decrease and thus could be easier for the model to learn. This explains why algorithms could perform better with less data in certain cases.

Besides, as mentioned in A1.1, there also exists the issue of overfitting. In cases of overfitting, comparing the weak performance lacks statistical meaning.

Q1.3: Similarly, Wiki-RL is reported to perform worse than the same model (ChibiT) in its original paper on Atari. Whereas DT is reported to perform better than in the wiki-RL paper. E.g. in Breakout 350 DT and 130 Wiki-RL in this paper vs 267 DT and 280 Wiki-RL in the Wiki-RL paper.

A1.3: The scores reported in Wiki-RL are different from those in our paper because we use different datasets. Wiki-RL runs its experiments on the 1% DQN-replay Atari dataset, while we run our experiments on the D4RL-Atari medium dataset, of different data ratios.

Q2: The ablation studies are presented each time on a different task & dataset-ratio combination and the full ablation results are not in the appendix (only the pre-training ablation is given). This can lead to a biased evaluation of LaMo if only good results are presented. The combination of components of LaMo is a major contribution of the paper, claimed to "unleash" the potential of pre-trained LMs, and ablation experiments are thus crucial to assess the validity of this claim and the significance of the paper.

A2: We enhance the ablations during the rebuttal phase and we hope that these full and consistent ablation experiments could show the consistent advantages of our proposed techniques in leveraging the pre-trained language model for offline RL tasks, as shown in Appendix G.

I thank the authors for their complete answers to all my concerns. Most of them have been resolved, improving the presentation of the paper and increasing my confidence that the experiments support the claims made in the paper (although I cannot validate the numbers themselves without running the code. In the future, I would strongly encourage the authors to share their code or maybe share tracking results from an experiment tracking service like Weights & Biases).

I have one concern remaining regarding the size of the models compared which I'd like the authors to address.

In their reply, the authors say that (A12:) "The numbers of parameters of DT, Wiki-RL, and our method are 7.3M, 125M, and 128M, respectively. And for non-transformers (value-based) methods, the numbers are relatively small, such as 0.41M for CQL."
Revealing a large gap in the number of parameters of the models compared, even within the same family of transformers DT vs LaMo. However, the authors also point out that the comparison should not be made on the basis of the total number of parameters but on the basis of the number of trainable parameters, which in that case would be 3.5M for LaMo and comparable to DT.

I appreciate this distinction, but believe that it could benefit from an explicit presentation in the paper, so that future readers do not extrapolate or misunderstand from the paper that using LaMo would result in better performance than training a DT model of the same size. It would also be beneficial to state the number of (trainable) parameters of all models and baselines in the appendix.

Regarding value-based models I believe that the following claim would need a reference:
"For value-based methods, when the model parameters are relatively large, training becomes extremely challenging, leading to results worse than smaller models. Therefore, for value-based baselines, we adhere to widely accepted model sizes, which are much smaller than Transformer-based methods."

With both of these resolved (making it clear that it's trainable sizes that are compared, listing them, and justifying the size of the value-based ones), I would be happy to significantly increase my score.

Minor concern. A10: "To determine the hyperparameter lora-dim that governs the parameter scale for LoRA, we explored values of 16, 32, and 64, ultimately choosing 32 for Atari, and 16 for MuJoCo, Kitchen."
Can this be reflected in table 7?