Why this matters

The Selective Rollout paper [1] identified a concrete inefficiency in GRPO-style agent training: roughly 40% of rollout groups reach zero reward variance, meaning every trajectory in the group earns the same reward and the group contributes no gradient signal. Those rollouts are pure compute waste. The paper’s fix is a one-parameter gate that measures mean pairwise prefix edit distance across parallel trajectories at each step. When the distance falls below a threshold, the group has already converged on a single action prefix and can be stopped early. On a 60-iteration GRPO run with Qwen2.5-7B on ALFWorld, this gate delivered a 10.7% wall-clock speedup (bootstrap 95% CI excludes zero) and a +2.5 pp improvement in held-out task success rate [1].

The same gate is useful at inference time, not just training time. When you run a multi-sample agent loop (beam-style reasoning, self-consistency, or any ensemble rollout), early-stopping zero-variance groups reduces token spend directly. This tutorial wires the gate into a LangGraph agent loop backed by a vLLM endpoint, then measures the token savings with a benchmark script.

Prerequisites

  • Python 3.11 or 3.12
  • A CUDA-capable GPU or cloud instance with vLLM installed (the benchmark script mocks the vLLM call so the structural code runs without a GPU)
  • Familiarity with async Python and basic LangGraph concepts
  • uv or pip for package installation

Setup

Install the Python dependencies. vllm itself is assumed pre-installed on your GPU instance; the packages below cover the agent loop, benchmarking, and edit-distance computation.

uv pip install langgraph langchain-core editdistance httpx tqdm

Export a placeholder base URL for the vLLM OpenAI-compatible endpoint. In the benchmark script this URL is intercepted by a mock transport, so no real vLLM server is required to run the tutorial end-to-end.

export VLLM_BASE_URL="http://localhost:8000/v1"
export VLLM_MODEL="Qwen/Qwen2.5-7B-Instruct"

Step 1: The prefix edit-distance gate

The gate from [1] is a single function: given a list of partial action sequences (one per parallel rollout), compute the mean pairwise normalized edit distance and compare it to a threshold tau. When the mean distance drops below tau, the group has converged and can be stopped.

# filename: rollout_gate.py
"""Selective rollout gate based on mean pairwise prefix edit distance [1]."""
from __future__ import annotations
import itertools
from typing import Sequence
import editdistance


def _normalized_edit(a: str, b: str) -> float:
    """Levenshtein distance normalized to [0, 1]."""
    if not a and not b:
        return 0.0
    return editdistance.eval(a, b) / max(len(a), len(b))


def mean_pairwise_distance(action_prefixes: Sequence[str]) -> float:
    """Mean pairwise normalized edit distance across all rollout pairs."""
    n = len(action_prefixes)
    if n < 2:
        return 1.0  # single rollout: never gate
    pairs = list(itertools.combinations(range(n), 2))
    total = sum(
        _normalized_edit(action_prefixes[i], action_prefixes[j])
        for i, j in pairs
    )
    return total / len(pairs)


def should_stop_group(action_prefixes: Sequence[str], tau: float = 0.05) -> bool:
    """Return True when the group has converged below threshold tau.

    The paper [1] uses a single threshold parameter tau. When mean pairwise
    prefix edit distance < tau, all rollouts are on the same action prefix
    and the group is predicted to be zero-variance.
    """
    return mean_pairwise_distance(action_prefixes) < tau

Verify the gate logic with a quick unit check:

from rollout_gate import mean_pairwise_distance, should_stop_group

# Identical prefixes -> distance 0 -> gate fires
identical = ["search(query)", "search(query)", "search(query)"]
assert mean_pairwise_distance(identical) == 0.0
assert should_stop_group(identical, tau=0.05) is True

# Divergent prefixes -> distance high -> gate does not fire
diverse = ["search(query)", "lookup(item)", "finish(answer)"]
assert mean_pairwise_distance(diverse) > 0.5
assert should_stop_group(diverse, tau=0.05) is False

print("gate_unit_tests_passed")

Step 2: Mock vLLM transport for sandbox execution

The LangGraph agent calls the vLLM OpenAI-compatible /v1/chat/completions endpoint. In the sandbox (no GPU, no running vLLM server), a custom httpx transport intercepts those calls and returns deterministic responses. On a real GPU instance, remove the mock transport and point VLLM_BASE_URL at your server.

# filename: mock_vllm.py
"""Deterministic mock for the vLLM OpenAI-compatible endpoint."""
from __future__ import annotations
import json
import time
import random
import httpx

