Decomposed Prompt Decision Transformer for Efficient Unseen Task Generalization

Thank you very much for your careful review of our work. We'll answer your questions one by one in the following, including some misunderstandings and some essential academic questions worth exploring.

W1: About the novelty of the paper

Please check Author Rebuttal AR1.

W2: About the prompt decomposition
We visualized the cross-task prompt matrix $P_c$ and the task-specific prompts $p_k$ for 10 tasks on MT10, as shown in Figures 1 and 2 in the PDF of Author Rebuttal. From the visual results, $P_c$ exhibits an irregular data distribution, suggesting weighted representations of common knowledge across tasks. In contrast, we found that some $p_k$ matrices for specific tasks show sparsity (values approaching 0), which may explain the essence of prompt decomposition: during training, $P_c$ extracts general features that are then integrated into task prompts $P_k^*$ through element-wise multiplication with sparse $p_k$.

To demonstrate that $P_c$ is more than just a generic scaling factor, we conducted experiments on three datasets: Cheetah-vel, MW ML45, and MT50. We performed fine-tuning without prompt distillation, directly applying $P_c$ element-wise multiplied with the frozen teacher prompts $P_k^*$ for each task. The TTA process remained consistent across experiments. The results are shown in the table below:

If $P_c$ in MPDT is merely a generic scaling factor, we would expect the results in the first row to be similar to MPDT's performance. However, MPDT demonstrates a significant performance advantage, indicating that the prompt decomposition process effectively separates knowledge, aligning with our expectations.

We fully agree that an effective loss function can guide the prompt decomposition process and help the model converge faster. An intuitive idea is to consider adding a sparse loss concerning $p_k$. We incorporate the sum of absolute values of elements in $P_k$ as a penalty term into the loss function, with a weight $\lambda_1$ of 0.2, termed as $L_{sparse}$.

$L_{sparse}=\lambda_1 \sum_{k=1}^{|S|} \sum_{i=1}^m \sum_{j=1}^n |P_{k_{ij}}|$

We validated the model's performance using $L_{sparse}$ on the Cheetah-vel and MW ML45 datasets.

We found that MPDT+$L_{sparse}$ slightly outperforms the original MPDT in terms of performance, implying that the sparsity level of $P_k$ is correlated with prompt decomposition and the model's final performance. Due to rebuttal time constraints, exploration of more instructive loss function designs will be part of our future work. Besides loss design, in our experiments, we tried different learning rates for $P_c$ and $P_k$ but observed optimal performance when both had the same learning rate. We speculate that this occurs because, over training iterations, both prompts converge to their optimal values, and differing learning rates disrupt their joint convergence, leading to poorer performance under similar runtime conditions. Below, we present experimental results on the ML45 dataset where different learning rates were applied to $P_c$ and $P_k$.

W3: About the test time adaptation
We have already validated similar ideas in the 'Impact of adaptation method' (Section 5.3). We investigated (1) combining $P_c$ with the average of all task-specific prompts $P_k$ from the training set (we call $\hat{P_k}$) for TTA and (2) freezing the $P_c$, randomly initializing a new task-specific prompt $P_r$ combined with $P_c$ for TTA. Following the reviewer's suggestion, we added (3) selecting one $P_k$ from a training task combined with the $P_c$ for TTA, where we randomly sampled three $P_k$ from training tasks and computed the average model performance. We found that both of these initialization methods resulted in suboptimal outcomes. We attribute this primarily to the limited information that TTA can provide [1], where introducing task-agnostic gradients (e.g., additional $p_k$) may significantly degrade model performance. Using $P_c$ for initialization proves to be the optimal approach. Nonetheless, we observed (4) that omitting TTA entirely, i.e., not utilizing any information from test tasks, and instead combining $P_c$ with the average of $p_k$ from training tasks directly for model inference, yields better results than using $P_c$ alone.

[1] https://link.springer.com/article/10.1007/s11263-024-02181-w

In the case of any follow-up questions, we will in turn provide with further clarifications.

