Context and Diversity Matter: The Emergence of In-Context Learning in World Models

## Abstract

Abstract The capability of predicting environmental dynamics underpins both biological neural systems and general embodied AI in adapting to their surroundings. Yet prevailing approaches rest on static world models that falter when confronted with novel or rare configurations. We investigate in-context environment learning (ICEL), shifting attention from zero-shot performance to the growth and asymptotic limits of the world model. Our contributions are three-fold: (1) we formalize in-context learning of a world model and identify two core mechanisms: environment recognition and environment learning; (2) we derive error upper-bounds for both mechanisms that expose how the mechanisms emerge; and (3) we empirically confirm that distinct ICL mechanisms exist in the world model, and we further investigate how data distribution and model architecture affect ICL in a manner consistent with theory ♠ \spadesuit . These findings demonstrate the potential of self-adapting world models and highlight the key factors behind the emergence of ICEL, most notably the necessity of long context and diverse environments.

## 1 Introduction

The ability to predict future environmental states is crucial for reasoning and decision-making in both animals and humans. Inspired by this principle, constructing predictive models, especially the world model [1], to forecast environmental dynamics and outcomes forms the foundation for enabling agents to plan effective decisions and behaviors [2, 3, 4, 5, 6]. Consequently, world models are widely applied in fields such as navigation [7, 8, 9, 10, 11], autonomous driving [12, 13, 14, 15, 16], robotics [17, 18, 19, 20, 21, 22], and are considered a cornerstone of embodied artificial intelligence.

Despite their proven effectiveness across various applications, previous prediction frameworks largely rely on static world models optimized for zero-shot, few-shot, or instantaneous performance. In contrast, humans and animals achieve real-time adaptation [23] through predictive coding—a process where prediction errors drive attention, generate feedback, and motivate learning and adjustment [24, 25]. For instance, when confronted with rare environments, humans experience surprise yet rapidly recalibrate their predictions for that setting, whereas static models continue to fail unless explicitly retrained on the relevant data.
The ability to dynamically modify predictive mechanisms based on observational evidence, rather than relying solely on fixed parametric memory and external mapping modules, can effectively enable the model to adapt to environments unseen during training.
This capability can be effectively addressed by In-Context Learning (ICL), as evidenced by recent advances in Large Language Models [26]. However, the potential of ICL for resolving plasticity in world models, particularly In-Context Environment Learning (ICEL), remains largely underexplored in current literature. Addressing this gap could further improve the generalization scope for embodied AI.

We isolate two canonical mechanisms that let world models adapt in context.
In-Context Environment Recognition (ICER) taps the model’s parametric memory of previously seen environments and uses the context only to identify which one is active.
ICEL ignores prior settings and instead extracts the full environment from the context itself.
By deriving error bounds for both modes, we show how their emergence depends on data distribution, including environment diversity, complexity, and trajectory length.
To test the theory, we introduce Linear-attention Long-context world model, *L2World*, which effectively model long sequences with high computational efficiency.
Across cart-pole control and vision-based indoor navigation, we find that distributional properties and long-context capacity govern whether the model invokes ICER or ICEL.
Despite the light image encoders and decoders of L2World, it sets a new state-of-the-art for long-sequence observation prediction in navigation, outperforming approaches that rely on heavy diffusion-based image backbones.
These results underscore the importance and the promise of deliberately incentivizing ICEL through carefully designed datasets and model architectures to strengthen the asymptotic generalization capability of world models, rather than focusing solely on instant or zero-shot frame-level performance.

## 2 Related Work

### 2.1 Dynamic Prediction and World Models

Dynamic models, also referred to as world models [1, 27], encompass probabilistic, physical, or generative frameworks that formalize an AI system’s environmental understanding [28, 29, 30, 31]. These models predict future states by leveraging historical observations and play a pivotal role in advancing reinforcement learning (RL) methodologies and related domains. Specifically, they constitute foundational components in model-based RL [32, 33], enable simulations that facilitate agent learning through virtual experiences (thereby reducing reliance on direct environmental interaction) [34, 2, 5, 35, 36], and serve as auxiliary tasks to augment model supervision [37, 38, 39].

World models demonstrate robustness in integrating multi-modal raw sensor data, including visual, textual, inertial, and tactile inputs. To mitigate challenges posed by high-dimensional sensory inputs, representation learning paradigms such as Generative Adversarial Networks (GANs) [40] and Variational Autoencoders (VAEs) [41] are widely utilized to compress raw data into compact latent spaces. Subsequent temporal modeling in these reduced dimensions is achieved through latent world model architectures [2, 42, 4, 5, 43, 39, 44, 45], which capture sequential dependencies and enable coherent long-term predictions.

In navigation, early systems relied on traditional, hand-crafted pipelines such as SLAM. Recent work replaces these modules with generative world models, including diffusion [7], VAE [9], and RL-enhanced variants [46, 10], which reconstruct dynamics or simulate semantics [47, 48]. However, most existing methods disregard continual adaptation, especially across episodes, leaving a persistent gap between zero-shot performance and lifelong operation.

### 2.2 In-Context Learning and Meta-Learning

The approach based on parametric memory, which predominantly relies on gradient-based optimization or In-Weight Learning (IWL), has faced criticism for its lack of plasticity and the associated challenges it poses in continual learning scenarios [49]. Conversely, ICL has emerged as a pivotal capability in large language models [26], facilitating generalization to novel tasks without the necessity of parameter fine-tuning. ICL leverages contextual memory for task solutions, rather than depending on parametric memory. The concept of ICL is not novel. Meta-learning [50, 51], which focuses on acquiring learning capabilities rather than mastering specific skills, utilizes well-curated environments or data instead of relying on large-scale, uncurated pre-training data. Nevertheless, the lack of well-structured, task-rich, and cost-efficient datasets continues to present a significant challenge.

ICL has been utilized to encode a diverse array of learning mechanisms, including language learning [52], regression [53], Bayesian inference [54], reinforcement learning [55, 56, 57], and world models [58, 59], highlighting its resemblance to biological plasticity [60].
Yet existing work concentrates on few-shot in-context adaptation, overlooking ICL’s potential as contexts grow indefinitely.
Moreover, little insight has been offered into its inner process and mechanisms.
Concurrent studies identify at least two ICL modes in LLMs: task recognition (TR) and task learning (TL) [61].
Key factors conducive to ICL emergence have also been investigated [62, 58], such as “burstiness” and diversity, and so is the relation between IWL and ICL [63], but theory in this area remains under-developed.

## 3 Methodologies

We consider an environment $e$ specified by a Partially Observable Markov Decision Process (POMDP) $e:\langle O,S,A,T_{e},Z_{e}\rangle$, where $S$ is the state space, $A$ is the action space, $O$ is the observation space, $T_{e}(s,a,s^{\prime})=p_{\tau,e}(s^{\prime}|s,a)$ is the transition model, and $Z_{e}(s,o)=p_{z,e}(o|s)$ is the observation model.
(^1^11While most prior work integrates the reward model into the world model, our analyses omit explicit consideration of rewards. Nonetheless, since rewards can typically be derived from the state or observation, our framework allows for straightforward extension to incorporate reward-related considerations.)
A fully observable MDP is denoted with $e:\langle S,A,T_{e}\rangle$ with $o\equiv s$. We denote a world model with the following equation:

$$ $\displaystyle\textbf{World Model: }\hat{o}_{t+1}\sim\hat{p}_{\theta}(\cdot|q_{t})=f_{\theta}(q_{t}),\text{with }q_{t}$ $\displaystyle=(s_{t},a_{t})\text{ (MDP)}$ $\displaystyle\text{or }q_{t}$ $\displaystyle=(o_{t-\Delta t},a_{t-\Delta t},\dots,o_{t},a_{t})\text{ (POMDP) }.$ (1) $$

Let $\theta$ denote the model parameters; values marked with a hat denote predictions, and unmarked values denote the ground truth. Consider extra contexts of observations and actions, $C_{T}=(o^{(C)}_{1},a^{(C)}_{1},\dots,o^{(C)}_{T},a^{(C)}_{T})$, where $T$ indexes the context length; the ICL capability of the world model is then characterized by the following condition:

$$ $\displaystyle\forall T_{1}>T_{2},D[\hat{p}_{\theta}(\cdot|q_{t},C_{T_{1}})||p_{e}(\cdot|q_{t})]<D[\hat{p}_{\theta}(\cdot|q_{t},C_{T_{2}})||p_{e}(\cdot|q_{t})],$ (2) $$

Here, $D$ represents a metric measuring the error between two distributions (lower values indicate better performance). Notably, the ICL of the world model fundamentally relies on cross-episode contexts rather than intra-episode state estimation. To distinguish these concepts clearly, we use $q_{t}$ to denote short-term state estimation and $C_{T}$ to represent long-term ICL. While these are typically aligned in a single sequence in practice, maintaining this distinction facilitates rigorous theoretical analysis. Building on prior work that partitions in-context learning into task-recognition (TR) and task-learning (TL) modes [61, 64], theoretically, we are able to identify two analogous modes within world-model ICL: *Environment Recognition* (ER) and *Environment Learning* (EL). We then derive error bounds that characterize the conditions under which each mode emerges.

### 3.1 Environment Recognition (ER)

