Pre-Trained Multi-Goal Transformers with Prompt Optimization for Efficient Online Adaptation
Authors: Haoqi Yuan
Yuhui Fu
Feiyang Xie
Zongqing Lu
Venue: NeurIPS 2024 poster (2024)
Keywords: reinforcement learning
pre-training
Transformers
Primary area: reinforcement_learning

Abstract: Efficiently solving unseen tasks remains a challenge in reinforcement learning (RL), especially for long-horizon tasks composed of multiple subtasks. 
Pre-training policies from task-agnostic datasets has emerged as a promising approach, yet existing methods still necessitate substantial interactions via RL to learn new tasks.
We introduce MGPO, a method that leverages the power of Transformer-based policies to model sequences of goals, enabling efficient online adaptation through prompt optimization.
In its pre-training phase, MGPO utilizes hindsight multi-goal relabeling and behavior cloning. This combination equips the policy to model diverse long-horizon behaviors that align with varying goal sequences.
During online adaptation, the goal sequence, conceptualized as a prompt, is optimized to improve task performance. We adopt a multi-armed bandit framework for this process, enhancing prompt selection based on the returns from online trajectories.
Our experiments across various environments demonstrate that MGPO holds substantial advantages in sample efficiency, online adaptation performance, robustness, and interpretability compared with existing methods.

