Перейти к содержимому
ПереводыDatawhale и китайские курсы

Happy-LLM, глава 8.3: Search-R1 — агент учится искать

Полный перевод раздела 8.3 книги Datawhale «Строим большую модель с нуля»: обучение модели самой решать, когда и что искать — инструмент search, поисковые бэкенды как среда, многоходовой rollout, непрерывность токенов, маска наблюдений, микробатчи и оценка на NQ и HotpotQA.

Источник кода: раздел составлен по полной многофайловой реализации из открытого репозитория автора agentic-rl-lab/03-search-r1. Текст следует настоящей цепочке вызовов: prepare_data.py, protocol.py, search.py, rollout.py, reward.py, train.py, eval.py и analyse.py.

Большая языковая модель хранит внутри огромный объём знаний, но эти знания ограничены временны́м охватом обучающих данных, ёмкостью параметров и покрытием «длинного хвоста» редких фактов. В генерации с дополнением найденным (Retrieval-Augmented Generation, RAG) система обычно заранее находит документы и передаёт результат модели для ответа. Search-R1 передаёт решения о поиске самой модели: модель может решить, когда искать, что искать, нужно ли искать дальше с учётом новых данных и когда закончить и дать ответ.

Из-за этого один ответ становится многоходовой траекторией:

τ=(x,a1,o1,a2,o2,…,aT),\tau=(x,a_1,o_1,a_2,o_2,\ldots,a_T),

где ata_t — вызов поиска или итоговый ответ, сгенерированные моделью, а oto_t — observation, которое вернула поисковая среда. Награду за результат модель получает только после всей траектории, а GRPO затем отличает успешные пути поиска от неудачных.

Сопутствующая реализация — не отдельный скрипт обучения, а цепочка исполнения из 9 файлов. Читая код, сначала ухватите эту основную линию:

prepare_data.py / data.py
↓ читаем вопросы и эталонные ответы
protocol.py
↓ строим prompt, разбираем вызовы search
search.py
↓ выполняем поиск и возвращаем observation
rollout.py
↓ генерируем полную многоходовую траекторию, reward и групповой advantage
train.py
↓ observation mask, Datum, importance_sampling, optim_step
eval.py / analyse.py
↓ оценка в фиксированной среде и сводка результатов

Обучаются только токены Assistant, сгенерированные языковой моделью. Поисковый сервис, за который отвечает search.py, относится к среде; protocol.py соединяет действия модели и observation среды в непрерывную траекторию; rollout.py управляет обоими; и только в конце train.py превращает траекторию в Datum, который PyTRIO может использовать для обновления. Ниже разберём код по этой настоящей цепочке вызовов.

8.3.1 От фиксированного поиска к самостоятельному

Заголовок раздела «8.3.1 От фиксированного поиска к самостоятельному»

Обучающие примеры Search-R1 содержат только вопрос и эталонный ответ — ни поисковых запросов, написанных людьми, ни образцовых траекторий поиска. prepare_data.py приводит исходные поля NQ и HotpotQA к такой структуре JSONL:

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

Логика очистки оставляет только поля, нужные для награды за результат:

def normalize_row(row: dict[str, Any]) -> dict[str, Any] | None:
question = str(row.get("question") or "").strip()
answers = extract_answers(row)
if not question or not answers:
return None
return {
"id": str(row.get("id") or ""),
"question": question,
"answers": answers,
"data_source": str(
row.get("data_source") or "unknown"
),
}

data.py затем превращает каждую строку в SearchExample:

@dataclass(frozen=True)
class SearchExample:
id: str
question: str
answers: list[str]
data_source: str

Эта граница данных очень важна. Поисковые запросы, число поисков и способ использования найденного текущая политика должна исследовать сама во время rollout; если бы обучающий набор уже давал фиксированные траектории поиска, задача была бы ближе к подражанию траекториям, а не к Search-R1, о котором идёт речь.

Различия фиксированного RAG и поискового агента видны в таблице 8.3.

ЭтапФиксированный RAGSearch-R1
Искать лирешает внешний процессрешает модель по текущему состоянию
Поисковый запросзадан заранее или генерируется один разможно многократно переписывать по observation
Число раундов поискаобычно фиксированорешает политика, с ограничением на максимум
Что обучаетсямодель ответа или поисковик (ретривер)единая политика поисковых действий и итогового ответа
Наградаобычно потеря обучения с учителем на ответенаграда за результат всей траектории

Таблица 8.3. Фиксированный RAG и Search-R1.

Например, вопрос «В какой стране родился автор „Маленького принца“?» требует сначала установить автора, а затем найти его место рождения. Модель может построить такую траекторию:

User: What country was the author of The Little Prince born in?
Assistant -> search("The Little Prince author")
Tool: Antoine de Saint-Exupéry ...
Assistant -> search("Antoine de Saint-Exupéry birthplace country")
Tool: He was born in Lyon, France ...
Assistant: Answer: France

Второй запрос зависит от первого observation, поэтому его нельзя определить заранее, до начала rollout. Код обучения должен хранить состояние среды и разбирать действие модели после каждого хода Assistant.

flowchart LR
Q["Вопрос:<br/>где родился автор<br/>«Маленького принца»?<br/>(общий корень)"] --> S1["Вызов поиска 1:<br/>search «автор Маленького принца»<br/>(действие ассистента)"]
S1 --> O1["Результат 1:<br/>Антуан де Сент-Экзюпери<br/>(наблюдение)"]
O1 --> S2["Вызов поиска 2:<br/>search «место рождения<br/>Сент-Экзюпери»"]
S2 --> O2["Результат 2:<br/>Лион, Франция<br/>(цепочка доказательств замкнута)"]
O2 --> A["Итоговый ответ:<br/>Answer: France"]
Рисунок 8.4 (верхняя часть). Настоящая многоходовая траектория поиска; эталон — France (схема перерисована переводчиком)

