Research map / Reinforcement learning / Visual and representation RL
Meta-Reinforcement Learning Robust to Distributional Shift Via Performing Lifelong In-Context Learning
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.
Proposes PSBL, a transformer (LILTrans) meta-trained with posterior-sampling-style losses to amortize the predictive posterior of the optimal policy, enabling lifelong in-context learning at test time that stays robust under task distribution shift in gridworld and MuJoCo tasks.
Problem
Meta-RL agents generalize poorly to out-of-distribution (OOD) tasks when test task distribution shifts from training, and frozen-parameter adaptation fails under such shift.
Approach
Train a transformer variant (LILTrans) with a supervised next-observation/reward loss, an optimal-policy KL loss, and a TD critic loss to directly approximate the posterior predictive distribution of the optimal action; at test time, use a rolling dynamic window so the frozen network keeps improving in context over an extended horizon.
Key results
- On OOD Gridworld and Dark-Key-to-Door, PSBL keeps improving across episodes beyond the training horizon while baselines such as VariBAD decline sharply under distribution shift.
- On HalfCheetah-Vel and Walker-2D-Param OOD tasks, PSBL adapts within the first one or two episodes and outperforms AMAGO, VariBAD, RL2 and PEARL, with the gap growing as shift increases.
- Ablation on Gridworld 6x6 shows a memory length of 10 with horizon 20 outperforms memory 20 (window not advancing) and horizon 12 / memory 6.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use PSBL instead of standard meta-RL baselines | Meta-RL agents · Gridworld and Dark-Key-to-Door OOD tasks | average returns | Helped |
| Use PSBL instead of standard meta-RL baselines | Meta-RL agents · HalfCheetah-Vel and Walker-2D-Param OOD tasks | average returns | Helped |
| Use a dynamic window that rolls forward (ablation) | PSBL · Gridworld 6x6 OOD | average returns | Helped |
| Increase memory length from 6 to 10 episodes while using horizon length 20 (ablation) | PSBL · Gridworld 6x6 OOD | average returns | Helped |
Surprising findings
- VariBAD adapts faster than PSBL on in-distribution tasks but degrades dramatically under distribution shift.
Limitations
- Evaluation limited to small gridworld, Dark-Key-to-Door and two MuJoCo tasks.
- Some components (critic, annealing, dynamic window) add hyperparameters; ablation is limited to gridworld 6x6 OOD.
- Theoretical justification relies on prior work (Ortega et al., Mikulik et al.) and an appendix proof rather than new guarantees.
Experimental setup
Datasets: Gridworld (5x5 training, 6x6-8x8 OOD) Dark-Key-to-Door (9x9 training, 10x10-12x12 OOD) HalfCheetah-Vel Walker-2D-Param
Benchmarks: Gridworld Dark-Key-to-Door HalfCheetah-Vel Walker-2D-Param
Metrics: average return in ID and OOD tasks
Baselines: AMAGO VariBAD RL2 PEARL
Builds on: PSRL PEARL VariBAD AMAGO Ortega et al. 2019 memory-based meta-learning
Closest papers by meaning
- Towards Large-Scale In-Context Reinforcement Learning by Meta-Training in Randomized Worlds (NeurIPS 2025 · similarity 0.89)
- MetaCARD: Meta-Reinforcement Learning with Task Uncertainty Feedback via Decoupled Context-Aware Reward and Dynamics Components (AAAI 2024 · similarity 0.88)
- Decoupling Meta-Reinforcement Learning with Gaussian Task Contexts and Skills (AAAI 2024 · similarity 0.88)
- Meta-Reinforcement Learning with Universal Policy Adaptation: Provable Near-Optimality under All-task Optimum Comparator (NeurIPS 2024 · similarity 0.88)
- DynaMITE-RL: A Dynamic Model for Improved Temporal Meta-Reinforcement Learning (NeurIPS 2024 · similarity 0.88)
- CERTAIN: Context Uncertainty-aware One-Shot Adaptation for Context-based Offline Meta Reinforcement Learning (ICML 2025 · similarity 0.88)
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.