Agent-Editing World Model: Rethinking World Modeling for LLM Agents

AI summary

Overview

• Research area: Natural Language Processing, specifically large language model (LLM) agents, long-horizon decision making, and "world models" for agentic systems.
• Technical level: Advanced. The paper assumes familiarity with LLM agent loops (ReAct-style reasoning–action–observation cycles), supervised fine-tuning, mid-training, rejection sampling fine-tuning, and macro-F1 evaluation.
• Scope in one sentence: The paper proposes the Agent-Editing World Model (AEWM), which replaces the conventional "predict the next environment observation" objective with judging and directly editing an agent's proposed reasoning–action pair before execution, and shows gains on Action Judge, six agent benchmarks, and offline agent training.

What This Paper Is About

Most "language world models" for LLM agents are trained to predict what the environment will return next — search results, terminal output, or test results. The authors argue this target is poorly suited to task-oriented agents, because such observations are high-entropy and depend on execution, and real tools already supply grounded feedback. Meanwhile, the paper identifies a recurring failure mode called task-state contamination: unsupported assumptions and outdated plans persist in the agent's history and distort later decisions. The goal of AEWM is therefore to model how an agent's reasoning and actions shape future task progress, and to edit noisily contaminated reasoning–action continuations before they are executed.

Key Contributions

1. A reframed world-modeling objective. Instead of predicting environment observations, AEWM predicts decision effects and directly edits the agent's pre-execution state, defined as the history plus the proposed reasoning–action pair. The resulting inference loop, EditAct, retains productive decisions, revises noisy ones, and executes the selected action in the real environment.
2. Two coupled capabilities with a cross-domain training framework. Action Judge classifies a proposed decision as Critical, Exploratory, or Noisy from its pre-execution context; State Revision generates an edited reasoning–action pair for noisy proposals. Both are learned from synthesized data via mid-training and supervised fine-tuning (SFT).
3. A 3,000-decision Action Judge benchmark and a transfer method called AEWM-RFT, which uses rejection sampling fine-tuning on verified EditAct trajectories grounded in real environment feedback to move AEWM-guided decision patterns back into the agent itself.
4. Demonstrated effectiveness and transferability across Search, Terminal, and Software Engineering (SWE) domains, including 10.6 macro-F1 points over the strongest baseline on action judgment, 3.2–6.7 average-point improvements across six benchmarks and three backbones, and 2.2–2.6 points over Self-RFT without online AEWM guidance.

Main Findings

