Title: Decomposed Prompt Decision Transformer for Efficient Unseen Task Generalization
Venue: NeurIPS 2024
Accepted: yes
Keywords: Offline Reinforcement Learning

Prompt Tuning
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

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 .

1 Introduction

The purpose of offline reinforcement learning (Offline RL) [1] is to 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, time-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 led to increasing research interest in multi-task reinforcement learning (MTRL) [2, 3, 4]. The goal of MTRL is to develop a 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 of sequence offline RL. Additionally, they provide possibilities for integrating advancements [7] from language modeling into MTRL methodologies. Prompt-tuning DT [8] uses a gradient-free method to introduce prompts, retaining context-specific information and catering to specific preferences. These existing sequence-based offline RL methods are primarily trained and tested on the same task or perform fine-tuning using a

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 a Meta-RL setting [10].

Leveraging MTRL to extract general knowledge offers a promising approach for facilitating crosstask 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. A feasible solution might be to enhance the model's prior knowledge.

To remedy these drawbacks, we propose a prompt-based MTRL method named Decomposed Prompt Decision Transformer (DPDT), inspired by some works in natural language processing where knowledge transfer in multi-task learning scenarios is achieved through prompting strategies [11, 12]. We first employ pre-trained parameters from GPT to initialize a DPDT architecture. Incorporating the parameters of PLMs is motivated by several studies in RL [13, 14, 15]. 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 reinforcement learning tasks. Then, we design a 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 a 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 to each task's unique characteristics. By isolating the cross-task prompt from the task-specific prompts, the model ensures that updates related to general knowledge do not conflict with those related to specific tasks. Compared to [12], 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.

