Search-R1

Code source and cost

• The complete implementation comes from KMnO4-zx/llm-agent-rl-lab/03-search-r1.
• The Evaluation and Training sessions shown in the source repository cost ¥13.30 on PyTRIO in total. This is a measured reference run, not a fixed price; actual usage depends on the number of steps, trajectory lengths, search calls, and evaluation size.
• See the SwanLab run for training curves. This is a multi-file project, so this page explains the main code and complete training logic while the source repository remains the runnable reference.

Introduction

Search-R1 trains a model to decide:

• when to search;
• what query to issue;
• how to use search results in later reasoning;
• when to stop searching and produce the final answer.

Unlike standard RAG, which retrieves once before generation, a Search-R1 trajectory can repeatedly alternate between model actions and environment observations:

assistant reasoning
→ search(query)
→ tool observation
→ more assistant reasoning
→ ...
→ Answer: <short answer>

This reproduction uses Qwen/Qwen3.5-4B, PyTRIO, and the Zhihu global-search API. It preserves the multi-turn tool environment, outcome reward, group-relative advantage, observation-token masking, and policy update while replacing the original local Wikipedia retrieval infrastructure with an online search service.

The key boundary is:

Only the model's LoRA weights are trained. The search API and its returned evidence are a fixed environment.

The model learns a tool-use and answer-generation policy; it does not train the search engine.

Configuration

Project layout

Search-R1 is not a single-file example:

03-search-r1/
├── prepare_data.py   # Download and prepare train/evaluation data
├── data.py           # Read local JSONL files
├── protocol.py       # Tool schema, prompt, and tool-call parsing
├── search.py         # Zhihu search client and call statistics
├── rollout.py        # Multi-turn tool-use state machine
├── reward.py         # Exact Match + Format reward
├── train.py          # PyTRIO training, packing, updates, checkpoints
├── eval.py           # Shared evaluation for base/checkpoint models
└── analyse.py        # Aggregate checkpoint evaluation results

Use the complete source directory to run the example. The following sections focus on the code that defines the algorithm.

Environment and data

Only a networked CPU machine is needed locally. Sampling, LoRA forward/backward, and optimizer updates run on the remote PyTRIO service.

git clone https://github.com/KMnO4-zx/llm-agent-rl-lab.git
cd llm-agent-rl-lab

uv sync
trio login
swanlab login

The data originates from PeterJinGo/nq_hotpotqa_train. The preparation script downloads a pinned ModelScope mirror:

uv run python 03-search-r1/prepare_data.py

It creates:

03-search-r1/datasets/
├── train.jsonl   # 169,615 NQ + HotpotQA training questions
├── dev.jsonl     # 10 questions from each of 7 benchmarks
└── test.jsonl    # Full evaluation pool

Each training record contains the question and reference answers, but no annotated query or expert search trajectory:

{
  "id": "...",
  "question": "...",
  "answers": ["..."],
  "data_source": "nq"
}

After obtaining API keys from the Zhihu Data Platform, copy the environment template:

cp 03-search-r1/.env.example 03-search-r1/.env

Add one or more keys:

ZHIHU_SEARCH_KEYS=your_first_key,your_second_key,your_third_key

Before training, verify that the keys have sufficient quota. During training, monitor search/success_rate and search/error_rate; an unstable tool environment directly contaminates reward.

Core logic

1. Declare search as a model tool

protocol.py defines the tool schema:

SEARCH_TOOL = {
    "type": "function",
    "function": {
        "name": "search",
        "description": "Search Zhihu for evidence. Use a concise English query.",
        "parameters": {
            "type": "object",
            "properties": {"query": {"type": "string"}},
            "required": ["query"],
        },
    },
}

prompt_tokens = tokenizer.apply_chat_template(
    messages,
    tools=[SEARCH_TOOL],
    tokenize=True,
    add_generation_prompt=True,
    enable_thinking=False,
)

Qwen3.5 emits a structured <tool_call>. protocol.py parses its query, search.py performs the real search, and the environment appends titles, passages, sources, and URLs as a role="tool" observation.

Tool intent is parsed from the structured call rather than inferred from a stop word. Stop sequences only terminate the current generation.

2. Advance multi-turn trajectories after they diverge

All eight trajectories for one question share the first prompt, so the first request can use:

sample_async(
    prompt=shared_prompt,
    num_samples=8,
)

After the first search, each trajectory has a different query and observation:

First turn:
1 shared prompt × num_samples=8

After the first search:
8 independent prompts × num_samples=1
multiple sample_async calls run concurrently

Each trajectory remains causally sequential:

assistant generation
→ search(query)
→ tool observation
→ next assistant generation

Only different trajectories run concurrently. The complete state machine is in rollout.py.

3. Score only the final answer

The final assistant turn must end with:

Answer: <short answer>

reward.py uses three outcomes:

def score_answer(text: str, references: list[str]) -> RewardResult:
    answer = extract_answer(text)
    if answer is None:
        return RewardResult(-0.1, False, False, None)
    exact_match = any(
        normalize_answer(answer) == normalize_answer(reference)
        for reference in references
    )
    return RewardResult(float(exact_match), True, exact_match, answer)

