Train with Tinker

Use the MATH harness to collect continuations and update the policy through Tinker. FOCUS selects which continuations to collect and computes an advantage for each segment. Tinker then trains on those segments, using their original tokens and generation probabilities. Each actor update overlaps with collection of the next batch, keeping the sampling policy at most one optimizer step behind the learner.

With Tinker, your script owns the training loop. focus.integrations.tinker provides functions to convert data, update the actor, and run the one-step asynchronous loop. You supply a function that builds a collector for one problem. Train with Slime uses the same function, so both trainers can share the task code below.

Before you begin

Use the environment and context manager from Quick start: MATH. In the FOCUS checkout, run python -m pip install -e '.[tinker]', install tinker==0.35.0 for this example, and set TINKER_API_KEY for your account.

Save these files together in a working directory and open them in your editor:

• math_quickstart.py: the environment and rollout harness.
• math_training.py: search configuration and collector construction.
• train_tinker.py: data loading, training, and checkpoint saving.

The steps below explain the code you will customize. The files include the imports and command-line entry points.

1. Describe the task: one collector per problem

Each training item is a Problem with a prompt for the policy and a label for grading. Here, the label is the reference solution; the policy never sees it. For each problem, the trainer integration passes its current sampling policy to your collector function.

Start with the search configuration in math_training.py:

math_training.py

def search_config():
    """A FIFO warm-up, then selective allocation by expected critic-error reduction."""
    return Selective(budget=32, warmup=16, max_concurrency=1)

Selective starts with a first-in, first-out phase with a budget of 16 task transitions and at most two branches per state. The next phase has a budget of 32 transitions and prioritizes expected critic-error reduction. Here max_concurrency=1 runs one continuation at a time across the tree. A transition is one generated MATH segment; each phase lets its running continuation finish, so actual usage can exceed the launch budget.

Connect the search to the task with a starting critic:

math_training.py

def make_collector(policy, solution, *, critic=None, search=None, execution=None):
    if critic is None:
        # A uniform starting prediction for the first collection round.
        critic = Critic.constant(0.5, context=math_context(policy))
    return Collector(
        make_agent(policy, solution, policy.decode, execution=execution),
        critic=critic,
        search=search if search is not None else search_config(),
    )

Critic.constant(0.5) predicts 0.5 for every unfinished state, giving the first round the same starting estimate at each checkpoint. math_context() decodes the recorded prompt and generated tokens at each checkpoint. At the root, it renders the initial chat prompt. It preserves a partial solution’s thinking and whitespace without inserting another chat turn. The same text is sent to the critic and saved in its training records.

How the critic context preserves the solution prefix

math_training.py

def math_context(policy):
    """Give the critic the exact MATH prefix, including unfinished reasoning."""
    def render(rollout):
        if rollout.turns:
            last = rollout.turns[-1]
            tokens = last.prefix + last.tokens
        else:
            tokens = policy.prepare(rollout.messages).tokens
        return policy.decode(tokens)
    return render

The trainer calls collector_for(policy, problem) with the policy for the current sampler:

math_training.py

def collector_for(policy, problem):
    """What both trainers ask for: the collector for one MATH problem under a policy.

    ``problem.prompt`` is the question the policy sees; ``problem.label`` is the
    reference solution, which stays inside the environment for grading.
    """
    return make_collector(policy, problem.label)

The function creates a fresh environment with that problem’s reference solution, keeping grading separate across concurrent problems. The integration uses each collector once.

2. What the integration collects

For each problem, the integration wraps the current Tinker sampling client in a TinkerPolicy. The sampling client is a frozen snapshot of the learner’s weights, so the learner can update while a FOCUS tree keeps using its original sampler. It then calls collector_for and collects one tree. profile="token_chunk" resolves the sampler’s tokenizer for a model used without tools; Segments in the shared make_agent() reads each reply as an exact 256-token segment.

