Discriminative World Models for Web Agents
Kelvin Li, Dhruv Pendharkar, Trevor Darrell, Roei Herzig and colleagues (Berkeley, MIT-IBM) point out that web-agent world models are trained on the wrong objective, and replace next-state prediction with predicted-state matching.
Ask this paper
The objective mismatch: World models are trained to reconstruct HTML or AXTree snapshots, but the downstream ranker needs predicted states to be discriminative across candidate actions. Reconstruction fidelity and discriminability are not the same target.
Predicted-state matching: The predicted representation must distinguish the true resulting state from the states reached by alternative actions, which is a contrastive objective aligned with how the prediction is actually consumed.
Branching data from Go-Browse: Training uses a branching WebArena dataset where every decision point carries multiple alternative actions and their resulting states, which is what makes the contrastive objective trainable at all.
Gains at three levels: Better on the held-out predicted-state matching benchmark, better PRM-style action ranking on WebPRMBench than action-only PRMs and supervised-world-model PRMs, and higher end-to-end success on WebArena-Lite.
Why it matters: A clean instance of a general lesson: train the auxiliary model on the loss its consumer actually needs, not on the loss that was easy to write.
Abstract
Recent web agents use world models for test-time action selection by sampling candidate actions, predicting the resulting web states, and ranking them with a ranker model or a Process Reward Model (PRM). These world models are typically trained via supervised next-state prediction to generate fixed representations like HTML or AXTree snapshots. However, this objective is misaligned with the downstream ranker, which relies on predicted states being discriminative across candidates to accurately score them. To address this, we introduce predicted-state matching, a training objective where the predicted representation must distinguish the true resulting state from those reached by alternative actions. We train these models using a branching web-agent dataset derived from WebArena Go-Browse trajectories, where every decision point contains multiple alternative actions and their resulting states. Experiments on our held-out predicted-state matching benchmark show that our approach outperforms world models trained with supervised next-state prediction. We further show that our approach improves PRM-style action ranking on WebPRMBench compared with action-only PRMs and PRMs augmented with supervised-next-state world models. Finally, on WebArena-Lite, using our world model for test-time action selection improves end-to-end task success. Our project page is available at: https://dhruvpendharkar.github.io/dwm/.