Decomposed Prompt Decision Transformer for Efficient Unseen Task Generalization
Authors: Hongling Zheng
Li Shen
Yong Luo
Tongliang Liu
Jialie Shen
Dacheng Tao
Venue: NeurIPS 2024 poster (2024)
Keywords: Offline Reinforcement Learning
Prompt Tuning
Primary area: reinforcement_learning
TL;DR: A novel offline reinforcement learning method for efficient generalization of unseen tasks

Abstract: Multi-task offline reinforcement learning aims to develop a unified policy for diverse tasks without requiring real-time interaction with the environment. Recent work explores sequence modeling, leveraging the scalability of the transformer architecture as a foundation for multi-task learning. Given the variations in task content and complexity, formulating policies becomes a challenging endeavor, requiring careful parameter sharing and adept management of conflicting gradients to extract rich cross-task knowledge from multiple tasks and transfer it to unseen tasks. In this paper, we propose the Decomposed Prompt Decision Transformer (DPDT) that adopts a two-stage paradigm to efficiently learn prompts for unseen tasks in a parameter-efficient manner. We incorporate parameters from pre-trained language models (PLMs) to initialize DPDT, thereby providing rich prior knowledge encoded in language models. During the decomposed prompt tuning phase, we learn both cross-task and task-specific prompts on training tasks to achieve prompt decomposition. In the test time adaptation phase, the cross-task prompt, serving as a good initialization, were further optimized on unseen tasks through test time adaptation, enhancing the model's performance on these tasks. Empirical evaluation on a series of Meta-RL benchmarks demonstrates the superiority of our approach. The project is available at https://github.com/ruthless-man/DPDT.

Decomposed Prompt Decision Transformer for Efficient Unseen Task Generalization
Hongling Zheng'
Shen 2f
Yong Luo't
Tongliang Liu? Jialie Shen Dacheng Tao 'The University of Sydney
WWhhan University
Shenzhen Campus of Sun Yat-sen University
*City, University of London *Nanyang Technological University
(hlzheng, luoyong) edu. en {mathshenli, jialie, dacheng tgo com tongliang. legg..ldudyyyyy edu. au
Abstract
Multi-task offline reinforcement learning aims to develop a unified policy for di- tasks without requiring real-time interaction with the environment. Recent work explores sequence modeling, leveraging the scalability of the transformer architecture foundation for multi-task learning. Given the variations in task content and complexity, formulating policies becomes a challenging endeavor, re- quiring careful parameter sharing and adept management of conflicting gradients extract rich cross-task knowledge from multiple tasks and transfer it to unseen tasks. In this paper, propose the Decomposed Prompt Decision Transformer (DPDT) that adopts a two-stage paradigm to efficiently learn prompts for unseen tasks in a marameter-efficieen manner. We incorporate parameters from pre-trained language models (PLMs) to initialize DPDT, thereby providing rich prior knowledge encoded language models. During the decomposed prompt tuning phase, learn both cross-task and task-specific prompts on training tasks to achieve prompt decompo- sition. In the test time adaptation phase, the cross-task prompt, serving as a good initialization, were further optimized on unseen tasks through test time adaptation, enhancing the model's performance on these tasks. Empirical evaluation on series of Meta-RL benchmarks demonstrates the superiority of our approach. The project is available /tttt::///// :ttp::://ttiiii httpg//hbb...o//uuueeeeeeeeeeeee man/DPDT
Introduction
The purpose of offline reinforcement learning (Offline RL) (T] is develop a reward-maximizing RL strategy using offline data. This approach is of highly valuable in real-world scenarios where online data collection is expensive, tine-consuming, or impractical. Existing offline RL algorithms typically perform well in single task but often struggle for multiple tasks with similar conditions and objectives, as they lack the ability to separate common knowledge from conflicting task gradients. In contrast, humans can leverage knowledge from existing tasks to excel in new ones, which has The goal
led increasing research interest in multi-task reinforcement learning (MTRL) 2 3 of MTRL is to develop universal strategy applicable to tasks with certain similarities, thereby enhancing adaptability and performance in multi-task environments.
Decision transformer (DT) [5] and Prompt-DT 6 introduce the transformer architecture to the field of RL, demonstrating the powerful data modeling capability sequence offline RL. Additionally, they provide possibilities for integrating advancements [7] from language modeling into MTRL methodologies. Prompt-tuning DT [8] uses gradient-free method introduce prompts, retaining context-specific information and catering to specific preferences. These existing sequence-basd based offline RI methods are primarily trained and tested on the same task or perform fine-tuning using a
Corresponding authors.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

small portion of labeled test data (9] As a result, these algorithms often perform poorly when faced with testing tasks that are unseen and unlabeled in Meta-RL setting [0]
Leveraging MTRL to extract general knowledge offers a promising approach for facilitating cross- task knowledge transfer in Meta-RL. scenarios. However, as the number of tasks increases, gradient conflicts become more pronounced, hindering MTRL performance due to unregulated parameter sharing. Additionally, while transformer architectures can capture extensive relationships in offline sequential data, their data-hungry nature means that insufficient training data for RL tasks significantly diminishes model performance. feasible solution might be to enhance the model's 's prior knowledge.
To remedy these drawbacks, we propose a prompt-based based MTRL method named Decomposed Prompt Decision Transformer (DPDT), inspired by some works in natural language processing where knowl- edge transfer in multi-task learning scenarios is achieved through prompting strategies [1] 12] We first employ pre-trained parameters from GPT initialize DPDT architecture. Incorporating the parameters of PLMs is motivated by several studies in RL. [3
Leveraging the rich prior knowledge encoded in Pre-trained Language Models (PLMs) effectively addresses the data hunger challenge of transformer architectures, providing ample semantic information for reinforce- ment learning tasks. Then, we design two-stage training and testing framework for DPDT. (1) Decomposed prompt tuning phase: At this stage, we use prompt decomposition to avoid gradient conflicts between different tasks and to extract common knowledge. Specifically, we decompose the task prompt for each task into cross-task prompt and a task-specific prompt. The cross-task prompt remains consistent across all training tasks, while the task-specific prompt is tailored each task's 's unique characteristics. By isolating the cross-task prompt from the task-specific prompts, the model ensures that updates related to general knowledge not conflict with those related to specific tasks. Compared to [2), our structured decomposition enables more regulated and harmonious parameter updates, thereby enhancing parameter efficiency and facilitating the extraction of general knowledge more effectively. (2) Test time adaptation phase: The cross-task prompt, serving as a strong initialization, is further optimized on unlabeled unseen tasks by incorporating the Test Time Adaptation (TTA) [16] to our model. TTA dynamically adjusts the cross-task prompts during the testing phase based on task characteristics, enhancing the model's ' adaptability to unseen tasks features.
14] 15].
Our DPDT has been empirically validated in the Meta-RL. setting, and the results demonstrate its superiority compared with many recent and competitive counterparts [3| 6 17 Furthermore, we conducted ablation experiments covering aspects such as prompt length, scalability, and model variants to establish the superiority of the model. The main contributions of this work are follows:
We reconsider the problem of knowledge extraction in MTRL and propose the method of Decomposed Prompt Decision Transformer.
We propose a two-stage paradigm, which includes decomposing the task prompts for training tasks into tross-task prompts and task task-specific prompts, and aligning the cross-task prompts in unlabeled unseen tasks.
We demonstrate the effectiveness of DPDT through intensive experiments on a broad spectrum of benchmarks, highlighting its competitive performance in Meta-RL scenarios.
2 Related works
In this section, we summarize the most related works as two-fold: offline RL and multitask RL.
Offline RL. In contrast to traditional RL methods [18 19, offline RL focuses on training models and performing trial-and-error using offline data without environmental interaction to arrive at appropriate strategies. These methods primarily address the issue of -fiit--diitriuution (OOD) through strategies such as constraining the learning policies [20] or bounding the overestimated policy values (21). The integration of transformer architecture in sequence modeling has emerged as a prominent approach for addressing offline RL tasks [22 23 further demonstrating the advantages of data- ddiiaaddvennnnn policy learning. Decision Transformer (DT) involves encapsulating rewards, states, and actions into triples and training them using autoregressive supervision on offline data. Owing to the transformer's 's proficiency capturing and fitting long time series features, it has achieved remarkable results in various offline RL tasks. HDT [17] generalizes new tasks by designing an adaptation module initialized by a hypernetwork. To alleviate the tuning burden while preserving performance, prompt tuning [24 25 in NLP focuses on optimizing only the input parameters while keeping the majority

---

of the PLMs parameters frozen. However, integrating prompt tuning into RL field poses challenge, as RL prompts lack semantic information and are challenging to optimize. While Prompt-DT [6] selects pre-defined expert trajectories combined with inputs to guide model training, it primarily relies on the quality of these trajectories for improvement, rather than enhancing prompts directly. On the other hand, Prompt-tuning DT [8] stands out for introducing prompt-tuning techniques using gradient-free approach, with the goal of preserving environment-specinee details and accommodating specific preferences.
Multitask RL. Multi- Task Reinforcement Learning (MTRL) 26 27 aims to address multiple similar reinforcement learning tasks using a unified model. A straightforward approach involves
developing a lask-conditional multi-task model, akin to those utilized goal-conditional RL [2] and
uisual-language grounding [30]. While this method has demonstrated success in certain scenarios, it often encounters challenges stemming from negative interference among tasks. PaCo [3 delves into a compositional structure within the parameter space, distinguishing between task-agnostic and task-specific components. This approach significantly enhances the efficiency and robustness of the MTRL process, leading to more effective training outcomes. Building upon the MTRL paradigm, multi-task prompt [32] aims to acquire transferable, cross-task prompts from multiple tasks, guiding outputs for unseen downstream tasks. Several studies have approached multi-task prompt design from the perspective of prompt decomposition, yielding notable results across various tasks [3
The most relevant work to ours is Prompt-DT [6], which utilizes carefully selected prompts for Meta-RL tasks training. Our approach differs in (1) employing trainable prompt decomposition to avoid gradient conflicts among multiple tasks, (2) further optimizing general prompts with TTA without using any test data labels, and (3) extending the model architecture by incorporating PLMS for initialization. To the best of our knowledge, we are the first to implement multi-task prompt tuning based on PLMs parameters the reinforcement learning domain.
Preliminary
In this section, we provide several concepts and terminologies that will be used in this work.
3.1 Prompt Decision Transformer
Prompt-DT 6 examines the beneficial effects of incorporating trajectory prompts on the DT in few-shot scenarios. Specifically, the form of trajectory prompts is the same as that of training trajectories, consisting of triplets made up of state s", action a", and return-to-go f* However, these trajectory prompts are significantly shorter than the training trajectories. During training, for each task the trajectory prompt Ti concatenated with the corresponding training trajectory T to form T; imput = (Ti Ti,K), which is then inputted into the model f for training. The concrete form of k. and TiKK are as follows:
T-K- = (ft, si ai...... Pi., ag.)) TrKK = =f, (i,,,,, SK, ax)
where K" represents the number of environment steps stored in the prompt, and K is the nearest steps of the training trajectory. The prediction head associated with the state token s is designed to predict the corresponding action It is noteworthy that during training, the model does not predict the actions of the trajectory prompts. The loss function formulated as follows: a,mm denotes the actual action the m-th timestep of the i-th task, while 7s.m-1 encompasses all data up and including the (m th timestep in the training trajectory of the 5th task.
K
Lm E Sippe To:
1
Σ
(aim - - 11))
(2)
K
m=
3.2 Test-Time Adaptation
Test time adaptation (TTA) 16 aims to minimize the gap between training data and testing data
distribution during the testing phase. Test-time prompt tuning (TPT) (35 leverages the extensive
knowledge in transformer architecture to enhance its generalization capabilities in zero-shot scenarios.
3

---

Decomposed Prompt Tuning
Test Time Adaptation
Linear Layer
Pre-trained Blocks
Embedding
10 PitP), dlP) o(P).ott))
task prompt
Distiliation
task specifis prompt Ps
prompt Pinacher
prompt P
M2
MISDIS
Figure I: The architecture diagram of DPDT. For simplicity and clarity, the figure only displays the decomposition and integration process during training on a single task. Snowflake icons represent the frozen parts of the model that are not subject to training updates, while flame icons indicate components the model that remain trainable. Left: Decomposed Prompt Tuning. The prompt decomposition is trained then combined with the training samples that include steps, inputted into DPDT, and outputs the corresponding action Right: Test Time Adaptation. When test samples fed into the model, we calculate the mean and variance of all samples at each layer and compute the loss by comparing them with the mean and variance of the corresponding layer from the training samples. The losses from all layers are summed to obtain the alignment loss, denoted as Lalign-
the entire dataset with the assistance of the teacher prompt Pk
teacher Pi
During the inference phase, multiple randomly augmented views created from the provided testing sample Xtest- Predictions with entropy below a specified threshold are retained, while other views filtered out using confidence selection criterion. The average entropy of the filtered predictions then used to unsupervisedly update the prompts p. Some methods [3 37 38), considering the data instability caused by augmentation techniques, alignment of feature values in attention layers or feedforward layers to achieve test-time adaptation. Our method falls into this category.
4 Decomposed Prompt Decision Transformer
In this section, provide comprehensive description of the proposed DPDT. Task prompts and training data are combined to calculate the loss through action prediction, which is then used to optimize the cross-task and task-specific prompts in backward pass. Once optimized, the cross-task prompts serve as good initialization for use in unseen tasks during test-time adaptation. The objective of decomposed prompt tuning is to learn cross-task prompts and task-specific prompts through prompt decomposition. Test-time adaptation provides an alignment approach for the cross-task prompts. Below, we will describe each module of DPDT in detail.
4.1 Decomposed Prompt Tuning
Initialization. Given that Transformers are intenniteesivve and require pre-training on substantial datasets to achieve satisfactory performance, integrating PLMs from the same architectural family into
offline RL is natural progression. Some existing work has already explored this avenue [13 39 40) Taking inspiration from this, we employ GPT2-SMAL these parameters frozen throughout training. It worth noting that, 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. These representations encode broad spectrum of patterns, relationships, and contexts that can transcend purely linguistic tasks. This taps into the reasoning and few-shot capabilities of language models, addressing challenging scenarios like data scarcity and sparse rewards.
[4]] to initialize our DPDT and maintain