Preparation computes the tree values once. For each expanded state with a critic prediction, the mean direct-branch backup $\bar G$ becomes a critic target. Combining that mean with the critic’s prediction gives the refined value $\tilde V$ that DirectBranchTd uses to compute segment advantages. A collection with failed expansions or critic predictions raises unless you pass allow_partial=True.

The export labels each record by its training role, or lane:

Lane	Learning signal	Training input
Actor	Segment advantage, stored in reward	Exact tokens, loss mask, and behavior log probabilities
Critic	Backup target, stored in target	The checkpoint’s serialized context

Each generation records its sampling snapshot’s log probabilities. These stay with the exported tokens after the learner updates its weights. Both lanes carry metadata["training"]: the batch index, the problem’s id, and behavior_step, the number of completed optimizer updates when the snapshot was created. Critic records also carry learner_step, identifying the update that trained on their batch.

To change the phases, priorities, or budgets, build a different collector in collector_for. To select another recipe or inspect a tree, use the collection APIs.

3. Turn actor segments into next-token targets

A language model predicts the next token. For an exported sequence [prompt, segment], Tinker’s input is tokens[:-1] and its target is tokens[1:]; the behavior log probabilities and mask shift by the same position. focus.integrations.tinker.to_datum(record) performs this conversion and returns a Tinker Datum.

Suppose a segment has advantage $0.3$. Every generated token in that segment receives advantage $0.3$; every prompt token receives zero. A segment with advantage $-0.2$ receives negative weights on its generated tokens. This teaches the actor which steps improved the expected outcome while leaving conditioning text masked. See Tinker’s data format.

The exported reward is the actor weight, used directly. The recorded log probabilities supply the behavior-policy term in the learner-to-behavior probability ratio.

4. Compute gradients and apply the update

update_actor(training_client, records, learning_rate=...) converts the actor records and calls forward_backward_async() for the clipped policy-gradient loss. After that finishes, it calls optim_step_async() and waits for the optimizer update before returning. A sampler published afterwards therefore holds the updated weights. Collection on an existing sampler continues during both operations.

The default loss is Tinker’s ppo with clipping thresholds $1 \pm 0.2$, matching the paper’s token-level PPO objective; pass loss_fn and loss_fn_config to change it. See Tinker’s PPO loss.

5. Export the critic targets for the next round

Both trainer integrations export critic targets for a separate value-model training loop. They pass each batch’s critic records to a critic sink, such as JsonlDirectory(path). This sink writes one JSONL file per batch, with a context and target for each expanded nonterminal state. Fit your value model to these targets with mean squared error:

$$\mathcal L_{\mathrm{critic}} = \frac{1}{N}\sum_{j=1}^{N} \left(V_\phi(\bar c_j)-\bar G(\bar c_j)\right)^2.$$

For example, a record with target 0.6 asks the critic to predict remaining return 0.6 from that partial solution. The actor’s signed advantage has a different purpose and stays in the actor lane.

Train the scalar critic in your value-model loop, then expose it through a Critic. Render its context the same way during training and collection. Each tree keeps fixed actor and critic snapshots; updated models apply to newly launched collections. This example uses Critic.constant(0.5) until you supply a trained critic.

Connect a trained scalar critic

For a value model served through vLLM pooling, create one client for the run and close it when training finishes. This optional helper can replace the train(...) call in main(); pass the training client, model name, problems, and your critic endpoint settings:

from focus import Critic
from focus.backends import VLLMCritic
from focus.integrations import JsonlDirectory
from focus.integrations.tinker import train
from math_training import make_collector, math_context

async def train_with_critic(training, model, problems, critic_url, critic_model):
    client = VLLMCritic(base_url=critic_url, model=critic_model, reduce="last")

    def collector_for(policy, problem):
        critic = Critic.from_fn(client.predict, context=math_context(policy))
        return make_collector(policy, problem.label, critic=critic)

    try:
        await train(
            training, model=model, problems=problems,
            collector_for=collector_for, steps=3, learning_rate=1e-5,
            profile="token_chunk", critic_sink=JsonlDirectory("critic-targets"),
        )
    finally:
        await client.aclose()

