Large language models (LLMs), when acting as agents, are expected to take observed data in context, infer the latent state space underlying the world, and leverage it for downstream prediction. However, prior work demonstrated that LLMs struggle to use representations learned in context on a graph tracking task, where the model needs to construct a representation of the graph governing data generation process and use it for subsequent predictions. In this paper, we show that extending this to few-shot settings, where each demonstration is generated from a different world with either the same or different graph topologies, enhances its prediction on 6 models from 4 model families. To understand this improvement, we linearly probe a low-dimensional world representation that encodes graph information in the hidden states. Notably, we find that few-shot demonstrations relocate the world representation and increase its predictive use. Specifically, for each model, these world representations shift in directions nearly orthogonal to their original subspace, and interventions on these representations selectively impair performance more than interventions on other subspaces. Consistent with this insight, we show that few-shot demonstrations with observations from different worlds improve performance on ARC-AGI-1&2, web agent tasks, and Othello. Our findings elucidate the role and internal mechanisms of few-shot demonstrations in in-context world modeling. More broadly, our work advances our understanding of how LLM agents learn from in-context observations and provides implications for their further improvement.
Figures & tables
Figure 1: (a) Conceptual Framework of In-Context World Modeling. An LLM is provided with random walks generated on a graph topology, where each node is assigned a fruit, and instructed to move two steps to the right. We probe the internal representations to determine whether the model has constructed the latent graph. A linearly decodable subspace of these representations is termed a world representation. See Figure 11 for the actual prompt. In this paper, we treat cucumbers as fruits. (b) LLMs Learned to Use World Representation with Few Shot Demonstrations. Few shot demonstrations consisting of random walks, task instructions, and answers from worlds with the same topology but different assignments of fruits to nodes improve task performance. Interventions on the world representation subspace cause larger performance degradation in the few shot setting than in the zero shot setting relative to interventions in random directions. The results are from Qwen3-4B-Base on a 4×4 Lattice Grid. (c) Few-Shot Demonstrations Shift World Representations. As the number of shots increases, the graph representation shifts orthogonally to the representation subspace. The left shows a 3D PCA of hidden states of the nodes, and the right shows a 2D PCA of mean hidden states over the random walks in the test set. The results are from Figure 6 .
Figure 2: Illustration of Few-Shot Prompt. See Figure 11 for the prompt example.
Figure 3: Few-Shot Performance. Accuracy of three Qwen3 models on the two-step-right prediction task across five graph topologies. Each point reports accuracy on 300 test samples with 0, 1, 2, 4, 8, 16, or 32 few-shot demonstrations. Dotted gray lines indicate random guessing over the candidate nodes for each topology. Lepori et al. (2026) corresponds to the zero-shot case.
Figure 4: Probe Performance and Normalized Dirichlet Energy (DE). Probe AUROC and normalized DE of hidden states across shots and five graph topologies for layer 20 of Qwen3-4B-Base. Gray lines show random feature probe baselines.
Figure 5: (a) Cross Probe Performance. AUROC of probes trained at one shot count and evaluated at another, averaged over five topologies. (b) World Representation Subspace Overlap. Overlap between probe subspaces across shot counts, averaged over the five topologies. (c) Angle Between Shot Vectors and the World Representation Subspace. Angles to the corresponding probe subspaces, shown separately for each topology. (d) Distance between Representation Subspaces across Shots. Normalized projection distance between probe subspaces at consecutive tested shot counts, shown separately for each topology. All results are from layer 36 of Qwen3-4B-Base.
Figure 6: (Left) PCA of Shift in Hidden States of a Random Walk. PCA of 4×4 Lattice Grid at layer 36 of Qwen3-8B-Base (Right) PCA of Shifts in Hidden-State Centroids across Random Walks. PCA of 300 test samples per shot for 4×4 Lattice Grid at layer 20 of Qwen3-8B-Base. Points represent individual random walks and triangles indicate their means. In each panel, PCA is fit jointly to all activations shown.
Figure 7: Intervention Results. Graph subspace erasure at iteration 20 in three Qwen3 models. Interventions target node subtokens in the query’s random walk at every Transformer block output. The random intervention uses an orthogonal subspace with matched rank and perturbation norm per node. Points show the mean difference in the probability of correct answer token reduction between graph and random interventions, normalized by each query’s unperturbed probability, across 300 test samples. Probabilities are normalized over candidate answer sequence likelihoods. Positive values indicate greater degradation from graph erasure. Horizontal bars show 95% bootstrap confidence intervals obtained by resampling test samples 5000 times.
Figure 8: (a) ARC-AGI and Othello with Different Worlds across Shots. In ARC-AGI, the world is the rule that maps inputs to outputs. In Othello, the world is the board state. (b) Performance across Different Numbers of Shots on ARC-AGI-1&2, Mind2Web, and Othello.
Appendix figures & tables25 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 9: Illustration of World Models in LLMs.
Figure 10: Illustration of Five Graph Topologies.
Figure 15: Few-Shot Performance where Demonstrations are from Different Graph Topologies. Accuracy of three Qwen3 models on the two-step-right prediction task across five graph topologies. The demonstrations are randomly sampled from four topologies other than the query topology, with each topology represented equally when the number of shots is 4, 8, 16, or 32. We report performance separately for each graph topology in the query. Each point reports accuracy on 300 test samples with 0, 1, 2, 4, 8, 16, or 32 few-shot demonstrations. Dotted gray lines indicate random guessing over the candidate nodes for each topology.
Figure 16: Probe Performance and Normalized Dirichlet Energy (DE). Probe AUROC (Top) and normalized DE of hidden states (Bottom) across shots and five graph topologies for layer 20 of three Qwen3 models. Gray lines show random feature probe baselines. Each condition is evaluated on 300 test samples.
Figure 17: Node Distance Geometry Relative to Zero-Shot Setting. Node distance geometry relative to zero-shot setting using probe projected representations (Top) and full hidden states (Bottom) across shots and five graph topologies for layer 20 of three Qwen3 models. Colored lines compare the same query across shots, while gray lines compare different query families with the same topology but different node vocabulary assignments. Lower values indicate more similar relative node geometry. Each condition uses 300 test samples.
Figure 18: Cross Probe Performance. AUROC of probes trained at one shot count and evaluated at another. Results are from three Qwen3 models on five graph topologies.
Figure 19: Representation Subspace Overlap Results. Overlap between probe subspaces across shot counts. Results are from three Qwen3 models on five graph topologies.
Figure 20: PCA of Shift in Hidden States of a Random Walk Across Different Number of Shots. In each panel, PCA is fit jointly to all activations shown. Results are from three Qwen3 models on five graph topologies.
Figure 21: PCA of Shifts in Hidden-State Centroids across Random Walks Across Different Number of Shots. PCA of 300 test samples per shot. Points represent individual random walks and triangles indicate their means. In each panel, PCA is fit jointly to all activations shown. Results are from three Qwen3 models on five graph topologies.
Figure 22: Angle Between the World Representation Subspace and the Shot Vector. Angles to the corresponding probe subspaces, shown separately for each topology. Results are from three Qwen3 models on five graph topologies.
Figure 23: Angle Between the Shot Vector and the Task Vector. The angle between the two vectors is in degrees from 0∘ to 180∘ . Results from Qwen3-4B-Base on five topologies.
Figure 24: Representation Subspace Distance Results. Normalized projection distance between probe subspaces at consecutive tested shot counts, shown separately for each topology. Results are from three Qwen3 models on five graph topologies.
Figure 25: Intervention Results across Intervention Iterations. Plotted points represent the average paired difference in this degradation between the 16-shot and 0-shot settings across 300 samples in test set. The top is iterative nullspace projection targeted to graph structure, and the bottom is random orthogonal subspace intervention matched in rank. For each prompt condition, relative degradation is computed per test sample as the drop in the model’s assigned probability for the correct answer after intervention, normalized by the probability before intervention. Positive values indicate that intervention causes a larger relative performance drop under 16-shot prompting than under 0-shot setting.
Figure 26: Probe Performance during Intervention Iterations. At each cumulative intervention step (4, 8, 12, 16, and 20), a new linear probe is trained on the intervened residual stream activations and evaluated on a held-out test set. Plotted points indicate the average of test AUROC across the layers 2–36 for Qwen3-4B-Base and Qwen3-8B-Base, and layers 2–40 for Qwen3-14B-Base.
Figure 27: Few-Shot Performance. Accuracy of Llama-3.2-3B, Gemma-4-12B, and Ministral-3-8B-Base-2512 on the two-step-right prediction task across five graph topologies. Each point reports accuracy on 300 test samples with 0, 1, 2, 4, 8, 16, or 32 few-shot demonstrations. Dotted gray lines indicate random guessing over the candidate nodes for each topology.
Figure 28: (a) Probe Performance. Probe AUROC. Gray lines show random feature probe baselines. (b) Cross Probe Performance. AUROC of probes trained at one shot count and evaluated at another. (c) World Representation Subspace Overlap. Overlap between probe subspaces across shot counts. (d) Angle Between Shot Vectors and the World Representation Subspace. Angles to the corresponding probe subspaces. (e) Distance between Representation Subspaces across Shots. Normalized projection distance between probe subspaces at consecutive tested shot counts. Results are from layer 28 of Llama-3.2-3B, layer 48 of Gemma-4-12B, and layer 34 of Ministral-3-8B-Base-2512 on 4×4 Lattice Grid.
Figure 29: Intervention Results. Graph subspace erasure at iteration 20 in Llama-3.2-3B, Gemma-4-12B, and Ministral-3-8B-Base-2512 on 4×4 Lattice Grid. Interventions target node subtokens in the final query’s walk at every Transformer block output. The random intervention uses an orthogonal subspace with matched rank and perturbation norm per node. Points show the mean difference in correct answer probability reduction between graph and random interventions, normalized by each query’s unperturbed probability, across 300 test samples. Probabilities are normalized over candidate answer sequence likelihoods. Positive values indicate greater degradation from graph erasure. Horizontal bars show 95% bootstrap confidence intervals obtained by resampling test samples 5000 times. Note that the wide confidence intervals for 0-shot Gemma-4-12B are because of near-zero baseline probabilities in the denominator.
Figure 30: (a) Probe Performance. Probe AUROC. Gray lines show random feature probe baselines. (b) Cross Probe Performance. AUROC of probes trained at one shot count and evaluated at another. (c) World Representation Subspace Overlap. Overlap between probe subspaces across shot counts. (d) Angle Between Shot Vectors and the World Representation Subspace. Angles to the corresponding probe subspaces. (e) Distance between Representation Subspaces across Shots. Normalized projection distance between probe subspaces at consecutive tested shot counts. Results are from Qwen3-4B-Base on 4×4 Lattice Grid, where r=8 .
Figure 31: (a) Probe Performance. Probe AUROC. Gray lines show random feature probe baselines. (b) Cross Probe Performance. AUROC of probes trained at one shot count and evaluated at another. (c) World Representation Subspace Overlap. Overlap between probe subspaces across shot counts. (d) Angle Between Shot Vectors and the World Representation Subspace. Angles to the corresponding probe subspaces. (e) Distance between Representation Subspaces across Shots. Normalized projection distance between probe subspaces at consecutive tested shot counts. Results are from Qwen3-4B-Base on 4×4 Lattice Grid, where r=16 .
Figure 32: Intervention Results. Graph subspace erasure at iteration 4 in Qwen3-4B-Base on 4×4 Lattice Grid, where r=8,16 . Interventions target node subtokens in the final query’s walk at every Transformer block output. The random intervention uses an orthogonal subspace with matched rank and perturbation norm per node. Points show the mean difference in correct answer probability reduction between graph and random interventions, normalized by each query’s unperturbed probability, across 300 test samples. Probabilities are normalized over candidate answer sequence likelihoods. Positive values indicate greater degradation from graph erasure. Horizontal bars show 95% bootstrap confidence intervals obtained by resampling test samples 5000 times.
Setting
Seed 1
Seed 2
Seed 3
Seed 4
Seed 5
Mean ± Std
L4, 0-Shot
0.6903
0.7072
0.7060
0.7048
0.7013
0.7019±0.0069
L4, 16-Shot
0.6867
0.6824
0.7046
0.6763
0.6941
0.6888±0.0109
L20, 0-Shot
0.7812
0.7761
0.7792
0.7784
0.7777
0.7785±0.0019
L20, 16-Shot
0.7664
0.7573
0.7591
0.7688
0.7722
0.7647±0.0064
L36, 0-Shot
0.7675
0.7508
0.7629
0.7591
0.7656
0.7612±0.0066
L36, 16-Shot
0.7407
0.7359
0.7419
0.7406
0.7412
0.7401±0.0024
Appendix
Table 1: Probe Performance. Results from Qwen3-4B-Base evaluated on 300 test samples on 4×4 Lattice Grid. “L4, 0-Shot” denotes the probe performance at Layer 4 with 0-shot prompting.
Setting
Seed 1
Seed 2
Seed 3
Seed 4
Seed 5
Mean ± Std
L4, 0→16 -Shot
0.6406
0.6487
0.6525
0.6483
0.6425
0.6465±0.0049
L4, 16→32 -Shot
0.6857
0.6821
0.7036
0.6755
0.6942
0.6882±0.0109
L20, 0→16 -Shot
0.7191
0.7112
0.7028
0.7136
0.7061
0.7106±0.0064
L20, 16→32 -Shot
0.7669
0.7580
0.7592
0.7692
0.7731
0.7653±0.0065
L36, 0→16 -Shot
0.6329
0.6727
0.6551
0.6641
0.6538
0.6557±0.0148
L36, 16→32 -Shot
0.7428
0.7375
0.7439
0.7436
0.7415
0.7419±0.0026
Appendix
Table 2: Cross Probe Performance. Results from Qwen3-4B-Base on 4×4 Lattice Grid. “L4, 0→16 -Shot” denotes the Layer 4 probe trained on 500 0-shot samples and evaluated on 300 16-shot samples.
Setting
Seed 1
Seed 2
Seed 3
Seed 4
Seed 5
Mean ± Std
L4, 0↔16 -Shot
0.6295
0.8563
0.7353
0.7794
0.7285
0.7458±0.0826
L4, 16↔32 -Shot
0.9804
0.9728
0.9681
0.9238
0.9928
0.9676±0.0262
L20, 0↔16 -Shot
0.3329
0.4381
0.3081
0.3230
0.2904
0.3385±0.0580
L20, 16↔32 -Shot
0.7970
0.7931
0.8709
0.8691
0.8247
0.8310±0.0376
L36, 0↔16 -Shot
0.1776
0.2669
0.1812
0.1554
0.1601
0.1882±0.0453
L36, 16↔32 -Shot
0.8654
0.8654
0.8993
0.8901
0.8861
0.8813±0.0153
Appendix
Table 3: Representation Subspace Overlap. Results from Qwen3-4B-Base on 4×4 Lattice Grid. “L4, 0↔16 -Shot” denotes the overlap between the rank-2 representation subspaces of probe layers trained on 500 samples.
Setting
Seed 1
Seed 2
Seed 3
Seed 4
Seed 5
Mean ± Std
0-Shot
-0.401
-0.649
-0.380
-0.359
-0.648
−0.487±0.148
16-Shot
1.010
1.293
1.069
1.246
1.635
1.251±0.245
Appendix
Table 4: Mean Relative Degradation Difference after Interventions. Results from Qwen3-4B-Base on the 4×4 Lattice Grid. Each row reports the mean difference in correct answer probability reduction between graph and random interventions after 20 iterations for the corresponding shot setting, normalized by the corresponding clean probability before interventions. Positive values indicate greater degradation under the graph intervention.
Setting
Seed 1
Seed 2
Seed 3
Seed 4
Seed 5
Mean ± Std
L4, 0-Shot
0.5862
0.6337
0.5985
0.6015
0.5806
0.6001±0.0207
L4, 16-Shot
0.5798
0.5897
0.5543
0.5758
0.5582
0.5716±0.0149
L20, 0-Shot
0.7158
0.6900
0.6920
0.7095
0.7000
0.7015±0.0111
L20, 16-Shot
0.6658
0.6783
0.6947
0.6891
0.6845
0.6825±0.0111
L36, 0-Shot
0.6905
0.7208
0.7263
0.7059
0.7109
0.7109±0.0139
L36, 16-Shot
0.6795
0.7056
0.6901
0.6915
0.6914
0.6916±0.0093
Appendix
Table 5: Probe Performance after Graph Interventions. Results from Qwen3-4B-Base evaluated on 300 test samples on 4×4 Lattice Grid after graph intervention for 20 iterations. “L4, 0-Shot” denotes the probe performance at Layer 4 with 0-shot prompting.
How do LLM agents come to both understand environments they act in and master tasks set within them? Through controlled experiments combining world-model training (next-state prediction) and policy training (reward maximization), we investigate this question. We dissect the resulting models through their additive parameter updates. Geometrically, we find effective world-model updates are low-rank and share an input-feature subspace with policy updates while writing to nearly orthogonal output directions, whether trained separately or sequentially. However, we find that, in projection interventions, the sequential update induces more robustness than separate policy RL when removing the world model's leading input directions, suggesting that it has learned alternative input pathways. Behaviorally, we find the sequentially trained agent explores a wider range of states and actions. Based on this, we ask: does policy training preserve world knowledge as well as it could? We probe this with training-free merging built on the geometrically motivated input basis plus an online world-model loss during policy RL, and show both improve over the untreated baseline. Our findings suggest world knowledge and task-directed ability can be learned in geometrically complementary forms, and that future post-training pipelines should consider how best to engineer the interface between them.
Large language models (LLMs) are being used as policies for autonomous decision-making and planning in many domains. Despite their strong reasoning capabilities, LLMs struggle with long-horizon tasks, especially under partial observability. World models are a promising way to enhance policy performance, both during training and inference. During inference, agents currently use world models to simulate the consequences of candidate actions before choosing an action, which can improve decision-making. However, we argue that simulation alone is an incomplete interface for decision-making under partial observability: simulation does not adequately capture uncertainty about the current state, which agents may need for accurate decision-making. We address this limitation with Belief-Based World Models (BB-WMs), which maintain a belief that LLMs can query to access information on what is known and uncertain about the current state. Before developing methods to learn accurate BB-WMs, this paper focuses on a more fundamental question: does exposing a world model's belief directly to an LLM policy improve decision-making? Our results show that giving LLM agents access to beliefs improves task performance under partial observability, while remaining complementary to existing simulation-based world models. Code: https://github.com/skumar-ml/belief-world-models.
A world model predicts environment dynamics based on current observations and actions, serving as a core cognitive mechanism for reasoning and planning. In this work, we investigate how world modeling based on language models can further push the boundaries of general agents. (i) We first focus on building foundation models for agentic environment simulation. We introduce Qwen-AgentWorld-35B-A3B and Qwen-AgentWorld-397B-A17B, the first language world models capable of simulating agentic environments covering 7 domains via long chain-of-thought reasoning. Leveraging more than 10M environment interaction trajectories of 7 domains in real-world environments, we develop Qwen-AgentWorld through a three-stage training pipeline: CPT injects general-purpose world modeling capabilities from the state transition dynamics and augmented professional corpora, SFT activates next-state-prediction reasoning, and RL sharpens simulation fidelity through a tailored framework with hybrid rubric-and-rule rewards. To evaluate language world models, we present AgentWorldBench, a comprehensive benchmark constructed from real-world interactions of 5 frontier models on 9 established benchmarks. Empirical results demonstrate that Qwen-AgentWorld significantly outperforms existing frontier models. (ii) Beyond foundation models, we further investigate two complementary paradigms through which world modeling enhances general agents. First, as a decoupled environment simulator, Qwen-AgentWorld supports scalable and controllable simulation of thousands of real-world environments for agentic RL, yielding gains that surpass real-environment training alone. Second, as a unified agent foundation model, world-model training acts as a highly effective warm-up that improves downstream performance across 7 agentic benchmarks. Code: https://github.com/QwenLM/Qwen-AgentWorld