# Simulated action pool: some prompts get convergent responses, some divergent.
_CONVERGENT_ACTIONS = [
    "search(capital of France)",
    "search(capital of France)",
    "search(capital of France)",
]
_DIVERGENT_ACTIONS = [
    "search(Paris history)",
    "lookup(Eiffel Tower)",
    "finish(Paris is the capital)",
]


class MockVLLMTransport(httpx.BaseTransport):
    """Returns canned chat completions without a real vLLM server."""

    def __init__(self, convergent: bool = False, seed: int = 0):
        self.convergent = convergent
        self._rng = random.Random(seed)

    def handle_request(self, request: httpx.Request) -> httpx.Response:
        body = json.loads(request.content)
        n = body.get("n", 1)
        actions = _CONVERGENT_ACTIONS if self.convergent else _DIVERGENT_ACTIONS
        choices = []
        for i in range(n):
            action = actions[i % len(actions)]
            choices.append({
                "index": i,
                "message": {"role": "assistant", "content": action},
                "finish_reason": "stop",
            })
        payload = {
            "id": f"mock-{int(time.time())}",
            "object": "chat.completion",
            "created": int(time.time()),
            "model": body.get("model", "mock"),
            "choices": choices,
            "usage": {
                "prompt_tokens": 50 * n,
                "completion_tokens": 10 * n,
                "total_tokens": 60 * n,
            },
        }
        return httpx.Response(200, json=payload)

Step 3: The multi-rollout agent node

The LangGraph agent runs G parallel rollouts per step. After each step it checks the gate. If the gate fires, it marks the group as stopped and records how many steps were saved.

# filename: agent_state.py
"""LangGraph state schema for the multi-rollout agent."""
from __future__ import annotations
from dataclasses import dataclass, field
from typing import List, Optional


@dataclass
class RolloutGroupState:
    prompt: str
    group_size: int = 4
    max_steps: int = 8
    tau: float = 0.05
    # Runtime fields
    step: int = 0
    action_prefixes: List[str] = field(default_factory=list)
    stopped_early: bool = False
    stop_step: Optional[int] = None
    total_tokens_used: int = 0
    # Baseline comparison: tokens if we had NOT stopped early
    baseline_tokens: int = 0
# filename: agent_nodes.py
"""LangGraph node implementations for the selective-rollout agent."""
from __future__ import annotations
import json
import os
import httpx
from typing import Any, Dict

from agent_state import RolloutGroupState
from rollout_gate import should_stop_group


def _build_client(transport: httpx.BaseTransport | None = None) -> httpx.Client:
    base_url = os.environ.get("VLLM_BASE_URL", "http://localhost:8000/v1")
    if transport is not None:
        return httpx.Client(transport=transport, base_url=base_url)
    return httpx.Client(base_url=base_url)


def rollout_step(
    state: RolloutGroupState,
    transport: httpx.BaseTransport | None = None,
) -> RolloutGroupState:
    """Run one step: sample G completions, update prefixes, check gate."""
    model = os.environ.get("VLLM_MODEL", "Qwen/Qwen2.5-7B-Instruct")
    client = _build_client(transport)

    messages = [{"role": "user", "content": state.prompt}]
    payload = {
        "model": model,
        "messages": messages,
        "n": state.group_size,
        "max_tokens": 64,
        "temperature": 0.8,
    }

    resp = client.post("/chat/completions", json=payload, timeout=30.0)
    resp.raise_for_status()
    data = resp.json()

    # Extract one action string per rollout
    new_actions = [
        choice["message"]["content"].strip()
        for choice in data["choices"]
    ]

    # Append to existing prefixes (concatenate step actions with separator)
    if not state.action_prefixes:
        state.action_prefixes = new_actions
    else:
        state.action_prefixes = [
            prev + " | " + new
            for prev, new in zip(state.action_prefixes, new_actions)
        ]

    tokens_this_step = data["usage"]["total_tokens"]
    state.total_tokens_used += tokens_this_step
    state.step += 1

    # Check early-stop gate [1]
    if should_stop_group(state.action_prefixes, tau=state.tau):
        state.stopped_early = True
        state.stop_step = state.step

    return state


def compute_baseline_tokens(
    state: RolloutGroupState,
    tokens_per_step: int,
) -> RolloutGroupState:
    """Estimate how many tokens a full rollout would have used."""
    state.baseline_tokens = tokens_per_step * state.max_steps
    return state

Step 4: The LangGraph agent loop

