zixianma/action-selection-heads

OpenWebRL-4B selection heads

Also here: three 4B LoRA selectors (full vision-language models that read the screenshot and the five candidates and answer with their pick): decoder-4b-lora-recipe2-2.6k/, trained on 2,612 Luna lookahead-labelled states, and decoder-4b-lora-gpt55-k3-p16/ / decoder-4b-lora-gpt55-pre-p16/, trained on the same 18,786 panels with GPT-5.5 lookahead or pre-execution labels. See their sections at the end.

Two 22.6 M-parameter candidate-set selection heads for the web agent OpenWebRL/OpenWebRL-4B-SFT (revision 15e777db2ddba2e0e82080ebccd3ad8d215b7f0a). At each browser step the actor samples five candidate responses; the head reads the frozen actor's own hidden states for the prompt and each candidate and picks the one to execute. The actor's weights are unchanged.

Subfolder	Training data	Target	Live mean success (OM2W judge terra / o4-mini)
piotr-data/	2,692 states / 35,916 five-candidate panels from Piotr's SelectionARM release	GPT-5.5's pick, made before execution from the screenshot and candidates	44.3% / 46.3%
lookahead-k3-1380/	1,380 OpenWebRL-4B-SFT branching states (347 tasks), one panel of the state's 2–5 distinct proposed actions	Luna lookahead vote share: Luna judges each action after executing it and following 3 more steps (two calls per state, two orders)	42.3% / 44.0%

Both heads are level within noise: piotr-data − lookahead-k3-1380 = +2.1 pp [−1.3, +5.4] on the three-benchmark mean.

Intended use

Best-of-5 action selection for OpenWebRL-4B-SFT only, as in the evaluation below:

1. At each step, sample 5 responses from OpenWebRL-4B-SFT for the same prompt (evaluated at temperature 1.0, top-p 0.9, max 1,024 new tokens). Keep duplicates; the head is permutation-equivariant.
2. Score the five candidates with the head and execute the argmax (ties to the lowest index).
3. Repeat for up to 30 steps.

The features are hidden states of the 4B-SFT checkpoint at decoder layers 18/27/36, so the heads are not valid for another actor or checkpoint (including RL-trained OpenWebRL actors) unless that actor's candidates are re-encoded by this exact SFT checkpoint; candidate token ids then transfer only if the tokenizer and chat template are identical. They are not reward models: logits are only comparable within one panel, not across states. They were never trained or evaluated with N ≠ 5 at inference except the piotr-data head at N = 3 on Online-Mind2Web, where it matched N = 5.

Usage

Both versions use the single file openwebrl_selection_head.py (torch, transformers, safetensors).

Quick start: you already run the OpenWebRL agent loop

The loop already has the actor's prompt text, the screenshot and the five sampled candidates; the head only re-scores them.

import importlib.util
from pathlib import Path
from huggingface_hub import snapshot_download

root = Path(snapshot_download("zixianma/action-selection-heads", allow_patterns=["*.py", "*.txt"]))
spec = importlib.util.spec_from_file_location("osh", root / "openwebrl_selection_head.py")
osh = importlib.util.module_from_spec(spec); spec.loader.exec_module(osh)

scorer = osh.SelectionHeadScorer.from_pretrained("zixianma/action-selection-heads", "lookahead-k3-1380")  # or "piotr-data"
choice = scorer.select(prompt_text, screenshot, candidates)   # index of the candidate to execute

prompt_text is the exact text the actor was sampled from (ending <|im_start|>assistant\n<think>\n, one image placeholder), screenshot a PIL image of the current page and candidates the 5 responses (strings, or preferably the generated token ids). scorer.score(...) returns the five logits instead.

Full step: building the actor prompt yourself

Without the agent loop, build_prompt reconstructs the actor's prompt from the task and the steps so far, using the evaluation's system prompt (system_prompt.txt).

obs = dict(screen_size=(1280, 1000),
           tabs=[dict(index=0, active=True, url="https://www.google.com/travel/flights", title="Google Flights")])
prompt_text = osh.build_prompt(
    task="Find the cheapest one-way flight from Seattle to Boston next Friday.",
    observations=[obs, obs],                       # one per step so far; the last is the current page
    responses=['I need to enter the origin city first.\n</think>\n\n<tool_call>\n'
               '{"name": "click", "arguments": {"point_2d": [310, 412]}}\n</tool_call>'],   # executed responses
    tool_responses=['Succeed: `click` on <input> "Where from?" at (397, 412) executed.'])     # environment feedback

