Agent editing improves reasoning accuracy in large language model agents
Agent-Editing World Model: Rethinking World Modeling for LLM Agents
Computation and LanguageArtificial IntelligenceMachine Learning
Summary
When language-based AI agents try to complete long and complex tasks, they often get stuck using outdated plans or incorrect assumptions, hurting their decisions. The authors introduce a new model called Agent-Editing World Model (AEWM) that focuses on how an agent's reasoning and actions affect progress, rather than just simulating tool outputs. AEWM can identify important decisions, exploratory moves, or mistakes, and then edit the agent's task progress to fix errors. This approach leads to better agent performance across several task types and different AI backbones.
What this means in practice
- •For software engineers: Improve AI assistants to better manage and revise their step-by-step reasoning when building software components.
- •For search engine developers: Enhance AI agents that interact with search tools by enabling them to correct faulty plans and reasoning in real time.
Authors
Shuang Sun, Guoxin Chen, Fanzhe Meng, Jia Deng, Huatong Song, Jinhao Jiang, Wayne Xin Zhao, Hongteng Xu, Ji-Rong Wen
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.