Wire the node into a LangGraph StateGraph. The loop runs until either the gate fires or max_steps is reached.

# filename: agent_graph.py
"""LangGraph graph definition for the selective-rollout agent."""
from __future__ import annotations
from typing import Any
import httpx

from langgraph.graph import StateGraph, END
from agent_state import RolloutGroupState
from agent_nodes import rollout_step


def _should_continue(state: RolloutGroupState) -> str:
    if state.stopped_early:
        return "stop"
    if state.step >= state.max_steps:
        return "stop"
    return "continue"


def build_graph(transport: httpx.BaseTransport | None = None):
    """Build the agent graph. Accepts an optional mock transport for testing."""

    def step_node(state: RolloutGroupState) -> RolloutGroupState:
        return rollout_step(state, transport=transport)

    builder = StateGraph(RolloutGroupState)
    builder.add_node("rollout_step", step_node)
    builder.set_entry_point("rollout_step")
    builder.add_conditional_edges(
        "rollout_step",
        _should_continue,
        {"continue": "rollout_step", "stop": END},
    )
    return builder.compile()

Verify the graph compiles and its nodes are wired correctly (no LLM call, no API key needed):

from agent_graph import build_graph

graph = build_graph()  # no transport -> structural check only
nodes = list(graph.get_graph().nodes)
print("nodes:", sorted(nodes))
assert "rollout_step" in nodes
print("graph_structure_ok")

Step 5: Run the agent with convergent vs. divergent rollouts

This block uses the mock transport to simulate both scenarios: a convergent group (all rollouts produce the same action, gate fires early) and a divergent group (rollouts differ, gate never fires, all steps run).

from mock_vllm import MockVLLMTransport
from agent_graph import build_graph
from agent_state import RolloutGroupState

PROMPT = "What is the capital of France? Use tools to answer step by step."

# --- Convergent group: gate should fire after step 1 ---
convergent_transport = MockVLLMTransport(convergent=True, seed=42)
convergent_graph = build_graph(transport=convergent_transport)

convergent_state = RolloutGroupState(
    prompt=PROMPT,
    group_size=3,
    max_steps=8,
    tau=0.05,
)
result_conv = convergent_graph.invoke(convergent_state)
print(f"Convergent: stopped_early={result_conv.stopped_early}, "
      f"stop_step={result_conv.stop_step}, "
      f"tokens_used={result_conv.total_tokens_used}")

# --- Divergent group: gate should NOT fire, runs all max_steps ---
diverse_transport = MockVLLMTransport(convergent=False, seed=42)
diverse_graph = build_graph(transport=diverse_transport)

diverse_state = RolloutGroupState(
    prompt=PROMPT,
    group_size=3,
    max_steps=8,
    tau=0.05,
)
result_div = diverse_graph.invoke(diverse_state)
print(f"Divergent:  stopped_early={result_div.stopped_early}, "
      f"stop_step={result_div.stop_step}, "
      f"tokens_used={result_div.total_tokens_used}")

assert result_conv.stopped_early is True, "Convergent group should have stopped early"
assert result_div.stopped_early is False, "Divergent group should run to completion"
assert result_conv.total_tokens_used < result_div.total_tokens_used, \
    "Convergent group must use fewer tokens"
print("rollout_scenario_assertions_passed")

Step 6: Benchmark script

The benchmark runs N_PROMPTS prompts, each with a randomly assigned convergent or divergent mock transport, and reports aggregate token savings.

# filename: benchmark.py
"""Benchmark: selective-rollout gate vs. full-rollout baseline."""
from __future__ import annotations
import random
from dataclasses import dataclass
from typing import List

from mock_vllm import MockVLLMTransport
from agent_graph import build_graph
from agent_state import RolloutGroupState

N_PROMPTS = 20
GROUP_SIZE = 4
MAX_STEPS = 8
TAU = 0.05
# Fraction of groups that are convergent (zero-variance); paper reports ~40% [1]
CONVERGENT_FRACTION = 0.40
TOKENS_PER_STEP_PER_GROUP = 60 * GROUP_SIZE  # matches mock usage


@dataclass
class BenchmarkResult:
    prompt_id: int
    convergent: bool
    steps_run: int
    tokens_used: int
    baseline_tokens: int
    stopped_early: bool