Thanks for the careful review of our work.

W1: How can the authors guarantee that they use the same amount of information in the test set for all baselines?
In methods not involving prompts, we indeed fine-tuned these methods on the test set. To highlight the superiority of MPDT as much as possible, we allowed for some unfairness in comparison. Therefore, we sampled $|X|$ labelled samples for fine-tuning MT-BC, MT-ORL, and HDT methods, making the experiments in Table 1 persuasive for non-prompt methods.

For methods involving prompts (Soft-Prompt and Prompt-DT),to ensure the absolute fairness of our method, we add experiments combining Soft-Prompt and Prompt-DT with TTA. The experimental results, labelled as Soft-Prompt-TTA and Prompt-DT-TTA, are shown in the table below.

Soft-Prompt-TTA showed performance improvements across all tasks, whereas Prompt-DT-TTA experienced performance declines in some tasks. The main reason for this is that Prompt-DT relies on high-quality trajectory data for prompts during testing, and applying TTA on unlabeled data may have adversely affected prompt optimization. Our MPDT method continues to demonstrate advantages across the majority of tasks.

W2: Did the authors reproduce the results of baselines?
We have replicated the experimental results of the baselines. In Table 6 of Appendix B, we provide the configuration of hyperparameters used. For Prompt-DT, MT-BC, and MT-ORL, we use the code provided by the Prompt-DT, adjusting most parameters to ensure a fair comparison with MPDT. Soft-Prompt, based on MPDT, training a universal prompt directly on datasets from all tasks. HDT followed all the details from the original paper. To ensure a fair comparison with MPDT, we aligned some general hyperparameters with those used by MPDT (See Author Rebuttal AR4).

For experiments in the few-shot generalization scenario, we did not adjust any hyperparameters and maintained the configuration from Table 6 because the chosen baselines either had architectures designed for few-shot scenarios or transferable insights applicable to our study.

[1] https://arxiv.org/abs/2304.08487

W3: Removing the claim of "SOTA".
Thanks for the suggestions and comments. We will remove the statement regarding SOTA terminology and clarify that our results are competitive.

W4: The technical novelty and performance problem.
About the technical novelty, please check Author Rebuttal AR1. About the performance problem in Table 3, please check Q2.

W5: The misleading problem of unseen tasks.
We are considering changing "Unseen Tasks Generalization" in Section 5.2 to "zero-shot generalization," so that our whole experiments can be referred to as "Unseen Tasks Generalization."

W6: The name might not be suitable.
Our main innovation lies in the implementation of prompt decomposition, distillation and subsequent TTA alignment in prompts. It can be modified to "Decomposed Prompt Decision Transformer Enables Unseen Task Generalization. This title may be more aligned with the current algorithm.

Q1: It looks unnatural that only the weights of prompts are trainable while the embedding layer and GPT blocks are all frozen.
Please check Author Rebuttal AR3.

Q2: Why simply using one or two of components won't achieve improvements?
The experimental results shown in Table 3 (fourth row) are misleading. The cross-task prompt $P_c$ contains rich inter-task information, providing a strong initialization for adapting to test tasks. To achieve optimal performance, fine-tuning task-specific information with TTA is necessary.

$P_c$ performs poorly in directly generalizing to unseen tasks. However, this does not imply that the prompt decomposition component cannot be used independently. To verify this, we performed experiments where $P_c$ was element-wise multiplied with the average of $P_k$ from all training tasks (referred to as $P_{ka}$ ), directly used for test tasks without TTA. As shown in the table below, our performance still surpasses that of the first row. The original intention of the fourth row in Table 3 was to demonstrate the poor performance of using $P_c$ alone. However, we recognize that this should not be construed as an experiment on the prompt decomposition component. Therefore, we propose replacing the fourth row in Table 3 with the results obtained by using $P_c \circ P_{ka}$ as the prompt for test tasks to avoid confusion.