The implementation does not reward the number of searches or intermediate query text, which avoids directly encouraging redundant tool calls. The policy explores the search path and receives outcome reward from the final answer.

4. Compute advantage over the complete question group

After all eight trajectories for a question finish:

$A_i = r_i - \operatorname{mean}(r_1, r_2, \ldots, r_8)$

mean_reward = sum(item.reward for item in group) / len(group)
for item in group:
    item.advantage = item.reward - mean_reward

If every reward in a group is identical, all advantages are zero and the group is skipped. The group mean must be computed before micro-batching; recomputing it inside arbitrary micro-batches changes the algorithm.

5. Keep search results in context but out of the loss

A trajectory mixes model actions and environment observations:

system / user / tool observation   → advantage = 0
assistant tool call / final answer → advantage = trajectory_advantage

build_datum() in train.py joins all turns into one continuous sequence. Observation tokens remain in context but receive zero old logprob and zero advantage:

full_tokens.extend(delta_observation)
full_tokens.extend(turn.completion_tokens)

old_logprobs_by_token.extend([0.0] * len(delta_observation))
old_logprobs_by_token.extend(turn.logprobs)

advantages_by_token.extend([0.0] * len(delta_observation))
advantages_by_token.extend(
    [trajectory.advantage] * len(turn.completion_tokens)
)

input_tokens = full_tokens[:-1]
target_tokens = full_tokens[1:]
old_logprobs = old_logprobs_by_token[1:]
advantages = advantages_by_token[1:]

datum = trio.Datum(
    model_input=trio.ModelInput.from_ints(input_tokens),
    loss_fn_inputs={
        "target_tokens": np.asarray(target_tokens, dtype=np.int64),
        "logprobs": np.asarray(old_logprobs, dtype=np.float32),
        "advantages": np.asarray(advantages, dtype=np.float32),
    },
)

All four fields use the same autoregressive shift and must have identical lengths. Non-zero old logprobs must come from the Student sampler used for the rollout, not from a later model state.

6. Finish each group before micro-batch updates

One logical step can contain:

8 questions × 8 trajectories = 64 trajectories

With trajectories of up to 8,192 tokens, they cannot be sent in one request. The implementation therefore follows:

complete rollout groups
→ reward
→ group-relative advantage
→ one Datum per complete trajectory
→ dynamic micro-batch packing
→ accumulated forward/backward calls
→ one optimizer step for the logical batch

The current limits are 8,192 tokens per Datum, 32 Datums per micro-batch, and at most 64,000 padded tokens (items × max_sequence_length) per micro-batch.

trajectories = rollout_batch(...)
datums = build_training_datums(trajectories)
micro_batches = pack_micro_batches(datums)

for micro_batch in micro_batches:
    training_client.forward_backward(
        weight_micro_batch_for_global_mean(
            micro_batch,
            total_samples=len(trajectories),
        ),
        loss_fn="importance_sampling",
    ).result()

if micro_batches:
    training_client.optim_step(adam_params).result()

Because the remote service averages samples within each forward_backward call, the code scales each micro-batch advantage by n_k / N. Accumulated gradients then remain equivalent to the global sample mean of the complete logical batch.

Run training

Start with 20 steps to validate the complete path:

uv run python 03-search-r1/train.py \
    --max-steps 20 \
    --questions-per-batch 8 \
    --group-size 8 \
    --save-every 5 \
    --swanlab-mode online

Every five steps, the script saves:

*-state    # Optimizer state for resuming training
*-weights  # Sampler weights for inference and evaluation

At 20 steps, check whether:

• reward/format starts increasing;
• the model ends multi-turn search with Answer:;
• reward/correct changes;
• search/success_rate remains stable;
• rollout/degenerate_group_rate is not excessive.

Evaluation

Evaluate the base model under the same environment:

uv run python 03-search-r1/eval.py \
    --batch-size 16 \
    --output 03-search-r1/eval_result/eval_results.jsonl

Then pass a saved sampler path to the evaluator:

uv run python 03-search-r1/eval.py \
    --batch-size 16 \
    --model-path 'trio://YOUR_STEP_20_SAMPLER_WEIGHTS' \
    --output 03-search-r1/eval_result/eval_results_rl_step_20.jsonl

When search quota is limited, add --limit 20 for a pipeline check before running the fixed 70-question evaluation.

The reference results on the fixed set were:

Base Model:  Macro EM 28.57% · Format 58.57%
RL Step 20:  Macro EM 31.43% · Format 87.14%
RL Step 50:  Macro EM 45.71% · Format 94.29%

Step 20 came from an earlier small run under a different live-search condition, and 70 questions are better suited to behavior validation than paper-level claims.

Metrics to monitor

If reward/correct falls while search/error_rate rises, inspect the environment before concluding that the policy regressed.

Reproduction boundaries

This example reproduces the core Search-R1 training loop rather than every part of the original infrastructure and reported scores:

• the model, search backend, training framework, and some settings differ;
• the fixed development evaluation contains only 70 questions;
• online search is affected by quotas, timeouts, and changing results;
• the search backend stays fixed while the model's tool-use and answer policy is trained.

For complete code, experiment figures, and a longer discussion, see the source repository's Search-R1 README.