Title: Towards Large-Scale In-Context Reinforcement Learning by Meta-Training in Randomized Worlds
Venue: NeurIPS 2025
Accepted: yes
Keywords: Meta-Learning

In-Context Reinforcement Learning

In-Context Learning
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.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn this paper the authors introduce AnyMDP, a novel generation framework that is able to create a wide range of tabular Markov Decision Processes (with distinct state-action spaces, transition functions and reward functions) to support research in in-context reinforcement learning (ICRL). Furthermore, the authors propose a training methodology for meta-learners that employs AnyMDP. The resulting model, \"OmniRL\" is evaluated in unseen tasks, against online RL approaches (Tabular Q-Learning with UCB action selection, PPO) trained on those tasks. The authors highlight how OmniRL performs on par with these approaches while requiring significantly less interactions/data. Subsequently, the authors highlight how OmniRL outperforms other in-context learning methods. Subsequently, the authors evaluate the role of the number of tasks, training iterations and context lenght in the performance of their method.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**:\n- This paper presents a strong contribution for in-context reinforcement learning: a novel generation framework for tabular MDPs and a novel training methodology for ICRL.\n- Extensive experimental results, showing strong performance and novel insights into the performance of ICRL.\n- Well written and clear paper.\n\n**Weaknesses**:\n- Some missing details and minor typos (See \"Comments\" below)\n\n\n**Comments**:\n- Overall, I thoroughly enjoyed this work. The authors did a good job motivating the need for this benchmark, and how it differs from previous approaches. The writing is clear and the paper is well structured.\n- Some variables are not defined in the text: $\\epsilon_P, \\epsilon_R$ (line 132), $w_t$ (line 4).\n- The captions of the figures could be improved and expanded, in particular Figure 1 and Figure 2. I understand that the authors were constrained by the space available, but I highly incentivize the authors to expand the captions in a revised version. Furthermore, Figure 2 is of low quality, contrasting with the rest of the paper: I would incentivize the authors to make a more complete figure, even if in Appendix, that explains more in detail the training scheme. The legend of Figure 5 is also too small and unable to be read in the printed version).\n- The banded transition kernel and composite reward functions introduced in Section 3.1 need to be discussed (at least in part) in the main paper. I understand the constrains with available space in the paper, but these are fundamental components to understand the sampling process of tasks in the AnyMDP and, as such, should not be fully relegated to Appendix.\n- The authors also use the concept of \"burstiness\" throughout the paper (line 76, line 289, line 290, line 316) yet do not define it. For clarity, it would be better to define it in the text, instead of relegating the definition to a reference.\n- The acronym ICRL is defined in two different ways: \"in-context reinforcement learning\" and \"inverse contextual reinforcement learning\" (line 309). It is better to be consistent.\n- There is an extra \" in line 256 that should be removed.\n\nQUESTIONS:\n1) Can the authors elaborate on what algorithms were used to generate the behavior policies (line 181)? For example, what model-based algorithm was employed? It would benefit the paper if the authors could describe how they were trained in Appendix. Also, are the Q-learning and model-based behavior policies always optimal, or they have different degrees of optimality?\n2) In the sampling process of MDPs how do you make sure that the MDPs are ergodic?\n3) Have the authors attempted to evaluate their approach in out-of-distribution tasks in regards to number of states and number of actions? Especially in tasks with more states and actions.\n4) I found the results with the multi-agent task quite interesting but the discussion in the text was quite unclear. Can the authors elaborate on what does \"configuring observation spaces\" mean (line 236)? Are the policies of the other agents fixed, or learned as well? If fixed, then the task essentially reduces to a single agent task, with the actions of the other agents part of the \"environment\" transition function. A short description of these tasks in Appendix would also clarify other future readers of the paper.\n5) In Section 4.2. the authors evaluate their approach without a priori knowledge, as an ablated version of their method. What about the performance of OmniRL without SS, using for example AD, AD$^\\epsilon$ and DPT, but with priori information?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-27"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper introduces the AnyMDP framework, which utilizes randomized Markov Decision Processes (MDPs) to generate diverse tasks for meta-training in In-Context Reinforcement Learning (ICRL) at a large scale. By incorporating step-wise supervision and prior knowledge induction, the framework improves efficiency and generalization, demonstrating that OmniRL, trained on these randomized tasks, performs well across unseen environments. The work highlights the importance of task diversity and world randomization in scaling ICRL and achieving robust performance across different tasks. The experiments show that larger task sets lead to better generalization but at the cost of increased adaptation time.\n\nSTRENGTHS AND WEAKNESSES:\n### Strengths\n\nThe paper's major strengths include its scalability, with AnyMDP enabling large-scale task generation, and the world randomization approach that reduces structural bias in task sets, leading to strong generalization across new environments. The use of step-wise supervision and prior knowledge induction also enhances the efficiency of meta-training.\n\n### Weaknesses\n\nThe paper’s main weaknesses include its limited applicability to real-world continuous or partially observable environments, as it focuses on discrete tasks. The computational costs of generating and training on large-scale tasks are not fully discussed. While task diversity is emphasized, the performance trade-offs with smaller task sets are not explored, and the longer adaptation times with larger task sets may limit real-time applications. Additionally, the paper doesn't sufficiently address the challenges of long-term generalization and knowledge accumulation in dynamic environments.\n\nQUESTIONS:\nQ1: The paper predominantly focuses on discrete tasks; however, many real-world environments are continuous or partially observable. How do the authors intend to extend the proposed framework to handle such complex, continuous, or partially observable environments? \n\nQ2:  The paper proposes the AnyMDP framework, but its effectiveness has yet to be comprehensively evaluated across a wider range of in-context reinforcement learning methods. It would be valuable for the authors to assess and demonstrate its applicability in this broader context to strengthen the validity of their claims.\n\nQ3: The proposed framework has not been tested across a broad range of in-context reinforcement learning tasks. How do the authors plan to evaluate the generalizability of the framework across diverse in-context RL environments, particularly those with varying complexity and different types of disturbances? Would additional benchmarks or real-world applications be considered to assess its robustness in such settings?\n\nQ4: The paper mentions that perturbations can lead to a decline in model performance, but it does not delve into the specific impacts of different types of perturbations (such as observation noise, action perturbations, and reward errors) on the model. How are these diverse perturbations simulated in AnyMDP, and how is the model's robustness ensured across various environmental changes?\n\nQ5: The computational costs of generating large-scale tasks and training the model are not fully discussed. Can you elaborate on the computational resources required for training on a large number of tasks, especially when task complexity increases?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 5, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper studies in-context learning for reinforcement learning, contributing a task generation framework called AnyMDP that generates tabular MDPs with randomized transition and reward functions. It proposes several techniques to help with meta-training in such randomized worlds.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths\n1. Highly relevant and timely topic\n2. Good positioning in related work\n3. Promising experimental results but with some caveats\n\nWeaknesses\n1. Some discussion on the procedural generation algorithm would be valuable, focusing on why the objective as presented on lines 143/144 was selected (for instance by detailing some of the analysis referred to on line 140).\n2. Reward shaping was applied in standard benchmark domains (see Questions). The selection of domains is unclear.\n3. The exact definition of the 7 types of behavior policies should be clarified, as well as the exact impact of diversity.\n4. On line 132, the two epsilon terms appear to be undefined.\n\nQUESTIONS:\n1. Can the authors shed some more light on why their sampling procedure (as detailed in Appendix B.1) results in high-quality MDPs by providing more intuition on why these are considered to be of high quality?\n\n2. Did the authors consider using the Deep Sea environment, which is a well-known sparse reward benchmark? More generally, how were the presented Gymnasium environments selected?\n\n3. What was the impact of the reward shaping? Specifically, what is the performance of the model is if is not aided by reward shaping?\n\nLIMITATIONS:\nAn obvious limitation that the authors readily acknowledge in Section 5 is the fact that the framework is limited to discrete action and state spaces. Some more discussion about how the work could be extended to continuous domains would be appreciated.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper studies the in-context reinforcement learning (ICRL) problem. It proposes AnyMDP, a class of procedurally generated tabular MDPs. Additionally, it proposed to scale meta-training via step-wise supervision and prior knowledge in the form of action markers corresponding to the behavior policy used. Through this, the paper empirically demonstrates that when the number of AnyMDP tasks used during meta-training is large, the model can generalize to unseen tasks.\n\nSTRENGTHS AND WEAKNESSES:\nDisclaimer: I'm not deeply familiar with ICRL literature, so please take my comments with that in mind.\n\n**Strengths**\n- The paper studies ICRL, which is a problem of current interest to the community.\n- Some empirical results are quite strong (e.g., models trained on AnyMDP tasks generalizing to (discretized) Pendulum from Gymanisum).\n\n**Weaknesses**\n- My major concern with this paper is the writing. As someone who is not particularly familiar with the ICRL literature, the paper felt a bit hard to follow at times. For instance, in-context learning is not formally defined anywhere, unusual notations: \\tau for MDPs, {Causal}_{\\theta} defined nowhere, hard to understand what Fig. 1 is trying to convey, etc.\n- Garnet MDPs is a widely used class of procedurally generated MDPs in the RL community [1, 2, 3]. The paper neither mentions nor compares AnyMDPs with Garnet MDPs. \n- While the paper claims many results on unseen tasks (e.g., Section 4.3), all of the test tasks considered in those setups are in-distribution, which is not particularly interesting. The results showcasing generalization on Gymnasium tasks (out of distribution) are interesting.\n- (Minor) At times, the paper would reference a figure in the main results but the figure would be in the appendices (e.g., Fig. 15, line 6) . This is not a good writing habit.\n\nOverall, the paper studies an interesting problem and provides some good results. However, the writing requires a lot more polish (along with proper comparison with prior work on procedurally generating MDPs) to be apt for publication at NeurIPS.\n\n[1] Bhatnagar et al. Natural Actor–Critic Algorithms, Automatica (Journal of IFAC) 2009.\n\n[2] Castro et al. MICo: Improved representations via sampling-based state similarity for Markov decision processes, NeurIPS 2021.\n\n[3] Castro et al. A Kernel Perspective on Behavioural Metrics for Markov Decision Processes, TMLR 2023.\n\nQUESTIONS:\n1. Section 4.2 demonstrates the effectiveness of incorporating prior information (action markers). Do the AD and DPT baselines have access to this information too? How realistic is the access to this prior information? Also, do the authors have any intuition as to why it helps significantly?\n2. The unseen tasks in Section 4.3 are basically from the AnyMDP class, and hence, in-distribution. Isn't generalization in this setting sort-of guaranteed?\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 3, "CLARITY": 1, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-02"}

