Quick start: MATH

A partial solution can lead to several different answers. By saving that prefix and trying new continuations, we can see how often the model succeeds from the same starting point.

This quick start builds a rollout harness for that experiment on MATH problems. The model writes a solution in 256-token segments. Each segment becomes an edge in a rollout tree, and its endpoint becomes a state from which FOCUS can resume generation. The final segment can be shorter when the model finishes.

Segments(256) keeps the original prompt and generated token IDs, and makes each segment boundary a checkpoint. Branches resume from the same partial solution without compacting it.

Problem
   └─ 256 tokens → partial solution
                     ├─ 256 tokens → another step → final answer
                     └─ alternative continuation → final answer

The two outcomes help estimate the value of the shared partial solution. The harness below connects a model to a verifier and preserves the prefix for each continuation.

Before you begin

Use Python 3.12 or newer and the FOCUS checkout from Installation. Install the math verifier with python -m pip install 'math-verify[antlr4_13_2]==0.9.0'.

Save math_quickstart.py in your working directory and open it in your editor. This is a complete program. The excerpts below explain how it works and which parts to change for your own task.

Start with the segment size and rollout limit:

math_quickstart.py

CHUNK_TOKENS = 256
MAX_CHUNKS = 16

The file includes a scripted mode for checking the harness. The complete source is also available at the end of this page.

1. Understand the segment action

Each generation becomes a TokenChunk: the original token IDs and a finished flag. A length finish leaves the response open; a stop finish completes it. The environment uses finished to decide when to grade the answer.

Segments(256) tells the agent to use token-chunk actions. The agent selects TokenChunkParser to read the generated IDs and finish reason; the parser needs no length setting. Set the segment size once, in Segments.

2. Write the environment and verifier

A state is a point in the rollout from which generation can resume. Its context is denoted $\bar c$. In this task, the context is the prompt followed by all tokens generated so far. The environment also stores an immutable MathState containing those generated tokens and the segment count.

Env[MathState, TokenChunk] declares the state and action types. Pass the tokenizer’s decode method as decode_tokens so the environment can turn the completed token sequence into text for grading:

math_quickstart.py

@dataclass(frozen=True)
class MathState:
    tokens: tuple[int, ...] = ()
    chunks: int = 0


class MathEnv(Env[MathState, TokenChunk]):
    def __init__(self, solution: str, decode_tokens: Callable[[Sequence[int]], str]):
        self.gold = parse_answer(solution)
        if not self.gold:
            raise ValueError("The reference solution needs a parseable \\boxed{answer}")
        self.decode_tokens = decode_tokens

    async def reset(self) -> MathState:
        return MathState()

    async def step(
        self, state: MathState, action: TokenChunk, *, workspace=None,
    ) -> StepResult[MathState]:
        next_state = MathState(state.tokens + action.tokens, state.chunks + 1)
        correct = False
        if action.finished:
            # Decode the complete token sequence; a character can span chunks.
            answer = parse_answer(self.decode_tokens(next_state.tokens))
            correct = bool(answer and verify(self.gold, answer))
        return StepResult(
            next_state=next_state, done=action.finished,
            reward_outcome=float(correct),
        )

MathEnv.step() appends one segment and returns a new state. Sibling branches therefore share the same starting prefix and accumulate their own continuations.

A length finish reason means the model reached the 256-token boundary: save a checkpoint and continue. A stop finish reason means the solution ended: extract the last boxed answer and compare it with the reference using Math-Verify. The verifier checks mathematical equivalence, so $\frac{1}{2}$ and $0.5$ can receive the same reward.

The environment returns reward 1 for a correct final answer and 0 otherwise. Intermediate segments have zero reward. A missing or incomplete box receives zero final reward. Grading decodes the entire generated sequence, so an answer can cross a segment boundary.

The file includes two answer-parsing helpers. boxed_expression() finds the final balanced box, including nested LaTeX braces. parse_answer() converts its contents into the symbolic form accepted by Math-Verify. Keep these helpers when using MATH; replace the verifier when adapting the harness to another task.

Answer-parsing helpers

math_quickstart.py

def boxed_expression(text):
    """Extract the last boxed answer, including nested LaTeX braces."""
    start = text.rfind(r"\boxed{")
    if start < 0:
        return ""
    start += len(r"\boxed{")
    depth = 1
    for i in range(start, len(text)):
        if text[i] == "{" and (i == 0 or text[i - 1] != "\\"):
            depth += 1
        elif text[i] == "}" and (i == 0 or text[i - 1] != "\\"):
            depth -= 1
            if depth == 0:
                return text[start:i]
    return ""


def parse_answer(text):
    expression = boxed_expression(text)
    if not expression:
        return []
    return parse(
        f"${expression}$", extraction_config=[LatexExtractionConfig()],
        fallback_mode="no_fallback",
    )

