Getting to know the library
Start by defining your task and connecting a model. FOCUS handles the rollout loop, branching, and training-data preparation. You provide the task behavior and can replace the search or learning rules for your experiment.
Use the supplied implementations where they fit. The sections below explain which interface to extend when you need different behavior.
Task actions, state, rewards, and completion
Policy and parserModel connection and action syntax
ContextManagerWhat context the model sees
BranchPoints and PriorityWhich states are eligible and which to continue next
CriticCheckpoint context and remaining-return predictions
EstimatorActor advantages from collected evidence
Env and tools: define your task
Section titled “Env and tools: define your task”Env defines what the agent can do, how the task state changes, how rewards are assigned, and when the task finishes. Implement reset() to create the initial state and step() to apply an action and return the next state and observations.
For a tool-using agent, write the tools and connect them through ToolEnv. Keep mutable task data in the branch’s state or workspace so alternative continuations can evolve independently. A custom Workspace is useful when your tools need storage or fork behavior beyond the built-in implementations.
The MATH quick start builds an environment and verifier. BrowseComp-Plus builds browser tools and a workspace that keeps each branch’s navigation independent.
Policy and parser: connect your model
Section titled “Policy and parser: connect your model”Choose an existing policy backend for your model service. Its PolicyFormat pairs a codec, which renders the model input, with a parser, which turns output into environment actions.
If your model uses a different prompt format or action syntax, supply a matching ChatCodec or ActionParser. Implement a policy backend when connecting a service the library does not already support. The parser’s output must match the actions your environment accepts.
ContextManager: choose what the model sees
Section titled “ContextManager: choose what the model sees”ContextManager controls the context presented at each step. The supplied implementations retain the conversation (AppendOnly), continue an exact token prefix (Segments), or summarize the live context (Summarize).
Write a custom manager when your agent needs another way to retain, select, or rewrite its context. This choice is separate from which checkpoints search selects and how the resulting trajectories are used for training.
BranchPoints and Priority: choose where to continue
Section titled “BranchPoints and Priority: choose where to continue”Two interfaces control different parts of branching. BranchPoints decides whether a state is eligible: implement eligible(state) to inspect the task state, context, or workspace. Priority ranks eligible checkpoints: implement __call__(node) to score them using the available evidence.
Use the supplied rules unless your task or experiment needs another selection strategy. EveryDecision and SparseCheckpoints provide common eligibility rules. Selective supplies the paper’s search configuration, collecting initial evidence before prioritizing continuations by expected critic-error reduction. Configure rollout budgets and branch limits through its search phases.
Critic: define the value prediction
Section titled “Critic: define the value prediction”Critic maps checkpoint contexts to predictions of the remaining return. You supply the value model and decide which information from a rollout it should see. Search uses these predictions to allocate continuations; preparation combines them with observed outcomes.
Use an existing critic backend when it fits your model, or implement the interface for another scoring service or context representation. Keep the actor and critic fixed while collecting each tree; your training system updates them between collections.
Estimator: choose the actor learning signal
Section titled “Estimator: choose the actor learning signal”Estimator assigns actor advantages to collected segments. The default, DirectBranchTd, uses the paper’s segment TD advantages over refined values. Keep it for the standard FOCUS method; implement advantages(edges, values) when experimenting with another actor learning rule.
Critic targets are computed separately, so changing the actor estimator does not redefine them. The prepared batch retains generated tokens, behavior probabilities, masks, and learning signals for your trainer. Train with Tinker and Train with Slime show how to connect those records to an update loop.
Assemble the components
Section titled “Assemble the components”Agent runs your policy in the environment using the chosen context manager and branch-point rule. Collector connects that agent to a search, critic, and estimator. Call collect() to obtain a training batch, or use search() and prepare() separately to inspect a tree before preparing its learning signals.
Use focus to assemble a run and focus.backends for model clients. The extension interfaces and their implementations live in focus.branching, focus.context, focus.priority, and focus.estimation; focus.extend also exports the extension interfaces. FOCUS includes implementations for execution services, checkpoint storage, and batch export. See the Python API for their configuration options.
Set budgets at the scope you intend: a search phase limits task transitions, Agent(max_steps=...) limits rollout steps, and SamplingParams.max_tokens limits one model response. Compaction consumes a rollout step without a task transition. Running continuations finish after a phase stops launching new work, so inspect batch statistics for actual usage. See Phase and Execution for the budget controls.