Towards Large-Scale In-Context Reinforcement Learning by Meta-Training in Randomized Worlds

TL;DR

This work tackles the scalability challenge of In-Context Reinforcement Learning (ICRL) by introducing AnyMDP, a procedurally generated, low-bias MDP suite with ergodic and banded-transition properties, plus a Composite Reward design to ensure ascending value. Building on this, the authors propose OmniRL, a scalable meta-training framework using Decoupled Policy Distillation and prior-information induction to learn long-context policies from massive context sequences (up to $512{,}000$ steps per sequence, $6{,}000{,}000{,}000$ total steps). Empirical results show strong generalization to unseen tasks, improved offline/online learning performance, and emergent general-purpose ICRL as task diversity grows, while highlighting trade-offs between generalization and adaptation time. The findings emphasize the importance of large, diverse task distributions and long-context modeling for robust, scalable ICRL, and point to future work extending these ideas to continuous spaces and broader evaluation frameworks.

Abstract

In-Context Reinforcement Learning (ICRL) enables agents to learn automatically and on-the-fly from their interactive experiences. However, a major challenge in scaling up ICRL is the lack of scalable task collections. To address this, we propose the procedurally generated tabular Markov Decision Processes, named AnyMDP. Through a carefully designed randomization process, AnyMDP is capable of generating high-quality tasks on a large scale while maintaining relatively low structural biases. To facilitate efficient meta-training at scale, we further introduce decoupled policy distillation and induce prior information in the ICRL framework. Our results demonstrate that, with a sufficiently large scale of AnyMDP tasks, the proposed model can generalize to tasks that were not considered in the training set through versatile in-context learning paradigms. The scalable task set provided by AnyMDP also enables a more thorough empirical investigation of the relationship between data distribution and ICRL performance. We further show that the generalization of ICRL potentially comes at the cost of increased task diversity and longer adaptation periods. This finding carries critical implications for scaling robust ICRL capabilities, highlighting the necessity of diverse and extensive task design, and prioritizing asymptotic performance over few-shot adaptation.

Figures (24)

• Figure 1: An ablation study comparing AnyMDP tasks with Garnet MDP and AnyMDP without composite reward demonstrates that the procedural generation algorithm of AnyMDP produces tasks of higher learning difficulty.
• Figure 2: Comparison of the stationary distributions (SDs) of the oracle policy and the uniform random policy across four classes of environments: AnyMDP - $\tau(16/64,5)$, and Garnet MDP-$(16/64,5,2/4)$ demonstrates that AnyMDP exhibits uniquely exponentially-decaying SDs. The results are averaged over 64 randomly sampled tasks, with the SDs for each task re-ranked in descending order.
• Figure 3: Comparison of the learning pipelines and data formulation of different ICRL methods (AD, AD$^{\varepsilon}$,DPT) and Decoupled Policy Distillation (DPD).
• Figure 4: OmniRL model structure and training losses.
• Figure 5: Left: Normalized episodic return of TQL-UCB, PPO and OmniRL on 64 AnyMDP evaluation tasks (mean $\pm95\%$ CI). Hyper-parameters of PPO and TQL-UCB were separately tuned by grid-search inside every task family to maximize final-episode performance. Right: Per-step validation loss of OmniRL on a held-out static dataset, mirroring the online-RL trend.

Theorems & Definitions (1)

• Theorem 1