Рисунок 8.4. Многоходовый поиск и маска наблюдений в Search-R1 (изображение в оригинале). Подзаголовок: «агент сам решает, что искать, сколько раундов поиска делать и когда остановиться; награду за результат получает после окончания всей траектории». Нижняя левая часть — «награда за траекторию и групповое преимущество»: одна задача, 4 полные траектории поиска с наградами τ0: 1\tau_0{:}\ 1, τ1: 0\tau_1{:}\ 0, τ2: 1\tau_2{:}\ 1, τ3: 0\tau_3{:}\ 0; Ai=ri−mean(rgroup)A_i = r_i - \mathrm{mean}(r_{group}); «в пределах одного бюджета поиска все ходы ассистента делят один advantage траектории». Нижняя правая часть — «маска наблюдений», надзор по токенам одной траектории: вопрос пользователя 00, вызов поиска AiA_i, результат поиска 00, продолжение поиска или ответ AiA_i; «токены действий ассистента участвуют в обучении; токены наблюдений инструмента — только контекст следующего хода; исходные сэмплированные токены и старые logprob должны быть непрерывно выровнены». Внизу: «на одну задачу первый ход общий (корень); получив свои наблюдения, каждая траектория продолжает генерацию из собственного состояния».

В статье Search-R1 действия и наблюдения обозначены текстовыми тегами <search>... </search>, <information>... </information> и т. п. В этой главе используется модель Qwen3.5-4B и её собственный (native) шаблон function calling:

SEARCH_TOOL = {
"type": "function",
"function": {
"name": "search",
"description": "Search the web for evidence.",
"parameters": {
"type": "object",
"properties": {
"query": {"type": "string"},
},
"required": ["query"],
},
},
}

Оба протокола несут одинаковый смысл для обучения с подкреплением: Assistant выдаёт обучаемое действие, поисковый бэкенд возвращает observation среды. Собственный шаблон модели уменьшает конфликт между базовыми способностями модели и форматом инструментов и к тому же сохраняет структурированные сообщения инструмента.

Примечание переводчика. Описание инструмента ("Search the web for evidence." — «Искать в сети подтверждения») оставлено по-английски, как и все промпты этого проекта: модель обучают на английских вопросах NQ и HotpotQA. Системный промпт из protocol.py в переводе: «Ты отвечаешь на фактические вопросы с помощью поискового инструмента. Ищи, когда нужны подтверждения. Можно вызывать search несколько раз с краткими английскими запросами. Вызывай search ровно один раз за ход ассистента и жди результата, прежде чем искать снова. Когда будешь готов, закончи ровно одной непустой строкой в формате Answer: <короткий ответ>. Не вызывай инструмент и не давай итоговый ответ в одном ходе» — и дальше разобран пример про «Маленького принца» в пять шагов.

Определив инструмент, нужно ещё строго разбирать вывод модели на три состояния: допустимый поиск, итоговый ответ или недопустимый вывод. Парсер в protocol.py требует, чтобы в одном ходе был ровно один полный вызов search и чтобы после вызова инструмента не было «подмешанного» итогового ответа:

def parse_assistant(text: str) -> ParsedAssistant:
matches = list(TOOL_CALL_PATTERN.finditer(text))
if not matches:
kind = (
"invalid" if "<tool_call>" in text else "answer"
)
return ParsedAssistant(
kind=kind,
content=text.strip(),
)
if len(matches) != 1 or text[matches[0].end():].strip():
return ParsedAssistant(
kind="invalid",
content=text.strip(),
)
query = matches[0].group(1).strip()
if not query or "<" in query or ">" in query:
return ParsedAssistant(
kind="invalid",
content=text.strip(),
)
content = text[:matches[0].start()].strip()
return ParsedAssistant(
kind="tool",
content=content,
query=query,
)

Отдельного действия «стоп» здесь нет. Если в выводе нет допустимого вызова инструмента, ход считается итоговым ответом и траектория заканчивается; потом reward.py проверит, соответствует ли он формату Answer: <short answer>. Вывод с обрывком <tool_call> помечается как invalid — он тоже заканчивает траекторию и получает штраф за формат.

Примечание переводчика. TOOL_CALL_PATTERN в файле ищет XML-подобную разметку, в которой Qwen3.5 записывает вызовы инструментов: <tool_call><function=search><parameter=query>…</parameter></function></tool_call>. Это и есть «родной формат»: вызов инструмента для модели — просто текст с особыми тегами, который шаблон чата превращает в структуру. Пример трёх исходов: «Ищу.<tool_call>…query…</tool_call>» → tool; «Answer: Paris» → answer; «<tool_call><function=search>» (обрыв) → invalid.

8.3.2 Как поисковый бэкенд становится observation среды

Заголовок раздела «8.3.2 Как поисковый бэкенд становится observation среды»

search.py предоставляет три бэкенда: DeepSeek Search, Wikipedia и Чжиху (Zhihu). Rollout не зависит ни от какого формата HTTP-ответа — только от единого SearchResult:

@dataclass(frozen=True)
class SearchItem:
title: str
content: str
source: str | None = None
url: str | None = None
@dataclass(frozen=True)
class SearchResult:
ok: bool
items: list[SearchItem]
latency: float
status: int | None = None
error: str | None = None

Точка входа обучения создаёт бэкенд по аргументам командной строки, а дальнейший конечный автомат всегда вызывает один и тот же интерфейс search(query):

def create_search_client(
backend: str,
env_path: str | Path | None = None,
*,
model: str = "deepseek-v4-flash",
timeout: float | None = None,
) -> SearchClient:
resolved_timeout = resolve_search_timeout(
backend, timeout
)
if backend == "deepseek":
return DeepSeekSearchClient.from_env(
env_path,
model=model,
timeout=resolved_timeout,
)
if backend == "wikipedia":
return WikipediaSearchClient(
timeout=resolved_timeout
)
if backend == "zhihu":
return ZhihuSearchClient.from_env(
env_path,
timeout=resolved_timeout,
)
raise ValueError(f"Неподдерживаемый поисковый бэкенд: {backend}")

