## Discriminative World Models for Web Agents

Discriminative World Models for Web Agents

Recent web agents rely on world models to pick actions at test time.

The typical pipeline samples many candidate actions, predicts the web state after each action, and then ranks the outcomes with a ranker or a Process Reward Model (PRM).

These world models are usually trained by supervised next‑state prediction.

The new approach replaces the generative predictor with a discriminative model that directly scores action–state pairs.

The discriminative model learns to tell whether a given action will lead to a desired state, instead of trying to reconstruct the whole next state.

This change makes prediction faster and more accurate for ranking.

A simple implementation can look like this:

import torch
from torch import nn

class DiscriminativeWorldModel(nn.Module):
    def __init__(self, state_dim, action_dim, hidden=128):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(state_dim + action_dim, hidden),
            nn.ReLU(),
            nn.Linear(hidden, 1),          # score for (state, action)
            nn.Sigmoid()
        )
    def forward(self, state, action):
        x = torch.cat([state, action], dim=-1)
        return self.net(x)

# Example usage:
score = model(current_state, candidate_action)
# higher score → better action

The paper shows that this discriminative world model improves ranking quality while using fewer samples.

It also reduces the need for a separate ranker or PRM, because the model’s score can be used directly for action selection.

Experiments on several web navigation tasks report higher success rates and lower latency.

Read the full paper for more details: Discriminative World Models for Web Agents (arXiv).