Skip to content

Priority

Rank eligible checkpoints using collected evidence. Implement Priority.__call__(node) or use FIFO, expected critic-error reduction, or token entropy. Search schedules continuations and owns phase budgets and branching caps.

Rank candidate checkpoints; the search continues the highest one. Implement __call__ and return a score, or None when the checkpoint should not be continued.

from focus import Phase

from focus.priority import NodeInfo, Priority

class Shallowest(Priority):
    def __call__(self, node: NodeInfo) -> float | None:
        return -node.depth

phase = Phase(Shallowest(), 16)

This priority continues the shallowest checkpoints first. Ties go to the earlier node. Set requires_critic = True when the score reads node.prior, so a collector without a critic rejects the search.

Definition

class Priority(ABC)

Fields

name: str = ''
requires_critic: bool = False

Required implementation.

@abstractmethod
def __call__(self, node: NodeInfo) -> float | None

Score one candidate checkpoint. Higher scores are continued first; None, NaN, or negative infinity mark the checkpoint ineligible for now.

Continue where one more outcome is expected to reduce critic error most. The gain combines the critic’s estimated squared error at the checkpoint with the variance of its outcomes, and shrinks as observations accumulate.

from focus import ErrorReduction, Phase

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

For a checkpoint with NN observations, including pending ones, the gain is g=(2κ2e^2+(N−κ)σ^2)/(κ+N)3g = (2\kappa^2 \hat e^2 + (N-\kappa)\hat\sigma^2)/(\kappa+N)^3, where κ\kappa is RewardConfig.value_prior_strength, e^2\hat e^2 the estimated squared critic error, and σ^2\hat\sigma^2 the outcome variance. Checkpoints with no positive gain are not continued. calibration estimates e^2\hat e^2 and σ^2\hat\sigma^2; the default BetaBernoulli expects returns in [0, 1], and Gaussian handles real-valued returns. Requires a critic.

Definition

class ErrorReduction(Priority)
def __init__(self, calibration: CalibrationModel | None=None) -> None
def __call__(self, node: NodeInfo) -> float | None

Continue checkpoints in the order they became eligible: first in, first out. Use it to collect initial evidence before selective allocation.

from focus import Fifo, Phase

phase = Phase(Fifo(), 16, max_children=2)

The priority decides the order; the phase supplies the budget and branching caps.

Definition

class Fifo(Priority)
def __call__(self, node: NodeInfo) -> float

Rank checkpoints by uncertainty in their children’s generated tokens. Higher token-entropy scores receive priority.

from focus import Phase, TokenEntropy

phase = Phase(TokenEntropy(window=8), 32, max_pending=2)

Each child edge is scored by its highest sliding-window mean of −plog⁡p-p \log p over generated tokens, and a checkpoint by the mean over its children. A checkpoint without scored children ranks first. tags chooses which turns count. The backend must return per-token log probabilities.

Definition

class TokenEntropy(Priority)
def __init__(self, window: int=8, tags: Sequence[TurnTag] | None=('task',)) -> None
def __call__(self, node: NodeInfo) -> float

What a priority may read about one candidate checkpoint: its depth, critic prior, observed outcomes, running continuations, and children.

from focus.priority import NodeInfo

Definition

@dataclass(frozen=True)
class NodeInfo

Fields

node: Node[AnyRollout]
tree: SchedulerView[AnyRollout]
reward: RewardConfig
observations: Observations
arrival: int
@property
def id(self) -> int

The node’s identifier in the tree.

@property
def depth(self) -> int

Edges from the root.

@property
def prior(self) -> float | None

The critic’s prediction at this checkpoint, or None when it was not scored.

@property
def outcomes(self) -> RunningStats

Running statistics of the returns observed from this checkpoint.

@property
def pending(self) -> int

Continuations from this checkpoint still running.

@property
def children(self) -> list[Node[AnyRollout]]

The checkpoint’s child nodes.