Возьмём для примера бэкенд Wikipedia, которому не нужен ключ: сам поиск происходит в _request(). Один запрос к API выполняет полнотекстовый поиск и сразу возвращает заголовки, начальные фрагменты текста и URL трёх первых страниц:

def _request(
self,
query: str,
started: float,
) -> SearchResult:
params = urllib.parse.urlencode(
{
"action": "query",
"format": "json",
"formatversion": 2,
"generator": "search",
"gsrsearch": query,
"gsrlimit": 3,
"prop": "extracts|info",
"explaintext": 1,
"exintro": 1,
"exchars": 1200,
"inprop": "url",
"redirects": 1,
"utf8": 1,
}
)
request = urllib.request.Request(
f"{WIKIPEDIA_SEARCH_ENDPOINT}?{params}",
headers={
"Accept": "application/json",
"Accept-Encoding": "gzip",
"User-Agent": WIKIPEDIA_USER_AGENT,
},
)
with urllib.request.urlopen(
request,
timeout=self.timeout,
) as response:
body = response.read()
if response.headers.get(
"Content-Encoding"
) == "gzip":
body = gzip.decompress(body)
payload = json.loads(body.decode("utf-8"))
pages = payload.get("query", {}).get(
"pages", []
)
if not isinstance(pages, list):
raise TypeError(
"В ответе поиска Wikipedia поле pages — не список"
)
if any(
not isinstance(page, dict)
for page in pages
):
raise TypeError(
"В ответе поиска Wikipedia page — не объект"
)
ordered_pages = sorted(
pages,
key=lambda page: int(
page.get("index", 1_000_000)
),
)
items = [
SearchItem(
title=str(
page.get("title") or "Untitled"
).strip(),
content=str(
page.get("extract") or ""
).strip(),
source="Wikipedia",
url=str(
page.get("fullurl") or ""
).strip(),
)
for page in ordered_pages
if str(page.get("extract") or "").strip()
]
return SearchResult(
ok=True,
items=items,
latency=time.perf_counter() - started,
status=response.status,
)

Внешний search() ещё отвечает за ограничение частоты запросов, timeout, ограниченное число повторов при ответах 429 и 5xx и копит статистику успешности и задержек. Если после повторов запрос всё равно не удался, он не обрушивает весь rollout исключением, а возвращает SearchResult(ok=False, error=...). rollout.py передаст модели и текст ошибки как observation — модель сможет ответить или переписать запрос.

Читаем параметры (от переводчика). Это MediaWiki API: generator=search + gsrsearch — полнотекстовый поиск, gsrlimit=3 — три страницы, prop=extracts с exintro и explaintext — вводный раздел статьи простым текстом, exchars=1200 — не длиннее 1200 символов, inprop=url — добавить ссылку. Ответ сортируется по полю index, то есть по релевантности поиска.

Результаты разных бэкендов в итоге форматируются как читаемые доказательства:

[1] Title: ...
Content: ...
Source: ...
URL: ...

Возвращать только заголовки или ссылки обычно мало: модели на самом деле нужен content, на который можно опереться в следующем шаге рассуждения. max_tool_response_tokens и max_trajectory_tokens ещё раз ограничивают длину доказательств перед тем, как observation попадёт в траекторию.

Одноходовому GRPO хватает одного sample(). В Search-R1 каждая траектория может пройти разное число поисковых вызовов, и тренеру приходится раз за разом выполнять «сгенерировать — разобрать — вызвать среду — дописать observation».

В этой главе траектория хранится как Trajectory:

@dataclass
class AssistantTurn:
prompt_tokens: list[int]
completion_tokens: list[int]
logprobs: list[float]
text: str
@dataclass
class Trajectory:
example: SearchExample
group_index: int
messages: list[dict[str, Any]]
next_prompt_tokens: list[int] | None = None
question_index: int = 0
turns: list[AssistantTurn] = field(default_factory=list)
search_calls: int = 0
final_text: str = ""
reward: float = -0.1
advantage: float = 0.0
valid_format: bool = False
exact_match: bool = False
done: bool = False

Каждый AssistantTurn хранит полный prompt перед сэмплированием этого хода, токены, сгенерированные моделью, и соответствующие logprob старой политики. Результаты поиска лежат в messages и в prompt следующего хода, но logprob от rollout у них нет: эти токены пришли из среды.

Запросы к PyTRIO на самом деле отправляет sample_requests_async(). В одном ходе разные вопросы или уже разошедшиеся траектории можно сэмплировать параллельно:

async def sample_requests_async(
sampling_client: Any,
requests: list[SampleRequest],
config: RolloutConfig,
tokenizer: Any,
) -> list[Any]:
tasks = []
for request in requests:
params = trio.SamplingParams(
max_tokens=request.max_tokens,
seed=request.seed,
stop=stop_sequences(tokenizer),
temperature=config.temperature,
top_p=config.top_p,
)
tasks.append(
sampling_client.sample_async(
prompt=trio.ModelInput.from_ints(
request.prompt_tokens
),
num_samples=request.num_samples,
sampling_params=params,
return_text=True,
)
)
return list(await asyncio.gather(*tasks))

Когда сэмплирование вернулось, consume_assistant() сначала сохраняет исходные токены и old logprob, затем вызывает парсер протокола из прошлого раздела. Допустимый поиск становится PendingSearch; итоговый ответ, недопустимый формат или исчерпанный бюджет заканчивают текущую траекторию:

def consume_assistant(
trajectory: Trajectory,
prompt_tokens: list[int],
sequence: Any,
tokenizer: Any,
config: RolloutConfig,
) -> PendingSearch | None:
tokens, logprobs, text = read_sequence(
sequence, tokenizer
)
trajectory.turns.append(
AssistantTurn(
prompt_tokens,
tokens,
logprobs,
text,
)
)
parsed = parse_assistant(text)
can_search = (
parsed.kind == "tool"
and trajectory.search_calls
< config.max_search_calls
and len(trajectory.turns)
< config.max_assistant_turns
)
if not can_search:
trajectory.messages.append(
{"role": "assistant", "content": text}
)
trajectory.final_text = text
trajectory.done = True
return None
call_id = (
f"search-{trajectory.question_index}-"
f"{trajectory.group_index}-"
f"{trajectory.search_calls + 1}"
)
messages_before_assistant = list(
trajectory.messages
)
trajectory.messages.append(
{"role": "assistant", "content": text}
)
return PendingSearch(
trajectory=trajectory,
messages_before_assistant=(
messages_before_assistant
),
assistant_text=text,
prompt_tokens=prompt_tokens,
completion_tokens=tokens,
call_id=call_id,
query=parsed.query or "",
)

resolve_searches() пулом потоков параллельно вызывает поисковый бэкенд для разных траекторий и затем в исходном порядке сопоставляет результаты ожидающим траекториям:

def resolve_searches(
pending_searches: list[PendingSearch],
search_client: SearchClient,
tokenizer: Any,
config: RolloutConfig,
) -> int:
if not pending_searches:
return 0
workers = min(
config.search_concurrency,
len(pending_searches),
)
with ThreadPoolExecutor(
max_workers=workers
) as pool:
results = list(
pool.map(
search_client.search,
[
pending.query
for pending in pending_searches
],
)
)
return sum(
finish_search(
pending,
result,
tokenizer,
config,
)
for pending, result in zip(
pending_searches,
results,
strict=True,
)
)

Каждый результат затем записывается обратно в свою траекторию функцией finish_search():

def finish_search(
pending: PendingSearch,
result: SearchResult,
tokenizer: Any,
config: RolloutConfig,
) -> bool:
fitted = fit_tool_content(
tokenizer,
pending.messages_before_assistant,
pending.assistant_text,
pending.prompt_tokens,
pending.completion_tokens,
pending.call_id,
result,
config,
)
if fitted is None:
pending.trajectory.final_text = (
pending.assistant_text
)
pending.trajectory.done = True
return True
content, next_prompt_tokens = fitted
pending.trajectory.messages.append(
tool_message(pending.call_id, content)
)
pending.trajectory.next_prompt_tokens = (
next_prompt_tokens
)
pending.trajectory.search_calls += 1
return False

fit_tool_content() добавляет найденные доказательства по одному. Как только observation превышает бюджет одного вызова инструмента или бюджет всей траектории, следующие доказательства перестают добавляться; если не помещается даже один результат, траектория сразу заканчивается. Так «сколько текста вернул поисковый сервис» и «сколько текста реально увидела модель» становятся двумя явно разделёнными границами.

Примечание переводчика. Обратите внимание: траектория, у которой не поместился результат, заканчивается на ходе с вызовом инструмента. Её final_text — это текст с <tool_call>, а не строка Answer: …, поэтому она получит награду −0,1-0{,}1 за неверный формат. Бюджет контекста тем самым тоже «учит» модель: слишком долгий поиск невыгоден.

Группа траекторий GRPO на первом ходе имеет полностью одинаковый prompt. Код одним запросом генерирует group_size ветвей:

first_requests: list[SampleRequest] = []
for index, trajectory in enumerate(roots):
request = make_request(
tokenizer,
trajectory,
index,
config.group_size,
config.seed + index,
config,
)
if request:
first_requests.append(request)

Каждая последовательность (sequence) из ответа сэмплирования глубоко копируется из корневой траектории в независимую ветвь, после чего её поисковое действие разбирается отдельно:

responses = asyncio.run(
sample_requests_async(
sampling_client,
first_requests,
config,
tokenizer,
)
)
pending_searches: list[PendingSearch] = []
for request, response in zip(
first_requests,
responses,
strict=True,
):
root = roots[request.trajectory_index]
if len(response.sequences) != config.group_size:
raise ValueError(
"Число выборок первого хода не совпадает с group_size"
)
for group_index, sequence in enumerate(
response.sequences
):
branch = copy.deepcopy(root)
branch.group_index = group_index
pending = consume_assistant(
branch,
request.prompt_tokens,
sequence,
tokenizer,
config,
)
trajectories.append(branch)
if pending is not None:
pending_searches.append(pending)
resolve_searches(
pending_searches,
search_client,
tokenizer,
config,
)

После поиска у разных ветвей уже разные observation и история действий. На следующих ходах для каждой незаконченной траектории отдельно ставится num_samples=1:

while any(not item.done for item in trajectories):
requests: list[SampleRequest] = []
for index, trajectory in enumerate(trajectories):
if trajectory.done:
continue
request = make_request(
tokenizer,
trajectory,
index,
1,
(
config.seed
+ index
+ len(trajectory.turns) * 10_000
),
config,
)
if request:
requests.append(request)
if not requests:
break
responses = asyncio.run(
sample_requests_async(
sampling_client,
requests,
config,
tokenizer,
)
)
pending_searches = []
for request, response in zip(
requests, responses, strict=True
):
if len(response.sequences) != 1:
raise ValueError(
"На следующих ходах сэмплируется ровно одна ветвь"
)
trajectory = trajectories[
request.trajectory_index
]
pending = consume_assistant(
trajectory,
request.prompt_tokens,
response.sequences[0],
tokenizer,
config,
)
if pending is not None:
pending_searches.append(pending)
resolve_searches(
pending_searches,
search_client,
tokenizer,
config,
)

Такая реализация «общий корень, затем ветвление» одновременно выполняет два условия: траектории одной задачи можно сравнивать внутри группы GRPO, а каждая ветвь может действовать дальше по своим результатам поиска. Несколько запросов сэмплирования к PyTRIO и поисковых запросов выполняются параллельно, и это сокращает ожидание, которое приносит многоходовое взаимодействие со средой.