3. Preserve the exact continuation prefix

Each request generates up to 256 tokens, then returns control to the agent. The agent saves a checkpoint before requesting the next segment, making that prefix available for branching during collection.

Use Segments(256) to keep generating the same assistant response. This context manager prepares each request: it renders the initial chat prompt, then appends the original generated token IDs for later segments:

$$\bar c_i = \bar c_0 \mathbin{\Vert} y_1 \mathbin{\Vert} \cdots \mathbin{\Vert} y_i,$$

where $y_i$ is a generated segment and $\Vert$ denotes concatenation. A fork at $\bar c_i$ reuses that exact prefix and samples a new next segment.

The saved IDs preserve the unfinished answer and the context in which each token’s probability was measured. A fork copies this history exactly.

AppendOnly() also retains history, but prepares complete conversation turns, including tool results. Segments preserves one unfinished assistant response and rejects intervening messages or edits to its generated text. The segment size caps each request; the model’s stop condition can end it earlier. BrowseComp-Plus uses Summarize to replace long conversations with notes after rounds of tool calls.

For models such as Qwen3.5, rendering a new conversation turn can remove earlier thinking from the prompt. Exact continuation avoids re-rendering: even if a segment stops inside <think>, the next request receives the same prompt and all tokens generated so far, including whitespace. This keeps the reasoning prefix intact across segment boundaries and forks.

4. Connect the policy to vLLM

Set model to the served model identifier and base_url to the vLLM API address, such as http://127.0.0.1:8000/v1:

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained(model)
format = PolicyFormat("math-chunks", ChatCodec(tokenizer), TokenChunkParser())
policy = VLLMPolicy(
    model=model, format=format, base_url=base_url, api_key="EMPTY",
)

Use the same tokenizer for prompt encoding and answer decoding. ChatCodec(tokenizer) renders the model’s chat template, and TokenChunkParser() reads each generated segment. VLLMPolicy sends the prepared token IDs to the server and returns the generated IDs and log probabilities. The model must provide a chat template, and the server must support exact token input and output.

The same Segments context manager works with TinkerPolicy, SlimePolicy, SGLangPolicy, or a custom exact token-input Policy. The scripted mode uses a small character tokenizer in place of the Hugging Face tokenizer.

5. Assemble the rollout harness

math_quickstart.py

def make_agent(policy, solution, decode_tokens, *, execution=None):
    return Agent(
        policy=policy, env=MathEnv(solution, decode_tokens),
        system_prompt=r"Solve the problem step by step. Put your final answer in \boxed{}.",
        context=Segments(CHUNK_TOKENS),
        sampling_params=SamplingParams(temperature=0.7, logprobs=0),
        max_steps=MAX_CHUNKS, execution=execution,
    )

Agent generates a segment, applies it to the environment, and records the resulting checkpoint. Segments(CHUNK_TOKENS) sets the defaults this task needs: each request is capped at CHUNK_TOKENS, every settled segment boundary is available for branching, and no finishing instruction is appended when the horizon is reached. In this environment, one decision is one chunk of reasoning.

To change the segment size, change CHUNK_TOKENS. The execution service combines policy defaults with the request’s sampling parameters and validates the returned token count against that limit. A length result must contain exactly the requested number of tokens; a stop result can contain fewer. A response that exceeds the limit raises an error.

temperature=0.7 allows alternative continuations, while logprobs=0 requests the sampled tokens’ probabilities without additional alternatives. max_steps=16 limits a root-to-leaf path to 16 segments: at most $16 \times 256 = 4{,}096$ generated tokens. Forks inherit the remaining horizon.

Segments disables the extra final-answer request by default. If the model reaches the horizon before finishing, the agent marks the leaf truncated. It remains nonterminal and receives no final-answer reward. A naturally completed solution has done=True. Inspect both flags when recording results.

6. Generate a prefix, then try two continuations

math_quickstart.py

async def rollout(policy, tokenizer, problem, solution):
    agent = make_agent(policy, solution, tokenizer.decode)
    async with agent:
        root = await agent.start(problem)
        checkpoint = await agent.advance(root)
        print(f"First segment: {len(checkpoint.state.tokens)} tokens")
        if checkpoint.done:
            endpoints = [checkpoint]
        else:
            # Both branches inherit the same context and remaining horizon.
            left = await agent.run(checkpoint.fork())
            right = await agent.run(checkpoint.fork())
            endpoints = [left, right]
        for i, endpoint in enumerate(endpoints, 1):
            print(f"Branch {i}: chunks={endpoint.state.chunks}, "
                  f"reward={endpoint.reward_outcome:g}, truncated={endpoint.truncated}")
            print(tokenizer.decode(endpoint.state.tokens, skip_special_tokens=True))
    return checkpoint, endpoints

