AEWM rethinks what a language World Model should predict: instead of guessing tool outputs, it judges the agent’s own proposed reasoning-plus-action as Critical, Exploratory, or Noisy, and rewrites Noisy ones before real execution, addressing what the authors call task-state contamination.
You’ve built a long-horizon LLM agent that loops through reasoning, calling a tool, reading the result, and deciding the next step. It runs for dozens of turns across a browser, a shell, or a code repo. The failure mode the authors focus on: partway through, the agent latches onto an unsupported assumption (a wrong entity, a stale plan, a half-passing test read as “done”). That bad belief sits in the history and every later turn is conditioned on it. Grounded observations from real tools don’t fix this, because the agent keeps interpreting them through the contaminated state.
The standard research response has been to build a language world model that predicts what the tool will return, so the agent can plan against a simulator (this is the MuZero lineage adapted to text, and recent work like Code world model (FAIR CodeGen)). The authors’ objection: search results, shell output, and test logs are high-entropy and depend on live state, so predicting them well is hard and, when the real tool is available anyway, low-value. Worse, a wrong prediction fabricates evidence. So they change the prediction target entirely.
The agent still runs a normal ReAct loop, but before each proposed reasoning-action pair executes, AEWM intercepts it. Two capabilities, one model:
•
Action Judge looks at the history plus the proposed (reasoning, action) and outputs one of three labels. Critical means it closes a needed gap. Exploratory means it reduces uncertainty or tests a branch. Noisy means repetition, irrelevance, or a wrong direction. The three-way split (rather than good/bad) is deliberate: exploration is preserved.
•
State Revision fires only on Noisy. It regenerates both the reasoning and the action from the same history, and the new action is what actually runs in the real environment. Both the revised reasoning and the revised action get written into history, so future turns condition on the corrected interpretation, not just a corrected next move.
Contrast with a conventional world model: the old formulation maps state to a predicted observation. AEWM maps state to an edited state, and lets the real environment supply the observation. The authors call the integrated inference procedure EditAct.
for t in range(max_steps):
r_hat, a_hat = agent.propose(history)
label = aewm.action_judge(history, r_hat, a_hat)
if label in {"critical", "exploratory"}:
r, a = r_hat, a_hat
else: # noisy
r, a = aewm.state_revision(history, r_hat, a_hat)
o = env.execute(a) # real tool, no simulation
history.append((r, a, o))
Training is two-stage on a single Qwen3.5-35B-A3B backbone. Mid-training on ~52B tokens mixes raw trajectories with synthesized judge/revision examples. SFT on 120K curated examples (60K judge, 60K revision, balanced across Search, Terminal, SWE) calibrates the two heads. Judge labels come from a strong annotator (DeepSeek-V4-Pro) looking at the full trajectory retrospectively, then filtered so the student only sees pre-decision context (no future leakage). Revision examples come from a weaker proposer, an earlier Judge checkpoint flagging Noisy proposals, and a stronger reviser generating replacements that must produce real downstream progress.
Separately, they use verified EditAct trajectories as SFT data for the base agent itself, called AEWM-RFT (Rejection-sampling fine-tuning (RFT)). The agent then runs without AEWM at inference time.
Judgment quality. On a held-out 3,000-decision benchmark the authors built (1,000 each in Search, Terminal, SWE), AEWM hits 70.5% macro-F1 overall vs 59.9% for the strongest frontier baseline (DeepSeek-V4-Pro). Gains are 10+ points in every domain. Frontier models tested include Gemini 3 Pro, GPT-5.5, GLM-5.2, and Qwen3.7-Max. Note this benchmark is the authors’ own construction, using their labeling pipeline, so it defines “correct judgment” the way their training data does.
Task performance with EditAct. Six benchmarks across three domains: BrowseComp and DeepSearchQA for search, Terminal-Bench 2.0, SWE-bench Pro, Doc2Repo, and NL2Repo. Compared to the strongest baseline (either step-level or trajectory-level Best-of-3 with a DeepSeek-V4-Pro verifier), EditAct improves the six-benchmark average by +6.7, +5.2, +3.2 points on Qwen 4B, 9B, and 35B-A3B respectively. Gains shrink as the base agent gets stronger, and the authors note they saw limited gains on Terminal/SWE when pairing AEWM with an even stronger Qwen3.5-Plus agent, attributing it to a capacity gap.
Ablations are the interesting part. Holding the 35B agent fixed on BrowseComp, Terminal-Bench 2.0, and Doc2Repo:
•
Random Gate (revise at the same frequency but pick turns randomly) is worse than learned judgment: judgment matters, not just intervention rate.
•
Agent Resampling (let the base agent try again on Noisy turns) and AEWM Hint (give the agent AEWM’s reasoning as a suggestion) both lose to direct replacement. Editing the state beats critiquing it.
•
Editing only the action or only the reasoning both underperform editing both.
•
Using DeepSeek-V4-Pro as the world model (no AEWM training) is close on BrowseComp but noticeably worse on Terminal-Bench and Doc2Repo. So both the intervention design and the domain-specific training carry weight.
AEWM-RFT. Agents fine-tuned on EditAct trajectories, then run without AEWM, beat Self-RFT (same recipe, own trajectories) by 2.2 to 2.6 points on three benchmarks, and use fewer turns on BrowseComp and Terminal-Bench. This is the transfer claim: EditAct rollouts contain patterns the agent can absorb.
•
If you’re running a long-horizon agent and observing the “latches onto wrong hypothesis and never recovers” failure, the paper’s diagnosis (contamination lives in reasoning text carried forward, not just in bad actions) is worth taking seriously. A practical implication: don’t only rewrite failed actions, rewrite the reasoning that produced them, because the next turn conditions on both.
•
Before building an observation-predicting world model for a tool-using agent, ask whether the real tool is cheap and reliable. If it is, the paper’s argument suggests predicting tool outputs mostly buys you fabrication risk. Predicting decision quality is a smaller, better-posed target.
•
The full recipe assumes you can train a 35B model on ~52B tokens plus 120K SFT examples with heavy annotation from a frontier model. That’s not casual. The lighter takeaway worth testing: a learned classifier that gates when to invoke a revision step is meaningfully better than triggering it randomly or every turn, even before you get to joint reasoning-action rewriting.
•
If you can only do one thing, AEWM-RFT is the cheapest downstream artifact: it’s ordinary SFT on trajectories your instrumented agent already produced, and the resulting agent needs no inference-time overhead. The paper only demonstrates this on one backbone (Qwen3.5-35B-A3B), so treat cross-backbone transfer as untested.
•
Code and datasets are released on GitHub and HuggingFace.
•
The Action Judge benchmark is authored by the same team, using the same annotator (DeepSeek-V4-Pro) that provides training labels. The 10-point gap over frontier models partly measures alignment with the authors’ labeling scheme, not a domain-independent notion of decision quality.
•
Gains from EditAct shrink as the base agent gets stronger (6.7 -> 5.2 -> 3.2 points from 4B to 35B), and the authors report that pairing AEWM with an even stronger agent gave limited gains outside search. The intervention helps most when the world model is comparable to or stronger than the agent.
•
Every experiment uses Qwen3.5 agents and Qwen3.5-35B-A3B as the AEWM backbone. Transfer to other agent families isn’t shown.
•
Several dataset names in the setup (the internal deep-search corpus, CalibForge, DeNovoSWE) are cited to concurrent or internal work; the paper doesn’t fully describe them, so reproducing the training pipeline exactly from public sources will require chasing citations.