Q3: How much efforts did the authors spend in tuning the hyper-parameters?
Almost all hyperparameters are referenced from Prompt DT, and we did not excessively tune our hyperparameters, which facilitates a fair comparison with other methods. The primary hyperparameter we focused on tuning is the prompt length $l$. The ablation experiments of the optimal value for $l$ are shown in Figure 3. In the table below, we present the results of ablation experiments on another hyperparameter $r$ for Cheetah-vel and MW ML45. Overall, MPDT is relatively sensitive to the selection of hyperparameters, which is a potential advantage of our work.

I thank the authors for the detailed responses to my questions. On reading the responses, most of my concerns have resolved.

Two concerns remain:

• (minor) W1. Prompt DT cannot utilize the same amount of test data due to its design. To alleviate this(?), the authors conduct experiments on Prompt-DT+TTA for fair comparison, but TTA isn't suitable for Prompt-DT, making it not fair enough. (Anyway, this is not a big problem.)
• (major) Q1. The question that "why the GPT blocks are all frozen" is not answered yet.

Another concern:

• (significant) The initial hyperparameter of Prompt-DT is #layer=3, #heads=1. The rebuttal states that "To ensure a fair comparison with MPDT, we aligned some general hyperparameters with those used by MPDT (See Author Rebuttal AR4)", and it hence seems that the experiments of Prompt-DT in this paper use #layer=12, #heads=12. If so, that would be unfair, since the other hyperparameters of Prompt-DT must be tuned extensively due to a significant change in the model size.

We agree with the reviewer's point and sincerely thank the reviewer for raising the score.

From the perspective of computational cost, the advantages of MPDT are clear (1.42M vs. 125.5M fine-tune) or (1.42M vs. 3.5M LoRA). Fully fine-tuning a model is undoubtedly costly, and LoRA's 3.5M parameters are based on the adjustable parameter count in a single-task scenario, as reported in [2]. When extending to a multi-task scenario, the parameter count will increase proportionally to the number of tasks.

From the perspective of scenario requirements, using prompts offers a natural advantage for knowledge transfer. Prompts have sufficient capacity to guide pre-trained models to adapt to reinforcement learning tasks with smaller data sizes. Additionally, achieving zero-shot generalization with prompts results in significantly less disruption to the model's knowledge and a lower error tolerance rate than fully fine-tuning the model.

[1] Can Wikipedia Help Offline Reinforcement Learning
[2] Unleashing the Power of Pre-trained Language Models for Offline Reinforcement Learning

Significant concern: The other hyperparameters of Prompt-DT must be tuned due to a change in the model size.

We did indeed use the 12-layer, 12-head version of Prompt-DT. To ensure a fair comparison, we reran the experiments and carefully adjusted some hyperparameters (setting the learning rate to 2e-5 and the learning rate decay weight to 1e-5). Other hyperparameters are independent of the model size. The model performance did not change significantly under the optimal hyperparameter configuration. The table below shows the performance of Prompt-DT and Prompt-DT-TTA across all tasks after adjusting the hyperparameters. We will replace the corresponding results in the main text with these updated results.

Table 1: Zero-shot generalization

Table 2: Few-shot generalization

We sincerely appreciate the reviewers' thorough examination. Below, we address the concerns one by one.

W1: Prompt DT cannot utilize the same amount of test data, and the TTA is unsuitable.
The Prompt-DT+TTA method selects the prompt $p$ that allows the model to perform optimally on the test task from the high-quality trajectories of the training tasks and then fine-tunes $P$ on the same amount of unlabeled test data as MPDT for TTA. This approach ensures zero-shot generalization (without using high-quality trajectories from the test task as prompts) and effectively combines Prompt-DT with TTA. We believe this is the optimal way to adapt Prompt-DT to zero-shot scenarios. The performance fluctuations of Prompt-DT+TTA across different tasks are mainly because Prompt-DT is not inherently a prompt-tuning method (despite using frozen prompts), as it requires high-quality labeled trajectories as prompts during both testing and training, which is a very strong prior condition. It would be more accurate to say that Prompt-DT might not be suitable for zero-shot generalization scenarios rather than stating it is unsuitable for TTA. In few-shot scenarios, MPDT still outperforms Prompt-DT, demonstrating MPDT's superior performance in existing scenarios and its ability to extend to new scenarios that existing algorithms struggle to handle effectively. From a rigor perspective, Prompt-DT could be removed from Table 1, which also indirectly highlights the limited exploration of prompt-based offline RL methods in zero-shot generalization scenarios.

