Research map / Reinforcement learning / Visual and representation RL
Generalization to New Sequential Decision Making Tasks with 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.
Shows that multi-trajectory transformers trained on bursty sequences of expert demonstrations can perform in-context learning of entirely new MiniHack and Procgen tasks from a handful of demonstrations without weight updates.
Problem
In-context learning of new sequential decision making tasks with different states, actions, dynamics and rewards from only a few expert demonstrations, without weight updates, which prior single-trajectory decision transformers fail to achieve.
Approach
Train a causal transformer on multi-trajectory sequences with trajectory burstiness (at least two trajectories from the same level per sequence), so the context contains full expert trajectories relevant to the query, enabling cross-task in-context imitation on unseen tasks.
Key results
- On unseen MiniHack tasks, the multi-trajectory transformer with one expert demonstration achieves episodic return of 0.780 ± 0.07 on Labyrinth versus -0.323 ± 0.05 for the single-trajectory transformer.
- The model's performance on unseen MiniHack and Procgen tasks increases with the number of demonstrations and outperforms BC baselines on Procgen test tasks.
- One-shot performance improves with dataset size up to about 30K levels and with model size up to about 30M parameters, then plateaus.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Train on multi-trajectory sequences instead of single-trajectory sequences (ablation) | Transformer · MiniHack-Labyrinth | episodic return: -0.323 \pm 0.05 -> 0.780 \pm 0.07 | Helped |
| Increase trajectory burstiness (ablation) | Multi-trajectory transformer · Unseen MiniHack tasks | one-shot episodic return | Helped |
| Train with stochastic rather than deterministic environment dynamics (ablation) | Multi-trajectory transformer · Unseen MiniHack tasks | episodic return | Helped |
| Increase dataset size from 2k to 30K levels (ablation) | Multi-trajectory transformer · Unseen MiniHack tasks | one-shot episodic return: 2k to 30K levels | Helped |
| Increase model size up to 30M parameters (ablation) | Multi-trajectory transformer · Unseen MiniHack tasks | one-shot episodic return: 30M parameters | Helped |
| Increase training task diversity from 1 to 8 tasks (ablation) | Multi-trajectory transformer · Unseen MiniHack tasks | one-shot episodic return: 1 to 8 | Helped |
| Remove reward tokens from the context (ablation) | Multi-trajectory transformer · Five unseen MiniHack tasks | episodic return | No clear effect |
| Condition on more expert demonstrations (ablation) | Multi-trajectory transformer · Unseen MiniHack tasks | episodic return | Helped |
Surprising findings
- Reward tokens in the context contribute negligibly to performance, since training uses expert data only.
- Models trained on stochastic dynamics generalize better to unseen tasks, whether test dynamics are deterministic or stochastic.
Limitations
- Performance is poor on tasks with unrecoverable states (unforgiving environments) and on tasks far from training distribution (distributional drift).
- Comparison to prior online meta-RL methods is not fair given the offline few-shot setting, per the authors.
- Evaluation uses human expert demonstrations on 20 levels per test task and rollouts on 100 levels for Procgen, a limited evaluation scope.
Experimental setup
Datasets: MiniHack expert trajectories from E3B policies (12 training tasks) Procgen expert trajectories from PPO policies (12 training tasks) Human expert trajectories for MiniHack evaluation tasks (20 levels each)
Benchmarks: MiniHack Procgen
Metrics: episodic return in-context action accuracy
Baselines: Hashmap BC-0 BC-1 Single-trajectory transformer
Builds on: Decision Transformer Burstiness in in-context learning (Chan et al.) E3B PPO
Closest papers by meaning
- In-Context Decision Transformer: Reinforcement Learning via Hierarchical Chain-of-Thought (ICML 2024 · similarity 0.89)
- When can in-context learning generalize out of task distribution? (ICML 2025 · similarity 0.88)
- In-Context Reinforcement Learning From Suboptimal Historical Data (ICML 2025 · similarity 0.88)
- Context is Environment (ICLR 2024 · similarity 0.88)
- Transformers Can Learn Temporal Difference Methods for In-Context Reinforcement Learning (ICLR 2025 · similarity 0.87)
- Towards Provable Emergence of In-Context Reinforcement Learning (NeurIPS 2025 · 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.