critic_url is the vLLM server root, which serves the pooling endpoint, and critic_model is its served model name. reduce="last" selects the last pooled score. Train this score to predict remaining task return: on our binary-outcome task, 0.7 means an estimated 70% chance of eventual success. Each collector wraps the shared client’s prediction method with its own math_context(policy), preserving the text used for critic training. Keep the served critic weights fixed for this run. Adjust the training options to match your experiment; the helper uses three updates and the same critic-output directory as the default launch.

6. Overlap collection and training with one-step lookahead

First collect batch 0 with the initial policy $\pi_{\theta_0}$. Then collect batch 1 from that same snapshot while the learner updates on batch 0. After the update, publish $\pi_{\theta_1}$ for the next collection. Batch 1 still carries the probabilities from $\pi_{\theta_0}$.

For a three-update run, the schedule is:

Phase	Learner	Collection
Warm-up	Wait for the first batch	Batch 0 with $\pi_{\theta_0}$
Update 0	Train on batch 0: $\theta_0 \to \theta_1$	Batch 1 with $\pi_{\theta_0}$
Update 1	Train on batch 1: $\theta_1 \to \theta_2$	Batch 2 with $\pi_{\theta_1}$
Update 2	Train on batch 2: $\theta_2 \to \theta_3$	All requested batches collected

focus.integrations.tinker.train(...) runs this schedule with at most one upcoming collection. If collection finishes first, its batch waits for the learner; if training finishes first, the learner waits for the batch. Before every update, the loop checks the batch index and computes lag = step - batch.behavior_step. With async_depth=1, it accepts a lag of zero or one.

Every batch trains exactly once, including the last, and no unused collection launches afterwards. Once an optimizer step finishes, the loop publishes a new sampler if another collection remains. It runs the SDK’s blocking snapshot call in a worker thread.

Set async_depth=0 to collect each batch after the preceding update, using the current learner weights. With either schedule, keep the actor and critic snapshots fixed within each tree, including when you add critic updates.

If the loop stops early, it cancels and awaits any outstanding local collection. An update that has already reached Tinker can still complete remotely; a failed run stops rather than retrying that update.

7. Assemble the entry point and save training state

In train_tinker.py, main() reads the MATH records and creates a LoRA training client. Its central call connects that client to your collector:

train_tinker.py

await train(
    training, model=args.model, problems=problems, collector_for=collector_for,
    steps=args.steps, async_depth=args.async_depth, learning_rate=args.learning_rate,
    # The model is used without tools; Segments reads each reply as a token chunk.
    profile="token_chunk",
    critic_sink=JsonlDirectory(args.critic_output),
    on_step=lambda report: print(json.dumps(asdict(report))),
)

read_problems() accepts one JSON record, a JSON array, or JSONL. Each record supplies problem and solution, plus an optional id, and becomes a Problem. Each training batch collects one problem; the loop cycles through them until it reaches steps updates. on_step receives a report per update: batch index, behavior step, learner step, policy lag, and Tinker’s metrics.

After the final batch, save the learner’s model and optimizer and await the returned future:

train_tinker.py

checkpoint = await training.save_state_async(name=f"after-{args.steps}-updates")
saved = await checkpoint.result_async()
print(f"Training checkpoint: {saved.path}")
print(f"Critic targets: {args.critic_output}")

The surrounding try/finally closes the service after success, failure, or interruption. See Tinker’s checkpoint API and session lifecycle.

Data loading and session cleanup

train_tinker.py

