Back to AI Research

AI Research

Discriminative World Models for Web Agents | AI Research

Key Takeaways

  • Discriminative World Models for Web Agents Recent web agents often use "world models" to plan their next moves.
  • 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.
Paper AbstractExpand

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: this https URL .

Discriminative World Models for Web Agents
Recent web agents often use "world models" to plan their next moves. By predicting what a website will look like after a specific action, these agents can compare different choices before committing to one. Traditionally, these models are trained to recreate a specific, fixed format—like a snapshot of the website’s code or a text summary. However, this paper argues that this approach is flawed because it focuses on reproducing a format rather than capturing the specific details that actually matter for making a decision. The authors introduce a new training objective called "predicted-state matching," which shifts the focus toward creating representations that help an agent distinguish between the outcomes of different possible actions.

A New Way to Train World Models

Instead of forcing a model to generate a rigid, predefined string, the authors train it to produce a representation that acts as a unique identifier for a specific outcome. To do this, they created a "branching" dataset from existing web-agent trajectories. In this dataset, every decision point includes multiple possible actions and the different resulting states they lead to. During training, the model is rewarded when its predicted representation allows a "judge" to correctly identify the true outcome of an action when presented with an alternative, incorrect outcome. This makes the model's output flexible, allowing it to prioritize the most relevant information needed to differentiate between choices.

Improving Action Selection

The authors tested whether these discriminative representations actually help agents perform better. They integrated their world model into a Process Reward Model (PRM), which is a system used to rank candidate actions. By providing the PRM with these more informative, discriminative representations, the agent could better evaluate the consequences of its potential moves. The results showed that this approach outperformed traditional models that were trained to simply predict fixed-format snapshots, as those snapshots often contained too much irrelevant information or missed the subtle changes that distinguish one action from another.

Better Performance on Web Tasks

The researchers evaluated their method on WebArena-Lite, a benchmark for web-based tasks. They found that using their discriminative world model for test-time action selection led to higher task success rates. By focusing on the differences between potential future states rather than just trying to "draw" a picture of the next page, the agent gained a clearer understanding of how its actions would impact the web environment. This suggests that for complex, multi-step navigation problems, it is more effective for an agent to learn what makes an outcome unique rather than trying to perfectly replicate the entire state of a webpage.

Comments (0)

No comments yet

Be the first to share your thoughts!