Research map / Reinforcement learning / Visual and representation RL
Towards Large-Scale In-Context Reinforcement Learning by Meta-Training in Randomized Worlds
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 AnyMDP, a procedurally generated tabular MDP task suite with banded transitions and composite rewards, and OmniRL, an in-context RL model meta-trained with decoupled policy distillation and prior-tag step-wise supervision on it.
Problem
Scaling in-context reinforcement learning (ICRL) is limited by a lack of large, low-bias task collections, and by the cost and train/inference trajectory mismatch of meta-training ICRL models.
Approach
Procedurally generate tabular MDPs with banded transition matrices, ascending value functions, and composite rewards so tasks are high-quality and low-bias; meta-train a linear-attention sequence model with decoupled policy distillation (diverse behavior policies, oracle reference labels per step) plus a behavior-policy tag as prior information.
Key results
- OmniRL trained only on AnyMDP reaches 95.3% normalized performance on T_tst(16,5) with 29 episodes cost, versus 92.0% for TQL-UCB at 4.7K episodes and 90.6% for PPO at 9.7K episodes.
- OmniRL reaches 100% on CliffWalking with 65 episodes cost, versus 100% for TQL-UCB at 35 episodes and 95.9% for PPO at 2.7K episodes.
- Increasing the number of meta-training tasks (100 to 128K) increases generalization to unseen tasks but lengthens the steps needed to reach 80% normalized ICL gain.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Meta-train OmniRL on AnyMDP rather than GarnetMDP | OmniRL · AnyMDP test (16, 32, 64, and 128 states) | normalized episodic performance: 47.8% vs 95.3%; 42.0% vs 90.3%; 47.1% vs 91.3%; 32.3% vs 80.2% | Helped |
| Meta-train OmniRL on AnyMDP rather than GarnetMDP | OmniRL · Garnet (16 and 64 states) | normalized episodic performance: 99.0% vs 85.9%; 87.3% vs 80.4% | Mixed |
| Increase training task count from limited coverage to 10K or more (ablation) | ICRL meta-training · unseen AnyMDP tasks | validation loss / generalization | Helped |
| Use fewer than or equal to 10K training tasks (ablation) | ICRL meta-training · unseen AnyMDP tasks | generalization during over-training | Hurt |
| Increase context length for adaptation (ablation) | OmniRL · unseen tasks | asymptotic performance | Helped |
| Add prior policy information to the trajectory (ablation) | OmniRL · AnyMDP online RL, offline RL, and imitation learning | performance | Helped |
| Use step-wise supervision instead of one-step action supervision (ablation) | ICRL meta-training · AnyMDP online RL, offline RL, and imitation learning | performance | Helped |
| Use composite reward sampling (ablation) | AnyMDP task generation · AnyMDP tasks | learning difficulty | Mixed |
| Increase the number of states (ablation) | AnyMDP tasks · AnyMDP | training steps to convergence | Hurt |
| Increase the number of actions (ablation) | AnyMDP tasks · AnyMDP | training steps to convergence | Hurt |
| Use longer sequences for generalization (ablation) | ICRL · unseen tasks | adaptation period | Hurt |
| Evaluate on increasing state-space sizes (ablation) | OmniRL (AnyMDP) · AnyMDP test | normalized episodic performance: 95.3% -> 90.3% -> 91.3% -> 80.2% | Mixed |
Surprising findings
- Generalization to unseen tasks comes at the cost of worse zero-shot and few-shot performance, so the authors argue for judging ICRL by asymptotic performance.
- OmniRL on Garnet-trained variants performs poorly on AnyMDP-style and other environments, showing that Garnet-trained models do not transfer, while AnyMDP-trained models adapt to most unseen tasks.
Limitations
- Restricted to discrete tabular MDPs; extension to continuous state-action spaces is not addressed.
- Continuous observation spaces (e.g., Pendulum) must be manually discretized, and environments are limited to n_s <= 128 and n_a <= 5.
- Generalization can reduce zero-shot and few-shot performance, and reward shaping is important for OmniRL to work.
Experimental setup
Datasets: AnyMDP tasks (procedurally generated) Gymnasium tasks (evaluation) DarkRoom (evaluation) Garnet MDP (comparison)
Benchmarks: AnyMDP test tasks Garnet MDP Gymnasium (CliffWalking, FrozenLake, Discrete-Pendulum, Switch) DarkRoom CQL comparison on Gymnasium tasks
Metrics: normalized episodic return (0% random to 100% oracle) steps cost to 99% best performance episodes cost to 99% best performance validation loss normalized ICL gain
Baselines: TQL-UCB PPO Algorithm Distillation AD-epsilon DPT Conservative Q-Learning OmniRL (w/o a priori)
Builds on: Algorithm Distillation Decision Pre-Training Transformer Garnet MDP Gated Slot Attention RWKV-7
Closest papers by meaning
- XLand-100B: A Large-Scale Multi-Task Dataset for In-Context Reinforcement Learning (ICLR 2025 · similarity 0.91)
- Vintix: Action Model via In-Context Reinforcement Learning (ICML 2025 · similarity 0.90)
- Towards an Information Theoretic Framework of Context-Based Offline Meta-Reinforcement Learning (NeurIPS 2024 · similarity 0.90)
- Behavior-agnostic Task Inference for Robust Offline In-context Reinforcement Learning (ICML 2025 · similarity 0.89)
- AMAGO-2: Breaking the Multi-Task Barrier in Meta-Reinforcement Learning with Transformers (NeurIPS 2024 · similarity 0.89)
- In-Context Decision Transformer: Reinforcement Learning via Hierarchical Chain-of-Thought (ICML 2024 · similarity 0.89)
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.