This post is a hands-on introduction to reinforcement learning through training a search agent with group-relative policy optimization (GRPO). Search is a fun place to learn RL because it has so many levers, and each of them visibly changes how the model searches. It is also a domain where, in my experience, a well-designed reward function can influence how a model searches more effectively than system prompt changes or harness engineering.
Our approach is based loosely on the Harness-1 paper from Jiang et al. We use their training dataset and a much-simplified version of their search harness.
Training is done using the Tinker API from Thinking Machines. Many thanks to the team for the training credits to support this work.
The dataset
We use the RL split of the Harness-1 dataset for all experiments. The dataset contains 3,453 queries over 37,113 SEC filings split into 2,115,106 chunks.
Each query includes a preamble followed by a question. The agent must both retrieve evidence for facts in the preamble and the answer to the question. In the canonical version of this task, retrieving the answer is more important than evidence for facts. However, for simplicity we will treat both the same.
To keep initial ablations inexpensive, we start with just 256 training queries and 32 evaluation queries over a reduced corpus of 124,395 chunks. For our scaled runs, we expand to 1,024 training queries over 237,533 chunks.
The harness
Our agent operates as a search subagent. Given a query, it is meant to retrieve all relevant docs. We keep the harness deliberately simple, only providing lexical search tools. The agent’s output is a curated set it builds up as it goes rather than a single submission at the end. We chose this because it works well for a fixed turn budget: since the set exists from the start, an episode that runs out of turns will will return something gradable.
| Tool | What it does |
|---|---|
| bm25_search | Keyword search over every chunk. Returns an id, title, score and snippet per hit, plus a running count of chunks seen and curated so far. |
| grep_corpus | A short Python regex scan across the corpus. Returns an id, title, and snippet per hit. |
| read_document | The title and text of one chunk. Only ids the agent has already seen in a result list can be read. |
| curate | Add chunk ids to the result set. |
| drop_curated | Remove ids from the result set. |
| finish | End the episode and return the set. |
The system prompt is similarly plain. We name the tools and give a general strategy for search. We don’t try to encourage any specific search characteristics (e.g. high recall or high precision) so that the agent can discover the right behaviors during training.
Retrieval agent system prompt
You are a retrieval subagent in a multi-agent system. Your specific role is to identify and retrieve the most relevant documents from a large corpus to help another agent answer questions. You do NOT answer questions yourself - you only find and retrieve relevant documents.
The user message contains the query you need to find documents for.
Available Tools:
- bm25_search: Keyword search over the corpus
- grep_corpus: Text pattern matching with a short Python regex
- read_document: Read a specific document that looks promising but incomplete
- curate: Add relevant documents to your curated result set
- drop_curated: Remove curated documents that turned out to be irrelevant
- finish: End the search and return your curated set
Your Process:
- Break down the query into its key concepts and information needs (list each one explicitly)
- For each key concept, develop a specific search strategy that targets that concept
- Consider what types of documents and evidence would be most helpful for answering this query
- Plan several distinct, non-overlapping search strategies that approach the question from different angles
- Then execute your searches using multiple parallel tool calls.
Your Thinking:
After each round of searches, consider the following
- What do I know?: List the key topics, themes, or aspects of the question that your curated documents address. What specific information do you have?
- What should I search for next?: Systematically consider what search approaches, keywords, or document types you haven't yet tried that might yield valuable information.
- What should I curate or drop?: Curate documents as soon as you have verified they are relevant; drop curated documents that later look redundant or off-topic.
- Do I have enough information?: Given the question's complexity and requirements, do you have sufficient information to help answer it, or are there critical gaps?
- Decide if additional searches are needed (and if so, ensure they use genuinely different approaches and do not duplicate prior searches)
- Avoid getting stuck on a single search strategy - if one approach isn't yielding results, backtrack and try different approaches
Tactics to Consider:
- When queries fail, try different approaches or keywords to improve the results
- Avoid duplicate or redundant searches
- Execute multiple tool calls in parallel when possible
- Focus on gathering as much relevant information as possible; it is useful to get multiple perspectives on the same topic to confirm the information you have found is correct
- Follow explicit textual evidence rather than speculation
Output (IMPORTANT):
- YOU MUST use the curate tool. This is the ONLY way to return documents.
- Every response must contain at least one tool call. Never reply with plain text, and never end a response while still planning a search - emit that search as a tool call instead.
- As soon as you verify a document is relevant, call curate with its ID.
- When your curated set covers the query's information needs, call finish. Your curated set is scored even if you run out of turns, but finishing cleanly is always better than timing out.
We use gpt-oss-20b for all training runs. We chose this model because it is cheap to sample and train, allowing us to do many ablations.
Quick GRPO recap
We train with Dr. GRPO, a variant of GRPO. The loop is as follows
- Take a query and run the agent on it N times using the current policy (in this setting, essentially the model’s weights). Each run is a rollout, and the N tries together form a group.
- Score each rollout. This score is its reward.
- Subtract the group’s mean reward from each rollout’s reward. The resulting difference is that rollout’s advantage. In vanilla GRPO, we would also divide each advantage by the group’s standard deviation. Dr. GRPO drops this term to remove a bias: a query where the eight rollouts nearly agree has a small deviation. Dividing by standard deviation scales up the advantages, so the update can end up dominated by queries the policy is already consistent on.
- Update the policy (model weights) proportionally to the rollout’s advantage, so that the tokens in higher-advantage rollouts become more likely and tokens in lower-advantage rollouts become less likely.
We would also typically use ratio clipping during the policy update to help stabilize training. However, to simplify training for the sake of pedagogy, we will train in a fully synchronous setting with zero training-inference mismatch (which Tinker provides by default) and only do one update per batch.
We won’t go deeper into the math behind here, but there are many other great resources out there for understanding GRPO.
The reward
Each task here consists of a query, a corpus, a set of facts, and a set of supporting documents for each fact. The agent has to curate a set of documents it determines to be relevant to those facts, and success is evaluated on how well this curated set covers the target facts.
A reward function is essentially a kind of evaluation scheme for a rollout. A few metrics common to search retrieval work naturally here.
Recall measures the fraction of relevant documents our agent has curated. Since our search problem is fact-based, this will be the fraction of relevant facts the curated chunks cover. Precision measures the fraction of curated documents that are actually relevant.
We can’t use just recall or precision alone in our reward function. High recall is desirable in our problem space, but can be trivially maximized by curating every document in the corpus. This is useless for an end user! Why use a search agent at all? High precision can be achieved just as trivially by curating only a single correct chunk per query, which is useless in the other direction.
We need a metric that punishes both forms of degenerate behavior. F1, the harmonic mean of precision and recall, is the usual choice here:
F1=P+R2PR
This metric is only high when both precision and recall are high. This is a totally reasonable metric to choose for our reward function. But thinking about our problem space again… it’s likely that a missing document here costs far more to an end user than an extra, useless one. So, we might want to bias the reward towards recall. With this in mind, we will use F-beta as our reward function. This is a weighted version of F1 with a coefficient for setting how much recall matters more than precision:
Fβ=β2P+R(1+β2)PR
We set β = 4 for our reward function, which weighs recall sixteen times more than precision. I chose this number based on my prior experience and notes from the Harness-1 paper rather than a formal ablation. Under this metric, which we call F4, a set with recall 1.0 and precision 0.25 scores 0.85; under F1 the same set would score 0.40.
One more thing we add to our reward function is a punishment for not curating anything—an episode that curates nothing at all gets a flat -0.2. This helps with differentiating between rollouts where the model never gets to curate anything and ones where it only curates incorrect documents.
Putting it together, the reward for a rollout with curated set C is
r(C)={−0.2F4if ∣C∣=0otherwise
Initial explorations
The core idea behind reinforcement learning is to let models explore and reward them for behaviors that lead to good results. This means that RL works specifically by reinforcing behaviors the model already produces. If the model either always reaches for the same actions or always fails, then there is nothing for it to learn from.
This is especially true in GRPO, where learning depends on variation in rewards within each group. If all rollouts receive the same reward, their advantages are zero and the group provides no gradient.
This is why it can be helpful to gauge the RL-ability of a task before training. Best-of-N evals are a good proxy for this. Just sample N runs over the same task and report the highest score. A good candidate for RL is a model that has noticeably higher best-of-N scores than best-of-1 — this gap is the headroom RL can close by just sharpening the existing policy.
Below, we sample gpt-oss-20b four times on each of the eval queries. All rollouts are viewable.
The gap between best-of-4 and single trial F1 is large: best-of-4 F1 is 0.323 vs 0.195 for a single try. The F1 score is also precision-dominated. The model is likely being over-conservative in its curation.
Close inspection of the rollout supports this. On seven of the 32 queries, none of the four tries ever calls curate. The model stalls, exhausts all its turns, or runs out of context before curating a single document. Failures like these signal more of a harness problem than a search problem, and that is the kind of behavior reinforcement learning is good at fixing.
Learning rate sweep
Learning rate is a critical hyperparameter to get right. Just like in supervised fine-tuning, a bad learning rate can destabilize training. In reinforcement learning, a learning rate too high can cause the policy to collapse; one that is too small means the model never learns anything during training. And even when a run looks fine, I’ve found that further tuning the learning rate can often lead to incrementally better performance.
That said, if you’re an experienced practitioner and know directionally what the right hyperparameters are, you can skip this step for initial experiments. My opinion is that once training works stably, tuning the learning rate only starts to matter when you’re trying to squeeze out extra performance from your data.
All runs in this section and the next use the 256-query subset of our larger dataset, with 32 groups of 8 rollouts per step. Here are the results of a sweep over five different learning rates, each run for one epoch:
train reward
train reward | step | 5e-6 | 1e-5 | 5e-5 | 1e-4 | 5e-4 |
|---|---|---|---|---|---|
| 0 | -0.02 | -0.02 | -0.02 | -0.02 | -0.02 |
| 1 | -0.03 | -0.03 | -0.02 | -0.02 | 0.01 |
| 2 | -0.04 | -0.04 | -0.04 | -0.03 | -0.03 |
| 3 | -0.02 | -0.03 | -0.02 | 0.00 | 0.05 |
| 4 | 0.01 | -0.01 | 0.02 | 0.05 | 0.12 |
| 5 | -0.01 | 0.02 | 0.03 | 0.09 | 0.16 |
| 6 | -0.01 | 0.01 | 0.04 | 0.12 | 0.12 |
| 7 | -0.04 | -0.02 | 0.02 | 0.15 | 0.01 |
| 8 | -0.03 | -0.03 | 0.06 | 0.14 | 0.08 |
valid train rollouts
valid train rollouts | step | 5e-6 | 1e-5 | 5e-5 | 1e-4 | 5e-4 |
|---|---|---|---|---|---|
| 0 | 43% | 43% | 43% | 43% | 43% |
| 1 | 44% | 39% | 42% | 44% | 46% |
| 2 | 36% | 34% | 34% | 42% | 53% |
| 3 | 40% | 39% | 39% | 45% | 73% |
| 4 | 48% | 45% | 53% | 54% | 91% |
| 5 | 45% | 47% | 51% | 66% | 100% |
| 6 | 40% | 43% | 52% | 63% | 100% |
| 7 | 37% | 41% | 55% | 79% | 75% |
| 8 | 40% | 38% | 63% | 88% | 100% |
eval f1
eval f1 | step | 5e-6 | 1e-5 | 5e-5 | 1e-4 | 5e-4 |
|---|---|---|---|---|---|
| 0 | 0.19 | 0.19 | 0.19 | 0.19 | 0.19 |
| 4 | 0.18 | 0.21 | 0.22 | 0.20 | 0.12 |
| 8 | 0.16 | 0.14 | 0.16 | 0.19 | 0.03 |
Training results over one epoch with different learning rates. The valid rollouts chart shows the percentage of training rollouts that curate at least one doc.
We can see in here the effects of both a too large and a too small learning rate:
- 1e-5 and5e-6 barely move train reward. We clearly need a larger gradient signal to meaningfully update the policy.
- 5e-4 improves reward rapidly at first and then quickly destabilizes.
- 1e-4 and5e-5 train steadily, with1e-4 resulting in a higher final train reward and a higher eval F1.
Since lr=1e-4 performs the best in the sweep, we’ll use it for the rest of our experiments.
If you want a better idea of what the training collapse for lr=5e-4 looks like, expand the explorer below to see its eval rollouts at steps 0 and 8.
Batch size ablation
Conventional wisdom when training with GRPO is to use large batch sizes, where batch size is groups per step x group size. To study the effects of batch size on training here, we fix group size at 8 and ablate over groups per step.
The charts below are plotted against prompts seen rather than total updates, since all runs have a different total number of steps.
train reward
train reward | step | 8 groups/step | 16 groups/step | 32 groups/step | 64 groups/step |
|---|---|---|---|---|
| 0 | -0.03 | -0.03 | -0.03 | -0.03 |
| 8 | -0.01 | | | |
| 16 | 0.04 | 0.02 | | |
| 24 | 0.04 | | | |
| 32 | -0.06 | -0.07 | -0.01 | |
| 40 | -0.06 | | | |
| 48 | -0.00 | -0.06 | | |
| 56 | 0.04 | | | |
| 64 | 0.25 | 0.04 | -0.04 | -0.05 |
| 72 | 0.03 | | | |
| 80 | 0.06 | -0.01 | | |
| 88 | 0.11 | | | |
| 96 | 0.14 | 0.09 | -0.01 | |
| 104 | 0.03 | | | |
| 112 | 0.07 | 0.04 | | |
| 120 | 0.13 | | | |
| 128 | 0.21 | 0.20 | 0.01 | -0.01 |
| 136 | 0.24 | | | |
| 144 | 0.16 | 0.21 | | |
| 152 | 0.17 | | | |
| 160 | 0.02 | 0.10 | 0.10 | |
| 168 | 0.10 | | | |
| 176 | 0.13 | 0.14 | | |
| 184 | 0.19 | | | |
| 192 | 0.09 | 0.14 | 0.08 | 0.07 |
| 200 | 0.16 | | | |
| 208 | 0.14 | 0.16 | | |
| 216 | 0.07 | | | |
| 224 | 0.08 | 0.11 | 0.12 | |
| 232 | 0.13 | | | |
| 240 | 0.06 | 0.07 | | |
| 248 | 0.05 | | | |
| 256 | 0.17 | 0.12 | 0.16 | 0.05 |
valid train rollouts
valid train rollouts | step | 8 groups/step | 16 groups/step | 32 groups/step | 64 groups/step |
|---|---|---|---|---|
| 0 | 41% | 41% | 41% | 41% |
| 8 | 36% | | | |
| 16 | 55% | 45% | | |
| 24 | 60% | | | |
| 32 | 39% | 43% | 44% | |
| 40 | 66% | | | |
| 48 | 72% | 48% | | |
| 56 | 85% | | | |
| 64 | 98% | 61% | 42% | 40% |
| 72 | 92% | | | |
| 80 | 97% | 62% | | |
| 88 | 100% | | | |
| 96 | 98% | 72% | 45% | |
| 104 | 100% | | | |
| 112 | 100% | 75% | | |
| 120 | 100% | | | |
| 128 | 100% | 90% | 54% | 40% |
| 136 | 100% | | | |
| 144 | 100% | 94% | | |
| 152 | 100% | | | |
| 160 | 100% | 92% | 66% | |
| 168 | 100% | | | |
| 176 | 100% | 95% | | |
| 184 | 100% | | | |
| 192 | 100% | 94% | 63% | 51% |
| 200 | 100% | | | |
| 208 | 100% | 98% | | |
| 216 | 100% | | | |
| 224 | 98% | 95% | 79% | |
| 232 | 100% | | | |
| 240 | 100% | 97% | | |
| 248 | 100% | | | |
| 256 | 100% | 95% | 88% | 56% |
eval f1
eval f1 | step | 8 groups/step | 16 groups/step | 32 groups/step | 64 groups/step |
|---|---|---|---|---|
| 0 | 0.19 | 0.19 | 0.19 | 0.19 |
| 64 | 0.20 | 0.18 | 0.24 | 0.17 |
| 128 | 0.13 | 0.20 | 0.20 | 0.21 |
| 192 | 0.12 | 0.11 | 0.20 | 0.29 |
| 256 | 0.14 | 0.18 | 0.19 | 0.24 |
At a glance, the eval results are in line with expectations: at the end of training, 64 groups/step has the highest eval F1, and each smaller batch size monotonically decreases in final eval F1. Larger batches also give smoother reward and eval curves.
There’s an interesting contradiction in the charts, though. Although 64 groups/step ends with the highest eval score, it has the lowest train reward at the end of training. Also, 8 groups/step has the highest train reward, but the lowest F1 score.
This has to do with our reward setup and training dynamics — the reward on a rollout that curates nothing is -0.2, which far outweighs the average reward on most valid rollouts. So, the fastest way to raise train reward early on is to just curate any document. Because 8 groups/step takes more updates than any other run, the model starts to produce valid rollouts sooner. The middle chart above shows this: as batch size decreases, the model reaches >75% valid train rollouts faster. The problem from there is that noisy updates (due to small batch size) keep the policy’s curation quality from ever meaningfully improving. What we’re seeing here is actually a form of reward hacking. The next section will explore this more.
Since 64 groups/step gives the most stable training and has the highest eval score, we will use it for our full training runs. We could technically increase groups/step further, but for the sake of budget we will stop here.
Scaling up training
For our next set of training runs, we will move to the larger dataset of 1,024 queries and 237,533 chunks. This is 4x the data samples and 2x the corpus size of the dataset we used for ablations. We train for 1.5 epochs, which at 64 groups/step is 24 total updates.
Full training configuration
| model |
| model | gpt-oss-20b |
| adapter | LoRA, rank 32 |
| optimisation |
| learning rate | 1e-4 |
| optimizer | Adam, β₁ 0.9, β₂ 0.95, ε 1e-8 |
| weight decay | 0 |
| KL penalty to reference | 0.0 |
| updates per batch | 1 |
| sampling temperature | 1.0 |
| harness |
| max turns | 40 |
| max new tokens per turn | 2,048 |
| context length limit | 30,720 tokens |
| curated set cap | 30 |
We first look at training metrics over the full run. In the spirit of openness, every training rollout can be viewed in the explorer below.
train reward
train reward | step | train reward |
|---|---|
| 0 | -0.06 |
| 1 | 0.02 |
| 2 | 0.01 |
| 3 | 0.03 |
| 4 | 0.04 |
| 5 | 0.10 |
| 6 | 0.17 |
| 7 | 0.15 |
| 8 | 0.14 |
| 9 | 0.15 |
| 10 | 0.13 |
| 11 | 0.17 |
| 12 | 0.15 |
| 13 | 0.16 |
| 14 | 0.18 |
| 15 | 0.21 |
| 16 | 0.25 |
| 17 | 0.17 |
| 18 | 0.22 |
| 19 | 0.23 |
| 20 | 0.21 |
| 21 | 0.27 |
| 22 | 0.24 |
| 23 | 0.21 |
valid rollouts
valid rollouts | step | valid rollouts |
|---|---|
| 0 | 30% |
| 1 | 42% |
| 2 | 47% |
| 3 | 64% |
| 4 | 79% |
| 5 | 92% |
| 6 | 96% |
| 7 | 98% |
| 8 | 98% |
| 9 | 98% |
| 10 | 98% |
| 11 | 97% |
| 12 | 94% |
| 13 | 95% |
| 14 | 94% |
| 15 | 90% |
| 16 | 92% |
| 17 | 92% |
| 18 | 95% |
| 19 | 96% |
| 20 | 96% |
| 21 | 94% |
| 22 | 94% |
| 23 | 96% |
entropy (nats)
entropy (nats) | step | entropy (nats) |
|---|---|
| 0 | 1.1 |
| 1 | 1.0 |
| 2 | 1.0 |
| 3 | 1.0 |
| 4 | 1.0 |
| 5 | 0.9 |
| 6 | 0.8 |
| 7 | 0.8 |
| 8 | 0.8 |
| 9 | 0.8 |
| 10 | 0.8 |
| 11 | 0.8 |
| 12 | 0.8 |
| 13 | 0.8 |
| 14 | 0.9 |
| 15 | 0.9 |
| 16 | 0.9 |
| 17 | 0.9 |
| 18 | 0.9 |
| 19 | 0.9 |
| 20 | 0.9 |
| 21 | 0.9 |
| 22 | 0.9 |
| 23 | 0.8 |
Train reward moves the way we’d expect it to here: up rapidly at first, and then slowing down as training matures. The training shape also mirrors what we see in ablations: the reward first teaches the model to curate. After that, reward moves around as the model starts to learn how to curate well instead of just how to curate. You can think of the -0.2 floor on no curations as jump-starting curation behavior.
Entropy is another useful metric to watch during reinforcement learning. Entropy is measured in nats here. One nat is 1 / ln 2, or about 1.44 bits. It measures how uncertain the policy’s output distribution is over the course of a rollout. High entropy means at several points in the rollout, the model is considering several possible actions (next tokens). As we’ve discussed earlier, this is good for RL because it means the model will explore more and hopefully stumble into behaviors that get rewarded. Low entropy isn’t inherently bad, but a rapid drop during training is a sign the model is becoming overconfident or committing to a single set of behaviors. This is known as entropy collapse and is particularly bad for GRPO — a policy that always takes the same actions produces groups with no variance, and groups with no variance produce no gradient.
During healthy training, entropy typically decreases gradually as reward improves. That is exactly what we see here: entropy starts at 1.1 nats, gradually goes to ~0.8 a quarter of the way through training, and then stays there.
Let’s see how our policy performs on held-out data now.
eval f1
eval f1 | step | eval f1 |
|---|---|
| 0 | 0.19 |
| 2 | 0.22 |
| 4 | 0.17 |
| 6 | 0.19 |
| 8 | 0.21 |
| 10 | 0.17 |
| 12 | 0.16 |
| 14 | 0.24 |
| 16 | 0.23 |
| 18 | 0.29 |
| 20 | 0.31 |
| 22 | 0.28 |
| 24 | 0.33 |
eval precision
eval precision | step | eval precision |
|---|---|
| 0 | 0.30 |
| 2 | 0.38 |
| 4 | 0.30 |
| 6 | 0.29 |
| 8 | 0.32 |
| 10 | 0.27 |
| 12 | 0.17 |
| 14 | 0.35 |
| 16 | 0.30 |
| 18 | 0.36 |
| 20 | 0.40 |
| 22 | 0.36 |
| 24 | 0.40 |
eval recall
eval recall | step | eval recall |
|---|---|
| 0 | 0.16 |
| 2 | 0.17 |
| 4 | 0.13 |
| 6 | 0.16 |
| 8 | 0.17 |
| 10 | 0.14 |
| 12 | 0.18 |
| 14 | 0.21 |
| 16 | 0.20 |
| 18 | 0.27 |
| 20 | 0.28 |
| 22 | 0.26 |
| 24 | 0.32 |
Progress is uneven early on as the model learns our harness, and eval performance only starts to improve around halfway through training. Mean eval F1 rises from 0.193 at the start to 0.332 at update 24, driven almost entirely by recall, which doubles from 0.16 to 0.32. Precision goes up as well, moving from 0.31 to 0.40. This is a good result because it means we are getting an all around improvement. That is, we aren’t trading off precision for recall in any way.
The model performance doesn’t come at the cost of test-time compute, either. As the following charts show, the turn count actually falls slightly over training and output tokens go down, which means our final model is finding more relevant docs with less searching.
turns per episode
turns per episode | step | turns per episode |
|---|---|
| 0 | 18 |
| 1 | 18 |
| 2 | 18 |
| 3 | 16 |
| 4 | 15 |
| 5 | 13 |
| 6 | 12 |
| 7 | 12 |
| 8 | 12 |
| 9 | 12 |
| 10 | 12 |
| 11 | 12 |
| 12 | 14 |
| 13 | 14 |
| 14 | 16 |
| 15 | 15 |
| 16 | 16 |
| 17 | 16 |
| 18 | 15 |
| 19 | 15 |
| 20 | 14 |
| 21 | 14 |
| 22 | 16 |
| 23 | 17 |
output tokens per episode
output tokens per episode | step | output tokens per episode |
|---|---|
| 0 | 3.998k |
| 1 | 3.47k |
| 2 | 3.05k |
| 3 | 2.621k |
| 4 | 2.202k |
| 5 | 1.754k |
| 6 | 1.408k |
| 7 | 1.438k |
| 8 | 1.414k |
| 9 | 1.497k |
| 10 | 1.698k |
| 11 | 1.709k |
| 12 | 1.972k |
| 13 | 1.954k |
| 14 | 2.328k |
| 15 | 2.461k |
| 16 | 2.527k |
| 17 | 2.898k |
| 18 | 2.693k |
| 19 | 2.654k |
| 20 | 2.673k |
| 21 | 2.649k |
| 22 | 2.825k |
| 23 | 2.74k |
So, we have a full training run that works and pushes eval scores noticeably higher. This is a great result! In all likelihood, we haven’t saturated our training setup, either—train reward is still steadily going up around update 24. We’d probably continue to see eval gains by scaling up the dataset or running for another epoch…
But that’s a boring next step to explore here. There are many other fun ways to squeeze out performance from our dataset.
Aside: What happened with 8 groups/step?
In our ablations, we saw that the 8 groups/step run quickly learned to curate, but this failed to translate over meaningfully into eval improvements. We extended that run over the larger dataset to see if more training steps would help. The charts below show the results of training over 8 groups/step vs 64 groups/step over the same first 384 queries of the larger dataset.
train reward
train reward | step | 64 groups/step | 8 groups/step |
|---|---|---|
| 0 | -0.06 | 0.06 |
| 8 | | 0.03 |
| 16 | | 0.06 |
| 24 | | 0.03 |
| 32 | | 0.13 |
| 40 | | 0.07 |
| 48 | | 0.23 |
| 56 | | 0.16 |
| 64 | 0.02 | 0.15 |
| 72 | | 0.21 |
| 80 | | 0.30 |
| 88 | | 0.18 |
| 96 | | 0.27 |
| 104 | | 0.12 |
| 112 | | 0.09 |
| 120 | | 0.20 |
| 128 | 0.01 | 0.11 |
| 136 | | 0.12 |
| 144 | | 0.13 |
| 152 | | 0.06 |
| 160 | | 0.10 |
| 168 | | 0.15 |
| 176 | | 0.08 |
| 184 | | 0.17 |
| 192 | 0.03 | 0.09 |
| 200 | | 0.14 |
| 208 | | 0.09 |
| 216 | | 0.13 |
| 224 | | 0.09 |
| 232 | | 0.09 |
| 240 | | 0.21 |
| 248 | | 0.30 |
| 256 | 0.04 | 0.10 |
| 264 | | 0.18 |
| 272 | | 0.23 |
| 280 | | 0.03 |
| 288 | | 0.10 |
| 296 | | 0.14 |
| 304 | | 0.25 |
| 312 | | 0.19 |
| 320 | 0.10 | 0.19 |
| 328 | | 0.25 |
| 336 | | 0.36 |
| 344 | | 0.13 |
| 352 | | 0.15 |
| 360 | | 0.08 |
| 368 | | 0.21 |
| 376 | | 0.16 |
| 384 | 0.17 | |
curated documents
curated documents | step | 64 groups/step | 8 groups/step |
|---|---|---|
| 0 | 1 | 1 |
| 8 | | 1 |
| 16 | | 1 |
| 24 | | 1 |
| 32 | | 1 |
| 40 | | 1 |
| 48 | | 2 |
| 56 | | 3 |
| 64 | 1 | 2 |
| 72 | | 3 |
| 80 | | 4 |
| 88 | | 4 |
| 96 | | 5 |
| 104 | | 5 |
| 112 | | 6 |
| 120 | | 7 |
| 128 | 1 | 6 |
| 136 | | 9 |
| 144 | | 8 |
| 152 | | 7 |
| 160 | | 7 |
| 168 | | 7 |
| 176 | | 7 |
| 184 | | 7 |
| 192 | 1 | 9 |
| 200 | | 12 |
| 208 | | 9 |
| 216 | | 9 |
| 224 | | 11 |
| 232 | | 11 |
| 240 | | 9 |
| 248 | | 13 |
| 256 | 1 | 18 |
| 264 | | 13 |
| 272 | | 12 |
| 280 | | 17 |
| 288 | | 16 |
| 296 | | 17 |
| 304 | | 21 |
| 312 | | 21 |
| 320 | 2 | 24 |
| 328 | | 22 |
| 336 | | 23 |
| 344 | | 24 |
| 352 | | 24 |
| 360 | | 24 |
| 368 | | 25 |
| 376 | | 25 |
| 384 | 2 | |
entropy (nats)
entropy (nats) | step | 64 groups/step | 8 groups/step |
|---|---|---|
| 0 | 1.1 | 1.1 |
| 8 | | 1.1 |
| 16 | | 1.0 |
| 24 | | 0.9 |
| 32 | | 0.9 |
| 40 | | 0.9 |
| 48 | | 0.8 |
| 56 | | 0.7 |
| 64 | 1.0 | 0.8 |
| 72 | | 0.7 |
| 80 | | 0.7 |
| 88 | | 0.7 |
| 96 | | 0.7 |
| 104 | | 0.7 |
| 112 | | 0.7 |
| 120 | | 0.7 |
| 128 | 1.0 | 0.8 |
| 136 | | 0.7 |
| 144 | | 0.8 |
| 152 | | 0.7 |
| 160 | | 0.7 |
| 168 | | 0.7 |
| 176 | | 0.7 |
| 184 | | 0.8 |
| 192 | 1.0 | 0.7 |
| 200 | | 0.8 |
| 208 | | 0.7 |
| 216 | | 0.7 |
| 224 | | 0.7 |
| 232 | | 0.6 |
| 240 | | 0.6 |
| 248 | | 0.7 |
| 256 | 1.0 | 0.5 |
| 264 | | 0.6 |
| 272 | | 0.6 |
| 280 | | 0.5 |
| 288 | | 0.5 |
| 296 | | 0.5 |
| 304 | | 0.4 |
| 312 | | 0.5 |
| 320 | 0.9 | 0.4 |
| 328 | | 0.4 |
| 336 | | 0.4 |
| 344 | | 0.4 |
| 352 | | 0.3 |
| 360 | | 0.3 |
| 368 | | 0.3 |
| 376 | | 0.3 |
| 384 | 0.8 | |
eval f1
eval f1 | step | 64 groups/step | 8 groups/step |
|---|---|---|
| 0 | 0.19 | 0.19 |
| 128 | 0.25 | 0.12 |
| 256 | 0.15 | 0.16 |
| 384 | 0.17 | 0.16 |
eval precision
eval precision | step | 64 groups/step | 8 groups/step |
|---|---|---|
| 0 | 0.30 | 0.30 |
| 128 | 0.38 | 0.11 |
| 256 | 0.30 | 0.14 |
| 384 | 0.29 | 0.13 |
eval recall
eval recall | step | 64 groups/step | 8 groups/step |
|---|---|---|
| 0 | 0.16 | 0.16 |
| 128 | 0.17 | 0.14 |
| 256 | 0.13 | 0.23 |
| 384 | 0.16 | 0.29 |
The 8 groups/step run reaches higher training reward, but it gets there by learning to just curate everything — although eval recall rises from 0.16 to 0.29, precision drops from 0.31 to 0.13. Entropy also falls from 1.1 to 0.34 nats over the course of training. This is a clear sign of entropy collapse. It’s likely that learning will outright stop if we were to continue training.
Part of this is our reward choice. We chose β = 4 to push the model towards recall over precision, and F4 weights recall 16x. Small-batch updates are also noisier, so once the policy stumbles onto the “curate every document you see” strategy, there is little pushing the model back to something more reasonable.
One way to mitigate the reward bias here would be to anneal β over the course of training, as is done in Chroma Context-1.
Reward shaping
Every run so far has used the plain F4 score as the reward. It worked, in the sense that eval scores went up. But reading the rollouts directly and exploring a different set of training metrics shows two other ways we can improve the policy.
trajectory vs curated recall
trajectory vs curated recall | step | trajectory recall | curated recall |
|---|---|---|
| 0 | 0.46 | 0.12 |
| 1 | 0.51 | 0.16 |
| 2 | 0.40 | 0.13 |
| 3 | 0.34 | 0.11 |
| 4 | 0.29 | 0.09 |
| 5 | 0.32 | 0.14 |
| 6 | 0.39 | 0.22 |
| 7 | 0.38 | 0.24 |
| 8 | 0.36 | 0.23 |
| 9 | 0.33 | 0.21 |
| 10 | 0.26 | 0.19 |
| 11 | 0.33 | 0.24 |
| 12 | 0.29 | 0.21 |
| 13 | 0.30 | 0.21 |
| 14 | 0.27 | 0.18 |
| 15 | 0.38 | 0.26 |
| 16 | 0.40 | 0.28 |
| 17 | 0.30 | 0.20 |
| 18 | 0.39 | 0.28 |
| 19 | 0.37 | 0.28 |
| 20 | 0.36 | 0.27 |
| 21 | 0.41 | 0.32 |
| 22 | 0.39 | 0.29 |
| 23 | 0.36 | 0.27 |
Trajectory recall measures how many relevant chunks the model encountered while searching. tool-call format drift
tool-call format drift | step | of calls | of episodes |
|---|---|---|
| 0 | 27.3% | 37.9% |
| 1 | 13.7% | 29.3% |
| 2 | 15% | 26.4% |
| 3 | 9.1% | 16.2% |
| 4 | 11.5% | 18.9% |
| 5 | 12.4% | 18.8% |
| 6 | 14.8% | 17.8% |
| 7 | 25% | 32.2% |
| 8 | 21.2% | 32.2% |
| 9 | 19.7% | 32.4% |
| 10 | 14.6% | 33.6% |
| 11 | 26.8% | 41.6% |
| 12 | 28.3% | 55.9% |
| 13 | 47.2% | 67.2% |
| 14 | 49.6% | 74.2% |
| 15 | 62.6% | 79.9% |
| 16 | 53.8% | 78.3% |
| 17 | 58.5% | 89.3% |
| 18 | 67.3% | 90.2% |
| 19 | 67.9% | 92.6% |
| 20 | 63% | 92.4% |
| 21 | 62.6% | 91% |
| 22 | 67.7% | 96.7% |
| 23 | 75% | 98.8% |
Measures % of tool calls or episodes with incorrect harmony format.
- Curation recall is bounded by trajectory recall. The model can only curate what it sees as it searches, and in the latter half of training, curation recall consistently sits 0.1 lower than trajectory recall. Trajectory recall also drifts down over training, indicating that the model is getting better at picking relevant docs from searches, but slightly worse at actually searching.
- Tool-call format adherence degrades over training. As training progresses, the share of incorrectly formatted tool calls nearly triples. Luckily, our parser is lenient, so this doesn’t seem to impact training performance too much.
These observations suggest two ways we can improve the model further: reward the model for encountering more evidence during search, and punish it for bad tool calls.
Rewarding discovery
The trajectory recall chart tells us that the model either isn’t searching well or isn’t searching enough. There are two ways we can influence this behavior.
First, we can reward the behavior directly: a bonus per search or grep call, or for more varied queries, and hope that together with the F4 reward, this leads to better coverage. However, a bonus just for searching is easy to hack. The model could just spam irrelevant calls and collect the bonus without learning any performance-improving behavior. We’d need to add a turn penalty to hold it back. This adds a second knob to tune and complicates the reward.
Another route is to add an explicit trajectory recall term to our reward. This is a better route because we’re rewarding the outcome we want rather than the activity it takes to get there. As our charts show, trajectory recall is directionally aligned with curated recall, and it’s also harder to game — extra searches don’t earn anything unless they are helpful. This is the route we’ll experiment with.
Concretely, we add a term for the fraction of the query’s facts that have a supporting chunk anywhere in the trajectory, curated or not:
r(C)={−0.2F4+0.2⋅Rtrajif ∣C∣=0otherwise
The explorer below shows the impact of different reward shapes on a group. Seven of the eight rollouts score zero F4 and are indistinguishable under the older reward, but separable under either bonus.
Apart from the reward, every setting in this run matches the full run above.
eval f1
eval f1 | step | w/o trajectory reward | w trajectory reward |
|---|---|---|
| 0 | 0.19 | 0.19 |
| 2 | 0.22 | 0.21 |
| 4 | 0.17 | 0.21 |
| 6 | 0.19 | 0.18 |
| 8 | 0.21 | 0.18 |
| 10 | 0.17 | 0.27 |
| 12 | 0.16 | 0.27 |
| 14 | 0.24 | 0.25 |
| 16 | 0.23 | 0.28 |
| 18 | 0.29 | 0.27 |
| 20 | 0.31 | 0.32 |
| 22 | 0.28 | 0.36 |
| 24 | 0.33 | 0.33 |
eval curated recall
eval curated recall | step | w/o trajectory reward | w trajectory reward |
|---|---|---|
| 0 | 0.16 | 0.16 |
| 2 | 0.17 | 0.15 |
| 4 | 0.13 | 0.16 |
| 6 | 0.16 | 0.14 |
| 8 | 0.17 | 0.14 |
| 10 | 0.14 | 0.22 |
| 12 | 0.18 | 0.23 |
| 14 | 0.21 | 0.20 |
| 16 | 0.20 | 0.24 |
| 18 | 0.27 | 0.24 |
| 20 | 0.28 | 0.29 |
| 22 | 0.26 | 0.34 |
| 24 | 0.32 | 0.30 |
eval trajectory recall
eval trajectory recall | step | w/o trajectory reward | w trajectory reward |
|---|---|---|
| 0 | 0.43 | 0.43 |
| 2 | 0.35 | 0.37 |
| 4 | 0.30 | 0.30 |
| 6 | 0.29 | 0.26 |
| 8 | 0.28 | 0.26 |
| 10 | 0.22 | 0.33 |
| 12 | 0.22 | 0.33 |
| 14 | 0.28 | 0.36 |
| 16 | 0.30 | 0.38 |
| 18 | 0.36 | 0.47 |
| 20 | 0.37 | 0.48 |
| 22 | 0.33 | 0.53 |
| 24 | 0.44 | 0.49 |
Eval is run 2x, and plotted points are the averages. Reward plots omitted because they are noncomparable.
Trajectory recall is clearly better here: after a dip around step 8, the shaped run recovers sooner and stays above the old, plain-F4 run. Eval F1 and curated recall also trend above the old run and even peak higher. Unfortunately, this doesn’t translate directly to a better final score under the same training budget.
Harmony, the format gpt-oss-20b uses, has a specific header for tool calls. At the start of training, the model gets it wrong on about a quarter of calls.
Because we use a lenient parse that accepts these slightly off-form calls, the model still gets rewarded whenever the episode they’re in happens to get a good score. This habit ends up compounding. By the final step of training, 99% of episodes contain at least one bad tool call. The model isn’t being told, in any way that reaches the gradient, that those calls are wrong, so it ends up learning that those calls are acceptable.
The following experiment fixes that directly. Every setting is the same as our full, plain F4 run, except for a slight change to the reward: subtract 0.1 if any tool call in the episode has an off-form header. If two rollouts in a group reach the same F4 but one uses a bad header, the clean one now gets more advantage.
f1
f1 | step | w/o format penalty | w format penalty |
|---|---|---|
| 0 | 0.19 | 0.19 |
| 2 | 0.22 | 0.20 |
| 4 | 0.17 | 0.27 |
| 6 | 0.19 | 0.21 |
| 8 | 0.21 | 0.24 |
| 10 | 0.17 | 0.20 |
| 12 | 0.16 | 0.25 |
| 14 | 0.24 | 0.27 |
| 16 | 0.23 | 0.34 |
| 18 | 0.29 | 0.33 |
| 20 | 0.31 | 0.39 |
| 22 | 0.28 | 0.36 |
| 24 | 0.33 | 0.36 |
recall
recall | step | w/o format penalty | w format penalty |
|---|---|---|
| 0 | 0.16 | 0.16 |
| 2 | 0.17 | 0.16 |
| 4 | 0.13 | 0.20 |
| 6 | 0.16 | 0.17 |
| 8 | 0.17 | 0.20 |
| 10 | 0.14 | 0.20 |
| 12 | 0.18 | 0.26 |
| 14 | 0.21 | 0.26 |
| 16 | 0.20 | 0.36 |
| 18 | 0.27 | 0.35 |
| 20 | 0.28 | 0.43 |
| 22 | 0.26 | 0.42 |
| 24 | 0.32 | 0.43 |
off-form tool headers
off-form tool headers | step | w/o format penalty (episodes) | w format penalty (episodes) | w/o format penalty (calls) | w format penalty (calls) |
|---|---|---|---|---|
| 0 | 37.9% | 39.8% | 27.3% | 27.8% |
| 1 | 29.3% | 23.4% | 13.7% | 10.3% |
| 2 | 26.4% | 21.7% | 15% | 6.1% |
| 3 | 16.2% | 9.2% | 9.1% | 2.7% |
| 4 | 18.9% | 6.8% | 11.5% | 2.3% |
| 5 | 18.8% | 3.9% | 12.4% | 0.7% |
| 6 | 17.8% | 3.5% | 14.8% | 1% |
| 7 | 32.2% | 7% | 25% | 2.1% |
| 8 | 32.2% | 4.9% | 21.2% | 0.5% |
| 9 | 32.4% | 9.6% | 19.7% | 2.8% |
| 10 | 33.6% | 12.7% | 14.6% | 3.6% |
| 11 | 41.6% | 8.6% | 26.8% | 1.1% |
| 12 | 55.9% | 8% | 28.3% | 0.8% |
| 13 | 67.2% | 7.2% | 47.2% | 0.6% |
| 14 | 74.2% | 8.2% | 49.6% | 0.5% |
| 15 | 79.9% | 8.6% | 62.6% | 0.4% |
| 16 | 78.3% | 5.9% | 53.8% | 0.4% |
| 17 | 89.3% | 5.1% | 58.5% | 0.2% |
| 18 | 90.2% | 4.9% | 67.3% | 0.2% |
| 19 | 92.6% | 3.1% | 67.9% | 0.1% |
| 20 | 92.4% | 3.1% | 63% | 0.1% |
| 21 | 91% | 4.3% | 62.6% | 0.2% |
| 22 | 96.7% | 4.7% | 67.7% | 0.2% |
| 23 | 98.8% | 8.2% | 75% | 0.4% |
turns per episode
turns per episode | step | w/o format penalty | w format penalty |
|---|---|---|
| 0 | 18 | 18 |
| 1 | 18 | 17 |
| 2 | 18 | 17 |
| 3 | 16 | 13 |
| 4 | 15 | 12 |
| 5 | 13 | 12 |
| 6 | 12 | 12 |
| 7 | 12 | 13 |
| 8 | 12 | 15 |
| 9 | 12 | 14 |
| 10 | 12 | 15 |
| 11 | 12 | 16 |
| 12 | 14 | 20 |
| 13 | 14 | 20 |
| 14 | 16 | 21 |
| 15 | 15 | 23 |
| 16 | 16 | 23 |
| 17 | 16 | 25 |
| 18 | 15 | 26 |
| 19 | 15 | 27 |
| 20 | 14 | 27 |
| 21 | 14 | 28 |
| 22 | 16 | 30 |
| 23 | 17 | 30 |
In the tool headers chart, solid lines show the share of episodes with off-form calls; dashed lines show the share of individual calls.
The extra penalty works, fast. Off-form tool calls fall below 1% within a few steps and stay there for the rest of training. Eval metrics and recall also improve against the baseline. Held-out F1 ends at 0.36 vs 0.33, and recall at 0.43 vs 0.32.
This result surprised me. I expected the penalty to fix tool call formatting, but not necessarily retrieval quality. Yet the model starts taking longer trajectories (does more searches) and achieves higher recall. I don’t have a concrete explanation for this. One possibility is that the penalty changes which search behaviors get reinforced, and well-formed tool calls are happen to be correlated with better search behaviors. Another is that the model just has more coherent long-context understanding when all of its preceding tool calls are well-formed.
Either hypothesis would require more training runs to test, but I’ll stop here for the sake of my training budget. What we should take away from this is that a seemingly narrow reward change can sometimes lead to big shifts in the trained policy. We need to be careful and intentional when working on reward functions.
Some closing thoughts
We started with a 20B model that struggled to properly execute over our harness, then ran ablations and experiments to drive its eval F1 score from 0.19 up to a peak of 0.39.
There’s still a lot more we could do here to improve performance. For one, we only trained on a fraction of the total SEC Harness-1 dataset. We could expand the training set or test more learning-rate and batch-size combinations. Tuning one hyperparameter at a time was the right call under our specific budget, but we could have carried over a few configurations to the full run.
A few other, more interesting experiment directions
- Curriculum learning During reinforcement learning, we typically want to give the model tasks that are on the frontier of the model’s capabilities. One way to do this is by getting an empirical pass rate of each task before training. We can use this pass rate to bucket tasks into different difficulty tiers and then sample with different weights over the course of training.
- SFT Warmup In our scaled up training runs, we end up devoting about a third of our data budget just to teaching the model how to use our harness. We can speed this up by adding an initial warmup stage with SFT over some easier tasks.
- Harness Engineering We generally don’t want to be too prescriptive in the system prompt because we want the model to discover behaviors during training. However, that the model initially does so poorly is a sign that we can improve either our system prompt or harness signals a bit.
- Other reward shapes So far, we’ve kept our reward functions deliberately simple. There’s a lot more we can do here. For example: a turn penalty to encourage shorter rollouts, rewards for tool diversity, etc.