Export

Inspect a collected tree and convert prepared training data into records for your trainer.

CriticSampleFactory

Define a callable that converts a checkpoint context and critic target into the trainer’s critic-record type. Accept context, target, group_index, and additional metadata keywords.

from focus.export import CriticSampleFactory

from focus.export import to_samples

def critic_record(*, context, target, group_index, **extra):
    return {"input": context, "target": target, "group": group_index}

records = batch.to_records(0, critic_factory=critic_record)

The callable maps each serialized checkpoint context and direct-branch backup target into your value trainer’s input format. It handles critic records; the actor sample factory handles token sequences and advantages.

Definition

class CriticSampleFactory[T](Protocol)

CriticSampleFactory.call

def __call__(self, *, context: str, target: float, group_index: int, **extra: Any) -> T

SampleFactory

Define a callable that converts a packed actor span into a trainer-specific sample. Accept its behavior log probabilities, prepared advantage, lane, group index, and additional metadata keywords.

from focus.export import SampleFactory

from focus import ActorSample, Batch

from focus.extend import PackedSequence

def actor_record(*, span, logprobs, reward, lane, group_index, **extra):
    return {
        "tokens": span.tokens,
        "loss_mask": span.trainable,
        "behavior_logprobs": logprobs,
        "advantage": reward,
        "group": group_index,
    }

span = PackedSequence(
    tokens=(10, 20), trainable=(False, True),
    logprobs=(0.0, -0.5),
)
edge = ActorSample(
    node_id=1, spans=(span,), group_id=0,
    estimator="DirectBranchTd", weight=0.4,
)
batch = Batch(actor=(edge,))
records = batch.to_records(0, actor_factory=actor_record)
assert records[0]["advantage"] == 0.4

The first token is conditioning text and the second is generated. The factory preserves that mask, the behavior log probability, and the prepared advantage of 0.4. Its **extra parameter accepts edge and estimator metadata supplied by the exporter.

Definition

class SampleFactory[T](Protocol)

SampleFactory.call

def __call__(self, *, span: PackedSequence, logprobs: tuple[float, ...], reward: float, lane: Lane, group_index: int, **extra: Any) -> T

format_edges

Format actor weights and critic targets from a prepared batch for inspection.

from focus.export import format_edges

print(format_edges(batch))

Inspect which segments receive positive or negative weights and which targets supervise the critic.

Definition

def format_edges(result: Batch) -> str

format_tree

Format the search tree, showing parent-child relationships, terminal states, rewards, and critic predictions.

from focus.export import format_tree

print(format_tree(SchedulerView(tree), batch))

Include the prepared batch to show actor weights beside their tree nodes.

Definition

def format_tree(view: SchedulerView[AnyRollout], result: Batch | None=None) -> str

to_samples

Convert a prepared batch into actor and critic records for a trainer. The default output is a list of dictionaries; factories can produce a trainer’s own record types.

from focus.export import to_samples

records = to_samples(batch, group_index=0)

This is what Batch.to_records() calls. The example assigns the records to prompt group 0. Each generated span becomes an actor record carrying its advantage, loss mask, and behavior log probabilities; each critic target becomes a critic record.

Definition

def to_samples(result: Batch, group_index: int, *, sample_factory: SampleFactory[Any] | None=None, critic_sample_factory: CriticSampleFactory[Any] | None=None, min_abs_advantage: float | None=None) -> list[Any]