Estimation

Assign actor advantages to collected trajectory segments. Implement Estimator.advantages(edges, values) or use DirectBranchTd or Grpo. Training-data preparation calls the estimator and computes critic targets separately.

Estimator

An actor recipe: assign a coefficient to every searched edge of a settled tree. Implement advantages(); the tree’s values are computed once and passed in.

from focus.estimation import Estimator

class RewardOnly(Estimator):
    def advantages(self, edges, values):
        return [edge.reward for edge in edges]

This recipe uses each edge’s immediate reward as its coefficient. Return a float per edge, or an Advantage to attach diagnostics. Critic targets do not depend on the recipe.

Definition

class Estimator(ABC)

Fields

name: str = ''
requires_completed_paths: bool = False

Estimator.advantages

Required implementation.

@abstractmethod
def advantages(self, edges: Sequence[Edge], values: TreeValues) -> Sequence[float | Advantage]

Return one coefficient per edge, in order. values holds each checkpoint’s prior, backup, and refined value.

DirectBranchTd

Compute the paper’s segment TD advantages from refined state values. Preparation supplies the shared values; this recipe assigns an actor weight to each searched edge.

from focus import DirectBranchTd

batch = collector.prepare(tree, estimator=DirectBranchTd())

With $n$ direct continuations, $\tilde V = (\kappa V_\phi + n \bar G)/(\kappa + n)$, where $\kappa$ is RewardConfig.value_prior_strength. An edge of $d$ steps with reward $R$ gets advantage $R + \gamma^d \tilde V(c') - \tilde V(c)$. With $\kappa = 1$, a prediction of 0.8 and two continuations returning 0 and 1 give a critic target of 0.5 and a refined value of 0.6. This is the collector’s default recipe. Critic targets are computed separately from the same tree values, so changing the actor recipe leaves those targets unchanged.

Definition

class DirectBranchTd(Estimator)

DirectBranchTd.advantages

def advantages(self, edges: Sequence[Edge], values: TreeValues) -> Sequence[Advantage]

Return the segment TD advantage of each edge from the tree’s refined values. Requires a critic prediction at every parent checkpoint.

Grpo

Group-relative advantages: each edge’s mean finished return, centered and scaled within the prompt’s tree. Needs no critic prediction.

from focus import Grpo

batch = collector.prepare(tree, estimator=Grpo())

Use it as a baseline recipe on the same searched tree. Critic targets are still the direct-branch backups, so the critic trains the same way under either recipe.

Definition

class Grpo(Estimator)

Grpo.advantages

def advantages(self, edges: Sequence[Edge], values: TreeValues) -> Sequence[Advantage]

Return each edge’s normalized mean finished return.

Advantage

An actor coefficient with optional diagnostics. Diagnostics appear in exported records and Batch.describe() but never enter the loss.

from focus.estimation import Advantage

Definition

@dataclass(frozen=True)
class Advantage

Fields

value: float
diagnostics: Mapping[str, float] = field(default_factory=lambda: MappingProxyType({}))

Edge

One searched continuation, from a parent checkpoint to a child, with its weighted reward and its length in steps.

from focus.estimation import Edge

Definition

@dataclass(frozen=True)
class Edge

Fields

node: Node[AnyRollout]
parent: Node[AnyRollout]
reward: float
duration: int