---

Prompt Decomposition. Given a set of training tasks S = (S1, S2....,,,, our objective is to learn general prompt P. that encapsulates common knowledge shared across all tasks in S and can efficiently adapt to unseen tasks. Extracting general task information from tasks with different distributions is often challenging. as gradient conflicts between tasks can lead to suboptimal convergence of information. We adopt prompt decomposition approach to address this issue. As shown in Figure 1 let Pe E IR and Ps E R The task prompt Pt for the k-th task is obtained by taking the dlement-wiss, product of P. and Pt. The goal of prompt decomposition is to enable efficient knowledge sharing across S, while still allowing each task to maintain its own parameters to encode task-specific knowledge. The Pe aims acquire general knowledge from S, while the task -peecffc prompt Pk allows task to retain its unique knowledge. The task prompts parameterization of the k-th training task is expressed as:
PE = P. Ps Ps Po (UR up).
In the specific implementation process, inspired by LORA [42].
Algorithm 1 Decomposed Prompt Tuning
we further decompose each task- specific prompt into two low- rank vectors 2kk E Rixr and us E R'Xs using low-rank method. We obtain Ps through vector multiplication. Here / rep- resents the prompt length, r rep- resents the hidden layer dimen- sion, and represents the prompt dimension. The hyperparameter r is a manually specified low- rank parameter. Its introduction is crucial for designing prompts for all tasks in the dataset, significantly maintaining model superiority while reducing computational load. We use standard normal distribution to initialize Pe. us and eg. Given the DPDT M, the mean squared loss between the model's predicted actions and the true actions is calculated as:
Input: Training task set S, Offline datasets DA, Batch size M, Learning rale a, training iterations N, teacher task prompts pi pinncher,
Initialize: Initialize a 12-layer, 12-head DPDT M using GPT2-SMALL, randomly initialize cross-task prompts P. and low-rank vectors uks Mg.
for / = 1 to N do
for k in S do
Select trajectory t that contains M samples task k. Calculate Pt by Equation
Calculate LMSE and Laka according to Equations and Computed loss function by Equation
0 O-- 0 a VoLoolll-
end for end for
LMSE = (a - M (P)
Prompt distillation. Due to the lack of explicit constraints, directly implementing prompt decom- position on the multitask dataset S may lead to overlap in the information learned by P. and Ps. potentially undermining their ability to capture distinct intended details. We employed knowledge distillation techniques compel the cross-task and task-specific prompts to learn their respective information. We obtained teacher task prompts pk tuning methods individually. During training, the mean squared error is calculated directly between teacker and pT: Pk
teacher for each task & by using traditional prompt-
Laie
Pk
teacher
pk
(5)
KE|S!
The total loss function for training task prompts for obtaining cross-task prompt to be transferred to the target side is then:
Drotal = LMSE + ALdis
(6)
where À is a weight to balance the impact of distillation loss terms. our experiments, we set À 0.5. The overall summary of the multitask training algorithm is presented in Algorithm []
to
4.2 Test Time Adaptation
During the test time adaptation (TTA) phase, we address distribution bias by aligning the distribution of unlabeled test samples with the training samples. For each test task in the test task set T, randomly select a subset X of unlabeled test samples, combine them with the cross-task prompts P. and input them into the model.
Here, we introduce the data collection method for X. The model's testing phase usually occurs in simulated environment where predefine our expected reward values r. The environment
5

---

provides the initial state of the environment, consistent with the settings during inference in prompt DT methods. However, unlike training tasks where ground-trune truth labels exist, for action a, we assign value sampled randomly from the action space (which is typically consistent between training and testing tasks). We feed this sequence of Markov chains into the environment, obtaining rewards and the next environment states iteratively, assigning randomly sampled value to action 02 in subsequent iterations. This process is repeated [X] times, resulting in data of the form (f0, So, do. fl,, $1, di,..... N S|N], A[N]).
In each layer of the model, we calculate the alignment loss based on the means and variances of the training and test samples. Our goal is to update the Pe for the given test task through this alignment loss. For each test task t, we denote the distribution of the test samples as and the distribution of the training samples as D. Specifically, we calculate the aligned token mean and variance via:
X
X
m (T)
X
HA oi
- X
Σ E11
His - (T))
Here, Hi.; represents the state of the gth sample at the 1th hidden layer, while (T) and of (T) denote the mean and variance of all test samples the 7th hidden layer, spectively. Similarly, for each hidden layer of the model, we pre-calculate the statisti- cal measures of mean yu (D) and variance of (D) of the training samples, which are uniformly sampled across all tasks in the training set, in an offline setting to reduce parallel computing costs, since both train- ing samples and labels are accessible. The formula for calculating the alignment loss
Algorithm 2 Test Time Adaptation
Input: Test samples set X, Cross-task prompts Po. (D()), of (D), The number of layers L.
I: for | = 1 to L do for in X do
2:
Calculate Hi,s obtained by inputting the concate- nation of Pa and into DPDT.
4: end for
5: end for
6: for 1 do
Compute (T(T) and of by Equation
end for
9: Compute token distribution alignment loss by Equation
10: Optimize Lanign to update P..
function is as follows:
1 Σ (T) (D)) + (
Lalign
T
(8)
1=1
The test time adaptation phase process is illustrated in Algorithm
5
Experiment
In this section, we present extensive evaluation of our proposed DPDT using widely recognized benchmarks. Additionally, we conduct empirical ablation studies to dissect and understand the individual contributions of the core components of our methodology.
5.1 Environments and Baselines
Environments. To ensure a fair comparison with existing multi-task offline reinforcement learn- ing algorithms, we conducted verification of DPDT using the MuJoCo 43 and MetaWorld (30 benchmarks, which serve as standard tasks in the domain of sequence offline RL, offering sufficient diversity and representing common challenges in classical RI,, such sparse rewards, complex state spaces, and precise control robotic systems. Our experiments on the Chectah-dir, Cheetah-vel. and Ant-dir environments in the MuJoCo benchmark meticulously adhere to the datasets and method- ologies outlined Prompt-DT. These tasks penalize agents for using excessive control signals. In the MetaWorld benchmark, we used the ML10, ML45, MT10 and MT50 environments for Meta- RL. A detailed description the datasets and the division of training and test tasks is provided in Appendix A
Baselines. We compared DPDT to the following offline RL baselines: (1) Multi-task Behaviour Cloning (MT-BC) [44) MT-BC optimizes multi-task learning by exclusively simulating trajectories from the original dataset, dispensing with the need for prompts and reward-to-go tokens. This approach emphasizes the utilization of intrinsic task-specific information, adopting a behavior cloning

---

Table I: Results for Meta-RL control tasks (zero-shot scenarios). The best mean accumulated returns are highlighted in bold. For each prompt-needing environment, prompts of length K=30 are utilized. Each experiment
run three times to ensure stability and reproducibility of the results. We report
the average returns and standard deviations for these three runs (the higher, the better).
MT-BC (44)
MT-ORL
[
Soft-Promph (45)
HDT 17
Prompi-DT (6
DPDT-WF
DPIT
Trainable Params
125.5M
125.5M
3.94M
12.94M
125.5M
1.42M
1.22M
Percentuge
100%
100%
3%
10.31%
100%
1.14%
1.14%
Cheetah-dir
24.711122
86.92 16.92
45.32
-7.92.2
11,73a128
50.3221147
Cheetab-vel
-201. 6622262
-148.24 18
162.75
-192,3821150
-143.14.21 40
-139.88
Ant-dir
131 8911299
109.21
119.45a142
115.43 115.43a1022
123,46,10.7
101.49217.4
121.84a801
MW ML.10
256.77a1183
343. 16.9
246 146.42e24ra do
292, 14x8.1
317.3
204.88 884.88
271.01
MW ML455
287.37gt1.s
266.74425.8
91.97214:11
274.88g19.74
294.55
300.71g15,74
347.21.11.52
MW MT 10
547.8341104
1064.58.1130
201 11
964.57.x8
1087.54.
1015.91
1317,52,g.22
MW MT 50
929.74a22.1
400.71.20.0
820.45.27.19
994.6325599
1131,01
1559,94,2,1
Average
225.76
354.04
130.62
309.79
373.88
374.66
518.28
strategy streamline the learning process. (2) Multi-Task Decision Transformer (MT-ORL) (5]
We train a decision transformer to learn multiple tasks from the training set. To construct the MT-DT, we exclude prompt augmentation present in DPDT, while retaining the rest of the training process identical to that of DPDT. (3) Soft-Prompt [45]: Soft-prompt is trained using a universal prompt across all tasks. (4) Hyper-decision transformer (HDT) [17]: HDT efficiently adapts DT to new tasks by augmenting them with an adaptation module, whose parameters are initialized by a hyper-network, enabling quick and efficient adaptation with minimal data. (5) Prompt-DT [6]: Prompt-DT builds on DT, leveraging trajectory prompts and reward-to- go for multi-task learning and generalization to unseen tasks.
Implementation details. All experiments were carried out on a server with 8 NVIDIA 3090 GPUs, each with 24GB of memory, using Py PyTorch (46) and Hugging Face Transformers libraries (47. The experimental hyperparameter configurations are shown in Appendix B The computer resources utilized by all methods are shown in Table 12
5.2 Main Results and Analysis
Zero-shot Generalization. In Table []] we compare the zero-shot generalization ability of DPDT and the baselines to investigate the overall performance of DPDT. For evaluation, we use the average episode cumulative returns in the test task set as the evaluation metric. Additionally, we introduce variant of DPDT that does not use GPT-SMALL parameters for initialization, referred to as DPDT- WP (DPDT-Without Pretrained). Soft prompts adapt to tasks in a parameter-efficien way. However, training a universal prompt across multiple tasks suffers from significant gradient interference, as demonstrated by the experimental results. The prompt-DT performs well in few-shot scenarios due to its utilization of test data for fine-tuning. However, in downstream tasks where data is scarce, the prompt fails to provide sufficient task-specific guidance to the model, resulting in suboptimal performance. Importantly, our proposed DPDT exhibits significant performance improvements over fine-tuning and prompt-based methods. This vividly demonstrates the distinct advantages offered by our innovative multitask training techniques. It is worth noting that without initialization with PLM, the performance of DPDT-WP inferior to most methods. We believe this is mainly due to the model lacking sufficient prior knowledge for effective multi-task prompt tuning. resulting in suboptimal performance of DPDT. In Figure 2 we illustrate the accumulated returns curves of DPDT and other baselines across the Cheetah-vel. MW ML45, and MW MT50 environments. Additional curves for other environments can be found in Figure 4 We also conducted experiments on Soft-Prompt and Prompt-DT combined with TTA (shown in Figure [T]. Soft- Poft-Prompt-TTA showed performance improvements across all tasks, whereas Prompt-D's -TT-TA 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.
Few-shot Generalization. We explored the performance of DPDT in few-shot scenarios and further investigated whether the prompt decomposition mechanism successfully isolated general knowledge. In this scenario. DPDT does not use TTA to align cross-task prompts P.. Instead, P. is fine-tuned
7

---

Table 2: Results for Meta-RL control tasks (few-shot scenarios). The best accumulated returns are highlighted in bold. For each prompt-needing environment, prompts of length K=30 are utilized. Each experiment
three times to ensure stability and reproducibility of the results. We report the average returns and standard deviations for these three runs (the higher, the better).
MT-ORL [S
Poft-Prompl (4S)
HDT 17
Prompi-DT (6)
DPDT-WP
DPOT
DPDI-F
Trainable Params
125.5M
3.94
12.94
125.5M
1.4MM
1.42M
125.5M
Percentage
100%
3%
10.31%
100%
1.14%
1.14%
100%
Cheetah-dir
46.22.xx
940.24.1
875.23
934.78,835
116.8
1037.85.5.
Cheetah-vel
1246.642222
-41.81a210
-63.81
-37.80g2mm
+18.07
-30.73
-29.85
Ant-dir
110.51a22
37
361.4925.1
411.96,uas
308. 10a5.12
384.29 a:000
400.01a839
MW MI.10
421.2249.11
379.82 414.76
467. 8 k
17
485.27 281.27
535.52217.39
670.24.3.8k
MW ML45
265.74
448.72211.38
477.1922.16
473.34 2A.12
519.28 2h.22
579.09 18.12
600.44
Average
120.60
421.204
423.56
419.47
442.27
484.66
535.74
MT-BC
MT-ORL
Soft-hrompt
HDT
Prompt-DT
NPOT-MP
MPDT
Cheetah-vel
ML 45
50
150
1500
Episode Reeur
-200
Epicocin Return
250
Egisadl Return
1250
-250
200
1000
-300
150
750
100
350
250
-400
50-
15000
1250
1000 1250
Training lteration
Training
Training Reration
Figure 2: Episodic accumulated returns in three tasks of MTBC, MT-ORL, Soft-Prompt, HDT, Prompt-DT, DPDT-WP and DPDT. Each method restricted to 1500 rounds of runs in each environ- ment.
directly small number of labeled test samples through self- self-super paradigm. Specifically, we randomly selected only one trajectory from the test dataset for fine-tuning. Other methods adhered to the same amount of fine-tuning data. Furthermore, in few-shot scenarios, we fully fine-tuned DPDT on the complete training and testing data, denoted as DPDT-F. The performance of the DPDT-F method represents the upper bound of all model performances in the current environment. Table 2 shows that the DPDT method, even after fine-tuning, still significantly outperforms or matches the baseline algorithms, demonstrating the effectiveness of prompt decomposition in few- shot environments. Moreover, in some datasets, DPDT approaches the performance of fully fine-tuned models on the test set using only 1.14% of the parameters, as observed in the cheetah-ve vel environment. It is worth noting that in the Ant-dir environment, the performance of DPDT is slightly inferior to that of Prompt-DT in both zeot and few-shot settings. We believe the primary reason for this is the significant domain difference between the language model and environment.
5.3 Further Analysis
In the ablation studies, we conducted research on the components of DPDT well the impact of prompt length and model parameters on convergence speed and performance.
Impact of model components. As shown in Table 3 the impact of prompt decomposition was evaluated. Compared to the soft-prompt method in the unseen task settings (first row), substituting with decomposed prompts P& and P. without distillation (third row) resulted in performance improve- ments across all three tasks. This ablation highlights the significance of the prompt decomposition strategy in DPDT, demonstrating that the shared component adeptly captures the diverse cross-tast knowledge essential for enhancing target downstream tasks.
To evaluate the impact of prompt distillation, we trained standard prompt shared across all training tasks using the same training loss as DPDT. The teacher prompts for each task remained consistent in DPDT. Compared to the basic bascline (first row), incorporating prompt distillation (second row)

---

Table Ablation:
The impact of prompt decomposition, prompt distillation and test time adaptation.
Decomposition
Distillation
TTA
Cheetah-vel
MW ML45
MW MTS0
-171.23
91.97
400.71
/
-63.00
108.01
709.81
V
-160.10
273.99
1137.39
-167.80
149.21
824.07
-139888
347.21
1559.94
Table 4: Ablation: The impact model size. The elements of the triplet represent, in order, the
number of transformer blocks, the count of attention heads, and the size of the hidden layers.
Model size
Cheetah-vel
Ant-dir
MW ML45
MW MT50
13,1,128)
-164.88
129.34
288.14
749.18
(12.12,768)
-139.88
121.84
347.21
1559.94
(24,16,768)
-210.35
165.99
292.48
1527.34
resulted in modest improvement in average performance across the three tasks. This emphasizes that prompt distillation from separately trained source prompts is an effective strategy for acquiring high-quality decomposable prompts.
To compare the impact of using TTA on model performance, we examined the results the fourth and fifth rows. We found that the use of TTA affects the model's final performance. The reason is intuitive: cross-task prompts provide a good initialization environment for TTA, but relying solely on the general information from cross-task prompts is insufficient for the model to perform well on unseen tasks. Incorporating TTA allows the model adapt to the specific nuances of each task during testing, resulting in substantial performance gains.
Impact of prompt length. We examined the influence of prompt length the performance of DPDT by investigating five distinct prompt lengths (3, 6, 30, 60, 90). Specifically, we explored the effect of prompt length variation the convergence behavior and generalization capability of the model. It is widely recognized that prompt lengths that are excessively short may impede model convergence, while prompt lengths that are overly long can result in slow convergence rates and potential overfitting. Ablation experiments revealed that prompt length of 30 is optimal. Further increasing the prompt length 60 90, however, results minor performance fluctuations but increases the convergence time. Therefore, we used prompt length of 30 for all our experiments.
Impact of model size. We explored the performance of DPDT under three model size configurations. The (3.1,128) configuration uses the pretrained model provided in the original Prompt-DT [6] to initialize DPDT, while the (24.16,768) configuration employs GPT-MIDDLE. Ta- ble shows that the size of the pretrained model param- eters correlated with the performance improvement of DPDT. When the model size is expanded to a certain extent (12,12,768), efficient parameter fine-tuning can extract ad- equate prior knowledge for downstream tasks. However, as the model complexity increases, such as when reaching the size of GPT-MIDDLEE, the model exhibits performance improvement on some tasks (Ant-dir) but decrease in per- formance others. This phenomenon can be attributed to the significant gap between the size of the dataset and the complexity of the model. Fine-tuning the model in such scenarios may encounter challenges in appropriately converging for reinforcement learning tasks, potentially leading to overfitting.
Sconss
2000
Prompt Length
Figure 3: Ablation: The effect of prompt length on DPDT's zero-shot generaliza- tion ability.
Impact of data quality. The quality of data used for fine-tuning cross prompts does indeed affect the final model performance. We conducted experiments focusing on data quality. the Cheetah-vel and MI.45 environments, we differentiated the quality of datasets into expert, medium, random, and mixed datasets. Each dataset consists of 200 time steps, which aligns with the setup for few-shot scenarios relative to the size of the training set. As shown in Table 5 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 DPDT method can extract information from suboptimal datasets to ensure model performance.
Impact of adaptation method. In addition to only utlizing cross-task prompts P. for TTA in zero-shot scenarios, we also investigated (1) combining cross- task prompts P. with the average of all task-specific prompts Ps from the training set for TTA, (2) freez- ing the cross-task prompts Pe. we initialized a new task-specific prompt combined with the tross-task prompts for TTA and (3) randomly selecting one Ps from a training task and combining it with Pa for TTA. However, we found that all of these initializa- tion methods resulted in suboptimal outcomes, shown in Table T0
Table 5: Ablation: The impact of data quality.
Cheetah-vel
ML44
expert datasets -30.10
586.84
medium datasets
.41.73
502.64
random datasets
.935.66
37.91
mixed datasets
-30.73
579.09
Impact of learning rate in prompt decompo- sition. As shown in Table 6 we present ex- perimental results on the ML45 dataset where different learning rates were applied to Pa and Ps in prompt decomposition. We observed opti- mal performance when both had the same learn- ing 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.
Table 6: Ablation: The impact of learning rate in prompt decomposition.
1rp22le-2
1rp.=lc-3
Irp.=lo-4
=le-2
310.74
307.36
311.40
Irp,=le-3 Irn,=le-4
198.17
350.99
338.21
204.94
104.07
347.21
Impact of low-rank parameter r. Table 7 presents the results of our ablation experiments focusing on the low-rank parame- ter r of prompt decomposition for Cheetah-vel and MW ML45. DPDT is relatively insensitive to selecting hyperparameters, a potential advantage of our work. As observed, the model's performance varies with different values of r. These findings suggest that the performance of DPDT remains stable across different values of r. This characteristic can be advantageous, as it allows users to implement the model without extensive tuning, streamlining the deployment process while still achieving competitive results across diverse tasks.
Table 7: Ablation: The impact of
low-rank parameter r.
Cheetah-vel
MW MI.55
1=1
-139.88
347.21
-138.08
344.10
r=10 10
-135.51
350.33
Conclusion, Limitation and Broader Impact
We have introduced novel approach, the Decomposed Prompt Decision Transformer (DPDT), aimed at efficient generalization to unseen tasks. Through the utilization of PLMs for parameter initialization and the implementation of parameter-efffiiiet multi-task prompt tuning techniques, we have successfully extracted cross-task general knowledge and further fine-tuned it on previously unseen tasks. Our experiments across various Meta-RL environments demonstrated the effectiveness of our components, achieving superior performance with significantly fewer task-specific parameters compared to fully fine-tuned methods. This approach offers robust framework for future research to further explore and optimize multi-task learning and generalization capabilities.
Limitation. Currently, our work primarily involves using large language models to initialize DPDT. While we' we've ve utilized parameter-efficient techniques to fine-tune the model and mitigate inter-domain variances, focusing optimizing these differences could potentially enhance model performance.
Broader Impact. Overall, the application of PEFT methods based on PLMs can help users obtain high-quality reinforcement learning decision models at minimal cost. However, this approach may lead to the misuse of language models when users are unaware of the inter-domain differences. potentially resulting in unforeseen negative outcomes in the RL pecision-mmking process.
Acknowledgements
This work supported by the STI 2030-Major Projects (No. 20212D0201405) the National Natural Science Foundation of China (Grant No. U23A20318 and 62276195), the Fundamental
10

---

Research Funds for the Central Universities (No. 2042024kf00399, the Science and Technology Major Project of Hubei Province under Grant 2024BAB046, and the Innovative Research Group Project of Hubei Province under Grant 2024AFA017. Tongliang Liu is partially supported by the following Australian Research Council projects: FT220100318 DP220102121, LP220100527, LP220200949, IC1901 00031. Dr. Tao's research is partially supported by NTU RSR and Start Up Grants. The numerical calculations in this paper have been done on the supercomputing system the Supercomputing Center of Wuhan University.
References
[1] Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint arXiv:2005.0 2020.
[2] Lingfeng Sun, Haichao Zhang, Wei Xu, and Masayoshi Tomizuka. Efficient multi-task and transfer reinforcement learning with parameler-compppitionnl framework. IEEE Robotics and Automation Letters, 2023.
[3] Haoran He, Chenjia Bai, Kang Xu, Zhuoran Yang, Weinan Zhang, Dong Wang, Bin Zhao, and Xuelong Li. Diffusion model is effective planner and data synthesizer for multi-task reinforcement learning. Advances in neural information processing systems, 36, 2024.
[4] Shengchao Hu, Ziqing Fan, Li Shen, Ya Zhang, Yanfeng Wang, and Dacheng Tao. Harmodt: Harmony multi-task decision transformer for offline reinforcement learning. arXiv preprint arXiv:2405. /8080, 2024.
[5] Lili Chen, 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:1 5584-150977 2021.
[6] Mengdi Xu, Yikang Shen, Shun Zhang, Yuchen Lu, Ding Zhao, Joshua Tenenbaum, and Chuang Gan. Prompting decision transformer for few-shot policy generalization. In international conference on machine learning, pages 24631-24644 PMLR, 2022.
[7] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
[8] Shengehao Hu, Li Shen, Ya Zhang, and Dacheng Tao. Prompt-mnnin decision transformer with preference ranking. arXiv preprint arXiv:2305.009644 2023.
[9] Shengchao Hu, Li Shen, Ya Zhang, Yixin Chen, and Dacheng Tao. On transforming reinforce- ment learning with transformers: The development trajectory. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2024.
[10] Suyoung Lee and Sac- Young Chung. Improving generalization in meta-rl with imaginary tasks from latent dynamics mixture. Advances in Neural Information Processing Systems, 34:27222-27235. 2021.
[11] Akari Asai, Mohammadreza Salehi, Matthew E Peters, and Hannaneh Hajishirzi. Attentional mixtures of soft prompt tuning for parameter-effiiiet multi-task knowledge sharing. arXiv preprint arXiv:2205. 1961,,1,33 3, 2022.
[12] Tu Thanh Vu, Daniel Matthew Cer, Noah Constant, Brian David Lester, and Rami Al-Rfou. Frozen model adaptation through soft prompt transfer, January 18 2024. US Patent App. 17/863,840.
[13] Michael Ahn, Anthony Brohan, Noah Brown, Yevgen Chebotar, Omar Cortes, Byron David, Chelsea Finn, Chuyuan Fu, Keerthana Gopalakrishnan, Karol Hausman, et al. Do as i can, not say: Grounding language in robotic affordances. arXiv preprint arXiv:2204001691, 2022.
[14] Ruizhe Shi, Yuyao Liu, Yanjie Ze, Simon S Du, and Huazhe Xu. Unleashing the power of pre- trained language models for offline reinforcement learning. arXiv preprint 10. 20580587 2023.
11

---

[15] Yao Wei, Yanchao Sun, Ruijie Zheng, Sai Vemprala, Rogerio Bonatti, Shuhang Chen, Ratnesh Madaan, Zhongjie Ba, Ashish Kapoor, and Shuang Ma. Is imitation all you need? generalized decision-making with dual-phase training. In Proceedings of the IEEF/CVE International Conference on Computer Vision, pages 16221-16231, 2023.
[16] Dian Chen, Dequan Wang, Trevor Darrell, and Sayna Ebrahimi. Contrastive test-time adaptation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Puttern Recognition, pages 295-305, 2022.
[17] Mengdi Xu, Yuchen Lu, Yikang Shen, Shun Zhang, Ding Zhao, and Chuang Gan. Hyper- decision transformer for efficient online policy adaptation. arXiv preprint arXiv:2304.08887, 2023.
[18] Hado Hasselt. Double q-learning. Advances in neural information processing systems, 23, 2010.
[19] Hao-nan Wang, Ning Liu, Yi-yun Zhang, Da-wei Feng, Feng Huang, Dong-sheng Li, and Yi-ming Zhang. Deep reinforcement learning: a survey. Frontiers of Information Technology & Electronic Engineering, 21(12): 2020.
[20] Dibya Ghosh, Anurag Ajay, Pulkit Agrawal, and Sergey Levine. Offline rl policies should be trained to be adaptive. In International Conference on Machine Learning, pages 7513-7530. PMLR, 2022.
[21] Haoran Xu, Xianyuan Zhan, and Xiangyu Zhu. Constraints penalized q-learning for safe offline reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pages 8753-8760, 2022.
[22] Shengchao Hu, Ziqing Fan, Chaoqin Huang. Li Shen, Ya Zhang, Yanfeng Wang, and Dacheng Tao. Q-value regularized transformer for offline reinforcement learning. arXiv preprint arXiv:2405.17098, 2024.
[23] Yueh-Hua Wu, Xiaolong Wang, and Masashi Hamaya. Elastic decision transformer. Advances in Neural Information Processing Systems, 36, 2024.
[24] Menglin Jia, Luming Tang, Bor-Chun Chen, Claire Cardie, Serge Belongie, Bharath Hariharan, and Ser-Nam Lim. Visual prompt tuning. In European Conference on Computer Vision, pages 709-727. Springer, 2022.
[25] Zhen Wang, Rameswar Panda, Leonid Karlinsky, Rogerio Feris, Huan Sun, and Yoon Kim. Multitask prompt tuning enables parameter-efficient transfer learning. arXiv preprint arXiv:2303.0288. 022661 2023.
[26] Ruihan Yang, Huazhe Xu, Yi Wu, and Xiaolong Wang. Multi-task reinforcement learning with soft modularization. Advances in Neural Information Processing Systems, 33:4767-4777, 2020.
[27] Shagun Sodhani, Amy Zhang, and Joelle Pineau. Multi-task reinforcement learning with context- based representations. In International Conference on Machine Learning, pages 9767-9779. PMLR, 2021.
[28] Shengchao Hu, Li Shen, Ya Zhang, and Dacheng Tao. Learning multi-agent communication from graph modeling perspective. arXiv preprint MrXiv:2405.08550 2024.
[29] Benjamin Eysenbach, Tianjun Zhang, Sergey Levine, and Russ R Salakhutdinov. Contrastive learning as goal-conditioned reinforcement learning. Advances in Neural Information Process- ing Systems, 35:35603-35620, 2022.
[30] Tianhe Yu, Deirdre Quillen, Zhanpeng He, Ryan Julian, Karol Hausman, Chelsea Finn, and Sergey Levine. Meta-world: A benchmark and evaluation for multi-task and meta reinforcement learning. In Conference on robot learning, pages (094-1100. PMLR, 2020.
[31] Lingfeng Sun, Haichao Zhang, Wei Xu, and Masayoshi Tomizuka. Paco: Parameter- compositional multi-task reinforcement learning. Advances in Neural Information Processing Systems, 35:21495-21555 2022.
12

---

[32] Roei Herzig, Ofir Abramovich, Elad Ben Avraham, Assaf Arbelle. Leonid Karlinsky, Ariel Shamir, Trevor Darrell, and Amir Globerson. Promptonomyvit: Multi-task prompt learning improves video transformers using synthetic scene data. In Proceedings of the IEEF/CVF Winter Conference on Applications of Computer Vision, pages 6803-6815, 2024.
[33] Tianxiang Sun, Zhengfu He, Qin Zhu, Xipeng Qiu, and Xuan-Jing Huang. Multitask pre-training of modular prompt for chinese few-shot learning. In Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Longg Papers), pages 11 1116--1117. 2023.
[34] Sheng Shen, Shijia Yang, Tianjun Zhang. Bohan Zhai, Joseph E Gonzalez, Kurt Keutzer, and Trevor Darrell. Multitask vision-language prompt tuning. In Proceedings of the IEEF/CVF Winter Conference on Applications of Computer Vision, pages 5656-5667, 2024.
[35] Manli Shu, Weili Nie, De-An Huang, Zhiding Yu, Tom Goldstein, Anima Anandkumar, and Chaowei Xiao. Test-time prompt tuning for zero-shot generalization in vision- language models. Advances in Neural Information Processing Systems, 35: 4274-14229, 2022.
[36] Shuai Wang, Daoan Zhang, Zipei Yan, Jianguo Zhang, and Rui Li. Feature alignment and uniformity for test time adaptation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 20050-20066 2023.
[37] Sanghun Jung, Jungsoo Lee, Nanhee Kim, Amirreza Shaban, Byron Boots, and Jaegul Choo. Cafa: Class-aware feature alignment for test-time adaptation. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 19060-19071, 2023.
[38] Jameel Abdul Samadh, Mohammad Hanan Gani, Noor Hussein, Muhammad Uzair Khattak, Muhammad Muzammal Naseer, Fahad Shahbaz Khan, and Salman H Khan. Align your prompts: Test-time prompting with distribution alignment for zero-shot generalization. Advances in Neural Information Processing Systems, 36, 2024.
[39] Scott Reed, Konrad Zolna, Emilio Parisotto, Sergio Gomez Colmenarejo, Alexander Novikov, Gabriel Barth-Maron, Mai Gimenez, Yury Sulsky, Jackie Kay, Jost Tobias Springenberg, et al. A generalist agent. arXiv preprint arXiv:2205.06 06175, 2022.
[40] Yujin Tang, Wenhao Yu, Jie Tan, Heiga Zen, Aleksandra Faust, and Tatsuya Harada. Saytap: Language to quadrupedal locomotion. arXiv preprint arXiv:2306.07580, 2023.
[41] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, al. Language models are unsupervised multitask learners. OpenAl blog. 1(8):9, 2019.
[42] Edward Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models. arXiv preprint arXiv:2106.
2021.
[43] Emanuel Todorov. Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In 2012 JEEE/RSJ international conference on intelligent robots and systems, pages 5026-5033 IEEE, 2012.
[44] Faraz Torabi, Garrett Warnell, and Peter Stone. Behavioral cloning from observation. arXiv preprint 1805.00 2018
[45] Brian Lester, Rami Al-Rfou, and Noah Constant. The power of scale for parameter-efficient prompt tuning. arXiv preprint arXiv:2104. 200::2200400691,,, 2021.
[46] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32, 2019.
[47] Thomas Wolf. Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, et al. Huggingface's 's transform- ers: State-of-the-art natural language processing. arXiv preprint arXiv: 19100037,77 2019.
13

---

Appendix
The appendix organized into several sections, each providing additional insights and details related to different aspects of the main work.
Detailed Environment
14
Hyperparameters configuration
15
Supplementary experiment
15
A
Detailed Environment
There are two tasks in Cheetah-dir with goal directions as forward and backward, respectively. The cheetah agent is rewarded with high velocity along the goal direction. The training and testing set are equal, and both contain the two tasks.
Cheetah-dir:
There are tasks Cheetah-vel with different goal velocities. The target velocities are uniformly sampled from the interval [0,3]. The agent is penalized with 12 errors to the target velocity. We hold out tasks to construct the testing set and train with the remaining 35 tasks.
Cheetah-vel:
Ant-dir:
There are 50 tasks in Ant-dir with different goal directions uniformly sampled in 21 space. The 8-joints ant is rewarded with high velocity along the goal direction. We sample 5 tasks for testing and leave the rest for training.
Meta-Worke MI.10: Meta-World ML10, the task is to control a Sawyer robot's end- effector to reach a target position in 3D space. The agent directly controls the XYZ location of the end-effector. Each task has a different goal position. We train in 10 tasks and test in unseen tasks.
Meta-World ML45: In Meta-World ML45, each task has different goal position. We train in 45 tasks and test in unseen 5 tasks.
Meta-World MT10: Meta-Worke MT10 comprises 10 distinct robot manipulation tasks with shared dynamics for training and 3 unseen robot manipulation tasks for testing. These tasks exhibit greater variability compared to standard meta-learning environments, posing greater challenges for extracting generalizable knowledge.
Meta- Woold MT50: The task design of Meta-War World MT50 is similar to MT10, with the key difference being an expansion of the training tasks to 45, while the number of unseen test tasks remains at 5.
Adhering to the experimental setup of Prompt-Do -T, we illustrate the task distribution for both training and testing sets in each dataset, detailed in Table Our experiments are meticulously crafted in accordance with the specifications outlined in this table.
14

---

B Hyperparameters configuration
We show the hyperparameter of DPDT and other baslines in Table 8 and Table ?
Table 8: Common Hyperparameters configuration of DPDT and DPDT-WP.
Hyperparameters
Value
Pretraining model
GPT2-small
K (length of context)
20
Prompt Length
30
training batch size for each task number of evaluation episodes for each task learning rate learning rate decay weight number of layers number of attention heads embedding dimension activation
16
5
lc-4
le-4
12
12
768
ReLU
Table 9: Common Hyperparameters configuration of MT-BC.MT-DT, Soft-prompt, HDT and Prompt- DT.
Hyperparameters
Value
K (length of context P)
20
training batch size for each task
16
number of evaluation episodes for each task
5
learning rate
le-4 (2e-5 for Prompt-DT)
learning decay weight number of layers number of attention heads embedding dimension activation
le-4 (le-5 for Prompt-DT)
12
12
768
ReLU
C Supplementary experiment
Table 10: Ablation: The impact of adaptation method. (1) Combining cross-task prompts P. with the average of all task-specific prompts Pa from the training set for TTA, (2) freezing the cross-task prompts Pc. initialized new task-specific prompt combined with the cross-task prompts for TTA and (3) randomly selecting one Ps from training task and combining it with P. for TTA.
Method
Cheetah-vel
MW ML45
MW MT50
(1)
-148.50
332.80
1482.75
(2)
-171.75
320.30
1418.50
(3)
-159.26
325.02
1079.82
DPDT
-139.88
347.21
1559.94
15

---

Table 11: Results for Meta-RL control tasks (zero-shot scenarios).
Soft-Prompt-TTA (45)
Prompt-DT--TAA (6
DPDT
Cheetah-dir
=.74
8.33+
40.32e1147
Chectah-vel
-100.5
-204.57+8.6
-139.88+1955
Ant-dir
119.14
120.07
121.8428.01
MW MT.10
251.37=617
316.74+12.05
371.01
MW ML455
301.74-18.55
299.4748.97
347.21+11.52
MW MT 10
141.99
1027.8015.58
1317.5248.2
MW MT
519.78g2097
1134.724792
1559.9442.99
Average
226.35
386.08
518.28
Table 12: Computer resources (memory, time of execution)
Method
Batch Size (each)
GPU usage
Wall Time
MT-BC
16
10.4GB
2 5-6 hours
MT-ORL
16
10.4GB
2 day
Soft-Prompt
16
20.1GB
2 hours
HDT
16
40.8GB
2 hours
Prompt-DT
10.9GB
hours
DPDT-WP
20.5GB
2-3 hours
DPDT
20.5GB
hours
MT-BC
MT-ORL
Soft-Prompt
HDT
Prompt-DT
MPDT-MP
MPDT
Cheetah-dir
Ant-dir
150
100
-50
Episode
Episode
-100
-150
50
250
500
750
1000
1250
1500
250
500
750
1000
1250
1500
Training Iteration
Training Iteration
ML 10
MT 10
Episode Return
Episode Return
1000
500
100
250
500
750
1000
1250
1500
250
500
750
1000
1250
1500
Training Iteration
Training Iteration
Figure 4: Episodic accumulated returns in four unseen tasks of MTBC, MT-ORL, Soft-Prompt, HDT, Prompt-DT, DPDT-WP and DPDT. Each method is restricted to 1500 rounds of runs each environment.
16

---

Table 13: Training and testing task indexes when testing the generalization ability in unseen tasks.
Cheetah-dir
Training set of size 2
[0, 1] [0.1]
Testing of size 2
Cheetah-vel
Training set of size 35
[0 -0-18 1,3 6,8 14, 16
24 -2,24-25,7
25,
-9]
Testing set size 5
[2, 15, 26]
Ant-dir
Training set of size 45
00- 16, 18 22, 24 29, 40, 42 49
Testing set size 5
[6, 23, 41]
Meta-Worde ML10
Training set of size
[0, 9, 19,29, 29, 33, 36, 40, 48, 49]
Testing set size 3
[1,24,,41
Meta-Worke ML45
Training set of size - 10-10, 12 16, 24, 26 35, - -40,4440 42 -99]
Testing set of size 5
41]
Meta-World MT10
("assembly-v2", "hatton-pres-toodoonn22"" "coffee-pssh---"""
Training set of size 10
"dial-turn-v2", "disassemble-v2", "door-open-v2", "hand-insert-v22,
Testing set of size 3
[" hwwpppuigsiddee222."hhemeeee- "hannnee"""""""addeeeeeee""""""" "handle-press-v2"]
Meta-World MT50
("basketball-v2,, "bin-picking-v2" "button-press-oopdown="""""
"button-presss-v2" bloffee-button-v2... "coffee-pull-v2",
"coffee-push-v2", "dial-tum-v22, "disassemble-v2",
"door-close-v2",
"door-lock-v2", "door-open-v2",
"hand-insert-v2""
"drawer-close-v2", "drawer-open-v2",
"faucct-close-v2". "handle-press-v2". "handlc-pull-sidex22",
"handlc-pull-v2",
"leverpulll2""""pee-ieettssee
Training set of size 45
"pick-place-walll-2" ", "reach-v2",
ppush-back-v2" "push-v2", "pick-place-v2",
s2'dee- seate--slie-bacc--y baaee--lld--bbc-ssiiie
"soccer-v2", "push-wall-v2", "shelf-place-v2",
"sweep-into-v22, "sweep-v2", "window-open-v2",
"hammer-v2", "peg-unplug-ssie-vv2" "reach-wall-v2",
"stick push-v2,, "stick-pull-v2", "box-close-v2"]
["
"handle-pree press---e...--s "buttonpress walll---
Testing set of size 5
"door-unlock-v2" "faucet-open-v2"]
17

---

NeurIPS Paper Checklist
1. Claims
Question:
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We introduce the scope of this paper in the abstract and the Introduction section, and summarize the contribution the Introduction section.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We state the limitation in Section 6
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory Assumptions and Proofs
Question: If you are including theoretical results, did you state the full set of assumptions of all theoretical results, and did you include complete proofs of all theoretical results?
Answer: [NA]
18

---

Justification:
Not applicable.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: If the contribution is dataset model, what steps did you take to make your results reproducible or verifiable?
Answer: [Yes]
Justification: Yes, we provide a detailed experimental description in 5.1 All the datasets, code, and model checkpoints are publicly available.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully specific model and empirical evaluation, may
might suffice, or if the contribution be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, other means that appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm. (b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibilityy the case of closed-source models, may be that access to the model limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: you ran experiments, did you include the code, data, and instructions needed to reproduce the main experimental results (cither in the supplemental material or a URL)?
19

---

Answer:
[Yes]
Justification:
All the datasets, code, and model checkpoints are publicly available.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips... cc/
cessinnnniccyl for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
/ cppppi.ccppulli//uuuiu wwwpc//ggue////eee CoodeSubeiiie ssi wsionnolli onnolicy)))) for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: If you experiments, did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
Answer: [Yes]
Justification: Both in main paper and Appendices A and
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer:
[Yes]
Justification: We provide standard deviations in numerical form in Tables and Tables Additionally, Figure includes standard deviations represented as shaded areas.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, at least for the experiments that support main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation the standard error of the mean.
20

---

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of enrors is not verified.
For asymmetric distributions, the authors should be careful to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of exccution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We provide the information in Table 12
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't make into the paper).
9. Code Of Ethics
Question:
Have you read the NeurIPS Code of Ethics https://hourips.cc/pllee EthicsCuiddlines and ensured that your research conforms to it?
Answer: [Yes]
Justification: This paper does not involve human subjects, and we have not used any personal data.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is special consid- cration due to laws or regulations their jurisdiction..
10. Broader Impacts
Question: If appropriate for the scope and focus of your paper, did you discuss potential negative societal impacts of your work?
Answer: [Yes]
Justification:
We discuss the broader impact in Section 6
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
21

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility MI..
11. Safeguards
Question: Do you have safeguards in place for responsible release of models with a high risk for misuse (e.g., pretrained language models)?
Answer: [NA]
The paper uses standard benchmark datasets and does not involve any high-risk data or models. The pre-trane pre-trained models and datasets are publicly available.
Justification:
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses
Question: If you are using existing assets (e.g., code, data, models), did you cite the creators and respect the license and terms of use?
Answer: [Yes]
Justification: We use standard benchmark datasets, which are properly credited and have open licenses.
Guidelines:
The answer NA means that the paper does existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, papprssiiiiataee thcode. com com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
22

---

If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. Assets
Question: If you are releasing new assets, did you document them and provide these details alongside the assets?
Answer: [NA]
Justification: Not applicable.
Guidelines:
The answer NA means that the paper does not release new assets,
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: If you used crowdsourcing conducted research with human subjects, did you include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: Not applicable.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. IRB Approvals
Question: Did you describe any potential participant risks and obtain Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements your institution), if applicable?
Answer: [NA]
Justification: Not applicable.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
23


Reviews:
[{"review_id": "gPZCJgFpXU", "reviewer": "Reviewer_urHM", "summary": "This paper proposes a novel Multi-Task Prompt Decision Transformer (MPDT), which leverages pre-trained language models as the initialization and adopts test-time adaptation. This approach achieves efficient generalization to unseen tasks through the prior knowledge from the pre-trained language model and by decomposing the task prompt. In the multi-task training stage, the prompt is decomposed into a cross-task prompt and a task-specific prompt, which can reduce gradient conflicts and computational load. Besides, the task-specific prompt is further decomposed into two low-rank vectors. Prompt distillation is also used to improve the quality of the prompt decomposition. In test-time adaptation, the method further optimizes the cross-task prompts on unseen tasks via alignment loss. An empirical study on Meta-RL benchmarks demonstrates the superior performance of MPDT compared to existing methods.", "questions": "- Can you explain why you use the pre-trained embedding layer for word tokens to encode the RL tokens?\n- Can you explain more about why you use prompt decomposition and prompt distillation?\n- Can you explain why you need to use a long prompt to achieve better performance?", "limitations": "See weaknesses and questions.", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- This paper leverages GPT as the initialization to provide the prior knowledge for RL tasks, which is reasonable. The empirical study demonstrates that MPDT with the pre-trained language model initialization outperforms MPDT without the language model initialization.\n- This paper utilizes prompt decomposition and only updates the prompt parameter in the multi-task training, which significantly reduces the trainable parameters. Compared with baselines, MPDT with fewer trainable parameters achieves superior performance. \n- The empirical study demonstrates that MPDT outperforms the baselines, and further analysis demonstrates the effect of the components in this method.", "weaknesses": "- This paper proposed utilizing the pre-trained language model as initialization but lacks an explanation of why the language pre-trained model could contribute to the RL tasks. \n- This method needs a long prompt to perform well compared with the Prompt DT. \n- This paper lacks detailed explanation on why using the cross-task prompt, task-specific prompt, and the prompt distillation.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n**W1: why the language model could contribute to the RL tasks?**  \nIn 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.\" \n\nAdditionally, 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.\n\n[1] https://arxiv.org/abs/2310.20587  \n[2] https://arxiv.org/abs/2201.12122\n\n**W2 & Q3: Why you need to use a long prompt to achieve better performance?**  \nIn 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.\n\nThe 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. \n\nIn 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.\n\n| prompt length | Cheetah-vel | ML45 |\n|---------------|-------------|------|\n| 3             | -299.20 (about 8.5 h) | 110.83 (about 11.0 h) |\n| 9             | -141.49 (about 4.8 h) | 348.03 (about 6.7 h) |\n| 30            | **-139.88** (about 2.4 h) | **347.21** (about 4.9 h) |\n\n[1] https://arxiv.org/abs/2206.13499\n\n**W3 & Q2: This paper lacks detailed explanation on why using the cross-task prompt, task-specific prompt, and the prompt distillation.**  \nIn 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.\n\nThe 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. \n\nWe verify the performance of the proposed component in Table 3 of the main text.\n\n**Q1: Why you use the pre-trained embedding layer for word tokens to encode the RL tokens?**  \nPlease check Author Rebuttal AR3.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "zjltqPSf7B", "reviewer": "Reviewer_zhrS", "summary": "This paper proposes a new method called Multi-Task Prompt Decision Transformer (MPDT) for efficient generalization to unseen tasks in offline reinforcement learning. MPDT involves two stages. First, the multitask training phase: MPDT is initialized with parameters from a pertained LM, which is GPT2. It decomposes the task prompt into a cross-task prompt shared across tasks and task-specific prompts. Second, the test time adaptation: The cross-task prompt is further optimized on unlabeled test tasks using test time adaptation by aligning the distributions of the test samples and training samples. The paper evaluates MPDT on seven meta RL environments from both MuJoCo and MetaWorld, showing its superior performance over baselines in generalizing to unseen tasks.", "questions": "1. How do tasks for each meta-environment differ? Do they have different reward functions $R(s_t, a_t)$ or transition function $T(s_t, a_t) \\to s_{t+1}$. It question is important for the proposed method MPDT. Because I wonder what kind of distribution shift that test-test can handle and what it cannot.\n\n2. How does the quality of the few shot prompt effect the model’s final performance?  Specifically, using expert chunk, random chunks, medium policy chunks, and mixed quality chunks. It would be better to have some analysis here.\n\nPresentation Questions:\n\n3. In algorithm 1, L_{dis} is calculated without being used?\n\n4. In Eq.(4) What is $M((P_k^*, \\tau))$? It is used without any explanation? Also, how many input does M take? Why there is a double layer bracket? \n\n5. In Line 176, the authors state that  $P_k\\in \\mathbf{R}^{l\\times s}$ is the vector multiplication of two low rank vectors $v_k\\in R^{l \\times r}$ and $u_k \\in R^{l\\times s}$. First low-rank is a matrix property, there is no low-rank vector. Second, putting the wording issue aside, what is the vector multiplication here? If it is the matrix multiplication, which is used in the low-rank decomposition methods, then the shape of the output should be $r\\times s$. Better to clarify here.", "limitations": "See my above comments for weakness and questions", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. Based on the review's knowledge, combining decision transformer and test time adaptation together for efficient multi-task RL solving is novel.\n\n2. Keeping the weights of the pretrained LM frozen is a natural idea to leverages its rich prior knowledge.\n\n3. Extensive experiments on seven Meta-RL environments demonstrate MPDT's effectiveness over DT-based baselines.\n\n4. Ablation studies individually analyze the impact of different components including prompt decomposition, distillation, and test time adaptation.\n\nOverall, MPDT appears to be a promising approach for multi-task offline RL by combining prompt-based techniques with test time adaptation in a novel way.", "weaknesses": "1. The paper appears to be hastily written and the presentation is hard to follow. (see question section)\n\n2. The authors claim that a major improvement compared with other PDTs is the decomposition of prompt into the common part and task-specific part. I was expecting to see the common part will be trained across tasks and the task-specific part is within a specific task. However, in line 186, the authors say ‘We use standard normal distribution to initialize Pc, uk and vk.’ Besides, in Eq(4), they are equally optimized for each task. I didn’t see why the common part is the slow weights and the task-specific part is the fast weight, which is claimed by authors in line 167. Possible to explain here?\n\n3. In line 192, the sentence ‘we obtain teacher task prompt $p_k^{teacher}$ for each task by using traditional prompt-tuning method individually. ‘ What method is used? How to learn it?  Is there a separate learning phase for $p_k^{teacher}$? If so, the method requires three training phases instead of two.  Also its quality matters. Can authors explain the details here?\n\n4. In line 201, for test time adaptation, ‘ we randomly select a subset X of unlabeled test samples,….’ Can authors clarify what is this subset? A sequence in the DT is organized  as (r_t, s_t, a_t, r_{t+1}, s_{t+1}, a_{t+1}…). I recommend being specific what you select?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for reviewing our work attentively. We will answer the reviewer's questions one by one in the following. \n\n**W1: The paper appears to be hastily written and the presentation is hard to follow**  \nWe will revise all unclear and erroneous statements highlighted by the reviewer further to improve the logical structure and clarity of the paper.\n\n**W2: The problem about prompt decomposition**  \nThe 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]. \n\n$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$.\n||$lr_{P_c}$=0.01|$lr_{P_c}$=0.001|$lr_{P_c}$=0.0001|\n|-|-|-|-|\n|$lr_{P_k}$=0.01|310.74|307.36|311.40|\n|$lr_{P_k}$=0.001|198.17|350.99|338.21|\n|$lr_{P_k}$=0.0001|204.94|104.07|347.21|\n\n[1] https://arxiv.org/abs/2109.04332  \n[2] https://arxiv.org/abs/2210.02390  \n\n**W3: How to obtain teacher task prompt?**  \nHere 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.\n\n**W4: Can authors clarify what is subset $X$?**  \nPlease check Author Rebuttal AR2.\n\n**Q1: How do tasks for each meta-environment differ?**  \nEach task within the environment has distinct reward functions and state transition functions.\n\nThe 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.\n\nML10, 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.\n\n**Q2: How does the quality of the few shot prompt effect the model’s performance?**  \nIntuitively, 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.\n||Cheetah-vel|ML45|\n|-|-|-|\n|expert datasets|-30.10|586.84|\n|medium datasets|-41.73|502.64|\n|random datasets|-935.66|37.91|\n|mixed datasets |-30.73|579.09|\n\nWe 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.\n\n**Q3: $L_{dis}$ is calculated without being used?**  \nWe are sorry for this mistake. The parameter update rule is $\\theta \\leftarrow \\theta-\\alpha \\nabla_\\theta \\mathcal{L}_{Total}$.\n\n**Q4: Why there is a double layer bracket?**  \n$\\mathcal{M}(P_k^*,\\tau)$ here denotes the concatenation of task prompt $P_k^*$ and trajectory $\\tau$ inputted into MPDT $\\mathcal{M}$.\n\n**Q5: What is the vector multiplication here?**  \nThe 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$.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "FL6nf6SmWp", "reviewer": "Reviewer_6jy3", "summary": "Multi-task learning is a critical pursuit in decision-making, and Decision Transformer (DT) is a popular framework in solving various decision-making problems. The author observe the suboptimal performance of prior work utilizing DT for multi-task learning, and propose a new method, Multi-Task Prompt Decision Transformer (MPDT), to alleviate this issue. MPDT is mainly composed of these components:\n- GPT2 pre-trained weights for initialization.\n- Prompt decomposition: cross-task prompt + task-specific prompt to prevent gradient conflicts.\n- Test time adaptation (TTA): dynamically optimizes the cross-task prompts during testing time.\n\nExperiments are conducted on standard Meta-RL tasks, demonstrating prominent improvements compared to prior baselines. Ablation experiments are also extensively conducted.", "questions": "- It looks unnatural to the reviewer that only the weights of prompts are trainable while the embedding layer and GPT blocks are all frozen, as shown in Figure 1. The ablation experiments of [1][2] show that parameter-efficient tuning like LoRA achieves better performance than freezing the whole block. Can the author provide some ablation experiments or reasonable explanations for this?\n- Notably, the ablation experiments in Table 3 and performance of MPDT-WP show that, only putting all these 3 components together can achieve best performance. If these components are really orthogonal and helpful, then why simply using one or two of them won't achieve overall improvements? For example, on MW ML45, only using all 3 components can beat Prompt-DT. Could the author share some insights on this?\n- How much efforts did the authors spend in tuning the hyper-parameters of MPDT?\n\n[1] Can Wikipedia Help Offline Reinforcement Learning? arXiv preprint arXiv:2201.12122.\n\n[2] Unleashing the Power of Pre-trained Language Models for Offline Reinforcement Learning. arXiv preprint arXiv:2310.20587.", "limitations": "The limitations are discussed in the paper.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Originality:\n- The MPDT framework combining pre-trained weights, prompt decomposition and TTA together is original.\n- The prompt decomposition technique to prevent gradient conflicts is orginal and well-motivated.\n\nClarity:\n- This paper is well-written and well-organized.\n- The method is easy to follow.\n\nSignificance:\n- Though pre-training DT is a popular trend in decision-making, this is the first work to successfully apply it in multi-task learning, to the knowledge of the reviewer.\n- The experimental results are prominent, compared with a set of strong baselines.", "weaknesses": "*Significant*\n- It is unclear to the reviewer that whether it is fair to compare MPDT with baselines like MT-BC, Soft-Prompt, and Prompt DT. since MPDT adopts much test-time information. How can the authors guarantee that they use the same amount of information in the test set for all baselines? If it is truly unfair, it would make the experiments less convincing. Please provide explanations for this.\n- Did the authors reproduce the results of baselines by themselves? If so, how did they pick the hyper-parameters? And for few-shot generalization, which appears to be a new benchmark designed by the authors, did the authors extensively tune the hyper-parameters of baselines?\n\n*Major*\n- It is always hard to claim the so-called \"SOTA\", which demands very rigorous statistical analysis. And it is rarely true in nowadays AI community, see https://rl-conference.cc/review_guidelines.html. Running 3 times is acceptable to present the performance, but not enough to support \"SOTA\". And the variances in Table 1,2 are very large (some [$\\mu-\\sigma$, $\\mu+\\sigma$] intervals are overlapped), thus making it hard to establish statistical improvements. The reviewer recommends removing this claim.\n- The framework is a bit complicated and engineering. Specifically, the 3 components, namely pre-training+prompt decomposition+TTA, are orthogonal and not well-connected. And thus the technical novelty of the method is a bit lacking. The novelty concentrates on prompt decomposition+distillation, while solely applying them doesn't achieve good performance (Table 3).\n\n\n*Minor*\n- The wording \"Unseen Tasks Generalization\" and \"Few-shot Generalization\" in Section 5.2 are misleading, as few-shot generatlization is also for unseen tasks.\n- The name of the proposed method, *Multi-Task Prompt DT*, might not be suitable, since it cannot reflect the differences from Prompt DT, which is already a well-developed algorithm in muti-task RL.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for the careful review of our work. \n\n**W1: How can the authors guarantee that they use the same amount of information in the test set for all baselines?**  \nIn 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.\n\nFor 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.\n||Soft-Prompt-TTA|Prompt-DT-TTA|MPDT\n|-|-|-|-|\n|Cheetah-dir|$2.91\\pm0.74$|$8.33\\pm0.41$|$50.32\\pm 11.47$ \n|Cheetah-vel|$-160.51\\pm10.13$|$-204.57\\pm8.36$|$-139.88\\pm19.65$\n|Ant-dir|$119.14\\pm10.72$|$120.07\\pm2.67$|$121.84\\pm8.01$\n|MW ML10|$251.37\\pm6.17$|$316.74\\pm12.05$|$371.01\\pm9.41$ \n|MW ML45|$301.74\\pm18.55$|$299.47\\pm8.97$|$347.21\\pm11.52$\n|MW MT 10|$541.99\\pm10.46$|$1027.80\\pm15.58$|$1317.52\\pm8.22$\n|MW MT 50|$519.78\\pm20.97$|$1134.72\\pm7.92$|$1559.94\\pm2.49$\n|Average|$226.35$|$386.08$|$518.28$ \n\nSoft-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.\n\n**W2: Did the authors reproduce the results of baselines?**  \nWe 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).\n\nFor 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.\n\n[1] https://arxiv.org/abs/2304.08487\n\n**W3: Removing the claim of \"SOTA\".**  \nThanks for the suggestions and comments. We will remove the statement regarding SOTA terminology and clarify that our results are competitive.\n\n**W4: The technical novelty and performance problem.**  \nAbout the technical novelty, please check Author Rebuttal AR1. About the performance problem in Table 3, please check Q2.\n\n**W5: The misleading problem of unseen tasks.**  \nWe 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.\"\n\n**W6: The name might not be suitable.**  \nOur 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.\n\n**Q1: It looks unnatural that only the weights of prompts are trainable while the embedding layer and GPT blocks are all frozen.**  \nPlease check Author Rebuttal AR3.\n\n**Q2: Why simply using one or two of components won't achieve improvements?**  \nThe 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.\n\n$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.\n|Decomposition|Distillation|TTA|Cheetah-vel|MW ML45|MW MT50| remark \n|-|-|--|-|-|-|-\n|$\\times$|$\\times$|$\\times$|-171.23|91.97|400.71|\n|$\\checkmark$|$\\checkmark$|$\\times$|-145.27|337.80|1304.07|$P_c\\circ P_{ka}$\n|$\\checkmark$|$\\checkmark$|$\\checkmark$|$\\mathbf{-139.88}$|$\\mathbf{347.21}$|$\\mathbf{1559.94}$|\n\n**Q3: How much efforts did the authors spend in tuning the hyper-parameters?**  \nAlmost 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.\n||Cheetah-vel|MW ML45\n|-|-|-\n|r=1|-139.88|347.21\n|r=4|-138.08|344.10\n|r=10|-135.51|350.33", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "2SY6JN2dTE", "reviewer": "Reviewer_d5sD", "summary": "This paper proposes the Multi-Task Prompt Decision Transformer (MPDT) algorithm for zero-shot multi-task offline reinforcement learning (RL). Leveraging a pre-trained language model (PLM) with prompt tuning, the MPDT innovatively decomposes multi-task prompts into task-specific and cross-task components. It also achieves zero-shot generalization through test time prompt alignment. Evaluations across various benchmarks demonstrate that MPDT outperforms prior multi-task (meta) offline RL methods.", "questions": "Please refer to the weakness part.", "limitations": "The limitations and social impact are provided in the conclusion.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is well-written and easy to follow.\n- The idea of decomposing prompts is straightforward for MTL and easy to understand.\n- The experiments are comprehensive.", "weaknesses": "- My main concerns come from the novelty of the paper. The major contribution lies in the use of cross-task prompts and task-specific prompts, specifically the prompt decomposition and prompt distillation in Section 4.1.\n- Regarding prompt decomposition, although the authors aim for cross-task prompts to contain common knowledge across tasks, the model and loss design do not ensure this. Specifically, since the authors use element-wise multiplication on $P_c$ and $P_k$, $P_c$ functions more as a common scaling factor for different task-specific prompts. After training, $P_c$ could become a constant scalar, and $P_k$ could simply scale as $1/N$ of $P_k^{teacher}$. The authors could provide the distribution of $P_c$ values to verify if it truly encapsulates common knowledge. And it is better to have a loss function to guide the common knowledge extraction during training.\n- About the test time adaptation, if the authors believe that $P_c$ contains the common knowledge, I suppose we should train a randomly initialized $P_k$ (or select one from a training task according to task similarity) for the test task using the alignment loss. This way, the model will still take $P_k \\cdot P_c$ as prompt input.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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. \n\n**W1: About the novelty of the paper**  \n\nPlease check Author Rebuttal AR1.\n\n**W2: About the prompt decomposition**  \nWe 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$.\n\nTo 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:\n\n|Method|Cheetah-vel|MW ML45|MW MT50|\n|-|-|-|-|\n|$P_c \\circ P_k^*$| -169.50|92.88|421.34|\n|MPDT | -139.88 |347.21|1559.94|\n\nIf $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.\n\nWe 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}$. \n\n$ L_{sparse}=\\lambda_1  \\sum_{k=1}^{|S|} \\sum_{i=1}^m \\sum_{j=1}^n |P_{k_{ij}}|$  \n\nWe validated the model's performance using $ L_{sparse}$ on the Cheetah-vel and MW ML45 datasets.\n|Method|Cheetah-vel|MW ML45|\n|-|-|-|\n|MPDT+$L_{sparse}$| -138.24.|350.49|\n|MPDT | -139.88 |347.21|\n\nWe 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$.\n\n| | $lr_{P_c}$=0.01 |$lr_{P_c}$=0.001 |$lr_{P_c}$=0.0001 |\n|-|-|--|--|\n|  $lr_{P_k}$=0.01| 310.74 |307.36 | 311.40|\n| $lr_{P_k}$=0.001| 198.17 | 350.99 | 338.21|\n| $lr_{P_k}$=0.0001| 204.94| 104.07 | 347.21 |\n\n**W3: About the test time adaptation**  \nWe 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.\n\n|Method|Cheetah-vel|MW ML45|MW MT50 | Is Pc frozen |Using TTA|\n|-|-|-|-|-|-|\n|(1)$P_c \\circ \\hat{P_k}$|-148.50|332.80|1482.75|×| ✓ |\n|(2) $P_c \\circ P_r$|-171.75|320.30|1418.50|✓ | ✓ |\n|(3) $P_c \\circ P_k$|-159.26 | 325.02| 1079.82 |×| ✓ |\n|(4) $P_c \\circ P_k$ | -167.80|149.21|824.07|×|×|\n|MPDT|$\\mathbf{-139.88}$ |$\\mathbf{347.21}$|$\\mathbf{1559.94}$|×|✓ |\n\n[1] https://link.springer.com/article/10.1007/s11263-024-02181-w\n\nIn the case of any follow-up questions, we will in turn provide with further clarifications.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank all the reviewers for their helpful feedback. Here, we address three main comments: innovation (AR1), the collection method of unlabeled data for TTA (AR2), and the purpose of using a pre-trained embedding layer (AR3). AR4 is the list of hyperparameters used during training and testing. The remaining questions and concerns are addressed in individual responses. In the attached PDF, we provide visualizations of the cross-task prompt and the task-specific prompts for MT10, along with introductions to 14 tasks in the MT50 dataset. All of these additional experiments and suggestions have been added to the updated main text. 

**AR1:**
The MPDT framework proposed by us encompasses the following innovations:
- Building upon prior knowledge from pre-trained models, we innovatively propose prompt decomposition, accelerating convergence through prompt distillation. Our idea of learning a transferable cross-task prompt by decomposing and distilling knowledge from multitask datasets is unique, which not only makes the prompt learning more performant but results in fewer parameters.

- Current TTA feature alignment techniques primarily focus on directly optimizing model feature layers. In scenarios with extremely limited unlabeled test samples, we innovatively use cross-task prompts as a robust initialization. By utilizing alignment-based reverse optimization of prompts, we preserve the language model's capabilities and condense task-specific information into cross-task prompts. Disruptive experimental results in the table below demonstrate that using randomly initialized prompts for TTA initialization severely compromises model performance.

- Our innovations interlock seamlessly, tightly integrating into an efficient and practical framework, paving the way for future paradigms in multi-task offline reinforcement learning.

|Decomposition| Distillation| TTA|Cheetah-vel| MW ML45 |MW MT50|remark|
|-|-|-|-|-|-|-|
| $\times$ | $\times$ | $\times$|-171.23| 91.97 | 400.71||
| $\times$ |$\times$ |$\checkmark$ | -180.39|99.82 | 644.33 |  Randomly initialized prompt|
| $\checkmark$ | $\checkmark$| $\checkmark$ | $\mathbf{-139.88}$ | $\mathbf{347.21}$ | $\mathbf{1559.94}$ ||

**AR2:**
Here, we introduce the data collection method for $X$. The model's testing phase usually occurs in a simulated environment where we predefine our expected reward values $\hat{r}$. The environment provides the initial state $s$ of the environment, consistent with the settings during inference in prompt DT methods. However, unlike in training tasks where ground-truth labels exist, for action $a_1$, we assign a value sampled randomly from the action space (which is typically consistent between training and testing tasks). Action dimension and action distribution are shown in the table below. We feed this sequence of Markov chains into the environment, obtaining rewards and the next environment states iteratively, assigning a randomly sampled value to action $a_2$ in subsequent iterations. This process is repeated $|X|$ times, resulting in data of the form $(\hat{r}_0, s_0, a_0, \hat{r}_1, s_1, a_1, \ldots, \hat{r}_{|N|}, s_{|N|}, a_{|N|})$. We consider the method of randomly sampling action values similar to assigning pseudo-labels to the data for TTA, helping the prompt understand the characteristics of the current task.

||Action dimension| Action distribution| 
|-|-|-|
|Cheetah-dir|6 |[-1,1]|
|Cheetah-vel|6 |[-1,1]|
|Ant-dir|8|[-1,1]|
|MW ML10|4|[-3,3]|
|MW MT45|4|[-3,3]| 
|MW MT10|4|[-1,1]|
|MW MT50|4|[-1,1]|

**AR3:** The purpose of using a pre-trained embedding layer is to maximize the retention of the internal information of the language model. Considering the differences between our input $(\hat{r},s,a)$ and natural language inputs, part of the function of the prompt is also to guide the language model in understanding inputs from RL tasks. To validate the effectiveness of using the pre-trained embedding layer, we adopted the RL tokens embedding layer design method from Prompt DT. During multi-task training, we concurrently trained the embedding layer, denoted as MPDT-E. Each experiment was run three times to ensure stability and reproducibility of the results.

|| Cheetah-vel|MW ML45|MW MT50|
|-|-|-|-|
|MPDT-E |-186.07| 301.49| 380.36|
|MPDT|-139.88| 347.21| 1559.94 |

We found that training a new RL tokens embedding layer resulted in performance improvements inferior to using the pre-trained embedding layer. We attribute this mainly to the retraining of the embedding layer, which disrupts the language model's understanding of input knowledge. Additionally, our inputs themselves treat $(\hat{r},s,a)$ as a single token, with temporal relationships between tokens, resembling traditional language input structures to some extent. Simultaneously, guidance from the prompt maximizes bridging the gap between these aspects, preserving the performance gains from the pre-trained embedding layer more effectively than retraining.

**AR4:**
|Hyperparameters|Value
|-|-
|K|20
|demonstration length|20
|training batch size for each task $M$|16
|number of layers|12
|number of attention heads|12
|number of gradient updates in each iteration|5
|number of evaluation episodes for each task|5
|learning rate|1e-4
|learning rate decay weight|1e-4
|activation|RELU

Meta-review: The paper originally received mixed reviews, and concerns regarding:

1. concerns with regard to the comparison to Prompt DT - the main baseline for the paper
2. reliance on long prompts and LLM for initialization
3. clarity
4. fairness of use of test-time adaptation with regards to baselines that do not use it
5. novelty
6. baselines reproduction and choice of HP
7. statistical significance of the gains

the authors seem to have been able to address the concerns of 2/3 reviewers who raised their score to Accept and Weak Accept. However, the reviewer urHM remained unconvinced and provided additional reasons to reject the paper, albeit after the reviewer-authors discussion period concluded, not allowing the authors to respond to these concerns. In light of the above, and observing some overlap between addressed concerns (as confirmed by other reviewers) and new concerns by urHM, the AC is inclined to accept the paper. Yet, the AC urges the authors to seriously consider all concerns raised by the reviewers, and incorporate all the feedback and responses into the revised version of their paper. The AC strongly suggests the authors pay special attention to additional concerns raised by urHM after the discussion period, also addressing them and incorporating them into the final revision.
Decision: Accept (poster)