Q1: Why the embedding layer and GPT blocks are all frozen.
In Author Rebuttal AR3, we conducted an ablation study on whether to freeze the embedding layer. Also, based on the ablation study conclusions from [1][2], we speculate that the embedding layer and GPT blocks are highly correlated. To achieve sufficient performance, one should either (1) freeze both the embedding layer and GPT blocks simultaneously and add external components (e.g., prompt) or (2) train both the embedding layer and GPT blocks together (either full fine-tuning or using LoRA). The choice of approach involves a trade-off between model performance, computational cost, and scenario requirements.

From the perspective of model performance, the primary purpose of using prompts is to preserve the model's inherent prior knowledge as much as possible. Both [2] and we believe that the full model fine-tuning used in [1] may lead to model overfitting, further disrupting the internal knowledge of the model. We considered adapting the method from [2] to our task (training the embedding layer and LoRA), but the results obtained based on the code provided by [2] differed significantly from those in the original paper. Moreover, when we applied [2] to our datasets, ML45 and MT50 (as shown in the table below), the performance on the training set was far inferior to MPDT, and even careful hyperparameter tuning could not alleviate the significant fluctuations in the reward curve. We speculate that language models using the LoRA structure may still face convergence difficulties on RL data. Furthermore, the experiments in [2] primarily focus on single-task scenarios, avoiding the challenges of multi-task scenarios where task gradients may mix with the internal knowledge of the model.

Furthermore, [2] introduced a language prediction auxiliary loss to ensure the model retains its language task memory during training. However, if the introduction of LoRA is entirely suitable and does not disrupt the internal knowledge of the model, there should be no need for an auxiliary loss. In other words, we must fully understand whether full fine-tuning or LoRA fundamentally introduces new knowledge by disrupting existing knowledge (replacement) or guides the model to use existing knowledge to achieve RL tasks (integration). The distinction between these two approaches results in differences in the model's performance ceiling and floor. The extent of adjustments within the model may need to be carefully controlled. Additionally, current work needs to explore convergence and generalization adequately. MPDT combines transferable prompts with the pre-trained model in a more harmonious and adaptable way. Most importantly, freezing the embedding layer and GPT blocks preserves the model's complete knowledge, which we believe is a better choice.

Thanks for reviewing our work attentively. We will answer the reviewer's questions one by one in the following.

W1: The paper appears to be hastily written and the presentation is hard to follow
We will revise all unclear and erroneous statements highlighted by the reviewer further to improve the logical structure and clarity of the paper.

W2: The problem about prompt decomposition
The significance of prompt decomposition lies in splitting the task prompt $P_k^*$ into cross-task prompts $P_c$ and task-specific prompts $P_k$. Both prompts are trained concurrently during the training process, where data from different tasks are sequentially inputted into the model. For the current task $k$, the model computes $P_k^*$ as the element-wise product of $P_k$ and $P_c$, which is then concatenated with the input data and optimized on the current task. Using standard normal distribution to initialize $P_c$, $u_k$ and $v_k$ is a common practice for prompt initialization [1,2].

$P_c$ is termed "slow" because it captures universal knowledge shared among all tasks in the task set $S$. Its learning rate should be slow to prevent overfitting to specific tasks. On the other hand, $P_k$ is termed "fast" because it compensates by quickly adapting to the characteristics of the current task, aiding better model convergence. In our experiments, we attempted different learning rates for $P_c$ and $P_k$ but found that performance was optimal when both had the same learning rate. We speculate that this is because, over training iterations, both prompts converge to their optimal values, and setting different learning rates disrupts their joint convergence, resulting in poorer performance under similar runtime conditions. Below, we present experimental results on the ML45 dataset where different learning rates were applied to $P_c$ and $P_k$.

