Research map / Reinforcement learning / Visual and representation RL
In-context Exploration-Exploitation for Reinforcement 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 ICEE, a Transformer that learns in-context exploration-exploitation from offline multi-episode sequences by training on an unbiased action objective and cross-episode return-to-go, matching GP-EI on discrete Bayesian optimization and solving grid-world RL tasks within tens of episodes.
Problem
In-context policy learning for RL is computationally expensive because it requires gathering long learning trajectories and training large Transformers; existing in-context RL needs hundreds of episodes to solve new tasks.
Approach
Train a sequence model with pure supervised learning on multi-episode sequences from cheap non-learning data-collection policies, using an importance-weighted objective that removes data-policy bias and a cross-episode return-to-go (indicator of beating all previous episodes) so exploration-exploitation emerges at inference from epistemic uncertainty in predictions.
Key results
- On 16 discrete BO benchmark functions, ICEE matches GP-based EI search efficiency per function evaluation and is orders of magnitude faster in elapsed time (all on a single A100).
- On Dark Room at the 10th episode, ICEE achieves average return 8.15 versus 3.74 for AD and 0.05 for AD-sorted, reaching 13.61 by the 50th episode.
- On Dark Key-To-Door, ICEE reaches 1.50 average return at the 10th episode versus 1.02 for AD, supporting that new RL tasks can be solved within tens of episodes.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use ICEE instead of AD-sorted | in-context policy learning · Dark Room (50th Episode) | average return: 12.48 (1.37) -> 13.61 (0.86) | Helped |
| Use ICEE instead of AD-sorted | in-context policy learning · Dark Room (Hard) (50th Episode) | average return: 0.75 (0.08) -> 0.79 (0.08) | Helped |
| Use ICEE instead of AD-sorted | in-context policy learning · Dark Key-To-Door (20th Episode) | average return: 1.77 (0.09) -> 1.84 (0.08) | Helped |
| Apply the unbiased action-bias correction objective (ablation) | ICEE · discrete Bayesian optimization | search efficiency | Helped |
| Apply the unbiased action-bias correction objective (ablation) | ICEE · Dark Room (Biased) | average return | Helped |
| Apply the unbiased action-bias correction objective (ablation) | ICEE · grid-world RL without action-collection bias | average return | No clear effect |
| Use ICEE instead of MGDT | in-context policy learning · Dark Room (50th Episode) | average return: 1.86 (0.93) -> 13.61 (0.86) | Helped |
| Use ICEE instead of MGDT | in-context policy learning · Dark Room (Hard) (50th Episode) | average return: 0.09 (0.06) -> 0.79 (0.08) | Helped |
| Use ICEE instead of MGDT | in-context policy learning · Dark Key-To-Door (20th Episode) | average return: 0.34 (0.11) -> 1.84 (0.08) | Helped |
Surprising findings
- AD-sorted reaches higher returns than ICEE at the 50th episode on Dark Room (12.48 vs 13.61 nearly matched, and on Dark Room Hard 0.75 vs 0.79), though it is much slower early on.
- The original Algorithm Distillation performs slightly below the data collection policy, since it learns to imitate that policy.
Limitations
- Experiments are limited to small discrete grid worlds and 2D discrete BO problems.
- ICEE-biased variant is not compared on the standard RL tasks because the data collection policy shows no clear bias there.
- The unbiased objective holds only when the data collection policy's probabilities are known and the action space is discrete; the epistemic-uncertainty argument ignores hyper-parameter uncertainty.
Experimental setup
Datasets: Synthetic GP-sampled BO functions Dark Room grid world (generated) Dark Key-to-Door grid world (generated) 2D benchmark functions (Branin, Beale, etc.)
Benchmarks: 16 2D benchmark functions for discrete Bayesian optimization Dark Room Dark Room (Hard) Dark Room (Biased) Dark Key-To-Door
Metrics: distance between current best estimate and true minimum average return over 100 sampled games action entropy wall-clock time
Baselines: Gaussian process expected improvement random search Algorithm Distillation AD-sorted Multi-game Decision Transformer data collection policy (source) ICEE-biased
Builds on: Decision Transformer Algorithm Distillation Multi-game Decision Transformer nanoGPT
Closest papers by meaning
- In-Context Decision Transformer: Reinforcement Learning via Hierarchical Chain-of-Thought (ICML 2024 · similarity 0.90)
- XLand-100B: A Large-Scale Multi-Task Dataset for In-Context Reinforcement Learning (ICLR 2025 · similarity 0.89)
- In-Context Reinforcement Learning From Suboptimal Historical Data (ICML 2025 · similarity 0.88)
- Towards Large-Scale In-Context Reinforcement Learning by Meta-Training in Randomized Worlds (NeurIPS 2025 · similarity 0.88)
- Distilling Reinforcement Learning Algorithms for In-Context Model-Based Planning (ICLR 2025 · similarity 0.88)
- OCEAN-MBRL: Offline Conservative Exploration for Model-Based Offline Reinforcement Learning (AAAI 2024 · 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.