Конечному автомату ещё нужны чёткие границы:

  • max_search_calls: сколько раз максимум одна траектория может вызвать поиск;
  • max_assistant_turns: сколько максимум ходов Assistant генерируется;
  • max_assistant_tokens: предел токенов на одно действие (один ход);
  • max_tool_response_tokens: предел токенов на одно observation;
  • max_trajectory_tokens: общий предел токенов на весь контекст;
  • timeout и concurrency (параллелизм) поиска: ограничивают ожидание внешнего сервиса и нагрузку на него.

Эти ограничения и контролируют стоимость обучения, и определяют, к какой части среды агент имеет доступ. При оценке нужно сохранять ту же конфигурацию — иначе больший бюджет поиска сам по себе может повысить долю правильных ответов.

Примечание переводчика. Значения по умолчанию в train.py: 8 вопросов на шаг, группа из 8 траекторий, до 4 поисков, до 6 ходов ассистента, до 1024 токенов на ход и на один результат поиска, до 8192 токенов на траекторию; бэкенд по умолчанию — DeepSeek Search (нужен ключ в .env), модель — Qwen/Qwen3.5-4B с LoRA ранга 32. Значит, на один шаг обучения — 64 траектории и до 256 поисковых запросов.

8.3.4 Протокол инструментов и непрерывность токенов

Заголовок раздела «8.3.4 Протокол инструментов и непрерывность токенов»

Самая частая ошибка многоходового обучения — токены после повторной отрисовки шаблона чата не совпадают с реально сэмплированными. Токенизатор может почистить пробелы, нормализовать текст вызова инструмента или сам добавить символ конца хода. Если на этапе обучения заново кодировать текст модели, старые logprob перестанут один к одному соответствовать целевым токенам.

В этой главе соблюдаются два принципа:

  1. для действий Assistant всегда хранятся исходные токены, которые вернул sampler;
  2. через шаблон чата вычисляются только вновь добавляемые символ конца хода и токены observation инструмента.

Полная задача build_next_prompt() — извлечь из шаблона чата только приращение «конец хода Assistant + новое observation» и приклеить его после настоящего ответа, который вернул sampler:

def build_next_prompt(
tokenizer: Any,
messages_before_assistant: list[dict[str, Any]],
assistant_text: str,
previous_prompt_tokens: list[int],
completion_tokens: list[int],
next_tool_message: dict[str, Any],
) -> list[int]:
canonical_prompt = build_prompt(
tokenizer,
messages_before_assistant,
)
empty_assistant_end = _render_chat(
tokenizer,
[
*messages_before_assistant,
{"role": "assistant", "content": ""},
],
add_generation_prompt=False,
)
if empty_assistant_end[
:len(canonical_prompt)
] != canonical_prompt:
raise ValueError(
"По пустому сообщению assistant не удалось выделить границу конца хода в chat template"
)
assistant_closing_tokens = empty_assistant_end[
len(canonical_prompt):
]
assistant_message = {
"role": "assistant",
"content": assistant_text,
}
messages_with_assistant = [
*messages_before_assistant,
assistant_message,
]
canonical_assistant_end = _render_chat(
tokenizer,
messages_with_assistant,
add_generation_prompt=False,
)
canonical_next_prompt = build_prompt(
tokenizer,
[*messages_with_assistant, next_tool_message],
)
if canonical_next_prompt[
:len(canonical_assistant_end)
] != canonical_assistant_end:
raise ValueError(
"После добавления tool observation chat template переписал историю сообщений"
)
observation_tokens = canonical_next_prompt[
len(canonical_assistant_end):
]
overlap = _suffix_prefix_overlap(
completion_tokens,
assistant_closing_tokens,
)
return [
*previous_prompt_tokens,
*completion_tokens,
*assistant_closing_tokens[overlap:],
*observation_tokens,
]

overlap обрабатывает случай, когда sampler уже вернул часть символов конца хода Assistant, — чтобы не дописать токены дважды. messages нужны для записи протокола и отладки, а next_prompt_tokens — это то, что действительно отправляется sampler на следующем ходе; путать их нельзя.

Пример (от переводчика). Пусть конец хода в шаблоне — токены [<|im_end|>, \n], а sampler остановился на стоп- токене и вернул ответ, оканчивающийся на <|im_end|>. Тогда самый длинный «хвост ответа = начало закрывающих токенов» — это один токен, overlap = 1, и дописывается только \n. Если sampler вернул ответ без <|im_end|>, overlap = 0 и дописываются оба токена.

При построении обучающих данных для следующего хода выполняется ещё одна проверка:

if turn.prompt_tokens[:len(full_tokens)] != full_tokens:
raise ValueError(
"Prompt следующего хода не является продолжением уже собранной траектории"
)

Эти два уровня проверок рано выявляют переписывание шаблона, дублирование конца хода и ошибки склейки observation. В Agentic RL непрерывность токенов — часть правильности алгоритма: траектория может выглядеть разумно на уровне текста и всё равно давать неверный градиент из-за сдвига границ токенов.

Примечание переводчика. Типичный пример «переписывания истории» — шаблоны рассуждающих моделей, которые при отрисовке удаляют блоки <think>…</think> из прошлых ходов ассистента. Тогда префикс следующего prompt не совпадёт с тем, что модель видела при генерации, и первая проверка выдаст ошибку. Вот почему код сравнивает токены, а не текст.

8.3.5 Награда, групповое преимущество и маска наблюдений

Заголовок раздела «8.3.5 Награда, групповое преимущество и маска наблюдений»

В этой главе используются задачи с коротким ответом из NQ и HotpotQA. В конце модель должна вывести ровно одну строку:

Answer: <short answer>

Награда по правилам — три уровня:

