Distilling Reinforcement Learning Algorithms for In-Context Model-Based Planning
Not on the map. Why?
Abstract (arXiv)
Recent studies have shown that Transformers can perform in-context reinforcement learning (RL) by imitating existing RL algorithms, enabling sample-efficient adaptation to unseen tasks without parameter updates. However, these models also inherit the suboptimal behaviors of the RL algorithms they imitate. This issue primarily arises due to the gradual update rule employed by those algorithms. Model-based planning offers a promising solution to this limitation by allowing the models to simulate potential outcomes before taking action, providing an additional mechanism to deviate from the suboptimal behavior. Rather than learning a separate dynamics model, we propose Distillation for In-Context Planning (DICP), an in-context model-based RL framework where Transformers simultaneously learn environment dynamics and improve policy in-context. We evaluate DICP across a range of discrete and continuous environments, including Darkroom variants and Meta-World. Our results show that DICP achieves state-of-the-art performance while requiring significantly fewer environment interactions than baselines, which include both model-free counterparts and existing meta-RL methods.
Key facts. Cited by 12 papers in the map (data through 7 Oct 2026). Builds on 52 of its 56 references that are in the map. First posted on arXiv 26 Feb 2025.
Builds on
52 of 56 references.
References · 52
- 2017Attention Is All You Needcited by 59,932
- 2020Language Models are Few-Shot Learnerscited by 32,199
- 2020An Image is Worth 16x16 Words: Transformers for Image Recognition at Scalecited by 29,915
- 2017Proximal Policy Optimization Algorithmscited by 14,441
- 2022Chain-of-Thought Prompting Elicits Reasoning in Large Language Modelscited by 12,302
- 2017Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networkscited by 4,861
- 2018Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actorcited by 4,635
- 2023Sparks of Artificial General Intelligence: Early experiments with GPT-4cited by 2,579
- 2015Trust Region Policy Optimizationcited by 2,503
- 2020Offline Reinforcement Learning: Tutorial, Review, and Perspectives on Open Problemscited by 1,873
- 2021Decision Transformer: Reinforcement Learning via Sequence Modelingcited by 1,528
- 2019Mastering Atari, Go, Chess and Shogi by Planning with a Learned Modelcited by 1,327
38 more
- 2023Mastering Diverse Domains through World Modelscited by 1,193
- 2019Meta-World: A Benchmark and Evaluation for Multi-Task and Meta Reinforcement Learningcited by 1,168
- 2018On First-Order Meta-Learning Algorithmscited by 1,003
- 2022A Generalist Agentcited by 802
- 2018Recurrent World Models Facilitate Policy Evolutioncited by 761
- 2019When to Trust Your Model: Model-Based Policy Optimizationcited by 753
- 2021Offline Reinforcement Learning as One Big Sequence Modeling Problemcited by 636
- 2024TinyLlama: An Open-Source Small Language Modelcited by 604
- 2019Model-Based Reinforcement Learning for Ataricited by 495
- 2016RL2: Fast Reinforcement Learning via Slow Reinforcement Learningcited by 490
- 2019Efficient Off-Policy Meta-Reinforcement Learning via Probabilistic Context Variablescited by 372
- 2016Learning to reinforcement learncited by 369
- 2017Neural Network Dynamics for Model-Based Deep Reinforcement Learning with Model-Free Fine-Tuningcited by 288
- 2018Learning to Adapt in Dynamic, Real-World Environments Through Meta-Reinforcement Learningcited by 270
- 2021Mastering Atari Games with Limited Datacited by 235
- 2022Multi-Game Decision Transformerscited by 212
- 2019VariBAD: A Very Good Method for Bayes-Adaptive Deep RL via Meta-Learningcited by 210
- 2022In-context Reinforcement Learning with Algorithm Distillationcited by 186
- 2023Supervised Pretraining Can Learn In-Context Reinforcement Learningcited by 137
- 2018Model-Based Reinforcement Learning via Meta-Policy Optimizationcited by 99
- 2020Offline Meta-Reinforcement Learning with Advantage Weightingcited by 92
- 2021Online and Offline Reinforcement Learning by Planning with a Learned Modelcited by 90
- 2018Meta Reinforcement Learning with Latent Variable Gaussian Processescited by 44
- 2022Robust Task Representations for Offline Meta-Reinforcement Learning via Contrastive Learningcited by 33
- 2023Emergent Agentic Transformer from Chain of Hindsight Experiencecited by 33
- 2020Model-based Adversarial Meta-Reinforcement Learningcited by 32
- 2021Procedural Generalization by Planning with Self-Supervised World Modelscited by 30
- 2020Exploration in Approximate Hyper-State Space for Meta Reinforcement Learningcited by 28
- 2023Emergence of In-Context Reinforcement Learning from Noise Distillationcited by 26
- 2024In-Context Decision Transformer: Reinforcement Learning via Hierarchical Chain-of-Thoughtcited by 25
- 2023In-Context Reinforcement Learning for Variable Action Spacescited by 23
- 2024MAMBA: an Effective World Model Approach for Meta-Reinforcement Learningcited by 13
- 2023When Meta-Learning Meets Online and Continual Learning: A Surveycited by 12
- 2023Recasting Continual Learning as Sequence Modelingcited by 11
- 2023Meta-Reinforcement Learning Based on Self-Supervised Task Representation Learningcited by 10
- 2023Offline Meta Reinforcement Learning with In-Distribution Online Adaptationcited by 10
- 2024Learning to Continually Learn with the Bayesian Principlecited by 8
- 2020Meta-Model-Based Meta-Policy Optimizationcited by 7
50 most cited of 52 All 52 (JSON)
Built on by
Citing papers per year: 2025: 5; 2026: 7.
12 papers.
- 2025Mixture-of-Experts Meets In-Context Reinforcement Learningcited by 7
- 2025Scalable In-Context Q-Learningcited by 5
- 2025How Should We Meta-Learn Reinforcement Learning Algorithms?cited by 5
- 2025Filtering Learning Histories Enhances In-Context Reinforcement Learningcited by 3
- 2025Safe In-Context Reinforcement Learningcited by 2
- 2026Large Language Model-Powered Evolutionary Code Optimization on a Phylogenetic Treecited by 2
- 2026MDForge: Agentic Molecular Dynamics Pipeline Design under Sparse Simulator Feedbackcited by 2
- 2026IntentScore: Intent-Conditioned Action Evaluation for Computer-Use Agentscited by 1
- 2026In-Context Planning with Latent Temporal Abstractionscited by 0
- 2026Latent Q-Barrier Shielding for Safe In-Context Reinforcement Learningcited by 0
- 2026Reinforcement Learning Foundation Models Should Already Be A Thingcited by 0
- 2026Scalable In-Context Reinforcement Learning with Recurrent Algorithm Distillationcited by 0
- 2026Scalable In-Context Reinforcement Learning with Recurrent Algorithm Distillation28 Sep 2026
- 2026Reinforcement Learning Foundation Models Should Already Be A Thing17 Jun 2026
- 2026MDForge: Agentic Molecular Dynamics Pipeline Design under Sparse Simulator Feedback11 Jun 2026
- 2026Latent Q-Barrier Shielding for Safe In-Context Reinforcement Learning24 May 2026
- 2026IntentScore: Intent-Conditioned Action Evaluation for Computer-Use Agents6 Apr 2026
- 2026In-Context Planning with Latent Temporal Abstractions21 Feb 2026
- 2026Large Language Model-Powered Evolutionary Code Optimization on a Phylogenetic Tree20 Jan 2026
- 2025Safe In-Context Reinforcement Learning29 Sep 2025
- 2025How Should We Meta-Learn Reinforcement Learning Algorithms?23 Jul 2025
- 2025Mixture-of-Experts Meets In-Context Reinforcement Learning5 Jun 2025
- 2025Scalable In-Context Q-Learning2 Jun 2025
- 2025Filtering Learning Histories Enhances In-Context Reinforcement Learning21 May 2025