Paper: Unleashing the Power of Pre-trained Language Models for Offline Reinforcement Learning (ICLR-2024)

Reviewer comment (Theory):
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.

Author response:
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**.

Reviewer comment (Experiments):
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?

Author response:
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.

Reviewer comment (Experiments):
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?

Author response:
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.

Reviewer comment (Experiments):
It would be nice to see how different scales of GPT-2 could affect the performance on your benchmarks.

Author response:
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.

Reviewer comment (Experiments):
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.

Author response:
**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.

*However, my main concern regarding the applicability of the framework to other LLMs and RL models still remains.*

Our framework (sequentially pre-training+deeper embedding+LoRA+language loss) is adaptable to transformer-based RL methods such as **[1][2]**. We have explored applying LLaMA to our framework, and we anticipate that our work could motivate the community to delve into further applications of stronger LMs in offline RL.

Again, we sincerely thank you for the constructive suggestions.

[1] Janner, et al. “Offline Reinforcement Learning as One Big Sequence Modeling Problem” arXiv preprint arXiv:2106.02039

[2] Liu & Abbeel. “Emergent Agentic Transformer from Chain of Hindsight Experience” arXiv preprint arXiv:2305.16554

Reviewer comment (Experiments):
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 [1] and DD [2]. The authors should include these as stronger baselines.

Author response:
**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.