Our DPDT has been empirically validated in the Meta-RL setting, and the results demonstrate its superiority compared with many recent and competitive counterparts [5, 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 as 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 cross-task prompts and 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 out-of-distribution (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-driven policy learning. Decision Transformer (DT) [5] involves encapsulating rewards, states, and actions into triples and training them using autoregressive supervision on offline data. Owing to the transformer's proficiency in 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 a 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 a gradient-free approach, with the goal of preserving environment-specific details and accommodating specific preferences.

Multitask RL. Multi-Task Reinforcement Learning (MTRL) [26, 27, 28] aims to address multiple similar reinforcement learning tasks using a unified model. A straightforward approach involves developing a task-conditional multi-task model, akin to those utilized in goal-conditional RL [29] and visual-language grounding [30]. While this method has demonstrated success in certain scenarios, it often encounters challenges stemming from negative interference among tasks. PaCo [31] 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 [33, 34].

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 in the reinforcement learning domain.

3 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 [5] 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 ˆ r ∗ . However, these trajectory prompts are significantly shorter than the training trajectories. During training, for each task i , the trajectory prompt τ ∗ i,K ∗ is concatenated with the corresponding training trajectory τ i,K to form τ input i = ( τ ∗ i,K ∗ , τ i,K ) , which is then inputted into the model f for training. The concrete form of τ ∗ i,K ∗ and τ i,K are as follows:

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 a . It is noteworthy that during training, the model does not predict the actions of the trajectory prompts. The loss function is formulated as follows: a i,m denotes the actual action at the m -th timestep of the i -th task, while τ i,m -1 encompasses all data up to and including the ( m -1) th timestep in the training trajectory of the i th task.

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.

cross-task prompt

𝑷𝒄

…

ො

𝒓

𝒔

𝒂

Offline calculation of mean and variance for training samples

Pre-trained GPT Blocks

𝑷𝐜

…

…

……

…

𝑳𝒂𝒍𝒊𝒈𝒏

……

……

……

……

……

Pre-trained GPT Blocks

ො

𝒓𝒕

𝒔𝒕

𝒂𝒕

…

𝒂𝒕

…

……

Pre-trained GPT Blocks

𝒂𝒕

…

……

……

……

……

……

Pre-trained GPT Blocks

Linear Layer

ො

𝒓

𝒔

𝒂

task prompt

𝑷𝒌

∗

𝒕 - 𝑲 + 𝟏

𝒕 - 𝑲 + 𝟏

𝒕 - 𝑲 + 𝟏

. . .

𝒕

𝒕

𝒕

𝒂 ∗

𝒂

𝒂

𝟏𝟐  ×

𝒕 - 𝑲 + 𝟏

𝒕

task specific prompt

𝑷𝒌

𝒗𝒌

𝒖𝒌

teacher prompt

𝑷𝒌

𝒕𝒆𝒂𝒄𝒉𝒆𝒓

Distillation

Pre-trained GPT Blocks

𝒂 ∗

𝒂 ∗

Test Time Adaptation

𝒂 ∗

𝒂𝒕

𝒂 ∗

𝒂𝒕

ො

𝒓𝒕

𝒔𝒕

𝒂𝒕

During the inference phase, multiple randomly augmented views are created from the provided testing sample X test . Predictions with entropy below a specified threshold are retained, while other views are filtered out using a confidence selection criterion. The average entropy of the filtered predictions is then used to unsupervisedly update the prompts p . Some methods [36, 37, 38], considering the data instability caused by augmentation techniques, use 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, we provide a 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 a backward pass. Once optimized, the cross-task prompts serve as a 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 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. Some existing work has already explored this avenue [13, 39, 40]. Taking inspiration from this, we employ GPT2-SMALL [41] to initialize our DPDT and maintain these parameters frozen throughout training. It is 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 a 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.

Prompt Decomposition. Given a set of training tasks S = { S 1 , S 2 , . . . , S n } , our objective is to learn a general prompt P c 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 P c ∈ R l × s and P k ∈ R l × s . The task prompt P ∗ k for the k -th task is obtained by taking the element-wise product of P c and P k . 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 P c aims to acquire general knowledge from S , while the task-specific prompt P k allows task k to retain its unique knowledge. The task prompts parameterization of the k -th training task is expressed as:

Algorithm 1 Decomposed Prompt Tuning

Input : Training task set S , Offline datasets D M , Batch size M , Learning rate α , training iterations N , teacher task prompts p teacher k .

Initialize : Initialize a 12-layer, 12-head DPDT M using GPT2-SMALL, randomly initialize cross-task prompts P c and low-rank vectors v k , u k .

for t = 1 to N do for k in S do

Select a trajectory τ that contains M samples in task k .

Calculate P ∗ k by Equation 3.

Calculate L MSE and L dis according to Equations 4 and 5. Computed loss function by Equation 6.

θ

end for end for

In the specific implementation process, inspired by LORA [42], we further decompose each taskspecific prompt into two lowrank vectors v k ∈ R l × r and u k ∈ R l × s using a low-rank method. We obtain P k through vector multiplication. Here l represents the prompt length, r represents the hidden layer dimension, and s represents the prompt dimension. The hyperparameter r is a manually specified lowrank 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 P c , u k and v k . Given the DPDT M , the mean squared loss between the model's predicted actions and the true actions is calculated as:

Prompt distillation. 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 employed knowledge distillation techniques to compel the cross-task and task-specific prompts to learn their respective information. We obtained teacher task prompts p teacher k for each task k by using traditional prompttuning methods individually. During training, the mean squared error is calculated directly between p teacher k and p ⋆ k :

The total loss function for training task prompts for obtaining a cross-task prompt to be transferred to the target side is then:

where λ is a weight to balance the impact of distillation loss terms. In our experiments, we set λ to 0.5. The overall summary of the multitask training algorithm is presented in Algorithm 1.

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 t in the test task set T , we randomly select a subset X of unlabeled test samples, combine them with the cross-task prompts P c and input them into the model.

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 ˆ r . The environment

θ

-

α

∇

θ

L

Total

.

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). 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 (ˆ r 0 , s 0 , a 0 , ˆ r 1 , s 1 , a 1 , . . . , ˆ r | 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 P c for the given test task through this alignment loss. For each test task t , we denote the distribution of the test samples as T and the distribution of the training samples as D . Specifically, we calculate the aligned token mean and variance via:

Here, H l,i represents the state of the i th sample at the l th hidden layer, while µ l ( T ) and σ 2 l ( T ) denote the mean and variance of all test samples at the l th hidden layer, respectively. Similarly, for each hidden layer of the model, we pre-calculate the statistical measures of mean µ l ( D ) and variance σ 2 l ( 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 training samples and labels are accessible. The formula for calculating the alignment loss function is as follows:

Algorithm 2 Test Time Adaptation

Input : Test samples set X , Cross-task prompts P c , µ l ( D ) , σ 2 ( D ) , The number of layers L .

l 1: for l = 1 to L do 2: for i in X do 3: Calculate H l,i obtained by inputting the concatenation of P c and i into DPDT. 4: end for 5: end for 6: for l = 1 to L do 7: Compute µ l ( T ) and σ 2 l ( T ) by Equation 7. 8: end for 9: Compute token distribution alignment loss by Equation 8. 10: Optimize L align to update P c .

The test time adaptation phase process is illustrated in Algorithm 2.

5 Experiment

In this section, we present an 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 learning 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 RL, such as sparse rewards, complex state spaces, and precise control of robotic systems. Our experiments on the Cheetah-dir, Cheetah-vel, and Ant-dir environments in the MuJoCo benchmark meticulously adhere to the datasets and methodologies outlined in 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 MetaRL. A detailed description of 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

strategy to 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 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 1, 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 a variant of DPDT that does not use GPT-SMALL parameters for initialization, referred to as DPDTWP (DPDT-Without Pretrained). Soft prompts adapt to tasks in a parameter-efficient 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 is 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 11). Soft-Prompt-TTA showed performance improvements across all tasks, whereas Prompt-DT-TTA experienced performance declines in some tasks. The main reason for this is that Prompt-DT relies on high-quality trajectory data for prompts during testing, and applying TTA on unlabeled data may have adversely affected prompt optimization.

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 c . Instead, P c is fine-tuned

0

250

500

750

1000

1250

1500

Training Iteration

450

400

350

300

250

200

150

Cheetah-vel

0

250

500

750

1000

1250

1500

Training Iteration

50

0

50

100

150

200

250

300

350

ML 45

0

250

500

750

1000

1250

1500

Training Iteration

0

250

500

750

1000

1250

1500

MT 50

MT-BC

MT-ORL

Soft-Prompt

HDT

Prompt-DT

MPDT-MP

MPDT

directly on a small number of labeled test samples through a self-supervised 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 fewshot 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-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 zero-shot 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 as well as 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 it with decomposed prompts P k and P c without distillation (third row) resulted in performance improvements 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-task knowledge essential for enhancing target downstream tasks.

To evaluate the impact of prompt distillation, we trained a 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 baseline (first row), incorporating prompt distillation (second row)

resulted in a 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 in 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 to 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 on the performance of DPDT by investigating five distinct prompt lengths (3, 6, 30, 60, 90). Specifically, we explored the effect of prompt length variation on 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 a prompt length of 30 is optimal. Further increasing the prompt length to 60 or 90, however, results in minor performance fluctuations but increases the convergence time. Therefore, we used a 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. Table 4 shows that the size of the pretrained model parameters is 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 adequate prior knowledge for downstream tasks. However, as the model complexity increases, such as when reaching the size of GPT-MIDDLE, the model exhibits performance improvement on some tasks (Ant-dir) but a decrease in performance on others. This phenomenon can be attributed to the significant gap between the size of the dataset and

3

9

30

60

90

Prompt Length

800

600

400

200

0

200

400

Average Scores

Cheetah-vel

Ant-dir

MW ML45

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.

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. In the Cheetah-vel and ML45 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, we found that models fine-tuned

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

Impact of adaptation method. In addition to only utilizing cross-task prompts P c for TTA in zero-shot scenarios, we also investigated (1) combining crosstask prompts P c with the average of all task-specific prompts P k from the training set for TTA, (2) freezing the cross-task prompts P c , we initialized a new task-specific prompt combined with the cross-task prompts for TTA and (3) randomly selecting one P k

from a training task and combining it with P c for TTA. However, we found that all of these initialization methods resulted in suboptimal outcomes, shown in Table 10.

Impact of learning rate in prompt decomposition. As shown in Table 6, we present experimental results on the ML45 dataset where different learning rates were applied to P c and P k in prompt decomposition. We 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.

Impact of low-rank parameter r . Table 7 presents the results of our ablation experiments focusing on the low-rank parameter 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.

6 Conclusion, Limitation and Broader Impact

We have introduced a 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-efficient 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 a 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've utilized parameter-efficient techniques to fine-tune the model and mitigate inter-domain variances, focusing on 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 decision-making process.

References

Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint arXiv:2005.01643 , 2020.

Lingfeng Sun, Haichao Zhang, Wei Xu, and Masayoshi Tomizuka. Efficient multi-task and transfer reinforcement learning with parameter-compositional framework. IEEE Robotics and Automation Letters , 2023.

Haoran He, Chenjia Bai, Kang Xu, Zhuoran Yang, Weinan Zhang, Dong Wang, Bin Zhao, and Xuelong Li. Diffusion model is an effective planner and data synthesizer for multi-task reinforcement learning. Advances in neural information processing systems , 36, 2024.

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.18080 , 2024.

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:15084-15097, 2021.

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-24645. PMLR, 2022.

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 are few-shot learners. Advances in neural information processing systems , 33:1877-1901, 2020.

Shengchao Hu, Li Shen, Ya Zhang, and Dacheng Tao. Prompt-tuning decision transformer with preference ranking. arXiv preprint arXiv:2305.09648 , 2023.

Shengchao Hu, Li Shen, Ya Zhang, Yixin Chen, and Dacheng Tao. On transforming reinforcement learning with transformers: The development trajectory. IEEE Transactions on Pattern Analysis and Machine Intelligence , 2024.

Suyoung Lee and Sae-Young Chung. Improving generalization in meta-rl with imaginary tasks from latent dynamics mixture. Advances in Neural Information Processing Systems , 34:27222-27235, 2021.

Akari Asai, Mohammadreza Salehi, Matthew E Peters, and Hannaneh Hajishirzi. Attentional mixtures of soft prompt tuning for parameter-efficient multi-task knowledge sharing. arXiv preprint arXiv:2205.11961 , 3, 2022.

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.

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 as i say: Grounding language in robotic affordances. arXiv preprint arXiv:2204.01691 , 2022.

Ruizhe Shi, Yuyao Liu, Yanjie Ze, Simon S Du, and Huazhe Xu. Unleashing the power of pretrained language models for offline reinforcement learning. arXiv preprint arXiv:2310.20587 , 2023.

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 IEEE/CVF International Conference on Computer Vision , pages 16221-16231, 2023.

Dian Chen, Dequan Wang, Trevor Darrell, and Sayna Ebrahimi. Contrastive test-time adaptation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 295-305, 2022.

Mengdi Xu, Yuchen Lu, Yikang Shen, Shun Zhang, Ding Zhao, and Chuang Gan. Hyperdecision transformer for efficient online policy adaptation. arXiv preprint arXiv:2304.08487 , 2023.

Hado Hasselt. Double q-learning. Advances in neural information processing systems , 23, 2010.

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):1726-1744, 2020.

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.

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.

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.

Yueh-Hua Wu, Xiaolong Wang, and Masashi Hamaya. Elastic decision transformer. Advances in Neural Information Processing Systems , 36, 2024.

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.

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.02861 , 2023.

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.

Shagun Sodhani, Amy Zhang, and Joelle Pineau. Multi-task reinforcement learning with contextbased representations. In International Conference on Machine Learning , pages 9767-9779. PMLR, 2021.

Shengchao Hu, Li Shen, Ya Zhang, and Dacheng Tao. Learning multi-agent communication from graph modeling perspective. arXiv preprint arXiv:2405.08550 , 2024.

Benjamin Eysenbach, Tianjun Zhang, Sergey Levine, and Russ R Salakhutdinov. Contrastive learning as goal-conditioned reinforcement learning. Advances in Neural Information Processing Systems , 35:35603-35620, 2022.

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 1094-1100. PMLR, 2020.

Lingfeng Sun, Haichao Zhang, Wei Xu, and Masayoshi Tomizuka. Paco: Parametercompositional multi-task reinforcement learning. Advances in Neural Information Processing Systems , 35:21495-21507, 2022.

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 IEEE/CVF Winter Conference on Applications of Computer Vision , pages 6803-6815, 2024.

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: Long Papers) , pages 11156-11172, 2023.

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 IEEE/CVF Winter Conference on Applications of Computer Vision , pages 5656-5667, 2024.

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:14274-14289, 2022.

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-20060, 2023.

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.

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.

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.06175 , 2022.

