Research map / Reinforcement learning / Visual and representation RL
Pre-Trained Multi-Goal Transformers with Prompt Optimization for Efficient Online Adaptation
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Pre-trains a small causal Transformer on task-agnostic data with hindsight multi-goal relabeling and behavior cloning, then adapts to unseen long-horizon tasks by optimizing the goal-sequence prompt with a multi-armed bandit (UCB/BPE) rather than online RL.
Problem
Efficient online adaptation to unseen long-horizon tasks composed of multiple subtasks, where existing pre-trained policy methods need many RL interactions to learn goal switching.
Approach
Pre-train a goal-sequence-prompted Transformer policy via hindsight multi-goal relabeling and behavior cloning so it can switch goals autonomously, then treat the goal prompt as the only optimization variable online, selecting prompts with a multi-armed bandit over prompts sampled from the best online trajectories.
Key results
- MGPO-UCB reaches -0.41 on MazeRunner-15 and -14.21 on MazeRunner-30, versus -0.96 (PTGM) and -2.62 (SPiRL) on MazeRunner-15.
- On Kitchen, MGPO-UCB achieves 0.63 success-style return versus 0.25 for PTGM and 0.22 for SPiRL; on GridWorld 0.58 versus 0.25 for PTGM.
- Under noisy observations or actions in Kitchen, trajectory-sampled MGPO prompts degrade less than BBT and GRIPS prompts (e.g., BBT drops by 0.55 under noisy actions versus 0.37 for MGPO-BPE).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use MGPO-UCB instead of the strongest baseline PTGM | MGPO · MazeRunner-15 | average return: -0.96 ± 0.08 -> -0.41 ± 0.10 | Helped |
| Use MGPO-UCB instead of the strongest baseline PTGM | MGPO · MazeRunner-30 | average return: -26.74 ± 1.71 -> -14.21 ± 0.19 | Helped |
| Use MGPO-UCB instead of the strongest baseline PTGM | MGPO · Kitchen | average return: 0.25 ± 0.03 -> 0.63 ± 0.02 | Helped |
| Use MGPO-UCB instead of the strongest baseline PTGM | MGPO · GridWorld | average return: 0.25 ± 0.03 -> 0.58 ± 0.03 | Helped |
| Use MGPO-UCB instead of PTGM | MGPO · Crafter | average return: 15.72 ± 0.12 -> 15.66 ± 0.14 | Hurt |
| Use a maximal prompt length of 5 rather than 1 (ablation) | MGPO-UCB · MazeRunner-15, MazeRunner-30, and Kitchen | average return: -2.63±0.00 -> -0.41±0.10; -27.06±0.00 -> -14.21±0.19; 0.09±0.00 -> 0.63±0.02 | Helped |
| Set maximal prompt length to 2 rather than 1 (ablation) | MGPO-UCB · MazeRunner-15, MazeRunner-30, and Kitchen | average return: -2.63±0.00 -> -1.13±0.10; -27.06±0.00 -> -19.57±1.21; 0.09±0.00 -> 0.28±0.00 | Helped |
| Set maximal prompt length to 3 rather than 2 (ablation) | MGPO-UCB · MazeRunner-15, MazeRunner-30, and Kitchen | average return: -1.13±0.10 -> -0.77±0.04; -19.57±1.21 -> -17.12±0.47; 0.28±0.00 -> 0.48±0.02 | Helped |
| Set maximal prompt length to 5 rather than 3 (ablation) | MGPO-UCB · MazeRunner-15, MazeRunner-30, and Kitchen | average return: -0.77±0.04 -> -0.41±0.10; -17.12±0.47 -> -14.21±0.19; 0.48±0.02 -> 0.63±0.02 | Helped |
| Set UCB exploration hyperparameter c to 0 rather than 1 (ablation) | MGPO-UCB · MazeRunner-15, MazeRunner-30, and Kitchen | average return: -0.41±0.10 -> -0.44±0.07; -14.21±0.19 -> -15.86±0.51; 0.63±0.02 -> 0.63±0.05 | Mixed |
| Set UCB exploration hyperparameter c to 10 rather than 1 (ablation) | MGPO-UCB · MazeRunner-15, MazeRunner-30, and Kitchen | average return: -0.41±0.10 -> -0.40±0.04; -14.21±0.19 -> -14.70±0.66; 0.63±0.02 -> 0.66±0.02 | Mixed |
| Use UCB or BPE prompt selection instead of the most-recent-prompt MGPO-explore strategy (ablation) | MGPO · most environments | average return | Helped |
Surprising findings
- BBT, a continuous prompt optimizer, achieves the best Kitchen result among prompt methods (0.74) but its prompts are out-of-distribution and less robust to noise.
- In Crafter, goal-conditioned BC, PTGM and MGPO perform similarly (about 15.7 to 11.8), because the dataset lacks diversity and the policy is less sensitive to the prompt.
Limitations
- Performance depends on the quality of the pre-training dataset; low-quality or non-diverse data reduces gains.
- Experiments are confined to simulated environments; real-world transfer is untested.
- Prompt-based policies can behave unpredictably under out-of-distribution or slightly altered prompts, raising robustness and safety concerns.
Experimental setup
Datasets: MazeRunner task-agnostic dataset (collected with A* handcrafted policy) Kitchen dataset collected with PPO-trained subtask policies GridWorld dataset (rule-based expert) Crafter dataset (AD policies, Achievement Distillation)
Benchmarks: MazeRunner-15 MazeRunner-30 Kitchen GridWorld Crafter
Metrics: average return after 100 online episodes standard deviation across 3 seeds
Baselines: Goal-conditioned BC BC-finetune SPiRL PTGM GRIPS BBT
Builds on: Decision Transformer AMAGO SPiRL PTGM GRIPS BBT
Closest papers by meaning
- DaCe-DT: Data-Centric Offline Multi-Task Reinforcement Learning via Adaptive Prompts and Trajectory Correction for Heterogeneous Tasks (NeurIPS 2026 · similarity 0.88)
- Prompt Tuning Decision Transformers with Structured and Scalable Bandits (NeurIPS 2025 · similarity 0.87)
- AMAGO-2: Breaking the Multi-Task Barrier in Meta-Reinforcement Learning with Transformers (NeurIPS 2024 · similarity 0.87)
- Adaptive Advantage-Guided Policy Regularization for Offline Reinforcement Learning (ICML 2024 · similarity 0.87)
- Online Reinforcement Learning in Non-Stationary Context-Driven Environments (ICLR 2025 · similarity 0.87)
- From Weak Data to Strong Policy: Q-Targets Enable Provable In-Context Reinforcement Learning (NeurIPS 2026 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.