[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

Reviewer comment (Writing):
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.

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

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

Author response:
**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**.

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

Reviewer comment (Evaluation):
Figure 1, in Reid et al paper, they show that a pretrained DT tend to give improved performance, why in your figure DT gets better performance in Kitchen and Atari?

Author response:
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**.

Reviewer comment (Experiments):
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 perforrmance gap between proposed method and baseline when other tasks are also tested.

Author response:
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.

Reviewer comment (Novelty):
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.

Author response:
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.

Reviewer comment (Theory):
It is unclear why language pretraining can help RL tasks which have a large domain gap?

Author response:
Your question is indeed insightful. Exploring this point is the main goal and highlight of our work. In **Section 5.5** “Language pre-training v.s. visual pre-training”, we attribute it to “**proficiency in sequential modeling is the key to unleashing the potential of cross-domain pre-trained models**”. We have conducted extensive experiments, replacing the pre-trained model with image pre-trained, early-stopped pre-trained, random corpus pre-trained, and random weights initialized. Empirically, all of those results indicate that: two properties, 1. high-quality pre-trained and 2. pre-trained on sequential modeling tasks, are actually required to prominently unleash the power of pre-trained models in our method. As recent works **[5][6]** proposed,  language model could be experts in motion control if we view motion control as a downstream task in sequential modeling, the same as language modeling, and we emphasize again that, our method can effectively use the prior knowledge acquired in language pre-training to deal with offline RL.

Reviewer comment (Evaluation):
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. For example, in Breakout, DT scores 350 and Wiki-RL scores 130 in this paper, versus DT scoring 267 and Wiki-RL scoring 280 in the Wiki-RL paper.

Author response:
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.

Reviewer comment (Experiments):
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.

Author response:
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**.

Reviewer comment (Evaluation):
I don't find LoRA to be well-motivated, or at least that its benefits are. The authors mention early overfitting but this is not clear from Figure 5 (LoMa also experiences small drops in performance).

Author response:
For the motivations,  1. As we view motion control as a downstream task in sequential modeling, LoRA, as a well-known parameter-efficient fine-tuning technique for adapting downstream tasks, is expected to leverage the pre-trained model well.  2. There are works **[1]** showing that fine-tuning only a small portion of parameters could outperform frozen and full fine-tuning. 

We have provided more ablation results on LoRA compared with full finetuning and freezing, as shown in **Appendix G, Fig 13**.  We observe that in our settings, freezing the pre-trained parameters and adapting with LoRA could not only improve training efficiency but also address the issue of overfitting that occurs in full finetuning. **Notably, our experiments show that full finetuning could not work at all in certain tasks**, and freezing all parameters in Transformer could restrict the learning ability of the mode.

Reviewer comment (Novelty):
Several elements of the paper can be mistakenly interpreted as original contributions of the paper, such as the language modeling loss (also used in Wiki-RL), the ablation on ImageGPT (also used in Wiki-RL), and the selection of the baseline (similar to Wiki-RL). While not being novel was not the issue, it was often not clear if those contributions were original.

Author response:
The language modeling loss term was proposed in **[2]**,  however, its ablation experiments in **Section 5.6 Table 6** of that paper have not prominently shown the improvement brought by this term. Built on their work, we further demonstrate its role in preventing overfitting in certain tasks by ablation study, and show that the pre-trained language model could master the knowledge of two domains: language and motion at the same time. To prevent being mistakenly interpreted, we have added a citation of Wiki-RL in the method part for this implementation.

As for the ablation on imageGPT and the choice of baselines, we did not regard them as the novelty of our work. We actually use different baselines from Wiki-RL, like IQL and BCQ, which are the most commonly used powerful baselines in offline RL, and notably, we also compare with Wiki-RL itself, as a baseline utilizing pre-trained models.

Reviewer comment (Presentation):
It’s hard to assess the significance of an improvement when curves are shown with shaded areas in [μ − 0.5σ, μ + 0.5σ].

Author response:
Thank you for your suggestion. We have provided more ablation results with shaded areas in [μ − σ, μ + σ] in **Appendix G**.

Reviewer comment (Writing):
Notation is not ambiguous for readers familiar with the topics but would benefit from better presentation. Some variables are undefined in some equations, e.g., the expectation distributions are underspecified in equation 1, $a'$ in equation 3, and $T$ in equation 4.

Author response:
Thank you for your advice, we have fixed these undefined notations in equation 1,3,4 in our updated paper.

Reviewer comment (Presentation):
Citations don’t include the journal/proceedings details which makes it hard to identify which version of a paper was used in some cases.

Author response:
We have revised the format of the reference to reflect the journal/proceedings details in our updated paper.

Reviewer comment (Reproducibility):
Are the hyperparameters in Appendix E used with all transformers in the paper? (LoMa, Wiki-RL, DT)?

Author response:
Yes, the context length, return-to-go, training steps, batch size, dropout are shared by all these DT-based methods. And our method utilizes the approximately same model size as Wiki-RL, including the same hidden size, number of layers, except for the MLP embedding part stated in Method. 

Besides, the learning rate and weight decay in Appendix E are not directly used in Wiki-RL and DT as our approach adopts a pre-train + parameter-efficient fine-tune approach, while DT and Wiki-RL both adopt full supervised training. We have tuned DT across different tasks to obtain good convergence performance among explored configurations, as shown in **Appendix F, Fig 11**. As for Wiki-RL, we sincerely follow the authors’ statements in their paper **[2] Section 4.1**, to use the same learning hyperparameters as DT.

Reviewer comment (Experiments):
Do the transformers and non-transformers used in the paper have a comparable number of parameters to ensure a fair comparison of performance?

Author response:
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. 

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. 

It is important to emphasize that simply increasing the size of the Transformer will not boost the performance, as shown in the results of Wiki-RL and our ablation studies. Moreover, although our method involves a relatively large model, the number of trainable parameters is fairly small, which is 3.5M. So the difference of the number of parameters between transformers and value-based methods does not compromise the fairness of the comparisons of performance.

Reviewer comment (Evaluation):
Figure 6(a): why do the curves not start from the same point at training step 0? How can the authors explain that the cross-entropy loss decreases significantly for an already pre-trained model (red curve) and also eventually decreases for the ablation (blue curve)?

Author response:
The curve does not start from the same point because we did not immediately evaluate after initialization. Actually, the starting point reflects the results after training for 2500 training steps. To enhance clarity, we have added a point at the beginning of the curve. 

The continued decrease in loss is attributed to the jointly trained model utilizing a small text dataset, WikiText, which introduces a distribution shift compared to the pretraining dataset of GPT-2, which is a very large corpus. This shift allows the model to further adapt, thus the entropy loss could decrease. 

It is important to note that the phenomenon of loss eventually decreasing in the ablation is not consistent. We provide an additional example that demonstrates a different result, as shown in **Appendix G, Fig 14(a)**.