[1] https://arxiv.org/abs/2109.04332
[2] https://arxiv.org/abs/2210.02390

W3: How to obtain teacher task prompt?
Here exists a separate training process for $P^{teacher}_k$, which has dimensions $\mathbb{R}^{l\times s}$. Prompt tuning is used to learn $P^{teacher}_k$ independently for each task $k$. Since there are no gradient conflicts in training on individual tasks, this process is fast and straightforward. We set the batch size to 256 and train for 100 epochs, typically converging in about 1 hour. There are indeed three training stages, but we consider the training of $P^{teacher}_k$ as a preparatory data phase. Once learned, $P^{teacher}_k$ does not require retraining.

W4: Can authors clarify what is subset $X$?
Please check Author Rebuttal AR2.

Q1: How do tasks for each meta-environment differ?
Each task within the environment has distinct reward functions and state transition functions.

The Cheetah-dir environment is one-dimensional, rewarding agents based on the angular difference between their movement direction and target direction. The Cheetah-vel environment is also one-dimensional, rewarding agents based on the difference between their velocity and a target velocity. The Ant-dir environment is two-dimensional, rewarding agents based on the angular difference between their movement direction and uniformly sampled target directions within 360 degrees.

ML10, ML45, MT10 and MT50 are trained and tested in a three-dimensional physical space. Each task in these datasets is a goal-conditioned environment, where the state space across all tasks shares the same dimensions. The action space remains identical across different tasks, though specific dimensions in the state space represent different semantic meanings. The task divisions are detailed in Table 9 of Appendix C. Due to space limitations, we describe 14 robotic manipulation tasks in Table 1 of the attached PDF for the Author Rebuttal.

Q2: How does the quality of the few shot prompt effect the model’s performance?
Intuitively, the quality of data used for fine-tuning cross prompts does indeed affect the final model performance. We conducted additional experiments focusing on data quality. In the Cheetah-vel and ML45 environments, we differentiated the quality of datasets into expert, medium, random, and mixed datasets. Specifically, we randomly selected labelled test set trajectories and partitioned the first 30% as random, the middle 30% as a medium, and the last 30% as expert datasets. The mixed datasets are an average blend of these three types. Each dataset consists of 200-time steps, which aligns with the setup for few-shot scenarios relative to the size of the training set.

We found that models fine-tuned using expert datasets perform the best, which aligns with our intuition. Additionally, the performance of models fine-tuned on mixed datasets is close to that on expert datasets, suggesting implicitly that the MPDT method can extract information from suboptimal datasets to ensure model performance.

Q3: $L_{dis}$ is calculated without being used?
We are sorry for this mistake. The parameter update rule is $\theta \leftarrow \theta-\alpha \nabla_\theta \mathcal{L}_{Total}$.

Q4: Why there is a double layer bracket?
$\mathcal{M}(P_k^*,\tau)$ here denotes the concatenation of task prompt $P_k^*$ and trajectory $\tau$ inputted into MPDT $\mathcal{M}$.

Q5: What is the vector multiplication here?
The dimension of $u_k$ should indeed be $\mathbb{R}^{r\times s}$. Furthermore, we should refer to both $v_k$ and $u_k$ as low-rank matrices rather than vectors, as these parameters exhibit matrix properties. Therefore, $P_k\in\mathbb{R}^{l\times s}$ is obtained through matrix multiplication $P_k = v_k\otimes u_k$.

Dear Reviewer zhrS,

We sincerely appreciate the time and effort you have already dedicated to reviewing our submission. We have carefully considered and addressed your initial concerns regarding our paper. Given that the discussion phase is nearing its end, we would greatly appreciate it if you could share any further feedback so we can respond promptly. Additionally, we welcome any new suggestions or comments you may have.

Thank you once again for your valuable insights and continued support.

Best regards,
The authors of Submission 9319

