Web agents · World models · 2026

Discriminative World Models
for Web Agents

Kelvin Li* Dhruv Pendharkar* Anish Pahilajani Chuyi Shang Leon Oks Leonid Karlinsky Rogerio Feris Trevor Darrell Roei Herzig

UC Berkeley  ·  MIT-IBM Watson AI Lab  ·  Cal Poly San Luis Obispo  ·  Xero
* Equal contribution

World models should predict what helps an agent choose—not simply what is easiest to reconstruct.

Paper · Coming soon Code · Coming soon Data · Coming soon
Qualitative comparison on a WebArena Reddit page. The proposed method correctly predicts that a combobox expands while two baselines fail to describe the action-induced change.
Say what changed. Our model captures the specific outcome of the queried action, while prior approaches predict a different action or repeat page structure that remains unchanged.
01 · Motivation

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.

80.80% State-matching accuracy
+16.90 Best-of-N points vs. direct PRM
28.48% WebArena-Lite success
02 · Approach

Train for the decision
the model must support.

Standard world models learn to imitate a fixed state format. Predicted-state matching instead rewards a representation when it makes the queried action’s true outcome distinguishable from alternative outcomes.

Comparison between supervised next-state prediction and predicted-state matching. The proposed method uses a judge to compare the target and a negative state.
From reconstruction to discrimination. The matching judge never sees the instruction, history, current state, or action; the generated representation itself must identify the correct resulting state.
1

Branch the data

Merge repeated browser states across Go-Browse trajectories into a graph with multiple executable actions and observed outcomes.

2

Predict the consequence

Given the task, history, current accessibility tree, and queried action, the world model generates a flexible textual state representation.

3

Match, don’t imitate

A judge must match the representation to the true next state over an alternative state reached by a competing action.

03 · Branching data

One state.
Multiple futures.

Linear trajectories hide the counterfactuals an agent needs for local decision making. Our state-action graph exposes them: 7,730 branching decision points from 2,839 trajectories yield 30,920 pairwise examples.

Construction of a branching state-action graph from repeated states in linear trajectories, followed by target and negative state matching.
Repeated states are merged into a state-action graph. Each outgoing action provides a true next state for one query and a hard negative for another.
04 · Results

More discriminative.
More useful.

Predicted-state matching improves both the intrinsic quality of world-model representations and the downstream decisions made from them.

A

Held-out predicted-state matching

Our Qwen3-8B model achieves the strongest overall matching accuracy, outperforming existing web world models and the data-matched supervised full-AXTree baseline.

Predicted-state-matching accuracy (%) using Qwen3-32B as the matching judge. Overall is micro-averaged.
ModelShoppingCMSRedditGitLabMapOverall
General language models
GPT-4o44.4447.8048.2247.0055.8149.40
GPT-4o-mini48.1546.7047.7247.0053.4948.80
Qwen3-4B59.2654.9552.7952.0053.9553.94
Qwen3-8B67.9055.5068.0056.5068.3062.86
Existing world models
WebDreamer-7B72.8070.8075.6375.0076.7074.51
WebWorld-8B79.0167.5877.6654.0077.2170.17
Data-matched SFT baseline
Full AXTree target45.6846.1547.7246.0051.6347.77
Predicted-state matching Ours77.7877.4779.7082.5084.1980.80
B

WebPRMBench action ranking

Under the same Qwen2.5-7B reward-model backbone and answer-only SFT setup, our state representations improve the controlled Best-of-N average by 16.90 points over direct action ranking.

Controlled Qwen2.5-7B reward models on WebPRMBench. Pair = Pairwise accuracy; BoN = Best-of-N accuracy.
Model Mind2WebWebArenaAssistantBenchWorkArenaAverage
PairBoNPairBoNPairBoNPairBoNPairBoN
Direct (no state)85.1460.9180.8552.7382.5056.6779.5752.8882.0255.80
+ WebWorld-8B states89.7473.2084.3265.6086.6667.4181.2164.3285.4867.63
+ State-matching states Ours96.7189.3988.5669.1588.3370.0083.8462.2689.3672.70
C

End-to-end WebArena-Lite

Adding predicted state representations to Best-of-5 action selection more than doubles GPT-4o ReAct-style task success in the same evaluation harness.

05 · Qualitative example

Say what changed.

When a combobox is clicked, our model names the exact consequence: the menu expands and displays available forums. Fixed-format or supervised baselines either predict the wrong action or repeat irrelevant page structure.

Current state Reddit submission page with the forum dropdown closed and highlighted.
Resulting state Reddit submission page after the forum dropdown has expanded to show options.

Takeaway

Optimize world models to preserve the action-relevant differences needed for downstream decisions—not to reproduce a predefined representation of the next state.
06 · Cite

Citation

@article{li2026discriminative,
  title   = {Discriminative World Models for Web Agents},
  author  = {Li, Kelvin and Pendharkar, Dhruv and Pahilajani, Anish and
             Shang, Chuyi and Oks, Leon and Karlinsky, Leonid and
             Feris, Rogerio and Darrell, Trevor and Herzig, Roei},
  year    = {2026}
}