make_agent passes tokenizer.decode to MathEnv for grading completed answers. The saved prefixes remain token IDs.

start() combines the system instruction with the problem and resets the environment. advance() generates one segment and saves its resulting state. If the solution has finished, that state is already the endpoint.

Otherwise, fork() creates an independent continuation from the saved checkpoint. Each run() completes its branch using the remaining segment allowance. The original checkpoint stays unchanged, so both branches start from the same partial solution.

Inspect endpoint.state.tokens for the complete generated solution, endpoint.reward_outcome for correctness, and endpoint.truncated for a run that reached its horizon. endpoint.edge(checkpoint) returns only the turns generated after the fork. For a single unbranched attempt, use await agent.run_from(problem).

async with agent closes its execution services when the block ends, including after an error or cancellation. The policy client is closed separately in main() because it can be shared by several agents.

Call the rollout from your program’s async entry point, and close the model client when the work finishes:

try:
    await rollout(policy, tokenizer, problem, solution)
finally:
    await policy.aclose()

7. Run the program you built

Run the file’s scripted mode to check your environment before connecting a model:

python math_quickstart.py --offline

This runs the built-in task, “Solve for x: 3x + 7 = 22,” using a character tokenizer and scripted continuations with synthetic log probabilities. It saves a checkpoint after 256 tokens, then resumes it twice. The two branches finish with answers 5 and 4, receiving rewards 1 and 0.

To generate solutions with a model, connect to a running vLLM server that supports token-ID prompts and returns generated token IDs and log probabilities:

python math_quickstart.py \
  --model "$FOCUS_MODEL" \
  --base-url http://127.0.0.1:8000/v1

Set FOCUS_MODEL to the model served by that endpoint; its tokenizer must match the server’s tokenizer and provide a chat template. The model context window must fit the task prompt plus 4,096 generated tokens.

For a MATH problem, pass a JSON file containing the dataset’s problem and solution fields:

problem.json format

{
  "problem": "Solve for x: 3x + 7 = 22.",
  "solution": "Subtract 7 and divide by 3 to get \\boxed{5}."
}

python math_quickstart.py \
  --model "$FOCUS_MODEL" \
  --problem-file problem.json

The reference solution must contain a final \boxed{...} answer. The script sends only problem to the policy. solution stays inside the environment for grading. Use the built-in algebra problem to check the harness, then supply a dataset problem with --problem-file.

Complete source

Show math_quickstart.py

math_quickstart.py

"""MATH collection with 256-token edges and exact, uncompacted prefixes.

Run with --offline for a scripted check, or --model and --base-url for vLLM.
Requires FOCUS and math-verify[antlr4_13_2]==0.9.0.
Live runs connect to a vLLM server.
"""
import argparse
import asyncio
from collections.abc import Callable, Sequence
from dataclasses import dataclass
import json
from pathlib import Path

from math_verify import LatexExtractionConfig, parse, verify

from focus import Agent, Env, Segments, StepResult
from focus.backends import (
    ChatCodec, GenerateResult, Policy, PolicyFormat, SamplingParams,
    TokenChunk, TokenChunkParser, VLLMPolicy,
)

CHUNK_TOKENS = 256
MAX_CHUNKS = 16


def boxed_expression(text):
    """Extract the last boxed answer, including nested LaTeX braces."""
    start = text.rfind(r"\boxed{")
    if start < 0:
        return ""
    start += len(r"\boxed{")
    depth = 1
    for i in range(start, len(text)):
        if text[i] == "{" and (i == 0 or text[i - 1] != "\\"):
            depth += 1
        elif text[i] == "}" and (i == 0 or text[i - 1] != "\\"):
            depth -= 1
            if depth == 0:
                return text[start:i]
    return ""


def parse_answer(text):
    expression = boxed_expression(text)
    if not expression:
        return []
    return parse(
        f"${expression}$", extraction_config=[LatexExtractionConfig()],
        fallback_mode="no_fallback",
    )


@dataclass(frozen=True)
class MathState:
    tokens: tuple[int, ...] = ()
    chunks: int = 0


class MathEnv(Env[MathState, TokenChunk]):
    def __init__(self, solution: str, decode_tokens: Callable[[Sequence[int]], str]):
        self.gold = parse_answer(solution)
        if not self.gold:
            raise ValueError("The reference solution needs a parseable \\boxed{answer}")
        self.decode_tokens = decode_tokens

    async def reset(self) -> MathState:
        return MathState()

    async def step(
        self, state: MathState, action: TokenChunk, *, workspace=None,
    ) -> StepResult[MathState]:
        next_state = MathState(state.tokens + action.tokens, state.chunks + 1)
        correct = False
        if action.finished:
            # Decode the complete token sequence; a character can span chunks.
            answer = parse_answer(self.decode_tokens(next_state.tokens))
            correct = bool(answer and verify(self.gold, answer))
        return StepResult(
            next_state=next_state, done=action.finished,
            reward_outcome=float(correct),
        )