Thanks for reviewing our work attentively. We'll answer your questions one by one in the following. Considering the word limitation, we combine Weaknesses and questions with similar meanings to answer.

W1: why the language model could contribute to the RL tasks?
In Section 4.1, the Initialization part, we explained the reasons for using language models for initialization，"Given that Transformers are data-intensive and require pre-training on substantial datasets to achieve satisfactory performance, integrating PLMs from the same architectural family into offline RL is a natural progression. Incorporating PLMs into RL is not predicated on the direct applicability of language data to RL tasks. Instead, the advantage lies in leveraging the deep, nuanced representations acquired by PLMs from a variety of datasets. Theserepresentations encode a broad spectrum of patterns, relationships, and contexts that can transcend purely linguistic tasks."

Additionally, incorporating the parameters of pre-trained language models is also motivated by several studies in RL[1,2]. Leveraging the rich prior knowledge encoded in PLMs effectively addresses the data hunger challenge of transformer architectures, providing ample semantic information for RL tasks. These papers demonstrate the advantages of using language models for RL. Our approach differs from previous work by freezing the core architecture of the language model and using trainable prompts to guide the model's knowledge for RL tasks. This approach reduces computational demands (Table 7) while enhancing transferability. Experimental results in Table 1 comparing MT-ORL, MPDT-WP, and MPDT further validate the benefits of initializing MPDT with PLMs.

[1] https://arxiv.org/abs/2310.20587
[2] https://arxiv.org/abs/2201.12122

W2 & Q3: Why you need to use a long prompt to achieve better performance?
In Prompt DT, prompts used during training and testing consist of carefully selected high-quality trajectory segments. The original paper of Prompt DT does not specify the exact method for selecting high-quality prompts, and since the main focus of training is the model itself rather than the prompt, the requirements for prompt length are relatively low. Furthermore, the performance of Prompt DT requires sufficient support from prior information (labeled test data), as illustrated in Figure 3 of the Prompt DT paper, where prompt quality significantly impacts model performance, especially in scenarios with poor data quality.

The key difference in the MPDT method lies in freezing the core model, significantly reducing parameter count compared to Prompt DT (1.42M vs 125.5M). More than directly comparing prompt lengths between Prompt DT and MPDT may be unfair. On the other hand, MPDT optimizes prompts initialized completely at random on training data, thus requiring almost no prior knowledge.

In Figure 3 (with fixed training duration), we find that a prompt length of around 30 performs best. Without considering computational costs and training time constraints, we reintroduce performances of prompt lengths 3 and 9 on two tasks (trained until model convergence), finding that a prompt length around 9 also achieves highly competitive results. Moreover, prompt lengths between 9 and 30 are feasible in scenarios using a language model, as they do not noticeably increase the computational burden.

[1] https://arxiv.org/abs/2206.13499

W3 & Q2: This paper lacks detailed explanation on why using the cross-task prompt, task-specific prompt, and the prompt distillation.
In the context of multitask reinforcement learning, conflicting gradients arising from different tasks due to their objectives and environments are key factors leading to poor model performance. Let us separate common features and conflicting features between tasks. In that case, it not only facilitates training the model on existing tasks but also allows us to use common features (manifested as prompts in our work) for prompt initialization in downstream tasks. The cross-task prompt remains consistent across all training tasks, while the task-specific prompt is tailored to each task's unique characteristics. Our structured decomposition enables more regulated and harmonious parameter updates, thereby enhancing parameter efficiency and facilitating the extraction of general knowledge more effectively.

The primary purpose of prompt distillation is to aid in further separating two types of prompts. Due to the lack of explicit constraints, directly implementing prompt decomposition on the multitask dataset $S$ may lead to an overlap in the information learned by $P_c$ and $P_k$, potentially undermining their ability to capture distinct intended details. We found that employing knowledge distillation from prompts trained separately on individual training tasks was a successful approach for obtaining well-decomposable prompts.

We verify the performance of the proposed component in Table 3 of the main text.