Yujin Tang, Wenhao Yu, Jie Tan, Heiga Zen, Aleksandra Faust, and Tatsuya Harada. Saytap: Language to quadrupedal locomotion. arXiv preprint arXiv:2306.07580 , 2023.

Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAI blog , 1(8):9, 2019.

Edward J 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.09685 , 2021.

Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In 2012 IEEE/RSJ international conference on intelligent robots and systems , pages 5026-5033. IEEE, 2012.

Faraz Torabi, Garrett Warnell, and Peter Stone. Behavioral cloning from observation. arXiv preprint arXiv:1805.01954 , 2018.

Brian Lester, Rami Al-Rfou, and Noah Constant. The power of scale for parameter-efficient prompt tuning. arXiv preprint arXiv:2104.08691 , 2021.

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.

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 transformers: State-of-the-art natural language processing. arXiv preprint arXiv:1910.03771 , 2019.

Appendix

The appendix is organized into several sections, each providing additional insights and details related to different aspects of the main work.

A Detailed Environment

Cheetah-dir : 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-vel : There are 40 tasks in Cheetah-vel with different goal velocities. The target velocities are uniformly sampled from the interval [0,3]. The agent is penalized with l2 errors to the target velocity. We hold out 5 tasks to construct the testing set and train with the remaining 35 tasks.

