Recent web agents select actions at test time by sampling candidate actions, predicting the resulting web states, and ranking them with a ranker or a Process Reward Model (PRM). Existing world models are trained with supervised next‑state prediction, outputting fixed representations such as HTML or AXTree snapshots. This objective is misaligned with the downstream ranker, which requires the predicted states to be discriminative across candidates for accurate scoring.
To address this, we introduce predicted‑state matching as a training objective. The model must distinguish the true resulting state from the states reached by alternative actions. We build a branching web‑agent dataset from WebArena Go‑Browse trajectories, where each decision point contains multiple alternative actions and their resulting states.
Experiments on a held‑out predicted‑state matching benchmark show that models trained with this objective outperform those trained with supervised next‑state prediction. On WebPRMBench, a PRM augmented with our world model outperforms action‑only PRMs and PRMs combined with supervised‑next‑state world models. Finally, on WebArena‑Lite, using the world model for test‑time action selection improves end‑to‑end task success.
Project page: https://dhruvpendharkar.github.io/dwm/ Review