Q1: Why you use the pre-trained embedding layer for word tokens to encode the RL tokens?
Please check Author Rebuttal AR3.

We sincerely appreciate the reviewers' thorough examination and hold the reviewers' opinions in the highest regard. We strive to do our utmost to address reviewers' concerns, not only about the acceptance of the paper but also in providing meaningful insights for the RL community.

The reviewer's concern can be understood as the possibility that not retraining certain layers of the language model may result in an inadequate understanding of RL data by the model. We believe preserving the language model's knowledge is crucial, and prompts can bridge the gap between the two tasks without retraining the embedding layer or GPT blocks. In AR3, we conducted an ablation study on whether to freeze the embedding layer. The experimental results are intuitive, as training only the embedding layer while freezing the GPT blocks still leads to a suboptimal understanding of input information. Ablation study conclusions from [1][2] also demonstrate that the embedding layer and GPT blocks are highly correlated ([1] performs full fine-tuning of a language model on RL tasks. In contrast, [2] trains a language model on RL tasks using LoRA). Comparing [1][2] with MPDT can further address the reviewer's concerns. The reviewer's concerns can be extended to whether to (1) freeze both the embedding layer and GPT blocks simultaneously and add external components (e.g., prompt) or (2) train both the embedding layer and GPT blocks together (either full fine-tuning or using LoRA). Only these two approaches can achieve better performance. The choice of approach involves a trade-off between model performance, computational cost, and scenario requirements. Our decision to adopt approach (1) is based on the following reasons.

From the perspective of model performance, the primary purpose of using prompts is to preserve the model's inherent prior knowledge as much as possible. Both [2] and we believe that the full model fine-tuning used in [1] may lead to model overfitting, further disrupting the internal knowledge of the model. We considered adapting the method from [2] to our task (training the embedding layer and LoRA), but the results obtained based on the code provided by [2] differed significantly from those in the original paper. Moreover, when we applied [2] to our datasets, ML45 and MT50 (as shown in the table below), the performance on the training set was far inferior to MPDT, and even careful hyperparameter tuning could not alleviate the significant fluctuations in the reward curve. We speculate that language models using the LoRA structure may still face convergence difficulties on RL data.

Furthermore, [2] introduced a language prediction auxiliary loss to ensure the model retains its language task memory during training. However, if the introduction of LoRA is entirely suitable and does not disrupt the internal knowledge of the model, there should be no need for an auxiliary loss. In other words, we must fully understand whether retraining the model internally (including the embedding layer or GPT blocks) fundamentally introduces new knowledge by disrupting existing knowledge (replacement) or guides the model to use existing knowledge to achieve RL tasks (integration). The distinction between these two approaches results in differences in the model's performance ceiling and floor. The extent of adjustments within the model may need to be carefully controlled. Most importantly, freezing the embedding layer and GPT blocks preserves the model's complete knowledge, which we believe is a better choice, and the experimental results also support this.

From the perspective of computational cost, the advantages of MPDT are clear (1.42M vs. 125.5M fine-tune) or (1.42M vs. 3.5M LoRA). Fully fine-tuning a model is undoubtedly costly, and LoRA's 3.5M parameters are based on the adjustable parameter count in a single-task scenario, as reported in [2]. When extending to a multi-task scenario, the parameter count will increase proportionally to the number of tasks.

From the perspective of scenario requirements, prompts have sufficient capacity to guide pre-trained models to adapt to reinforcement learning tasks with smaller data sizes. Additionally, achieving zero-shot generalization with prompts results in significantly less disruption to the model's knowledge and a lower error tolerance rate than fully fine-tuning the model.

We demonstrated from both experimental results and theoretical principles whether it is necessary to freeze the embedding layer and GPT blocks. We sincerely welcome further discussion with the reviewer and will try to address any concerns.

[1] Can Wikipedia Help Offline Reinforcement Learning? arXiv preprint arXiv:2201.12122.
[2] Unleashing the Power of Pre-trained Language Models for Offline Reinforcement Learning. arXiv preprint arXiv:2310.20587.