Agent-Editing World Model: Rethinking World Modeling for LLM Agents
Shuang Sun, Guoxin Chen, Fanzhe Meng, Jia Deng, Huatong Song, Jinhao Jiang, Wayne Xin Zhao, Hongteng Xu, Ji-Rong Wen
cs.CL, cs.AI, cs.LG
2026-09-24
AEWM judges and rewrites the agent's reasoning-action pair instead of simulating tools; Action Judge hits 70.5 macro-F1, EditAct adds 3.2-6.7 points.
Language world models for LLM agents usually predict the next tool payload: search hits, terminal dumps, test logs. Those observations are high-entropy and execution-dependent, so the predictor is both hard and prone to fabricating evidence. Long-horizon agents fail in a more specific way: task-state contamination. Unsupported assumptions harden into history, stale plans survive contradictory feedback, and partial progress is treated as done. Later steps then look locally sensible on a wrong mental model.
AEWM, from Renmin University's Gaoling School, changes the prediction target. It does not ask what the environment will look like. It asks whether the current reasoning-action pair is Critical, Exploratory or Noisy for future task progress. Noisy pairs are rewritten from the same observed history, then executed in the real environment.
The pre-execution state is \(st = ht \oplus (\hat rt, \hat at)\): task and history plus the agent's proposed reasoning and action. Action Judge emits one of three labels. Critical closes a necessary gap or performs a required state change. Exploratory reduces uncertainty or tests a plausible branch. Noisy is repetition, irrelevance, a constraint break or a wrong direction. Only Noisy triggers State Revision, which writes \((\tilde rt, \tilde at)\). Both the next action and the sentence that will enter history change. EditAct wires judge, rewrite and real execution: keep Critical and Exploratory, rewrite Noisy, always take observations from real tools.
Training spans Search, Terminal and SWE. An annotator labels successful trajectories at turn level. Supervision is prospective: the input is only the history visible at time \(t\) plus the proposal, while the label comes from the finished trajectory. Revision samples are stricter: after a Noisy call, the rewritten action must produce later real progress or it is dropped. Two stages: about 52B tokens of mid-training mixing raw trajectories with synthetic judge/revise data, then SFT on 120K examples (60K each task, balanced across three domains). The backbone is Qwen3.5-35B-A3B.
AEWM-RFT rejection-samples verified EditAct trajectories into the agent so inference no longer needs the world model.
The Action Judge benchmark has 3,000 decisions, split evenly. AEWM reaches 70.5 macro-F1 against 59.9 for the strongest baseline, DeepSeek-V4-Pro, a 10.6 point gap. Per domain: 60.9 / 72.1 / 77.8 on Search / Terminal / SWE, each more than 10 points above the best domain baseline. Search remains the hardest.
EditAct beats ReAct and both step-level and trajectory-level Best@3 (DeepSeek-V4-Pro as selector) on all six benches and three backbones:
| Backbone | ReAct | Best Best@3 | EditAct |
| Qwen3.5-4B | 28.5 | 35.1 | 41.8 |
| Qwen3.5-9B | 34.5 | 38.9 | 44.1 |
| Qwen3.5-35B-A3B | 42.2 | 45.6 | 48.8 |
Average lifts over the strongest baseline are 6.7 / 5.2 / 3.2 points. On 4B, BrowseComp moves from 31.5 to 41.2 and Terminal-Bench 2.0 from 17.6 to 37.1. Qwen3.5-9B with EditAct averages 44.1, above the 35B-A3B ReAct baseline of 42.2. Out-of-distribution DeepSearchQA and SWE-Bench Pro rise as well.
AEWM-RFT versus Self-RFT: BrowseComp 45.4 vs 43.2, Terminal-Bench 2.0 43.4 vs 40.8, Doc2Repo 48.6 vs 46.1, with 30.0% and 16.7% fewer turns on the first two. Ablations that random-gate at the noisy rate, resample the agent, hint only, edit reasoning or action alone, or swap in Self / DeepSeek-V4-Pro as the world model, all sit below full EditAct. Mid-training and SFT are complementary; dropping either hurts Terminal and Doc2Repo more.
In the annotation corpus, Search is 60.0% noisy, Terminal is 43.2% exploratory, SWE is 42.1% critical. Online predictions keep the same shape.
The prediction target of a language world model is the wrong object if real tools are already in the loop. What needs modeling is whether the agent is contaminating its own task state. EditAct is not extra samples plus a picker; it replaces the current mental continuation before the tool runs. For search, terminal and repo agents, this is middleware that plugs into ReAct and can be distilled away with RFT. Gains shrink as the backbone grows, which reads as a patch for small-model state maintenance, not an unbounded new skill.
Judge labels come from retrospective annotation of successful trajectories; noisy patterns on failed traces may differ. Search is still weak at 60.9 F1, and a 60% noisy rate means dense intervention with a real risk of clipping useful exploration. On stronger Qwen3.5-Plus, BrowseComp moves from 44.1 to 52.7 while Terminal and SWE barely move, which the paper attributes to a capacity gap between AEWM and the agent. Extra latency and tokens for EditAct are not reported. 52B tokens of mid-training is not cheap. The world model edits the agent's continuation and does not learn environment dynamics, so a new tool protocol can break the judge.