Skip to content

Allocation

Track the evidence a search gathers at each checkpoint: observed returns, purchases, and in-flight continuations.

Model returns in [0, 1] with a Beta prior centered on the critic prediction. This is the default calibration model for ErrorReduction, including tasks with binary correctness rewards.

from focus.allocation import BetaBernoulli
from focus.statistics import RunningStats

stats = RunningStats.of([0.0, 1.0])
error, variance = BetaBernoulli().posterior(value=0.8, stats=stats, kappa=1.0)
assert round(error, 2) == 0.10
assert round(variance, 2) == 0.24

One failed and one successful continuation revise the prior mean from 0.8 to 0.6. The model estimates squared critic error as 0.10 and next-outcome variance as 0.24. Critic values or return means outside [0, 1] are clamped with a warning; choose Gaussian for unbounded returns.

Definition

class BetaBernoulli(CalibrationModel)
def __init__(self) -> None
def posterior(self, value: float, stats: RunningStats, kappa: float) -> tuple[float, float]

Use the critic prediction and observed returns to form a Beta posterior. Return the posterior expected squared critic error and Bernoulli predictive variance.

Define the statistical model that turns a critic prediction and observed returns into an estimated squared value error and the predictive variance of another return. ErrorReduction uses these two quantities to estimate the benefit of another continuation.

from focus import ErrorReduction

from focus.allocation import CalibrationModel
from focus.statistics import RunningStats

class KnownVariance(CalibrationModel):
    def __init__(self, variance: float):
        if variance <= 0:
            raise ValueError("variance must be positive")
        self.variance = variance

    def posterior(self, value: float, stats: RunningStats, kappa: float):
        mean = (kappa * value + stats.n * stats.mean) / (kappa + stats.n)
        mean_variance = self.variance / (kappa + stats.n)
        return (value - mean) ** 2 + mean_variance, self.variance + mean_variance

priority = ErrorReduction(KnownVariance(variance=0.25))

This example assumes normally distributed returns with known variance 0.25. The first result estimates squared critic error, including uncertainty in the return mean. The second is the predictive variance of one additional return. value is the frozen critic prediction, stats summarizes observed suffix returns, and positive kappa is the critic prior strength. ErrorReduction consumes these statistics; training targets come from a separate recursive backup over the completed tree.

Definition

class CalibrationModel(ABC)

Required implementation.

@abstractmethod
def posterior(self, value: float, stats: RunningStats, kappa: float) -> tuple[float, float]

Return (expected_squared_error, predictive_variance). The squared error compares the supplied critic prediction with the unknown expected return. Predictive variance includes uncertainty about that expected return as well as variation between individual outcomes. Implement this method for your return distribution; RewardConfig.value_prior_strength supplies positive kappa.

Model real-valued returns with a normal approximation. Use the observed sample variance once two returns are available; v_prior supplies the variance before then, and var_floor bounds it away from zero.

from focus import ErrorReduction

from focus.allocation import Gaussian

priority = ErrorReduction(Gaussian(v_prior=0.25, var_floor=0.001))

Use this model for real-valued returns. Before two observations, v_prior=0.25 supplies the outcome variance. Afterward, the sample variance is floored at 0.001. The predictive variance also includes uncertainty in the return mean.

Definition

class Gaussian(CalibrationModel)
def __init__(self, *, v_prior: float=0.25, var_floor: float=0.001) -> None
def posterior(self, value: float, stats: RunningStats, kappa: float) -> tuple[float, float]

Return the squared difference between the critic and refined mean plus posterior mean variance, together with the predictive variance of the next return.

Track the continuation outcomes observed from each checkpoint across the passes of one search. Allocation uses these observations to decide where another continuation is useful.

from focus.allocation import RolloutLedger

Definition

class RolloutLedger(AllocationLedger[AnyRollout])
def __init__(self, reward_config: RewardConfig, observations: dict[int, Observations] | None=None) -> None
def bind(self, view: SchedulerView[AnyRollout]) -> None
def record(self, parent: Node[AnyRollout], chain: list[Node[AnyRollout]]) -> None

The root-metadata key under which a search stores its rollout evidence ledger.

from focus.allocation.evidence import EVIDENCE_KEY

Definition

EVIDENCE_KEY = 'rollout_evidence'

Store the suffix-return observations and reward configuration shared by allocation and training preparation.

from focus.allocation.evidence import RolloutEvidence

Definition

@dataclass(frozen=True)
class RolloutEvidence

Fields

reward_config: RewardConfig
observations: dict[int, Observations]

Read the rollout evidence ledger a search recorded on a tree. Raises ValueError for a tree without one.

from focus.allocation.evidence import evidence_of

Definition

def evidence_of(view: SchedulerView[AnyRollout], config: RewardConfig) -> RolloutEvidence

Compute the marginal reduction in the blended value estimate’s mean squared error from squared critic error, predictive return variance, effective evidence count, and prior strength. A nonpositive result declines the continuation.

from focus.allocation import retraining_gain

Definition

def retraining_gain(e2: float, outcome_var: float, n_eff: int, kappa: float) -> float