• Action judgment: AEWM reaches 70.5% macro-F1 overall on the paper's Action Judge benchmark, outperforming the strongest baseline, DeepSeek-V4-Pro at 59.9%, by 10.6 percentage points. Domain-level macro-F1 is 60.9% (Search), 72.1% (Terminal), and 77.8% (SWE), exceeding the strongest domain-specific baselines by 10.3, 10.5, and 13.4 points respectively. The largest gain is on SWE; Search remains the most challenging domain.
• Agent performance: Across six benchmarks and three backbones, EditAct improves average scores by 6.7, 5.2, and 3.2 points over the strongest baseline for Qwen3.5-4B, Qwen3.5-9B, and Qwen3.5-35B-A3B respectively. Compared with ReAct, the average-score gains are 13.3, 9.6, and 6.6 points.
• Cross-scale effect: Qwen3.5-9B with EditAct reaches an average of 44.1 versus 42.2 for Qwen3.5-35B-A3B with ReAct, showing that state editing can narrow gaps across agent scales.
• Out-of-distribution generalization: Gains hold on both in-distribution and out-of-distribution benchmarks. DeepSearchQA and SWE-Bench Pro are out-of-distribution; BrowseComp, Terminal-Bench 2.0, Doc2Repo, and NL2Repo are in-distribution.
• Ceiling effect with a stronger agent: With Qwen3.5-Plus, EditAct improves BrowseComp from 44.1% to 52.7%, but gains on Terminal and SWE are limited, which the authors attribute to the capacity gap between AEWM and the stronger agent.
• Offline transfer: AEWM-RFT improves over the base agent by 4.5, 6.4, and 5.8 points on BrowseComp, Terminal-Bench 2.0, and Doc2Repo, and exceeds Self-RFT by 2.2, 2.6, and 2.5 points. Average turns fall by 30.0% and 16.7% on the first two benchmarks, while Doc2Repo gains 2.5 points over Self-RFT with more turns.
• Editing beats critiquing: EditAct outperforms Random Gate (random intervention at AEWM's domain-specific noisy rate), Agent Resampling, and AEWM Hint (reasoning guidance), supporting direct replacement of the reasoning–action continuation over critique-guided regeneration.
• Both components matter: Reasoning-only Revision and Action-only Revision each underperform full EditAct, supporting joint editing of reasoning and action. Self-WM (using the inference agent as the world model) does worse on BrowseComp and Terminal-Bench 2.0, and DeepSeek-V4-Pro WM remains below AEWM on all three benchmarks.
• Both training stages matter: Mid-training followed by SFT is best on all three domains, improving over SFT-only by 0.8, 6.4, and 4.3 points on BrowseComp, Terminal-Bench 2.0, and Doc2Repo, and over mid-training-only by 0.3, 4.1, and 1.9 points.
• Domain decision patterns differ: In the annotation corpus, Search contains 60.0% noisy actions; Terminal has the largest exploratory share at 43.2% with 25.1% noisy; SWE has 42.1% critical and 28.0% noisy. Online AEWM predictions retain the same broad pattern.
• Behavioral shift in search: After editing, the mean longest search-only sequence drops from 16.9 to 13.2, while distinct webpages requested per question rises from 7.7 to 10.1, consistent with broader evidence gathering.

Methodology in Plain English

The researchers start from a diagnosis rather than a model: by examining agent trajectories, they find that agents treat unverified assumptions as facts, keep outdated plans despite contradictory feedback, and mistake partial progress for completion — errors that persist in history and compound.

To intervene before execution, they define the agent's "pre-execution state" as the task, the interaction history, and the proposed reasoning–action pair. A single model, AEWM, is trained to do two things with that state. Action Judge labels the proposal as Critical (closes a key gap or performs a required state change on a compact solution path), Exploratory (meaningfully reduces uncertainty or tests a plausible branch), or Noisy (little expected progress, repetition, irrelevance, constraint violation, or wrong direction). State Revision rewrites the reasoning and action for noisy proposals, grounded in the same observed history.

Training data is synthesized. Action Judge data comes from decomposing verified successful trajectories into turns, where a strong annotation agent labels each turn using the matched observation and the rest of the trajectory, producing a forward-looking rationale and an action type. Importantly, the student model sees only the history and proposal, not the future evidence used to create the label — the paper describes this as using retrospective labels as outcome-based supervision for a prospective prediction. State Revision data is built by having a proposal agent attempt tasks, using an Action Judge checkpoint to flag noisy proposals, and having a revision agent generate replacements that are actually executed; only replacements that lead to substantive progress are kept. Data is filtered with structural checks and a rubric where each dimension is scored 0–4 and every dimension must score at least 3.

Training has two stages: mid-training on roughly 52B tokens (reported as 52.16B, comprising 20.56B Action Judge, 25.01B State Revision, and 6.59B trajectory tokens) to build broad interaction knowledge, then supervised fine-tuning on 120K curated examples — 60K Action Judge and 60K State Revision, with 40K examples each from Search, Terminal, and SWE.

At inference, EditAct runs a simple loop: the agent proposes a reasoning–action pair; Action Judge labels it; Critical and Exploratory proposals are kept, Noisy proposals are replaced by State Revision; the selected action is executed in the real environment and the resulting observation enters history. Finally, verified EditAct trajectories are used for rejection sampling fine-tuning (AEWM-RFT) so the agent can internalize these patterns without AEWM present at inference time.

Why This Matters

Impact on research. The paper challenges a default assumption in language world modeling — that the target should be the next observation. It argues that for task-oriented agents, predicting an action's downstream contribution and editing the agent's own state is a more useful and more learnable target. It also provides a labeled action-judgment benchmark of 3,000 decisions and a reproducible data-synthesis recipe spanning three domains, and it connects the idea to existing lines of work on process supervision, revision learning, and step-rating judges.

Real-world applications.

• Deep research and web-search assistants: reducing premature conclusions so that retrieval covers the evidence the task actually requires.
• Software engineering agents: catching flawed dependency assumptions, preserving existing interfaces, and preventing local edits from overriding earlier requirements.
• Terminal, DevOps, and workspace agents: stopping intermediate success signals from being mistaken for verified task completion.
• Training pipelines for agent products: AEWM-RFT shows that edited trajectories can serve as offline supervision, transferring gains without running a world model at inference time.

Industry relevance. The reported gains are largest on smaller backbones (6.7 points for Qwen3.5-4B), and Qwen3.5-9B with EditAct surpasses Qwen3.5-35B-A3B with ReAct on average. Combined with the reduced turn counts under AEWM-RFT (30.0% and 16.7% fewer average turns on two benchmarks), this suggests a route to better agent reliability per unit of compute. The caveat reported by the authors — limited Terminal and SWE gains with the stronger Qwen3.5-Plus agent — also signals a practical limit: the world model needs to keep pace with the agent it supervises.

Future Directions

• Closing the capacity gap. The paper's reported limitation with Qwen3.5-Plus — BrowseComp improving from 44.1% to 52.7% while Terminal and SWE gains stay limited — raises the question of how to train or scale AEWM so it remains useful for stronger agents.
• Search as the hardest domain. Search has the lowest Action Judge macro-F1 (60.9%) and the highest noisy share in the annotation corpus (60.0%), so improving judgment in high-entropy retrieval environments is a clear open problem.
• When to revise versus when to explore. Terminal has the largest exploratory share (43.2%), which makes the boundary between useful exploration and waste a natural target for further study.
• Reducing dependence on verified trajectories. AEWM-RFT currently relies on verified EditAct trajectories grounded in real feedback; extending the approach to settings where such verification is expensive or unavailable is an open question.

Target Audience

Researchers and engineers working on LLM agents, agentic reinforcement learning, and tool-use systems will get the most from this paper, particularly those focused on long-horizon reliability, search agents, terminal agents, and software engineering agents. It is also relevant to readers interested in world models, process supervision, and rejection-sampling fine-tuning. Because the paper uses formal notation for state definitions and reports detailed ablation tables, it is best suited to readers with prior exposure to agent scaffolding and post-training methods.

Authors’ abstract

Recent advances in large language models (LLMs) have enabled agents to tackle long-horizon tasks across diverse environments. To further improve agent performance, existing language world models typically predict environment observations, yet reconstructing high-entropy, execution-dependent tool responses offers limited value when real feedback is available. Meanwhile, agents suffer from \emph{task-state contamination}, where unsupported assumptions and outdated plans persist in history and distort subsequent decisions. We propose the \textbf{Agent-Editing World Model (AEWM)}, which models how reasoning and actions shape future task progress rather than simulating tool responses. AEWM combines \textbf{Action Judge} to distinguish \textsc{Critical}, \textsc{Exploratory}, and \textsc{Noisy} decisions with \textbf{State Revision} to edit noisy reasoning--action continuations from the same observed history. \textbf{EditAct} integrates these capabilities with real execution, directly changing the state underlying subsequent decisions rather than merely providing critiques. We train AEWM across Search, Terminal, and Software Engineering through mid-training and supervised fine-tuning. AEWM achieves 70.5\% macro-F1 on our Action Judge benchmark, exceeding the strongest frontier baseline by 10.6 points. Across six benchmarks and three agent backbones, EditAct improves average scores by 3.2--6.7 points over the strongest baseline. Furthermore, rejection sampling fine-tuning on verified EditAct trajectories, termed \textbf{AEWM-RFT}, improves over Self-RFT by 2.2--2.6 points across three domains without online AEWM guidance.