# sample 5 candidates from OpenWebRL-4B-SFT on prompt_text + the current screenshot
# (evaluated with temperature 1.0, top-p 0.9, max 1,024 new tokens), then:
choice = scorer.select(prompt_text, screenshot, candidates)

Append the chosen response and its environment feedback to responses / tool_responses for the next step. Tab titles are passed as OpenWebRL reports them (truncated to 50 characters); an observation without tabs can pass url= instead. Rebuilt from their parsed contents, 360 saved actor prompts from the evaluation runs (all three benchmarks, 4,644 history steps) match build_prompt byte for byte.

Options. By default the five candidates are processed one after another, which fits a small (18 GB) GPU. On a large GPU, batch_size=5 processes them together: same scores, faster. If the actor is already loaded for sampling, pass it as model= and processor= to avoid a second copy. The head itself is tiny and runs on CPU in about 1 ms.

How it works. The scorer runs the 4B actor over the prompt and screenshot once, then reads each candidate's text through the same model (no new text is generated). From the actor's internal activations at three layers (18, 27, 36) it takes one summary of the page (at the end of the prompt) and, per candidate, summaries of its reasoning, of its action and of its final token. A small transformer compares the five candidates side by side and gives each a score; the highest score is executed. Candidate order does not matter.

In the live evaluation these activations were recorded while the actor sampled, instead of recomputed afterwards. The two routes give near-identical activations (similarity at least 0.999); on 34–37 of 40 saved states they chose the same candidate, and the remaining choices were between near-tied candidates.

Evaluation

OpenWebRL-4B-SFT actor, 5 candidates per step, temperature 1.0 / top-p 0.9, 30 steps, one episode per task, every task of each benchmark. Success counts invalid episodes as failures; "Mean" weights the three benchmarks equally. WebVoyager and DeepShop are judged by GPT-4o. Online-Mind2Web is reported under two judges on the same episodes: the original o4-mini judge and gpt-5.6-terra (re-judged from the saved judge requests).

Selector	Online-Mind2Web (300) o4-mini	Online-Mind2Web (300) gpt-5.6-terra	WebVoyager (595)	DeepShop (150)	Mean o4-mini OM2W	Mean terra OM2W
piotr-data head	43.7%	37.7%	51.3%	44.0%	46.3%	44.3%
lookahead-k3-1380 head	38.3%	33.0%	51.1%	42.7%	44.0%	42.3%
Piotr SelectionARM (4B LoRA selector, same Piotr data)	41.0%	39.0%	52.4%	41.3%	44.9%	44.3%
None: one OpenWebRL-4B-SFT sample	36.3%	31.7%	44.4%	34.0%	38.2%	36.7%

Paired differences on the terra mean, pp [95% bootstrap, tasks resampled within each benchmark]: vs one SFT sample, piotr-data +7.6 [+4.0, +11.3] and lookahead-k3-1380 +5.6 [+1.8, +9.2]; piotr-data − SelectionARM +0.1 [−3.5, +3.6]; piotr-data − lookahead-k3-1380 +2.1 [−1.3, +5.4].

Selector cost per step: about 0.06 PFLOP (actor sampling 5 candidates plus the head) vs 0.098 for the 4B SelectionARM; the head adds about 0.02 s per decision.

Verification. openwebrl_selection_head.py + model.safetensors reproduce the original checkpoints' logits exactly (max difference 0.0 on 250 stored training panels), and re-encoding 5 saved states from scratch on an H200 reproduces the stored training features (cosine 1.0) and the logits of both heads to three decimals.

Training

	piotr-data	lookahead-k3-1380
States / panels	2,692 / 35,916 (5 sampled candidates each, duplicates kept)	1,380 / 1,380 (2–5 distinct actions, reshuffled each epoch)
Target	soft one-hot of GPT-5.5's pick	Luna K=3 lookahead vote share over the actions
Loss	soft cross-entropy	soft cross-entropy
Optimizer	AdamW, lr 3e-4, weight decay 0.05, batch 256	batch 64 (same optimizer settings)
Epochs / checkpoint	20; epoch 6 by pick agreement on Piotr's validation split (3,759 panels)	20; last epoch
Seed	0	0

