mlsanityData to 7 Oct 2026

Distilling Reinforcement Learning Algorithms for In-Context Model-Based Planning

arXiv 2502.19009 ↗First posted 26 Feb 2025cs.LG cs.AIcited by 12

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

  1. 2017Attention Is All You Needcited by 59,932
  2. 2020Language Models are Few-Shot Learnerscited by 32,199
  3. 2020An Image is Worth 16x16 Words: Transformers for Image Recognition at Scalecited by 29,915
  4. 2017Proximal Policy Optimization Algorithmscited by 14,441
  5. 2022Chain-of-Thought Prompting Elicits Reasoning in Large Language Modelscited by 12,302
  6. 2017Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networkscited by 4,861
  7. 2018Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actorcited by 4,635
  8. 2023Sparks of Artificial General Intelligence: Early experiments with GPT-4cited by 2,579
  9. 2015Trust Region Policy Optimizationcited by 2,503
  10. 2020Offline Reinforcement Learning: Tutorial, Review, and Perspectives on Open Problemscited by 1,873
  11. 2021Decision Transformer: Reinforcement Learning via Sequence Modelingcited by 1,528
  12. 2019Mastering Atari, Go, Chess and Shogi by Planning with a Learned Modelcited by 1,327
38 more
  1. 2023Mastering Diverse Domains through World Modelscited by 1,193
  2. 2019Meta-World: A Benchmark and Evaluation for Multi-Task and Meta Reinforcement Learningcited by 1,168
  3. 2018On First-Order Meta-Learning Algorithmscited by 1,003
  4. 2022A Generalist Agentcited by 802
  5. 2018Recurrent World Models Facilitate Policy Evolutioncited by 761
  6. 2019When to Trust Your Model: Model-Based Policy Optimizationcited by 753
  7. 2021Offline Reinforcement Learning as One Big Sequence Modeling Problemcited by 636
  8. 2024TinyLlama: An Open-Source Small Language Modelcited by 604
  9. 2019Model-Based Reinforcement Learning for Ataricited by 495
  10. 2016RL2: Fast Reinforcement Learning via Slow Reinforcement Learningcited by 490
  11. 2019Efficient Off-Policy Meta-Reinforcement Learning via Probabilistic Context Variablescited by 372
  12. 2016Learning to reinforcement learncited by 369
  13. 2017Neural Network Dynamics for Model-Based Deep Reinforcement Learning with Model-Free Fine-Tuningcited by 288
  14. 2018Learning to Adapt in Dynamic, Real-World Environments Through Meta-Reinforcement Learningcited by 270
  15. 2021Mastering Atari Games with Limited Datacited by 235
  16. 2022Multi-Game Decision Transformerscited by 212
  17. 2019VariBAD: A Very Good Method for Bayes-Adaptive Deep RL via Meta-Learningcited by 210
  18. 2022In-context Reinforcement Learning with Algorithm Distillationcited by 186
  19. 2023Supervised Pretraining Can Learn In-Context Reinforcement Learningcited by 137
  20. 2018Model-Based Reinforcement Learning via Meta-Policy Optimizationcited by 99
  21. 2020Offline Meta-Reinforcement Learning with Advantage Weightingcited by 92
  22. 2021Online and Offline Reinforcement Learning by Planning with a Learned Modelcited by 90
  23. 2018Meta Reinforcement Learning with Latent Variable Gaussian Processescited by 44
  24. 2022Robust Task Representations for Offline Meta-Reinforcement Learning via Contrastive Learningcited by 33
  25. 2023Emergent Agentic Transformer from Chain of Hindsight Experiencecited by 33
  26. 2020Model-based Adversarial Meta-Reinforcement Learningcited by 32
  27. 2021Procedural Generalization by Planning with Self-Supervised World Modelscited by 30
  28. 2020Exploration in Approximate Hyper-State Space for Meta Reinforcement Learningcited by 28
  29. 2023Emergence of In-Context Reinforcement Learning from Noise Distillationcited by 26
  30. 2024In-Context Decision Transformer: Reinforcement Learning via Hierarchical Chain-of-Thoughtcited by 25
  31. 2023In-Context Reinforcement Learning for Variable Action Spacescited by 23
  32. 2024MAMBA: an Effective World Model Approach for Meta-Reinforcement Learningcited by 13
  33. 2023When Meta-Learning Meets Online and Continual Learning: A Surveycited by 12
  34. 2023Recasting Continual Learning as Sequence Modelingcited by 11
  35. 2023Meta-Reinforcement Learning Based on Self-Supervised Task Representation Learningcited by 10
  36. 2023Offline Meta Reinforcement Learning with In-Distribution Online Adaptationcited by 10
  37. 2024Learning to Continually Learn with the Bayesian Principlecited by 8
  38. 2020Meta-Model-Based Meta-Policy Optimizationcited by 7

50 most cited of 52 All 52 (JSON)

Built on by

Cited by year (2026 to 7 Oct)

Citing papers per year: 2025: 5; 2026: 7.

12 papers.