Preparation

Convert a searched tree into actor advantages and critic targets. DirectBranchTd is the paper’s recipe: critic targets are direct-branch backups, and actor advantages are segment TD errors over refined values.

ActorSample

Store one segment’s training tokens, masks, and scalar advantage. prepare_samples() creates these records; the trainer uses weight as the segment’s learning signal.

from focus import ActorSample

Definition

@dataclass
class ActorSample

Fields

node_id: int
spans: tuple[PackedSequence, ...]
group_id: int
estimator: str
weight: float = 0.0
diagnostics: Mapping[str, float] = _NO_DIAGNOSTICS
search_pass: int | None = None

ActorSample.edge_tokens

@property
def edge_tokens(self) -> int

Count trainable tokens across the entire edge. Use this to normalize by edge when a chat template splits one edge into several training spans.

Batch

The actor and critic samples one prompt produced, with collection statistics and failures. Convert it to trainer records with to_records().

from focus import Batch

Definition

@dataclass(frozen=True)
class Batch

Fields

actor: tuple[ActorSample, ...] = ()
critic: tuple[CriticSample, ...] = ()
reward_config: RewardConfig | None = None
stats: Mapping[str, float] = field(default_factory=lambda: _NO_STATS)
failures: tuple[tuple[int, str], ...] = ()

Batch.len

def __len__(self) -> int

Batch.to_records

def to_records(self, group_index: int, *, actor_factory: Any=None, critic_factory: Any=None, min_abs_advantage: float | None=None) -> list[Any]

Flatten the batch into trainer records: one per generated actor span and one per critic target. Records are dictionaries unless you pass factories for your trainer’s types. group_index identifies the prompt group; min_abs_advantage drops actor records whose advantage is too small to matter.

Batch.describe

def describe(self) -> str

Render each actor edge with its advantage and diagnostics, then each critic target, one per line.

CriticSample

Pair a serialized checkpoint context with its target remaining return. The critic learns to predict target from context.

from focus import CriticSample

Definition

@dataclass(frozen=True)
class CriticSample

Fields

node_id: int
context: str
target: float
diagnostics: Mapping[str, float] = field(default_factory=dict)

ESTIMATOR_KEY

Name the metadata field identifying the estimator that prepared a training record.

from focus.export import ESTIMATOR_KEY

Definition

ESTIMATOR_KEY = 'estimator'

LANE_KEY

Name the metadata field distinguishing actor records from critic records.

from focus.export import LANE_KEY

Definition

LANE_KEY = 'lane'

Lane

Restrict a training record’s destination to "actor" or "critic".

from focus.export import Lane

Definition

type Lane = Literal['actor', 'critic']

critic_targets

Critic training samples: the direct-branch backup $\bar G$ at every expanded checkpoint that was scored, paired with the context the critic read.

from focus.preparation import critic_targets

Definition

def critic_targets(values: TreeValues) -> list[CriticSample]

prepare_samples

Turn a searched tree into actor advantages and critic targets with a chosen recipe. Collector.prepare() calls this with the collector’s reward settings.

from focus import DirectBranchTd, RewardConfig

from focus.preparation import prepare_samples

batch = prepare_samples(
    tree,
    reward_config=RewardConfig(),
    estimator=DirectBranchTd(),
)

Critic targets do not depend on the recipe: every expanded checkpoint the critic scored regresses on the mean of its direct-branch backups. DirectBranchTd measures actor advantages against refined values; estimator=None prepares only critic targets.

Definition

def prepare_samples(state: SchedulerState[AnyRollout], *, reward_config: RewardConfig, estimator: Estimator | None=DirectBranchTd()) -> Batch