Each subfolder holds model.safetensors (weights), config.json (head architecture under head_config, plus the sha256 of the original checkpoint and its saved metadata) and head.pt (the original PyTorch checkpoint, kept for provenance: piotr-data sha256 755d4cd4…, lookahead-k3-1380 sha256 11d5eff8…).

4B LoRA selector: decoder-4b-lora-recipe2-2.6k/

A separate selector model, not a head: OpenWebRL/OpenWebRL-4B-SFT fine-tuned with a rank-32 LoRA to read the task, the action history, the current screenshot and five numbered candidate actions, and to reply {"selection": N} (N = 1–5). It uses the same prompt and output format as Piotr Teterwak's SelectionARM and is a drop-in replacement for it.

Path	Contents
decoder-4b-lora-recipe2-2.6k/	merged weights (bf16, 2 safetensors shards, 8.9 GB), tokenizer, processor and chat template; load like any Qwen3-VL checkpoint
decoder-4b-lora-recipe2-2.6k/adapter/	the LoRA adapter alone (rank 32, alpha 64, all linear layers of the language model), for peft on top of OpenWebRL/OpenWebRL-4B-SFT

Usage

1. At each step, sample 5 candidate responses from the actor.
2. Build the selection prompt with build_catts_vision_prompt_v2 from inference/selection_prompt.py of piotr-teterwak/action-reward-models (revision 02276b0, file sha256 043ab7e9…), with these environment flags set: CATTS_VISION_PROMPT_V2=1 CATTS_VISION_COLORED=1 VISION_NO_SOM=1 VISION_ABLATE_DOM=1 VISION_ABLATE_VOTES=1 NORMALIZE_COORDS=1 CLUSTER_NO_DOM=1 VISION_NO_COT=1. Pass each candidate as one cluster (vote_count 1) with its action and its reasoning, in sampled order.
3. Generate greedily (temperature 0, up to 128 new tokens), parse {"selection": N} and execute candidate N (1-indexed).

Training

States	2,612 OpenWebRL-4B-SFT branching states; one panel per state of its 2–5 distinct proposed actions
Target	the action with the highest Luna K=3 lookahead vote share (each action executed and followed for 3 more steps, then judged); tied best actions each give one example
Orderings	each panel shown in 10 candidate orderings: 28,150 training examples
Recipe	LLaMA-Factory SFT, qwen3_vl template, LoRA rank 32 / alpha 64 / dropout 0.05 on all linear layers, vision tower and projector frozen, lr 1e-4 cosine (5% warmup), effective batch 16, 2 epochs, bf16, seed 42, cutoff 8,192 tokens, images up to 262,144 pixels

Evaluation

Live, every task of Online-Mind2Web (300, judge gpt-5.6-terra), WebVoyager (595, GPT-4o) and DeepShop (150, GPT-4o); 5 candidates per step at temperature 1.0 / top-p 0.9, 30 steps, three runs per row pooled. Per-benchmark columns are success among valid episodes; "Mean overall" counts invalid episodes as failures. Differences are in mean overall success, pp [paired 95% bootstrap over tasks].

With the OpenWebRL-4B RL checkpoint as actor (4,096 response tokens; October 31 – November 1):

Selector	Online-Mind2Web	WebVoyager	DeepShop	Mean overall	Mean valid-only	− actor alone
None (one RL sample)	30.2%	56.9%	44.6%	42.0%	43.9%	—
SelectionARM	43.4%	62.7%	50.9%	49.0%	52.3%	+7.0 [+4.8, +9.2]
decoder-4b-lora-recipe2-2.6k	42.4%	64.7%	51.7%	48.6%	52.9%	+6.6 [+4.1, +8.9]

decoder-4b-lora-recipe2-2.6k − SelectionARM: −0.4 [−2.6, +1.7].

With OpenWebRL-4B-SFT as actor (1,024 response tokens; October 29–30):

Selector	Online-Mind2Web	WebVoyager	DeepShop	Mean overall	Mean valid-only	− actor alone
None (one SFT sample)	32.3%	45.2%	29.0%	33.6%	35.5%	—
SelectionARM	43.1%	56.0%	36.0%	41.4%	45.0%	+7.7 [+5.6, +9.8]
decoder-4b-lora-recipe2-2.6k	38.8%	52.5%	38.3%	40.1%	43.2%	+6.5 [+4.3, +8.7]