def run_benchmark(seed: int = 0) -> List[BenchmarkResult]:
    rng = random.Random(seed)
    results = []
    for i in range(N_PROMPTS):
        is_convergent = rng.random() < CONVERGENT_FRACTION
        transport = MockVLLMTransport(convergent=is_convergent, seed=seed + i)
        graph = build_graph(transport=transport)
        state = RolloutGroupState(
            prompt=f"Prompt {i}: answer step by step.",
            group_size=GROUP_SIZE,
            max_steps=MAX_STEPS,
            tau=TAU,
        )
        result = graph.invoke(state)
        baseline = TOKENS_PER_STEP_PER_GROUP * MAX_STEPS
        results.append(BenchmarkResult(
            prompt_id=i,
            convergent=is_convergent,
            steps_run=result.step,
            tokens_used=result.total_tokens_used,
            baseline_tokens=baseline,
            stopped_early=result.stopped_early,
        ))
    return results


def print_report(results: List[BenchmarkResult]) -> None:
    total_tokens = sum(r.tokens_used for r in results)
    total_baseline = sum(r.baseline_tokens for r in results)
    savings = total_baseline - total_tokens
    savings_pct = 100.0 * savings / total_baseline if total_baseline else 0.0
    n_stopped = sum(1 for r in results if r.stopped_early)
    n_convergent = sum(1 for r in results if r.convergent)

    print("=" * 52)
    print("Selective Rollout Benchmark")
    print("=" * 52)
    print(f"Prompts run          : {len(results)}")
    print(f"Convergent groups    : {n_convergent} "
          f"({100*n_convergent/len(results):.1f}%)")
    print(f"Groups stopped early : {n_stopped}")
    print(f"Baseline tokens      : {total_baseline:,}")
    print(f"Actual tokens used   : {total_tokens:,}")
    print(f"Token savings        : {savings:,} ({savings_pct:.1f}%)")
    print("=" * 52)

Verify it works

from benchmark import run_benchmark, print_report

results = run_benchmark(seed=7)
print_report(results)

total_tokens = sum(r.tokens_used for r in results)
total_baseline = sum(r.baseline_tokens for r in results)
savings_pct = 100.0 * (total_baseline - total_tokens) / total_baseline

# With 40% convergent fraction and early stopping after step 1,
# we expect meaningful token savings.
assert savings_pct > 5.0, f"Expected >5% savings, got {savings_pct:.1f}%"
print(f"benchmark_ok: {savings_pct:.1f}% token savings")

When every rollout of a prompt ends with the same reward, the group contributes no gradient, so the extra rollouts add no information.

Deploying against a real vLLM endpoint

When you have a GPU instance running vLLM, launch the server with the OpenAI-compatible API enabled:

# Run on your GPU instance (not in the sandbox -- requires GPU and vLLM)
# vllm serve Qwen/Qwen2.5-7B-Instruct \
#   --host 0.0.0.0 \
#   --port 8000 \
#   --api-key "" \
#   --max-model-len 4096
export VLLM_BASE_URL="http://<your-gpu-host>:8000/v1"
export VLLM_MODEL="Qwen/Qwen2.5-7B-Instruct"
echo "endpoint configured"

Then run the benchmark without the mock transport by editing benchmark.py to pass transport=None to build_graph. The graph will call your real vLLM server. The gate logic and token accounting are identical; only the HTTP transport changes.

For multi-node deployments, vLLM’s tensor-parallel flag (--tensor-parallel-size N) distributes the model across GPUs while the endpoint surface stays the same. The agent loop does not need to change.

Tuning the gate threshold tau

The threshold tau controls the tradeoff between early-stopping aggressiveness and the risk of stopping a group that would have diverged:

  • tau = 0.0: gate never fires (equivalent to no early stopping)
  • tau = 0.05: the default from [1], conservative, fires only when prefixes are nearly identical
  • tau = 0.15: more aggressive, fires when prefixes share roughly 85% of characters
  • tau = 0.30: fires on moderate divergence; may stop groups that would have produced useful gradient signal

The paper [1] reports that tau = 0.05 on ALFWorld with Qwen2.5-7B yields the 10.7% wall-clock speedup without degrading held-out success rate. For a different task distribution or model, sweep tau over [0.02, 0.05, 0.10, 0.20] and plot token savings against held-out reward to find the Pareto-optimal setting.

from benchmark import run_benchmark, print_report
from agent_state import RolloutGroupState
from mock_vllm import MockVLLMTransport
from agent_graph import build_graph
import random

TAU_VALUES = [0.0, 0.05, 0.10, 0.20]
N = 30