To clarify [Equation˜2](https://arxiv.org/html/2509.22353v1#S3.E2), we consider a world model optimized on the finite environment set $\mathcal{E}=\{e_{1},\dots,e_{|\mathcal{E}|}\}$. Assume the system possesses an environment-specific model $\hat{p}_{\theta,e}$ for every environment $e$, then $\hat{p}_{\theta}$ decomposes as follows:

$$ $\displaystyle\hat{p}_{\theta,ER}(o_{t+1}|q_{t},C_{T})$ $\displaystyle=\sum_{e\in\mathcal{E}}\underbrace{\hat{p}_{\theta}(e|q_{t},C_{T})}_{\text{Environment Recognition}}\cdot\underbrace{\hat{p}_{\theta,e}(o_{t+1}|q_{t})}_{\text{Environment-Specific World Model}}$ (3) $\displaystyle\hat{p}_{\theta,e}(o_{t+1}|q_{t})$ $\displaystyle=\int\underbrace{\hat{p}_{\theta,s,e}(s_{t}|q_{t})}_{\text{State Estimation}}\cdot\underbrace{\hat{p}_{\theta,\tau,e}(s_{t+1}|s_{t},a_{t})}_{\text{Dynamics}}\cdot\underbrace{\hat{p}_{\theta,z,e}(o_{t+1}|s_{t+1})}_{\text{Observation Model}}\,ds_{t}$ $$

[Equation˜3](https://arxiv.org/html/2509.22353v1#S3.E3) implies that in-context learning in world models arises mainly from the continual refinement of environment recognition, since the context yields no improvement to the environment-specific world-model term $\hat{p}_{\theta,e}=\{\hat{p}_{\theta,s,e},\hat{p}_{\theta,\tau,e},\hat{p}_{\theta,z,e}\}$. Therefore, in the ER regime, the model first acquires world models for the entire environment set through IWL or parametric memory, and then uses the context solely to identify the current environment.

### 3.2 Environment Learning (EL)

[Equation˜3](https://arxiv.org/html/2509.22353v1#S3.E3) is efficient when the environment set $\mathcal{E}$ is small, yet its effectiveness diminishes rapidly as the size and diversity of $\mathcal{E}$ grow or when the system faces open worlds. However $\hat{p}_{\theta}$ can also be approximated without estimating $e$ at all, by directly accumulating the evidence for $(q_{t},o_{t+1})$ across all contexts:

$$ $\displaystyle\hat{p}_{\theta,EL}(o_{t+1}|q_{t},C_{T})$ $\displaystyle=\frac{p(q_{t},o_{t+1}|C_{T})}{p(q_{t}|C_{T})}$ (4) $$

An intuitive observation of [Equation˜4](https://arxiv.org/html/2509.22353v1#S3.E4) is that EL functions, at minimum, as a in-context memorizer. For highly complex environments, its performance might degrade sharply, because accurate estimation of the transition of $q_{t}$ demands contexts that scale with the environment’s complexity.

### 3.3 Theoretical Analyses of ER and EL

Although the training process and data distribution play a key role in effectively incentivizing ICL [62], how does the data distribution determine whether EL or ER emerges? If training consistently minimizes predictive error, the error bounds of EL and ER become the decisive factor in selecting the emergent mode. To investigate the conditions governing the emergence of the two modes (ER and EL), we analyze the error upper bounds for each paradigm. For tractable analysis, we introduce the following simplifying assumptions:
(1) The observation, state, and action spaces are discrete;
(2) Both modes achieve ideal state estimation $p_{e}(s|q)$;
(3) The context $C_{T}$ has a uniform state-action distribution.
Under these assumptions, we derive an upper bound on the error of the world model optimized over environment set $\mathcal{E}$, measured by the total variation (TV) distance, when deployed in an unseen environment $e_{0}$ at context horizon $T$.
Formally, the TV error is bounded by:

###### Theorem 1 .

For Environment Recognition and Environment Learning whose predictive models $\hat{p}_{ER/EL}$ have been sufficiently optimized on the training environments $\mathcal{E}$, the upper bound of the total-variation (TV) distance between the predicted and the ground-truth transition, given a context $C_{T}$ of length $T$, can be estimated as:

$$ $\displaystyle\text{TV}(\hat{p}_{ER},p_{e_{0}})\leq$ $\displaystyle\min\{\underbrace{\alpha(|\mathcal{E}|-1)T^{-1/2}/3}_{\text{Identification Error}}+\underbrace{\min_{e\in\mathcal{E}}TV(\hat{p}_{\theta,e},p_{e_{0}})}_{\text{Best Matching Error}},\quad\underbrace{\max_{e\in\mathcal{E}}TV(\hat{p}_{\theta,e},p_{e_{0}})}_{\text{Worst Matching Error}}\}$ $\displaystyle\text{TV}(\hat{p}_{EL},p_{e_{0}})\leq$ $\displaystyle\underbrace{\sqrt{2|O||S||A|log(4|O|/\delta)}}_{\text{Environment Complexity}}\cdot T^{-1/2},$ $\displaystyle\text{with probability }1-\delta\text{, and }T>4|S|^{2}|A|^{2}\log(4|S||A|/\delta)$ (5) $$

Proofs and detailed assumptions for the above theorems are deferred to [Appendix˜A](https://arxiv.org/html/2509.22353v1#A1).
An immediate observation from [Theorem˜1](https://arxiv.org/html/2509.22353v1#Thmtheorem1) is that ICEL enjoys an ideal error upper bound that decays as $T^{-1/2}$, whereas ICER carries a non-decaying residual term (the best-matching error) that becomes the dominant obstacle to generalizing across unseen environments.
To enhance generalization, we therefore ask: under what condition is ICEL preferred to ICER?
For the entire training set $\mathcal{E}$, ICEL dominates whenever
$\mathbb{E}_{e\in\mathcal{E}}\bigl[\text{TV}(\hat{p}_{\text{EL}},p_{e})\bigr]\ll\mathbb{E}_{e\in\mathcal{E}}\bigl[\text{TV}(\hat{p}_{\text{ER}},p_{e})\bigr];$
The opposite inequality favors ICER.
Although the errors themselves are intractable to evaluate directly, the following insights are obtained by comparing their upper bounds:

(1) Lower environment complexity and greater environment diversity favor EL over ER:
Note that the best-matching error is effectively zero, because the model is always evaluated on environments seen during training.
The cardinality of the training set, $|\mathcal{E}|$, enters only the ER bound, whereas the environment complexity, $|O||S||A|$, enters only the EL bound.
Consequently, lower complexity and a larger training set push the EL bound below the ER bound.
Moreover, because ER error is limited by the worst-matching error, scarce diversity in the training environments can keep the ER error small and prevent EL from ever overtaking it.

(2) Long context and environment diversity are key to both ER and EL:
Insufficient diversity not only blocks the emergence of ICEL but also suppresses ICER, because a small worst-matching error removes the pressure to reduce the identification error.
Once the training set is sufficiently diverse, both ER and EL obtain an upper error bound that decays as $T^{-1/2}$, demonstrating that long context is indispensable for either mechanism.

(3) Over-training and powerful IWL facilitate ER over EL: we hypothesize that IWL perfectly models the environment-specific dynamics ($\hat{p}_{\theta,e}$) in the training set, so the best-matching error is nearly zero during training; however, this is not always true. Early in training, IWL can still incur large errors, transiently pushing the model toward EL instead of ER. This transiency is also investigated by prior works in ICL [65]. As training proceeds and IWL becomes increasingly accurate, the model may revert to ER.

In the following section, we empirically confirm that these insights hold not only for discrete settings but also for continuous MDPs and POMDPs. Because theory predicts that large environmental diversity and low task complexity are required for incentivizing EL, and most of the closed benchmarks can not satisfy those requirements. We construct our dataset from randomly sampled *cart-poles* and procedurally generated *mazes* to evaluate the performance of ICER and ICEL.

Figure: Figure 1: The world model structure for the empirical study.
Refer to caption: x1.png

###### Theorem 1 .

### 3.4 L2World: Long-Context and Linear-Attention World Models

Prior work relies on multi-token representations or diffusion models to deliver high-fidelity single-frame reconstructions specifically for images.
While these approaches set the state-of-the-art for static images, they introduce prohibitive memory and computational bottlenecks when sequences grow to the length required for ICEL.
We therefore introduce L2World that trade per-frame fidelity for temporal scalability: for image observations, we compress each frame $o_{t}$ into a latent state $s_{t}$ with a lightweight variational auto-encoder (VAE) [66] whose encoder $f_{I}$ and decoder $g_{I}$ are ResNet stacks; for low-dimensional observations, we simply apply a small multi-layer perceptron encoder/decoder pair.
For computational efficiency, we do not model the full state estimation with $\hat{p}(s_{t}|q_{t})$. Instead, we construct a pseudo-state that depends solely on the instant observation and leaves all temporal-related encoding to transition modeling [4].
We then construct the adaptive world model $\hat{p}_{\theta}(o_{t+1}|q_{t},C_{T})$ using an efficient sequence decoder $f_{\theta}$. Here we implement gated slot attention layers [67, 68, 69] with chunk-wise parallelization during the training phase, while retaining the recurrent form during the inference phase. The predictor of transition at first yields the output $h_{t}$, which is further processed by the decoder $g_{S}$ to produce the predicted state $\hat{s}$, corresponding to predicting a Gaussian distribution over the latent space $\hat{p}\sim\mathcal{N}(\hat{s}_{t},\sigma^{2}_{s})$. Although this assumption could lead to significant loss of accuracy in stochastic environments, for the navigation tasks we consider, it is acceptable and greatly increases computational efficiency. The model and the target function are listed as follows (see details in [Section˜B.2](https://arxiv.org/html/2509.22353v1#A2.SS2)):

$$ $\displaystyle\text{Observation Encoder}:s_{t},\sigma_{s,t}=f_{I}(o_{t})\quad\text{Observation Decoder}:\hat{o}_{t}=g_{I}(\hat{s}_{t})$ $\displaystyle\text{Latent Decoder}:\hat{s}_{t},\hat{\sigma}_{s,t}=g_{S}(h_{t})\quad\text{Action Encoder}:a_{t}=f_{A}(Action[t])$ $\displaystyle\text{Chunk-wise Temporal Modeling}:\phi_{t},h_{1},...,h_{t}=f_{\theta}(s_{1},a_{1},...,s_{t},a_{t})$ $\displaystyle\text{Recurrent Temporal Modeling (Evaluating)}:\phi_{t},h_{t+1}=f_{\theta}(\phi_{t-1},s_{t},a_{t})$ $\displaystyle\text{Observation Reconstruction Loss}:\mathcal{L}_{o}=||o_{t}-g_{I}(\hat{s}_{t})||+\lambda KL(\mathcal{N}(s_{t},\sigma_{s,t})||\mathcal{N}(0,1))$ $\displaystyle\text{State Transition Loss:}\mathcal{L}_{s}=-\sum_{t}KL(\mathcal{N}(s_{t},\sigma_{s,t})||\mathcal{N}(\hat{s}_{t},\hat{\sigma}_{s,t}))$ $$

When the observation is an image, we first pre-train the image encoder $f_{O}$ and decoder $g_{O}$ on pre-sampled observations; after this stage, their parameters are frozen while the temporal model is trained. For lower-dimensional observations, all encoders/decoders are updated jointly with the temporal model in a single end-to-end phase.

## 4 Experiments

Although prior studies [62, 65, 53] have validated the influence of data distribution on ICL, they have largely concentrated on simplified tasks such as regression and classification. To examine how data distribution, model architecture, and training procedure jointly affect the emergence of ICEL/ICER, we select two canonical benchmarks. First, the cart-pole, a classical continuous-control problem, in which ICEL primarily targets the acquisition of varying embodiments and physical constants. Second, indoor navigation, a widely recognized POMDP, in which ICEL focuses on learning and memorizing spatial coefficients. These two experiments confirm not only the impact of each factor on ICEL but also demonstrate that ICEL spans a broad generalization spectrum, extending from spatial reasoning and memorization to adaptation of embodiment and physical constants.

### 4.1 Random Cart-Poles

Experiment Setting. To investigate ICEL and ICER in cart-pole environments, we randomized four variables in the environment settings: gravity $g$, cart mass $m_{c}$, pole mass $m_{p}$, and pole length $l$. We focus on two different scopes of the configurations to investigate the impact of the diversity issue: Scope 1 remains close to the original task, whereas Scope 2 covers a larger region and excludes Scope 1 (details are left to [Figure˜5](https://arxiv.org/html/2509.22353v1#A2.F5)). For each environment, we first trained an RL agent and then collected trajectories with an expert policy perturbed by uniform noise spanning $[0.3,0.7]$ to ensure adequate coverage of the state-action space. We trained five comparison models; all share the same data scale (128K trajectories $\times$ 200 step/trajectory) but differ in the number and scope of the environments.

- •
1-Env: 96K trajectories are sampled from the original environment ($g=9.8,m_{c}=1.0,m_{p}=0.1,l=0.5$), with 200 steps per trajectory.
- •
4-Envs: 4 environments sampled from Scope 1+2, each with 24K trajectories.
- •
16-Envs: 16 environments sampled from Scope 1+2, each with 8K trajectories.
- •
8K-Envs (Scope 1/1+2): 8,000 environments sampled from Scope 1 or 1+2, each with 16 trajectories.

We evaluate with three test sets: (1) Seen 4-Envs, but the trajectories are independently sampled; (2) 256 environments from Scope 1; and (3) 256 environments from Scope 2. We evaluate mainly with the average prediction error, which is averaged over each of the context lengths $T$.

Figure: Figure 2: Comparison of models trained on different datasets (color-coded) in Cart-Poles. Performance varies markedly with the training data, revealing distinct tendencies toward ICER, ICEL, or an inability to perform ICL.
Refer to caption: x2.png

We plot the evaluation results in [Figure˜2](https://arxiv.org/html/2509.22353v1#S4.F2), where the following insight is worth noting:

Importance of both environment scope and environment number: A comparison between the model trained on 1 environment (1 Env) and 4 environments (4 Envs) with the other group demonstrates that an insufficient number of environments leads to the absence of ICL and generalization, except in the tasks that the model has already seen. The 4 Envs group exhibits clear ICER characteristics, with a substantial performance gap between seen and unseen tasks. The 16 Envs (Scope 1 + 2) and 8K Envs (Scope 1) groups display similar capabilities; however, they lag significantly behind the 8K Envs (Scope 1 + 2) group, indicating that both the scope of tasks and the number of tasks are crucial.

Divergence between few-shot and many-shot performances: Another insight gleaned from the comparison between 4-Envs and 8K-Envs is that the latter, which has a broader generalization scope, also requires more context to learn. This is evidenced by the fact that the performance of the latter group does not surpass that of the former group until $T>10$. This also validates the theoretical analysis, indicating that a longer context is a cost for achieving better generalization.

Over-training reduces generalization when training environments are insufficient: To isolate the effect of over-training, we extract an early-stage checkpoint from the 4-Env group and evaluate it across all environments. Although its performance in seen environments is sub-optimal, it generalizes to unseen environments with a considerable margin over the over-trained model, confirming the shift from ICL- to IWL-based reliance.

### 4.2 Navigation

Experiment settings. Procedurally generated mazes are a demanding test-bed for transition prediction [70, 71].
By stripping observations of semantic cues, the maze framework keeps task complexity low while still exposing models to stochastic, partially observable dynamics.
We amplify diversity through fully randomized configurations that vary topology, textures, object placement, and agent embodiment (^1^11Procedural Maze environments: [https://github.com/FutureAGI/Xenoverse/tree/main/xenoverse/mazeworld](https://github.com/FutureAGI/Xenoverse/tree/main/xenoverse/mazeworld)).
Room-tour data are collected in a way similar to cart-pole: we perturb the oracle object-navigation policy [72, 70] with random noise levels in $[0.05,0.95]$ and record the complete trajectory of an agent in each environment.
The oracle policy was derived using Dijkstra’s algorithm based on the ground truth 2-D occupancy maps. Observations are RGB images standardized to 128 × 128 pixels; The action space comprises 17 discrete actions, each corresponding to a unique offset and rotational movement.
Table [1](https://arxiv.org/html/2509.22353v1#S4.T1) lists the resulting training sets, each drawn from a different number of environments and exhibiting varying trajectory lengths.
To isolate the effect of data distribution, all maze datasets contain the same number of frames while differing in the coverage of trajectories and environments.
To further investigate its transferability to more realistic environments, we also collect trajectories from the semantically rich ProcTHOR simulation, which offers a wide variety of assets [73, 74]. Specifically, we curate two datasets: a larger one with 40K trajectories and a smaller one with 5K trajectories, each trajectory having a length of 2K.

**Table 1: A summary of data distribution across the training datasets.**
| Training DataSet | # envs ($|\mathcal{E}|$) | Len. Traj. | # Traj. | # frames | Indoor area |
| --- | --- | --- | --- | --- | --- |
| Maze-32K-L | 32K | 10K | 32K | 320M | $380\sim 3422m^{2}$ |
| Maze-32K-S | 32K | 100 | 3.2M | 320M | $380\sim 3422m^{2}$ |
| Maze-128-L | 128 | 10K | 32K | 320M | $380\sim 3422m^{2}$ |
| Maze-128-S | 128 | 100 | 3.2M | 320M | $380\sim 3422m^{2}$ |
| ProcTHOR-5K | 5K | 2K | 5K | 10M | $40\sim 600m^{2}$ |
| ProcTHOR-40K | 40K | 2K | 40K | 80M | $40\sim 600m^{2}$ |

For training, inspired by overshooting [43] and Hu et al. [38], we randomly mask the $s_{t}$ of the input sequences at some positions to enhance the model’s capability to predict the distant future (see training details in [Section˜B.1](https://arxiv.org/html/2509.22353v1#A2.SS1)). Evaluation is conducted on both seen and unseen tasks. By default, we use an evaluation set scale of $|\mathcal{E}|=256$. The evaluation process involves encoding a context of length $t$ for ICEL and then predicting future $k$-step transitions using auto-regression and the ground truth action records $a_{t+1},a_{t+2},...,a_{t+k}$. We assess the error in both the latent spaces and the decoded images, which we refer to as $k$-step prediction with context length $T=t$. For example, $T=10$ and $k=4$ predict future $4$ steps with a context length of $10$.

We evaluate two additional baselines alongside the proposed long-context world model:
(1) Navigation World Model (NWM) [7], which employs diffusion layers to predict the next frame from the preceding four frames;
(2) Dreamer-v3 [34], which uses LSTM layers for temporal encoding.
For NWM, we retain its original pre-trained image encoders and re-train only the diffusion layers on the target dataset.
For Dreamer-v3, we remove the policy components and train only the world-model module to ensure a fair comparison.

**Table 2: Comparison of the performances (PSNR $\uparrow$) of 1-step future prediction in Mazes.**
| Model | Seen | Unseen |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| T=1 | T=10 | T=100 | T=1000 | T=10000 | T=1 | T=10 | T=100 | T=1000 | T=10000 |  |
| L2World (Maze-32K-L) | 16.80 | 20.97 | 23.11 | 24.65 | 25.05 | 16.37 | 21.24 | 23.17 | 24.66 | 24.65 |
| L2World (Maze-32K-S) | 18.57 | 19.28 | 19.67 | 20.21 | 20.48 | 18.45 | 19.24 | 19.63 | 20.29 | 20.31 |
| L2World (Maze-128-S) | 19.47 | 20.39 | 20.58 | 22.02 | 21.77 | 18.01 | 18.63 | 19.00 | 19.67 | 19.63 |
| L2World (Maze-128-L) | 18.54 | 20.86 | 23.32 | 25.65 | 26.00 | 17.54 | 19.43 | 20.96 | 21.54 | 21.52 |
| Dreamer (Maze-32K-L) | 16.40 | 21.82 | 19.24 | 21.26 | 21.89 | 16.81 | 20.48 | 21.40 | 22.65 | 22.12 |
| Dreamer (Maze-128-L) | 17.13 | 20.64 | 21.83 | 22.20 | 22.43 | 14.26 | 14.54 | 14.09 | 13.46 | 13.50 |
| NWM (Maze-32K-L) | 20.84 | 20.21 | 19.19 | 22.32 | 21.06 | 16.20 | 16.71 | 17.00 | 17.37 | 17.85 |

Figure: Figure 3: Comparison of k-step autoregressive PSNR in Mazes(Unseen).
Refer to caption: x3.png

Impact of data distribution and model architecture on ICEL.
[Table˜2](https://arxiv.org/html/2509.22353v1#S4.T2) reports the next-frame prediction quality estimated by PSNR ($k=1$) for models trained on different datasets. Three empirical findings corroborate [Theorem˜1](https://arxiv.org/html/2509.22353v1#Thmtheorem1): (1) The 32K-L dataset yields the best generalization to unseen environments, whereas the 128-L dataset excels on seen ones; in both cases, peak performance occurs at the asymptotic stage, not at the beginning of the context. (2) Long-context training consistently produces stronger ICL than short-context training, confirming that extensive context is necessary for ICL to emerge. (3) Dreamer and NWM fall short even with a long-context dataset: Dreamer’s LSTM backbone and NWM’s 4-frame horizon show that architectures incapable of fully leveraging long contexts cannot achieve many-shot ICEL.
[Figure˜3](https://arxiv.org/html/2509.22353v1#S4.F3) further presents the $k=\{1,2,4,8\}$-step prediction performances on unseen Mazes, measuring how far ahead the world models can reliably foresee. For $k>1$, the performance–context-length curve largely tracks the next-frame trend across models, except for Dreamer (Maze 32K-L): at $k=8$, its performance plateaus once $T>100$, whereas $k=1$ keeps improving, revealing a larger compound-error accumulation than in our method.

Figure: Figure 4: Example predictions produced by L2World and the baselines in Mazes and ProcTHOR for $T=\{1,100,10000\}$ and $k=10$, together with the corresponding ground-truth sequences.
Refer to caption: x4.png

[Figure˜4](https://arxiv.org/html/2509.22353v1#S4.F4) illustrates 10-step-ahead predictions at $T=\{1,100,10K\}$ for our method, Dreamer, and NWM trained on Maze-32K-L. NWM produces visually convincing frames with fine textures and a plausible layout, yet frame-wise fidelity alone is insufficient for correct long-range forecasting because the model lacks long-term memory and spatial reasoning.

**Table 3: Comparison of the PSNR of 1-step future prediction in ProcTHOR (Unseen)**
| Model | Pre-train | Post-train | T=1 | T=10 | T=100 | T=1000 | T=10000 |
| --- | --- | --- | --- | --- | --- | --- | --- |
| L2World | - | ProcTHOR-5K | 15.49 | 18.22 | 19.02 | 19.74 | 19.81 |
| L2World | Maze-32K-L | 16.46 | 20.23 | 21.05 | 21.89 | 22.04 |  |
| L2World | Maze-32K-S | 19.80 | 19.45 | 19.86 | 20.57 | 20.61 |  |
| L2World | Maze-128-L | 19.16 | 19.60 | 20.20 | 20.94 | 16.46 |  |
| Dreamer | - | ProcTHOR-40K | 19.82 | 22.61 | 23.99 | 23.51 | 22.76 |
| NWM | - | 18.30 | 21.41 | 21.11 | 21.02 | 20.08 |  |
| L2World | - | 21.57 | 22.67 | 23.39 | 24.92 | 22.98 |  |
| L2World | Maze-32K-L | 17.21 | 22.81 | 24.32 | 25.40 | 23.94 |  |

Transferability of ICEL and ICER: To verify whether ICEL emerges in semantically richer, realistic environments and whether skills acquired in simpler environments transfer to more complex ones, we fine-tune the models on ProcTHOR trajectories and evaluate them in varied contexts, as shown in [Table˜3](https://arxiv.org/html/2509.22353v1#S4.T3). The model pre-trained on Maze-32K-L not only excels in unseen mazes but also retains its advantage when fine-tuned on the small ProcTHOR-5K set; its transferability markedly exceeds that of Maze-128-L and other baselines, evidencing ICEL’s domain-generality. Further increasing ProcTHOR data yields continual improvement in our model while maintaining a large margin over Dreamer and NWM.
However, the performance from $T=1K$ to $T=10K$ begins to deteriorate once the training data switches from ProcTHOR-5K to ProcTHOR-40K, implying that ProcTHOR trajectories with $T\leq 2K$ fail to incentivize ICEL at longer contexts. Yet, this phenomenon does not appear with ProcTHOR-5K, verifying that the ICEL capability acquired in L2World transfers effectively to ProcTHOR despite the large gap between the two classes of environments.

## 5 Conclusions and Limitations

Conclusions: This work investigates in-context learning of world models, specifically dynamic models, focusing on the possible modes of ICEL and ICER in MDP and POMDP. Theoretically, we analyze error upper bounds for both modes to characterize their properties and identify the conditions under which each excels. Empirically, we introduce L2World and validate these insights in cart-pole and navigation tasks. Our results underscore that both high environment diversity and sufficient context length of the world model are essential to elicit ICEL.

Limitations: At present, our analysis is confined to the dynamic model; the reward and policy models can be addressed subsequently. This work constitutes a first step toward broader ICL mechanisms such as In-Context Reinforcement Learning. More sophisticated validations on real-world datasets and environments are desirable in the future.

## Acknowledgments and Disclosure of Funding

This project is supported by Longgang District Shenzhen’s “Ten Action Plan” for supporting innovation projects (under Grant LGKCSDPT2024002).

## References

- Ha and Schmidhuber [2018a]
David R Ha and Jürgen Schmidhuber.
World models.
*ArXiv*, abs/1803.10122, 2018a.
URL [https://api.semanticscholar.org/CorpusID:4807711](https://api.semanticscholar.org/CorpusID:4807711).
- Hafner et al. [2025]
Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap.
Mastering diverse control tasks through world models.
*Nature*, pages 1–7, 2025.
- Zhang et al. [2023a]
Weipu Zhang, Gang Wang, Jian Sun, Yetian Yuan, and Gao Huang.
Storm: Efficient stochastic transformer based world models for reinforcement learning.
*Advances in Neural Information Processing Systems*, 36:27147–27166, 2023a.
- Mazzaglia et al. [2024]
Pietro Mazzaglia, Tim Verbelen, Bart Dhoedt, Aaron Courville, and Sai Rajeswar.
Genrl: Multimodal-foundation world models for generalization in embodied agents.
*Advances in Neural Information Processing Systems*, 37:27529–27555, 2024.
- Samsami et al. [2024]
Mohammad Reza Samsami, Artem Zholus, Janarthanan Rajendran, and Sarath Chandar.
Mastering memory tasks with world models.
In *The Twelfth International Conference on Learning Representations*, 2024.
- Alonso et al. [2024]
Eloi Alonso, Adam Jelley, Vincent Micheli, Anssi Kanervisto, Amos J Storkey, Tim Pearce, and François Fleuret.
Diffusion for world modeling: Visual details matter in atari.
*Advances in Neural Information Processing Systems*, 37:58757–58791, 2024.
- Bar et al. [2025]
Amir Bar, Gaoyue Zhou, Danny Tran, Trevor Darrell, and Yann LeCun.
Navigation world models.
In *Proceedings of the Computer Vision and Pattern Recognition Conference (CVPR)*, pages 15791–15801, June 2025.
- Mendonca et al. [2021]
Russell Mendonca, Oleh Rybkin, Kostas Daniilidis, Danijar Hafner, and Deepak Pathak.
Discovering and achieving goals via world models.
In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan, editors, *Advances in Neural Information Processing Systems*, volume 34, pages 24379–24391. Curran Associates, Inc., 2021.
URL [https://proceedings.neurips.cc/paper_files/paper/2021/file/cc4af25fa9d2d5c953496579b75f6f6c-Paper.pdf](https://proceedings.neurips.cc/paper_files/paper/2021/file/cc4af25fa9d2d5c953496579b75f6f6c-Paper.pdf).
- Koh et al. [2021]
Jing Yu Koh, Honglak Lee, Yinfei Yang, Jason Baldridge, and Peter Anderson.
Pathdreamer: A world model for indoor navigation.
In *Proceedings of the IEEE/CVF International Conference on Computer Vision*, pages 14738–14748, 2021.
- Duan et al. [2024]
Yuanlin Duan, Wensen Mao, and He Zhu.
Learning world models for unconstrained goal navigation.
In *The Thirty-eighth Annual Conference on Neural Information Processing Systems*, 2024.
URL [https://openreview.net/forum?id=aYqTwcDlCG](https://openreview.net/forum?id=aYqTwcDlCG).
- Liu et al. [2025]
Zeyuan Liu, Ziyu Huan, Xiyao Wang, Jiafei Lyu, Jian Tao, Xiu Li, Furong Huang, and Huazhe Xu.
World models with hints of large language models for goal achieving.
In Luis Chiruzzo, Alan Ritter, and Lu Wang, editors, *Proceedings of the 2025 Conference of the Nations of the Americas Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume 1: Long Papers)*, pages 50–72, Albuquerque, New Mexico, April 2025. Association for Computational Linguistics.
ISBN 979-8-89176-189-6.
URL [https://aclanthology.org/2025.naacl-long.3/](https://aclanthology.org/2025.naacl-long.3/).
- Hu et al. [2023]
Anthony Hu, Lloyd Russell, Hudson Yeo, Zak Murez, George Fedoseev, Alex Kendall, Jamie Shotton, and Gianluca Corrado.
Gaia-1: A generative world model for autonomous driving.
*ArXiv*, abs/2309.17080, 2023.
URL [https://api.semanticscholar.org/CorpusID:263310665](https://api.semanticscholar.org/CorpusID:263310665).
- Russell et al. [2025]
Lloyd Russell, Anthony Hu, Lorenzo Bertoni, George Fedoseev, Jamie Shotton, Elahe Arani, and Gianluca Corrado.
Gaia-2: A controllable multi-view generative world model for autonomous driving.
*arXiv preprint arXiv:2503.20523*, 2025.
- Gao et al. [2024]
Shenyuan Gao, Jiazhi Yang, Li Chen, Kashyap Chitta, Yihang Qiu, Andreas Geiger, Jun Zhang, and Hongyang Li.
Vista: A generalizable driving world model with high fidelity and versatile controllability.
In A. Globerson, L. Mackey, D. Belgrave, A. Fan, U. Paquet, J. Tomczak, and C. Zhang, editors, *Advances in Neural Information Processing Systems*, volume 37, pages 91560–91596. Curran Associates, Inc., 2024.
URL [https://proceedings.neurips.cc/paper_files/paper/2024/file/a6a066fb44f2fe0d36cf740c873b8890-Paper-Conference.pdf](https://proceedings.neurips.cc/paper_files/paper/2024/file/a6a066fb44f2fe0d36cf740c873b8890-Paper-Conference.pdf).
- Zhang et al. [2024a]
Lunjun Zhang, Yuwen Xiong, Ze Yang, Sergio Casas, Rui Hu, and Raquel Urtasun.
Copilot4d: Learning unsupervised world models for autonomous driving via discrete diffusion.
In *The Twelfth International Conference on Learning Representations*, 2024a.
URL [https://openreview.net/forum?id=Psl75UCoZM](https://openreview.net/forum?id=Psl75UCoZM).
- Wang et al. [2024a]
Yuqi Wang, Jiawei He, Lue Fan, Hongxin Li, Yuntao Chen, and Zhaoxiang Zhang.
Driving into the future: Multiview visual forecasting and planning with world model for autonomous driving.
In *Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition*, pages 14749–14759, 2024a.
- Zhang et al. [2023b]
Zhejun Zhang, Alexander Liniger, Dengxin Dai, Fisher Yu, and Luc Van Gool.
Trafficbots: Towards world models for autonomous driving simulation and motion prediction.
In *2023 IEEE International Conference on Robotics and Automation (ICRA)*, pages 1522–1529, 2023b.
doi: 10.1109/ICRA48891.2023.10161243.
- Wu et al. [2023]
Philipp Wu, Alejandro Escontrela, Danijar Hafner, Pieter Abbeel, and Ken Goldberg.
Daydreamer: World models for physical robot learning.
In *Conference on robot learning*, pages 2226–2240. PMLR, 2023.
- Hansen et al. [2024]
Nicklas Hansen, Hao Su, and Xiaolong Wang.
TD-MPC2: Scalable, robust world models for continuous control.
In *The Twelfth International Conference on Learning Representations*, 2024.
URL [https://openreview.net/forum?id=Oxh5CstDJU](https://openreview.net/forum?id=Oxh5CstDJU).
- Pang et al. [2025]
Jing-Cheng Pang, Nan Tang, Kaiyuan Li, Yuting Tang, Xin-Qiang Cai, Zhen-Yu Zhang, Gang Niu, Masashi Sugiyama, and Yang Yu.
Learning view-invariant world models for visual robotic manipulation.
In *The Thirteenth International Conference on Learning Representations*, 2025.
URL [https://openreview.net/forum?id=vJwjWyt4Ed](https://openreview.net/forum?id=vJwjWyt4Ed).
- Zhou et al. [2024]
Siyuan Zhou, Yilun Du, Jiaben Chen, Yandong Li, Dit-Yan Yeung, and Chuang Gan.
Robodreamer: learning compositional world models for robot imagination.
ICML’24. JMLR.org, 2024.
- Barcellona et al. [2025]
Leonardo Barcellona, Andrii Zadaianchuk, Davide Allegro, Samuele Papa, Stefano Ghidoni, and Efstratios Gavves.
Dream to manipulate: Compositional world models empowering robot imitation learning with imagination.
In *The Thirteenth International Conference on Learning Representations*, 2025.
URL [https://openreview.net/forum?id=3RSLW9YSgk](https://openreview.net/forum?id=3RSLW9YSgk).
- Vorhees and Williams [2014]
Charles V Vorhees and Michael T Williams.
Assessing spatial learning and memory in rodents.
*ILAR journal*, 55(2):310–332, 2014.
- Rao and Ballard [1999]
Rajesh PN Rao and Dana H Ballard.
Predictive coding in the visual cortex: a functional interpretation of some extra-classical receptive-field effects.
*Nature neuroscience*, 2(1):79–87, 1999.
- Salvatori et al. [2023]
Tommaso Salvatori, Ankur Mali, Christopher L Buckley, Thomas Lukasiewicz, Rajesh PN Rao, Karl Friston, and Alexander Ororbia.
A survey on brain-inspired deep learning via predictive coding.
*arXiv preprint arXiv:2308.07870*, 2023.
- Brown et al. [2020]
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 are few-shot learners.
*Advances in neural information processing systems*, 33:1877–1901, 2020.
- Forrester [1958]
Jay Wright Forrester.
Industrial dynamics.
*Harvard Business Review*, 36(4):37–66, 1958.
- Sutton [1990]
Richard S Sutton.
Integrated architectures for learning, planning, and reacting based on approximating dynamic programming.
In *Machine learning proceedings 1990*, pages 216–224. Elsevier, 1990.
- Battaglia et al. [2013]
Peter W Battaglia, Jessica B Hamrick, and Joshua B Tenenbaum.
Simulation as an engine of physical scene understanding.
*Proceedings of the National Academy of Sciences*, 110(45):18327–18332, 2013.
- Ha and Schmidhuber [2018b]
David Ha and Jürgen Schmidhuber.
Recurrent world models facilitate policy evolution.
*Advances in neural information processing systems*, 31, 2018b.
- LeCun [2022]
Yann LeCun.
A path towards autonomous machine intelligence.
*Open Review*, 2022.
- Finn and Levine [2017]
Chelsea Finn and Sergey Levine.
Deep visual foresight for planning robot motion.
In *2017 IEEE International Conference on Robotics and Automation (ICRA)*, pages 2786–2793, 2017.
- Schrittwieser et al. [2020]
Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, et al.
Mastering atari, go, chess and shogi by planning with a learned model.
*Nature*, 588(7839):604–609, 2020.
- Hafner et al. [2019a]
Danijar Hafner, Timothy Lillicrap, Jimmy Ba, et al.
Dream to control: Learning behaviors by latent imagination.
*arXiv preprint arXiv:1912.01603*, 2019a.
- Sutton [1991]
Richard S Sutton.
Dyna, an integrated architecture for learning, planning, and reacting based on approximating dynamic programming.
*ACM Sigart Bulletin*, 2(4):160–163, 1991.
- Kaiser et al. [2020]
Łukasz Kaiser, Mohammad Babaeizadeh, Piotr Miłos, Błażej Osiński, Roy H Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, et al.
Model based reinforcement learning for atari.
In *International Conference on Learning Representations*, 2020.
- Jaderberg et al. [2016]
Max Jaderberg, Volodymyr Mnih, Wojciech Marian Czarnecki, Tom Schaul, Joel Z Leibo, David Silver, and Koray Kavukcuoglu.
Reinforcement learning with unsupervised auxiliary tasks.
*arXiv preprint arXiv:1611.05397*, 2016.
- Hu et al. [2022]
Anthony Hu, Gianluca Corrado, Nicolas Griffiths, Zak Murez, Corina Gurau, Hudson Yeo, Alex Kendall, Roberto Cipolla, and Jamie Shotton.
Model-based imitation learning for urban driving.
In *Proceedings of the 36th International Conference on Neural Information Processing Systems*, NIPS ’22, Red Hook, NY, USA, 2022. Curran Associates Inc.
ISBN 9781713871088.
- Zhang et al. [2024b]
Yumeng Zhang, Shi Gong, Kaixin Xiong, Xiaoqing Ye, Xiao Tan, Fan Wang, Jizhou Huang, Hua Wu, and Haifeng Wang.
Bevworld: A multimodal world model for autonomous driving via unified bev latent space.
*arXiv preprint arXiv:2407.05679*, 2024b.
- Goodfellow et al. [2014]
Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, et al.
Generative adversarial nets.
In *Advances in Neural Information Processing Systems*, pages 2672–2680, 2014.
- Kingma and Welling [2013]
Diederik P Kingma and Max Welling.
Auto-encoding variational bayes.
*arXiv preprint arXiv:1312.6114*, 2013.
- Wang et al. [2024b]
ZiRui Wang, Yue Deng, Junfeng Long, and Yin Zhang.
Parallelizing model-based reinforcement learning over the sequence length.
*Advances in Neural Information Processing Systems*, 37:131398–131433, 2024b.
- Hafner et al. [2019b]
Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson.
Learning latent dynamics for planning from pixels.
In *International conference on machine learning*, pages 2555–2565. PMLR, 2019b.
- Garrido et al. [2024]
Quentin Garrido, Mahmoud Assran, Nicolas Ballas, Adrien Bardes, Laurent Najman, and Yann LeCun.
Learning and leveraging world models in visual representation learning.
*arXiv preprint arXiv:2403.00504*, 2024.
- Li et al. [2024]
Yingyan Li, Lue Fan, Jiawei He, Yuqi Wang, Yuntao Chen, Zhaoxiang Zhang, and Tieniu Tan.
Enhancing end-to-end autonomous driving with latent world model.
*arXiv preprint arXiv:2406.08481*, 2024.
- Poudel et al. [2023]
Rudra P. K. Poudel, Harit Pandya, Chao Zhang, and Roberto Cipolla.
Langwm: Language grounded world model.
*ArXiv*, abs/2311.17593, 2023.
URL [https://api.semanticscholar.org/CorpusID:265498747](https://api.semanticscholar.org/CorpusID:265498747).
- Nie et al. [2025]
Dujun Nie, Xianda Guo, Yiqun Duan, Ruijun Zhang, and Long Chen.
Wmnav: Integrating vision-language models into world models for object goal navigation, 2025.
URL [https://arxiv.org/abs/2503.02247](https://arxiv.org/abs/2503.02247).
- Liu et al. [2023]
Rui Liu, Xiaohan Wang, Wenguan Wang, and Yi Yang.
Bird’s-eye-view scene graph for vision-language navigation.
In *Proceedings of the IEEE/CVF International Conference on Computer Vision*, pages 10968–10980, 2023.
- Dohare et al. [2024]
Shibhansh Dohare, J Fernando Hernandez-Garcia, Qingfeng Lan, Parash Rahman, A Rupam Mahmood, and Richard S Sutton.
Loss of plasticity in deep continual learning.
*Nature*, 632(8026):768–774, 2024.
- Santoro et al. [2016]
Adam Santoro, Sergey Bartunov, Matthew Botvinick, Daan Wierstra, and Timothy Lillicrap.
Meta-learning with memory-augmented neural networks.
In *International conference on machine learning*, pages 1842–1850. PMLR, 2016.
- Duan et al. [2016]
Yan Duan, John Schulman, Xi Chen, Peter L Bartlett, Ilya Sutskever, and Pieter Abbeel.
Rl^2: Fast reinforcement learning via slow reinforcement learning.
*arXiv preprint arXiv:1611.02779*, 2016.
- Akyürek et al. [2024]
Ekin Akyürek, Bailin Wang, Yoon Kim, and Jacob Andreas.
In-context language learning: architectures and algorithms.
In *Proceedings of the 41st International Conference on Machine Learning*, pages 787–812, 2024.
- Raventós et al. [2023]
Allan Raventós, Mansheej Paul, Feng Chen, and Surya Ganguli.
Pretraining task diversity and the emergence of non-bayesian in-context learning for regression.
*Advances in neural information processing systems*, 36:14228–14246, 2023.
- Xie et al. [2024]
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*, 2024.
- Laskin et al. [2023]
Michael Laskin, Luyu Wang, Junhyuk Oh, Emilio Parisotto, Stephen Spencer, Richie Steigerwald, DJ Strouse, Steven Stenberg Hansen, Angelos Filos, Ethan Brooks, et al.
In-context reinforcement learning with algorithm distillation.
In *The Eleventh International Conference on Learning Representations*, 2023.
- Lee et al. [2023]
Jonathan Lee, Annie Xie, Aldo Pacchiano, Yash Chandak, Chelsea Finn, Ofir Nachum, and Emma Brunskill.
Supervised pretraining can learn in-context reinforcement learning.
*Advances in Neural Information Processing Systems*, 36:43057–43083, 2023.
- Wang et al. [2025]
Fan Wang, Pengtao Shao, Yiming Zhang, Bo Yu, Shaoshan Liu, Ning Ding, Yang Cao, Yu Kang, and Haifeng Wang.
Towards large-scale in-context reinforcement learning by meta-training in randomized worlds, 2025.
URL [https://arxiv.org/abs/2502.02869](https://arxiv.org/abs/2502.02869).
- Anand et al. [2022]
Ankesh Anand, Jacob C Walker, Yazhe Li, Eszter Vértes, Julian Schrittwieser, Sherjil Ozair, Theophane Weber, and Jessica B Hamrick.
Procedural generalization by planning with self-supervised world models.
In *International Conference on Learning Representations*, 2022.
URL [https://openreview.net/forum?id=FmBegXJToY](https://openreview.net/forum?id=FmBegXJToY).
- Gupta et al. [2024]
Sharut Gupta, Chenyu Wang, Yifei Wang, Tommi Jaakkola, and Stefanie Jegelka.
In-context symmetries: Self-supervised learning through contextual world models.
In A. Globerson, L. Mackey, D. Belgrave, A. Fan, U. Paquet, J. Tomczak, and C. Zhang, editors, *Advances in Neural Information Processing Systems*, volume 37, pages 104250–104280. Curran Associates, Inc., 2024.
URL [https://proceedings.neurips.cc/paper_files/paper/2024/file/bcad07d4bfab51243efaa08b8ed475b3-Paper-Conference.pdf](https://proceedings.neurips.cc/paper_files/paper/2024/file/bcad07d4bfab51243efaa08b8ed475b3-Paper-Conference.pdf).
- Lior et al. [2024]
Gili Lior, Yuval Shalev, Gabriel Stanovsky, and Ariel Goldstein.
Computation or weight adaptation? rethinking the role of plasticity in learning.
*bioRxiv*, pages 2024–03, 2024.
- Pan et al. [2023]
Jane Pan, Tianyu Gao, Howard Chen, and Danqi Chen.
What in-context learning “learns” in-context: Disentangling task recognition and task learning.
In *Findings of the Association for Computational Linguistics, ACL 2023*, Proceedings of the Annual Meeting of the Association for Computational Linguistics, pages 8298–8319. Association for Computational Linguistics (ACL), 2023.
doi: 10.18653/v1/2023.findings-acl.527.
Publisher Copyright: © 2023 Association for Computational Linguistics.; 61st Annual Meeting of the Association for Computational Linguistics, ACL 2023 ; Conference date: 09-07-2023 Through 14-07-2023.
- Chan et al. [2022]
Stephanie Chan, Adam Santoro, Andrew Lampinen, Jane Wang, Aaditya Singh, Pierre Richemond, James McClelland, and Felix Hill.
Data distributional properties drive emergent in-context learning in transformers.
*Advances in neural information processing systems*, 35:18878–18891, 2022.
- Chan et al. [2025]
Bryan Chan, Xinyi Chen, András György, and Dale Schuurmans.
Toward understanding in-context vs. in-weight learning.
In *The Thirteenth International Conference on Learning Representations*, 2025.
URL [https://openreview.net/forum?id=aKJr5NnN8U](https://openreview.net/forum?id=aKJr5NnN8U).
- Kirsch et al. [2022]
Louis Kirsch, James Harrison, Jascha Sohl-Dickstein, and Luke Metz.
General-purpose in-context learning by meta-learning transformers.
In *Sixth Workshop on Meta-Learning at the Conference on Neural Information Processing Systems*, 2022.
- Singh et al. [2023]
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:27801–27819, 2023.
- Kingma and Welling [2014]
Diederik P Kingma and Max Welling.
Auto-encoding variational bayes.
In *2nd International Conference on Learning Representations*, 2014.
- Wang et al. [2020]
Sinong Wang, Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma.
Linformer: Self-attention with linear complexity.
*arXiv preprint arXiv:2006.04768*, 2020.
- Yang et al. [2024]
Songlin Yang, Bailin Wang, Yikang Shen, Rameswar Panda, and Yoon Kim.
Gated linear attention transformers with hardware-efficient training.
In *International Conference on Machine Learning*, pages 56501–56523. PMLR, 2024.
- Zhang et al. [2024c]
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.
*Advances in Neural Information Processing Systems*, 37:116870–116898, 2024c.
- Pašukonis et al. [2023]
Jurgis Pašukonis, Timothy P Lillicrap, and Danijar Hafner.
Evaluating long-term memory in 3d mazes.
In *The Eleventh International Conference on Learning Representations*, 2023.
- Wang et al. [2024c]
Fan Wang, Chuan Lin, Yang Cao, and Yu Kang.
Benchmarking general-purpose in-context learning.
*arXiv preprint arXiv:2405.17234*, 2024c.
- Ehsani et al. [2024]
Kiana Ehsani, Tanmay Gupta, Rose Hendrix, Jordi Salvador, Luca Weihs, Kuo-Hao Zeng, Kunal Pratap Singh, Yejin Kim, Winson Han, Alvaro Herrasti, et al.
Spoc: Imitating shortest paths in simulation enables effective navigation and manipulation in the real world.
In *Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition*, pages 16238–16250, 2024.
- Kolve et al. [2017]
Eric Kolve, Roozbeh Mottaghi, Winson Han, Eli VanderBilt, Luca Weihs, Alvaro Herrasti, Matt Deitke, Kiana Ehsani, Daniel Gordon, Yuke Zhu, et al.
Ai2-thor: An interactive 3d environment for visual ai.
*arXiv preprint arXiv:1712.05474*, 2017.
- Deitke et al. [2022]
Matt Deitke, Eli VanderBilt, Alvaro Herrasti, Luca Weihs, Kiana Ehsani, Jordi Salvador, Winson Han, Eric Kolve, Aniruddha Kembhavi, and Roozbeh Mottaghi.
Procthor: Large-scale embodied ai using procedural generation.
*Advances in Neural Information Processing Systems*, 35:5982–5994, 2022.
- Wang et al. [2024d]
Yian Wang, Xiaowen Qiu, Jiageng Liu, Zhehuan Chen, Jiting Cai, Yufei Wang, Tsun-Hsuan Johnson Wang, Zhou Xian, and Chuang Gan.
Architect: Generating vivid and interactive 3d scenes with hierarchical 2d inpainting.
*Advances in Neural Information Processing Systems*, 37:67575–67603, 2024d.
- Paszke et al. [2019]
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf, Edward Yang, Zach DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala.
*PyTorch: an imperative style, high-performance deep learning library*.
Curran Associates Inc., Red Hook, NY, USA, 2019.
- Rousseeuw [1987]
Peter J Rousseeuw.
Silhouettes: a graphical aid to the interpretation and validation of cluster analysis.
*Journal of computational and applied mathematics*, 20:53–65, 1987.

## Appendix A Assumptions and Proofs for Theorem ˜ 1

Proof for the first part of [Theorem˜1](https://arxiv.org/html/2509.22353v1#Thmtheorem1): We first define

$$ $\displaystyle\Delta(e_{1},e_{2})$ $\displaystyle=\mathbb{E}_{q}D_{KL}(\hat{p}_{\theta,e_{1}}(\cdot|q),\hat{p}_{\theta,e_{2}}(\cdot|q))$ $\displaystyle\kappa(e_{1},e_{2})$ $\displaystyle=\max_{q}D_{KL}(\hat{p}_{\theta,e_{1}}(\cdot|q),\hat{p}_{\theta,e_{2}}(\cdot|q))$ $\displaystyle\hat{e}_{0}$ $\displaystyle=argmin_{e\in\mathcal{E}}\Delta(e,e_{0})$ $\displaystyle\kappa_{i}$ $\displaystyle=\inf_{e\in\mathcal{E},e\neq\hat{e}_{0}}\kappa(e,\hat{e}_{0})$ $\displaystyle\kappa_{s}$ $\displaystyle=\sup_{e\in\mathcal{E},e\neq\hat{e}_{0}}\kappa(e,\hat{e}_{0})$ $\displaystyle p_{e}(C_{T})$ $\displaystyle=\prod_{(q_{t},o_{t+1})\in C_{T}}p_{e}(o_{t+1}|q_{t})$ $$

We then make the following Assumptions:

- •
Queries $q_{t}=(s_{t},a_{t})$ are sampled i.i.d. from a distribution $\mu(s,a)$ with $\mu(s,a)\equiv\frac{1}{|S||A|}$.
- •
$\kappa(e_{1},e_{2})=\alpha(e_{1},e_{2})^{2}\Delta(e_{1},e_{2})$ with $\alpha(e_{1},e_{2})=\sqrt{\frac{\kappa(e_{1},e_{2})}{\Delta(e_{1},e_{2})}},\alpha>1$. We further define $\alpha=\max_{e_{1},e_{2}}\alpha(e_{1},e_{2})$, so that $\Delta(e_{1},e_{2})\geq\frac{\kappa(e_{1},e_{2})}{\alpha^{2}}$. Note that $\alpha$ can be interpreted as the measure of ”non-uniformity” between any two environments in the set: it attains maximum when the two environments are almost identical yet differ significantly at only a few positions, and approaches $1$ when the environments are either completely different or exactly the same.
- •
The environment recognizer selects closest task $\hat{e}$ by $\hat{e}=argmax_{e\in\mathcal{E}}p_{e}(C_{T})$

First, we gave that

$$ $\displaystyle TV(\hat{p}_{ER},p_{e_{0}})=TV(\hat{p}_{\theta,\hat{e}},p_{e_{0}})$ $\displaystyle\leq TV(\hat{p}_{\theta,\hat{e}},\hat{p}_{\theta,\hat{e}_{0}})+TV(\hat{p}_{\theta,\hat{e}_{0}},p_{e_{0}})$ $\displaystyle=TV(\hat{p}_{\theta,\hat{e}},\hat{p}_{\theta,\hat{e}_{0}})+\min_{e\in\mathcal{E}}TV(\hat{p}_{\theta,e},p_{e_{0}})$ $$

We then estimate the first term and use the Chernoff bound for derivation:

$$ $\displaystyle TV(\hat{p}_{\theta,\hat{e}},\hat{p}_{\theta,\hat{e}_{0}})$ $\displaystyle=\sum_{e\in\mathcal{E},e\neq\hat{e}_{0}}p(p_{e}(C_{T})>p_{\hat{e}_{0}}(C_{T}))TV(\hat{p}_{\theta,e},\hat{p}_{\theta,\hat{e}_{0}})$ $\displaystyle\leq\sum_{e\in\mathcal{E},e\neq\hat{e}_{0}}\alpha\cdot\underbrace{\exp(-T\cdot\Delta(e,\hat{e}_{0}))}_{\text{Chernoff bound}}\underbrace{\sqrt{1/2\Delta(e,\hat{e}_{0})}}_{\text{Pinsker's Inequality}}$ $\displaystyle<\sum_{e\in\mathcal{E},e\neq\hat{e}_{0}}\underbrace{\frac{\alpha}{2\sqrt{e\cdot T}}}_{\text{achieved maximum when }T\cdot\Delta(e,\hat{e}_{0})=1/2}$ $\displaystyle<\frac{\alpha(|\mathcal{E}|-1)}{3\sqrt{T}}$ $$

Proof for the second part of [Theorem˜1](https://arxiv.org/html/2509.22353v1#Thmtheorem1). We keep the aforementioned assumptions that the distribution of the context $C_{T}$ is uniform on the state and action space. We denote $n(s,a)$ as times of appearance of $(s,a)$ in $C_{T}$. It is first straightforward to prove with Hoeffding’s inequality that with probability of $1-\delta/2$,

$$ $\displaystyle\frac{T}{|S||A|}-\sqrt{\frac{T\log{(4|S||A|/\delta)}}{2}}\leq n(s,a\in C_{T})\leq\frac{T}{|S||A|}+\sqrt{\frac{T\log{(4|S||A|/\delta)}}{2}}\quad\forall s,a$ $$

Then, by add the constraint $T>4|S|^{2}|A|^{2}\log(4|S||A|/\delta)$, with at least probability of $1-\delta/2$,

$$ $\displaystyle n(s,a)>\frac{T}{2|S||A|}$ (6) $$

To estimate $\hat{p}_{EL}(s^{\prime}|s,a)$, we use [Equation˜4](https://arxiv.org/html/2509.22353v1#S3.E4) to acquire:

$$ $\displaystyle\hat{p}_{\theta,EL}(o_{t+1}|q_{t},C_{T})$ $\displaystyle=\frac{p(q_{t},o_{t+1}|C_{T})}{p(q_{t}|C_{T})}$ $\displaystyle=\frac{\sum_{s}p(s,a_{t},o_{t+1}|C_{T})p(s|q_{t})}{\sum_{s}p(s,a_{t}|C_{T})p(s|q_{t})}$ (7) $$

Now employ [Equation˜6](https://arxiv.org/html/2509.22353v1#A1.E6) and [Equation˜7](https://arxiv.org/html/2509.22353v1#A1.E7) in the calculation of TV:

$$ $\displaystyle TV(\hat{p}_{EL},p_{e_{0}})$ $\displaystyle=TV(\frac{\sum_{s}p(s,a_{t},o_{t+1}|C_{T})p(s|q_{t})}{\sum_{s}p(s,a_{t}|C_{T})p(s|q_{t})},\frac{\sum_{s}p(s,a_{t},o_{t+1})p(s|q_{t})}{\sum_{s}p(s,a_{t})p(s|q_{t})})$ $\displaystyle\leq\max_{s}TV(\frac{n(s,a_{t},o_{t+1}\in C_{T})}{n(s,a_{t}\in C_{T})},\frac{p(s,a_{t},o_{t+1})}{p(s,a_{t})})$ $\displaystyle\leq\underbrace{\sqrt{\frac{|O|log(4|O|/\delta)}{n(s,a_{t}\in C_{T})}}}_{\text{Hoeffding's inequality}}\text{ with probability at least }1-\delta/2$ $\displaystyle<\sqrt{\frac{2|O||S||A|log(4|S|/\delta)}{T}}\text{ with probability at least }1-\delta$ $$

This finishes the proof of [Theorem˜1](https://arxiv.org/html/2509.22353v1#Thmtheorem1)

## Appendix B Additional Experiment Settings

### B.1 Environments and Datasets

Figure: Figure 5: Configuration scopes and cases of random Cart-Poles (upper left), t-SNE visualization of the configuration distribution (upper right), and a list of training and evaluation datasets.
Refer to caption: x5.png

Cart-pole: [Figure˜5](https://arxiv.org/html/2509.22353v1#A2.F5) lists the scopes of the Cart-Pole environment variants, their t-SNE visualization, and details of the training and evaluation data. All datasets share a trajectory length of 200, which is sufficient for ICL in Cart-Pole variants. The training data comprises a total of 25.6M timesteps to avoid interference from data scale on performance, and the evaluation data contains 205K steps in total.

Figure: Figure 6: bird’s-eye view of 128 procedurally generated mazes in the Homogeneous dataset.
Refer to caption: x7.png

Figure: Figure 7: Distribution of the configurations (area, occupancy rate, and number of objects in each scene) in the $32K$ procedurally generated mazes
Refer to caption: x8.png

Mazes: Our maze environment is closely related to the settings described in Pašukonis et al. [70], Wang et al. [75], which are generated on a $15\times 15$ grid world. The distribution of the sampled mazes is shown in [Figure˜6](https://arxiv.org/html/2509.22353v1#A2.F6) and [Figure˜7](https://arxiv.org/html/2509.22353v1#A2.F7). The primary distinction between our mazes and those in previous work is the enhanced environmental diversity achieved through randomized configurations, which include the following:

- •
The textures of the ceiling, ground, and walls are randomly selected from a collection of $87$ real-world textures.
- •
The scale of each grid varies from $1.5\,\text{m}$ to $4.5\,\text{m}$, and the indoor height ranges from $2\,\text{m}$ to $6\,\text{m}$.
- •
The ground clearance and the field of view (FOV) of the camera are varied between $[1.6m,2.0m]$ and $[0.3\pi,0.8\pi]$, respectively.
- •
Two-wheeled dynamics are employed for the embodiment.
- •
Each environment involves $[5,15]$ objects, each marked with a crossable, translucent light wall of a different color.
- •
The agent receives a reward of 1.0 when reaching the goal, and a negative reward for collisions with walls, which is dependent on the agent’s speed.

ProcTHOR: We sample $336$ training and $256$ evaluating houses from the ProcTHOR-10K dataset. Unless otherwise specified, we keep both the trajectories and the environments of the validation datasets in ProcTHOR and Mazes separate from those of the training datasets. In seen-task validation, the environments have overlaps, but the trajectories are resampled.

Figure: (a) Procedurally Generated Mazes
Refer to caption: figures/dataset_environment_maze.png

### B.2 Details of Model Structures

The framework consists of the following modules:

- •
Observation Encoder (Image): A convolutional encoder that processes 128×128 images into a $D$-dimensional latent vector through 9 convolution layers and 8 residual blocks.
- •
Observation Decoder (Image): A deconvolutional decoder that reconstructs 128×128 images from a $D$-dimensional latent vector through 1 convolution layer, 8 residual blocks, and 3 transposed-convolution layers.
- •
Observation Encoder (Raw): A linear layer mapping from hidden size of $4$ to $D$.
- •
Observation Decoder (Raw): A linear layer mapping from hidden size of $D$ to $4$.
- •
Latent Decoder: A 1-layer MLP with input size $D$, hidden size $D$, layer normalization, and residual connections.
- •
Action Encoder: A 1-layer MLP that encodes discrete actions into a $D$-dimensional hidden state.
- •
Sequence Decoder: A gated self-attention architecture with $L$ layers, hidden size $D$, inner hidden size $D$, $d$ attention heads, memory length $l_{M}$, layer normalization, and block recurrence.

The hyper-parameters are specified as [Table˜4](https://arxiv.org/html/2509.22353v1#A2.T4).

**Table 4: Model architecture and hyper-parameters for the two classes of environments**
| Environments | Encoder & Decoder | $D$ | $L$ | $d$ | $l_{M}$ |
| --- | --- | --- | --- | --- | --- |
| Cart-pole | Raw | 128 | 4 | 4 | 64 |
| Navigation | Image | 1024 | 18 | 32 | 64 |

### B.3 Details of Training

All models were trained on NVIDIA A800 GPUs using the AdamW optimizer with the default settings in PyTorch [76].

Cart-pole training details.
We train the model with a per‑GPU batch size of 128, an epoch of 100, an initial learning rate of $\mathrm{1.0e}^{-4}$ and decayed to $\mathrm{2.04e}^{-5}$.

Maze pre-training details.
We first train the Image Encoder and Image Decoder with a per‑GPU batch size of 400, an epoch of 50, and an initial learning rate of $\mathrm{3e}^{-4}$. Subsequently, we train all model components with a per‑GPU batch size of 10, an epoch of 10, an initial learning rate of $\mathrm{2.0e}^{-4}$, and decayed to $\mathrm{8.8e}^{-5}$.

ProcTHOR training details.
We first train the Image Encoder and Image Decoder using the same settings as in Maze pretraining, except for the initial learning rate, which is set to $\mathrm{2.0e}^{-4}$. The VAE is then frozen, while the remaining modules are initialized from the weights trained on the Maze dataset and further fine-tuned on ProcThor data. For the 40k version of ProcThor, we train for 10 epochs with a per-GPU batch size of 10 and an initial learning rate of $\mathrm{2.0e}^{-4}$, decayed to $\mathrm{1.2e}^{-4}$. For the 5k version, we use the same settings but train for 20 epochs.

## Appendix C Additional Results

Additional reults of the performances of L2World in Mazes.
[Table˜5](https://arxiv.org/html/2509.22353v1#A3.T5) presents the prediction error (measured as the mean square error between observations) of L2World and the other baselines on the Maze datasets. The prediction errors are fully consistent with the PSNR evaluation; therefore, we report only the PSNR results in the remaining experiments.
[Table˜6](https://arxiv.org/html/2509.22353v1#A3.T6) details the K-step prediction performance, corresponding to the histograms in [Figure˜3](https://arxiv.org/html/2509.22353v1#S4.F3).

**Table 5: 1-step prediction error($\downarrow$) of different world models in Mazes.**
| Model | Seen | Unseen |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| T=1 | T=10 | T=100 | T=1000 | T=10000 | T=1 | T=10 | T=100 | T=1000 | T=10000 |  |
| L2World (Maze-32K-L) | 2.09e-02 | 8.01e-03 | 4.89e-03 | 3.43e-03 | 3.13e-03 | 2.30e-02 | 7.51e-03 | 4.82e-03 | 3.42e-03 | 3.42e-03 |
| L2World (Maze-32K-S) | 1.39e-02 | 1.18e-02 | 1.08e-02 | 9.53e-03 | 8.96e-03 | 1.43e-02 | 1.19e-02 | 1.09e-02 | 9.35e-03 | 9.31e-03 |
| L2World (Maze-128-S) | 1.13e-02 | 9.15e-03 | 8.74e-03 | 6.28e-03 | 6.66e-03 | 1.58e-02 | 1.37e-02 | 1.26e-02 | 1.08e-02 | 1.09e-02 |
| L2World (Maze-128-L) | 1.40e-02 | 8.20e-03 | 4.66e-03 | 2.72e-03 | 2.51e-03 | 1.76e-02 | 1.14e-02 | 8.02e-03 | 7.01e-03 | 7.04e-03 |
| Dreamer (Maze-32K-L) | 2.29e-02 | 6.57e-03 | 1.19e-02 | 7.47e-03 | 6.47e-03 | 2.08e-02 | 8.95e-03 | 7.25e-03 | 5.43e-03 | 6.14e-03 |
| Dreamer (Maze-128-L) | 1.94e-02 | 8.63e-03 | 6.56e-03 | 6.02e-03 | 5.72e-03 | 3.75e-02 | 3.51e-02 | 3.90e-02 | 4.50e-02 | 4.46e-02 |
| NWM (Maze-32K-L) | 8.22e-03 | 9.52e-03 | 1.20e-02 | 5.86e-03 | 7.83e-03 | 2.02e-02 | 2.26e-02 | 2.48e-02 | 1.70e-02 | 1.28e-02 |

**Table 6: Average PSNR of $k$-Step auto-regressive prediction in Mazes (Unseen).**
| Model | T=1 | T=10 | T=100 | T=1000 | T=10000 |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| $k=1$ | $k=2$ | $k=4$ | $k=8$ | $k=1$ | $k=2$ | $k=4$ | $k=8$ | $k=1$ | $k=2$ | $k=4$ | $k=8$ | $k=1$ | $k=2$ | $k=4$ | $k=8$ | $k=1$ | $k=2$ | $k=4$ | $k=8$ |  |
| L2World (Maze-32K-L) | 16.94 | 15.79 | 14.74 | 13.39 | 21.02 | 19.88 | 18.96 | 15.83 | 23.24 | 22.16 | 20.85 | 18.31 | 24.60 | 23.57 | 21.33 | 18.03 | 24.80 | 24.11 | 21.81 | 18.72 |
| L2World (Maze-32K-S) | 18.44 | 17.41 | 16.82 | 15.61 | 19.23 | 19.05 | 17.92 | 16.73 | 19.62 | 19.43 | 18.19 | 16.69 | 20.29 | 20.17 | 18.89 | 17.03 | 20.31 | 20.21 | 18.83 | 17.00 |
| L2World (Maze-128-S) | 18.00 | 16.83 | 16.23 | 15.17 | 18.62 | 18.18 | 17.06 | 15.96 | 18.98 | 18.53 | 17.39 | 15.93 | 19.65 | 19.09 | 17.86 | 16.10 | 19.62 | 19.12 | 17.91 | 16.18 |
| L2World (Maze-128-L) | 17.53 | 15.99 | 15.38 | 14.53 | 19.45 | 18.73 | 17.67 | 16.64 | 20.96 | 20.47 | 19.11 | 17.57 | 21.54 | 20.76 | 19.45 | 17.48 | 21.53 | 21.05 | 19.58 | 17.65 |
| Dreamer (Maze-32K-L) | 16.16 | 15.80 | 14.62 | 13.12 | 19.66 | 18.96 | 17.72 | 16.25 | 21.89 | 20.99 | 20.33 | 17.74 | 23.13 | 22.08 | 20.57 | 17.00 | 23.16 | 21.85 | 20.23 | 16.14 |
| NWM (Maze-32K-L) | 16.20 | 13.37 | 11.68 | 11.16 | 16.71 | 13.34 | 11.65 | 11.00 | 17.00 | 13.86 | 12.13 | 11.48 | 17.37 | 13.82 | 12.24 | 12.60 | 17.85 | 14.51 | 12.60 | 12.46 |

Figure: Figure 9: We investigated the correlation between the average prediction error and the performance change. The performance change was quantified by $\Delta$PSNR. $\Delta$PSNR is acquired by the loss of accuracy by replacing the input $s_{t}$ with predicted $\hat{s}_{t}$ and then measuring the loss of PSNR in future predictions with $k=1$ and $k=8$.
Refer to caption: x9.png

ICEL and ICER in navigation world models implicitly perform global mapping. We further show that predicting transitions alone, without any specifically designed tasks, can potentially capture the global map implicitly. We collect $12$ trajectories from $4$ unseen mazes ($3$ trajectories in each maze) and track the transformation of memory states ($\phi_{t}$) across each of the linear attention layers. Among the $18$ layers used in the transition model, we select layers $\{1,6,12,18\}$ for t-SNE visualization of the memory states ($\phi_{t}$).
To ensure that the $4$ unseen mazes are not trivially discriminable, for instance, by geometric and embodiment configurations in one-shot, we maintain identical configurations for most aspects of the evaluated $4$ mazes while varying only their topology, such as the arrangement of walls. We employ silhouette scores [77] to quantify clustering quality, where higher values indicate better environment separation.
As visualized in [Figure˜10](https://arxiv.org/html/2509.22353v1#A3.F10), we highlight two insights: first, training solely on transition prediction across diverse environments and long contextual windows implicitly learns a spatial map, potentially removing the need for auxiliary mapping modules; second, ICEL and ICER behave differently across layers, suggesting distinct underlying mechanisms.

Figure: Figure 10: t-SNE visualization of the memory states $\phi_{t}$ in state transition prediction for Maze-32K-L and Maze-128-L groups. We visualize the memory states from layers 1, 6, 12, and 18. The visualization encompasses $12$ trajectories across $4$ distinct environments that are similar yet exhibit slight differences. Trajectories originating from the same environment are indicated by similar colors. The Silhouette score is computed by treating trajectories from each environment as individual classes.
Refer to caption: x10.png

Investigating ICEL from perspectives of predictive coding. The concept of predictive coding has emerged as a foundational mechanism in both biological and artificial learning systems [24], in which the discrepancy between expected and observed outcomes drives attention and learning.
To investigate the correlation between ICEL and predictive coding, we conducted experiments to examine how prediction error influences learning progress. Specifically, we selected three positions $T=\{10,100,1000\}$ for each of $256$ evaluating sequences. At these positions, we replaced the ground-truth observations $s_{t}$ with model-generated predictions, thereby suppressing error-correction signals derived from real-world feedback. We then compared the accuracy loss in subsequent frames (including $k=1$ and $k=8$) with and without this replacement. This comparison quantifies the importance of ground-truth observations for learning progress.
As illustrated in [Figure˜9](https://arxiv.org/html/2509.22353v1#A3.F9), we observed a clear positive correlation between the mean performance difference ($\Delta$PSNR) and the prediction error (MSE loss) of the replaced frame. This finding suggests that ICEL progress is sensitive to prediction error, a phenomenon reminiscent of predictive coding in biological systems. These results further validate ICL as a prospective mechanism for adapting world models to variant environments.

Figure: Figure 11: 2 failed examples produced by L2World on Mazes with $T=\{1,100,10000$} and $k=8$.
Refer to caption: x11.png

Figure: (a) Maze
Refer to caption: x12.png

Additional Cases. [Figure˜11](https://arxiv.org/html/2509.22353v1#A3.F11) shows two failure predictions of L2World (Maze-32K-L) on unseen maze environments.
These failures appear to stem from excessive blurring in the predictions, which causes the compound error to explode rapidly as $k$ increases.
Building on this observation, a natural extension of our work is to inject additional overshooting during training so that the world model can better forecast distant futures.
[Figure˜12](https://arxiv.org/html/2509.22353v1#A3.F12) juxtaposes the performance of L2World trained with 32K-L and 32K-S, clearly evidencing the benefit of longer sequences in incentivizing ICEL.

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 6
presentation: 4
questions: I have no questions for the authors.
soundness: 4
strengths: The paper is clearly written. To my knowledge, the conditions under which ICL task recognition vs task learning appear in world-model-based tasks have not yet been explored though the results are not surprising. The theoretical bounds are consistent with the experiments.
summary: The authors consider in-context learning of a world model, in particular focusing on the dichotomy between environment-learning and environment-recognition observed in other in-context learning setups. Consistent with theoretical error bounds that they derive, their main contribution is to show that long contexts and environment diversity favor environment-learning in cart-pole and procedural maze navigation tasks.
weaknesses: I can find no major weaknesses but in my opinion the paper offers few new insights compared to what is already known in the literature. The experimental paradigms considered are sufficiently complicated that it is difficult to draw specific inferences into the mechanisms that lead to ICEL versus ICER. In short, though it is certainly valuable to characterize phenomenology as the authors have done here, it is not clear what question is being answered. While this is likely beyond the scope of the current work, the authors could consider setting up the problem in a simpler, symbolic RL setup that captures the dichotomy between ICEL and ICER, and which potentially lends itself to a deeper mechanistic analysis for what factors determine ICEL vs ICER. 

Highlighting a couple of references that could be included: the ICL regression task was originally introduced by Garg et al, Neurips, 2022 (and not Raventos et al, 2023, line 106) and a minimal theoretical model that explains the origin of the ICL/IWL tradeoff was proposed by Nguyen and Reddy, ICLR 2025 (lines 111-113).

---

confidence: 3
contribution: 2
initial_rating: 4
presentation: 1
questions: Maybe worth citing these papers which directly discuss ER vs EL:
https://arxiv.org/abs/2412.01003
https://arxiv.org/abs/2503.05631
https://arxiv.org/abs/2506.17859
soundness: 3
strengths: S1: This paper extends in-context learning studies to world model learning settings where the model should learn the inference time world, and this construct an adaptive world model. 

S2: The experimental design is thorough, it covers both continuous control (cart-pole) and partial observability (navigation). They authors also careful ablate data diversity, trajectory length and environment complexity.

S3: Verifying that ICL behaves as ER and EL in linear attention models is novel.

S4: The experimental grounding of ICEL for world model learning sets a nice expectation for future experimentalists.
summary: This paper investigates in-context learning (ICL) in world model training, focusing on how models adapt to novel environments through ICL rather than parameter updates. The authors distinguish two ICL mechanisms: Environment Recognition (ER) and Environment Learning (EL) and derive error bounds characterizing when each emerges. They introduce L2World, a linear-attention architecture for long-context world modeling, and validate their theoretical predictions on cart-pole control and vision-based navigation tasks. The key findings emphasize that environment diversity and long context windows are critical for enabling EL over ER.
weaknesses: W1: Most findings on ER vs EL seems already established in the literature, for example in https://arxiv.org/abs/2412.01003 and https://arxiv.org/abs/2503.05631. It doesn't seem like these are "identified" through this study.

W2: The paper's focus is a little bit unclear. there is a theoretical part, an introduction of a new architecture, and empirical experiments on ICL, it is quite unclear what is the main presented innovation/finding.

W3: The paper discusses larger context length to further enable ICEL. While this claim is likely true, it would be nice to see equal-flop experiments as longer context models get strictly more training.

W4: The presentation of the data is poor. The plots/tables could be enhanced so that, for example, the data diversity's effect can be more directly read out.

---

confidence: 3
contribution: 3
initial_rating: 4
presentation: 2
questions: - 1) Theoretical results assume “sufficiently optimized” models and discrete environment complexity. How realistic are these assumptions for continuous control or visual domains? 

-  2) The analysis of Theorem 1 relies on comparing the bounds in (5). Can the authors comment on the tightness on these bounds? I understand that's hard to do theoretically. But the bounds can be very loose and hence any analysis based on them might not be correct. I think even an empirical analysis of the tightness is needed to solidify the analysis and the claims. 


- 3) Have you conducted any controlled experiments (e.g., context-swapping or corrupted contexts) to isolate ICEL behavior from ICER? maybe gradually corrupt the context  and report how quickly performance collapse? 

- 4)  Why were primarily pixel-level reconstruction metrics (e.g., PSNR) used, and do these metrics accurately reflect the model’s ability to support downstream decision-making or planning tasks?
soundness: 2
strengths: - The paper is easy to follow and make clear and consistent concepts. as far as know, the explicit separation between In-Context Environment Recognition (ICER) and In-Context Environment Learning (ICEL) is novel and it made the analysis more intuitive and cleaner. 

-I particularly highly appreciate the theoretical analysis, which adds rigor and depth to the work. I skimmed through the proofs and they seem correct. 

- The proposed L2World framework demonstrates the practical relevance of the theory. Experiments on well-established control tasks (cart-pole and navigation) show that world models can indeed adapt to novel environments via context, supporting the theoretical claims.
summary: his paper examines how world models can learn and adapt through context, focusing on in-context learning (ICL) within both MDP and POMDP settings. The authors distinguish between two key processes, namely In-Context Environment Learning (ICEL) and In-Context Environment Recognition (ICER). Furthermore, analyze their theoretical properties by deriving error upper bounds that explain when and how each process is effective.

To validate these insights, the paper introduces L2World, a framework designed to study in-context world modeling, and evaluates it on cart-pole and navigation tasks. The experiments show that world models can adapt to new environments through context alone and that both environmental diversity and longer context windows are crucial for enabling ICEL.
weaknesses: **W1** ICER vs ICEL:  The distinction is clear conceptually, but the empirical analysis in my perspective does not produce direct causal evidence that the model is performing learning inside context vs recognition + recall. The observed improvement in prediction performance with increasing context length T could be explained by several alternative mechanisms. For instance, the model might simply be recognizing previously encountered environment configurations or recalling stored representations from training, rather than actively updating its internal dynamics to learn new environmental behaviors on the fly. This part is not clear from the experiments. 

**W2** While the derived bounds in Theorem 1 is correct and rigorous. I think the derived claims from the bounds is a bit loose: some of the paper’s main claims rely on comparing these bounds. Upper bounds can be loose; which bound is smaller may not reflect which method has smaller actual error. I think some analysis on the tightness of the bounds is needed in order to solidify the claims.

**W3** The authors mainly report PSNR on pixel-level. This raises questions about whether the proposed framework would scale to more complex, real-world settings or provide tangible gains in downstream task and decision making performance.