decoder-4b-lora-recipe2-2.6k − SelectionARM: −1.2 [−3.5, +1.1].

4B LoRA selector: decoder-4b-lora-gpt55-k3-p16/

Same model type, prompt and output format as decoder-4b-lora-recipe2-2.6k/ (a drop-in replacement for SelectionARM; same usage steps above), trained on more states and on panels drawn the way the actor proposes them at inference, with GPT-5.5 K = 3 lookahead labels.

Path	Contents
decoder-4b-lora-gpt55-k3-p16/	merged weights (bf16, 2 safetensors shards, 8.9 GB), tokenizer, processor and chat template; load like any Qwen3-VL checkpoint
decoder-4b-lora-gpt55-k3-p16/adapter/	the LoRA adapter alone (rank 32, alpha 64, all linear layers of the language model), for peft on top of OpenWebRL/OpenWebRL-4B-SFT

Training

States / panels	2,112 OpenWebRL-4B-SFT states (806 tasks); 18,786 panels, each 5 candidates drawn from the state's 80 sampled actor responses (duplicates kept, as at inference)
Target	GPT-5.5's pick after seeing each candidate executed and followed for 3 more actor steps (K = 3 lookahead, screenshots only); tied best candidates each give one example: 32,974 training examples
Recipe	identical to decoder-4b-lora-recipe2-2.6k: LLaMA-Factory SFT, qwen3_vl template, LoRA rank 32 / alpha 64 / dropout 0.05 on all linear layers, vision tower and projector frozen, lr 1e-4 cosine (5% warmup), effective batch 16, 2 epochs, bf16, seed 42, cutoff 8,192 tokens, images up to 262,144 pixels; final train loss 0.123

Evaluation

OpenWebRL-4B-SFT actor (1,024 response tokens), protocol as above; three runs per row pooled. Per-benchmark columns are success among valid episodes; "Mean overall" counts invalid episodes as failures; differences are in mean overall success, pp [paired 95% bootstrap over tasks].

Selector	Online-Mind2Web	WebVoyager	DeepShop	Mean overall	Mean valid-only	− actor alone
None (one SFT sample) ²	34.6%	45.6%	43.1%	39.6%	41.1%	—
SelectionARM	42.7%	52.3%	49.9%	45.2%	48.3%	+5.6 [+3.7, +7.6]
decoder-4b-lora-gpt55-k3-p16	41.1%	52.7%	47.3%	44.7%	47.0%	+5.1 [+3.0, +7.3]
decoder-4b-lora-gpt55-pre-p16	41.4%	53.7%	51.2%	46.2%	48.8%	+6.7 [+4.5, +8.9]

All three selectors ran in one same-day batch (November 4, three runs each). Same day: decoder-4b-lora-gpt55-k3-p16 − SelectionARM −0.5 [−2.4, +1.4], decoder-4b-lora-gpt55-pre-p16 − SelectionARM +1.0 [−1.0, +3.1], pre − K = 3 +1.5 [−0.6, +3.6]. ² The SFT-alone run is from the adjacent window (November 3–4), so the "− actor alone" column is adjacent-window; SelectionARM itself moved −0.5 [−2.4, +1.4] pp between the two windows. DeepShop success has drifted upward since late October, so these absolute numbers are not comparable with the decoder-4b-lora-recipe2-2.6k tables above.

4B LoRA selector: decoder-4b-lora-gpt55-pre-p16/

The pre-execution counterpart of decoder-4b-lora-gpt55-k3-p16/: same model type, prompt, output format, usage, panels and recipe; only the labels differ. Layout as above (merged weights at the top level, the LoRA adapter alone in adapter/).

States / panels	the same 2,112 states (806 tasks) and 18,786 panels as decoder-4b-lora-gpt55-k3-p16
Target	GPT-5.5's pick before execution, from the screenshot, history and the five candidates only (Piotr's SelectionARM labelling setup); tied best candidates each give one example: 34,498 training examples
Recipe	identical to decoder-4b-lora-gpt55-k3-p16; final train loss 0.112

Evaluation: same batch and table as decoder-4b-lora-gpt55-k3-p16 above (both LoRAs and SelectionARM, November 4).