World models let observers explore and interact with generated worlds. Recent work shows that when observers revisit a region, current models often fail to regenerate it consistently. We view this as a context retrieval problem. LLM agents manage context retrieval at two layers of their stack. At the model layer, the KV cache design determines how much context fits in memory. At the agent harness layer, harnesses such as Claude Code pick the context sent to the LLM. Prior work on world models has mostly focused only on the model layer. In this paper, we make three contributions. First, we build a benchmark to measure world model context retrieval. It generates new worlds for every episode, preventing a model from recalling rooms using training data. It also separates context retrieval from guesses based on nearby frames. Second, we build an agent harness that performs context compaction for world models, which to our knowledge no prior work has studied. Third, we use our benchmark to study four world model designs. Context compaction lets our world models retrieve rooms up to 4 times beyond their context window without retraining. However, world models with sliding-window attention lose most of their context retrieval when dreaming. Therefore, a dreaming world agent benefits more from extending the reach of its world model using context compaction than switching to sliding window attention at the model layer.
Figures & tables
Figure 1: Left: Impact of model architecture on context retrieval for our video game room generation world model. Decoder-only needs 28.3 MB of KV cache memory for a context window of 128 steps. In the same memory, DD fits a context of 1,020 steps. Middle: Context retrieval gain of each world model when it dreams (with context compaction enabled). DD, which saves its memory with SWA, loses most of its context retrieval. Right: Agent harness layer can improve context retrieval on top of model layer changes. Decoder-only retrieves a room revisited after 262 steps. The visit from 262 steps ago is no longer in Decoder-only’s 128-step context window. The model therefore cannot reproduce the room. However, context compaction allows the original visit to still be sent to the model.
Figure 2: A frame in our benchmark world with the observer placed at the center of the room. The observer is allowed to use the doors to transition to the neighboring rooms in the world’s 8×8 grid.
Figure 3: The architecture of the four world models we train and evaluate on. Each world model has 8 transformer blocks. We show the KV cache corresponding to a transformer block to the right of each respective block. From left to right, each architecture changes one property of the model incrementally. Switching from Decoder-only to CED, the upper blocks compute their keys and values from the output z of the last lower block. From CED to CED-SWA, the lower blocks switch to using SWA so they only attend over the 2 most recent steps. Finally, from CED-SWA to DD, the 4 upper KV caches become one shared KV cache that is read by all the upper blocks.
Figure 4: Our agent harness conducts context compaction before sending it to a world model. It records the entire session as the full context, which soon outgrows the model’s context window. Context compaction allows it to send the model a walk through every visited room instead of the full context.
Figure 5: Performing context compaction at the agent harness layer lets every world model retrieve rooms beyond its raw context window. Each bar shows the context retrieval gain for revisits with gaps in the range [128,511] steps.
Figure 6: Context retrieval gain of our world model designs when the memory budget is fixed at 2,304 vectors. Note that each vector has D=384 elements and we use ρ=32 bits per floating point.
Figure 7: Context retrieval gain of our world model designs when the raw context window size is fixed at 128 steps.
Figure 8: The effect of dreaming on context retrieval gain for our four world model designs.
Figure 9: Mean agent latency per step over a 512-step exploration session of the benchmark world.
Appendix figures & tables4 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 10: The context retrieval gain of all four of our world model designs over a range of gap bins. We use solid black lines to show the 128-step raw context window and use dashed gray lines to show a compacted context of about 80 steps.
Figure 11: Context retrieval gain for gap bins of 128 to 511 steps.
Figure 12: Context revisit gain for revisit loops of length 8 , 32 , and 128 steps respectively with and without letting the world model dream.
Decode
Model
Batch 1
Batch 16
Prefill
Decoder-only
5.74
17.42
14.53
CED
5.89
18.06
9.80
CED-SWA
5.30
11.37
9.56
DD
4.49
7.53
9.20
Appendix
Table 1: The latency (in milliseconds) for a call to our world model designs using a context of 128 steps. We report the median latency over 30 iterations.
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
World models (WMs) simulate the transition dynamics of environments, enabling agents to plan over the consequences of their actions. In text-based environments, fine-tuning a Language Model (LM) to serve as a WM has emerged as a dominant paradigm. However, despite the widespread success of non-parametric approaches such as Retrieval Augmented Generation (RAG), retrieval for LM-based world modelling remains underexplored. We conduct a systematic evaluation across five diverse environments spanning embodied, web navigation and social settings, comparing fine-tuning and RAG-based approaches for LM-based world modelling. Our study reveals that fine-tuning often outperforms RAG, with fine-tuned WMs enabling agents to obtain higher rewards on 15/20 settings. While both construction paradigms benefit from additional and more diverse exploration, RAG-based approaches prove more data-efficient, and fine-tuning approaches disproportionately benefit from scaling the amount of experience collected. With a focus on RAG-based WMs, we devise a procedure that uses counterfactual intervention to estimate the error rate of the retrieval stage, and show that retrievers consistently surface suboptimal transitions from the experience buffer. Hoping to address this failing, we study a variety of query reformulation strategies, demonstrating that a hierarchical approach outperforms the traditional retrieval pipeline. Finally, we compose our findings into a hybrid world modelling system that parametrically captures core environment dynamics, while learning to rely on retrieval from an actively maintained memory store. Our hybrid system consistently outperforms other methods across multiple environments and models, showcasing the robustness of the approach and the applicability of our findings.
Dhananjay Ashok, Shantanu Agarwal, Vivek Datla +2
Information Sciences Institute University of Southern California · CapitalOne · Capital One
Large Language Models (LLMs) have a bounded context window. The context window is the maximum input size an LLM can consume for a single inference. AI agents rely on a process called context compaction to fit their state within the context window when calling an LLM. Despite its ubiquity, context compaction has received essentially no formal analysis. In this paper, we initiate a formal study of context compaction. We first introduce a framework consisting of two games that capture the two algorithmic strategies for context compaction used by contemporary AI agents in practice. The Context Selection Game models context compaction algorithms that select a subset of an agent's accumulated state to retain. The Context Generation Game models context compaction algorithms that summarize an agent's state by an arbitrary message of bounded length. We then prove an equivalence between the Context Generation Game and one-way communication complexity. The minimum context compaction budget for answering a set of queries within a target error is equal to the one-way communication complexity of the induced communication problem at the same error. Known bounds from communication complexity therefore transfer directly to context compaction. We also show that the Context Selection Game corresponds to a restricted class of one-way communication protocols. Any gap between selection and generation is therefore a gap between two classes of communication protocols. We prove that there exists a set of queries for which generation needs strictly less budget than selection. The equivalence between the Context Generation Game and one-way communication also lets us measure how well a deployed context compaction algorithm performs on a query relative to the optimal strategy. As an example, we present a case study that evaluates Anthropic's context compaction endpoint on set membership queries.
Hayder Tirmazi, Sam Markelon, Allison Bishop +1
AllSpice Inc. · Boston University · Proof Trading +2