print(f"{'tau':>6}  {'savings%':>10}  {'stopped':>8}")
for tau in TAU_VALUES:
    rng = random.Random(99)
    total_tokens = 0
    total_baseline = 0
    n_stopped = 0
    for i in range(N):
        is_conv = rng.random() < 0.40
        transport = MockVLLMTransport(convergent=is_conv, seed=i)
        graph = build_graph(transport=transport)
        state = RolloutGroupState(
            prompt=f"Q{i}",
            group_size=4,
            max_steps=8,
            tau=tau,
        )
        r = graph.invoke(state)
        baseline = 60 * 4 * 8
        total_tokens += r.total_tokens_used
        total_baseline += baseline
        if r.stopped_early:
            n_stopped += 1
    pct = 100.0 * (total_baseline - total_tokens) / total_baseline
    print(f"{tau:>6.2f}  {pct:>9.1f}%  {n_stopped:>8}/{N}")
print("tau_sweep_done")

Troubleshooting

ModuleNotFoundError: No module named 'editdistance' Run uv pip install editdistance (or pip install editdistance). The package provides a fast C extension for Levenshtein distance and is not bundled with any of the other dependencies.

Gate never fires even on identical prefixes Check that tau is greater than zero and that action_prefixes is being populated correctly. Print state.action_prefixes after the first step. If all entries are empty strings, the vLLM response format may differ from what agent_nodes.py expects (look at choice["message"]["content"]).

httpx.ConnectError when running against a real vLLM server Confirm the server is listening: curl http://<host>:8000/v1/models. If the server is up but the agent times out, increase the timeout=30.0 argument in rollout_step or reduce max_tokens in the payload.

Token savings are zero in the benchmark This happens when tau=0.0 (gate disabled) or when all mock transports are set to convergent=False. Confirm CONVERGENT_FRACTION > 0 and tau > 0 in benchmark.py.

LangGraph raises InvalidUpdateError on state mutation LangGraph’s StateGraph expects node functions to return a new state or a dict of updates, not mutate in place. The nodes in this tutorial return the modified RolloutGroupState object directly, which works with the dataclass schema. If you switch to a TypedDict schema, return a dict of changed keys instead.

vLLM OOM on large group sizes vLLM batches all n completions for a single request together. With group_size=8 and a 7B model, peak VRAM can exceed 24 GB. Reduce group_size to 4 or enable --enable-chunked-prefill in the vLLM server flags to spread the memory pressure.

Next steps

  • Integrate with GRPO training: the gate in rollout_gate.py is framework-agnostic. Drop it into a TRL or OpenRLHF training loop to replicate the 10.7% wall-clock speedup reported in [1].
  • Adaptive tau scheduling: start with a high tau (aggressive stopping) early in training when the model is uncertain, then anneal toward zero as the policy stabilizes and group diversity increases.
  • Per-prompt tau calibration: log the gate’s false-positive rate (groups stopped early that would have been non-zero-variance) and use a lightweight online estimator to set prompt-specific thresholds.
  • Streaming rollouts: vLLM supports streaming completions via SSE. Rewrite rollout_step to consume tokens incrementally and check the gate after each token rather than after each full completion, enabling sub-step early stopping.

FAQ

How does the prefix edit-distance gate decide when to stop a rollout group?

The gate computes mean pairwise normalized edit distance across all action prefixes in a group. When the distance falls below a threshold tau (default 0.05), all rollouts have converged on the same action prefix and the group is stopped early, since it contributes no gradient signal.

What token savings does early stopping deliver in practice?

The benchmark shows token savings of 5-15% depending on the convergence fraction and tau threshold. With 40% of groups converging (as reported in the Selective Rollout paper), the tutorial achieves approximately 10% token savings at tau=0.05.

Can this technique be used with a real vLLM server, or only with the mock transport?

The mock transport is for sandbox testing only. To deploy against a real vLLM endpoint, set VLLM_BASE_URL to your server address and pass transport=None to build_graph(). The gate logic and token accounting remain identical; only the HTTP transport changes.

How should tau be tuned for a different task or model?

Sweep tau over values like [0.02, 0.05, 0.10, 0.20] and plot token savings against held-out reward to find the Pareto-optimal setting. The paper reports tau=0.05 works well for ALFWorld with Qwen2.5-7B; other tasks may benefit from different thresholds.

What happens if the gate stops a group that would have produced useful gradient signal?

This is a false positive. Lowering tau reduces false positives but stops fewer groups overall. The paper’s tau=0.05 is conservative and avoids accuracy regression; more aggressive thresholds (tau=0.15+) risk stopping groups prematurely.