R(τ)={1,формат верен и ответ точно совпал,0,формат верен, но ответ неверен,−0.1,итоговый формат недопустим.R(\tau)=\begin{cases}1, & \text{формат верен и ответ точно совпал},\\ 0, & \text{формат верен, но ответ неверен},\\ -0.1, & \text{итоговый формат недопустим}.\end{cases}

reward.py сначала извлекает единственную строку итогового ответа, затем приводит к одному виду регистр, пунктуацию, английские артикли и пробелы:

def normalize_answer(text: str) -> str:
lowered = text.lower()
without_punctuation = "".join(
char
for char in lowered
if not unicodedata.category(char).startswith("P")
)
without_articles = ARTICLE_PATTERN.sub(
" ", without_punctuation
)
return " ".join(without_articles.split())
def extract_answer(text: str) -> str | None:
matches = ANSWER_PATTERN.findall(text)
if len(matches) != 1:
return None
answer = matches[0].strip()
return answer or None
def score_answer(
text: str,
references: list[str],
) -> RewardResult:
answer = extract_answer(text)
if answer is None:
return RewardResult(
-0.1, False, False, None
)
normalized = normalize_answer(answer)
exact_match = any(
normalized == normalize_answer(reference)
for reference in references
)
return RewardResult(
float(exact_match),
True,
exact_match,
answer,
)

Этот код оценивает только итоговый результат задачи — никаких дополнительных баллов за «сколько раз искал» или «запрос выглядит хорошо». Для GG полных траекторий одного вопроса по-прежнему используется групповая база GRPO:

Ai=R(τi)−1G∑j=1GR(τj).A_i=R(\tau_i)-\frac{1}{G}\sum_{j=1}^{G}R(\tau_j).

В статье Search-R1 рассматриваются разные способы оптимизации, включая PPO и GRPO. В этой главе выбрана схема GRPO с группами по одной задаче, чтобы Search-R1 строил advantage тем же способом, что и раздел 8.1.

Примечание переводчика. ANSWER_PATTERN — ^\s*Answer:\s*(.*?)\s*$ с флагами «без учёта регистра» и «многострочный»; ARTICLE_PATTERN удаляет слова a, an, the. Пример: ответ Answer: The Eiffel Tower. и эталон eiffel tower после нормализации оба дают eiffel tower — награда 1. Две строки Answer: в выводе — уже недопустимый формат, −0,1-0{,}1. Такую метрику называют точное совпадениеточное совпадение (EM)EN exact matchМетрика «1, если нормализованный ответ совпал с эталоном, иначе 0». Всё решает нормализация и извлечение ответа..

rollout_batch() обязан дождаться окончания всех траекторий группы, и только потом выполнить такую функцию:

def assign_group_advantages(
trajectories: list[Trajectory],
) -> int:
groups: dict[int, list[Trajectory]] = {}
for trajectory in trajectories:
groups.setdefault(
trajectory.question_index, []
).append(trajectory)
degenerate = 0
for group in groups.values():
mean_reward = sum(
item.reward for item in group
) / len(group)
for item in group:
item.advantage = (
item.reward - mean_reward
)
if all(
item.advantage == 0.0 for item in group
):
degenerate += 1
return degenerate

Когда награды всей группы одинаковы, все advantage равны 0, и группа не даёт полезного градиента. Код обучения пропускает такие траектории и записывает долю вырожденных групп. Групповое среднее нельзя считать после разбиения на micro-batch (микробатчи) — иначе база сравнения уже изменится.

Этот advantage присваивается всем токенам, сгенерированным в каждом ходе Assistant траектории. Observation поиска — лишь состояние, нужное для следующего решения, а не действие политики, поэтому logprob и advantage у токенов observation равны 0:

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

Такая observation mask (маска наблюдений) делает две вещи. Во-первых, тренер не требует от модели предсказывать текст, который вернул поисковик; во-вторых, observation остаётся в контексте, и последующие токены Assistant могут опираться на эти доказательства.

Вызовы поиска и итоговый ответ делят один advantage на уровне траектории. Если итоговый ответ верен, поощряются и спланированные запросы, и поисковые слова, и токены ответа, которые привели к успеху; если неверен — подавляется вся цепочка действий. Именно поэтому награда за результат (outcome reward) способна обучить стратегию использования инструментов.

Пример (от переводчика). По рисунку 8.4: награды [1,0,1,0][1, 0, 1, 0], среднее 0,50{,}5, advantage [+0,5; −0,5; +0,5; −0,5][+0{,}5;\ -0{,}5;\ +0{,}5;\ -0{,}5]. В траектории τ0\tau_0 из двух поисков и ответа все три хода ассистента получат +0,5+0{,}5 на каждый токен; тексты из Википедии между ними — 0. Если бы модель «зазубрила» текст Википедии без маски, она училась бы генерировать чужие статьи — ровно то, чего маска не допускает.

build_datum() в train.py сводит вместе непрерывность токенов и маску наблюдений из предыдущих разделов. Вот полная основная реализация:

def build_datum(
trajectory: Trajectory,
) -> TrainingDatum:
if not trajectory.turns:
raise ValueError(
"Нельзя построить обучающий Datum из траектории без ходов assistant"
)
full_tokens: list[int] = []
old_logprobs_by_token: list[float] = []
advantages_by_token: list[float] = []
assistant_token_count = 0
for turn_index, turn in enumerate(
trajectory.turns
):
if len(turn.completion_tokens) != len(
turn.logprobs
):
raise ValueError(
f"В ходе assistant №{turn_index + 1} "
"длины токенов и logprob не совпадают"
)
if turn_index == 0:
delta_observation = turn.prompt_tokens
elif turn.prompt_tokens[
:len(full_tokens)
] == full_tokens:
delta_observation = turn.prompt_tokens[
len(full_tokens):
]
else:
raise ValueError(
f"Prompt хода assistant №{turn_index + 1} "
"не является продолжением уже собранной траектории, "
"logprob сэмплирования нельзя безопасно выровнять"
)
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)
)
assistant_token_count += len(
turn.completion_tokens
)
if assistant_token_count == 0:
raise ValueError(
"Нельзя построить обучающий Datum из траектории без токенов assistant"
)
if not (
len(full_tokens)
== len(old_logprobs_by_token)
== len(advantages_by_token)
):
raise ValueError(
"У полной траектории не совпадают длины токенов, logprob и advantage"
)
input_tokens = full_tokens[:-1]
target_tokens = full_tokens[1:]
old_logprobs = old_logprobs_by_token[1:]
advantages = advantages_by_token[1:]
if not (
len(input_tokens)
== len(target_tokens)
== len(old_logprobs)
== len(advantages)
):
raise ValueError(
"У Datum не совпадают длины input, target, logprobs и advantages"
)
if len(input_tokens) > MAX_TRAIN_CONTEXT_TOKENS:
raise ValueError(
f"Datum длиннее {MAX_TRAIN_CONTEXT_TOKENS} токенов"
)
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,
),
},
)
return TrainingDatum(datum, len(input_tokens))