def read_problems(path):
    """Read MATH records from one JSON object, a JSON array, or JSONL."""
    path = Path(path)
    text = path.read_text()
    rows = ([json.loads(line) for line in text.splitlines() if line.strip()]
            if path.suffix == ".jsonl" else json.loads(text))
    if isinstance(rows, dict):
        rows = [rows]
    if not isinstance(rows, list) or not rows or any(
        not isinstance(row, dict) or
        not all(isinstance(row.get(key), str) and row[key].strip() for key in ("problem", "solution"))
        for row in rows
    ):
        raise ValueError("Provide MATH records with nonempty problem and solution strings")
    return [Problem(row["problem"], row["solution"], row.get("id", i)) for i, row in enumerate(rows)]


async def main(args):
    import tinker
    problems = read_problems(args.problem_file)
    if args.steps < 1:
        raise ValueError("--steps must be positive")
    service = tinker.ServiceClient()
    status = "errored"
    try:
        training = await service.create_lora_training_client_async(base_model=args.model, rank=32)
        await train(
            training, model=args.model, problems=problems, collector_for=collector_for,
            steps=args.steps, async_depth=args.async_depth, learning_rate=args.learning_rate,
            # The model is used without tools; Segments reads each reply as a token chunk.
            profile="token_chunk",
            critic_sink=JsonlDirectory(args.critic_output),
            on_step=lambda report: print(json.dumps(asdict(report))),
        )
        checkpoint = await training.save_state_async(name=f"after-{args.steps}-updates")
        saved = await checkpoint.result_async()
        print(f"Training checkpoint: {saved.path}")
        print(f"Critic targets: {args.critic_output}")
        status = "success"
    except (asyncio.CancelledError, KeyboardInterrupt):
        status = "interrupted"
        raise
    finally:
        await service.close(status=status)

To start another run from that checkpoint, use service.create_training_client_from_state_with_optimizer_async(checkpoint_path) instead of creating a fresh LoRA client. The new run starts with an empty pipeline and a sampler from the restored learner; step labels start at zero. Write its critic records to a new directory to keep the runs separate.

A learner checkpoint stores model and optimizer state. A FOCUS rollout checkpoint stores an agent’s state for branching, as shown in BrowseComp-Plus.

8. Run three overlapping collection and update steps

Install FOCUS with its Tinker client and the math verifier:

python -m pip install -e '.[tinker]'
python -m pip install 'tinker==0.35.0' 'math-verify[antlr4_13_2]==0.9.0'

From the directory containing the three files, set TINKER_API_KEY for your account and FOCUS_MODEL to a model available to your Tinker service with a compatible chat tokenizer. Supply a MATH record as shown in the quick start, or a dataset of records:

python train_tinker.py \
  --model "$FOCUS_MODEL" \
  --problem-file problems.jsonl \
  --steps 3 \
  --async-depth 1 \
  --learning-rate 1e-5

This command performs three actor updates, saves the final training checkpoint, and writes critic targets from all three batches to the critic-targets directory. Each printed line reports one update; the expected lags are 0, 1, 1. Use --async-depth 0 for synchronous collection or --steps 1 to check a single update.

Complete source

The MATH task: search and collector

math_training.py

"""The MATH collector both trainer tutorials hand to their FOCUS integration."""
from focus import Collector, Critic, Selective

from math_quickstart import make_agent


def search_config():
    """A FIFO warm-up, then selective allocation by expected critic-error reduction."""
    return Selective(budget=32, warmup=16, max_concurrency=1)


def math_context(policy):
    """Give the critic the exact MATH prefix, including unfinished reasoning."""
    def render(rollout):
        if rollout.turns:
            last = rollout.turns[-1]
            tokens = last.prefix + last.tokens
        else:
            tokens = policy.prepare(rollout.messages).tokens
        return policy.decode(tokens)
    return render


def make_collector(policy, solution, *, critic=None, search=None, execution=None):
    if critic is None:
        # A uniform starting prediction for the first collection round.
        critic = Critic.constant(0.5, context=math_context(policy))
    return Collector(
        make_agent(policy, solution, policy.decode, execution=execution),
        critic=critic,
        search=search if search is not None else search_config(),
    )


