Towards Large-Scale In-Context Reinforcement Learning by Meta-Training in Randomized Worlds
Authors: Fan Wang
Pengtao Shao
Yiming Zhang
Bo Yu
Shaoshan Liu
Ning Ding
Yang Cao
Yu Kang
Haifeng Wang
Venue: NeurIPS 2025 poster ()
Keywords: Meta-Learning
In-Context Reinforcement Learning
In-Context Learning
Primary area: reinforcement_learning
TL;DR: The paper presents AnyMDP, a framework for procedurally generating diverse tasks to enhance In-Context Reinforcement Learning (ICRL) scalability, and explores the trade-off between generalization and adaptation efficiency.

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.

Towards Large-Scale In-Context Reinforcement Learning by Meta-Training in Randomized Worlds
+Fan Wang! 1.2 *Pengtao Shao? Yiming Zhang', Bo Yu², +Shaoshan Liu?,
Ning Ding'. Yang Caol. Yu Kangl and Haifeng Wang*
University of Science and Technology of China, Heifei, China Shenzhen Institute of Artificial Intelligence and Robotics for Society, Shenzhen, China Anhui Province Key Laboratory of Intelligent Low-Carbon Information Technology and Equipment "Baidu Inc, Beijing, China
Abstract
In-Context Reinforcement Learning (ICRL) enables agents learn automatically and on-the-fly from their interactive experiences. However, a major challenge in scaling up ICRL the lack of scalable task collections. To address this, we propose the procedurally generated tabular Markov Decision Processes, named AnyMDp* Through carefully designed randomization process, AnyMDP is capable gener- ating high-quality tasks on large scale while maintaining relatively low structural biases. To facilitate efficient meta- taaiatirainnn scale, we further introduce decoupled policy distillation and induce prior information in the ICRL framework Our results demonstrate that, with 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 ship 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.
thorough empirical investigation of the relation-
Introduction
In-Context Learning (ICL) [ has emerged pivotal paradigm for large pre-trained models 2 5), enabling adaptation to novel tasks without parameter adjustments. Unlike gradient- based in- weight learning (IWL), which modifies parameter weights through optimization. ICL facilitates skill acquisition through natural interaction. This is achieved by unifying versatile learning paradigms by sequence of interactions the context, including supervised learning 6 7 imitation learning I0] and reinforcement learning [1HK]. etc. By leveraging those contexts, ICL. eliminates reliance on manually engineered objective functions and the labor-intensive optimization infrastructure required by traditional IWL frameworks, providing versatile, self-directed manners of learning, ICL further aligns with black-box meta- letanlea principles [1] where task-agnostic adaptation occurs through computation and memorization rather than explicit parameter tuning. It can be regarded unifying
Equal Contribution
Conessonning to: fanwang wxw.lamppxx@mmmlccoo.cmo www.aiii..oo.hhhhhaaaaluuu......
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