class CharacterTokenizer:
    """Offline fixture only: one character per token."""
    chat_template = "offline"

    def encode(self, text, **kwargs):
        return [ord(c) for c in text]

    def decode(self, tokens, **kwargs):
        return "".join(chr(t) for t in tokens)

    def apply_chat_template(self, conversation, *, add_generation_prompt, **kwargs):
        text = "".join(f"{m['role']}: {m['content']}\n" for m in conversation)
        return self.encode(text + ("assistant: " if add_generation_prompt else ""))


class ScriptedCompletion:
    """Two segments per attempt; alternate correct and incorrect final answers."""
    def __init__(self, tokenizer):
        self.tokenizer = tokenizer
        self.calls = 0
        self.endings = 0

    async def __call__(self, prefix, params):
        self.calls += 1
        assert params.max_tokens == CHUNK_TOKENS
        text = self.tokenizer.decode(prefix)
        if text.endswith("assistant: "):
            output = "Subtract seven, then divide by three. ".ljust(CHUNK_TOKENS)
            reason = "length"
        else:
            self.endings += 1
            output = r"\boxed{5}" if self.endings % 2 else r"\boxed{4}"
            reason = "stop"
        tokens = tuple(self.tokenizer.encode(output))
        return GenerateResult(
            prefix_tokens=tuple(prefix), tokens=tokens, text=output,
            stop_reason=reason, logprobs=[-0.5] * len(tokens),
        )


def make_agent(policy, solution, decode_tokens, *, execution=None):
    return Agent(
        policy=policy, env=MathEnv(solution, decode_tokens),
        system_prompt=r"Solve the problem step by step. Put your final answer in \boxed{}.",
        context=Segments(CHUNK_TOKENS),
        sampling_params=SamplingParams(temperature=0.7, logprobs=0),
        max_steps=MAX_CHUNKS, execution=execution,
    )


async def rollout(policy, tokenizer, problem, solution):
    agent = make_agent(policy, solution, tokenizer.decode)
    async with agent:
        root = await agent.start(problem)
        checkpoint = await agent.advance(root)
        print(f"First segment: {len(checkpoint.state.tokens)} tokens")
        if checkpoint.done:
            endpoints = [checkpoint]
        else:
            # Both branches inherit the same context and remaining horizon.
            left = await agent.run(checkpoint.fork())
            right = await agent.run(checkpoint.fork())
            endpoints = [left, right]
        for i, endpoint in enumerate(endpoints, 1):
            print(f"Branch {i}: chunks={endpoint.state.chunks}, "
                  f"reward={endpoint.reward_outcome:g}, truncated={endpoint.truncated}")
            print(tokenizer.decode(endpoint.state.tokens, skip_special_tokens=True))
    return checkpoint, endpoints


async def main(args):
    problem = "Solve for x: 3x + 7 = 22."
    solution = r"Subtract 7 and divide by 3 to get \boxed{5}."
    if args.problem_file:
        row = json.loads(Path(args.problem_file).read_text())
        problem, solution = row["problem"], row["solution"]

    if args.offline:
        tokenizer = CharacterTokenizer()
        format = PolicyFormat(
            "math-chunks", ChatCodec(tokenizer), TokenChunkParser(),
        )
        policy = Policy(
            format=format, model="scripted", complete=ScriptedCompletion(tokenizer),
        )
    else:
        from transformers import AutoTokenizer

        tokenizer = AutoTokenizer.from_pretrained(args.model)
        format = PolicyFormat(
            "math-chunks", ChatCodec(tokenizer), TokenChunkParser(),
        )
        policy = VLLMPolicy(
            model=args.model, format=format,
            base_url=args.base_url, api_key=args.api_key,
        )

    try:
        await rollout(policy, tokenizer, problem, solution)
    finally:
        await policy.aclose()


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--offline", action="store_true")
    parser.add_argument("--problem-file", help="One MATH JSON record with problem and solution")
    parser.add_argument("--model", help="Actor model/tokenizer path or Hugging Face ID")
    parser.add_argument("--base-url", default="http://127.0.0.1:8000/v1")
    parser.add_argument("--api-key", default="EMPTY")
    args = parser.parse_args()
    if not args.offline and not args.model:
        parser.error("Supply --model, or use --offline for the scripted example")
    if args.offline and args.problem_file:
        parser.error("The offline fixture solves only the built-in algebra problem")
    asyncio.run(main(args))

Continue to collection and training

Train with Tinker and Train with Slime add selective rollout collection, prepare actor advantages and critic targets, and pass those records to the trainer. BrowseComp-Plus builds a tool-using agent that summarizes its context and branches at compaction boundaries.