Pre-Trained Multi-Goal Transformers with Prompt Optimization for Efficient Online Adaptation

Authors: haoqi  yuan
yuhui  fu
feiyang  xie
zongqing  lu
… … 𝑜
  𝑎
𝐻  𝑜
𝐻  𝑎
  𝑜
𝑖  𝑜
  𝑝~𝜏
𝑜  𝐻
ො 𝑎 𝐻
𝑔  𝑖
𝑔  𝐻
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-02 07:10:19

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.

Main paper: 
 introduction: 
 In the evolving landscape of deep learning, the paradigm of 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, 41] and natural language processing (NLP) [30, 32] . This paradigm has recently been explored in deep reinforcement learning (RL) [3, 4] , addressing the issue of sample efficiency in RL when solving unseen tasks by leveraging the acquired knowledge during pre-training. For example, offline meta-RL (OMRL) [33, 59, 56] studies pre-training a policy on multi-task datasets, which adapts to a 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 focus on pre-training policies [4, 54] and skills [38, 47] on task-agnostic datasets and adapt to unseen 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., MuJoCo [50] and MetaWorld [58] ). These approaches often fall short of tackling long-horizon tasks which are characterized by a sequence of sub-processes. In our study, we consider this challenging yet realistic scenario: given a task-agnostic dataset characterized by diverse behaviors of the agent, we aim to pre-train a policy to facilitate efficient online adaptation to new, long-horizon tasks.Recent studies [38, 60] tackle this problem by assuming that diverse short-term skills can be acquired from task-agnostic datasets, which are sufficient to 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 RL 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, it 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.Is it feasible to pre-train a policy conditioned on a sequence of goals, capable of autonomously 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 a 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 on a long history context.To tackle this challenge, we draw inspiration from the success of Transformers [51] in language modeling [10, 7] , which have demonstrated a remarkable capability to model long sequences and, when pre-trained on diverse datasets, adapt efficiently to downstream tasks through prompt optimization [25] . In our proposed Multi-Goal Transformers with Prompt Optimization (MGPO), we pre-train a Transformer-based policy that takes a goal sequence as a prompt and predicts actions in the following sequence of environment observations. During online adaptation, we only optimize for the sequence of goals in 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 a novel prompt optimization strategy, leveraging the returns of online trajectories to guide the selection of goal sequences. This process is formulated as a multi-armed bandit problem and we introduce two approaches through upper confidence bound (UCB) and Bayesian posterior estimation. Figure 1 provides an overview of our framework.We evaluate MGPO across diverse domains, including maze navigation, the robotic simulation environment 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. Comparisons 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 adaptation 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 interpretability 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. 
