Training Search Agents with GRPO

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 setting for learning RL because there are many different factors we can play around with that affect model performance. It is also a domain where, in my experience, a well-design 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 osiur 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 choose this setting because it works well for a fixed turn-budget setting: since the set exists from the very start, an episode that runs out of turns can still return relevant docs and be graded during training.

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.

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:

1. Take a query and run the agent on it N times using the current “policy” (in this context, essentially the model’s weights). Each run is called a “rollout”, the N tries together are a “group.”
2. Score every rollout. This is known as the reward.
3. Subtract the group’s mean reward from each rollout’s reward and assign the resulting difference as the rollout’s “advantage.” In vanilla GRPO, we would also divide each advantage by the group’s standard deviation. Dr. GRPO drops this step for training stability: 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.
4. Update the policy (model weights) so that every token in a rollout gets more likely in proportion to its advantage.

We would also typically use ratio clipping during the policy update to stabilize training further, but to simplify training further for the sake of pedagogy, we will train in a fully synchronous setting with zero training-inference mismatch (Tinker provides this by default) and only do one update per batch. We won’t go much deeper into the math behind GRPO in this blog post, but there are many great other resources out there for understanding GRPO.

The reward

Each task here consists of a query, a corpus, a set of facts, 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 trivial maximized by curating every document in the corpus. This is useless for an end user! Why use a search agent at all? High precision the model can be achieved just as trivially by curating only a single correct chunk per query, which is useless in the other duration.

We need a metric that punishes both forms of degenerate behavior. F1, the harmonic mean of precision and recall, is the usual choice here:

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:

We set β = 4 for our reward function, which weighs recall sixteen times more than precision. This number is chosen based on prior experience and notes from the Harness-1 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

Initial explorations

The core idea behind reinforcement learning is to allow models to explore and then reward it for finding good paths. This means that RL can only reinforce what the model already exhibits at least occasionally. 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. For each task, we sample a group of rollouts for each query, score them, and then push the policy towards the ones that score above the group’s mean. If all rollouts in a group fail, then there is nothing to push the model toward. The group is effectively dropped and the query contributes no gradient. The effect of this is that RL teaches by sharpening behaviors the base policy can already produce; it does not discover ones it can’t.

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.264 vs 0.120 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 nine 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 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:

We can see in here the effects of both a too large and a too small learning rate:

• 1e-5 and 5e-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 and 5e-5 train steadily, with 1e-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 the 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.

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 group/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 235,533 training samples. 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.

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.