One-sentence Summary

Researchers at Renmin University of China propose the Agent-Editing World Model (AEWM), which models how reasoning and actions shape future task progress instead of simulating tool responses and combines Action Judge and State Revision to classify and edit noisy reasoning-action continuations; its EditAct integrates these capabilities with real execution, improving average scores by $3.2\text{--}6.7$3.2–6.7 points across six benchmarks and three agent backbones.

Key Contributions

• The paper introduces the Agent-Editing World Model (AEWM), which models how reasoning and actions shape future task progress rather than simulating execution-dependent tool responses. AEWM combines Action Judge for classifying decisions as CRITICAL, EXPLORATORY, or NOISY with State Revision for editing noisy reasoning-action continuations, and EditAct applies these capabilities with real execution to alter the state behind subsequent decisions.

• The work develops a cross-domain training framework that uses trajectory synthesis, mid-training, and supervised fine-tuning across Search, Terminal, and Software Engineering. It also contributes a 3,000-decision Action Judge benchmark and AEWM-RFT, which performs rejection sampling fine-tuning on verified EditAct trajectories to transfer decision patterns without online AEWM guidance.

• AEWM reaches 70.5% macro-F1 on the Action Judge benchmark, exceeding the strongest frontier baseline by 10.6 points. EditAct improves average scores by 3.2 to 6.7 points across six benchmarks and three agent backbones, and AEWM-RFT surpasses Self-RFT by 2.2 to 2.6 points across three domains.

Introduction

World models are seen as a foundation for general intelligence and are increasingly important for long-horizon LLM agents that must understand task environments and sustain coherent decisions. Prior language world models typically follow an environment-centered objective, predicting action-conditioned next states or tool responses, but in open-ended agent tasks these observations are often high-entropy and execution-dependent, and reconstructing them provides limited value while risking simulated or fabricated evidence. The authors identify a recurring failure mode called task-state contamination, where agents accept unsupported assumptions, retain outdated plans, or mistake partial progress for completion. To address this, they propose the Agent-Editing World Model (AEWM), which shifts world modeling from reconstructing observations to modeling decision effects and editing agent states. AEWM includes Action Judge to classify proposed reasoning-action pairs as critical, exploratory, or noisy, and State Revision to replace noisy continuations before execution through EditAct, with training and transfer demonstrated across Search, Terminal, and Software Engineering domains.

Dataset

Dataset Description

The dataset is built from verified agent trajectories and is split into Action Judge and State Revision components.

Action Judge Data

• Composition: Turn-level annotations from successful, verified agent trajectories. Each example covers a single turn.
• Annotation and schema: An annotation agent labels each action as CRITICAL, EXPLORATORY, or NOISY based on the environment observation and its role in later task completion. It also writes a forward-looking reasoning trace. Inputs contain the visible history up to turn t and the proposed reasoning and action. Outputs contain the retrospective reasoning trace and action type.
• Filtering: The authors filter for label consistency, action-observation alignment, action validity, and reasoning grounded only in the pre-execution history.
• Benchmark: A held-out set of 3,000 decisions, equally distributed across Search, Terminal, and SWE. It is built from verified trajectories with repeated annotation, consistency filtering, model-based review, and diversity-aware sampling. Accuracy and macro-F1 are primary evaluation metrics.
• Size: The benchmark size is 3,000. The text does not specify the total training set size.

State Revision Data

• Composition: Created from tasks attempted by a proposal agent, with revisions generated by a separate revision agent and judged by an Action Judge checkpoint.
• Filtering: The authors retain only samples where the proposed action is labeled NOISY and the revised reasoning-action pair leads to substantive progress in the subsequent real-environment trajectory.
• Schema: Inputs contain the visible history and the noisy proposed reasoning-action pair. Outputs contain the revised reasoning-action pair.

Use in the Model

• Action Judge data trains the model to predict an action's downstream contribution from its pre-execution context.
• State Revision data trains the model to revise noisy proposed reasoning-action pairs into improved ones.
• The Action Judge benchmark is held out for evaluating pre-execution action judgment.

Method

The authors introduce the Agent-Editing World Model (AEWM), a framework designed to mitigate task-state contamination in long-horizon agent tasks. Task-state contamination occurs when agents treat unsupported assumptions as facts or retain outdated plans, allowing errors to persist and amplify through the interaction history. Rather than predicting high-entropy environment observations, AEWM directly models and edits the agent's reasoning and action states before execution.

