Search

Choose where to spend the rollout budget. A Search runs phases in order; each phase ranks checkpoints with a priority and stops at its budget. Selective is the paper’s configuration.

Phase

One stage of a search: a priority that ranks checkpoints, a budget of task transitions, and per-checkpoint branching caps.

from focus import ErrorReduction, Phase

phase = Phase(ErrorReduction(), 64, max_pending=2)

max_children caps total branches from a checkpoint, counting finished and running ones; max_pending caps only those running at once. Continuations still running when the budget is reached finish normally. sampling overrides the agent’s sampling parameters for this phase’s continuations.

Definition

@dataclass(frozen=True)
class Phase

Fields

priority: Priority
budget: int
max_children: int | None = None
max_pending: int | None = None
sampling: SamplingParams | None = None

Search

Run search phases in order over one tree. Each phase continues the checkpoints its priority ranks highest until its budget is spent; max_concurrency bounds simultaneous continuations across the tree.

from focus import ErrorReduction, Fifo, Phase, Search

search = Search(
    [
        Phase(Fifo(), 16, max_children=2),
        Phase(ErrorReduction(), 64, max_pending=2),
    ],
    max_concurrency=4,
)

The first phase collects initial evidence in arrival order, with at most two branches per checkpoint. The second continues where another outcome is expected to reduce critic error most, with at most two continuations from a checkpoint running at once. This is the configuration Selective(64, warmup=16) builds.

Definition

class Search

Search.init

def __init__(self, phases: Sequence[Phase], *, max_concurrency: int=4) -> None

Search.requires_critic

@property
def requires_critic(self) -> bool

Whether any phase ranks checkpoints with critic predictions. A collector without a critic rejects such a search.

Search.run

async def run(self, tree: SchedulerState[AnyRollout], agent: Agent[Any, Any], *, ledger: RolloutLedger, scorer: NodeScorer[AnyRollout] | None=None) -> None

Run every phase over tree, continuing checkpoints with agent.

Each phase regates the tree with its caps and spends its budget on top of what earlier phases already spent. A phase whose priority needs critic predictions is skipped when every prediction failed.

Selective

The paper’s search: a first-in, first-out warm-up, then selective allocation by expected critic-error reduction.

from focus import Selective

search = Selective(budget=64, warmup=16)

The warm-up spends 16 transitions with at most two branches per checkpoint (warmup_width). Selective allocation then spends 64 transitions; total width is uncapped, and each checkpoint runs at most two continuations at once (max_pending_per_node). calibration chooses the error model: the default Beta-Bernoulli model expects returns in [0, 1]; use Gaussian() for real-valued returns.

Definition

class Selective(Search)

Selective.init

def __init__(self, budget: int, *, warmup: int, warmup_width: int=2, max_pending_per_node: int=2, max_concurrency: int=4, calibration: CalibrationModel | None=None) -> None