Integrations

Hand FOCUS collections to a trainer. Every integration asks for the same thing: a CollectorFor that builds the collector for one Problem under the trainer’s current policy, plus a critic sink. focus.integrations.tinker supplies a loop you run; focus.integrations.slime supplies hooks Slime calls.

CollectorFor

What every trainer integration asks of you: a function from the trainer’s current policy and one Problem to the Collector that should collect it. The Tinker and Slime integrations take the same function.

from focus.integrations import CollectorFor

from focus import Agent, Collector, Critic, Segments, Selective

def collector_for(policy, problem):
    agent = Agent(
        policy=policy,
        env=MathEnv(problem.label, policy.decode),
        system_prompt="Solve the problem.",
        context=Segments(256),
    )
    return Collector(
        agent,
        critic=Critic.constant(0.5),
        search=Selective(budget=32, warmup=16),
    )

Build a fresh environment per problem so concurrent problems never share a reference answer. The integration uses each collector once inside async with, so the collector’s own setting decides whether its execution services close afterwards; policies and critics stay with their owner.

Definition

type CollectorFor = Callable[[Policy[Any], Problem], Collector]

CriticSink

Where critic records go. Neither trainer trains the critic, so each collection’s critic records are handed to this callable; your value-model loop reads them back.

from focus.integrations import CriticSink

Definition

type CriticSink = Callable[[Sequence[dict[str, Any]]], None]

JsonlDirectory

A critic sink that writes each collection’s critic records to its own JSONL file in a directory. One file per call, so concurrent collectors and worker processes never interleave lines; a collection without critic records writes nothing.

from focus.integrations import JsonlDirectory

sink = JsonlDirectory("critic-targets")

Pass it as critic_sink= to the Tinker loop or the Slime generate hook. Each line is one checkpoint context and its direct-branch backup target.

Definition

class JsonlDirectory

JsonlDirectory.init

def __init__(self, directory: str | Path) -> None

JsonlDirectory.call

def __call__(self, records: Sequence[dict[str, Any]]) -> None

Problem

One training item: the prompt the policy sees, a label that stays out of the prompt (typically a reference answer the environment grades against), an id recorded in exported metadata, and free-form metadata. Both trainer integrations hand problems to your CollectorFor.

from focus.integrations import Problem

problem = Problem("Solve 3x + 7 = 22.", r"\boxed{5}", id="linear")

The Tinker loop takes a list of problems; the Slime hooks build one from each Slime sample, with prompt and label filled by Slime’s --input-key and --label-key.

Definition

@dataclass(frozen=True)
class Problem

Fields

prompt: Any
label: Any = None
id: Any = None
metadata: Mapping[str, Any] = field(default_factory=lambda: MappingProxyType({}))

collect_records

Collect one problem with one collector and return its (actor, critic) records. Both integrations call this; use it directly for a trainer FOCUS does not integrate yet. A collection with failed expansions or critic predictions raises unless allow_partial=True.

from focus.integrations import collect_records

Definition

async def collect_records(collector: Collector, problem: Problem, *, group_index: int, workspace_for: Callable[[Problem], Workspace] | None=None, min_abs_advantage: float | None=None, allow_partial: bool=False) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]

GenerateHook

The signature Slime calls a custom generate function with: (args, sample, sampling_params), returning the samples to train on. sampling_params is accepted but does not override FOCUS’s sampling settings.

from focus.integrations.slime import GenerateHook

Definition

class GenerateHook(Protocol)

GenerateHook.call

async def __call__(self, args: Any, sample: Any, sampling_params: Any=None) -> Any

advantages

Advantage hook: broadcast each segment’s FOCUS advantage over its response tokens, overriding the estimator Slime’s argument parser selected. Refuses advantage normalization; returns are zero because critic targets are exported separately. Name it with --custom-advantage-function-path focus.integrations.slime.advantages.

from focus.integrations.slime import advantages

Definition

def advantages(args: Any, rollout_data: dict[str, Any]) -> None

make_evaluate

Build Slime’s evaluation generate function from the same CollectorFor: one unbranched rollout of the collector’s agent, returned as a Slime sample whose reward is the outcome rather than a signed advantage.

from focus.integrations.slime import make_evaluate

Definition

def make_evaluate(collector_for: CollectorFor, *, profile: str | None=None, workspace_for: Callable[[Problem], Workspace] | None=None) -> GenerateHook

make_generate

Build Slime’s custom generate function from your CollectorFor. For each sample it builds the router’s policy, wraps the sample as a Problem, collects one tree with a fresh collector, sends critic records to critic_sink, and returns the actor segments as Slime samples.

from focus.integrations import JsonlDirectory
from focus.integrations.slime import make_generate

generate = make_generate(
    collector_for,
    profile="token_chunk",
    critic_sink=JsonlDirectory("critic-targets"),
)