Refer to the framework diagram for a comprehensive overview of the AEWM architecture, data synthesis pipeline, and integration process.

Agent State Modeling and Editing Given a task $x$x, the interaction history before step $t$t is denoted as $h_t = (x, (r_i, a_i, o_i)_{i=0}^{t-1})$ht​=(x,(ri​,ai​,oi​)i=0t−1​). The agent proposes a reasoning and action pair $(\hat{r}_t, \hat{a}_t)$(r^t​,a^t​), forming the pre-execution state $s_t = h_t \oplus (\hat{r}_t, \hat{a}_t)$st​=ht​⊕(r^t​,a^t​). AEWM, denoted as $\mathcal{M}$M, intervenes on this state through two primary modules: Action Judge and State Revision.

The Action Judge module, $\mathcal{M}_{\mathrm{AJ}}$MAJ​, evaluates the proposed continuation and predicts a decision-effect label $\widehat{y}_t \in \{critical, exploratory, noisy\}$y​t​∈{critical,exploratory,noisy}. Critical decisions close key gaps or perform required state changes, exploratory decisions reduce uncertainty or test plausible branches, and noisy decisions offer little progress or promote repetition and constraint violations.

If the Action Judge classifies a proposal as noisy, the State Revision module, $\mathcal{M}_{\mathrm{SR}}$MSR​, generates a revised reasoning and action pair $(\widetilde{r}_t, \widetilde{a}_t)$(rt​,at​) to form an edited state $\widetilde{s}_t$st​. This grounded intervention contrasts with conventional observation-predictive world models:

$$\underbrace {s _ {t} \xrightarrow {\mathcal {M} _ {\mathrm{obs}}} \widehat {o} _ {t}} _ {\text {Conventional WM: observation prediction}} \quad \text {vs.} \quad \underbrace {s _ {t} \xrightarrow {\mathcal {M} _ {\mathrm{SR}}} \widetilde {s} _ {t} = h _ {t} \oplus (\widetilde {r} _ {t}, \widetilde {a} _ {t})} _ {\text {AEWM: agent - state editing}}$$Conventional WM: observation predictionst​Mobs​​ot​​​vs.AEWM: agent - state editingst​MSR​​st​=ht​⊕(rt​,at​)​​

EditAct: Integrating State Editing and Acting During inference, the authors integrate AEWM into the agent loop via a process called EditAct. At each step, the Action Judge produces the label $\widehat{y}_t = \mathcal{M}_{\mathrm{AJ}}(s_t)$y​t​=MAJ​(st​). The reasoning and action pair committed to the trajectory is determined by:

$$(r _ {t}, a _ {t}) = \left\{ \begin{array}{l l} (\hat {r} _ {t}, \hat {a} _ {t}), & \widehat {y} _ {t} \in \{\text {critical, exploratory} \}, \\ (\widetilde {r} _ {t}, \widetilde {a} _ {t}), & \widehat {y} _ {t} = \text {noisy}. \end{array} \right.$$(rt​,at​)={(r^t​,a^t​),(rt​,at​),​y​t​∈{critical, exploratory},y​t​=noisy.​

State Revision is invoked exclusively in the noisy case. The selected action $a_t$at​ is then executed in the real environment $\mathcal{E}$E, yielding an observation $o_t \sim \mathcal{E}(\cdot \mid h_t, a_t)$ot​∼E(⋅∣ht​,at​) and updating the history as $h_{t+1} = h_t \oplus (r_t, a_t, o_t)$ht+1​=ht​⊕(rt​,at​,ot​). This loop allows state editing to redirect problematic decisions while real environment observations ground subsequent evolution.

Learning and Internalizing Agent Editing To train AEWM, the authors synthesize specialized datasets and employ a two-stage training process. Action Judge data is created by decomposing successful agent trajectories and using an annotation agent to label each turn as critical, exploratory, or noisy based on its downstream contribution. State Revision data is generated by having a proposal agent attempt tasks; when its noisy proposals are revised by a revision agent and lead to substantive progress, the revised pairs are retained.

The unified AEWM is trained across Search, Terminal, and SWE domains. The first stage is Mid-Training on a corpus of approximately 52 billion tokens, combining original trajectories with the synthesized data to teach the model broad task knowledge and decision-effect judgment. The second stage is Supervised Fine-Tuning (SFT) on a curated set of 120,000 examples, evenly split between Action Judge and State Revision tasks, to calibrate both capabilities.

Finally, the authors internalize the editing capabilities into the base agent through AEWM-based Rejection Sampling Fine-Tuning (AEWM-RFT). By running the agent with EditAct in real environments, they collect high-quality trajectories that include AEWM's state corrections. Fine-tuning the base agent on these verified trajectories allows it to learn local state corrections and sustain task progress independently, effectively internalizing the ability to avoid and recover from task-state contamination.

Experiment

The experiments evaluate EditAct across Search, Terminal, and SWE domains using ReAct and Best@3 baselines on six benchmarks, while separately assessing Action Judge and State Revision components. AEWM outperforms frontier models in action judgment, and EditAct improves over baselines across model scales and out-of-distribution benchmarks, with AEWM-RFT providing transferable supervision that improves success and often reduces turns. Ablations show that learned intervention selection, direct joint reasoning-action revision, and the full mid-training plus SFT recipe are important, while analyses indicate AEWM repairs task-state contamination by rechecking evidence, verifying execution, and preserving requirements.

EditAct improves agent performance over ReAct and Best@3 baselines across benchmarks and model scales. The gains remain consistent on both in-distribution and out-of-distribution tasks, with especially large improvements in terminal and repository-related benchmarks. A mid-size model using EditAct can surpass a larger model using ReAct, indicating that the approach narrows performance gaps across agent scales. EditAct outperforms ReAct and both Best@3 variants on every benchmark and at each tested model scale. Average-score gains over the strongest baseline shrink but remain positive as model size increases. Compared with ReAct, EditAct yields larger average improvements than over the strongest Best@3 baseline, and a mid-size model with EditAct exceeds a larger model with ReAct. Improvements hold on both in-distribution and out-of-distribution benchmarks, with stronger relative gains on Terminal-Bench and Doc2Repo.

Fine-tuning on AEWM-guided trajectories outperforms both the base agent and an agent fine-tuned on its own trajectories across BrowseComp, Terminal-Bench 2.0, and Doc2Repo. Efficiency improves on the first two benchmarks, with higher scores and fewer average turns, while Doc2Repo improves with more turns. The pattern indicates that EditAct-based trajectories provide useful supervision beyond self-generated successful rollouts. AEWM-RFT achieves the highest scores across all three benchmarks, exceeding Self-RFT by about two to three points. It reduces average turns on BrowseComp and Terminal-Bench 2.0, while Doc2Repo improves with additional turns.

The full AEWM setup with EditAct achieves the best scores across BrowseComp, Terminal-Bench 2.0, and Doc2Repo, outperforming all tested inference interventions and training-stage combinations. Learned intervention selection and direct state revision contribute to this advantage, since random gating, agent resampling, hint guidance, single-component edits, and alternative world models all produce lower scores. The combined mid-training plus SFT recipe is stronger than either stage alone, with the largest gains on Terminal-Bench 2.0 and Doc2Repo. Full AEWM with EditAct achieves the highest score on BrowseComp, Terminal-Bench 2.0, and Doc2Repo. Random gating, agent resampling, and AEWM hint interventions all underperform the full method, supporting learned intervention selection and direct editing over resampling or critic-like guidance. Action-only and reasoning-only revisions each trail the full EditAct, indicating that joint reasoning-action editing is more effective. Using the inference agent as the world model underperforms the full method and is especially weak on Terminal-Bench 2.0; a DeepSeek-V4-Pro world model also trails full AEWM. Mid-training followed by SFT outperforms either training stage alone, with the full recipe producing larger gains on Terminal-Bench 2.0 and Doc2Repo.

Across BrowseComp, Terminal-Bench 2.0, and Doc2Repo, the experiments compare EditAct with ReAct and Best@3 baselines at multiple model scales, then evaluate fine-tuning on AEWM-guided trajectories and ablate the full AEWM pipeline. EditAct consistently improves over the baselines on both in-distribution and out-of-distribution tasks, and a mid-size EditAct model can surpass a larger ReAct agent. Fine-tuning on AEWM-guided trajectories outperforms self-generated successful rollouts, while the full AEWM setup with joint reasoning-action editing and the combined mid-training plus SFT recipe achieves the best results, confirming the value of learned intervention selection and direct state revision.