IMPORTANT: yes
Rethinking World Models for AI Agents: From Predicting Observations to Editing States
Large language model agents now tackle long-horizon tasks across search, terminal, and software engineering environments. These agents reason, act, observe, and reason again in a loop that can span dozens or hundreds of steps. A persistent challenge in this setting is task-state contamination: unsupported assumptions and outdated plans accumulate in the agent's interaction history, distorting subsequent decisions in ways that are locally plausible but globally wrong. An agent may pursue a search candidate contradicted by earlier results, build code on an unverified foundation, or interpret partial progress as completion. Grounded environment observations alone do not solve this problem, because the contamination lives in the agent's internal state, not in the external world.
Shuang Sun and colleagues from Renmin University of China propose a fundamental rethink of what a world model should do for an agent. Rather than predicting the next environment observation, their Agent-Editing World Model (AEWM) predicts how an agent's reasoning and actions shape future task progress, and directly edits the agent's state when it judges a proposed decision will be unproductive. Submitted on September 23, 2026, the paper introduces a framework that combines decision judgment, state revision, and execution integration, achieving substantial improvements across six benchmarks and three agent backbones.
The Core Insight: Task-State Contamination, Not Observation Prediction
Existing language world models typically learn to predict action-conditioned next observations, following the paradigm established in embodied and video-world modeling. For an agent navigating a text-based environment, however, the next observation is a terminal output, a search result, or a test report. These observations are high-entropy, execution-dependent, and often difficult to predict accurately. More importantly, reconstructing them offers limited value when the real tool is available and can provide grounded feedback. The paper argues that the more pressing modeling target is not what the environment will look like next, but whether the agent's current reasoning and action will advance or undermine task progress.
This shift in perspective reveals a failure mode that the authors call task-state contamination. An agent exploring an unfamiliar environment must maintain hypotheses, plans, and judgments about verified progress from partial observations. When these are wrong, they persist in the history and compound through locally plausible actions. The critical modeling question becomes: given the current task history and a proposed reasoning-action pair, what is the likely downstream effect of committing to that pair? AEWM answers this question with two complementary capabilities.
Action Judge: Separating Critical, Exploratory, and Noisy Decisions
Action Judge is the first component of AEWM. Given the agent's current pre-execution state, which consists of the task, the accumulated history, and the proposed reasoning-action pair, Action Judge classifies the proposal as one of three types. A Critical decision closes a key gap, obtains necessary evidence, or performs a required state change along a compact solution path. An Exploratory decision meaningfully reduces uncertainty or tests a plausible branch that is not yet the most direct path. A Noisy decision provides little expected progress, promotes repetition or irrelevance, violates constraints, or proceeds in an incorrect direction.
The explicit inclusion of the Exploratory class is a deliberate design choice. Not all useful actions are the shortest path to the answer. Information-gathering steps that do not immediately advance the solution but narrow the hypothesis space are valuable and should not be mistaken for noise. The three-way classification preserves this distinction, giving the editing system a finer-grained basis for intervention.
The paper constructs a training dataset for Action Judge by decomposing verified, successful agent trajectories into individual turns. A strong annotation agent examines each turn's history and proposed action, then labels it based on its contribution to subsequent task completion. The labels are forward-looking in substance: they capture whether committing to a particular reasoning-action pair at a particular point in history is likely to help or hinder. The benchmark consists of 3,000 decisions equally distributed across Search, Terminal, and Software Engineering, constructed through repeated annotation, consistency filtering, and diversity-aware sampling.
State Revision: Editing the Agent's Internal Continuation
When Action Judge labels a proposal as Noisy, State Revision takes over. Given the same pre-execution state, State Revision generates a new reasoning-action pair that is grounded in the task and observed evidence. The revised reasoning can reconcile conflicting information, revise an unsupported hypothesis, or update the plan. The revised action turns this correction into the next interaction with the real environment.
The crucial distinction between AEWM and conventional world models is captured in a simple diagram in the paper. A traditional observation-predictive world model takes the current state and produces a predicted next observation. AEWM takes the current state and produces an edited state, where the editing applies specifically to the reasoning and action continuation while preserving the accumulated history. The revised pair replaces the original in the agent's history, and subsequent observations are obtained through real execution.
State Revision training data is constructed using a proposal agent, a revision agent, and a pre-trained AEWM checkpoint specialized for Action Judge. The proposal agent attempts each task; when its proposed reasoning-action pair is judged Noisy, the revision agent generates a replacement from the same history, and the new action is executed in the real environment. Samples are retained only when the revised reasoning and action lead to substantive progress, as evidenced by the subsequent trajectory.
EditAct: Integrating Judgment and Revision with Execution
EditAct is the inference-time integration of Action Judge and State Revision into the agent's decision loop. At each step, the agent proposes a reasoning-action pair. AEWM judges it. If the label is Critical or Exploratory, the proposal is executed unchanged. If the label is Noisy, State Revision generates an edited pair, which is executed instead. The selected action runs in the real environment, and the resulting observation enters the history alongside the selected reasoning.
What makes EditAct more than a critique system is that it changes the state underlying subsequent decisions. A conventional critic might flag a problematic action and prompt the agent to regenerate, but the agent's history still contains the original noisy proposal. EditAct replaces it. The revised reasoning and action become part of the history that shapes all future reasoning, interrupting the propagation of task-state contamination at its source.
Training: From Synthetic Data to Real-World Transfer
AEWM is trained across three domains through a two-stage pipeline. Mid-training uses approximately 52 billion tokens combining original agent trajectories with synthesized Action Judge and State Revision data. The original trajectories provide broad interaction knowledge, while the synthesized data teach the model to judge decision effects and revise unproductive continuations. Supervised fine-tuning then activates and calibrates both capabilities using 120,000 carefully curated examples: 60,000 for Action Judge and 60,000 for State Revision, evenly divided across Search, Terminal, and SWE.
The paper introduces AEWM-RFT, a rejection sampling fine-tuning procedure that transfers AEWM-guided decision patterns back into the agent. Verified, high-quality EditAct trajectories grounded in real environment feedback are used to fine-tune the agent, with supervision covering both local state corrections and subsequent decisions. Crucially, this transfer happens without online AEWM guidance at inference time, meaning the agent internalizes the editing patterns and applies them independently.
Results: Substantial Gains Across Domains and Model Scales
On the Action Judge benchmark, AEWM achieves 70.5% macro-F1, exceeding the strongest frontier baseline, DeepSeek-V4-Pro, by 10.6 percentage points. The advantage is consistent across all three domains: 60.9% on Search, 72.1% on Terminal, and 77.8% on SWE, with gains above 10 points in each. This cross-domain consistency suggests that AEWM learns robust decision-level judgment rather than specializing to particular interaction dynamics.
EditAct improves agent performance across all six benchmarks and three backbones. Average score gains over the strongest baseline are 6.7 points for Qwen3.5-4B, 5.2 points for Qwen3.5-9B, and 3.2 points for Qwen3.5-35B-A3B. Compared with ReAct, the improvements are 13.3, 9.6, and 6.6 points respectively. Notably, Qwen3.5-9B with EditAct (44.1%) surpasses Qwen3.5-35B-A3B with ReAct (42.2%), indicating that AEWM can narrow performance gaps across agent scales.
The improvements hold on both in-distribution and out-of-distribution benchmarks, including BrowseComp and DeepSearchQA, which were not part of AEWM's training distribution. AEWM-RFT further improves performance beyond Self-RFT by 2.2 to 2.6 points across BrowseComp, Terminal-Bench 2.0, and Doc2Repo, without requiring online AEWM guidance at inference. The corrected trajectories also reduce average turns by 30% on BrowseComp and 16.7% on Terminal-Bench 2.0, suggesting that editing not only improves accuracy but also makes agents more efficient.
Ablations: What Makes Editing Work
The ablation studies confirm that EditAct's benefits come from its specific design rather than from generic intervention. Random gating at AEWM's noisy rate performs worse than learned judgment, supporting the value of trained decision-quality assessment. Agent resampling and AEWM hint both perform worse than direct state revision, indicating that editing the reasoning-action continuation is more effective than prompting the agent to regenerate or providing critic-like guidance. Neither reasoning-only nor action-only revision matches the full EditAct, supporting joint editing of both components. Using the inference agent itself as the world model (Self-WM) performs worse than AEWM, as does using DeepSeek-V4-Pro as a world model, confirming that specialized training matters.
The training ablation shows that both mid-training and SFT contribute. SFT alone improves ReAct by 0.8 to 4.3 points across benchmarks, while the full training recipe adds a further 0.3 to 4.1 points beyond SFT alone, with the largest gains on Terminal-Bench 2.0 and Doc2Repo.
How AEWM Edits Agents in Practice
The paper provides detailed case studies across all three domains. In Search, AEWM learns to check hypotheses against accumulated evidence rather than pursuing contradicted candidates. In Terminal, it learns to verify required behavior before committing to execution paths that may fail. In Software Engineering, it learns to diagnose faults while preserving compatibility with existing code. These patterns reflect a common theme: AEWM edits agents toward more disciplined information-gathering and verification, interrupting the drift into unsupported assumptions that characterizes task-state contamination.
Limitations and Open Directions
The paper acknowledges several limitations. The Action Judge benchmark, while cross-domain, is limited to three task types. Extending the framework to richer environments, including visual and embodied settings, would require adapting the state representation and editing mechanisms. The AEWM-RFT procedure relies on verified EditAct trajectories, and the quality of these trajectories depends on the underlying agent and environment. Additionally, the training data is synthesized using strong annotation and proposal agents; whether these capabilities scale to more complex domains remains open.
The paper also notes that EditAct's benefits vary with agent backbone capacity. Gains over the strongest baseline diminish at larger scales, likely because the base agent already possesses much of the reasoning quality that AEWM edits into smaller models. The Qwen3.5-Plus experiments show that gains can re-emerge when the base agent has different capacity characteristics, but the overall pattern suggests diminishing returns as base models improve.
Why This Matters
AEWM represents a conceptual shift in how world models should serve agents. Rather than trying to predict what the environment will do next, it predicts what the agent's decisions will do to the task, and intervenes directly when those decisions threaten to corrupt the agent's internal state. This shift aligns the world model's objective with the agent's actual needs: not a simulation of the environment, but a safeguard against self-inflicted reasoning errors.
The practical results are substantial. Across six benchmarks spanning search, terminal interaction, and software engineering, EditAct improves performance by 3 to 7 percentage points over the strongest baselines, with improvements that transfer across model scales and generalization settings. The AEWM-RFT procedure further demonstrates that the editing patterns learned by AEWM can be transferred to agents without online guidance, making the approach deployable in settings where AEWM would add inference latency.
The framework's emphasis on distinguishing Critical, Exploratory, and Noisy decisions provides a language for analyzing agent behavior that goes beyond simple success or failure. By separating information-gathering from solution-progression and from outright noise, AEWM offers both a practical intervention and a conceptual lens for understanding why long-horizon agents fail and how to prevent it.