На первом ходе delta_observation — это системный промпт, описание инструментов и вопрос пользователя; на следующих — конец предыдущего хода, результаты поиска и начало нового хода Assistant. В них сохраняются настоящие целевые токены, но old logprob и advantage равны 0. Только токены Assistant, реально сгенерированные sampler, несут old logprob из rollout и advantage траектории.

Примечание переводчика. Сравните со сдвигом в 8.1: там prompt «отрезали» вручную (len(prompt) - 1 нулей), здесь сначала собирают всю последовательность токенов, а потом сдвигают одним движением: input = full[:-1], target = full[1:], а logprob и advantage сдвигают так же, как target. Результат тот же, но способ работает для любого числа ходов. MAX_TRAIN_CONTEXT_TOKENS в файле равен 8192.

Награды всей группы должны быть посчитаны полностью — только тогда получится правильное групповое среднее. Обновлять политику сразу после окончания одной траектории нельзя: награды других ветвей той же группы ещё неизвестны.

Длины многоходовых траекторий сильно различаются, и если положить все данные в один батч с дополнением (padding), впустую уйдёт много токенов. Сопутствующий код раскладывает Datum по нескольким micro-batch по длине примеров и перевзвешивает потерю пропорционально доле траекторий каждого микробатча в полном rollout. Так удаётся держать под контролем объём padding в одном батче и при этом сохранить смысл градиента как среднего по всем примерам.

8.3.7 От rollout к одному обновлению параметров

Заголовок раздела «8.3.7 От rollout к одному обновлению параметров»

Внешний цикл train.py соединяет все предыдущие модули в полный шаг обучения. Ключевой порядок менять нельзя:

batch = take_batch(
examples,
step * args.questions_per_batch,
args.questions_per_batch,
)
sampling_client = (
training_client
.save_weights_and_get_sampling_client()
)
trajectories = rollout_batch(
sampling_client=sampling_client,
tokenizer=tokenizer,
search_client=search_client,
examples=batch,
config=rollout_config,
)
datums = build_training_datums(trajectories)
micro_batches = pack_micro_batches(datums)
trainer_results = []
for micro_batch in micro_batches:
weighted_datums = (
weight_micro_batch_for_global_mean(
micro_batch,
total_samples=len(trajectories),
)
)
result = training_client.forward_backward(
weighted_datums,
loss_fn="importance_sampling",
).result()
trainer_results.append(result)
if micro_batches:
training_client.optim_step(
adam_params
).result()

Первое: save_weights_and_get_sampling_client() в начале каждого шага выгружает текущие LoRA-веса, поэтому old logprob берутся у текущей политики до этого обновления. Второе: rollout_batch() внутри уже посчитал награды и advantage для целых групп, а build_training_datums() пропускает траектории с нулевым advantage. Третье: все micro-batch только накапливают градиенты, и для всего логического батча в конце вызывается один optim_step().

PyTRIO усредняет потери по примерам внутри одного forward_backward(). Если в kk-м микробатче nkn_k траекторий, а в полном батче rollout — NN траекторий, код сначала умножает advantage этого микробатча на nk/Nn_k/N:

∑knkNmean⁡(Lk)=mean⁡(Lglobal).\sum_k\frac{n_k}{N}\operatorname{mean}(\mathcal L_k)=\operatorname{mean}(\mathcal L_{\mathrm{global}}).

Реализация только заново упаковывает Datum, не меняя целевые токены и old logprob:

micro_batch_weight = np.float32(
len(micro_batch) / total_samples
)
weighted_datums = []
for item in micro_batch:
loss_inputs = item.datum.loss_fn_inputs
weighted_datums.append(
trio.Datum(
model_input=item.datum.model_input,
loss_fn_inputs={
"target_tokens": (
loss_inputs["target_tokens"]
.to_numpy()
),
"logprobs": (
loss_inputs["logprobs"]
.to_numpy()
),
"advantages": (
loss_inputs["advantages"]
.to_numpy()
* micro_batch_weight
),
},
)
)

Поэтому динамическое разбиение на батчи меняет только padding и расход видеопамяти при выполнении, но не меняет градиент логического батча как среднее по примерам.

Проверим формулу (от переводчика). N=8N = 8 траекторий с потерями L1,…,L8L_1, \ldots, L_8 разложены на микробатчи из 5 и 3. Каждый forward_backward усредняет: 15∑i≤5Li\frac{1}{5}\sum_{i \le 5} L_i и 13∑i>5Li\frac{1}{3}\sum_{i > 5} L_i. С весами 58\frac{5}{8} и 38\frac{3}{8} сумма равна 18∑i=18Li\frac{1}{8}\sum_{i=1}^{8} L_i — среднему по всем восьми. Без весов траектории из маленького микробатча весили бы больше. Умножать можно именно advantage, потому что потеря importance_sampling линейна по нему. Тот же приём — накопление градиентанакопление градиентовEN gradient accumulationГрадиенты складываются (+=), а не перезаписываются. Внутри графа так учитываются все пути к узлу; на уровне обучения — шаг оптимизатора делают раз в несколько микробатчей. с правильной нормировкой.

Код Search-R1 разделён по обязанностям: данные, протокол, среда, rollout и обучение:

ФайлНазначение
prepare_data.pyготовит данные NQ и HotpotQA
data.pyзагружает примеры в едином формате JSONL
protocol.pyописывает поисковый инструмент и протокол многоходовых сообщений
search.pyоборачивает поисковые бэкенды и приводит результаты к единому формату
rollout.pyвыполняет конечный автомат многоходового поиска
reward.pyизвлекает короткий ответ и считает награду по правилам
train.pyстроит маску наблюдений и микробатчи, обновляет политику
eval.pyоценивает в той же среде исходную модель (Base Model) или контрольную точку
analyse.pyсводит результаты оценки

Для оценки не написана отдельная упрощённая логика поиска — переиспользуется тот же rollout_batch(). Главное отличие — на каждый вопрос остаётся одна траектория: сравнение внутри группы, нужное при обучении, больше не требуется:

config = RolloutConfig(
group_size=1,
max_search_calls=args.max_search_calls,
max_assistant_turns=args.max_assistant_turns,
max_trajectory_tokens=(
args.max_trajectory_tokens
),
max_assistant_tokens=args.max_assistant_tokens,
max_tool_response_tokens=(
args.max_tool_response_tokens
),
search_concurrency=args.search_concurrency,
temperature=args.temperature,
top_p=args.top_p,
seed=args.seed,
)
batch_trajectories = rollout_batch(
sampling_client,
tokenizer,
search_client,
batch,
config,
)

Если model_path пуст, eval.py создаёт sampler исходной модели; если передать путь, который вернул save_weights_for_sampler(), оценивается обученная контрольная точка. Обе модели используют одинаковые протокол, поисковый бэкенд, бюджет, набор вопросов и проверяльщик ответов.

Сначала подготовьте данные:

Окно терминала
python docs/chapter8/search-r1/prepare_data.py

Бэкенду Wikipedia не нужен API Key — он подходит, чтобы проверить всю цепочку:

Окно терминала
python docs/chapter8/search-r1/train.py \
--max-steps 1 \
--questions-per-batch 1 \
--group-size 4 \
--max-search-calls 2 \
--search-backend wikipedia \
--search-concurrency 3 \
--swanlab-mode disabled

Оцените исходную модель в той же среде Wikipedia:

Окно терминала
python docs/chapter8/search-r1/eval.py \
--search-backend wikipedia \
--batch-size 8 \
--output docs/chapter8/search-r1/eval_result/base.jsonl

Для оценки обученных весов сэмплирования оставьте остальные параметры без изменений и добавьте только путь к контрольной точке:

Окно терминала
python docs/chapter8/search-r1/eval.py \
--search-backend wikipedia \
--batch-size 8 \
--model-path 'trio://runxxxxxxxxxx' \
--output docs/chapter8/search-r1/eval_result/checkpoint.jsonl

Убедившись, что процесс работает, постепенно увеличивайте число вопросов, group size, число поисков и бюджет траектории. В эксперименте нужно записывать как минимум:

  • exact match итогового ответа;
  • долю ответов в допустимом формате;
  • среднее число вызовов поиска и долю траекторий без поиска;
  • долю неудачных поисков, таймаутов и обрезанных observation;
  • среднюю длину траектории и число полезных обучающих токенов на шаг;
  • долю вырожденных групп;
  • время каждого этапа: rollout, поиск и обучение.

Результаты поиска зависят от внешних систем. Онлайн-страницы обновляются, ранжирование поисковых сервисов тоже может меняться. Сравнивая исходную модель и обученную контрольную точку, фиксируйте набор вопросов, поисковый бэкенд, бюджет поиска, параметры сэмплирования и временное окно оценки. Для строго воспроизводимых экспериментов уровня статьи корпус и индекс поиска можно зафиксировать как снимок только для чтения.

В этой главе обновляется только политика языковой модели, поисковый бэкенд остаётся неизменным. Если обучаться должен и сам поисковик, нужно дополнительно определить пространство поисковых действий, версию индекса и распределение заслуг (credit assignment) между модулями — это уже выходит за рамки примера этого раздела.

Search-R1 показал, как расширить групповое относительное преимущество GRPO на многоходовые поисковые траектории с внешними observation. Следующий раздел, ReTool, сохранит тот же многоходовый rollout и маску наблюдений, заменит среду — вместо поисковика будет интерпретатор кода — и будет обновлять политику встроенной в PyTRIO потерей PPO.


  1. Парсер и награда без GPU. Скопируйте parse_assistant (с TOOL_CALL_PATTERN из примечания) и score_answer и прогоните десяток строк: правильный вызов, два вызова в одном ходе, вызов + Answer:, обрыв <tool_call>, Answer: The Beatles против эталона beatles.
  2. Поиск как среда. Вызовите MediaWiki API с параметрами из _request() для запроса Antoine de Saint-Exupéry birthplace: что вернётся, сколько символов, сколько это токенов в токенизаторе Qwen?
  3. Маска на бумаге. Для траектории «prompt 50 токенов → ход 1 (20) → результат поиска (300) → ход 2 (8)» посчитайте длину Datum и сколько позиций получат ненулевой advantage. Ответ: 50+20+300+8−1=37750 + 20 + 300 + 8 - 1 = 377 позиций, из них 20+8=2820 + 8 = 28 с advantage.
  4. Локально. Тот же цикл без PyTRIO умеют verl (многоходовые rollout с инструментами) и slime; как китайские лаборатории строят агентный RL, разобрано в уроке 8.3.4.
EnglishРусский
agentic RLагентное обучение с подкреплением
autonomous searchсамостоятельный поиск
search backendпоисковый бэкенд
observationнаблюдение
observation maskмаска наблюдений
multi-turn trajectoryмногоходовая траектория
state machineконечный автомат
shared root / branchingобщий корень / ветвление
tool protocolпротокол инструментов
token continuityнепрерывность токенов
end-of-turn tokenсимвол конца хода
outcome rewardнаграда за результат
exact matchточное совпадение
budgetбюджет (предел)
micro-batchмикробатч
credit assignmentраспределение заслуг
checkpointконтрольная точка