def collector_for(policy, problem):
    """What both trainers ask for: the collector for one MATH problem under a policy.

    ``problem.prompt`` is the question the policy sees; ``problem.label`` is the
    reference solution, which stays inside the environment for grading.
    """
    return make_collector(policy, problem.label)

Tinker launch

train_tinker.py

"""Train a Tinker actor with one-step asynchronous FOCUS MATH collection."""
import argparse
import asyncio
from dataclasses import asdict
import json
from pathlib import Path

from focus.integrations import JsonlDirectory, Problem
from focus.integrations.tinker import train
from math_training import collector_for


def read_problems(path):
    """Read MATH records from one JSON object, a JSON array, or JSONL."""
    path = Path(path)
    text = path.read_text()
    rows = ([json.loads(line) for line in text.splitlines() if line.strip()]
            if path.suffix == ".jsonl" else json.loads(text))
    if isinstance(rows, dict):
        rows = [rows]
    if not isinstance(rows, list) or not rows or any(
        not isinstance(row, dict) or
        not all(isinstance(row.get(key), str) and row[key].strip() for key in ("problem", "solution"))
        for row in rows
    ):
        raise ValueError("Provide MATH records with nonempty problem and solution strings")
    return [Problem(row["problem"], row["solution"], row.get("id", i)) for i, row in enumerate(rows)]


async def main(args):
    import tinker
    problems = read_problems(args.problem_file)
    if args.steps < 1:
        raise ValueError("--steps must be positive")
    service = tinker.ServiceClient()
    status = "errored"
    try:
        training = await service.create_lora_training_client_async(base_model=args.model, rank=32)
        await train(
            training, model=args.model, problems=problems, collector_for=collector_for,
            steps=args.steps, async_depth=args.async_depth, learning_rate=args.learning_rate,
            # The model is used without tools; Segments reads each reply as a token chunk.
            profile="token_chunk",
            critic_sink=JsonlDirectory(args.critic_output),
            on_step=lambda report: print(json.dumps(asdict(report))),
        )
        checkpoint = await training.save_state_async(name=f"after-{args.steps}-updates")
        saved = await checkpoint.result_async()
        print(f"Training checkpoint: {saved.path}")
        print(f"Critic targets: {args.critic_output}")
        status = "success"
    except (asyncio.CancelledError, KeyboardInterrupt):
        status = "interrupted"
        raise
    finally:
        await service.close(status=status)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--model", required=True)
    parser.add_argument("--problem-file", required=True, help="MATH JSON object, JSON array, or JSONL")
    parser.add_argument("--steps", type=int, default=3, help="Actor updates; cycle through the input records")
    parser.add_argument("--async-depth", type=int, choices=(0, 1), default=1,
                        help="1 overlaps next collection with training; 0 collects synchronously")
    parser.add_argument("--critic-output", default="critic-targets",
                        help="Directory for the critic records, one JSONL file per batch")
    parser.add_argument("--learning-rate", type=float, default=1e-5)
    asyncio.run(main(parser.parse_args()))

Collection and execution with Tinker

Run train_tinker.py as one Python process. It owns the environments, rollout trees, and training loop. Tinker serves generation and optimizer requests remotely, so this tutorial needs no Ray cluster or separate collection service.

The run has two independent concurrency controls. --async-depth 1 overlaps one upcoming batch with the current update. Selective(max_concurrency=...) controls simultaneous branch continuations within that batch. To allow four continuations at once, return collectors with a wider search:

from focus import Selective

def collector_for(policy, problem):
    search = Selective(budget=32, warmup=16, max_concurrency=4)
    return make_collector(policy, problem.label, search=search)

This keeps the same rollout budget and one-batch lookahead. All branches in a tree use the same sampling client and fixed critic; increasing concurrency does not introduce another model snapshot. Throughput depends on the capacity available to your Tinker account.

Each collector closes its local execution services when its collection finishes. Your script owns the Tinker session: the loop cancels and awaits any pending collection before main() closes that session. For a trainer that manages its own inference workers, see Train with Slime.