problem formulation: 
 A task in an environment is formulated as a partially observable Markov Decision Process (POMDP) M = ⟨S, O, A, T, ρ, r, γ⟩ 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 a reward. At each timestep (π) = E at∼π(at|ht-1,ot) [R], where R = H t=1 γ t r(s t , a t )is the discounted cumulative return of an episode. We focus on long-horizon tasks that require executing a substantial number of subtasks sequentially. We define a goal function g = f G (o) that maps an observation to a goal, representing key information related to task completion, such as the agent's position in navigation tasks or the states of objects in robotic scenarios. G denotes the goal space. A goal g is said to be reached at timet if g t = f G (o t ).We assume that the task M in an unknown environment provides a task goal g M , indicating the final goal to be reached. To optimally solve the long-horizon task, the policy must sequentially reach multiple goals. As an example, consider an environment resembling rooms with unknown structures where the task is to navigate to a goal location g M . 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 (g 1 , ..., g k ) as necessary goals.During pre-training, we assume access to a task-agnostic dataset D = {τ i } N i=1 consisting of trajectories τ = (o 1 , a 1 , ..., o H , a H ), in which the agent performs diverse behaviors. For online adaptation, we aim to solve an unseen task M with its goal g M provided, within an online interaction budget of N episodes. The objective is to find the optimal policy π maximizing the expected return J(π).Some previous methods [4] pre-train a policy π θ (a t |h t-1 , o t ) parameterized by θ and then finetune with online RL. This approach is sample-inefficient due to discrepancies between the behaviors in the dataset and those required for downstream tasks. Other methods [38, 60] pre-train a policy π θ (a t |z, h t-1 , o t ) conditioned on a variable z to perform diverse short-term skills and then train a high-level policy π ϕ (z|h t , o t ) online using RL. As discussed earlier, this approach often results in limited sample efficiency because the high-level policy must learn when to switch skills z based on 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 our policy, known for its effectiveness in sequence modeling tasks. Transformers leverage an attention mechanism [51] , enabling the model to weigh different parts of the input sequence differently, thus effectively addressing long-term dependencies.In pre-training, we aim to develop a Transformer-based policy π θ (a t |p, h t-1 , o t ) parameterized by θ, taking an additional input variable p. This policy processes inputs in an auto-regressive manner, where the input is a sequence (p, o 1 , a 1 , ..., o H , a H ), and outputs an action a t at each observation o t . Here, the input variable p acts similarly to a prompt in language models, guiding the policy's behavior. During online adaptation, we keep the model parameters θ fixed and treat p as the optimization variable, transforming the problem into finding an optimal prompt p:EQUATIONThis approach mirrors the concept of prompt optimization in language models and eliminates the need to use online RL to train a high-level policy.3 Method 
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 1 , ..., o H ) in a trajectory τ , we construct a sequence of goals g = (g 1 , ..., g H ) with the same length H, representing the process of agent behavior in this episode. We first uniformly sample a number k ∼ U [0, K -1] where K is the maximal prompt length, then uniformly sub-sample the goal sequence to construct a prompt p = (g i1 , ..., g i k , g H ), where 0 ≤ i 1 ≤ ... ≤ i k < H. Since the trajectory ends at the last goal g H , we keep g H at the last position in the prompt. Using this sampling mechanism, the prompt can describe behavior in different granularities with its varying length. We denote the process of sampling a prompt from the trajectory as p ∼ P (p|τ ).Similar to Decision Transformers [8] , we adopt a causal Transformer to build the policy π θ (a t |p, h t-1 , o t ), which takes as input a sequence concatenating the prompt and the trajectory:(p, τ ) = (g i1 , ..., g i k , g H , o 1 , a 1 , ..., o H , a H ).At each input o t , the policy only sees the sub-sequence from g i1 to o t due to the causal attention mask and predicts the action distribution π θ (a t |p, h t-1 , o t ).We sample batches of trajectories in D and train the policy using behavior cloning:EQUATIONThis training scheme encourages the policy to utilize the information provided in the prompt to reduce uncertainty in 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 a trade-off between exploration and exploitation in online adaptation.AMAGO [15] also adopts a 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.Given an unseen task with the task goal g M , we start with an initial prompt p 0 = (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 an optimized prompt. In principle, we can employ any black-box optimization approach for this purpose, including discrete prompt search [9, 40] and continuous prompt-tuning [48, 27] methods. However, prompts optimized in the entire prompt space G K 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 a novel method that samples prompts from online collected trajectories, aligning more closely with the training distribution of prompts. We assume that, if we condition the policy on prompts sampled from a trajectory with high return, it 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 τ * of the highest return R * . Each iteration involves: (1) selecting a prompt p * from B most likely to yield high returns and collecting a trajectory with π(a|p * , h, o); (2) sampling a new prompt p ′ ∼ P (p|τ * ), which is then added to the buffer with the return obtained with π(a|p ′ , h, o). Prompt selection in (1) , which requires acting under uncertainty given observed returns in history, can be modeled as a multi-arm bandit (MAB) problem. We implement two optional solutions including upper-confidence bound (UCB) [2] and a method based on 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) Trajectorybased 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) Rewardguided exploration: We further refine our prompt search by selecting prompts from trajectories that have highest returns, thereby enhancing the likelihood of performance improvement. (3) Taskgoal 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 1 in Appendix E.1, where the implementations of UCB and BPE are also provided.Question: For each theoretical result, does the paper provide the full set of assumptions and a 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 crossreferenced. • All assumptions should be clearly stated or referenced in the statement of any theorems.• The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition. • Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material. • Theorems and Lemmas that the proof relies upon should be properly referenced. 
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?Our evaluation spans multiple domains featuring long-horizon tasks. We collect datasets to 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 a designated goal position (x, y). The agent's observation includes its position and nearby terrain, receiving a reward of +1 for reaching the goal and a -0.1 penalty for each timestep. We collect the dataset with a handcrafted 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 × 15 and 30 × 30, where the latter has a longer horizon of 500 steps.Kitchen and GridWorld: Kitchen [16] is a robotic environment with continuous observations and actions, where a 7-DoF robot arm manipulates diverse objects in a simulated kitchen scene. We define each long-horizon task as completing a sequence of n subtasks in a specific order, providing + 1 n reward when the next correct subtask is completed. The task goal is provided as a set of subtasks without revealing their order. Other goals are represented as 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 [18] : A simplified benchmark of the open-world game Minecraft, where the 2D world is procedurally generated. The agent receives 64 × 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 a unique, unexplored world map.Table 1 : Performance of MGPO compared with baseline 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. Goal-conditioned BC has no error bars since it does not perform online optimization. We compare MGPO with previous methods that focus on task-agnostic pre-training. Further details are provided in Appendix F.Goal-conditioned BC feeds the task goal g M 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-trained 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 uses online RL to train a 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 performance by the average return of the optimized policy or prompt. Results in Table 1 summarize the performance of MGPO against all baselines across different environments, where MGPO is implemented with UCB in online adaptation.• MGPO's superiority: MGPO demonstrates superior performance in all environments, showcasing its efficiency in adapting to new tasks.• Limitations of BC-finetune: This method underperforms others in most environments due to 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 GridWorld. These approaches are limited by their reliance on learning a high-level RL policy for goal switching at different states during online adaptation. For example, PTGM trains a high-level policy π(g|h t , o t ) to select goals, which is effectively searching policies in the joint space of O K × G.In contrast, MGPO, with its inherent ability to seamlessly switch between goals, optimizes in the prompt space G K , leading to more efficient adaptation.• Goal-conditioned BC: 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 Goalconditioned 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 environment information in the unseen task and can search for a prompt sequence to stitch different short-term behaviors learned from data.• Performance in Crafter: Performance of Goal-conditioned BC, PTGM, and MGPO in Crafter is similar compared with other environments. This is 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. In contrast, in other environments, the prompt substantially influences the policy's behavior and performance. 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 a 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 [48] is a black-box optimization method for continuous prompt-tuning. It utilizes a lowdimensional vector z, which is added to the initial prompt through random projection. This vector is then optimized using a CMA-ES [19] evolution strategy. Since it extends the prompt into a 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. UCB and BPE consistently outperform MGPO-explore in most environments. This result demonstrates 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-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 of deterministic and differentiable transitions in Kitchen, the continuous optimization method BBT may quickly find a 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 3 compares the performance of the optimized prompts BBT GRIPS UCBIn 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 4 .Maximal Prompt Lengths: We study the impact of different maximal prompt lengths K during online adaptation. Complex tasks often necessitate multiple goals within a prompt to specify varying behaviors in long-horizon trajectories. We observe that optimizing prompts with increasing length consistently improves performance. Specifically, a length of K = 5 yields the best results, and reducing K to 1 results in performance akin to that of Goal-conditioned BC. Appendix B.2 shows additional results when further increasing K.UCB Hyperparameters: UCB includes a hyperparameter c that balances exploration and exploitation (Appendix E.1). In our test with MGPO-UCB under varying values of c, we observe that the performance exhibits low sensitivity to this hyperparameter's selection, as shown in Table 4 .Dataset Quality: We also assess the influence of dataset quality on MGPO's performance. Detailed results are provided in Appendix B.3.In MazeRunner, we visualize the evolution of MGPO during online adaptation. Figure 3 displays the progression of both the prompt and policy behavior at various stages. At the beginning, the policy is 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 are iteratively refined based on returns from the environment, and the policy behavior is steered towards more efficient paths. More results are presented in Appendix B.4.Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of 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, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not. • If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable. • Depending on the contribution, reproducibility can be accomplished in various ways.For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed. • While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. , with an open-source dataset or instructions for how to construct the dataset). (d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results. 
related work: 
 Policy Pre-Training for RL is a 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] or 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, 54] , or hierarchical skill learning [38, 45, 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 [8, 22, 14, 57] recast RL as a sequence modeling problem conditioned on return-to-go. In contextbased meta-RL, Transformers are employed to handle multiple trajectories, modeling task adaptation as in-context learning [23, 56, 29] . In multi-task RL, large-scale Transformers are adopted with extensive datasets to address complex robotic domains [6, 5] and generalization to various tasks [44, 49, 15] . In our work, we utilize Transformers to address the demands of long-term memory in modeling long-horizon trajectories with multiple goals.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 a sequence of tokens input to the model specifying the task. In the recent large language models (LLMs) [7] , task adaptation can be easily achieved by prompt design, with the techniques of in-context learning [11] and chain-of-thought reasoning [52] . For earlier language models, prompt-tuning methods optimize prompts in a continuous space [25, 27] , while prompt search methods focus on optimizing discrete prompt tokens for interpretability [9, 40] . We adopt the concept of prompt optimization to address online adaptation in RL, focusing on optimizing the goal sequence for the pre-trained Transformer. Unlike in NLP tasks, our approach places a unique emphasis on sample efficiency, a crucial aspect given the necessity of online evaluation.Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?Answer: [Yes]Justification: We have released the code.Guidelines:• The answer NA means that paper does not include experiments requiring code. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
conclusion and limitations: 
 We propose MGPO, a novel framework for policy 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-horizon tasks opens new possibilities for real-world applications where efficient online adaptation is crucial. However, our 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 more complex, real-world environments.Like many offline RL approaches, the performance of MGPO is 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.A Task: unlock all achievements in an unseen world. Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to 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 in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 
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 deviations across 3 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, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors).• It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified. • For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). • If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [Yes] Justification: Details are presented in 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, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
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 societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
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 (e.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 a 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 a URL.• The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 
 
 Appendix: Figure 6 shows the learning curves of MGPO and baselines during online adaptation. Figure 7 shows the learning curves of MGPO with different prompt optimization methods during online adaptation.Since we use K = 5 in the pre-training phase and the model has not encountered prompts exceeding a 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, 20 and 40 to test online adaptation with larger prompt lengths. Figure 8 shows the results. In MazeRunner-15, setting K = 10 yields the optimal performance. Increase beyond this value leads to a decline in learning efficiency. In MazeRunner-30, performance exhibits a slight improvement as K extends to 40. This suggests that for more extended, long-horizon tasks, the exploration of longer prompts can be advantageous. We assess MGPO's performance with datasets of varying quality in MazeRunner. In Table 5 , we denote datasets collected using an A* algorithm with a maximum of n goal switches per episode as A*-n, and datasets from a random exploration policy as Random. Results show that MGPO achieves better performance trained on A*-2 datasets than A*-1, indicating its efficacy with data containing diverse long-horizon behaviors. The comparatively lower performance on the A*-4 dataset in MazeRunner-30 and Random datasets suggests MGPO's reliance on the quality of data collection policies. Future work could involve integrating offline RL methods to boost performance with lower-quality datasets. 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 . We observe that a shorter prompt, by providing less information to the policy, introduces higher uncertainty and thus encourages more exploratory behaviors. Conversely, a longer prompt, containing more detailed information from the trajectory, prompts more deterministic and exploitatory behavior.Figure 10 shows additional results on the progression of both the prompt and policy behavior at various stages. MazeRunner is a 2D maze navigation environment implemented in AMAGO [15] and is mainly based on Memory Maze [37] . The agent navigates in a grid world of size N × N , where we use N = 15 and 30 in our experiments. For each task, the unknown structure of walls in the maze is distinct, and the agent should move to the provided goal location g M = (x M , y M ) to accomplish the task. The agent receives 6-dimensional observations (x, y, d 1 , d 2 , d 3 , d 4 ), 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 use the agent's positions in observations to 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 a penalty of -0.1 is given at each step in addition. The task horizon is 64 steps for MazeRunner-15 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 a path to reach it with A* algorithm. The A* algorithm takes the partial observations of the agent, estimate a 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 4 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.The Kitchen [16] 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 to adjust the corresponding object to a target pose. The observation is a 59-dimensional 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 a set of n subtasks among the five subtasks and the order to complete them. The task horizon is set to 500 timesteps. When a subtask is completed in the correct order, the environment provides a reward of + 1 n . The goal space consists of 5-dimensional binary vectors representing each subtask in each bit. For the task goal revealed to the agent, bits of the selected subtasks are set to 1 and other bits are set to 0. 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 a policy for each sub-task using a shaped reward function and varied initial states from the Kitchen-mixed-v0 dataset [13] . 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 a test set of 25 tasks, which has no overlap to tasks used in data collection.This environment is a 10 × 10 grid world with 7 distinct objects at different locations. Each object has a binary state, which defines a subtask that is completed when the state turned to 1 from the initial state of 0. The agent observes a 9-dimensional vector representing its location and states of objects. The discrete action space consists 6 actions: moving to the four directions, flip the object state, and no operation. Since this environment is a simplified version of Kitchen, tasks and goals are defined similar to 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.Crafter [18] 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 × 3 rgb images showing the first-person view of the agent, which contains a 9 × 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, it receives a reward of +1. When the agent's health increases or decreases by one point, it receives a corresponding reward of +0.1 or -0.1. The target of the task is to complete all achievements. Different tasks are featured with different generated world maps, which can be specified with a 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 to 500 steps. We use the policy trained in Achievement Distillation [35] to collect a 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.We build our Transformer architecture based on Decision Transformers [8] , using a lightweight backbone of GPT-2 [42] with 0.6M parameters. Input tokens of o, a, and g are transformed into 128dimensional 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 a sequence with the same dimension and length. We use a linear layer to decode tokens at positions of the input tokens of observations to predict actions and compute loss. To train on batches of trajectories, we pad the input prompts and trajectories to lengths of K and H, respectively.In MazeRunner, we sample prompts from the agent's locations in the whole trajectory. To augment the diversity of task goals and trajectory lengths, we truncate the trajectory at a random timestep h for each sampled trajectory and use o h to represent its task goal. In other environments, goals in prompts are sampled from the completed subtasks in each trajectory and g H represents the actual set of subtasks completed at the end of each trajectory. All models are trained on a lab machine with a single NVIDIA RTX 4090 GPU and Intel i9 CPUs. For each environment, the pre-training stage takes about 12 hours.We summarize our proposed prompt optimization algorithm in Algorithm 1.GRIPS [40] is a gradient-free, search-based algorithm for prompt optimization. At each iteration, it 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 a random position. Del: Randomly delete a goal in the current prompt. Swap: Randomly select two goals and swap them in the current prompt.We initialize the prompt with p 0 = (g M ) and maintain g M at the last position of the prompt while editing the preceding goals. In each iteration, we randomly generate n = 5 edited prompts, online evaluate their performance, and preserve the prompt with the highest return.We summarize GRIPS in Algorithm 2.Input: pre-trained policy π θ ; task goal g M ; interaction budget N ; the number of edited prompts per iteration n Initialize: best return R * = -∞ and best prompt p * = (g M ) while not exceeding N episodes do Edit p * with random operations from {Add, Del, Swap} to generate prompts{p i } n i=1 . Online evaluate each p i with π θ (a|p i , h, o), get return R i . i * = argmax i R i . if R i * > R * then R * = R i * , p * = p i * . end if end while E.3 BBTBBT [48] is a black-box continuous optimization method for prompt-tuning. We use p = (g M ) to collect a trajectory τ and sample a initial prompt with p 0 ∼ τ . In the following iterations, we optimize z ∈ R d in a low-dimensional space (d < D, where D is the dimension of prompt) and use a random projection matrix A ∈ R D×d to project z and add to the initial prompt. Our objective is:EQUATIONwhere R is the expected return of the policy with prompt Az + p 0 .BBT uses CMA-ES for optimization. In each iteration t, we sample λ = 10 new vectors z (t+1) from the multivariate normal distribution z∼ m(t) + σ(t)N (0, C(t)), i = 1, ..., λ, where m(t) ∈ R d is the mean vector of the search distribution at iteration t, σ(t) ∈ R + is a standard deviation controlling the step length, and C(t) ∈ R d×d is a 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.Input: pre-trained policy π θ ; task goal g M ; interaction budget N ; the population size λ Initialize: best return R * = -∞; best prompt p * ; m(0), σ(0), C(0), and A Collect an episode τ with p = (g M ). Sample initial prompt p 0 ∼ τ . while not exceeding N episodes do Sample z i , i = 1, ...λ from the multivariate normal distribution. Collect τ i and get return R i with π θ (a|Az i + p 0 , h, o) for each i. Update R * and p * . Update m, σ, C with CMA-ES. end while of the relabeled goal. In online adaptation, we use the discrete goal space as the high-level action space. We use DQN [34] to train the high-level policy π H ϕ (g|h t , o t ). Its architecture is an MLP with 3 layers and hidden layer dimensions of 128. We also incorporate a KL regularization reward between the policy and the goal prior p ψ to accelerate online RL. Meta-RL / Offline Meta-RL: Meta-RL approaches [56, 21, 59 ] 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 on pre-training from task-agnostic, offline datasets and adapting to long-horizon tasks.Skill-based RL: Online skill discovery methods [12, 61, 31] primarily focus on learning short-term skills through online RL. In contrast, our work focuses on offline skill acquisition from diverse, task-agnostic datasets, as conceptualized in the literature [38, 47, 60] . These methods typically use hierarchical RL to learn a high-level policy during online adaptation. Our approach differs by avoiding the inefficiencies of online RL; instead, we utilize a pre-trained Transformer-based policy, enabling us to streamline the adaptation process through efficient prompt optimization.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, it also carries inherent risks such as 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 or autonomous driving, raises significant ethical considerations. Ensuring these systems are used responsibly is paramount to prevent negative societal impacts.NeurIPS Paper Checklist Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 


Decision: 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.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: See questions.
questions: 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?

- How to determine the length of the sequences of goals while online adaptation. Is there any terminal signal or pre-defined length?

- 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?

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

- 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?


- Can you provide analysis about why the BC-finetune perform worse than Goal-conditioned BC?

- Some typos: like the bold formatting of the citation text on line 233.
limitations: limitations: See questions.
paper_summary: paper_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.
strength_weakness: strength_weakness: - The process of prompt optimization with multi-arm bandit modeling is not given in detail, and seems confusing and unconvincing.
- The relationship between the proposed method and the hierarchical RL/skill-based RL is not discussed and analyzed.
- The superiority over existing offline meta-RL methods is not elaborated in the text and not demonstrated in experiments.
questions: 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?
- 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?
- 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.
- 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: limitations: - The multi-arm bandit modeling module is unclear with potentially high complexity and brittle theoretical guarantee.
- The task goal is needed at test time, which could further hinder its applicability to general RL problems.
paper_summary: paper_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.
strength_weakness: strength_weakness: Clarity and overall writing could be improved. Concretely, the introduction seems overly abstract. A running example might help to illustrate the intended motivation. 
Also, 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. 

Also, 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.

Regarding 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.
questions: questions: How does the proposed approach solve the exploration challenge typically coupled with long-horizon RL tasks? 

What problem is solved if the behavior collecting the data can already solve the intended task?
limitations: limitations: Limitations could have been discussed more extensively, especially regarding the volatility of a prompt-based architecture.

Comments: comment: Thank you! We appreciate your time and your constructive feedback.
Thank you for your extensive response and for providing additional explanations and clarifications. My concerns are mostly addressed, and I will raise my score accordingly.
Thank you very much for raising the score! We appreciate your constructive feedback.
Thank you for your time. Most of my concerns have been addressed, and I will increase my rating.
Dear Reviewers,

Thank you for your insightful comments. We have addressed your initial questions through our rebuttal and are eager to clarify any further points you might raise. Please feel free to provide additional feedback. We greatly appreciate your continued engagement.

Best regards,

Authors
Thanks again for your time reviewing our paper and for your constructive feedback!
Thank the authors for their detailed response, which has addressed most of my concerns. I will keep my score.