Research map / Reinforcement learning / Visual and representation RL
Distilling Reinforcement Learning Algorithms for In-Context Model-Based Planning
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 DICP, an in-context model-based RL framework where a Transformer jointly learns an in-context dynamics model (via a meta-model loss) alongside imitated policy and plans with beam/greedy search at test time.
Problem
In-context RL via algorithm distillation inherits the suboptimal, gradual-update behavior of the source RL algorithm, limiting sample efficiency on unseen tasks.
Approach
Add a dynamics/meta-model loss (reward, next observation, return-to-go prediction) to existing in-context RL architectures (AD, DPT, IDT) so the same sequence model learns environment dynamics in-context, then use MPC-style beam or greedy search over predicted outcomes to select actions, deviating from the source algorithm's suboptimal actions.
Key results
- On Meta-World ML1 (5 tasks), DICP-IDT reaches 100 on Reach, Push, Pick-Place, Sweep and 87 on Peg-Insert-Side with 200K steps, versus IDT at 100/94/95/96/76.
- On ML10, DICP-IDT achieves 46.9% success vs IDT 36.7% with 500K steps, exceeding RL^2 (35.8%) trained with 350M steps.
- On ML1 at 5K steps, DICP-IDT gets 87 on Sweep and 45 on Peg-Insert-Side versus IDT 71 and 41.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add model-based planning to the trained dynamics model (ablation) | DICP-AD on Darkroom and DICP-IDT on Reach-v2 · Darkroom and Reach-v2 | performance | Helped |
| Increase beam size or sample size (ablation) | DICP · Darkroom and Reach-v2 | performance | Helped |
| Increase beam size or sample size beyond 10 (ablation) | DICP · Darkroom and Reach-v2 | performance | No clear effect |
| Increase context length (ablation) | DICP-IDT · Reach-v2 meta-test tasks | dynamics prediction mean squared error | Helped |
| Use SAC instead of PPO as the source algorithm (ablation) | DICP-IDT compared with IDT · Meta-World ML1 tasks | success rate (%): 85 vs 85; 65 -> 69; 17 -> 20; 4 -> 3; 1 -> 3 | Mixed |
| Use immediate rewards alone instead of immediate reward plus return-to-go for planning (ablation) | DICP · Darkroom variants and Meta-World | performance | Helped |
| Add in-context dynamics learning and planning | IDT · Meta-World ML10 | success rate (%): 36.7 -> 46.9 | Helped |
| Add in-context dynamics learning and planning | IDT · Meta-World ML1 mean | success rate (%): 75 -> 80 | Helped |
Surprising findings
- The best model-free in-context method (AD, DPT, IDT) varies across environments, and DICP-AD underperforms on several ML1 tasks (e.g., Pick-Out-Of-Hole 2 vs AD 36), indicating planning does not uniformly help.
Limitations
- Additional computational cost at action selection compared to model-free methods.
- Relies on learning-history datasets from a source algorithm and may not adapt to changing dynamics.
- Some baselines results are taken from original papers rather than reproduced.
Experimental setup
Datasets: Darkroom Dark Key-to-Door Darkroom-Permuted Meta-World ML1 Meta-World ML10
Benchmarks: Darkroom Dark Key-to-Door Darkroom-Permuted Meta-World ML1 Meta-World ML10
Metrics: success rate mean episode reward (learning curves) test loss (MSE)
Baselines: AD DPT IDT RL^2 MAML PEARL MACAW FOCAL BOReL MoSS IDAQ MuZero
Builds on: Algorithm Distillation Decision-Pretrained Transformer In-context Decision Transformer Model Predictive Control
Closest papers by meaning
- Transformers Provably Implement In-Context Reinforcement Learning with Policy Improvement (NeurIPS 2026 · similarity 0.89)
- In-Context Decision Transformer: Reinforcement Learning via Hierarchical Chain-of-Thought (ICML 2024 · similarity 0.89)
- Vintix: Action Model via In-Context Reinforcement Learning (ICML 2025 · similarity 0.89)
- In-Context Reinforcement Learning From Suboptimal Historical Data (ICML 2025 · similarity 0.89)
- Transformers Can Learn Temporal Difference Methods for In-Context Reinforcement Learning (ICLR 2025 · similarity 0.88)
- In-context Exploration-Exploitation for Reinforcement Learning (ICLR 2024 · 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.