framework that bridges the gap between conventional meta-learning approaches and contemporary large-scale language model capabilities.
The historical conception of ICL has been predominantly associated with supervised few-shot learning paradigms. Recently, the frameworks for learning from dynamic experiences have achieved growing success, where the contextual information can encompass observations, actions, feedback, and reasoning, Among these emerging paradigms, In-Context Reinforcement Learning (ICRL) stands out as a critical approach. Unlike traditional ICL, which arises from self-supervised pre-training, ICRL typically requires explicit training via reinforcement learning (RL) algorithms [12], introducing substantially higher computational complexity and training costs for large-scale applications. While recent supervised learning approaches have shown promise enhancing ICRL efficiency [16]
current implementations remain constrained by limited context lengths (typically <IK tokens) and narrow task scales (often <IK distinct tasks). These limitations significantly restrict ICRL's capacity for broad task generalization, highlighting an urgent need for scalable architectures that can reconcile contextual flexibility with computational efficiency.
Additionally, existing benchmarks for ICRL predominantly fall into two categories: 1) oversimplified environments like multi-armed bandit tasks, or 2) procedurally generated seed tasks such as maze navigation [7711 or arcade-style games [20]. While the former lacks practical relevance due to their constrained problem spaces, the latter suffers from superficial parameter randomization diversifying only peripheral elements like visual textures or sub-objectives while preserving core environmental biases (e.g.. fixed transition dynamics and reward functions). This design paradigm confines ICRL agents to narrow problem distributions, yielding systems that demonstrate limited generalization scope. Consequently, scaling task complexity beyond a few hundred variants encounters diminishing returns, as models inherit fundamental limitations from their training environments' structural biases.
To advance the scalability and generalization capabilities of ICRL, this work presents two interrelated contributions: a novel benchmarking task set and a scalable ICRL framework. First, we introduce AnyMDP, a scalable task generation environment where Markov Decision Processes (MDPs) are systematically designed with fully randomized transition dynamics and reward functions. To bal- ance environmental richness with learner challenge, we propose the procedural generation process remarked with banded transition matrices. This design minimizes structural biases while preserving problem complexity, creating a benchmark that demands genuine contextual adaptation rather than exploiting environmental regularities. Second, we propose Decoupled Policy Distillation (DPD) to en- hance ICRL training efficiency. By disentangling reference policy from behavior policy, and inducing prior knowledge into the context, we achieve unprecedented training scales and ICRL versatility. Our implementation, OmniRL, was trained exclusively on AnyMDP tasks using context sequences of up to 512K steps per sequence and 6 billion steps overall. Empirical validation demonstrates OmniRL's capability of many-shot generalization to entirely novel tasks, confirming our models generalization scope and versatile learning ability.
Moreover, leveraging the scalable AnyMDP task set, we conduct ablations that reveal critical insights into ICRL's dependence on task diversity and long-context modeling. First, the boundary between IWL. and ICL is fundamentally determined by training-task coverage: sufficiently diverse distributions (at least 10K unique tasks in our case) elicit emergent task learning capabilities, whereas limited coverage reduces systems to task recognition regimes. Second, long-context modeling proves indispensable for interpreting complex task specifications and achieving broad generalization. These findings establish a new paradigm for ICRL research in which scalable benchmarks and context-aware training jointly yield systems that approach human-like adaptability in dynamic decision-making domains.
2
Related Work
2.1 Emergence of In-Context Learning
pertains to a category of approaches that prioritize the acquisition of generalizable adaptation skills across a spectrum of tasks. Large models pre-trained on massive, uncurated datasets naturally give rise to In-Context Learning (ICL), phenomenon analogous to model-based meta-learning [1] g 23. Investigations reveal that ICL ability is tightly linked to pre-training data distribution, with factors spanning burstiness and intra-sequence correlation [28 29 25 to overall diversity [30]. Complementary analyses show
Meta-leaming, also known as learning to learn 21
22 12

---

that computation- based ICL can emulate a rich spectrum of learning behaviors, including gradient descent [3], Bayesian inference [39], and broader learning paradigms [T3] Those findings suggest its potential as general-purpon purpose learning machine (3437] rather than a mere mechanism
for instruction following or few-shot adaptation. Building on the recent progress in ICI. and the principles of meta-learning, we adopt the term meta-training to denote the training procedure that explicitly equips the model with the capacity to learn in context, thereby distinguishing it from pre-training, which yields only zero-shot skills and incentivizes ICL as incidental by-product of exposure to uncurated data.
2.2 In-Context Reinforcement Learning
In-context reinforcement learning (ICRL) encompasses algorithms that dynamically adapt decision- making tasks by synthesizing self- gelf-generated trajectories and incorporating external feedback in the context [2[3] Unlike one-off ICL adaptation, ICRL stresses continual policy improvement driven
by accumulated experience, echoing the "third system" forged through ongoing interaction and experience accumulation, as articulated by Barabasi et al. [38] ICRL typically employs recurrent
12 [9] and attention-baset [13 40] neural structures that are capable of encoding the interactive
histories in the inner loop. The meta-training process that searches the parameters for learning
functionality is called outer loop. Common choices for the outer-loop optimizer for ICRL include
reinforcement learning [12 41 40), evolutionary strategies [42 17 and supervised learning [13]
16.. While supervised learning generally achieves higher sample efficiency compared to RL and
evolutionary strategies, it often suffers from the bottleneck of distribution shift [43 44 Additionally,
the absence of oracle policy can become critical bottleneck for supervised learning; consequently, alternatives replace the expert with an RL coach for data synthesis [8 [5]. A further obstacle to ICRL research is the lack of large-scale task suites, which hampers the emergence of ICRL in real-world decision-making problems and motivates the creation of procedurally generated tasks at scale.
2.3
Procedurally Generated Tasks
Commonly employed poocedural-generatinn techniques include randomizing a subset of domain parameters to create variant tasks by randomizing the rewards or targets while keeping the transitions
fixed [13
16
40.2
22 46, 47], randomizing the dynamics while keeping the targets unchanged [22
17
and randomizing the observations and labels without altering the underlying transitions and rewards [34 19 (5). These techniques is frequently applied to enhance the robustness of decision and policy model for sim-to-real transfer (51)53), namely domain randomization (DR). Although DR creates a class of tasks with a certain variety, is restricted by the original task setting and basically forms close and finite task set. Recently, generating open-ended tasks with enhanced diversity has emerged as a focal point in research, with a significant emphasis on game- (66] Such approaches are increasingly recognized as vital inductive biases for fostering adaptability. While procedurally generated discrete Markov Decision Processes (MDPs) provide promising avenue for minimizing structural bias, owing their configuration via constrained set of hyperparameters, naive sampling of these MDPs often results in a concentration of trivial tasks, which lack meaningful complexity (57]58) In this work, we address this limitation by strategically incorporating critical features during the procedural generation of discrete MDPs.
20
48
49
18
based frameworks [54] 52 36
55
54
52
3 AnyMDP: Procedural Generation of High-Quality MDP Tasks
3.1 Problem Setting and Definitions
We denote a task of Markov Decision Processes (MDPs) with - (S.A,P,R,, R., SE. TE. 7) where S denotes the state space, A the action space, P S A S the state-transition probability distribution, R. : S x A x S + R the reward function, Sn the set of initial states, and Sg the set of terminal states, 1 g the maximum steps in an episode, and 7 the discount factor. We consider only fully visible discrete MDPs with S = [I.....a... A = [1,...na).. denoted by T. We define average state transition (s'(ss = Ea-r)(a)=))- which the transition under uniform policy v: T- T'ass) = TAT For absorbing states & E SE, the transition was connected back to Sp. We then formally define the AnyMDP task collection T T(n,, Ra, 7), e, K. b_) as MDPs satisfying the following conditions:
3

---

Ergodic:
average state transition P. has unique positive stationary distribution.
Banded Transition Matrix: there exists a ranking of states: 5 Sn,, such that: P. satisfies the following condition:
P((8)ls))= Vj <i <i- > b+
Σ P((s)ls)) > if b_, Σ P((sylss) > >. -1 JEL
(I)
Ascending Value Function: given the aforementioned ranking of states, the value function v under the optimal policy satisfies the following condition:
V (sn.)) max #ESp
(*(s) -
(2)
Notice that T(n n.) represents the widely used multi-armed bandits benchmark for ICRL [2 As the number of states n, increases, tasks demand progressively more sophisticated reasoning over delayed rewards and complex state spaces. This transition evolves the problem from simple multi-armed bandit framework to a full reinforcement learning paradigm, where long-term planning and environmental interaction become critical. Furthermore, since the ground truth dynamics P and reward function are known within the simulation environment, this setup enables straightforward computation of the oracle solution through value iteration (59]
16
3.2 Motivations of AnyMDP
AnyMDP is motivated by the following visions: (1) Existing procedurally generated MDP benchmarks like PROCON [57] and Garnet [58] overlook terminal states (Sş) which is critical feature of real- world environments where agents enter terminal states (e.g., task completion or failure). Therefore, AnyMDP allows terminal states to match many real tasks. (2) Through banded transition matrix and ascending value function, AnyMDP has high-valued states that are exponentially less possible to reach by random exploration. In real-world tasks, especially long-horizon ones, the probability of reaching the goal by naïve random choices decays exponentially with task length. Consequently, solely relying random exploration can be extremely inefficient in solving those tasks; instead, continuous exploration-exploiitaion trade-off is required to approach the optimum gradually. Inspired by this necessity, impose the Markov Decision Process the constraint that states with higher values are reached with lower probability, thereby guiding the learner toward systematic, intelligent exploration rather than mere random search.
By defining the stationary distribution Pt with P. Pt. the upper and lower bound of the chances of arriving at any states is ensured by the following theorem:
Theorem 1. Given the condition of Equation [T) for > b. where b = max(b_, bo) + b4, there exists a value Pe has a unique positive stationary distribution (Ergodic). Specifically, for the state 8j, the ps(s)) satisfies the following bounds:
8 1/(b++++// 1). If > o,the sampling algorithm ensures that the transition
ciett> < C2(1 - -8)--
where Ci and C2 are positive constants.
Theorem guarantees that the probability of staying at sy by random choice of actions decreases least exponentially with respect to j. while also being bounded above by another exponentially decaying, yet non-negligible, probability to ensure ergodicity. Using banded transition matrix, we further equip the MDP with Composite Reward (CR) that decomposes reward generation into independent components, ensuring an ascending value function while preserving low structural bias when sampling. We have relegated the proof of Theorem and the details the sampling procedures to the Appendix B for brevity.
3.3 Comparison with Other Procedural MDPs and Empirical Validation
The proposed sampler guarantees high-quality MDPs three aspects: ergodicity, low structural bias, and high learning difficulty. To empirically verify the high learning difficulty and to validate

---

REgEE BASI
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.
Random Policy
Oracle Policy
Random Policy
Oracle Policy
Probability
Probability
043
090
AnyMDP(16,5) 2 GGrret(1665AA)
AnyMDP(64,55
Garmet(64.5.22 Garnet(64,5,4)
Figure 2: Comparison the stationary distributions (SDs) the oracle policy and the uniform random policy across four classes of environments: AnyMDP r16664, 5), and Garnet MDP- (16/64, 5, 2/4) demonstrates that AnyMDP exhibits uniquely exponentially-decayin SDs. The results averaged 64 randomly sampled tasks, with the SDs for each task re-ranked in descending order.
Theorem 1 we run two groups of experiments. Figure compares the performance of two well- known discrete MDP learning methods: Tabular Q-Learnig [60] with Upper Confidence Bound (TQL-UCB) [6]] and Proximal Policy Optimization (PPO) [62], on tasks sampled from AnyMDP Garnet uses parameters 11,, na, b, a, t to shape generated MDPs. Here, we set = 0.1, 0, b (2, 4), and keeps n., and equal AnyMDP for comparison. To isolate the impact of composite reward (CR) Equation (TO) sampling, we also include baseline without the composite reward sampling technique (AnyMDP w/o CR). The results show that AnyMDP tasks pose greater challenges for both RL methods. Figure 2 further illustrates the stationary distribution (SD) of states by tasks sampled from AnyMDP and Garnet MDPs. It demonstrates exponential decay in SD of Mny empirically validating Theorem D under the random policy.
and Garnet MDP
while Garnet MDPs exhibit flatter SDs, especially
4 The Scalable ICRL Framework of OmniRL
Tasks
EPI
PPI
RP1
B22
Reference Policy
Behavior Policy
AD
DPT
Decoupled Policy Distillation
Figure 3: Comparison of the learning pipelines and data formulation of different ICRL methods (AD, AD DPT) and Decoupled Policy Distillation (DPD).
Decoupled Policy Distillation (DPD): Meta-trann Meta-training for ICRL using RL or Evolution Strategies faces significant costs, including cumbersome infrastructure requirements and high computational costs. Recent advances supervised learning-based meta-training methods- such as Algorithm Distillation (AD) 13 AD [5]. and Decision Pre-Training Transformers (DPT) [66] have shown promise for scalable ICRL. meta-training. However, a critical challenge arises during inference: contextual trajectories are generated by the model itself, creating unavoidable gap between
5

---

training-time time and inference-timet time trajectories. This discrepancy can lead to catastrophic failures during deployment. While leveraging diverse behavior trajectories has been shown to mitigate this issue [16], it introduces a new challenge of maintaining training efficiency. To address these limitations, we propose Decoupled Policy Distillation (DPD), a framework inspired by DPT and Our approach hinges on two decoupled policies: The behavior policy refers to the policy executed to generate trajectories during training. The reference policy refers to the target policy to be imitated, which remains decoupled from direct execution. Decoupling the behavior and reference policies enables the introduction of diversity into the behavior policy, thereby reducing the discrepancy between training and inference trajectories while maintaining the optimality of the reference policy, as shown in Figure 3 Unlike DPT, which trajectory, our step-wise supervision framework is inherently designed to align with high-efficiency chunk-wise training pipelines for sequence models such Transformers [65] and their optimized variants [66 67], producing T times training efficiency (T is the sequence length).
data aggregation techniques for imitation learning [3] [4]
imitates only one-step action conditioned
Prior knowledge augmented ICRL: For SS, which employs diverse policies for trajectory generation, prior knowledge becomes crucial for interpreting actions derived from heterogeneous policies [68]. This motivates the incorporation of prior knowledge, specifically metadata indicating the policy used to generate each action. In this work, we implement a diverse set of behavior policies IT = (x(0), with (b) denoting the behavior policies. x(6) comprises seven distinct types, including myopic greedy, oracle, Q-learning, and model- based reinforcement learning, denoted by marker tag(z) E (0. 6]. To handle unseen or unclassified policies, we reserve an additional identifier "Unk" with m 7. The trajectory is denoted by IT((T = [(51, 91, 01, r ri),... (8t, gT, aT; rT)], where 8 E S, ar E A, az ~ (x(b) (st), Ft ~ R(8,, at, 11+1), ), and gt = g(a:) = tag n (b) denotes the prior knowledge for action.
Chunwise Training: We use ग² denote the reference policy, which is the oracle policy with > 0.99. is then used to label a list of the reference actions step-by-step as (T(7) = [a 42.... d7] with ot ~ pi- = x² (st). Notice that for most of the time at + at. We first collect the training and validation datasets by:
D(T) = {< hr (), T III
(3)
Scaling ICL to handle complex tasks necessitates the efficient modeling of extensive contexts, which in turn demands substantial computational resources and large-capacity hardware memories. To further break down the limitation in context length, we break long sequence hr into K segments [1,Ti], [Ti + 1. [Tx-1 T.] The forward pass is calculated recurrently across the segments, and the backward calculation is performed within each segment. The gradient for the memory states of the linear attention layer is blocked across the segments. Thus the sequence is encoded in the following chunkwise manner:
PT. Pr.+1:-..,Ph.. PTe, 0b+1 (oausalgSS((ii) S1,+1, ... S12
STa+1
with SG representing stopping gradient. the meta-training process, the gradients are calculated within each segment and accumulated in cache first. They are applied to the parameters only at the end of the trajectory. The final target as follows:
Minémize L
Σ Σ Eorlogp' (a))
(5)
NTJTED
The complete architecture of the model and its training pipeline are depicted in Figure We utilize a causal sequence model (where the prediction at each position relies solely on information from preceding positions) to encode the input sequence and forecast the action at the position aligned with the state inputs. In this paper, we mainly employ linear attention layers including gated slot attention (GSA) [67] layers, mamba-2 [69], and RWKV-7 [70] as the causal sequence models.
5 Experiments
5.1 Demonstration of Generalization and Scalability
We first validate the representational capability of AnyMDP tasks as universal MDPs. To this end, we collect dataset Dera (T (T(n,, n.)) comprising 512K sequences for training, where n,

---

Segment
Segment
Reference Policy
(Oracle)
Cre
188*.
Causal Sequence M Modeling
Backward
Forward
92-1
Behavior Policies
Figure 4: OmniRL model structure and training losses.
Table 1: Comparison of best average episodic performance within 10,000 episodes for each learner. The episode performances are normalized to scale of 0% (uniform random policy) to 100% (oracle policy). The minimum steps and episodes required achieve within at least 99% of the best episodic performances are also listed. The hyper-parameters parameters for Q-Learning and PPO are optimized under the evaluated task or task set. Results for AnyMDP and Garnet Tasks are averaged over 64 tasks, results for Gymnasium tasks are averaged over 3 independent runs for each task. Gymnasium, DarkRoom, and Bandits tasks are entirely absent from OmniRL's ' training regimen.
Environments
Performances/ / AVG. Steps cost AVG. Episodes cost
OmniRL (AnyMDP)
OmniRL (GametMDP)
TQL-UCB
PPO
(1)(1.5) (Bandits)
92.1% 1100/100010 110
95.6% 111.//1///2/ 11.2K
82.5% 11///113///1/ 10/3
46.6% /33
(1,(16,5)
92.0% /297////7
90.6%/476K] 79.7K
95.3%. /2.0%/ /29
47.8% /1.. 224
(32,5)
84.7% /71K//666/5666
/612K/
/9.7K
90.3% /6.5K /47
42.0% 55.0K/
Tis((4,55)
63.7%/
58.3%/\.11M/9.4F
91.3% /25
47.1%
(128128,
73.2%/
7..0%/1.M/8.66
80.2% /3/33/33110000
32.3%/
Gamet(16,5522
/2.8%/241K//.1
97.1%/57K /0.5
85.9% 85.9%
99.0%/10.8KK /95
(64,
98.7%/614K/1.77
98.1%/96K /0.26K
80.4%/8.0K/19
87.3%/7.AK/ ///////KK//33
ClittWakking
100%/3.1K//S
95.9% 3K/ /2.7K
100% /3 OK /65
63%, /29 /300
FrozenLake (non-slippery)
95.3%/23..6X3 /3.7K
96.8% 22/.// /2.1K
99.8% /0.3K /35
75.1% /L0KK /250
FrozenLake (slippery)
96%/208K 10.0K
95.6% 73 GK /4.7K
79.3% /7.7 7.7K /245
31.3% 11.88 1L88/8800 /800
Discrete-Pendulum (g=1)
94.9% 94.9%/22K// (110
99.3% 198 K/ /990
90.5% /SK /40
0.0%
Sescrete-eendulmm (g=S)
99.7%//26KK/
23K
99.8%/13255 K/660
91.8% /4K /170
0.0%
Discrete-Pendulle Pendullur (g=9.8)
20.2%%/2.0.///
98 08.3%/11611 11118866/33
73.4% /33/ / /165
0.0%
Switch2 (Multi-Agent) 7
98%/3.8K//10
80.4%/2.8K 8//00
Darkroom (6x6)
98.1%
97.6%/10..6. /560
95.2% /84///5
90.5%, 10.5%%/21.334 21/3///11//11 4440
Darkroom (8x8)
26.8%%/24.5kk K..5KK///KKK
96, 96.7%/15. 9 9/930
93 83.8% 1.5K /40
88 9%, /30 AK/ 4880
Darkroom (10x10)
80% /1.7K
92.3% 12.3% /570
/2.7% 1000
75.6%//2.88K, /2. /280
[16, 128). n, = 5. The length of each sequence T is 12K, resulting in a total of 6B time steps. For testing, we independently sample tasks Tis with n, {1, 16, 32, 64, 128]. ensuring each n, group contains 256 tasks.
The meta-training process is primarily conducted using 8 Nvidia Tesla A800 GPUs. We use a batch size of per GPU, divided into segments (chunks) of 2K steps each. We optimize using the Adam W algorithm with a learning rate that decays from a peak value of x 10-4, The average time cost per iteration is 8 seconds for trajectories with 7' 12K, and this cost increases linearly with sequence
La Gd
Ea Ga
Sape Crtec
Figure Left: Normalized episodic return of TQL-UCB, PPO and OmniRL on 64 AnyMDP evaluation tasks (mean +95% 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 RL trend.
7

---

length. For more details please check Appendix C.2 For the causal sequence model, we evaluate four architectures: RWKV-7 [70], Gated Delta-Net (GDN) 772). Gated Self-Attention (GSA) [ Mamba2 [73]. The test results are largely consistent with the conclusions reported in language processing (RWKV-7 R GDN > GSA > Mamba2, with details in Appendix [D.I), demonstrating the capability of AnyMDP to serve as a benchmark for long-term sequence modeling. Therefore, we select RWKV-7 for subsequent experiments.
Without any further parameter tuning, we evaluate our model, namely OmniRL, on both unseen AnyMDP tasks in Figure 5 Gymnasium tasks, and DarkRoom [3] in Figure and those perfor- mances are shown in Table I Notably, unlike previous ICRL works, our training set does not include any instances of DarkRoom. our experiments. the selected tasks are constrained to environments with observation spaces of dimension n, I 128 and action spaces of dimension n. I 5. For environments with continuous observation spaces, such as Pendulum-vl, we manually discretize the observation space into 60 discrete classes using a grid-bases discretization method. To adapt OmnIRL that is trained with na = 5 to environments with less actions (Na < 5), we reassign unused actions to valid ones. This further demonstrates the compatibility of OmniRL across environments with varying action space dimensions. We also found that proper reward shaping is important for OmniRL to work, shown in Figure 3 the details can be found at Appendix
In Table we compare the normalized performance, episode cost, and step costs of OmniRL, classical Tabular Q-learning (TQL) [60] with upper confidence bound (UCB) [61] (TQL-UCB for short), and Proximal Policy Optimization (PPO) 62). OmniRL, meta-trained solely on AnyMDP, adapts effectively to most unseen tasks, confirming AnyMDP's representational power; by contrast, OmniRL. trained on Garnet MDPs excels only on Garnet MDPs. The results also demonstrates OmniRL's superior sample efficiency, which aligns with prior ICRL findings. Notably, despite being trained solely on single-agent tasks, OmniRL adapts to multi-agent tasks like Switch2 by configuring observation spaces, enabling emergent inter-agent cooperation without explicit multi-agent interaction during training and thus decoupling cooperative behavior emergence from centralized mechanisms. Furthermore, in line with expectations, solving MnP tasks becomes more difficult with increased (n,)e(n,)) or action(n,)) space size, with PPO proving more sensitive to action space extension and TQL-UCB more vulnerable to state space growth, as illustrated in Figure
5.2 OmniRL Performs Both Offline and Online Learning Better
For the ablation study and comparison with the other methods including AD, AD and DPT, we collect smaller dataset with (ammat] comprising 128K sequences for training, where n1 16, n, 5, 8K, with a total of 1B time steps. Figure summarizes the performance of different methods trained on Dimall with identical training iterations. The comparison includes AD, AD DPT, OmniRL, and OmniRL (w/o priori) where the prior info 9: is removed from the sequence.
We examine the performance of different methods with different initial contexts: (1) Online- -L:nee-RL:TTe The agent starts with an empty trajectory (hn- = 0). (2) Offline-RL.: The agent starts with an existing context derived from imperfect demonstrations (e.g., disturbed oracle policy) (hg = h²). (3) Imitation Learning: The agent starts with an existing context derived from oracles(hg (exp)). For all three categories, the subsequent interactions are continually added to the context. Therefore, the models differ only in their initial memory cache. The evaluation assesses the agents' abilities in two key areas: their capacity to exploit existing information and their ability to explore and exploit continually. In the results in Figure OmniRL and OmniRL (w/o priori) surpass AD, AD and DPT with large gap, validating the effectiveness of Step-wise Supervision (SS). OmniRL (w/o a priori) lags behind OmniRI with noticeable gap in all three groups, demonstrating the effectiveness of integrating the prior information. Table and Figure 17 further demonstrate that the offline offline-leaming ability of OmniRL can generalize to Gymnasium environments.
5.3 Emergence of General-Purpose ICRL by Increasing Task Number
We validate task diversity's critical role in ICRL via independent meta-training across four datasets, each (Trro (16, 5)) containing 128K sequences but differing in task numbers (TTraal E {100, 1K, 10K, 128K), }). Note that different trajectories can be generated from a single task, arising from the diverse behavior policies and random sampling in both decision and transition. We examine how the validation losses C, on both seen and unseen tasks change with the number of meta-training
8

---

Orline-LL T,/16
OffineRRL Tyy
Inkation Learring r.
Figure 6: AD, AD DPT, and OmniRL on 32 AnyMDP tasks, each tested under three initial contexts tailored for online RI,, offline RL,, and imitation learning. offline-RI. and II. agents are seeded with 100 demonstration episodes before any interaction begins; results highlight the gains from DPD and prior-infon prior-ifformation integration.
Table 2: Offline-RL. and IL performance of meta-trained OmniRL versus Conservative Q-Learning (CQL) 774 on four unseen Gymnasium tasks, using demonstrations from both oracle and random policies as initial contexts. Returns in episodes (-200 measure offline RI. and II. efficacy; returns in episodes 180-200 show the benefit of subsequent online RI.
0--20/180-200 EPISODES PERFORMANCE
ENVIRONMENTS
CQL
OMNIRL
ORACLE (77. (77.80%
76.61% / 76.11%
70.12% / 77.27% 54.3% / 67.38%
FROZENLAKE (SLIPPeRy)
RANDOM (1.46%)
37.79% / 72.36%
ORACLE (-13)
-30.8 / -13
-13 /13
CLIFF
RANDOM (- (-109.84)
-560.2 -16.2
/-17
ORACLE (-153.81)
-605.89 -258.22
-180/-127
DISCRETE- DISCRETE-PENDULUYX (G=5)
RANDOM (-941.65)
-1022.40 -184.29
-646 / -208
ORACLE (0.22)
0.23 / 0.22
0.22 / 0.21
DARKROOM (10X10)
RANDOM (-15.07)
-4.05 0.21
0.09 0.19
iterations (outer-loop steps) and steps in context / (inner-loop steps) simultaneously; the results are shown in Figure 7 We remark the following observations:
Task scale and diversity are crucial the generalization of ICRL. The previous investigation on ICL [25] emphasizes the importance of "burstiness". Our results demonstrate for the first time that even when using "bursty" sequences alone, both the number of tasks and their overall diversity remain critically important. Specifically, in groups with |T S 10K, over-training leads to transiency of ICL 75 unseen tasks but a continued improvement in seen tasks. These findings confirm and extend the discovery of the "task identification" phase mentioned in Kirsch et al. [3), Pan et al. [6.. Drawing on the theories of ICL and IWL in Chan al. (77), possible explanation
that IWL dominates performance, with the model memorizing tasks and ICL selecting the correct one, leading to fast seen-task adaptation but poor unseen generalization. As the number of tasks increases continuously, the model becomes more dependent on ICL since memorizing task-specific information becomes less feasible. This is characterized by the improved generalization to unseen tasks and longer adaptation periods in both seen and unseen tasks, as shown in Figure 7
Long-context adaptation is a tax pay for generalization scope. Our results highlight a key insight on ICRL evaluation. Most previous ICRL works assess performance based on the average results over fixed, short context span. However, our findings indicate that more generalized in- context learners may actually perform worse in zero-shot and even few-shot evaluations, particularly when there significant overlap between the training and evaluation sets, i.e., when evaluation sets are closer to seen tasks. Therefore, we argue that it is more critical to focus on the asymplotic performance of learner. This can effectively evaluated by examining the performance at the final steps or episodes of sufficiently long context, rather than short-term metrics.

---

Tov0=000
[f//=1K
1o00=100
[=//=1K
1.50
1.50
1.50
0.50
0.50
0.50
0.50
1K
24
5K
3K
2K
1k
15K
43K
8K
1K"
2K
2K
5K
800
10K
20K
8K
65
2k
[/m/- ffn/-10K
128K
Tir |OK
Tra 128K
1.50
1.50
1.50
1.50
0.50
0.59
0.50
0.59
40K
Gr
2x
20k Faining Aerraiona
aU5
bok
Bg. Ge
*x 2x o Contect Length
20k
40k
bok
2K
20k Tairinn Barations
ack
6ok
Be Ge
Content Length
Bix
BOX
8K
(a) Seen Tasks
(b) Unseen Tasks
Validation
Tera
Metrics
set
100
IK
10K
128K
Seen
max(d:) min(t) s.t. de 2 80%
81.0%
> 65.44%
86.6%
84.5%
tasks
0.88K
2.4K
5.2K
(ax)
17.9%
38.0%
84.4%
2 284.8%
Uneen
min(l) s.t. de > 80%
3.9K
5.1K
tasks
max(de) achieved iteration
2K
14K
72K
2 80K
Figure 7: Position-wii wise validation losses (Lt, where lower values indicate better performance) and their properties on both seen and unseen tasks across meta-training iterations, varying context lengths, and variant number of tasks 'Teral- Each of the groups of training data had 128K sequences, which were generated from 100, 1K, 10K, and 128K tasks, respectively. Each dataset underwent meta-training for up to 80K iterations.
the table, the notation ">" indicates values that could not be fully determined due to training being stopped early when performance on unseen tasks began to deteriorate. The normalized gain of ICL is defined as de = 1 - Li/ Co. representing the improvement of performance the context (t) increases. The table summarizes key findings of this study: as the number of tasks increases, the minimum step cost required to achieve an 80% normalized gain of ICL also increases.
Conclusions and Discussions
We introduce AnyMDP, scalable, low-bias task suite for benchmarking In-Context Reinforcement Learning (ICRL). and a companion framework featuring two innovations: Decoupled Policy Distilla- tion and mnor-information induction. Trained solely on AnyMDP, our model OmniRL outperforms prior methods in generalization, mastering wider task families and supporting more diverse learning paradigms.
Broader impact: Complementing prior studies, our findings highlight that task diversity and long sequence length are key determinants of general-purpose ICRL. Our results also indicate that long context is a necessary cost for generalization, thus advocate shifting evaluation metrics toward asymptotic performance measures. This work further motivates the construction of carefully curated synthetic datasets specifically designed for large-scale meta-training.
Limitations and future work: While procedural Discrete MDPs provide clean benchmarks for ICRL and long-contes modeling, their extension to continuous state-action spaces remains bottlenecked. We recognize that central challenge in this direction is to balance task complexity, task diversity, and data integrity while simultaneously calibrating generalization scope against practical constraints.
10

---

Acknowledgments and Disclosure of Funding
This project is supported by Longgang District Shenzhen's "Ten Action Plan" for supporting in- novation projects (Grant LGKCSDPT2024002) Guangdong Provincial Leading Talent Program (Grant No. 2024TX08Z3199, National Natural Science Foundation of China (62033012), Guangdong S&T Program (Grant No. 20223000000000033) and the "Zhiguo" Action of Guangxi Science and Technology Program under Grant No.72G250398003.
References
[1] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models few-shot learners. Advances in neural information processing systems, 33:1 13:7177-19 2020.
[2] Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In International conference on machine learning, pages
PMLR, 2021.
[3] Danny Driess, Fei Xia, Mehdi SM Sajjadi, Corey Lynch, Aakanksha Chowdhery, Brian Ichter, Ayzaan Wahid, Jonathan Tompson. Quan Vuong, Tianhe Yu, et al. Palm-e: An embodied multimodal language model. arXiv preprint arXiv:2303.03378 2023.
[4] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timo- thée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.
2023.
[5] Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint arXiv:2303.08774 2023.
[6] Adam Santoro, Sergey Bartunov, Matthew Botvinick, Daan Wierstra, and Timothy Lillicrap. Meta-leamning with memory-augmentt neural networks. In International conference machine learming, pages 1842-1850. PMLR, 2016.
[7] Shivam Garg, Dimitris Tsipras, Percy Liang, and Gregory Valiant. What can transformers learn in-context? a study of simple function classes. Advances in Neural Information Processing Systems, 35:30583-30595 2022.
[8] Scott Reed, Konrad Zolna, Emilio Parisotto, Sergio Gómez Colmenarejo, Alexander Novikov, Gabriel marto, Mai Giménez, Yury Sulsky, Jackie Kay, Jost Tobias Springenberg, et al. generalist agent. Transactions on Machine Learning Research, 2022.
[9] Letian Fu, Huang Huang, Gaurav Datta, Lawrence Yunliang Chen, William Chung-Ho Panitch, Fangchen Liu, Hui Li, and Ken Goldberg. In-context imitation learning via next-token prediction. In Ist Workshop on X-Embodiment Robot Learning, 2024.
[10] Vitalis Vosylius and Edward Johns. Few-shot in-context imitation learning via implicit graph alignment. In 7th Annual Conference on Robot Learning, 2024.
[11] Jane X Wang. Zeb Kurth-Nelson, Dhruva Tirumala, Hubert Soyer, Joel Z Leibo, Remi Munos, Charles Blundell, Dharshan Kumaran, and Matt Botvinick. Learning to reinforcement learn. arXiv preprint ar 105763, 2016.
[12] Yan Duan, John Schulman, Xi Chen, Peter Bartlett, Ilya Sutskever, and Pieter Abbeel. RI2: Fast reinforcement learning via slow reinforcement learning. arXiv preprint arXiv: 161 1.02779, 2016.
[13] Michael Laskin, Luyu Wang, Junbyuk Oh, Emilio Parisotto, Stephen Spencer, Richie Steiger- wald, DJ Strouse, Steven Stenberg Hansen, Angelos Filos, Ethan Brooks, et al. In-context reinforcement learning with algorithm distillation. In NewrIPS 2022 Foundation Models for Decision Making Workshop, 2022.

---

[14] Gregory Koch, Richard Zemel, Ruslan Salakhutdinov, et al. Siamese neural networks for one-shot image recognition. In ICML deep learning workshop, volume 2. pages 1-30. Lille, 2015.
[15] Ilya Zisman, Vladislav Kurenkov, Alexander Nikulin, Viacheslav Sinii, and Sergey Kolesnikov. Emergence of in-context reinforcement learning from noise distillation. Forty-first Interna- lional Conference on Machine Learning. 2024.
[16] Jonathan Lee, Annie Xie, Aldo Pacchiano, Yash Chandak, Chelsea Finn, Ofrr Nachum, and Emma Brunskill. Supervised pretraining can learn in-context reinforcement learning. Advances in Neural Information Processing Systems, 36, 2024.
[17] Fan Wang, Hao Tian, Haoyi Xiong, Hua Wu, Jie Fu, Yang Cao, Yu Kang, and Haifeng Wang. Evolving decomposed plasticity rules for information--ottlenecck. meta-learning. Transactions on Machine Learning Research, 2022.
[18] Alexander Nikulin, Vladislav Kurenkov, Ilya Zisman, Artem Agarkov, Viacheslav Sinii, and Sergey Kolesnikov. Gland-minigrid: Scalable meta meta-reinfoon reinforcement learning environments in jax. arXiv preprint arXiv:2312.12044 2023.
[19] Steven Morad, Ryan Kortvelesy, Matteo Bettini, Stephan Liwicki, and Amanda Prorok. Popgym: Benchmarking partially observable reinforcement learning. arXiv preprint arXiv:2303.01859.
2023.
[20] Karl Cobbe, Chris Hesse, Jacob Hilton, and John Schulman. Leveraging procedural generation to benchmark reinforcement learning. In International conference on machine learning, pages 2048 2056. PMLR, 2020.
[21] Sebastian Thrun and Lorien Pratt. Learning to learn: Introduction and overview. In Learning learn, pages 3-17. Springer, 1998.
[22] Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adap- tation of deep networks. In International conference on machine learning, pages 1126-1135, PMLR, 2017.
[23] Yanda Chen, Ruiqi Zhong, Sheng Zha, George Karypis, and He He. Meta-learning via language model in-context tuning. arXiv preprint arXiv:21 10.07814, 2021.
[24] Julian Coda-Forno, Marcel Binz, Zeynep Akata, Matt Botvinick, Jane Wang, and Eric Schulz. Meta-in-context learning in large language models. Advances in Neural Information Processing Systems, 36:65189-6520 2023.
[25] Sharath Chandra Raparthy, Eric Hambro, Robert Kirk, Mikael Henaff, and Roberta Raileanu. Generalization to new sequential decision making tasks with in-context learning. In Proceedings of the 41st International Conference on Machine Learning, pages 42138-42158, 2024.
[26] Ezra Edelman, Nikolaos Tsilivis, Benjamin Edelman, Eran Malach, and Surbhi Goel. The evolution of statistical induction heads: In-context learning markov chains. Advances in Neural Information Processing Systems, 37:64273-64311, 2024.
[27] Toni J.b. Liu, Nicolas Boulle. Raphaël Sarfati, and Christopher Earls. Llms learn governing principles of dynamical systems, revealing an in-context neural scaling law. In Proceedings of the 2024 Conference on Empirical Methods in Natural Language Processing, pages 15097- 15117, Miami, Florida, USA, 2024. Association for Computational Linguistics.
[28] Stephanie Chan, Adam Santoro, Andrew Lampinen, Jane Wang, Aaditya Singh, Pierre Richemond, James McClelland, and Felix Hill. Data distributional properties drive emer- gent in-context learning in transformers. Advances in Neural Information Processing Systems, 8188818891, 2022.
[29] Yu Zhao, Yuanbin Qu, Konrad Staniszewski, Szymon Tworkowski, Wei Liu, Piotr Miloš, Yuxiang Wu, and Pasquale Minervini. Analysing the impact of sequence composition on language model pre-training. Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1; Long Papers), pages 7897-7912, 2024.
12

---

[30] A Raventós, M Paul, F Chen, et al. Pretraining task diversity and the emergence of non-bayesian in-context learning for regression. In Advances in Neural Information Processing Systems, volume 36, pages 14228-14244 14246, 2023.
[31] Stephanie CY Chan, Ishita Dasgupta, Junkyung Kim, Dharshan Kumaran, Andrew K Lampinen, and Felix Hill. Transformers generalize differently from information stored in context vs in weights. arXiv preprint arXivv2210.05660 2022.
[32] Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. In International Conference on Machine Learning, pages 35151-35174, PMLR, 2023.
[33] Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of in-context learning as implicit bayesian inference. In International Conference on Learning Representations, 2022.
[34] Louis Kirsch, James Harrison, Jascha Sohl-Dickstein, and Luke Metz. Gereral-purpose in- context learning by meta-learning transformers. arXiv preprint arXiv:22/2.04455, 2022.
[35] Louis Kirsch, James Harrison, Daniel Freeman, Jaseha Sohl-Dickstein, and Jürgen Schmidhuber. Towards general- peneral-purpose in-context learning agents. Workshop on Distribution Shifts, 37th Conference on Neural Information 2023.
[36] Fan Wang, Chuan Lin, Yang Cao, and Yu Kang. Benchmarking general purpose in-context learning. arXiv preprint arXiv:2405. 2024.
[37] Wang Fan and Shaoshan Liu. Putting the smarts into robot bodies. Commun. ACM, 68(3)6--, 2025. doi: 10.174577733761 /ttts://:://do org/10. 1145/3703761
[38] Dániel Barabási, André Ferreira Castro, and Florian Engert. Three systems of circuit formation: assembly, updating and tuning. Nature Reviews Neuroscience, pages 1-12, 2025.
[39] Jane Wang. Zeb Kurth-Nelson, Hubert Soyer, Joel Leibo, Dhruva Tirumala, Remi Munos, Charles Blundell, Dharshan Kumaran, and Matt Botvinick. Learning to reinforcement learn. In Proceedings of the Annual Meeting of the Cognitive Science Society, volume 39, 2017.
[40] Nikhil Mishra, Mostafa Rohaninejad, Xi Chen, and Pieter Abbeel. A simple neural attentive meta-learner. In International Conference on Learning Representattins, 2018
[41] Jake Grigsby, Linxi Fan, and Yuke Zhu. Amago: Scalable in-context reinforcement learning for adaptive agents. In The Twelfth International Conference on Learning Representations, 2024.
[42] Elias Najarro and Sebastian Risi. Meta-learning through hebbian plasticity in random networks. Advances in Neural Information Processing Systems, 33:20719-20733 2020.
[43] Mohammad Gheshlaghi Azar, Ian Osband, and Rémi Munos. Minimax regret bounds for reinforcement learning. In International conference on machine learning, pages 263-272. PMLR, 2017.
[44] Tian Xu, Ziniu Li, and Yang Yu. Error bounds of imitating policies and environments. Advances in Neural Information Processing Systems, 33:15737-15749, 2020.
[45] Zhejun Zhang, Alexander Liniger, Dengxin Dai, Fisher Yu, and Luc Van Gool. End-to-end urban driving by imitating reinforcement learning coach. In Proceedings of the IEEE/CVF international conference on computer vision, pages 15222-15232, 2021.
[46] Tianhe Yu, Deirdre Quillen, Zhanpeng He, Ryan Julian, Karol Hausman, Chelsea Finn, and Sergey Levine. Meta-world: A benchmark and evaluation for multi-task and meta reinforcement learning. In Conference on robot learning, pages 1094-1100. PMLR, 2020.
[47] Lucas N Alegre, Florian Felten, El-Ghazali Talbi, Grégoire Danoy, Ann Nowé, Ana LC Bazzan, and Bruno C da Silva. Mo-gym: A library of multi-objective reinforcement learn- ing environments. In Proceedings of the 34th Benelux Conference on Artificial Intelligence BNAIC/Benelearn, volume 2022, page 2, 2022.
13

---

[48] Joachim Winther Pedersen and Sebastian Risi. Evolving and merging hebbian learning rules: increasing generalization by decreasing the number of rules. In Proceedings of the Genetic and Evolutionary Computation Conference, pages 892-900, 2021.
[49] Ekin Akyürek, Bailin Wang, Yoon Kim, and Jacob Andreas. In-context language learning: Architectures and algorithms. International Conference on Machine Learning, pages 787-812.
PMLR, 2024.
[50] Viacheslav Sinii, Alexander Nikulin, Vladislav Kurenkov, Ilya Zisman, and Sergey Kolesnikov. In-context reinforcement learning for variable action spaces. In Proceedings of the 41 International Conference Machine Learning, pages 45773-45793, 2024.
[51] Xue Bin Peng, Marcin Andrychowicz, Wojciech Zaremba, and Pieter Abbeel. Sim-to-real transfer of robotic control with dynamics randomization. In 2018 IEEE international conference on robotics and automation (ICRA), pages 3803-3810. IEEE, 2018.
[52] Cultural General Intelligence Team, Avishkar Bhoopchand, Bethanie Brownfield, Adrian Collis- ter, Agustin Dal Lago, Ashley Edwards, Richard Everett, Alexandre Frechette, Yanko Gitahy Oliveira, Edward Hughes, al. Learning robust real-time cultural transmission without human data. arXiv preprint arXiv:2203.00715. 2022.
[53] Ilge Akkaya, Marein Andrychowicz, Maciek Chociej. Mateusz Litwin, Bob McGrew, Arthur Petron, Alex Paino, Matthias Plappert, Glenn Powell, Raphael Ribas, al. Solving rubik's cube with robot hand. arXiv preprint arXiv: 19100// 13, 2019.
[54] Open Ended Learning Team, Adam Stooke, Anuj Mahajan, Catarina Barros, Charlie Deck, Jakob Bauer, Jakub Sygnowski, Maja Trebacz. Max Jaderberg, Michael Mathieu, et al. Open- ended learning leads to generally capable agents. arXiv preprint arXiv:2107. 12808 2021.
[55] Alexander Nikulin, Vladislav Kurenkov, Ilya Zisman, Artem Agarkov, Viacheslav Sinii, and Sergey Kolesmikov. Xland -iind-iiiggri::: Scalable meta-reiinorremment learning environments in jax. Advances in Neural Information Processing Systems, 37:43809-433835. 2024.
[56] Jakob Bauer, Kate Baumli, Feryal Behbahani, Avishkar Bhoopchand, Nathalie Bradley-Schmieg, Michael Chang, Natalie Clay, Adrian Collister, Vibhavari Dasagi. Lucy Gonzalez, et al. Human- timescale adaptation in open-ended task space. In Proceedings of the 40th International Conference on Machine Learning, pages 1887-1935, 2023.
[57] Thomas Welsh Archibald, KIM McKinnon, and Lyn C Thomas. On the generation of markov decision processes. Journal of the Operational Research Sociery, 46(3):354-366 1995.
[58) Shalabh Bhatnagar, Richard S Sutton, Mohammad Ghavamzadeh, and Mark Lee. Natural actor-critic algorithms. Automatica, (5(11)2477-2482 2009.
[59] Richard Bellman. Dynamic programming. Chapter IX, Princeton University Press, Princeton, New Jersey, 1958.
[60] Richard S Sutton, Andrew G Barto, et al. Reinforcement learning: An introduction, volume 1. MIT press Cambridge, 1998.
[61] Chi Jin, Zeyuan Allen-Zhu, Sebastien Bubeck, and Michael Jordan. Is q-learning provably efficient? Advances in neural information processing systems, 31, 2018.
[62] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.0634f 2017.
[63] Stéphane Ross and Drew Bagnell. Efficient reductions for imitation learning. In Proceedings of the thirteenth international conference on artificial intelligence and statistics, pages 661-668. JMLR Workshop and Conference Proceedings, 2010.
[64] Stéphane Ross, Geoffrey Gordon, and Drew Bagnell. A reduction of imitation learning and structured prediction to no-regret online learning. In Proceedings of the fourteenth interna- tional conference on artificial intelligence and statistics, pages 627-635. JMLR Workshop and Conference Proceedings, 2011.
14

---

[65] A Vaswani. Attention is all you need. Advances in Neural Information Processing Systems, 2017.
[66] Yingcong Li, Muhammed Emrullah Ildiz, Dimitris Papailiopoulos, and Samet Oymak. Trans- formers as algorithms: Generalization and stability in in-context learning. In International Conference on Machine Learning, pages 19565- PMLR, 2023.
[67] Yu Zhang, Songlin Yang, Rui-Jie Zhu, Yue Zhang, Leyang Cui, Yiqiao Wang, Bolun Wang, Freda Shi, Bailin Wang, Wei Bi, et al. Gated slot attention for efficient linear-time sequence modeling. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
[68] Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas. and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems, 2021.
[69] Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00052, 2023.
[70] Bo Peng, Ruichong Zhang. Daniel Goldstein, Eric Alcaide, Haowen Hou, Janna Lu, William Merrill, Guangyu Song, Kaifeng Tan, Saiteja Utpala, et al. Rwkv-7" goose" with expressive dynamic state evolution. arXiv preprint arXiv:2503. 14456, 2025.
[71] Anurag Koul. ma-gym: Collection of multi-agent environments based on openai gym. //github.com /otiluubccom//oolauurga///oo gym 2019.
https
[72] Songlin Yang, Jan Kautz, and Ali Hatamizadeh. Gated delta networks: Improving mamba2 with delta rule. 2025.
[73] Tri Dao and Albert Gu. Transformers are ssms: generalized models and efficient algorithms through structured state space duality. Proceedings of the 41st International Conference on Machine Learning, pages 10041-10071, 2024.
[74] Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. Advances in neural information processing systems, 33: 1179-1191, 2020.
[75] Aaditya Singh, Stephanie Chan, Ted Moskovitz, Erin Grant, Andrew Saxe, and Felix Hill. The transient nature of emergent in-context learning in transformers. Advances in Neural Information Processing Systems, 36, 2024.
[76] Jane Pan, Tianyu Gao, Howard Chen, and Danqi Chen. What in-context learning" learns" in-context: Disentangling task recognition and task learning. In The 61 Annual Meeting Of The Association For Computational Linguistics, 2023.
[77] Bryan Chan, Xinyi Chen, András György, and Dale Schuurmans. Toward understanding in-context in-weight learning. In The Thirteenth International Conference Learning Representations, 2025.
[78) Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang, Xiao Bi, et al. Deepseek-rl: Incentivizing reasoning capability in Iims via reinforcement learning. arXiv preprint arXiv:2501.1294 2025.
[79] Peter C Wason and J St BT Evans. Dual processes in reasoning? Cognition, 3(2):141-154, 1974.
[80] Daniel Kahneman. Thinking, fast and slow, Farrar, Straus and Giroux, 2011.
15

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
The main contributions of this paper include: (1) a scalable task set; (2) improved ICRL framework; and (3) in-depth discussions on data distribution and the emergence of ICRL. These contributions are explicitly elaborated in both the abstract and the concluding section of the introduction.
Justification:
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long it is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
In Section
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
16

---

Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: The paper mainly provide empirical validations. Slight theoretical results and proof' are provided in Section
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: We provide all the information needed for the emprical results in Section and appendices.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access hosted model (e.g.. in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (c.g., to registered users), but it should be possible for other researchers to have some path to reproducing verifying the results.
17

---

reproducible, they 5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Answer: [Yes]
Justification:
We attach the codes and data necessary to reproduce the results.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.co cc/
public/guides
for more details.
While encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
7/nips cc/ wtc/publicc/ hoolllccggiiess. gwt.///ieee//eee// Codeeoomiiiieeee /esooeuuuude SanSSubbissss sebbbbisssiooi ssisssoicccyyy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: In Section 5 and appendices.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: We report most results with error bars. For results that are too computationally costly and less noisy, we do not.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
18

---

The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful figures symmetric error bars that would yield results that are out of range (e.g. negative rates).
show in tables or
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: In Section 5 and appendices.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't 't make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics
Answer:
[Yes]
Justification:
The research poses no risk at violating NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- cration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA].
Justification:
There no societal impact related to the codes, data, and model at the current
stage.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
19

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
The paper includes synthetic data only that poses risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: The code and assets used in this study either originally created by sourced from open-source repositories and properly referenced.
Guidelines:
The answer NA means that the paper does use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
20

---

If assets are released, the license, copyright information, and terms of use in the thcode. has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
package should be provided. For popular datasets, paperswi
latasets.paperssit
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: Readme file for the codes and meta-data for datasets are provided.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
Crowdsourcing or research with human subjects not involved.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: Crowdsourcing or research with human subjects not involved.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
21

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMS if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, originality of the research, declaration is not required.
Question:
Answer: [Yes]
Justification: Appendix D.5 uses LLM for comparison and clearly states its usage and sources.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non stansaaad components. wwt.o//oeeeenece///////////
Please refer to our LLM policy for what should or should be described.
A Notations
B
Additional Information on AnyMDP
B.1 Proof of Theorem
I
We began by establishing the upper bound for (s().. Notably, if SE 0 the upper bound would be further reduced for indices > b. as system termination resets states to indices j < bo. Consequently, it suffices to analyze the case where SE as non-empty SE only tightens the bound.
Under the conditions that Ejer Pr(sy]ss) (ss) n if > b_ and Vj > + bt. Pr(syls;) O, we construct worst case transition kernel. This kernel defined as follows:
Vs; with nı - b_ > b, P4 (8 = n, (4 (8:) = ==-11 17. P+ (-18:) = 0 for other cases.
It follows directly from this construction that (s,) < p+ (3;) for > b. as P+ maximizes transition probabilities later states under the given constraints.
By the definition of the SD, we have:
(8i-)) + (1
(+:),
(6)
where P+, is the SD of s; under transition P4. Then, define that ô p+ (s:) (s-))/ from which we derive the following:
p+ (8146.)//44 = 1-8-1
(1 11) -)
By applying the conditions 0 < < 1/1b++1111 +1) and > 1.0 5, we can further bound the expression in Equation as follows:
p+ (si4b, )/p+ (x:) < 5).+.
(8)
Equation proves that p+ (s:) decays a rate faster than (1 6) as increases. Consequently. the SD also decays faster than (1 8).
To establish the lower bound, we can utilize the constraint Σz>t Pr(sy]ss) and construct worst-case transition matrix follows:
Vs; with n, 1 i > P_ (si-o_) P. = P-(ss) = 0 for other cases.
By analyzing this construction, we can validate that pantpp(s)) (8:) decays rate slower than when increases. Consequently, the SD pe also decays slower than e. This completes the proof of Theorem
22

---

Table Default simplifications and notations used throughout the paper.
ICL
IN-CONTEXT LEARNING
IWL
IN-WEIGHT LEARNING
ICRL
IN CONTEXT REINFORCEMENT LEARNING
MC
MARKOV CHAIN
MDP
MARKOV DECISION PROCESS
SS
STEP-WISE SUPERVISION
SD
STATIONARY DISTRIBUTION
STATE OR OBSERVATION SPACE
So
STATES FOR RESET
SE
STATES TRIGGERING TERMINATION LENGTH
TE
EPISODE
a()
TRANSITION FUNCTION OF 8, 0,0 +
R(),a,8')
REWARD FUNCTION OF 8,
DISCOUNT FACTOR FOR REWARDS
A
ACTION SPACE
SET COLLECTION OF POLICIES
BES
STATE OR OBSERVATION
aEA TER
ACTION
REWARD OR FEEDBACK
PRIOR INFORMATION OF ACTION
TEll (" (*,a))
POLICY FUNCTION
STATE-ACTION VALUE FUNCTION
STATE VALUE FUNCTION, V E. ~-x (Q" ORACLE
T a' ~ oT Q",""
ACHIEVING HIGHEST EXPECTED EPISODIC
ACTION GENERATED BY ORACLE POLICY
VALUE FUNCTION WITH ORACLE POLICY
ma)
TASK WITH DISCRETE STATES AND DISCRETE ACTIONS A COLLECTION OF TASKS T(n,, na))
T(namaa) na)
D(T)
EXECUTION TN A TRAJECTORY [[81, 91, a1, (1), (St, (tt ae rt]]
DATASET RECORDED
h
4
STEP-WISE SUPERVISION LABELS FOR TRAJECTORY he
(*(*)
STATIONARY DISTRIBUTION (SD) OF MDP WITH THE POLICY n²
PARAMETERS MODEL, KEPT
INNER LOOP (ICL)
CACHES MEMORIES, STARTS FROM SCRATCH IN THE INNER LOOP
AGENT WITH ORACLE POLICY x"
Q-LEARNING AGENT
AGENT WITH RANDOM POLICY
MODEL-BASER REINFORCEMENT LEARNING AGENT AGENT O DISTURBED WITH A DECAYING NOISE E
B.2 Details of Procedural Generation
Algorithm eelaborates on the detailed procedural generation of AnyMDP tasks.
Transition Sampling. The generation process commences by randomly sampling ranking of states, denoted as $1,..., Sn,, The initial state set Sy are sampled from {sy < nu). while the absorbing state set (including pitfalls and goals) Sg is randomly sampled. The goals constrained only to 8u,- Subsequently. the average transition matrix is sampled, and based on this matrix, is randomly decomposed into state-action transition matrix. Notice that the state-action transition is constructed utilizing the average state transition kernel and the weights wi (ax(a)), as follows:
exp (w, /)2/
P(s (s,, ax) = Pr(s)lss) Σ, exp ( /)2/
Ergodicity in the average transition of AnyMDP is checked by computing the stationary distribution for all initial states and evaluating the variance among them. In practice, observed that the majority sampled transition matrices adhering to banded transition constraints exhibit ergodicity, with non-ergodic scenarios being exceedingly uncommon.
While the theoretical bound Theorem provides valuable non-triviality guarantees, it remains inherently conservative due the analytical challenges posed by randomly generated task structures.
23

---

Algorithm AnyMDP T Sampler
1: Input: n2, The Returns: 7 (R))
2: Randomly generate an ordered list arranged from low-valued to high-valued states S
[$1, $2,,8,,..,,,, by permute S (I,...,
3: Randomly sample states for reset So C (50. $1...,
4: Randomly sample states triggering termination SE S/So
P Sample banded state transition matrix under uniform random policy n, do 6: Randomly sample P.(-1s)). s.t. banded transition constraints (Equation 1
5: for
7: end for
Ensure the state transition has stationary distribution
8: Calculate P. = (P.) by considering So, Sg; if (V, > E, go back to resample P..
P Sample aate-action transition with the constraint the state transition
9:
for <- n, do
10: vk: E [1, nal. randomly sample scalars wy(ax) E [i - b_, b41. ok ~ Exp(1)
11:
Sample state-action transition P(s;) [ ( by applying Equation (9
12: end for
Sample composite reward function and ensure the value function distribution is expected
13: Randomly sample state-action-dependeri rewards rsa E IR X
Randomly sample state-dependent potential Us E I 15: repeat
Randomly sample state-dependent reward function r. ER", s.t. Vj (s)) 2 rs(s)
17: Set (s, a, s') by composite rewards (Equation (10))
Calculate value function (*)s) based on P, R. So. SB
19: until Ascending value function condition is satisfied (Equation Ensure the optimal trajectory not trivial
20:
Final validation: calculate oracle policy steady state distribution (Pa)
Eg Pa log Po/ log n, >> Hg then accept, else resample the task.
Empirically, we observe that setting b4 S n,,44 b. 2 n,/2. > 1.0e - 3, 17 > 0.5 is enough to consistently yield high-quality Markov chain formulations.
Composite Reward Sampling. To further enhance the quality of the task, rather than naively sampling the reward function R from Xg,,,n,,,, Xn, e independently sample the mean reward matrix J.R and the noises R. This ensures that the reward r(s, 0, 8') ~ NN ER s.) The mean reward matrix "'R.s,,s'' is sampled from composite reward functions, which is defined as follows:
R 1s...e - = rr' r*'(s, a) + ( v v'('))
(10)
We first sample the state-dependent reward rd based on the ranking of states 81,..., 8n, and whether each state corresponds to pitfalls or goals. Subsequently,
sample the disturbance reward function you and state-value function v", which can be interpreted as random state-action costs and random reward shaping terms, respectively.
To ensure non-trivial task formulations, we implement suite of validation checks, including: 1. Ergodicity check the average transition function; 2. Ascending value function check to guarantee monotonic improvement; 3. Minimum threshold check on the normalized entropy H(P.) of the state distribution (SD) under the oracle policy o. If any of these checks fail, we either resample the transition dynamics or readjust the sampling of composite rewards. Details of the sampling process provided in the Appendices.
B.3 Additional Properties of AnyMDP
Computational Efficiency of Task Sampling. Figure 8 illlsstrtte the computational cost of generat- ing AnyMDP tasks for various state space sizes n, E [8, 16, 32, 64, 128). Notably, n, 8 exhibits significantly higher computation times. This is primarily due to the frequent resampling required when final validation check fails. It important to note that the AnyMDP task generation process was executed on single CPU. Given this, the use of readily available parallelization techniques could significantly accelerate task generation.
24

---

Average Time Consumption Task Generation
16
05
Figure 8: Time consumption of AnyMDP task generation on an a Intel(R) Xeon(R) Platinum 8374C CPU.
TD4
Cuttk- Dt
300
FPO
Culk
Figure 9: Performance Tabular Q-Learnng and PPO on AnyMDP tasks of variant state space and action spaces, with respect to episodes and steps.
OmniRL- Online Evalution
OmniRL - Online Evalution
1.0
1.0
Tu:16.5)
fe<32.55
Tux6445)
Norma lized Score
Normalized Score
0.8
Fe(128,
Twi16,55
0.
6r132.5)
0.
16,6
f.1128, 5)
10F
10a
16'
104
107
32
Episodes
Steps
Figure 10: Performance of OmniRL on AnyMDP tasks of variant state space, with respect to steps.
25

---

Relations between solution difficulty and the state-action space. Figure 9 illustrates the perfor- mance Tabular Q-Learning and Proximal Policy Optimization (PPO) on AnyMDP tasks with varying state space and action space sizes. The results indicate that increasing either the state space size (ng) or the action space size (na) enhances the complexity of the task, as evidenced by the need for more training steps to achieve convergence. OmniRL's result on AnyMDP tasks with different state spaces also supports this phenomenon, shown Figure I0] Additionally, an increase in the state space size (n,) leads to higher number of steps per episode.
7116,3
(b) 732,, 5)
764,
(d) Garnet(44.5.2)
(c) FrozenLake non-slippery)
(f) FrozenLake (Slipper)))
(g) Discrete-Pendullmm (g=9.8)
(h) DarkRoom(S x 8, (7, 7)
(i) DarkRoom(10 x 10, [9,9])
Figure 11: Visualization three tasks sampled from AnyMDP, with the number of states n, varying across {16, 32, 64), Garnet(64, 2). and Gymnasium tasks including non-slippery and slippery and 2 tasks sampled from DarkRoom. States are reordered
FrozenLake and Discrete-Penduluue according the SD (pc), ordered from high low. Gray blocks indicate positive state transition kernels. Red and blue blocks mark pitfalls (S and goals (S+). respectively. which trigger episode termination. Green blocks mark Sg. The black line denotes the state value function under the optimal policy (s). Notably, AnyMDP is capable of generating a diverse range of tasks, including those with and without pitfalls and goals. The visualizations demonstrate that tasks generated by AnyMDP can be of comparable quality to those designed by humans. A common principle observed is that higher rewards are often associated with more challenging goals, akin the concept of "high-hanging fruit".
26

---

Table Summarizing the data synthesis strategies different methods under the Decoupled Policy
Distillation (DPD) framework.
DATA SYNTHESIS PIPELINE
BEHAVIOR POLICIES (II)
REFERENCE POLICY
AD
13
AD
15
DPT 16
q. 0. m, t, O
DECOUPLED POLICY DISTILLATION (OURS)
Visualizing discrete MDPs. Following the previous analysis, for any discrete Markov Decision Process (MDP), we can rearrange the states such that the SD decreases monotonically. We then plot the transition kernel P. in this rearranged order. In the visualization, we use varying opacity to represent the elements of P. and different colors to distinguish the initial states So, positieely rewarded terminal states (goals) S and negatively rewarded terminal states Sg. This visualization, shown in Figure 11 enables us to analyze both procedurally generated AnyMDP tasks and human-designed Gymnasium tasks. Several interesting observations can be made:
Higher rewards for higher effort. Both procedurally generated AnyMDP tasks and human- designed Gymnasium tasks exhibit a negative correlation between the SD Pe and the value function V This suggests common principle: states with lower SD probability tend to have higher value functions, akin to the concept of "high hanging fruit".
Banded transition kernel. When ordered by decreasing SD probability, the transition kernels of all Markov chains display the characteristics of banded matrix. This obser- vation further validates the effectiveness of the procedural generation method outlined in Algorithm I
Additional Information on Experiment Settings
C.1 Data Synthesis
Algorithm 2 Data Synthesis Pipeline
Input: T. NNample Collection of behavior policies II, reference policy x
set: (TT)
for [1, . Naample] do
sample: task T
set: f 0, ho = 0.
repeat
sample: behavior policy (0) ~ II
reset: and update 8t
repeat
sample: dt ~ n (b) (a))s) (0) x² (als)
sample:
execute: 0, in T and obtain S141, Tt
set: ht h4-1 + [St, It. a, rı], 4,11 1 aj, t++1
until Episode over
until > T
Set: D(T) = D(T) (hr.4}
end for Return: D(T)
The data synthesis pipeline of OmniRL involves generating diverse trajectories h using variety of behavior policies and creating step-wise labels 1 with an oracle policy σ². This pipeline is detailed in Algorithm
We incorporate least five distinct types of agents:
An agent with the oracle policy (o),
27

---

Table 5: Correpsondance of prompt IDs and the policies represents.
(g)
AGENT TYPE
DESCRIPTION
( 0)
SINGLE-STED GREEDY
o( Y = 0.5)
MYOPIC GREEDY
o Y = 0.93)
SHORT-TERM ORACLE
ORACLE WITH 0.99
MODEL-BASED REINFORCEMENT LEARNER
q
TABULAR Q LEARNER
6
t
KANDOMIZED POLICY (INCLUDING PERTURBED ACTION IN o²)
7
UNK
RESERVED ID
An agent with a randomized policy (t),
A tabular Q-Learning agent (q),
A model-based reinforcement learning agent (m),
An agent with the oracle policy perturbed by a decaying noise (of).
With these notations, Table can but also the previous imitation meta-training titatttrining-basee based ICRL methods, including AD, AD and DPT, as shown in Table Notably, the synthesis pipeline of OmniRL is most similar to that of DPT. However, there are key differences: OmniRL employs a more diverse set of behavior policies and incorporates step-wise supervision (SS).
can be used to represent not only the data synthesis pipeline of OmniRL
For the prior information g. we assign eight different IDs to the actions with y E [0. which originated from 8 types of different agents, as shown in Table Specifically, we exclude the actions generated by the seven types of different agents and reserve = 7. This reserved ID is used to replace the action ID approximately 15% of the time steps with the data synthesis pipeline of Algorithm 2
Terminal or Truncation
Episode k
E Eissode k+1
94-1
9 1
St
Ir.: Unk
a2
n2
Ft.
0
Scp+1
1ep+1
Tpt1
Figure 12: A sketch of the tokens used denote the intervention between two episodes. An extra action token is introduced, which is distinct from the normal actions (at, |A| na).
Addressing terminal states: The presence of terminal and truncation states necessitates special handling in reinforcement learning. In OmniRL, avoid explicitly adding terminal truncation token to the sequence. Instead, we encode terminal and truncation states by introducing an additional action which is maintained as distinct from the standard action space A. Additionally, assign reward of and set the prior information = 7 for these special steps, as depicted in Figure] 12
C.2 Meta-Treas Training Details
Before injection into causal models, the states (st) and actions (o.) encoded using embedding layers with a hidden size of 512. The rewards (rt) are treated as continuous features encoded by 1 x 512 linear layer. The sequence model has hidden size of 512, inner hidden size of 1024, hidden ratio of 2, and block number of 18 for RWKV-7. The model has approximately 43.6M total parameters (42.9M in RWKV-7 blocks), shown in T We employ the open-source implementation of Mlush--inear-attniii
Model structures:
Algorithm Aigorithmooutling Joutines the detailed process of the meta meta-traininn training procedure. Notably, we perform the backward pass segment- wise and accumulate the gradients. The gradients are not applied until the end of sequence. We utilize constant segment length T Tí 2K, which results in six backward passes for = 12K before applying the gradient.
Meta-training:
https///ithuib.om//tttooo/gggg/e linear-ateention
28

---

Table 6: The parameter settings of the sequence model. Note that for different models, the relationship between head dimension, head number, and hidden size varies. We follow the settings used in flash- linear-attention.
GDN
GSA
Mamba2
RWKV7
Block nums
18
18
18
18
Hidden size
512
512
512
512
Inner hidden size
x 512
512
x 512
512
Head nums
8
8
8
8
Head dim
64
128
64
Parameters
46.3M
42.9M
31.6M
42.9M
Algorithm 3 Meta-Training Process
Input: D(Tira), D(Test)
for epochs from 1 to maximum epochs do for hr, lr E D(Trra) do
set: segments K = T//Tseg gradients = 0, initial memory do
for k: E [0, K) do
forward: update 0k-1 + Φb based on Equation (4), Ox-1, TT:: and IT. Tr-1 backward: calculating 9k V VI WeCr by stopping gradient of Ø
accumulate gradient: 9 9 + 9k
end for
apply gradient: g to update 0
end for
validate: averaging L, and L (T
end for
Table provides an overview of the primary datasets used this study. For the Large dataset, the state space size n, is uniformly sampled from the range [16, 128] to ensure robustness across varying state spaces. To evaluate the extrapolation capability of the model trained on Larges we conducted a validation test with context length of million steps and observed that the loss began to gradually increase beyond 80 steps. Building upon this observation, we incorporated a post-training stage for long sequences with context length of 512K, the dataset is denoted Donn-
Algorithm 4 Evaluation Process
Input: Tist. collection of demonstration trajectories Ho {ho}. set: Seral = (f
for 7 E Tiat do
set: aaerraverage episodic reward of -verrge episodic reward of t
set: Seval [
repeat
retrieving: ho from Ho according to reset: and obtain $1, R repeat
sample: 0e ~ Pe with Equation 44
execute: 04 in 7 and obtain 8141, Tt
set: hy h4-1 ① [st. It, at, P2] with 9: ="Unk"
set: R + R 1rtt,, + ++1
until Episode
calculate: normalized performance Sceal +- Savel +
R- R Rmin
Rmaz Ruan
set: N Nepisodes f- N Nepissdes +
until Nepisodes > Nmar
Record: Seeal + Seval Seval
end for
Return: Seeaa
29

---

Table 7: Details of the meta-training dataset
DATASET
DESCRIPTION
TIME STEPS
n1 16, na =5 Tiral = [D(Ttra)] 128K. SequenceLength 8K
Dsmall
1B
n4 E [16, 128 The 512K, SequenceLength 12K
D Large
6
Timal (D(Tira)
Rh E [16, 128], Ra 5
D Long
6B
Teral
(D(rea))
12K, Sequence 512K
1 if reach goal
reach goal reach hole otherwise
reward
1
if reach hole otherwise
reward
1,
0,
-0.05,
(a) FrozenLake-V (slippery)
(b) FrozzenLake-v((nd) slippery)
1
if reach goal if reach cliff otherwise
reward
1,
reward
reward max
+ +.1,, 001
0.03,
30
(c) CliffWalking-vO
(d) Pendulum-vl
1,
reach goal
0.08,
if distance to goal decrease
agent reward
-0.12,
if distance to goal increase
-0.04,
if still
if finish
shared reward = Σ agent reward,
(e) Switch
Figure 13: Reward shaping
C33
Evaluation Details
As shown in Algorithm 4 since the episode length and baseline average episodic reward vary significantly across different tasks, normalize the episodic reward using the oracle policy (o) and the uniform random policy (t). This normalization represents the percentage achieved. For AnyMDP, the evaluation averages the performances over 64 variant unseen tasks. For Gymnasium tasks, the evaluation conducted by averaging the results over 3 runs on the same task.
oracle performance
By default, the normalized performance geval is averaged across tasks with identical Nopisodes. The deviation is estimated using the 95% confidence interval of the mean.
For Tabular Q-Learning and PPO, we conduct 5 episodes of testing after every 100 episodes of training for each run. We evaluate performance based solely these test episodes. In contrast, for ICRL, we do not differentiate between training and testing phases.
When performing online inference with OmniRL, we do not employ any additional exploration strategies. Instead, we maintain softmax sampling temperature of 0.5. For offline learning, where demonstrations from the teacher are encoded, we the original prior information 9e without
30

---

modification. In contrast, for online learning, where actions are generated by the agent itself, we set gi 7 (Unk).
We apply some reward shaping to Gymnasium tasks as shown Figure 13 OmniRL. supports ns S 128 and na 5. For environments with na < 5, we find directly setting mod na is enough, which also demonstrates the generalizability of OmniRL to variant action spaces.
D Additional Empirical Results
D.1 AnyMDP as long-context benchmark for procedural memory
Architectures Meta Training Curves
Architectures Stepwise Validation Loss
2.0
GDN
GDN
GSA
GSA
Mamba2
1.2
Mamba2
150
RWKV7
1.0
RWKV7
Loss
125
1.00
0.2
10000
20000
30000
40000
50000
201
Training Iteration
Steps in Context (t)
Figure 14: Comparison of meta-training dynamics across AnyMDP dataset of 6B time steps and step-wise loss (lower is better) on the validation ( [16, 128), n. 5)) for different linear-attention models.
Figure Figure 14 reports the performance of different linear-attention models throughout AnyMDP training and on the held-out AnyMDP evaluation data (T(n, [16, 128), n2 5)), including Gated Delta Net (GDN), Gated Slot Attention (GSA), Mamba2, and RWKV-7. Error decreases polynomially with context length up to 20k tokens, yielding an almost straight line when plotted against the logarithm of context length. Consequently, AnyMDP serves as a long-context benchmark for procedural memory, whereas prior benchmarks such as Needle-in- Hayhe-Haa primarily probe episodic memory.
D.2 OmniRL achieves automatic trade-off' between exploration and exploitation
Entropy of Decision
(16, 5))
(32,
12-
DUA,(64, 5))
1.0
Dlr,,128, 5))
Entropy
0.8
0.6
0.4
0.2
Steps in Context (t)
Figure 15: The position-wis wise entropy when validating RWKV-7 on different datasets.
31

---

Previous studies have noted that in-context reinforcement learning (ICRL) can automatically balance exploration and exploitation. This phenomenon has been theoretically linked to posterior sampling. In Figure T5
we illustrate the entropy of the decision-making process as function of steps within the context. When compared Section [.T] we observe that the decrease in loss (L,) is primarily driven by the reduction in the entropy of the policy. Specifically, the agent initially assigns equal probabilities to all actions, reflecting an exploratory phase. As more contextual information accumulates, the agent gradually converges its choices, thereby transitioning towards exploitation. This empirical finding suggests that imitating an optimal policy (oracle) is sufficient to achieve an automatic balance between exploration and exploitation.
D.3
Additional Evaluation on Gymnasium
Figure T6 and Figure 17 demonstrate OmniRL's online-RL, offline-RL, and imitation learning capabilities toward diverse unseen tasks.
Online Prozenlake vlislopery)
Online RL Fropeniaie viinot stepery)
Onlne Cittlaking 90
Soe oue CONARUA capel
Onine Dark6oom/13111
Online RL Pendufum vilg Invol
Ondine RL Pendulum viig Smis?
On the RL Penou vi 3.8ml
Nut: agene Onlin RL Switon2
Ep
Esoado
Gracle
OmniRt
TOL
PPO
Agent
Agent
Figure 16: Selected online evaluation results for TQL-UCB, PPO, and OmniRL across Gymnasium environments. Notably, despite never having been exposed to these environments during training, OmniRL demonstrates strong adaptability by achieving competitive performance on most tasks with high sample efficiency.
Offinne RL Inseentaie-vllsligeeer
Offine-RL CilfWalkingvd
Mult--gget OfflineRL Switch2
NVV
WT
1000
Epiiddes
Ephcder
Invitationn Learning
OffiineRR (Expert)
Offiin-RL (Random)
Agent
Agent
Teacher Demprattoetonn
Figure 17: Selected offline evaluation results for OmniRL across Gymnasium environments, demon- strating the model's offline-Ro and imitation learning capabilities toward unseen tasks.
D.4 Memory states in ICRL implicitly encode the task structure
ICRL with linear attention captures all the information required to solve the environment in its memories (6:). We perform comprehensive t-SNE analysis to examine how these memories transform across different tasks during Online-RL RL evaluation. As shown in Figure I8 the clustering patterns confirm the distinct task distributions Gym, Darkroom, and AnyMDP. Notably, Darkroom and Gym clusters are predominantly located in the top-left region, while AnyMDP occupies broader spatial area, reflecting greater diversity. This spatial differentiation emphasizes AnyMDP's unique characteristics and highlights OmniRL's strong generalization ability across diverse tasks.
32

---

Online-RL tSNE
Darkroom 6x6
Cliff
Darkroom 8x8
Darkroom 10x10
Gym Penclulum
Gym Lake
-25
-50
-75
-75
50
75
100
Figure 18: t-SNE visualization of the trajectory of the memory states (ø.) of OmniRL in online- evaluation with variant environments. The unboxed points correspond to T(16,5) 5). Trajectories originating from the same environment represented in the same color.
Table Comparison of the performance (success rate) of LLM and other methods discussed in the paper within the ake environment.
METHODS
NON-SLIPPERY
SLIPPERY
RANDOM
1.6%
1.3%
LLM W/ STATE ONLY
< 2%
< 2%
LLM W/ GLOBAL MAP
100%
5.6%
LLM W/ GLOBAL MAP & HINT
17%
OMNIRL W/ STATE ONLY
100%
60%
ORACLE
100%
75%
D.5 Comparison with Pre-trained LLMSs
We also investigate whether well weel-pnetr LLM can naturally solve the decision makinn tasks investigated in this paper. To cireumvent the lack of common sense in AnyMDP tasks, primarily conducted tests in the Frozemlakee task with DeepSeek-RI 78 in two modes:
1. Similar to the evaluation of standard ICRL, we do not provide the agent with the map. Instead, we report only the state ID and reward the agent. The initial prompts used to initiate the evaluation shown in Figure 19
33

---

You are playing the Frozen Lake game. The environment is a 4x4 grid where you need to maximize the success rate by reaching the goal (+1) without falling into holes (-1). You move in four directions: left, down, right, and up (represented as 0, 1, 2, respectively). You will receive the current state and need provide the optimal action based on your learning. When asked for the optimal action, your response must be an integer ranging from 0 to 3, and other context permitted. There are two kind of request types:
Linteger: the integer is the current state, and you need to provide the optimal action. 2.list: The list contains one or more tuples, where each tuple contains the last state, action taken, reward received, and next state. To save time, you don't 't need to respond when receiving a list. You will play the game multiple times. game ends when the reward is or try get higher success rate.
Note: am asking you to play this game, not to find a coding solution or method.
You will be provided with a conversation history. The latest prompt is the current state, and others are the list of sequential environment feedback history tuple type. Each tuple contains four values, the first one is state, the second one is action, the third one is reward and the fourth one is next state.
Your response must be an integer from 0 to 3 during the entire chat. If you find the last state is equal to the next state, your policy in the last state can't 'n be this action. If you find the reward in the tuple is -I, your policy in the last state can't an't be this action. You need to get to the goal as soon as possible.
Figure 19: Prompts for LLM to initialize the Lake4 x 4 (Slippery) task without a global map
There is a game with the following basic description and rules:
Frozen Lake involves crossing a frozen lake from the start the goal without falling into any holes by walking over the frozen lake. The player may not always move in the intended direction due the slippery nature of the frozen lake.
The game starts with the player at location [0,0] of the frozen lake grid world, with the goal located at the far extent of the world, for example, [3,3] for the 4x4 environment. Holes in the ice are distributed in set locations when using a pre-determined map or in random locations when a random map is generated.
The player makes moves until they reach the goal or fall into hole. The lake slippery, so the player may move perpendicular the intended direction sometimes. If the intended direction is to the left, the actual move may be to the left, up, or down, with the corresponding probability distribution: P(move left) = 1/3, P(move up) = 1/3, P(move down) = 1/3. If the intended direction is to the right, the actual move may be to the right, up. or down, with the corresponding probability distribution: P(move right) = 1/3, P(move up) = 1/3, P(move down) 1/3. If the intended direction is up, the actual move may be up. left, or right, with the corresponding probability distribution: P(move up) = 1/3, P(move left) = 1/3, P(move right) 1/3. If the intended direction is down, the actual move may be down, left, right, with the corresponding probability distribution: P(move down) = 1/3, P(move left) = 1/3, P(move right)
1/3. You given 4x4 map where: S represents the start.
F represents the frozen surface that can be walked on.
H represents a hole; falling into it will return the player to the start.
G represents the goal.
The map is as follows: The first row from left right "SFFF" The second row from left to right is "FHFH". The third row from left right "FFFH". The fourth row from left to right is "HFFG". Please determine the optimal policy that maximizes the success rate of safely reaching the goal from the start. The optimal policy is the intended direction at each map location, where actions 0, 1, 2, and represent moving left, down, right, and up, respectively. Note: You are not required to write code to solve this problem; instead, directly provide the optimal policy.
Figure 20: Prompts for LLM to solve Lake4 x 4 (Slippery) with global map
34

---

In this particular environment, you should focus more on strategies to avoid encountering pitfalls rather than concentrating reaching the goal. A reevaluation of the decision-making process in this context may yield more effective outcomes.
Figure 21: Additional guide (hint) for LLM to solve Lake4 x (Slippery) with global map
The optimal policy for each state (row, column) is follows, where each cell lists the intended direction (action) that maximizes the probability of reaching the goal:
**(0,0)**; Down (1) **011)**; Down (I) **(0,2)**8 Down (1) **(0,3)** Up (3)
**(1,0)*#: Down (1) *0(1))) Down (1) **(2,0)**; Right (2) - **(2,1)**; Down (1)-
**(2,2)**: Down (1)- 13,111)**; Right (2) ": Right (2)
**Exllnattoon*** **: *Toop row (0,0 to 0,3)**: Prioritize moving down to avoid the risky rightward path toward holes. M*dddee rows (1,0 and 1,2y*s; Move down progress toward safer lower rows Lower rows (2,0 to 3,2)**; Navigate rightward toward the goal, avoiding holes. For example, (3,1) and (3,2) move right reach the goal directly. **(0,))**; Move up (intended) to bias movement leftward (safer) and avoid falling into the hole at (1,3).
Figure 22: Responses from DeepSeek-RI for Lake4 x 4 (Slippery) with global map
The optimal policy, revised to prioritize avoiding holes, is as follows:
**(0,0)**; Down (1) * *Avoids (0,1)'s risks by favoring transitions to (1,0) or (0.1) without direct hole exposure.*
**(0,1)**; Right (2) *Moves toward (0,2) or (I,1) (hole). While risky, all actions here have 1/3
hole risk; right progresses toward safer (0,2).*
**(0,2)**: Right (2) "Directs to (0,3) instead of down to (1,2), which is high-risk.*
***003)**#: Up (3) "Biases movement left (to (0,2)) and avoids the hole (1,3).*
***11,0)** Left (0) *Avoids 1/3 chance of moving right into (1,1) (hole). Transitions to (0,0) or (2.0) safely.*
**112))** Right (2) *Despite 1/3 hole risk, it allows progress to (2,2) or (0,2). Left (action 0)
has equal risk but less progress.* *
***2,0)**; Up (3) *Avoids 1/3 chance of moving down into (3,0) (hole). Routes to (1,0) or (2.1).*
=**12,)**; Down (1) *Safely moves to (3.1) (frozen), progressing toward the goal.*
=**(2.2***; Down (1) *Reaches (3,2), safe step toward the goal.* *
***3,1)**; Right (2) *Directly progresses to (3.2). then (3,3) (goal).* *
**(3,2)**; Right (2) *Final step the goal (3.3).*
### Key Adjustments from Original Policy: 1. **110)***; Changed to **left** (0) to avoid 1/3
chance of falling into (1,1). 2. *1,2))*** Changed to **rgght** (2) instead of down, reducing
hole risk from 2/3 to 1/3. 3. ** Changed to **u * (3) to avoid 1/3 chance of entering
(3,0). 4. (0,1))** **igght** (2) minimizes repeated exposure to (1.1) compared to down.
### Path Safety Focus: Avoids actions with > 1/3 hole probability (e.g.. down from (1,2)). Prioritizes paths with lower immediate risk, even if longer (c.g., routing through (2,0) - (2.1) (3.1) goal).
Figure 23; Responses from DeepSeek-RI for Lake4 (Slippery) with global map and hint
Oracle - Non- Slippery
Oracle Slippery
LLM- Slippery
LLM (w/ hint)
Slippery
Figure 24: Comparison of the solutions of different methods in the FrozenLake environment. In the LLM w/ hint condition, we provide additional guidance to the agent, instructing prioritize avoiding holes over reaching the goal.
35

---

2. We initially provide the global map to the DeepSeek RI and then commence the interac- tion. In this mode, the LLM can leverage the global map to make decisions. The prompts are shown in Figure 20 and Figure 21
As shown in Table 8 and Figure (results excecuted by following the responses of DeepSeek-RI (version 2025/03) in Figure 22 and Figure 23). LLM agents only able to solve the FrozenLake (non-slippery) environment when provided with a global map. Without access to a global map, conducted extensive interactions between LLM agents and the environment, running up 500 episodes (100, 000 steps). Despite these efforts, the agents failed to solve even the non-slippery variant of the task, achieving scores that were comparable to those of random policy.
Even with the aid of global map, the performance of LLM agents on the FrozenLake (non-slippery) environment remains notably poor. To improve their performance, we introduced additional hints suggesting that a better solution should prioritize avoiding holes over reaching the goal. However, this intervention only marginally improved the agents' performance, raising it from 5.6% to 17%. This level of performance still significantly lower than that of the Oracle and OmniRL agents.
Notably, increasing the length of the chain of thought can potentially enhance performance when global map is available, but it has minimal impact on performance when only the current state is considered. The former scenario emphasizes System decision decision-making, which is characterized by based and analytical thinking. contrast, the latter scenario highlights the in-context adaptation which relies continual external feedback and represents rapid, intuitive decision-making 779 80]. We argue that future research should place greater emphasis on the latter approach.
of System 1 decision-makee
36


Reviews:
[{"review_id": "iOT69qt4GA", "reviewer": "Reviewer_1krJ", "summary": "* Proposes to procedurally generate tabular MDPs in order to study various properties of ICRL\n* Results demonstrate that with a sufficiently large scale of tasks, the proposed model can generalize to tasks unseen in training set", "questions": "* Which findings do the authors believe would(n't) scale to more complex domains? \n* 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", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 3, "originality": 4, "strengths_and_weaknesses": "Strengths\n* The paper is well-written and clearly presented; in fact, exceptionally so\n* The details and analysis in the main body are thorough but not overwhelming\n* OmniRL method produces very strong performance against baselines\n* Useful insights---especially that \"generalization of ICRL can be at odds with its zero-shot or few-shot performance\"\n\nWeaknesses\n* Only applicable to discrete state and action spaces (acknowledged as limitation)\n* 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", "was_revised": true, "final_justification": "I stand by my initial review and therefore maintain my current score.", "rebuttal": "# On the Value of Discrete Space MDPs\nWe appreciate the reviewer's positive feedback on the value of the seemingly toyishness of the discrete MDPs.  At this stage, partially pursuing the task complexity risks undermining the core advantages of procedurally generated discrete MDPs. Below, we explain the challenges of preserving key properties in continuous domains and contrast our approach with contemporary benchmarks.\n\n## Challenges of Continuous MDPs in Preserving Structural Integrity\nExtending procedurally generated MDPs to continuous spaces while maintaining low structural bias and sufficient diversity poses significant technical hurdles. At the current stage, efforts moving to continuous MDPs risk typically **prioritizing \"absolute task complexity\" over qualitative diversity**, introducing biases or homogenizing task structures—contrary to our framework's goals.\n\n\n## Inherent Long-Term Dependency and Qualitative Diversity\nExisting continuous MDP benchmarks rely on continuously randomized hyperparameters to simulate \"infinite tasks,\" prioritizing few-shot learning over long-horizon in-context learning (ICL). Our analysis reveals two insights:\n- Task Diversity ≠ Numerical Scale: True diversity depends on qualitative variation in task structures, not just quantity.\n- ICL efficacy depends on disparity between tasks in the training set, not individual task complexity. Existing continuous benchmarks tend to produce homogeneous task structures, limiting models to superficial pattern recognition via in-weight learning (IWL) rather than sustained reasoning. In contrast, discrete MDPs in AnyMDP are designed with procedural complexity and temporal coherence, enabling rigorous evaluation of a model’s ability to generalize across structurally distinct tasks and sustain reasoning over extended sequences.\n\n## Compatibility with Existing LLMs\nBeyond being a novel benchmark, AnyMDP bridges a critical gap in assessing experience-driven, holistic reasoning capabilities of pre-trained LLMs. Existing benchmarks predominantly measure declarative memory and one-shot circuits, while AnyMDP emphasizes procedural memory and experience-driven learning.\n\n## Low-cost Access to Oracle Policy\nThis property significantly enhances the efficiency and scalability of ICRL training, making AnyMDP a scalable benchmark.\n\n## Empirical analysis across state-of-the-art architectures\nWe also augment the empirical analysis with the training curves and ICL comparisons across state-of-the-art architectures, including Mamba2, RWKV-7, Gate Delta Net, and other competitive models. This will demonstrate AnyMDP’s utility as a rigorous benchmark for advancing long-context modeling capabilities.\n\n- Table 1. Validation Loss of Different Models on AnyMDP Datasets\n| CL \t  | 1  \t   | 100\t    | 1000    | 10000   | 30000 |\n|-------------|------------|------------|------------|------------|------------|\n|  Mamba2   | 1.417 ± 0.035 | 1.272 ± 0.051 | 0.822 ± 0.064 | 0.436 ± 0.053 | 0.294 ± 0.042 |\n|  GSA   | 1.406 ± 0.034 | 1.235 ± 0.053 | 0.765 ± 0.066 | 0.342 ± 0.050 | 0.235 ± 0.042 |\n|  Gate DeltaNet   | 1.356 ± 0.035 | 1.130 ± 0.051 | 0.599 ± 0.055 | 0.267 ± 0.049 | 0.173 ± 0.042 |\n|  RWKV-7   | 1.371 ± 0.034 | 1.151 ± 0.052 | 0.627 ± 0.055 | 0.269 ± 0.042 | 0.175 ± 0.029 |\n\n\nOverall, in our research on ICL, the benefits brought by the discrete MDPs far outweigh the drawbacks of its limited problem space. We have also started to explore ICL in the continuous space.\n\n# Question about more complex domains and continuous settings\n- For future work in continuous settings, there are two directions. First, we will focus on specific continuous problems and prioritize in-domain adaptation. Second, we will continue exploring how to generate tasks with high diversity and low structural bias in continuous spaces. Undoubtedly, the first direction holds more practical value, while the second poses greater theoretical challenges.\n- If the second direction remains unaddressed, cross-domain generalization in continuous spaces will be difficult to achieve, but in-domain adaptation is highly feasible. For instance, in frame prediction tasks, generalization manifests as adaptation to environments with distinct features and spatial structures. Similarly, in robot whole-body locomotion control, it manifests in adapting to different robot configurations (size, shape, weight).\n- Another more complex domain is **Multi-agent** problem. As agents need to collaborate and their performance cannot be guaranteed to be consistent, a new challenge arises: in-context adaptation across other agents' actions. If a single agent can observe and learn other agents' actions online, while adjusting its own real-time policy, it could potentially alleviate the fragility of existing multi-agent systems. We believe this scenario will uncover more unknown findings."}, {"review_id": "FznNeNE4Rf", "reviewer": "Reviewer_nuE2", "summary": "The 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.", "questions": "1. 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?", "limitations": "Yes.", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 1, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Disclaimer: 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.", "was_revised": true, "final_justification": "I would like to maintain my score as my concerns remain unaddressed.", "rebuttal": "First and foremost, we sincerely thank the reviewer for identifying the typos and providing valuable suggestions on the writing. We will revise accordingly to further improve the paper’s readability, especially on the explanation of ICL and Casual Blocks. In the following, we will elaborate in detail on the weaknesses and questions raised by the reviewer.\n\n# Garnet MDP and other related works\nWe appreciate the reviewer’s reference to Garnet MDP, an efficient MDP sampling methodology. Key parallels and distinctions include:\n- Similarities: \n     - Both prioritize MDP structural integrity\n     - Both focus on stopping \"fast mixing\" MDPs. (In our case, theorem 1 also guarantees this point)\n- Differences: \n    - Garnet MDP controls the transitions under an optimal policy, whereas AnyMDP controls the transitions under the uniformly random policy. In our case, we guarantees that a uniformly random policy induces a strictly positive probability of reaching every state from the initial state (ergodic). However, AnyMDP does not guarantee that the optimal policy is strictly ergodic---a better property aligns with most existing benchmarks. Nonetheless, the banded transition kernel and reward structure ensure that the optimal policy requires transitions across at least a substantial number of states (nearly ergodic).\n    - Garnet MDP is tailored for infinite-horizon MDPs, while AnyMDP supports both infinite- and finite-horizon settings (aligning with common RL benchmarks).\n\nThe paper, Natural actor–critic algorithms, presents four new natural actor–critic reinforcement learning algorithms, provides their convergence proofs, and improves upon prior work by being the first to offer such proofs and propose fully incremental algorithms for these methods. MICo, the learnt representations in MDP problems, demonstrated a gratifying performance improvement for DQN and SAC deep reinforcement learning agents. A Kernel Perspective on Behavioural Metrics defines a new metric equivalent to the reduced MICo distance, establishes new theoretical results, and demonstrates strong empirical effectiveness in deep reinforcement learning. \n\nWhile our work also focuses on MDP problems, the emphasis differs. The AnyMDP we propose aims to generate tasks with both diversity and quality. Using datasets generated by AnyMDP, we train a generalizable OmniRL model to investigate the in-context learning (ICL) mechanism of large models. \n\nWe appreciate the reviewer's suggestion. We will include a proper comparison with prior work on procedurally generating MDPs in the related works section and cite these papers.\n\n# AD, AD𝜖, DPT, and our AnyMDP\n## Explanation of Figure 1\nFigure 1 illustrates the differences between AD, DPT, and our AnyMDP. Blue arrows (Behavior) depict the trajectory fed to the model; yellow arrows (Reference) provide the corresponding supervisory signal. (Left) For AD, the reference action is identical to the behavioral action. (Middle) For a sequence of behaviors, DPT provides a reference action only at the end of the sequence. (Right) For our method, we provide a reference action at each step of the behavioral sequence, which in turn enhances the strength of the supervisory signal and improves training efficiency. The key contribution of Stepwise Supervision is that it tolerates step-level inconsistencies between behavior and reference, aiming at solving the issue of distribution shift in imitation learning.\n## Explanation of prior information\nThe reviewers raised questions regarding the inclusion of prior information in AD, AD𝜖, and DPT. To clarify, prior information introduces additional details about different behavioral policies—specifically, policy tags. Since the AD policy remains unchanged during generation, adding prior information has no theoretical effect. For AD𝜖, policy tags can be incorporated based on the noise level. Additionally, as DPT involves multiple policies, the inclusion of policy tags is feasible. We believe that adding policy tags is also highly likely to enhance performance for AD𝜖 and DPT.  \n\nIndeed, compared with DPT, our primary contribution lies in proposing the Stepwise Supervision method. As for policy tags, they serve as an additional benefit—analogous to high-quality prompts in LLMs—by informing the model of the actual level/type of the current action, thereby facilitating strategy distillation during training.\n\nThe prior information is accessible. During the training phase, the classification of prior information is manually designed, and we randomly mask some tags as \"Unknown\". In online testing, we allow the model to assess current performance with all tags set to \"Unknown\".\n\n# Experiment settings in Section 4.3\nSection 4.3 aims to examine how the number of tasks influences generalization. While all tasks in our experiments are drawn from AnyMDP, we ensured sufficient diversity among them—and our results support this observation. Specifically, models trained on fewer tasks showed limited performance on other AnyMDP tasks, whereas satisfactory generalization to additional AnyMDP tasks was only achieved when the number of training tasks was sufficiently large. This aligns with the key point we intended to convey: a larger number of tasks can foster generalization.\n\nFurthermore, our experimental observations suggest that if a model fails to generalize within AnyMDP, it is unlikely to generalize to Gym tasks either. We believe that out-of-distribution generalization is inherently more challenging than in-distribution generalization."}, {"review_id": "vTtQDmWCpH", "reviewer": "Reviewer_WkeX", "summary": "The 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.", "questions": "1. 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?", "limitations": "An 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.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths\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.", "was_revised": true, "final_justification": "The authors have satisfactorily answered my questions. I’m confident the responses can be easily incorporated in a next version of the paper, leading me to raise my score.", "rebuttal": "First and foremost, we sincerely thank the reviewer for identifying the typos. We will revise accordingly to further improve the paper’s readability. In the following, we will elaborate in detail on the weaknesses and questions raised by the reviewer.\n\n# On the Value of Discrete Space MDPs\nWe appreciate the reviewer's interest in extending the task to continuous domains. Before discussing how to extend to continuous spaces, we would first like to engage in a deeper discussion with the reviewers about why we chose discrete MDPs at this stage, as well as the benefits this choice brings.\n\nAt this stage, partially pursuing the task complexity risks undermining the core advantages of procedurally generated discrete MDPs. Below, we explain the challenges of preserving key properties in continuous domains and contrast our approach with contemporary benchmarks.\n\n## Challenges of Continuous MDPs in Preserving Structural Integrity\nExtending procedurally generated MDPs to continuous spaces while maintaining low structural bias and sufficient diversity poses significant technical hurdles. At the current stage, efforts moving to continuous MDPs risk typically **prioritizing \"absolute task complexity\" over qualitative diversity**, introducing biases or homogenizing task structures—contrary to our framework's goals.\n\n\n## Inherent Long-Term Dependency and Qualitative Diversity\nExisting continuous MDP benchmarks rely on continuously randomized hyperparameters to simulate \"infinite tasks,\" prioritizing few-shot learning over long-horizon in-context learning (ICL). Our analysis reveals two insights:\n- Task Diversity ≠ Numerical Scale: True diversity depends on qualitative variation in task structures, not just quantity.\n- ICL efficacy depends on disparity between tasks in the training set, not individual task complexity. Existing continuous benchmarks tend to produce homogeneous task structures, limiting models to superficial pattern recognition via in-weight learning (IWL) rather than sustained reasoning. In contrast, discrete MDPs in AnyMDP are designed with procedural complexity and temporal coherence, enabling rigorous evaluation of a model’s ability to generalize across structurally distinct tasks and sustain reasoning over extended sequences.\n\n## Compatibility with Existing LLMs\nBeyond being a novel benchmark, AnyMDP bridges a critical gap in assessing experience-driven, holistic reasoning capabilities of pre-trained LLMs. Existing benchmarks predominantly measure declarative memory and one-shot circuits, while AnyMDP emphasizes procedural memory and experience-driven learning.\n\n## Low-cost Access to Oracle Policy\nThis property significantly enhances the efficiency and scalability of ICRL training, making AnyMDP a scalable benchmark.\n\n## Empirical analysis across state-of-the-art architectures\nWe also augment the empirical analysis with the training curves and ICL comparisons across state-of-the-art architectures, including Mamba2, RWKV-7, Gate Delta Net, and other competitive models. This will demonstrate AnyMDP’s utility as a rigorous benchmark for advancing long-context modeling capabilities.\n\n- Table 1. Validation Loss of Different Models on AnyMDP Datasets\n| CL \t  | 1  \t   | 100\t    | 1000    | 10000   | 30000 |\n|-------------|------------|------------|------------|------------|------------|\n|  Mamba2   | 1.417 ± 0.035 | 1.272 ± 0.051 | 0.822 ± 0.064 | 0.436 ± 0.053 | 0.294 ± 0.042 |\n|  GSA   | 1.406 ± 0.034 | 1.235 ± 0.053 | 0.765 ± 0.066 | 0.342 ± 0.050 | 0.235 ± 0.042 |\n|  Gate DeltaNet   | 1.356 ± 0.035 | 1.130 ± 0.051 | 0.599 ± 0.055 | 0.267 ± 0.049 | 0.173 ± 0.042 |\n|  RWKV-7   | 1.371 ± 0.034 | 1.151 ± 0.052 | 0.627 ± 0.055 | 0.269 ± 0.042 | 0.175 ± 0.029 |\n\n\nOverall, in our research on ICL, the benefits brought by the discrete MDPs far outweigh the drawbacks of its limited problem space. We have also started to explore ICL in the continuous space. \n\nFor future work in continuous settings, there are two directions. First, we will focus on specific continuous problems and prioritize in-domain adaptation. For instance, in frame prediction tasks, generalization manifests as adaptation to environments with distinct features and spatial structures. Similarly, in robot whole-body locomotion control, it manifests in adapting to different robot configurations (size, shape, weight). Second, we will continue exploring how to generate tasks with high diversity and low structural bias in continuous spaces. Undoubtedly, the first direction holds more practical value, while the second poses greater theoretical challenges. \n\n# Explanation of AnyMDP\n## How to generate high-quality MDPs\n(To answer the reviewer's Weakness 1 and Question 1.)\n\nThe sampling of transition matrix and reward matrix in AnyMDP is not uniformly random, to avoid generating trivial MDP problems (e.g., MDPs with optimal episode length=1). Uniform randomization would result in generating a majority of trivial Markov Decision Processes (MDPs), such as those with optimal episode lengths of 1 or MDPs where most states are irrelevant to the optimal solution. The policy only needs to remember which short path has the best reward and will not learn the trade-off of risk assessment. Eventually, the environment will be more inclined towards the bandit problem. In most complex RL problems, delayed reward is quite common. Usually, a large positive reward is only obtained when the goal is achieved. Sometimes, there may even be a stage of deduction before getting the positive reward.\n\nErgodicity is ensured by Theorem 1 in the appendix, which guarantees that a uniformly random policy induces a strictly positive probability of reaching every state from the initial state. At the same time, AnyMDP does not guarantee that the optimal policy is strictly ergodic—a trait absent in most existing benchmarks. However, the banded transition kernel and reward structure ensure that the optimal policy requires transitions across a substantial number of states.\n\nAs we have presented above, AnyMDP can serve as a Long-Term Dependency Benchmark. For this reason, we must limit the step length—meaning that to reach high-value targets, the path cannot be short. As the number of states increases, the path to achieving a high score grows correspondingly longer. Thus, \"Under a uniform random policy, the probability of reaching high-valued states remains greater than 0 but decreases at least exponentially with respect to $n_s$.\"\n\n## Behavior policies\n(To answer the reviewer's Weakness 3.)\n\nOn page 20, we present Table 5 about \"Correspondence of prompt IDs and the policies it represents\". We thank the reviewer for pointing out the lack of clarity in this part, and will subsequently include them in the paper.\n- For IDs 0 to 3, the policies are value iteration policies with different gamma factors, corresponding to short-term to long-term oracles ($\\gamma$=0.5, 0.93, 0.994). \n- The model-based reinforcement learner comprises two distinct components: an environment model and a reward model, both conditioned on the current interaction history. A Bellman update is applied at every step to compute the mean value function, which is subsequently augmented with Optimistic Thompson Sampling to balance exploration and exploitation.\n-  Tabular Q-learner with UCB is a model-free method that maintains a Q-table of state-action values. These values are iteratively refined via value iteration and adjusted by an upper-confidence-bound bonus derived from rewards and successor-state information observed during interaction.\n\nRegarding the exact impact of the diversity of behavior policies in training data, DPT has demonstrated its benefit on model performance. Indeed, compared with DPT, our primary contribution lies in proposing the Step-wise Supervision method, which enhances the strength of supervisory signals and improves training efficiency. As for policy tags, they serve as an additional benefit—analogous to high-quality prompts in LLMs—by informing the model of the actual level/type of the current action, thereby facilitating strategy distillation during training.\n\n# Gymnasium environments and reward shaping\n(To answer the reviewer's Question 2.)\n\nWe select the Gym envs based on two points: 1. The env should have a delayed reward to make it more difficult. 2. The observation and action space can match our model. We will also cite the Deep Sea environment, and compare the difference with the Gymnasium.\n\n(To answer the reviewer's Question 3.)\n\nReward shaping is a common practice in reinforcement learning, and it is applied consistently across all baselines as well as our model—we consider this approach fair. The quality of such reward shaping is analogous to that of prompts in large models: it helps render the problem more intuitive. Notably, without reward shaping, both the context size required for our model’s online learning and the number of training iterations needed for baselines would increase."}, {"review_id": "ZQVmT5oagb", "reviewer": "Reviewer_zr1Z", "summary": "The 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.", "questions": "Q1: 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?", "limitations": "Yes", "rating": 3, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "### 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.", "was_revised": true, "final_justification": "In summary, I appreciate the paper’s scope of ICRL scalability and its task distribution construction with structured diversity. My main concern is that the empirical evidence might be insufficient to support the key claim on large-scale ICRL, which was not well addressed by the authors’ response. I tend to maintain my current score as the reasons to reject outweigh the reasons to accept. If with sufficient empirical evidence on large-scale problems, I believe this paper is well qualified for the next venue.", "rebuttal": "# On the Value of Discrete Space MDPs\n(To answer the reviewer's Question 1.)\n\nWe appreciate suggestions to extend our framework to continuous MDPs, but we feel that the value of discrete MDPs is underestimated. At this stage, partially pursuing the task complexity risks undermining the core advantages of procedurally generated discrete MDPs. Below, we explain the challenges of preserving key properties in continuous domains and contrast our approach with contemporary benchmarks.\n\n## Challenges of Continuous MDPs in Preserving Structural Integrity\nExtending procedurally generated MDPs to continuous spaces while maintaining low structural bias and sufficient diversity poses significant technical hurdles. At the current stage, efforts moving to continuous MDPs risk typically **prioritizing \"absolute task complexity\" over qualitative diversity**, introducing biases or homogenizing task structures—contrary to our framework's goals.\n\n\n## Inherent Long-Term Dependency and Qualitative Diversity\n(To answer the reviewer's Question 3.)\n\nThe reviewer suggests conducting tests on tasks with varying complexity and different types of disturbances. In reality, existing continuous MDP benchmarks rely on continuously randomized hyperparameters to simulate \"infinite tasks,\" prioritizing few-shot learning over long-horizon in-context learning (ICL), such as environments like maze navigation and open-world games. Our analysis reveals two insights:\n- Task Diversity ≠ Numerical Scale: True diversity depends on qualitative variation in task structures, not just quantity.\n- ICL efficacy depends on disparity between tasks in the training set, not individual task complexity. Existing continuous benchmarks tend to produce homogeneous task structures, limiting models to superficial pattern recognition via in-weight learning (IWL) rather than sustained reasoning. In contrast, discrete MDPs in AnyMDP are designed with procedural complexity and temporal coherence, enabling rigorous evaluation of a model’s ability to generalize across structurally distinct tasks and sustain reasoning over extended sequences.\n\n## Compatibility with Existing LLMs\nBeyond being a novel benchmark, AnyMDP bridges a critical gap in assessing experience-driven, holistic reasoning capabilities of pre-trained LLMs. Existing benchmarks predominantly measure declarative memory and one-shot circuits, while AnyMDP emphasizes procedural memory and experience-driven learning.\n\n## Low-cost Access to Oracle Policy\nThis property significantly enhances the efficiency and scalability of ICRL training, making AnyMDP a scalable benchmark.\n\n## Empirical analysis across state-of-the-art architectures\nWe augment the empirical analysis with the training curves and ICL comparisons across state-of-the-art architectures, including Mamba2, RWKV-7, Gate Delta Net, and other competitive models. This will demonstrate AnyMDP’s utility as a rigorous benchmark for advancing long-context modeling capabilities.\n\n- Table 1. Validation Loss of Different Models on AnyMDP Datasets\n| CL \t  | 1  \t   | 100\t    | 1000    | 10000   | 30000 |\n|-------------|------------|------------|------------|------------|------------|\n|  Mamba2   | 1.417 ± 0.035 | 1.272 ± 0.051 | 0.822 ± 0.064 | 0.436 ± 0.053 | 0.294 ± 0.042 |\n|  GSA   | 1.406 ± 0.034 | 1.235 ± 0.053 | 0.765 ± 0.066 | 0.342 ± 0.050 | 0.235 ± 0.042 |\n|  Gate DeltaNet   | 1.356 ± 0.035 | 1.130 ± 0.051 | 0.599 ± 0.055 | 0.267 ± 0.049 | 0.173 ± 0.042 |\n|  RWKV-7   | 1.371 ± 0.034 | 1.151 ± 0.052 | 0.627 ± 0.055 | 0.269 ± 0.042 | 0.175 ± 0.029 |\n\n\nOverall, in our research on ICL, the benefits brought by the discrete MDPs far outweigh the drawbacks of its limited problem space. We have also started to explore ICL in the continuous space.\n\nFor future work in continuous settings, there are two directions. First, we will focus on specific continuous problems and prioritize in-domain adaptation. For instance, in frame prediction tasks, generalization manifests as adaptation to environments with distinct features and spatial structures. Similarly, in robot whole-body locomotion control, it manifests in adapting to different robot configurations (size, shape, weight). Second, we will continue exploring how to generate tasks with high diversity and low structural bias in continuous spaces. Undoubtedly, the first direction holds more practical value, while the second poses greater theoretical challenges.\n\n# The \"Perturbation\" in Question 4\nThe reviewer mentions \"perturbations can lead to a decline in model performance.\" We believe there are no such claims in our paper. Two candidates appear in our manuscript, neither fitting this question:\n- Domain/World randomization (Line 132): This creates task diversity; AnyMDP already exposes the learner to maximal variations, so no need to delve into specific impacts of perturbations.\n- Stochasticity in the behavior policy (Line 620 & Table 5): This is necessary for data synthesis and seems irrelevant to the mentioned perturbation. Moreover, incorporating Domain/World randomization and stochasticity in the behavior policy has enhanced model performance, contradicting the reviewer's viewpoint.\n# About Long-Term Adaptation in Weaknesses\nThe reviewer views long-term adaptation as a disadvantage of our methods, but we believe there are misunderstandings. Our proposed methods, step-wise supervision (SS) and chunk-wise learning, aim to improve in-context sample efficiency, not intentionally extend context length. Our experiments reveal that **\"Long-Term Adaptation is a tax we must pay toward increasing the generalization scope for ICRL.\"** Characterizing increased context length as an inherent limitation seems unjustified, even from the standpoint of whether trading generalization for long-context adaptation is worthwhile.\n\n# About Computational Costs in Question 5\nIn Figure 7 on page 17, we show the time consumption of the AnyMDP task generation on an Intel(R) Xeon(R) Platinum 8374C\nCPU. After generating the task files, the time consumption of generating the sequences accordingly is within hours, which can be ignored compared to the training cost. To train an OmniRL model with [16,128] state space and action space of 5, we generate 512K sequences with length 12K. The total training iterations are 0.2M, 8 seconds per iteration on 8 GPUs. We have also considered conducting the training across task complexity to assess the exact consumption, but the computing resources are tight. In general, training a model with a larger number of states or actions will bring greater resource consumption. After the problem becomes more complex, more sequences are needed to learn generalization, and each sequence has to be longer to obtain sufficient information for a single task.\n\n# Evaluated More ICRL Methods in Question 2\nWe have conducted experiments on well-known ICRL frameworks, including AD, AD𝜖, and DPT. Another technical direction is RL². However, with existing large-scale task sets, model scales, and context lengths, the computational cost of running RL is substantial. Using fewer tasks, smaller models, or shorter context lengths would render the conclusions lacking in comparative significance. Furthermore, methods like DPT have already been compared with RL²; thus, we believe the results from comparisons with DPT are sufficiently compelling."}, {"review_id": "tmm7PUXyBS", "reviewer": "Reviewer_Dwj8", "summary": "In 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.", "questions": "1) 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?", "limitations": "yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**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.", "was_revised": true, "final_justification": "The authors introduce a novel generation framework that is able to create a wide range of tabular Markov Decision Processes (with discrete state/action spaces) and show that it enables in-context learning of unseen MDPs. I believe this contribution is both novel and relevant for the RL community. The authors also addressed positively my questions and concerns. As such I maintain my positive score.", "rebuttal": "First and foremost, we greatly appreciate the reviewer for pointing out the typos in the comments, as well as the constructive suggestions regarding paper writing. We will revise accordingly to further improve the paper’s readability and quality. In the following, we will elaborate in detail on the questions raised by the reviewer.\n\n# Questions on AnyMDP\n## Behavior policies\n(To answer the reviewer's Question 1.)\n### Definitions for each policy\nOn page 20, we present Table 5 about \"Correspondence of prompt IDs and the policies it represents\". We thank the reviewer for pointing out the issue of lacking exact definitions for each policy, and will subsequently include them in the paper. For IDs 0 to 3, the policies are value iteration policies with different gamma factors, corresponding to short-term to long-term oracles. For the model-based reinforcement learner, we implement the RL solver algorithm based on Optimistic Thompson Sampling (OTS) for episodic reinforcement learning problems in unknown environments. It balances exploration and exploitation by dynamically adjusting exploration noise (incorporating UCB heuristics) to learn the optimal policy. The Tabular Q Learner is a model-free reinforcement learning algorithm that maintains a Q-table to store Q-values for each state-action pair, updates these values via temporal-difference learning using rewards and next-state information from environmental interactions.\n\n### Purpose of using different learning methods\nWhile the Q-learning and model-based behavior policies are always optimal at the end of training, their learning methods and paths during training differ. In our data generation process, these two policies learn synchronously along the trajectory of the data, rather than using pre-trained models to generate data. Because what we aim for our OmniRL model to learn is how to converge online to a good policy in a new task through interaction with the environment. Thus, the learning processes of these two behavior policies are actually more important than the final optimal policy—after all, the oracle policy can be directly obtained via value iteration.\n\n## Random sampling with condition ensures both diversity and quality\n(To answer the reviewer's Question 2.)\n\nThe sampling of transition matrix and reward matrix in AnyMDP is not uniformly random, to avoid generating trivial MDP problems (e.g., MDPs with optimal episode length=1). Uniform randomization would result in generating a majority of trivial Markov Decision Processes (MDPs), such as those with optimal episode lengths of 1 or MDPs where most states are irrelevant to the optimal solution. The policy only needs to remember which short path has the best reward and will not learn the trade-off of risk assessment. Eventually, the environment will be more inclined towards the bandit problem. In most complex RL problems, delayed reward is quite common. Usually, a large positive reward is only obtained when the goal is achieved. Sometimes, there may even be a stage of deduction before getting the positive reward.\n\nErgodicity is ensured by Theorem 1 in the appendix, which guarantees that a uniformly random policy induces a strictly positive probability of reaching every state from the initial state. At the same time, AnyMDP does not guarantee that the optimal policy is strictly ergodic—a trait absent in most existing benchmarks. However, the banded transition kernel and reward structure ensure that the optimal policy requires transitions across a substantial number of states.\n\nThus, via the \"Banded transition kernel\" condition specified in Line 525, we ensure the path to reaching high-value targets is longer. Following the procedure outlined in Algorithm 1 (page 15), we: randomly generate states with varying values (Line 2 of Algorithm 1); randomly select states that trigger reset or termination (Lines 3–4); randomly sample transition matrices under the \"Banded transition kernel\" condition (Lines 5–12); and randomly sample composite reward functions that satisfy the value condition (Lines 13–19). This approach ensures both task diversity and quality simultaneously.\n\n## Priori knowledge\n(To answer the reviewer's Question 5.)\n\nThe reviewers raised questions regarding the inclusion of prior information in AD, AD𝜖, and DPT. To clarify, prior information introduces additional details about different behavioral policies—specifically, policy tags. Since the AD policy remains unchanged during generation, adding prior information has no theoretical effect. For AD𝜖, policy tags can be incorporated based on the noise level. Additionally, as DPT involves multiple policies, the inclusion of policy tags is feasible. We believe that adding policy tags is also highly likely to enhance performance for AD𝜖 and DPT.  \n\nIndeed, compared with DPT, our primary contribution lies in proposing the Step-wise Supervision method, which enhances the strength of supervisory signals and improves training efficiency. As for policy tags, they serve as an additional benefit—analogous to high-quality prompts in LLMs—by informing the model of the actual level/type of the current action, thereby facilitating strategy distillation during training.\n\n# Question on experiment settings\n## Testing task\n(To answer the reviewer's Question 3.)\n\nWe chose the RL benchmark Gym for our out-of-distribution tasks. In most suitable games within it, state and action spaces are fixed, making it infeasible to conduct a comprehensive experiment on the number of states and actions. However, in Figure 3 on page 6, we compare performance across AnyMDP tasks with different numbers of states, where results show that difficulty increases as the state space grows.\n\n## Multi-agent test\n(To answer the reviewer's Question 4.)\n\nBoth agents are independent OmniRL agents. They perform ICL in the same environment with distinct observations. Observations are partial: one agent can only observe the column number of the other agent. Thus, both agents operate dynamically. We were also pleasantly surprised that the model—trained on single-agent datasets with a single-agent architecture—could solve multi-agent problems. This also inspires us to explore the application of ICL in multi-agent scenarios."}]

Author rebuttal:

Meta-review: This paper introduces the AnyMDP framework, which leverages randomized MDPs to generate diverse tasks for large-scale meta-training in In-Context Reinforcement Learning (ICRL). By incorporating step-wise supervision and prior knowledge induction, the framework improves both efficiency and generalization. Experiments show that OmniRL, trained on these randomized tasks, achieves strong performance on unseen environments. The work highlights the importance of task diversity and world randomization in scaling ICRL and improving robustness, with results demonstrating that larger task sets enhance generalization, albeit with increased adaptation time.

All reviewers agree that the paper makes novel and valuable contributions, particularly in enabling scalable task generation and reducing structural bias through world randomization. Concerns remain regarding task complexity and presentation. While the writing is generally clear and well organized, the explanation of key results—especially in figure captions—could be improved. Moreover, the proposed task generation does not yet scale to complex or continuous domains, and these limitations should be explicitly acknowledged in the camera-ready.
Decision: Accept (poster)