Ant-dir : There are 50 tasks in Ant-dir with different goal directions uniformly sampled in 2D 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-World ML10 : In Meta-World ML10, the task is to control a Sawyer robot's endeffector 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 3 tasks.

Meta-World ML45 : In Meta-World ML45, each task has a different goal position. We train in 45 tasks and test in unseen 5 tasks.

Meta-World MT10 : Meta-World 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-World MT50 : The task design of Meta-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-DT, we illustrate the task distribution for both training and testing sets in each dataset, as detailed in Table 13. Our experiments are meticulously crafted in accordance with the specifications outlined in this table.

B Hyperparameters configuration

We show the hyperparameter of DPDT and other baslines in Table 8 and Table 9.

C Supplementary experiment

0

250

500

750

1000

1250

1500

Training Iteration

150

100

50

0

50

Cheetah-dir

0

250

500

750

1000

1250

1500

Training Iteration

50

0

50

100

150

Ant-dir

0

250

500

750

1000

1250

1500

Training Iteration

0

100

200

300

400

ML 10

0

250

500

750

1000

1250

1500

Training Iteration

0

500

1000

MT 10

MT-BC

MT-ORL

Soft-Prompt

HDT

Prompt-DT

MPDT-MP

MPDT

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 in the Introduction section.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in 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.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of 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 in 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 on 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 a few datasets or with a 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, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a 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 be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

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]

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 crossreferenced.

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 to 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 a dataset or 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 to 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 might suffice, or if the contribution is a specific model and empirical evaluation, it may 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 a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to 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: If you ran experiments, did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?

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 NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a 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: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to 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 a 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 as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: If you ran 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 B.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: We provide standard deviations in numerical form in Tables 1 and Tables 2. Additionally, Figure 2 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, confidence intervals, or statistical significance tests, at least for the experiments that support the 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, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

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 a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not 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

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: We provide the information in Table 12.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as 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 or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Have you read the NeurIPS Code of Ethics https://neurips.cc/public/ EthicsGuidelines 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 NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in 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 or 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.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a 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 is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are 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 to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

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]

Justification: The paper uses standard benchmark datasets and does not involve any high-risk data or models. The pre-trained models and datasets are publicly available.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a 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 to 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 a 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 not use 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 a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.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, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a 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.

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/model 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 or conducted research with human subjects, did you include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: Not applicable.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, 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 of your institution), if applicable?

Answer: [NA]

Justification: Not applicable.

Guidelines:

The answer NA means that the paper does not 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 NeurIPS 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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\nPlease refer to the weakness part.\n\nLIMITATIONS:\nThe limitations and social impact are provided in the conclusion.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-06-15"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nMulti-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.\n\nSTRENGTHS:\nOriginality:\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.\n\nWEAKNESSES:\n*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.\n\nQUESTIONS:\n- 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.\n\nLIMITATIONS:\nThe limitations are discussed in the paper.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-06-25"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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?\n\nQUESTIONS:\n1. 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.\n\nLIMITATIONS:\nSee my above comments for weakness and questions", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nSee weaknesses and questions.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-13"}