Point --custom-generate-function-path at the module attribute holding generate. One collector per sample keeps concurrent prompts’ reference answers apart. workspace_for gives each prompt its own workspace, min_abs_advantage drops near-zero actor edges after search, and allow_partial=True keeps a collection with failures.

Definition

def make_generate(collector_for: CollectorFor, *, profile: str | None=None, critic_sink: CriticSink | None=None, workspace_for: Callable[[Problem], Workspace] | None=None, min_abs_advantage: float | None=None, allow_partial: bool=False) -> GenerateHook

preserve_advantages

Reward post-process hook: return the FOCUS advantages unchanged, so Slime’s group centering does not alter them. Name it with --custom-reward-post-process-path focus.integrations.slime.preserve_advantages.

from focus.integrations.slime import preserve_advantages

Definition

def preserve_advantages(args: Any, samples: list[Any]) -> tuple[list[float], ...]

sum_policy_tokens

Policy-loss reducer hook: sum masked token losses instead of averaging them, so each generated token contributes its segment’s advantage once. Requires context parallel size 1. Name it with --custom-pg-loss-reducer-function-path focus.integrations.slime.sum_policy_tokens.

from focus.integrations.slime import sum_policy_tokens

Definition

def sum_policy_tokens(total_lengths: Any, response_lengths: Any, loss_masks: list[Any], calculate_per_token_loss: bool=False) -> Callable[[Any], Any]

to_sample

Turn one actor record into a response-aligned Slime sample. The response starts at the first generated token; conditioning tokens inside it, such as tool results, keep loss mask zero. reward carries the FOCUS advantage, and every segment of one tree shares the source prompt’s rollout_id.

from focus.integrations.slime import to_sample

Definition

def to_sample(record: Mapping[str, Any], source: Any, tokenizer: Any, sample_type: Any) -> Any

PPO_CLIP

Tinker’s PPO loss configuration for the paper’s ratio clip, $1 \pm 0.2$.

from focus.integrations.tinker import PPO_CLIP

Definition

PPO_CLIP: Mapping[str, float] = {'clip_low_threshold': 0.8, 'clip_high_threshold': 1.2}

StepReport

What one actor update trained on: the batch index, the behavior step its sampler came from, the learner step, the policy lag between them, and Tinker’s metrics.

from focus.integrations.tinker import StepReport

Definition

@dataclass(frozen=True)
class StepReport

Fields

batch: int
behavior_step: int
learner_step: int
policy_lag: int
metrics: Mapping[str, Any] = field(default_factory=dict)

to_datum

Turn one actor record into a Tinker Datum: the model input is tokens[:-1], and targets, behavior log probabilities, and advantages are read one position later. Conditioning tokens get advantage zero.

from focus.integrations.tinker import to_datum

Definition

def to_datum(record: Mapping[str, Any]) -> Any

train

Run the one-step asynchronous Tinker loop: collect one problem per update from a frozen sampler, train on it, and publish a new sampler, overlapping the next collection with the current update.

from focus.integrations import JsonlDirectory
from focus.integrations.tinker import train

reports = await train(
    training_client,
    model=model,
    problems=problems,
    collector_for=collector_for,
    steps=3,
    learning_rate=1e-5,
    profile="token_chunk",
    critic_sink=JsonlDirectory("critic-targets"),
)

With async_depth=1 (the default) each batch is at most one update behind the learner and every batch trains exactly once; async_depth=0 publishes a fresh sampler before each collection. profile resolves the sampler’s format as for TinkerPolicy. Critic records reach the sink with their learner step, and on_step receives one StepReport per update. An outstanding collection is cancelled and joined if the loop stops early.

Definition

async def train(training_client: Any, *, model: str, problems: Sequence[Problem], collector_for: CollectorFor, steps: int, learning_rate: float, critic_sink: CriticSink | None=None, async_depth: Literal[0, 1]=1, profile: str | None=None, loss_fn: str='ppo', loss_fn_config: Mapping[str, float]=PPO_CLIP, workspace_for: Callable[[Problem], Workspace] | None=None, min_abs_advantage: float | None=None, allow_partial: bool=False, on_step: Callable[[StepReport], None] | None=None) -> list[StepReport]

update_actor

Take one Tinker optimizer step on a batch’s actor lane. Gradients are awaited before the optimizer step is submitted, and the optimizer result before returning, so a sampler published afterwards holds the updated weights. The default loss is ppo with ratio clip $1 \pm 0.2$.

from focus.integrations.tinker import update_actor

Definition

async def update_actor(training_client: Any, records: Sequence[Mapping[str, Any]], *, learning_rate: float, loss_fn: str='ppo', loss_fn_config: Mapping[str, float]=PPO_CLIP) -> Any