comments: SUMMARY:
* Proposes to procedurally generate tabular MDPs in order to study various properties of ICRL
* Results demonstrate that with a sufficiently large scale of tasks, the proposed model can generalize to tasks unseen in training set

STRENGTHS AND WEAKNESSES:
Strengths
* The paper is well-written and clearly presented; in fact, exceptionally so
* The details and analysis in the main body are thorough but not overwhelming
* OmniRL method produces very strong performance against baselines
* Useful insights---especially that "generalization of ICRL can be at odds with its zero-shot or few-shot performance"

Weaknesses
* Only applicable to discrete state and action spaces (acknowledged as limitation)
* Procedural generation is difficult to scale to more complex domains; but I consider the toyishness of this work a strength, as it enabled in-depth analysis (hard to do in non-tabular settings, where PCG is not possible); and I personally believe it is likely the findings would be similar for more complex domains (if such thorough experiments could be run); but see related questions below

QUESTIONS:
* Which findings do the authors believe would(n't) scale to more complex domains? 
* How might future work approach the challenge of broadening the scope of the efforts in this work to continuous settings, and ones where procedural generation may not be possible?

LIMITATIONS:
Yes
TITLE: Official Review
CLARITY: 4
IMPACT: 3
ORIGINALITY: 4
DATE: 2025-07-07