Pre-Trained Multi-Goal Transformers with Prompt Optimization for Efficient Online Adaptation
Haoqi Yuan
Yuhui Fu'
Feiyang Xie
Zongqing Lu Lul.3+ 13+
School of Computer Science, Peking University
Yuanpei College, Peking University
Beijing Academy of Artificial Intelligence
Abstract
Efficiently solving unseen tasks remains a challenge in reinforcement learning (RL), especially for long-horizen tasks composed of multiple subtasks. Pre-training poli- cies from task-agnostic datasets has emerged as a promising approach, yet existing methods still necessitate substantial interactions via RL learn new tasks. We introduce MGPO, a method that leverages the power of Transformer-based policies model sequences of goals, enabling efficient online adaptation through prompt optimization. In its pre-training phase, MGPO utilizes hindsight multi-goal rela- beling and behavior cloning. This combination equips the policy to model diverse long-horizon behaviors that align with varying goal sequences. During online adaptation, the goal sequence, conceptualized as a prompt, is optimized improve task performance. We adopt a multi-armes bandit framework for this process, enhancing prompt selection based on the returns from online trajectories. Our experiments across various environments demonstrate that MGPO holds substantial advantages sample efficiency, online adaptation performance, robustness, and interpretability compared with existing methods.
Introduction
In the evolving landscape of deep learning, the paradigm of pre- pre-training followed by fine-tuning has become a dominant approach for improving learning downstream tasks, particularly in the domain
of computer vision 20 62] [4 and natural language processing (NLP) 30 32). This paradigm has recently been explored in deep reinforcement learning (RL)
addressing the issue of sample efficiency in RI. when solving unseen tasks by leveraging the acquired knowledge during pre-training. For example, offline meta-RL. (OMRL) multi-task datasets, which adapts to new task with limited interactions in this task. However, these approaches typically require extensive data collection for each specific task. In contrast, given the relative ease of acquiring large, task-agnostic datasets that contain diverse behaviors. other studies on task aask-agnostic datasets and adapt to unseen
56 studies pre-training policy on
focus on pre- pre-traiming policies 4 54 and skills 38 47 tasks with RL finetuning.
Most works [56] (55) especially in the context of OMRL, primarily focus on short-term tasks with shaped rewards, where information of the unknown task can be inferred within a few steps of online interactions (e.g., MuloCo (50) and Meta World (88). These approaches often fall short of tackling long- long-horizon tasks which are characterized by sequence of sub-processes. In our study, we consider this challenging yet realistic scenario: given a aggnosicc dataset characterized by diverse behaviors of the agent, we aim to pre-train policy to facilitate efficient online adaptation to new, long-horizon
tasks.
*Correspondence Zongqing Lu <zongging. @uw..u.du....cn.....
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Recent studies [8] 60] tackle this problem by assuming that diverse short-term skills can be acquired from task-agn datasets, which are sufficient compose the required behaviors for complex. long-horizon tasks. These methods pre-train a goal-conditioned policy focused on short-term skills. During online adaptation, they use RI. to train a separate high-level policy that selects goals based on current states and executes the pre-trained policy for multiple steps. While this approach avoids the extensive finetuning of the pre-trained policy, still requires substantial number of interaction steps for online RL. This limitation arises from the pre-trained policy's inability to sequentially achieve multiple goals during a single attempt. The capability to switch goals deciding at which state to transition to the next goal must be developed during the costly online adaptation phase.
it feasible to pre-train a policy conditioned on sequence of goals, capable of autonomousty transitioning between goals while executing long-horizon behaviors aligned with these goals? Achieving such a policy would pave the way for developing efficient online adaptation algorithms, eliminating the need for learning high-level policy for state-conditioned goal switching. However, the challenge lies in the substantial demand on the policy's long-term memory capabilities, which must effectively remember the goal sequence and base action prediction long history context.
To tackle this challenge, we draw inspiration from the success of Transformers [5] in language mod- eling [0 77. which have demonstrated a remarkable capability to model long sequences and, when pre-trained on diverse datasets, adapt efficiently downstream tasks through prompt optimization [25]. In our proposed Multi-Goal Transformers with Prompt Optimization (MGPO), we pre-train a Transformer-bast policy that takes a goal sequence as a prompt and predicts actions the following sequence of environment observations. During online adaptation, we only optimize for the sequence of goals the prompt, drawing parallels with the concept of prompt optimization in language models.
In the pre-training stage, we employ hindsight relabeling to construct prompts from goals visited in the trajectory and train the Transformer policy through behavior cloning. Thus, the policy learns to produce behaviors consistent with the goal sequence. The number of goals in the prompt modulates the deterministic or exploratory nature of the policy's behavior. For online adaptation, we propose novel prompt optimization strategy, leveraging the returns of online trajectories to guide the selection of goal sequences. This process formulated as a multi- armed bandit problem and we introduce two approaches through upper confidence bound (UCB) and Bayesian posterior estimation. Figure provides an overview of our framework.
We evaluate MGPO across diverse domains, including maze navigation, the robotic simulation envi- ronment Kitchen, and the open-world game Crafter. Our results demonstrate that MGPO significantly surpasses prior methods in terms of sample efficiency and performance during online adaptation. MGPO adapts to new tasks within a small budget of 100 online episodes in all environments. Com- parisons to existing prompt optimization methods highlight the interpretability and robustness of our method.
In summary, our main contributions are:
We propose pre-training multi-goal Transformers to address the challenge of online adap- tation in long-horizon tasks. Our approach combines the strengths of Transformers in sequence modeling and the advantages of goal-conditioned pre-training for task adaptation.
We introduce novel methods for optimizing goal sequences, offering enhanced interpretabil- ity and robustness compared to existing prompt optimization methods.
Our experimental results in diverse, challenging environments demonstrate that MGPO significantly enhances sample efficiency over existing methods.
2 Preliminaries
2.1 Problem Formulation
A task in an environment is formulated as a partially observable Markov Decision Process (POMDP) M = (S. O, T.P.P. r.7) representing the state space, the agent's observation space, the action space, the transition probability of the environment, the initial state distribution, the reward function, and the discount factor, respectively. Starting from the initial state, the agent takes an action at each timestep, and then the environment transitions to the next state and returns reward. At each timestep

---

max log Rg (alp.t)
Pre- Pre-Trained Transformer
T gnostic Dataset
a2
axy
Environment
Multi- Moal--Goal Transformer
R
0000
p-t
P
FS
Multi-Goal Relabeling
Pre-Training
Online Prompt Optimization
Figure An overview of MGPO. We pre-train Transformer- based policies on task-agnostic datasets, leveraging hindsight multi-goa relabeling and behavior cloning to endow policies with the capacity for modeling long-term behaviors. In unseen tasks, multi-goal Transformers adapt efficiently through prompt optimization, which searches for a sequence of goals with the aim of maximizing online returns.
the environment provides the agent an observation o, via an emission function or E() This process continues until the task either terminates or reaches a maximum timestep H. We denote the historical observations and actions as f = {(a,, a,) ) Reinforcement learning (RL) aims to learn policy π (ayl/4--1 o₂) maximize the expected return J (T) (a)hro)()))))))) [R]. where
R = E 7"r(( (2) is the discounted cumulative return of an episode.
We focus long- honghhorizon tasks that require executing substantial number of subtasks sequentially. We define a goal function g = (o) that maps an observation to goal, representing key information related to task completion, such the agent's position in navigation tasks the states of objects in robotic scenarios. G denotes the goal space. A goal g is said to be reached at time t = ff((o) We assume that the task M in unknown environment provides task goal gM, indicating the final goal to be reached. To optimally solve the long-horizon task, the policy must sequentially reach multiple goals. As example, consider an environment resembling rooms with unknown structures where the task is to navigate to goal location ga Under partial egocentric observations, the agent should learn from trial and error to find the shortest path to the goal location, which involves reaching several specific waypoints (g1....,90 Ik) as necessary goals.
During pre-trainin, assume access to task-agnostic dataset D {m.) consisting of trajecto- ries T = (01, 41, ..., оμ, az) ), in which the agent performs diverse behaviors. For online adaptation, we aim to solve unseen task M with its goal gM provided, within online interaction budget of N episodes. The objective is to find the optimal policy π maximizing the expected return J J(T).
2.2 Transformers and Prompt Optimization
Some previous methods [4 pre-train policy no (del/ti, 0:) parameterized by and then finetune with online RL. This approach is sample-inefficien due discrepancies between the behaviors in the dataset and those required for downstream tasks. Other methods [8] 60 pre-train a policy (o(a, h,-1, ot) conditioned
variable to perform diverse short-term skills and then train high-level policy TT(= (4, ot) online using RL. As discussed earlier, this approach often results in limited sample efficiency because the high-level policy must learn when switch skills based the observations during the online phase.
In this work, we explore pre-training a policy capable of performing diverse long-horizon behaviors, aiming to develop efficient online adaptation methods that do not require additional RL. To model such diverse long trajectories, we employ the Transformer architecture in policy, known for its effectiveness in sequence modeling tasks. Transformers leverage an attention mechanism [5T], en- abling the model to weigh different parts of the input sequence differently, thus effectively addressing long-term dependencies.
In mre-training, we aim to develop Transformer- based policy ( (p, he-1, parameterized by taking an additional input variable p. This policy processes inputs an auto-regressive manner, where the input is sequence (p, 01, (i,,.., OH, an)), and outputs action a, at each observation or- Here, the input variable acts similarly to prompt in language models, guiding the policy's 's behavior. During online adaptation, keep the model parameters 0 fixed and treat p the optimization

---

variable, transforming the problem into finding an optimal prompt
max J (mo (oelp, 01))
(I)
This approach mirrors the concept of prompt optimization in language models and eliminates the need to use online RL to train a high-level policy. policy.
3 Method
3.1 Pre-Training Multi-Goal Transformers
We implement a prompt as a sequence of goals, describing an ordered sequence of key states to be visited, which abstracts the long-term behavior of the agent. Inspired by recent works [28] 60] that employ hindsight relabeling [1] to generate goals from offline trajectories, we sample a sequence of goals within the trajectory to construct the prompt.
Using the observations (o,...., OH) in a trajectory т. we construct a sequence of goals = (g1 ... gH) with the same length H. representing the process of agent behavior in this episode. We first uniformly sample a number ~ U|00, K - where K the maximal prompt length, then uniformly sub-sample the goal sequence to construct prompt p- (gt, ........ ...::9 1 99 where 0 S i1 S ... is H. Since the trajectory ends at the last goal Ди. we keep 9H at the last position in the prompt. Using this sampling mechanism, the prompt describe behavior in different granularities with its varying length. We denote the process of sampling a prompt from the trajectory as P(plT).
we adopt causal Transformer to build the policy To(a; [p, it) which takes as input sequence concatenating the prompt and the trajectory: (p, T) (Di, .... : 9H 01, 01, .... 0z aH) ). At each input or, the policy only sees the sub-sequence from gi, to or due to the causal attention mask and predicts the action distribution Th p, h4- ot). We sample batches of trajectories in D and train the policy using behavior cloning:
Similar to Decision Transformers g
H
max EEDD, p(ple))
Σ
log To h o.)
This training scheme encourages the policy to utilize the information provided in the prompt to reduce uncertainty action prediction. After pre-training, it learns to perform behaviors that follow the sequence of goals in the prompt and seamlessly alter behavior between different goals. Prompts with different numbers of goals provide varying amounts of information to match behavior [14, thereby introducing different levels of uncertainty in action, enabling trade-off between exploration and exploitation in online adaptation.
AMAGO [5] also adopts similar multi-goal relabeling strategy to train RL policies. While it utilizes hindsight relabeling for better exploration, we are the first to use this scheme to pre-train policies capable of stitching different goals sequentially, thereby facilitating efficient task adaptation.
3.2 Online Prompt Optimization
Given unseen task with the task goal gMM, we start with an initial prompt Po (g*M) and aim to find an optimal prompt p" with a maximal length of K that maximizes the expected return. The algorithm can evaluate prompts for N episodes to return
optimized prompt. In principle, we can employ any black-box box optimization approach for this purpose, including discrete prompt search 19, 40) and continuous prompt-tuning [48] 27 methods. However, prompts optimized in the entire prompt space GK or even in a continuous space can contain uninterpretable goals and are out of training distribution. The policy with such prompts may produce unpredictable behaviors.
To enhance interpretability and robustness, we propose novel method that samples prompts from online collected trajectories. aligning more closely with the training distribution of prompts. We assume that, we condition the policy on prompts sampled from a trajectory with high return, is likely to yield high expected returns since it performs similar behavior of this trajectory. Thus, we propose a method that alternates between exploring online trajectories for high returns and sampling new prompts from the best seen trajectory of the highest return.

---

Formally, we maintain a buffer B of sampled prompts and their history of returns and the trajectory r* of the highest return R Each iteration involves: (1) selecting prompt p from B most likely to yield high returns and collecting a trajectory with x <alpp" o); (2) sampling a new prompt p' ~ P(plr"), which then added to the buffer with the return obtained with x(alp' h o). Prompt selection in (1), which requires acting under uncertainty given observed returns in history, can be modeled as multi-arm bandit (MAB) problem. We implement two optional solutions including upper-confidence bound (UCB) [2] and method based Bayesian posterior estimation (BPE) [36]. where the former selects the prompt with the highest UCB of expected return and the latter selects the prompt most possible to yield return higher than R based on the estimated posterior of its return distribution.
In summary, our method employs several key designs to optimize prompts effectively: (1) Trajectory- based sampling: Instead of exploring the combinatorial space of goal sequences, we restrict our search to prompts derived from collected trajectories, ensuring both feasibility and relevance. (2) Reward- guided exploration: We further refine our prompt search by selecting prompts from trajectories that have highest returns, thereby enhancing the likelihood of performance improvement. (3) Task- goal consistency: We maintain the final goal within each prompt as the task goal, ensuring that all exploration efforts are aligned with task completion.
Our prompt optimization method is detailed in Algorithm in Appendix tions of UCB and BPE are also provided.
E.1 where the implementa-
4 Experiments
In this section, we present experimental results obtained across various domains to evaluate the efficacy of MGPO. We aim to answer three questions: (1) Does MGPO improve sample efficiency in solving new tasks compared to previous methods? (2) How does our proposed prompt optimization method compare with existing methods? (3) How does each component in MGPO contribute to efficient online adaptation?
4.1 Environments and Datasets
Our evaluation spans multiple domains featuring long-horizon tasks. We collect datasets pre-train models and evaluate their online adaptation capabilities on test sets of unseen tasks or environment configurations. Detailed descriptions of these environments and datasets are available in Appendix C
MazeRunner: A 2D Maze navigation environment with partial observation as introduced in [15]
The maze has randomly generated walls, where the task is to reach designated goal position (z, y). The agent's 'ss observation includes its position and nearby terrain, receiving a reward of +1 for reaching the goal and a 1 penalty for each timestep. We collect the dataset with a handcrafied policy that explores various goals within each trajectory. Online adaptation requires finding optimal paths in unknown mazes with limited trials. We test MazeRunner with two maze sizes, 15 x 15 and 30 x 30, where the latter has a longer horizon of 500 steps.
Kitchen and World: Kitchen [66 is robotic environment with continuous observations and actions, where a 7-DoF robot arm manipulates diverse objects a simulated kitchen scene. We define each long-horizon task as completing a sequence of n subtasks in specific order, providing + reward when the next correct subtask is completed. The task goal is provided as a set subtasks without revealing their order. Other goals are represented one-hot vectors indicating the next subtask. We collect the dataset using policies trained with RL. We also introduce GridWorld, a simplified 2D version of Kitchen, where the agent navigates to switch states of 7 objects located in different positions in specified orders.
Crafter [88] A simplified benchmark of the open-world game Minecraft, where the 2D world is procedurally generated. The agent receives 64 x 64 egocentric image observations and takes discrete actions. The objective is to unlock 22 achievements, each providing a reward of +1. Goals are defined as one-hot vectors indicating the next achievement. The dataset is collected using policies from AD [23]. For online adaptation, each task features unique, unexplored world map.
5

---

Table 1: Performance of MGPO compared with baseline methods. Each result shows the average performance on all test tasks the environment and the standard deviation across 3 random seeds for online test. Gond-conditooned BC has error bars since it does not perform online optimization.
Method
MazeRunner-15
MazeRunner-30
Kitchen
GridWorld
Crafter
Gond-condd BC
-2.63
-27.06
0.09
0.05
11.78 1.88 = 0.19
BC-finetune
-3.09 = 0.12
-43.19 + 2.97
0.00 0.00
0.04 = 0.00
SPiRL
-2.62 = 0.35
-28,94 0.79
0,22 0.05
0.10
10,96 + 0.25
PTGM
-0,96 = 0.08
-26.74 1.71
0.25 - 0.03
0.25 0.03
15.72 = 0.12
MGPO
-0.41 + 0.10
-14.21 0.19
0.63 0.02
0.58 + 0.03
15.66 0.14
4.2
Baselines and Main Results
We compare MGPO with previous methods that focus on task-agnostic pre-training. Further details provided in Appendix F
Goal-conditioned BC feeds the task goal MM directly into the pre-trained goal-conditioned policy for new tasks. Its performance is akin to the initial prompt performance in MGPO without prompt optimization.
BC-finetune updates the parameters of the pre- treitrained model with RL based on online collected trajectories, which is a common method in previous work [4] Conditioned on the task goal, we finetune the Transformer parameters using PPO [46..
SPiRL 38
pre-trains short-term skills along with their latent representations and online RL to train high-level policy for online adaptation. PTGM [60] improves this approach by pre-training a goal-conditioned policy and discretizing the goal space for the high-level policy.
For each test task, we evaluated each method for 100 episodes of online rollout, measuring per- formance by the average return of the optimized policy prompt. Results in Table I Tummariiu the performance of MGPO against all baselines across different environments, where MGPO is implemented with UCB in online adaptation.
summarize
MGPO's superiority: MGPO demonstrates superior performance in all environments, showcasing its efficiency adapting to new tasks.
Limitations of BC-finetune: This method underperforms others in most environments due the instability and inefficiency of finetuning parameters in the entire model with online RL. It fails to solve tasks within the small budget of 100 online episodes.
Comparing SPiRL, PTGM, and MGPO: While SPiRL and PTGM exhibit better sample efficiency than BC-finetune, they largely underperform MGPO in MazeRunner, Kitchen, and World These approaches are limited by their reliance on learning high-level RL policy for goal switch- ing at different states during online adaptation. For example, PTGM trains high-level policy T(gl he, o,) select goals, which is effectively searching policies in the joint space of O x G. contrast, MGPO, with its inherent ability to seamlessly switch between goals, optimizes in the prompt space GK leading more efficient adaptation.
The performance of this method highlights MGPO's capability to refine and improve upon the initial prompt. Conditioned on a single goal, the performance of Goal- conditioned BC is sub-optimal due to limitations in dataset quality (e.g.. MazeRunner) and the partial observability in task specifications (Kitchen and GridWorld) and environment layouts (MazeRunner and Crafter). In contrast, through online exploration, MGPO gathers more task and cnvironment information in the unseen task and can search for a prompt sequence to stitch different short-term behaviors learned from data.
Goal-conditioned BC:
Performance in Crafter:
Performance of conditooni BC, PTGM, and MGPO in Crafter is similar compared with other environments. This because, the dataset does not feature great diversity since the data collection policy always aims at unlocking more achievements, thereby the pre-trained policy may be less sensitive to the prompt. contrast, in other environments, the prompt substantially influences the policy's behavior and performance.

---

Table 2: Performance of MGPO with different prompt optimization methods. Each result shows the average performance on all test tasks in the environment and the standard deviation across 3 random seeds for online test.
Method
MazeRunner-15
MazeRunner-30
Kitchen
GridWorld
Crafter
MGPO-GRIPS
-0.56 1 0.09
-13.01 +: 1.22
0.54 + 0.09
0.40 + 0.01
15.53 +: 0.02
MGPO-BBT
-0.80 + 0.16
-12.72 + 2.22
0.74 0.06
0.42 = 0.03
16.31 + 0.03
MGPO-explore
-0.71 + 0.03
-14.74 +: 0.49
0.47 0.05
0.48 + 0.01
15.82 +: 0.07
MGPO-UCB
-0.41 + 0.10
-14.21 + 0.19
0.63 0.02
0.58 + 0.03
15.66 0.14
MGPO-BPE
-0.38 + 0.05
-14.86 1.32
0.65 0.02
0.57 + 0.01
15.65 0.06
Table 3: Performance of prompts optimized with different methods in Kitchen with noisy observations or actions. Each result shows the average performance on all test tasks and the decrease compared with the environment without noise.
Method
MGPO-GRIPE
MGPO-BBT
MGPO-UCB
MGPO-BPE
Noisy observations
0.29 (-0.25)
0.42 (-0.32)
0.49
0.45 (-0.20)
Noisy actions
0.12 (-0.42)
0.19 (-0.55)
0.30 (-0.33)
0.28 (-0.37)
4.3 Prompt Optimization Methods
In exploring prompt optimization methods for the pre-trained multi-goal Transformer, we draw insights from the area of prompt optimization for language models, including two contemporary methods to implement MGPO:
GRIPS [40] is genetic algorithm designed for prompt search within the discrete prompt space. In our implementations, we start with an initial prompt and generate new prompts to evolve it via online evaluating all prompts. We implement operations of adding, deleting, and swapping goals for generating new prompts. The algorithm returns the best-evaluated prompt.
BBT 488] is black-box optimization method for continuous prompt-tuning It utilizes a low- dimensional vector z, which is added to the initial prompt through random projection. This vector is then optimized using CMA-ES [19] evolution strategy. Since it extends the prompt into continuous space, the optimized prompts may not be interpreted as sequences of discrete goals.
For our proposed method, we investigate the two options for the MAB algorithm: UCB and BPE. Additionally, we examined an ablation approach dubbed MGPO-explore, which differs from our MAB formulation by leveraging the most recent prompt sampled from the highest-return trajectory for further exploration.
Table 2 presents experimental results of MGPO implemented with different prompt optimization methods. Across all environments. all MGPO methods demonstrate superior performance compared to the baseline methods in Table showing MGPO's great compatibility with different prompt optimization methods.
UCB and BPE consistently outperform MGPO-explore in most environments. This result demon- strates the efficacy of our MAB formulation in the prompt optimization process. By exploring and exploiting existing prompts, MAB-based methods yield trajectories with higher returns, thus improving the overall performance.
UCB and BPE outperform the discrete search method GRIPS in four out of the five environments. This result highlights the strength of our proposed methods in discrete optimization for multi-goal prompts. While UCB and BPE outperform BBT in MazeRunner-I5 15 and GridWorld, a reversal occurred in the other three environments, particularly in Kitchen where BBT surpasses the performance of all other methods. We speculate that, due to the nature deterministic and differentiable transitions in Kitchen, the continuous optimization method BBT may quickly find local optimum.
The performance of BBT raises critical considerations about the nature of the optimized prompts. Despite its efficiency, the optimized continuous prompts deviate significantly from the training distribution, potentially harming the robustness of the policy. We conduct further studies to evaluate the robustness of the optimized prompts. Table compares the performance of the optimized prompts
7

---

LES
UCB
CRUPS
BET
Figure 2: Visualization of the optimized prompts in four MazeRunner-15 tasks. The prompt with each method is displayed in unique color. and 7 represent the start position and the task goal respectivelyy, and 91. 92, .-- represents goal positions in the prompt. We display the goals that exceed the maze boundaries in the gray bar on the right side.
Table 4: Results of ablation study on MGPO with varying maximal prompt length K and hyperpa- rameter è UCB.
Ablation
MazeRunner-15
MazeRunner-30
Kitchen
K = 1
-2.63++0.00
-27.06=0.00
0.09++0.00
K 2
-1.13++.10
-19.57=1.21
0.28++0.0
K
-0.77++0.04
-17.12=0.47
0.48+00.02
K 5
-0.4111.110
-14.21=0.19
0.63
UCB-0
-0.44+00.07
-15.86=0.51
0.63
UCB-1
-0.41++.10
-14.21=0.19
0.63+00.22
UCB-10
-0.40++004
-14.70=0.66
0.66+0.02
when the robotic environment Kitchen has noisy observations actions. reveals that prompts with BBT and GRIPS are particularly susceptible to environment perturbation, hinting at out-of- distribution prompts. In contrast, prompts with our proposed methods maintain robust performance facing perturbation.
Figure visualizes goals in prompts optimized by different methods MazeRunner-Ii. 15. It clearly showcases the interpretability of the prompt with our proposed methods, which represents meaningful waypoints toward the task goal. contrast, prompts with GRIPS and BBT lack this level of clarity and interpretability.
4.4 Ablation Study
In addition to MGPO-explore, we conduct more ablation studies to examine the impact of components in MGPO. The results of ablation studies are presented in Table
Maximal Prompt Lengths: We study the impact of different maximal prompt lengths during online adaptation. Complex tasks often necessitate multiple goals within a prompt specify varying behaviors in long-horizon trajectories. We observe that optimizing prompts with increasing length consistently improves performance. Specifically, length of K 5 yields the best results, and reducing K to results in performance akin to that of coal-conditioned BC. Appendix B.2 shows additional results when further increasing K.
UCB Hyperparameters: UCB includes hyperparameter c that balances exploration and exploitation (Appendix [). In test with MGPO- UCB under varying values of observe that the performance exhibits low sensitivity to this hyperparameter' 's selection, as shown in Table
Dataset Quality: We also assess the influence of dataset quality on MGPO's performance. Detailed results are provided in Appendix B.3
4.5 Visualization and Case Study
In MazeRunner, we visualize the evolution of MGPO during online adaptation. displays the progression of both the prompt and policy behavior at various stages. At the beginning, the policy is

---

Eptcoda 1-20
2p-0
Ppaodo-4-00
Ppaode
Epiood
Figure 3: Visualization of the optimized prompts and state visitation in the adaptation stage of MGPO-UCB. Red squares with S, T. 9: represent the start positions, task goals, and the optimized prompts, respectively. We purple display the visitation frequency of each location, with darker shades indicating higher frequencies. In this example, initially, the agent explores the left half of the maze, aiming for the front-left task goal but is hindered by walls. In exploration, it discovers rewarding routes to the right, MGPO-UCB adapts by sampling new prompts from these better paths. An optimized prompt is achieved after 40 episodes.
conditioned on a short prompt, resulting in high action uncertainty and exploratory behavior, where the policy explores for diverse trajectories toward the task goal. This exploratory phase is crucial, as it discovers trajectories with higher returns toward the goal and helps in understanding the unknown task environment. As the algorithm progresses, the prompts iteratively refined based returns from the environment, and the policy behavior is steered towards more efficient paths. More results are presented Appendix B.4
Related Work
Policy Pre-Traine Training for RL topic studying learning from datasets to enhance the efficiency of RL on new tasks. Offline meta-RL (OMRL) focuses on learning to adapt to new tasks with a few samples, employing context-based learning methods (26) 59 39
gradient-based meta-learning 33 However, OMRL necessitates multi-task datasets for training, requiring extensive trajectory
collection within each task. On the other hand, task-agnostic pre-training leverages rich behaviors in task-agnostic datasets, exploring imitation learning [43 4. offline RL [24] [4]], or hierarchical
skill learning 38 60 for policy pre-training. Our work addresses the challenging setting of
task-agnostic pre-training for long-horizon tasks. Unlike previous methods that rely on RL for online
adaptation, our proposed multi-goal Transformer enables RL-free optimization methods, enhancing
sample efficiency.
Transformers for RL. The Transformer architecture has become increasingly popular in RL for its ability in sequence modeling and long-term memory. In offline RL, Decision Transformers 57 recast RL as sequence modeling problem conditioned on return-to-go. In context- based meta-RL, Transformers are employed to handle multiple trajectories, modeling task adaptation as in-context learning 23 56 29 multi-task RL, large-scale Transformers are adopted with extensive datasets to address complex robotic domains 16, [ and generalization to various tasks our work, we utilize Transformers to address the demands long-term memory in modeling long-horizon trajectories with multiple goals.
22
49
15).
Prompt Optimization for pre-trained Transformer-based language models has demonstrated its effectiveness in adapting to downstream NLP tasks without tuning model parameters. It optimizes the prompt, which is sequence of tokens input the model specifying the task. the recent large language models (LLMs) (Z]. task adaptation can be easily achieved by prompt design, with the techniques of in-context learning [I]] and chain-of-thousee models, prompt-tuning methods optimize prompts in continuous space 25 27), while prompt search methods focus on optimizing discrete prompt tokens for interpretability the concept of prompt optimization to address online adaptation in RI, focusing on optimizing the goal sequence for the pre-trained Transformer. Unlike in NL.P tasks, our approach places a unique emphasis sample efficiency, crucial aspect given the necessity of online evaluation.
reasoning [52] For earlier language
9 40
We adopt

---

Conclusion and Limitations
We propose MGPO, novel framework for policy pre- pre-training to enhance online adaptation in unseen long-horizon tasks. By integrating the strengths of Transformer architectures and goal-conditioned policies during pre-training, MGPO enables an efficient prompt optimization process in the online adaptation phase. Our extensive experimental results across various environments show MGPO's superiority over existing methods. A comparative analysis of different prompt optimization techniques highlights the advantages on interpretability and robustness of our method over other contemporary approaches.
The effectiveness of MGPO in solving long-horizzn tasks opens new possibilities for real-world applications where efficient online adaptation is crucial. However, experiments have been confined to simulated environments so far. Future research can focus on scaling MGPO to larger pre-training datasets and testing it in complex, real-world environments.
Like many offline RL approaches, the performance of MGPO influenced by the quality of the dataset used during pre-training. To address this limitation, future work could explore several potential directions: incorporating online data collection to improve dataset quality, using offline RL methods such as DT to train the multi-goal Transformer, and integrating MGPO with finetuning methods to enhance its adaptability.
The prompt- based policy in MGPO, similar to language models, can exhibit unpredictable behavior when encountering out-of-distribution prompts. Furthermore, even minor alterations to the prompt may lead to unintended behaviors, raising safety and robustness challenges. Future efforts could explore enhancing the robustness of prompt-based policies.
Acknowledgments
This work was supported by NSFC under Grant 62450001 and 62476008. The authors would like to thank the anonymous reviewers for their valuable comments and advice.
References
[1] Marcin Andrychowicz, Filip Wolski, Alex Ray, Jonas Schneider, Rachel Fong. Peter Welinder, Bob McGrew, Josh Tobin, OpenAl Picter Abbeel, and Wojeiech Zaremba. Hindsight experience replay. Advances in neural information processing systems (NeurIPS), 2017.
[2] Peter Auer, Nicolo Cesa-Bianchi, and Paul Fischer. Finite-time analysis of the multiarmed bandit problem. Machine learning, 2002.
[3] Yusuf Aytar, Tobias Pfaff, David Budden, Thomas Paine, Ziyu Wang, and Nando De Fre- itas. Playing hard exploration games by watching youtube. Advances in neural information processing systems, 2018
[4] Bowen Baker, 1lge Akkaya, Peter Zhokov, Joost Huizinga, Jie Tang, Adrien Ecoffet, Brandon Houghton, Raul Sampedro, and Jeff Clune. Video pretraining (vpt): Learning to act by watching unlabeled online videos. Advances in Neural Information Processing Systems, 2022.
[5] Authony Brohan, Noah Brown, Justice Carbajal, Yevgen Chebotar, Xi Chen, Krzysztof Choro- manski, Tianli Ding, Danny Driess, Avinava Dubey, Chelsea Finn, al. Rt-2: Vision Vision-languuae- action models transfer web knowledge to robotic control. arXiv preprint arXiv:2307.15816 2023.
[6] Anthony Brohan, Noah Brown, Justice Carbajal, Yevgen Chebotar, Joseph Dabis, Chelsea Finn, Keerthana Gopalakrishnan, Karol Hausman, Alex Herzog, Jasmine Hsu, et al. Rt-1: Robotics transformer for real-world control at scale. arXiv preprint arXiv:22 22.00222 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 are few-shot learners. Advances in neural information processing systems, 2020.
[8] 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 (NeurIPS), 2021.
10

---

[9] Mingkai Deng, Jianyu Wang. Cheng-Ping Hsich, Yihan Wang, Han Guo, Tianmin Shu, Meng Song, Eric Xing. and Zhiting Hu. Riprompt: Optimizing discrete text prompts with reinforce- ment learning. arXiv preprint urXiv:2205. /2548, 2022.
[10] Jacob Devlin, Ming- Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04800 04805, 2018.
[11] Qingxiu Dong, Lei Li, Damai Dai, Ce Zheng, Zhiyong Wu, Baobao Chang, Xu Sun, Jingjing Xu, and Zhifang Sui. survey for in-context learning. arXiv preprint vrXiv:2301.0025: 2022.
[12] Benjamin Eysenbach, Abhishek Gupta, Julian Ibarz, and Sergey Levine. Diversity all you need: Learning skills without a reward function. arXiv preprint arXiv: wrXivv/802.00070 2018.
[13] Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning. arXiv preprint arXiv:2004.07219, 2020.
[14] Hiroki Furuta, Yutaka Matsuo, and Shixiang Shane Gu. Generalized decision transformer for offline hindsight information matching. arXiv preprint arXiv:21 2021.
[15] Jake Grigsby, Linxi Fan, and Yuke Zhu. AMAGO: Scalable in-context reinforcement learning for adaptive agents. The Twelfth International Conference on Learning Representations, 2024.
[16] Abhishek Gupta, Vikash Kumar, Corey Lynch, Sergey Levine, and Karol Hausman. Relay policy learning: Solving long-horizon tasks via imitation and reinforcement learning. Conference on Robot Learning (CORL), 2020.
[17] Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off- policy maximum entropy deep reinforcement learning with a stochastic actor. In International conference muchine learning. 2018.
[18] Danijar Hafner. Benchmarking the spectrum of agent capabilities. arXiv:2109.06780 2021.
arXiv preprint
[19] Nikolaus Hansen and Andreas Ostermeier. Completely derandomized self-adaptation in evolu- tion strategies. Evolutionary computation, 2001.
[20] Kaiming He, Ross Girshick, and Piotr Dollár. Rethinking imagenet pre-training. In Proceedings of the IEEF/CVF International Conference on Computer Vision, 2019.
[21] Biwei Huang, Fan Feng, Chaochao Lu, Sara Magliacane, and Kun Zhang. Adarl: What, where, and how to adapt in transfer reinforcement learning. In International Conference on Learning Representations, 2022.
[22] Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as one big sequence modeling problem. Advances in neural information processing systems, 2021.
[23] Michael Laskin, Luyu Wang. Junhyuk Oh, Emilio Parisotto, Stephen Spencer, Richie Steiger- wald, DJ Strouse, Steven Hansen, Angelos Filos, Ethan Brooks, al. In-context reinforcement learning with algorithm distillation. arXiv preprint arXiv:22 10. 142 2022.
[24] Seunghyun Lee, Younggyo Seo, Kimin Lee, Pieter Abbeel, and Jinwoo Shin. Offline-to-online reinforcement learning via balanced replay and pessimistic q-ensemble. In Conference on Robot Learning (CORL). 2022.
[25] Brian Lester, Rami A1-Rfou, and Noah Constant. The power of scale for parameter-efffiiient prompt tuning. arXiv preprint arXiv:2104. 10600,,, 2021.
[26] Lanqing Li, Rui Yang, and Dijun Luo. FOCAL: Efficient fully-offline meta-reiiforcement learning via distance metric learning and behavior regularization. In International Conference on Learning Representations, 2021.
(27) Xiang Lisa Li and Percy Liang. Prefix-tuning: Optimizing continuous prompts for generation. arXiv preprint 2rXhv:2101.0019, 2021.

---

[28] Shalev Lifshitz, Keiran Paster, Harris Chan, Jimmy Ba, and Sheila Mcllraith. Steve-1: A generative model for text-to-behavior minecraft. arXiv preprint 0xXi:23366000937 2023,
[29] Hao Liu and Pieter Abbeel. Emergent agentic transformer from chain of hindsight experience. arXiv preprint arXiv:2305. 16554, 2023.
[30] Pengfei Liu, Weizhe Yuan, Jinlan Fu, Zhenghao Jiang, Hiroaki Hayashi, and Graham Neubig. Pre-train, prompt, and predict: A systematic survey of prompting methods in natural language processing. ACM Computing Surveys, 2023.
[31] Shaoteng Liu, Haoqi Yuan, Minda Hu, Yanwei Li, Yukang Chen, Shu Liu, Zongqing Lu, and Jiaya Jia. RI-gpt: Integrating reinforcement learning and code-as-policy, arXiv preprini arXiv:2402. 19299, 2024.
[32] Bonan Min. Hayley Ross, Elior Sulem, Amir Pouran Ben Veyseh, Thien Huu Nguyen, Oscar Sainz, Eneko Agirre, Ilana Heintz, and Dan Roth. Recent advances in natural language processing via large pre-trained language models: A survey. ACM Computing Surveys, 2023.
[33] Eric Mitchell, Rafael Rafailov, Xue Bin Peng, Sergey Levine, and Chelsea Finn. Offline meta- reinforcement learning with advantage weighting. In International Conference Machine Learning, 2021.
[34] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wierstra, and Martin Riedmiller. Playing atari with deep reinforcement learning. arXiv preprint arXiv:1312.5602, 2013.
[35] Seungyong Moon, Junyoung Yeom, Bumsoo Park, and Hyun Oh Song. Discovering hi- erarchical achievements in reinforcement learning via contrastive learning. arXiv preprint arXiv:2307.03486, 2023.
[36] Kevin P Murphy. Conjugate bayesian analysis of the gaussian distribution. def, 2007.
[37] Jurgis Pasukonis, Timothy Lillicrap. and Danijar Hafner. Evaluating long-term memory in 3d mazes. arXiv preprint arXiv:22 13383, 2022.
[38] Karl Pertsch, Youngwoon Lee, and Joseph Lim. Accelerating reinforcement learning with learned skill priors. In Conference on Robot Learning (CORL), 2021.
[39] Vitchyr H Pong, Ashvin V Nair, Laura M Smith, Catherine Huang, and Sergey Levine. Offline meta-reinforcement learning with online self-supervision. In International Conference on Machine Learning, 2022.
[40] Archiki Prasad, Peter Hase, Xiang Zhou, and Mohit Bansal. Grips: Gradient-free, edit-based instruction search for prompting large language models. arXiv preprint arXiv:2203.07281, 2022.
[41] Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In International conference machine learning, 2021.
[42] Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models unsupervised multitask learners. OpenAl blog, 2019.
[43] Ram Ramrakhya, Dhruv Batra, Erik Wijmans, and Abhishek Das. Pirlnav: Pretraining with imitation and rl finetuning for objectnav. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 2023.
[44] Scott Reed, Konrad Zolna, Emilio Parisotto, Sergio Gomez Colmenarejo, Alexander Novikov, Gabriel Sarth-Maron Mai Gimenez, Yury Sulsky, Jackie Kay, Jost Tobias Springenberg, et al. A generalist agent. arXiv preprint arXiv:2205.06177 2022.
[45] Erick Rosete-Beat Beas, Oier Mees, Gabriel Kalweit, Joschka Boedecker, and Wolfram Burgard. Latent plans for task-agnostic offline reinforcement learning. In Conference on Robot Learning (CORL), 2023.
12

---

[46] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint orXiv:1707.06347, 2017.
[47] Lucy Xiaoyang Shi, Joseph J Lim, and Youngwoon Lee. Skill-based model-based reinforcement learning. In Conference Robot Learning (CORL), 2023.
[48] Tianxiang Sun, Yunfan Shao, Hong Qian, Xuanjing Huang, and Xipeng Qiu. Black-box tuning for saatlaage-mdel-ss-----e-
In International Conference on Machine Learning, 2022.
[49] Adaptive Agent Team, Jakob Bauer, Kate Baumli, Satinder Baveja, Feryal Behbahani, Avishkar Bhoopchand. Nathalie Bradley-Schmice Michael Chang, Natalie Clay, Adrian Collister. et al. Human-timescall adaptation in an open-ended task space. arXiv preprint arXiv:2301.000 076008 2023.
[50] Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: physics engine for model-based control. In 2012 IEEE/RSJ international conference on intelligent robots and systems, 2012.
[51] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems (NeurIPS), 2017.
[52] Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoe V Le, Denny Zhou, et al. Chain-of-thought prompting elicits reasoning in large language models. Advances in Neural Information Processing Systems, 2022.
[53] Ronald J Williams. Simple statistical gradient- foolowing algorithms for connectionist reinforce- ment learning. Machine learning, 1992.
[54] Zhihui Xie, Zichuan Lin, Deheng Ye, Qiang Fu, Yang Wei, and Shuai Li. unsupervised pretraining for decision transformer. In International Conference on Machine Learning, 2023.
[55] 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.0848; 2023.
[56] 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, 2022.
[57] Taku Yamagata, Ahmed Khalil, and Raul Santos-Rodriguer Q-learning decision transformer: Leveraging dynamic programming for conditional sequence modelling in offline rl. In Interna- tional Conference on Machine Learning, 2023.
[58] Tianhe Yu, Deirdre Quillen, Zhanpeng He, Ryan Julian, Karol Hausman, Chelsea Finn, and Sergey Levine. Meta-worid: A benchmark and evaluation for multi-task and meta reinforcement learning. In Conference on robot learning. 2020.
[59] Haoqi Yuan and Zongqing Lu. Robust task representations for offline meta-reiiforcemment learning via contrastive learning. In International Conference on Machine Learning, 2022.
[60] Haoqi Yuan, Zhaneun Mu, Feiyang Xie, and Zongqing Lu. Pre-training goal-based models for sample-efficient reinforcement learning. In The Twelfth International Conference on Learning Representations, 2024.
[61] Haoqi Yuan, Chi Zhang, Hongeheng Wang, Feiyang Xie, Penglin Cai, Hao Dong, and Zongqing Lu. Skill reinforcement learning and planning for worn-world long-horizon tasks, arXiv preprint arXiv:2303. 2023.
[62] Barret Zoph, Golnaz Ghiasi, Tsung-Y Yi Lin, Yin Cui, Hanxiao Liu, Ekin Dogus Cubuk, and Quoc Le. Rethinking pre-training and self-training. Advances in neural information processing systems, 2020.
13

---

A Illustrations of MGPO
(wood. sword, stone, coal, colcow?
(1) Pre-training: learn actions conditioned on goal sequence.
wood
labil
wood sword
stone
0cal
(drink, hindsight malt- goal relabeling
driic
oat DOW
skoleton
zombie
Online adaptatiion search for the optimal goal sequence for task completion.
Task unlock all achievements in an unseen world.
Agan: prompt optimization coal iron, iron sword, ...)
adaptation
Figure A running example illustrating MGPO. During the pre-tre pre-training stage, the agent learns to complete arbitrary sequences of goals using offline trajectories with rindeight-relabeled goals. In the online adaptation stage, the agent optimizes the goal sequence to maximize returns an unseen environment.
Task: bottom burner - slide cabinet - hinge cabinet microwave
Prompt
Behavior t= 63
Success
t=
t=
50
t= 110
t 215
bottom burner slide cabinet
hinge cabinet microwave
115
200
slide cabinet hinge cabinet microwave bottom burner
x
Figure 5: An illustration of prompts (goal sequences) and the agent's ' behavior in Kitchen task. The order goals within the prompt is crucial for successful task completion.
Additional Experimental Results
B.1 Learning Curves During Online Adaptation
Figure showw the learning curves of MGPO and baselines during online adaptation. Figure shows the learning curves of MGPO with different prompt optimization methods during online adaptation.
B.2 Ablation Study on the Maximal Prompt Length
Since we use = 5 the pre-training phase and the model has not encountered prompts exceeding length of 5, we present results with maximal prompt lengths K < 5 in our main paper. Here, we present additional results when further increasing K. We pre-train models with K 10, and 40 to test online adaptation with larger prompt lengths.
14

---

MazeRunner-15
MazeRunner-30
Kitchen
0.6
Performani atoomaaioo
Performance
Performance aeanance
0.4
-30
40
-50
0.0
20
100
Episodes
Episodes
Episodes
GridWorld
Crafter
0.6
Perforancc
0.4
0.2
12
10
20
40
80
100
20
40
80
100
Episodes
Episodes
MGPO
Goal-conditoned
BC-Anetune
SPILL
Figure 6: Performance of MGPO compared with baseline methods during online adaptation. The vertical axis indicates the task performance of the optimized policy and the horizontal axis indicates the number of online episodes. Each figure shows the average performance on all test tasks in each environment and the standard deviation across random seeds.
MazeRunner-15
MazeRunner-30
Kitchen
10
Porformance
Perfermance
Perforannce
0.6
20
-30
20
40
80
100
20
40
80
100
20
40
60
80
100
Episodes
Episodes
Episedes
GridWorld
Crafter
16
0.5
Sefformance
1.5
0.3
1.4
0.2
0.0
20
40
Episodes
Episodes
MGPO-UCB
MGPO-3PE
MGPO-explore
MGPO-GRUPS
MGPO-B8T
Figure Performance MGPO with different prompt optimization methods during online adapta- tion.
15

---

Figure 8 shows the results. In MazeRunner-155 15, setting K = 10 yields the optimal performance. Increase beyond this value leads decline in learning efficiency. In MazeRunner-30, performance exhibits slight improvement as K extends to 40. This suggests that for more extended, long-horizon tasks, the exploration of longer prompts can be advantageous.
MazeRunner-15
MazeRunner-30
-10
-0.5
Performance
-15
Performance
-1.5
-2.0
-25
30
20
60
100
Episodes
Episodes
K-2
K-3
K-5
K-10
K-20
K-40
Figure 8: Additional results of ablation study on the maximal prompt length K.
B.3 Ablation Study on Dataset Quality
We assess MGPO's performance with datasets of varying quality in MazeRunner. In Table 5 denote datasets collected using an A algorithm with a maximum of n goal switches per episode * and datasets from random exploration policy as Random. Results show that MGPO achieves better performance trained on A datasets than **-1, indicating its efficacy with data containing diverse long-horizon behaviors. The comparatively lower performance in MazeRunner-30 and Random datasets suggests MGPO's 'ss reliance on the quality of data collection policies. Future work could involve integrating offline RL methods to boost performance with Jower-qualit, datasets.
the A 4 dataset
Table 5: Results of ablation study on varying dataset quality, evaluated with MGPO-UCB.
Dataset
MazeRunner-13 15
MazcRunner-30
Random
-1.05110.01
-17.13
A+11
-0.97+40.03
-9.26+40.38
A+.2
-0.36++0.02
-8.66+40.5
A+.4
-0.41+0.10
-14.21:0.19
B.4 Additional Visualization Results
By prompting the pre-trained policy with prompts of varying lengths, MGPO can naturally trade-off exploration and exploitation during online adaptation. To empirically demonstrate this capability, we visualize prompts of different lengths and the corresponding behaviors of the pre-trained policy in Figure 9 higher uncertainty and thus encourages more exploratory behaviors. Conversely, longer prompt, containing more detailed information from the trajectory, prompts more deterministic and exploitatory behavior.
We observe that shorter prompt, by providing less information to the policy, introduces
Figure T0 shows additional results on the progression of both the prompt and policy behavior various stages.
16

---

LSG
Short
LSC S
Long
Figure 9: Visualization of prompts with different lengths and their state visitation counts in MazeRunner-155, Each image the upper row shows a short prompt with length of 1. Each image in the bottom row shows long prompt with length of 5.
Brisode--29 1-20
21-
61-60
Episode 61-80
81-10d8811100
Figure 10: Additional results visualization of the optimized prompts and state visitation in the adaptation stage MGPO-UCB. Each row represents a task MazeRunner-I1 15.
17

---

9
MazeRunne-15
MazeRummex-30
Kitchen
GridWorkd
Crafter
Figure 11: Environments in our experiments.
C Details for Environments and Data Collection
C.1 MazeRunner
MazeRunner is 2D maze navigation environment implemented in AMAGO [5] and is mainly based on Memory Maze [37] The agent navigates grid world of size N x N, where we use N 15 and 30 in experiments. For each task, the unknown structure of walls in the maze is distinct, and the agent should move to the provided goal location gM (x,,, .y to accomplish the task. The agent receives 6-dimensional observations (Ir,) y, d1, d2, d3, da). where the first two values are the agent's position and the last four values are its distances to walls in the four directions. We the agent's positions in observations represent goals. The discrete action space consists of moving left, right, forward, backward, and no movement. The reward is +1 when the agent reaches the task goal and penalty of -0.1 is given at each step in addition. The task horizon 64 steps for MazeRunner-Is and 500 steps for MazeRunner-30. The episode terminates when the agent reaches the task goal or reaches the task horizon.
For data collection, we implement a rule-based policy that repeatedly samples a goal position in maze and plan for path to reach it with A# algorithm. The A algorithm takes the partial observations of the agent, estimate score for each position in the known part of the maze, and returns a shortest path towards the position with highest score. In each episode, the maze is randomly generated and the policy samples a new goal when achieving the last goal for times, until achieving all goals or the environment terminates. We collect 10K trajectories for both MazeRunner-15 and MazeRunner-30. We sample 50 tasks of different mazes and task goals for test.
C.2 Kitchen
The Kitchen [66 environment is designed to simulate a kitchen scenario for indoor robotics with several interactive objects. The subtasks used in this paper include activating the bottom bunner, switching on the light, opening the microwave, sliding the slide cabinet, and opening the hinge cabinet. The target of each subtask is adjust the corresponding object to a target pose. The observation is a di-einnnsi vector, representing positions and velocities of all objects and the 7-DoF robot arm. The actions are 9-dimensional continuous vectors representing velocities to control the motor actuators in the robot arm. Each task is defined with set of n subtasks among the five subtasks and the order to complete them. The task horizon set to 500 timesteps. When subtask is completed in the correct order, the environment provides a reward of + 1. The goal space consists of 5-dimensional binary vectors representing each subtask in each bit. For the task goal revealed to the agent, bits the selected subtasks are set to and other bits are set to O. Other goals are one-hot vectors representing the current subtask to complete.
We verified that the Kitchen datasets provided in D4RL 13 do not contain diverse transitions between the five subtasks. To collect trajectories completing different sets of subtasks, we use PPO 46) to train policy for each sub-tas using a shaped reward function and varied initial states from Kitchen-miied-vd dataset [3] We sample 300 tasks for data collection. In each task, we chain the policies learned with PPO to sequentially solve subtasks and collect 30 episodes. The whole dataset consists of 9K trajectories. We sample test set of 25 tasks, which has no overlap to tasks used in data collection.
18

---

C.3 GridWorld
This environment is a 10 x 10 grid world with 7 distinct objects at different locations. Each object has a binary state, which defines subtask that is completed when the state turned I from the initial state of O. The agent observes 9-dimensional vector representing its location and states of objects. The discrete action space consists actions: moving to the four directions, flip the object state, and no operation. Since this environment is simplified version of Kitchen, tasks and goals are defindd similar Kitchen's The task horizon is 90 in this environment.
We use a rule-based expert policy to collect 13K trajectories and sample a test set of 50 unseen tasks.
C.4 Crafter
Crafter [8] is a research-friendly benchmark mirroring the popular open-world survival game Minecraft. The 2D world map with infinitely large size is procedurally generated. The observations are 64 x 3 rgb images showing the first-person view of the agent, which contains a 9 x 9 part of the world and the rendered information of inventory and life. The discrete space contains 17 actions, including movement, attack, tool use, and crafting. There are 22 achievements in the game, When the agent unlocks each achievement for the first time, receives a reward of +1. When the agent's health increases or decreases by one point, it receives corresponding reward of +0.1 or 0.11. The target of the task is to complete all achievements. Different tasks are featured with different generated world maps, which can be specified with random seed. Goals are 22 dimensional binary vectors representing each achievement with each bit. For hindsight relabeling, the task goal is set to the accomplished achievements in each trajectory. For online test, we set the task goal to all-one to encourage more achievements. Other goals are one-hot vectors indicating the current achievement to unlock.
We set the task horizon 500 steps. We use the policy trained in Achievement Distillation 35 to collect dataset of 2K trajectories. We use 50 test task specified with random seeds unused in data collection. The average episode return in the dataset is 12.3 and our method can increase this to 16 in online adaptation.
D Training Details
We build our Transformer architecture based on Decision Transformers g using lightweight backbone of GPT-2 42 with 0.6M parameters, Input tokens of o, and g are transformed into 128- dimensional embedding vectors using three different linear layers, respectively. We use a positional embedding that embeds the timestep of each token and adds it to the token embedding, enhancing the timestep information of inputs. The Transformer then processes the sequence of embeddings and outputs sequence with the same dimension and length. We use a linear layer decode tokens at positions of the input tokens of observations to predict actions and compute loss. To train on batches of trajectories, pad the imput prompts and trajectories to lengths of K and H, respectively.
In MazeRunner, we sample prompts from the agent's 'ss locations in the whole trajectory. To augment the diversity of task goals and trajectory lengths, we truncate the trajectory at random timestep h for each sampled trajectory and use on to represent its task goal. other environments, goals in prompts are sampled from the completed subtasks in each trajectory and 9и represents the actual set of subtasks completed at the end of each trajectory.
Table 6 lists the hyperparameters used in pre-training.
All models are trained on a lab machine with single NVIDIA RTX 4090 GPU and Intel i9 CPUs. For each environment, the pre-training stage takes about 12 hours.
E Details for Prompt Optimization
E.1 MGPO-UCB and MGPO-BPE
We summarize proposed prompt optimization algorithm in Algorithm
I
19

---

Table 6: Hyperparameters used in pre-training for all environments.
Name
Value
Embedding dimension
128
Number of layers
3
Number of attention heads
1
Activation
ReLU
Batch size
64
Learning rate
le-4
Learning rate decay weight
le-4
Dropout
0.1
Warmup steps
10000
Algorithm Prompt Optimization in MGPO-UCB and MGPO-BPE
Input:
pre-trained policy To; task goal gM; interaction budget N Initialize: prompt buffer B {(g M) 0]; best return R and best trajectory r²
while not exceeding N episodes do Select p" from B with MAB algorithms.
Collect and get return R with ne <alp"", h, o). Update B, R²,T".
Sample new prompt p' ~ Replace the last goal in p' with gh Collect and get return R with To (alp' h o).
Update R", -²
if R >> maxpEB RS then Insert (p': (R)) to end if end while
Here present the MAB algorithms used in methods for prompt selection. At each itera- tion, we have a buffer prompts B = p. where each prompt p maintains its historical returns (Rp,,,". R.,ng). We aim select p" to collect trajectories of high returns.
UCB: We estimate the upper confidence bound of the expected return for each prompt:
log(N)
UCB(p) = Rp +
v
(3)
2np
where Rp is the mean return in history, is hyperparameter which is set to 1 in our main results, and = =EpEB np. The prompt for online exploration is selected with p = argmax UCB
We assume a Gaussian prior distribution N(Hn, σо) for returns of each prompt. We set #40 (R2in + Rmax) and σa (RRmx Rmin), where Amin and Rmmax are the minimal and maximal possible returns in the environment, respectively. Given historical returns of each prompt, we estimate its posterior distribution P, R = (R) (p) with a Bayesian approach:
BPE:
1
n2 Rp
1p
oz
P40
Rp
(4)
00
0Z
R
where RS is the mean historical return and oR, the standard deviation of historical returns. We select the prompt which has the maximum probability to yield returns exceeding the current highest return: p" argmaxpe P2 > R*).
20

---

E.2 GRIPS
a gradient-free, search-based algorithm for prompt optimization. At each iteration, performs editing operations on the prompt and evaluates the edited prompts. We implement three operations to edit the prompt. Add: Randomly sample a goal and insert it into the current prompt at random position. Del: Randomly delete goal in the current prompt. Swap: Randomly select two goals and swap them in the current prompt.
GRIPS [40]
We initialize the prompt with Po (g and maintain gMو at the last position of the prompt while editing the preceding goals. In each iteration, we randomly generate edited prompts, online evaluate their performance, and preserve the prompt with the highest return.
We summarize GRIPS in Algorithm
Algorithm 2 GRIPS
Input:
pre-trained policy Te: task goal M, interaction budget N; the number of edited prompts per iteration
Initialize: best return R -00 and best prompt y" (g while not exceeding N episodes do
Edit pa with random operations from (Add, Del, Swap) to generate prompts 1= Online evaluate each p; with (a pi, h, o), get return R4.
i* argmax, Rg.
if R3 R" then
R p" = Pi
end if end while
E.3 BBT
is black-box continuous optimization method for prompt-tuning. We use p (g to collect a trajectory and sample a initial prompt with Po ~ т. In the following iterations, we optimize Rdd in low-dimensional space (d < D, where D is the dimension of prompt) and use random projection matrix E Dxd to project and add to the initial prompt. Our objective is:
BBT 48
argmax RA ZER
(5)
where is the expected return of the policy with prompt A= + Po-
BBT uses CMA-ES for optimization. In each iteration f, sample X = 10 new vectors (t+1))
from the multivariate normal distribution 2; 2((+1) m(t + (t) C(t)), 1.... X, where
m(t) d is the mean vector of the search distribution at iteration (t) R_ standard deviation controlling the step length, and C(t) Rdxd is covariance matrix that determines the shape of the distribution. m(t), (t). and C(t) are updated to maximize return.
We summarize BBT in Algorithm 3
Algorithm 3 BBT
Input: pre-trained policy тæ; task goal gM interaction budget N; the population size Initialize: best return R² oo; best prompt y": mn00, (0), C(O), and A Collect an episode 7 with p = M).
Sample initial prompt Po T. while not exceeding N episodes do
Sample 2,, i ... from the multivariate normal distribution.
Collect t, and get return Rt with ne (al A - Po, h, o) for each Update Rt and p".
Update m, a, C with CMA-ES.
end while
21

---

Table Hyperparameters in BC-finetune.
Name
Value
Discount factor Y
0.98
Learning rate a
le-4
Learning rate decay
le-4
Batch size
5
Table Performance the BC-finetune baseline with different RL algorithms.
Method
MazeRunner-I5 15
MazeRunner-30
Kitchen
GridWorld
Crafter
REINFORCE
-2.59 + 0.24
-44.22 + 2.75
0.00 = 0.00
0.06 + 0.02
0,48 + 1.09
PPO
-3.09 + 0.12
-43.19 + 2.97
0.00 0.00
0.04 + 0.00
1.88 + 0.19
F
Details for Baselines
F.
BC-finetune
We fix the prompt p (g and finetune the policy parameters using RL algorithms during online adaptation. We have experimented with the policy gradient method REINFORCE [53] and the modern actor-critic method PPO [46). However, we find that both approaches significantly underperform compared to MGPO. Thus, we only report results using PPO in our main paper. Table 7 lists the hyperparameters used in experiments. Table presents the results for BC-finetune using both REINFORCE and PPO.
F.2 SPiRL and PTGM
These methods pre-trains diverse short-term skills accelerate online RL by providing temporal abstractions.
SPiRL [88 pre-trains skill encoder gu (z(2/, or, (ex+s), skill decoder Te (0e:+4 he, Ot, 2), and a skill prior (2)/,,,,, ot). du and ma trained using the framework of conditional variational autoencoders, while Pe is optimized the short term actions a5.++k are encoded into 10-dimensional continuous latent skills z. For each downstream task, it trains a high-level RL policy THK (z]h, o) with SAC [17). The policy is regularized with KL divergence loss between the policy and the skill prior pv- Table lists the hyperparameters used SPiRL.
match the posterior distribution of 9.. In this approach,
Table 9: Hyperparameters in SPiRL.
Hyperparameters
Value
Weight for the KL loss
5e-4
Low-level steps
10
Discount factor
0.95
Learning rate
le-3
Batch size
256
Target network update interval
10
pre-trains a low-level, goal-conditioned policy no(alg, hr; oz) along with a goal prior Pe (g ).). We slightly modify MGPO to implement PTGM. We remove the prompt before the trajectory and associate each observation with a goal which is obtained via hindsight relabeling. my optimized via goal-conditioned behavior cloning and p& is optimized to maximize the log-liieldhood
PTGM 660
22

---

of the relabeled goal. In online adaptation, we use the discrete goal space as the high-level action space. We use DQN B4 to train the high-level policy =o ( 0t). Its architecture an MLP with layers and hidden layer dimensions of 128. We also incorporate KL regularization reward between the policy and the goal prior Pşj to accelerate online RL. Table I0 lists the hyperparameters used in PTGM.
Table 10: Hyperparameters in PTGM.
Hyperparameters
Value
Weight for the regularization reward
0.5
Low-level steps
10
Discount factor
0.95
e for e-greedy exploration
0.1
Learning rate
le-3
Batch size
64
Target network update interval
10
G
Comparison with Related Areas
typically require access to multiple trajectories per task and are often limited to short-term tasks due to the challenges of few-shot task inference. In contrast, our approach focuses pre-training from task-agnostic, offline datasets and adapting to long-horizon tasks.
Meta-RI. / Offline Meta-RL:: Meta-RI. approaches [56 21 59
primarily focus on learning short-term skills through online RL. In contrast, our work focuses on offline skill acquisition from diverse, task- datasets, conceptualized in the literature [38]
Skill-based RL:
Online skill discovery methods [12
61 31
47 600. These methods typically use hierarchical RL learn high-level policy during online adaptation. Our approach differs by avoiding the inefficiencies of online RI; instead, we utilize pre-trained Trranformer-base policy, enabling us to streamline the adaptation process through efficient prompt optimization.
H Broader Impacts
This research contributes to the field of pre-training for RL, with the goal of improving the efficiency of RL agents to learn new tasks. While MGPO offers promising advancements, also carries inherent risks such reward hacking, where agents exploit loopholes in the reward function to achieve unintended high rewards. This can lead to unsafe or undesirable behaviors, especially in real-world applications. To mitigate these risks, it is vital to design robust reward functions and incorporate safety measures during training. Furthermore, the deployment of RL systems, particularly in sensitive areas like healthcare autonomous driving, raises significant ethical considerations. Ensuring these systems are used responsibly is paramount to prevent negative societal impacts.
23

---

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: The paper's contribution and scope is reflected in the introduction and summa- rized at the end of the introduction.
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 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 expected to generalize to other settings.
It 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:
Limitations are discussed in Section
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those not discussed in the paper.
The authors are encouraged to create 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 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, a facial recognition algorithm may perform poorly when image resolution low images 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 be used by reviewers grounds for rejection, worse outcome might be that reviewers discover limitations that aren't tt acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play 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: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]

---

Justification:
No theoretical results.
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 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 short proof sketch to provide intuition.
Inversely, any informal proof provided 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: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: Details of environments, datasets, and methods are presented in Appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, 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 dataset and/or model, the authors should deseribe 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 novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, may be necessary to either make 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 hosted model (e.g., in the case of a large language model), releasing model checkpoint, or other means that are 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) 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 reprodueibility 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, may be that access to the model is limited 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: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
25

---

Answer: [Yes]
Justification:
We have released the code.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.cr cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is acceptable answer. Papers cannot be rejected simply for not including code, unless this central to the contribution (c.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed run to reproduce the results. See the NeurIPS code and data submission guidelines (https
oosissoooooiiiooooooooooocooo 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 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 ver- sions (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: Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Answer: [Yes]
Justification:
Details of experimental settings are provided in Appendix C
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented the core of the paper to a level of detail that 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 report error bars for all experimental results, showing the standard devia- tions across random seeds.
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, statistical significance tests, least for the experiments that support the main
claims of the paper.
The factors of variability that the enror 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 (c.g., Normally distributed errors).
26

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
It is OK to report I-sigma error bars, but one should state it. The authors should preferably report 2-sigma error bar than state that they have 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 figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures 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 execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
Details are presented Appendix D
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, 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: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics http:///wwur.p...pppaaaaccaccaaa
Answer: [Yes]
Justification:
Our submission conforms with the NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS 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 special consideration due to laws regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: Discussed in Appendix H 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 socictal 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.
27

---

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 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 used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that 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 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 addition to attacks, mecha- nisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility MI.).
11. Safeguards
Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (c.g., pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification:
This paper poses no such risks.
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 best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: Papers for the codebase, environments, and datasets are properly cited.
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 URL.
The name of the license (e.g.. CC-BY 4.0) should be included for each asset.
For scraped data from 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, paperswi chcode. 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.
If this information not available online, the authors are encouraged to reach out to the asset's creators.
28

---

13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
This paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets. Researchers should communicate the details of the dataset/code//modll as part of their sub- missions 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 anonymized URL include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: This paper does involve crowdsourcing nor research with human subjects.
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 much detail 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 the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
This paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. 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 the institution conducting the review.
29


Reviews:
[{"review_id": "w5HGYk50VI", "reviewer": "Reviewer_V4S2", "summary": "This paper addresses the fast adaptation of pre-trained policy from task-agnostic datasets. The authors propose to avoid RL interactions on new tasks through the combination of Transformer-based policies to model multiple goals and efficient online adaptation through prompt optimization. The experiments demonstrate MGPO’s superior performance and efficient adaptation.", "questions": "- What is the advantage of hindsight relabeling in this paper, as the environment already give feedback for goals. Why the hindsight relabeling is required? Could you please provide more discussion and explanations?\n\n- How to determine the length of the sequences of goals while online adaptation. Is there any terminal signal or pre-defined length?\n\n- While the prompt provide a way of fast online adaptation, I am curious about how advantage the transformer-based policy backbone provided, as there are already some RL works in fast adpation through optimizing a subset of network parameters, like optimizing the free parameter in [1]. How about the performance of this method if taking other neural network as the policy backbone, like cnn-based policies and lstm-based policies, where the prompt can be modeled as free parameters and learned through multi-task optimization, like [1]. Are prompt and transformer policy nessary for learning the goal switching?\n\n[1] Huang B, Feng F, Lu C, et al. AdaRL: What, Where, and How to Adapt in Transfer Reinforcement Learning[C]//International Conference on Learning Representations.\n\n- Could you please provide illustration of the learned prompts and the goals in the environment. How will the performance change if we permute the learned goals in prompt? If there is an obvious decrease, the proposed method indeed learn the sequetial relationship of the goals, right?\n\n\n- Can you provide analysis about why the BC-finetune perform worse than Goal-conditioned BC?\n\n- Some typos: like the bold formatting of the citation text on line 233.", "limitations": "See questions.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "It is straightforward and reasonable to model the sequence of goals through prompt and  apply prompt optimization for online adaptation and goal learning in unseen environment. Converting goal switching capacity learning into fast online adaptation is natural.\n\nThe experimental results are sound and demonstrates good performance of the proposed method. The paper is well-written, and the visual examples well verify the interpretability of the method.", "weaknesses": "See questions.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "## Thanks for your review! Here, we respond to your comments and address the issues. We hope to hear back from you if you have further questions!\n\n**Q1:** Why the hindsight relabeling is required? \n\n**A1:** We use hindsight relabeling in the pre-training stage to learn from an offline dataset without environment interactions. Hindsight relabeling is essential for associating each trajectory with a specific sequence of goals, thereby training our policy to be goal-conditioned. Specifically, within each offline trajectory, we extract a sequence of goals that the agent was implicitly aiming to achieve during data collection. This approach allows the pre-trained policy to learn the relationship between sequences of goals and the corresponding actions, enhancing its ability to adapt to varied goal sequences and execute appropriate behaviors in different contexts.\n\n\n**Q2:** How to determine the goal sequence length during online adaptation. \n\n**A2:** In online adaptation, we maintain the maximal prompt length $K$ used during pre-training. The prompt length for each iteration is determined by sampling a number of subgoals $k \\sim U[0, K-1]$ and sampling $k$ subgoals from the best trajectory collected. These details are provided in Section 4.1 and 4.2. We acknowledge and will correct a typo in line 199 regarding the sampling process, which should be $p' \\sim P(p|\\tau^*)$.\n\n**Q3:** Compare prompt optimization with the approach in AdaRL.\n\n**A3:** It's important to note that our work and AdaRL tackle distinct problems in fast adaptation. AdaRL focuses on **online meta-RL**, optimizing for fast adaptation to new tasks through multi-task online training. It leverages a compact context representation to facilitate few-shot adaptation primarily in **short-term tasks**, such as Cartpole and Atari Pong. Conversely, our approach focuses on **offline, task-agnostic** pre-training for **long-horizon** tasks. We aim at fast adaptation to new tasks with multiple subgoals or stages. To ensure a comprehensive evaluation, we included various baseline methods from this literature, including finetuning methods (BC-finetune) and skill-based methods (SPiRL and PTGM).\n\n**Q4:** Compare the Transformer policy to CNN-based and LSTM-based policies.\n\n**A4:** While CNNs and LSTMs are capable of sequence modeling, we choose the Transformer architecture due to its demonstrated superiority in offline RL [1,2] and imitation learning [3,4]. Its effectiveness in modeling trajectories surpasses other architectures, as evidenced by existing experiments [3,4]. This led us to adopt the Decision Transformer [1] architecture for both our method and all baseline comparisons.\n\n**Q5:** Illustrate the learned prompts and goals.\n\n**A5:** We maintain the order of sampled goals when constructing prompts through hindsight relabeling, ensuring the policy learns to achieve each goal sequentially. Permuting the goals within the prompt will cause the policy to attempt to reach subgoals in a different order. We illustrate our prompts and goals in the Kitchen environment in Figure 1 of our attached PDF. Permuting the goal sequence affects task execution, leading to failure in achieving the original task. This is critical for real-world tasks that require subtasks to be completed in a specific order.\n\n\n**Q6:**  Why does BC-finetune perform worse than Goal-conditioned BC?\n\n**A6:** BC-finetune requires substantial iterations to adapt to new tasks. Within the limited budget of 100 episodes in our experiment, it diverges from the original Goal-conditioned BC policy but fails to fully adapt, leading to underperformance. This is often due to the instability and initial performance decline of RL finetuning methods. To address this, we introduce an enhanced version, BC-finetune-KL, which integrates a KL-divergence loss with the initial policy to stabilize the finetuning process. As shown in Table 1 of the attached PDF, BC-finetune-KL shows some improvement over Goal-conditioned BC but still lags behind MGPO.\n\n**Q7:** Typos.\n\n**A7:** Thank you for pointing these out! We will fix the typos.\n\n\n[1] Chen et al., Decision Transformer: Reinforcement Learning via Sequence Modeling, 2021. \n\n[2] Li et al., A Survey on Transformers in Reinforcement Learning, 2023. \n\n[3] Dasari et al., Transformers for One-Shot Visual Imitation, 2020. \n\n[4] Kim et al., Transformer-based deep imitation learning for dual-arm robot manipulation, 2021.", "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": "pJAqxJz9ii", "reviewer": "Reviewer_JzeR", "summary": "The paper proposes a pretrain-and-prompt-tuning paradigm to tackle the generalization challenge in RL. It pretrains a goal-conditioned transformer from task-agnostic datasets, and during fine-tuning, it constructs a goal sequence as a prompt and tunes that prompt via multi-arm bandit algorithms.", "questions": "- The proposed method pretrains the “skills” using behavior cloning. If the dataset quality is not high enough, those skills could be sub-optimal temporal abstractions. Does this pretraining scheme further limits the performance, in the long term?\n- The multi-arm bandit algorithm is employed to optimize the goal sequence, which is a NP-hard, combinatorial optimization problem. How to ensure it can find a good goal sequence using only a few iterations?\n- In the multi-arm bandit modeling, what does one arm correspond to? It seems that one arm corresponds to an element from the prompt buffer. The prompt buffer expands as new prompts are added. The number of arms is infinite. The complexity of the multi-arm bandit algorithm could be high, and its theoretical guarantee can be hard to derive.\n- At test time, the learner is given an unseen task with the “task goal”. Does this part need domain knowledge, to some extent? If the agent knows the task goal at any time, we can just train a goal-conditioned policy to do the job, as have been investigated in the literature.", "limitations": "- The multi-arm bandit modeling module is unclear with potentially high complexity and brittle theoretical guarantee.\n- The task goal is needed at test time, which could further hinder its applicability to general RL problems.", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The proposed pretrain-and-prompt-finetuning paradigm matches the trend of AI research and the future RL systems in real-world deployments.\n- The conceptualization of the prompt and prompt-tuning in RL is interesting and promising.\n- The proposed method is well formulated and presented with comprehensive experiments.", "weaknesses": "- The process of prompt optimization with multi-arm bandit modeling is not given in detail, and seems confusing and unconvincing.\n- The relationship between the proposed method and the hierarchical RL/skill-based RL is not discussed and analyzed.\n- The superiority over existing offline meta-RL methods is not elaborated in the text and not demonstrated in experiments.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "## Thanks for your review! Here, we respond to your comments and address the issues. We hope to hear back from you if you have further questions!\n\n**Q1:** Using BC to train skills will limit the performance when the dataset quality is not high.\n\n**A1:** Indeed, the dataset quality impacts the performance of all imitation learning and offline RL methods. Our approach, however, does not rely solely on expert data but instead utilizes diverse, task-agnostic datasets containing meaningful behaviors, aligning with the literature on task-agnostic pre-training [1,2,3]. This enables the use of abundant data sources such as human gameplay videos for real-world tasks, without requiring expert labeling.\n\nIn Appendix A.3, we present how dataset quality influences performance. Despite not using an expert policy for data collection, our method effectively composes short-term skills to solve long-horizon tasks like MazeRunner. Variations in dataset quality do impact model performance. In section 6, we discuss potential future directions to mitigate this, including the integration of offline RL and online finetuning methods.\n\n**Q2:** Questions about the MAB algorithm.\n\n**A2:** In our MAB framework, each arm corresponds to a prompt of goal sequence. Due to the combinatorial nature of potential prompts in the space $G^K$, this problem is NP-hard. However, we introduce specific insights to reduce the search space effectively, as detailed in Section 4.2:\n\n- Trajectory-based sampling: Instead of sampling prompts in the combinatorial space, we only sample prompts within the collected trajectories, thus grounding the search in feasible and relevant goal sequences.\n- Reward-guided exploration: We further refine the prompt search to only include those from the collected trajectories with best returns, enhancing the likelihood of improving performance.\n- Task-goal consistency: The final goal in each prompt remains fixed as the task goal, ensuring that exploration efforts are aligned with task completion.\n\nThese strategies do not guarantee theoretical optimality due to the problem's complexity. Empirical results, as shown in Table 2 of our attached PDF, indicate that our method outperforms a vanilla MAB approach that utilizes the whole prompt space.\n\nWe acknowledge the need for clearer exposition of our MAB framework and will revise the relevant sections in our manuscript.\n\n**Q3:** Does giving a task goal in test need domain knowledge? \n\n**A3:** No, domain knowledge is not required beyond knowing the task goal, which provides only partial information about the task's objective. The necessary subtasks to complete the task and the environment dynamics are unknown and vary across test tasks, making the challenge distinct from what a traditional goal-conditioned policy might handle. For instance:\n- In Crafter, the task goal is always to unlock all achievements, while the agent must adapt its policy to an unseen world map and identify an optimal sequence of subtasks.\n- In MazeRunner, the task goal is simply to reach a specified position $(x,y)$. The agent faces the challenge of navigating an unfamiliar maze layout to find the most efficient path.\n\nOur experiments show that the Goal-conditioned BC baseline, which is conditioned on the task goal, performs poorly across all tasks. This is because it lacks the ability to adapt to varying environmental dynamics in the test scenarios.\n\n**Q4:** The relationship to HRL/skill-based RL is not discussed.\n\n**A4:** Skill-based RL is closely related to our work, particularly in the context of offline skill pre-training for long-horizon tasks. As detailed in Section 1 (third paragraph) and Section 2 (first paragraph), we discuss relevant literature such as SPiRL[1], PTGM[3], and SKiMo[4], which focus on learning short-term skills offline and adapting to long-horizon tasks through online RL. These methods generally suffer from inefficiencies during online adaptation, which we address through prompt optimization.\n\nHierarchical RL (HRL) also intersects with our study, as it commonly employs pre-trained skills for temporal abstraction in skill-based RL. Existing approaches [1,3,4] rely on HRL to learn a high-level policy during online adaptation. Unlike these methods, our approach sidesteps the inefficiencies associated with HRL by leveraging a pre-trained Transformer-based policy. This allows us to transform the online RL process into a more efficient prompt optimization process. We will add this discussion to our manuscript. \n\n**Q5:** Compare with existing offline meta-RL (OMRL) methods.\n\n**A5:** As outlined in Section 1 (second paragraph), OMRL and our approach target fundamentally different settings.\n-  OMRL necessitates task-specific training data with numerous trajectories per task, whereas our method leverages a **task-agnostic dataset** that captures diverse behaviors.\n- OMRL is typically confined to short-term tasks, enabling fast adaptation from suboptimal trajectories. In contrast, our focus is on **long-horizon tasks**, presenting a greater challenge that is distinct from the scope of OMRL.\n\nThese differences justify our choice of finetuning and skill-based methods for baseline comparison, rather than OMRL approaches.\n\n\n[1] Pertsch et al., Accelerating reinforcement learning with learned skill priors, 2021.\n\n[2] Rosete-Beas et al., Latent plans for task-agnostic offline reinforcement learning, 2023.\n\n[3] Yuan et al., Pre-training goal-based models for sample-efficient reinforcement learning, 2024.\n\n[4] Shi et al., Skill-based model-based reinforcement learning, 2023.", "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": "qOSNCNDqSI", "reviewer": "Reviewer_YjJH", "summary": "The authors apply the concept of prompt optimization to reinforcement learning by pre-training a transformer-based policy on a tasks-agnostic (offline) dataset, injected with subgoals, to fine-tune a meta-policy optimizing the trajectory of goals to accomplish a given task.", "questions": "How does the proposed approach solve the exploration challenge typically coupled with long-horizon RL tasks? \n\nWhat problem is solved if the behavior collecting the data can already solve the intended task?", "limitations": "Limitations could have been discussed more extensively, especially regarding the volatility of a prompt-based architecture.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "The authors propose an exciting approach combining tasks-agnostic pre-training for reinforcement learning with the training paradigms of transformer architectures and large language models.", "weaknesses": "Clarity and overall writing could be improved. Concretely, the introduction seems overly abstract. A running example might help to illustrate the intended motivation. \nAlso, the overall setting remains not entirely clear. Especially the concept of goal relabeling in a task-agnostic setup, requiring a function to map observations to goals, should be explained earlier, given that Fig. 1 is already referenced in the introduction. Furthermore, the overall structure could be improved, by switching sections 2 and 3. \n\nAlso, the definition of M could be improved, e.g., including H, to define a finite horizon MDP.  Also, g seems to map each observation to one specific goal, causing that neither multiple goals could be accomplished concurrently nor an observation with no accomplished goal could occur. Extending the definition of the goal function should by adding its type signature could enhance comprehensibility. Furthermore, basing the goal on the observation seems to restrict its capabilities vastly.\n\nRegarding the general setting, the proposed approach seems to require data that already accomplishes the intended task, vastly reducing the overall impact.  Also, I am missing comparisons to other tasks-agnostic pre-training approaches, like DIAYN (Eysenbach et al. 2018), and a precise delimitation of tasks and goals.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "## Thanks for your review! Here, we respond to your comments and address the issues. We hope to hear back from you if you have further questions!\n\n**Q1:** Paper writing issues.\n\n**A1:** Thank you for your constructive feedback! We've added a running example in Figure 2 of our attached PDF to clarify our approach: using a gameplay dataset from Crafter, we train a goal-sequence-conditioned policy by extracting sequences of unlocked achievements from each trajectory. This policy adapts to unseen environments by optimizing the sequence of achievements to unlock.\n\nRegarding the concept of goal relabeling in a task-agnostic setup, we will provide an early explanation in the text, illustrated by how we define goals using game achievements extracted from observations. We also acknowledge that swapping the order of Sections 2 and 3 would improve clarity, and we plan to adjust our manuscript accordingly in the revision.\n\nWe appreciate your suggestions on including the finite horizon $H$ in the definition of $M$. Given that $M$ represents a partially observable MDP, integrating $H$ will not change its properties, as the timestep $t$ can be regarded as a part of the hidden state. Regarding the goal function, we allow observations without any achieved goals, such as the Kitchen and Crafter environments. We will revise its definition by adding a type signature. Our current framework doesn't handle concurrent goals, which could be a potential area for future exploration.\n\n\n**Q2:** Can the behaviors in the dataset already solve the intended task?\n\n**A2:** The Goal-conditioned BC baseline fails to solve the test tasks in our experiments, indicating that the behavior policies used to collect the datasets cannot solve these tasks.\n\nIn our problem formulation, we do not require data accomplishing the test tasks. Instead, we use **task-agnostic** datasets containing diverse, meaningful behaviors in the environment. Our pre-training scheme allows the policy to learn how to efficiently adapt to **unseen tasks with unknown dynamics and subtasks**. As in the case of the Crafter game, the dataset includes **general gameplay behaviors** without targeted objectives. During testing, our model efficiently optimizes goal sequences for a new world map, determining the most effective order to unlock achievements. We highlight the capability of our approach in efficiently solving new tasks without the need for task-specific pre-training. This aligns with the literature [1,2,3].\n\n\n**Q3:** Comparisons to task-agnostic pre-training methods, like DIAYN.\n\n**A3:** DIAYN, an unsupervised RL method, addresses a different problem from ours by discovering diverse skills online without task-specific rewards. Our approach, in contrast, leverages **offline task-agnostic datasets** to pre-train skills for efficient adaptation to long-horizon tasks. We acknowledge DIAYN's relevance in skill discovery but note that our focus is on enhancing task adaptation efficiency using offline pre-trained skills. Our paper discusses closely related works such as SPiRL[1], PTGM[2], and SKiMo[3], which share similarities with our approach but often struggle with inefficiencies in online adaptation -- a challenge we address through our prompt optimization method.\n\n**Q4:** A precise delimitation of tasks and goals.\n\n**A4:** In our framework, a 'task' is defined with a POMDP and typically involves long-horizon features, requiring sequential execution of many sub-processes. A 'goal' is a description of the agent's achieved state, as defined by our goal function.\n\nFor example, as detailed in Section 3.1, we describe an environment with rooms of unknown layouts where the task is to navigate to a target location. The agent should find the most efficient path, which involves reaching several specific waypoints as necessary goals. Here, the 'task' is to efficiently reach the target location, while each 'goal' represents a waypoint the agent reaches.\n\n\n**Q5:** How does the proposed approach solve the exploration challenge in long-horizon RL tasks?\n\n**A5:**  Our approach addresses the exploration issue in long-horizon tasks with two main strategies:\n- **Temporal abstraction**: The pre-trained goal-conditioned policy leverages goal sequences to reduce exploration to high-level goal achievements, rather than detailed action sequences. This abstraction simplifies exploration, allowing the policy to execute extensive sequences towards each goal. Typically, fewer than 10 goals are required, although up to 500 environment steps are needed in each task in our environments.\n- **Prompt optimization:** Rather than training a high-level policy for goal-switching, our approach optimizes the sequence of goals directly. This avoids the necessity of learning goal selection based on observations during online adaptation, thus simplifying the task complexity and enhancing exploration efficiency.\n\n\n**Q6:** Limitations regarding the volatility of a prompt-based architecture.\n\n**A6:** We acknowledge the limitations of the prompt-based architecture. Just as with language models, our approach can perform unpredictably when faced with out-of-distribution prompts. Additionally, subtle changes to the prompt could lead to unintended behaviors, posing challenges in terms of safety and robustness. In our manuscript revision, we will further discuss these limitations.\n\n[1] Pertsch et al., Accelerating reinforcement learning with learned skill priors, 2021.\n\n[2] Yuan et al., Pre-training goal-based models for sample-efficient reinforcement learning, 2024.\n\n[3] Shi et al., Skill-based model-based reinforcement learning, 2023.", "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:
Thank you to all the reviewers for your insightful comments and constructive feedback. Here, we provide a summary of the reviews and our responses to the key points raised.

### Summary of positive feedback
- All reviewers: The proposed method (multi-goal pre-training and prompt optimization) is promising and exciting. 
- Reviewer JzeR: the method is well formulated.
- Reviewer V4S2 and JzeR: the experimental results are sound and good.
- Reviewer V4S2: the paper is well-written, with good visual examples.

### Common questions addressed
- Comparison with meta-RL / offline meta-RL:  Meta-RL typically assumes access to multiple trajectories for each task and often tackles short-term tasks due to the challenges of few-shot task inference. Our study focuses on pre-training from **task-agnostic, offline** datasets and adapting to **long-horizon** tasks.
- Comparison with other skill-based RL methods: Online skill discovery methods, such as DIAYN [1], learn short-term skills via online RL. Our work aligns with the literature on **offline skill pre-training** [2,3,4] that acquires diverse skills from task-agnostic datasets and composes them for online task adaptation.

For individual issues, we have responded to each reviewer. We are committed to further refining our work based on your feedback. Some figures and tables are provided in the **attached PDF**. Please refer to it if needed.


[1] Eysenbach et al., Diversity is all you need: Learning skills without a reward function, 2018

[2] Pertsch et al., Accelerating reinforcement learning with learned skill priors, 2021.

[3] Yuan et al., Pre-training goal-based models for sample-efficient reinforcement learning, 2024.

[4] Shi et al., Skill-based model-based reinforcement learning, 2023.

Meta-review: The paper presents a method for efficient online adaptation and long-horizon tasks through pretraining transformers for predicting actions based on sequences of goals. The method is evaluated across a diverse set of tasks, demonstrating its practical impact.
However, there are some concerns regarding the presentation and technical explanations that could be improved. The authors are advised to enhance the clarity of their work by broadening the comparative analysis. This can be achieved by including a more detailed comparison with other task-agnostic pre-training methods to better position MGPO within the existing literature. Additionally, the authors should expand the discussion on the limitations of the prompt-based architecture, particularly addressing its volatility and unpredictability. Exploring potential solutions or future research directions that could mitigate these challenges would further strengthen the paper.
Decision: Accept (poster)