diff --git a/docs/65-example-search-r1.md b/docs/65-example-search-r1.md
index a98325a1c..079b0f07f 100644
--- a/docs/65-example-search-r1.md
+++ b/docs/65-example-search-r1.md
@@ -15,6 +15,11 @@ The example supports two API styles:
1. **Chat Completions API:** the standard text-in/text-out API used by default.
2. **Token-in/token-out Completions API:** the agent sends prompt token IDs and receives generated token IDs while preserving the multi-turn token sequence.
+Both modes compute exact-match reward from the answer action selected by the agent.
+Retrieved passages and later, unused answer tags do not count as the agent's answer.
+The extra model response after the turn limit follows the same rule: its first action
+must be an answer to receive exact-match reward.
+
## Data Preparation
Prepare the Wikipedia corpus, E5 FAISS index, training data, and retriever environment from the repository root:
diff --git a/examples/search_r1/agents/search_r1_agent.py b/examples/search_r1/agents/search_r1_agent.py
index 0ea80617d..946e3e1f7 100644
--- a/examples/search_r1/agents/search_r1_agent.py
+++ b/examples/search_r1/agents/search_r1_agent.py
@@ -15,7 +15,7 @@
import httpx
-from examples.search_r1.agents.qa_em import compute_score_em
+from examples.search_r1.agents.qa_em import em_check
log = logging.getLogger(__name__)
@@ -235,7 +235,7 @@ async def run(self) -> None:
client = AsyncOpenAI(base_url=openai_base_url, api_key=agl_key, max_retries=6)
prompt = INSTRUCTION_FORMAT + question
messages = [{"role": "user", "content": prompt}]
- rollout_content = ""
+ final_answer: str | None = None
finished = False
invalid_action = False
start_time = time.time()
@@ -249,10 +249,9 @@ async def run(self) -> None:
max_tokens=max_tokens,
)
valid_turn_response = postprocess_response(turn_response)
- rollout_content += valid_turn_response
messages.append({"role": "assistant", "content": valid_turn_response})
- action, _ = extract_action(valid_turn_response)
+ action, content = extract_action(valid_turn_response)
if action is None:
invalid_action = True
log.info("turn=%d invalid response=%r", turn_id, valid_turn_response)
@@ -260,10 +259,10 @@ async def run(self) -> None:
turn_env_feedback = await execute_response(valid_turn_response, retrieval_url=retrieval_url, topk=topk)
if not turn_env_feedback:
+ final_answer = content
finished = True
log.info("turn=%d finished response=%r", turn_id, valid_turn_response)
break
- rollout_content += turn_env_feedback
messages.append({"role": "user", "content": turn_env_feedback})
log.info("turn=%d response=%r env_feedback_chars=%d", turn_id, valid_turn_response, len(turn_env_feedback))
@@ -274,11 +273,13 @@ async def run(self) -> None:
temperature=train_temperature,
max_tokens=max_tokens,
)
- rollout_content += turn_response
messages.append({"role": "assistant", "content": turn_response})
+ action, content = extract_action(turn_response)
+ if action == "answer":
+ final_answer = content
log.info("last_turn response=%r", turn_response)
- reward = 0.0 if invalid_action else float(compute_score_em(rollout_content, answer_list))
+ reward = 0.0 if final_answer is None else float(em_check(final_answer, answer_list))
reason = "invalid_action" if invalid_action else "em_match" if reward > 0 else "em_miss"
await post_reward(event_url, agl_key, reward, reason)
@@ -311,7 +312,7 @@ async def run(self) -> None:
client = AsyncOpenAI(base_url=openai_base_url, api_key=agl_key, max_retries=6)
prompt_ids = encode_text(tokenizer_model, INSTRUCTION_FORMAT + question)
- rollout_content = ""
+ final_answer: str | None = None
finished = False
invalid_action = False
start_time = time.time()
@@ -331,10 +332,9 @@ async def run(self) -> None:
temperature=train_temperature,
max_tokens=max_tokens,
)
- rollout_content += turn_response
next_prompt_ids = prompt_ids + response_ids
- action, _ = extract_action(turn_response)
+ action, content = extract_action(turn_response)
if action is None:
invalid_action = True
log.info("turn=%d invalid response=%r response_tokens=%d", turn_id, turn_response, len(response_ids))
@@ -342,12 +342,12 @@ async def run(self) -> None:
turn_env_feedback = await execute_response(turn_response, retrieval_url=retrieval_url, topk=topk)
if not turn_env_feedback:
+ final_answer = content
finished = True
log.info("turn=%d finished response=%r response_tokens=%d", turn_id, turn_response, len(response_ids))
break
feedback_ids = encode_text(tokenizer_model, turn_env_feedback)
- rollout_content += turn_env_feedback
prompt_ids = next_prompt_ids + feedback_ids
log.info(
"turn=%d response=%r response_tokens=%d env_feedback_tokens=%d",
@@ -365,10 +365,12 @@ async def run(self) -> None:
temperature=train_temperature,
max_tokens=max_tokens,
)
- rollout_content += turn_response
+ action, content = extract_action(turn_response)
+ if action == "answer":
+ final_answer = content
log.info("last_turn response=%r response_tokens=%d", turn_response, len(response_ids))
- reward = 0.0 if invalid_action else float(compute_score_em(rollout_content, answer_list))
+ reward = 0.0 if final_answer is None else float(em_check(final_answer, answer_list))
reason = "invalid_action" if invalid_action else "em_match" if reward > 0 else "em_miss"
await post_reward(event_url, agl_key, reward, reason)
diff --git a/tests/examples/test_search_r1_reward.py b/tests/examples/test_search_r1_reward.py
new file mode 100644
index 000000000..cdf1f2c58
--- /dev/null
+++ b/tests/examples/test_search_r1_reward.py
@@ -0,0 +1,111 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Search-R1 rewards must describe the answer action selected by the agent."""
+
+from __future__ import annotations
+
+import copy
+from typing import Any
+from unittest.mock import AsyncMock
+
+import openai
+import pytest
+
+from examples.search_r1.agents import search_r1_agent
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("api_type", ["chat", "completion"])
+@pytest.mark.parametrize(
+ ("responses", "max_turns", "retrieval_text", "expected_reward", "expected_reason"),
+ [
+ (["LondonParis"], 2, "", 0.0, "em_miss"),
+ (["ParisLondon"], 2, "", 1.0, "em_match"),
+ (["Paris"], 2, "", 0.0, "em_miss"),
+ (["capital", "I do not know."], 1, "Paris", 0.0, "em_miss"),
+ (["capital", "Paris"], 2, "Paris is the capital.", 1.0, "em_match"),
+ (["capital", "Paris"], 1, "Paris is the capital.", 1.0, "em_match"),
+ (["No action tags."], 2, "", 0.0, "invalid_action"),
+ (
+ ["capital", "againParis"],
+ 1,
+ "Paris is the capital.",
+ 0.0,
+ "em_miss",
+ ),
+ (
+ ["capital", "LondonParis"],
+ 1,
+ "Paris is the capital.",
+ 0.0,
+ "em_miss",
+ ),
+ ],
+ ids=[
+ "first-answer-miss",
+ "first-answer-match",
+ "nested-answer-is-literal",
+ "retrieval-tags-without-answer",
+ "search-then-answer",
+ "answer-after-turn-limit",
+ "invalid-action",
+ "search-before-unused-answer-at-turn-limit",
+ "first-answer-at-turn-limit",
+ ],
+)
+async def test_reward_uses_selected_answer(
+ monkeypatch: pytest.MonkeyPatch,
+ api_type: str,
+ responses: list[str],
+ max_turns: int,
+ retrieval_text: str,
+ expected_reward: float,
+ expected_reason: str,
+) -> None:
+ for name, value in {
+ "QUESTION": "What is the capital of France?",
+ "GOLDEN_ANSWERS": '["Paris"]',
+ "AGL_KEY": "test-key",
+ "AGL_EVENT_URL": "http://gateway/events",
+ "AGL_OPENAI_BASE_URL": "http://gateway/v1",
+ "SEARCH_R1_MAX_TURNS": str(max_turns),
+ }.items():
+ monkeypatch.setenv(name, value)
+ monkeypatch.setattr(openai, "AsyncOpenAI", lambda **_kwargs: object())
+ replies = iter(responses)
+ contexts: list[Any] = []
+
+ async def call_llm(_client: object, messages: list[dict[str, str]], **_kwargs: Any) -> str:
+ contexts.append(copy.deepcopy(messages))
+ return next(replies)
+
+ async def call_completion(_client: object, prompt_ids: list[int], **_kwargs: Any) -> tuple[str, list[int]]:
+ contexts.append(list(prompt_ids))
+ text = next(replies)
+ return text, list(text.encode())
+
+ monkeypatch.setattr(search_r1_agent, "call_llm", call_llm)
+ monkeypatch.setattr(search_r1_agent, "call_completion", call_completion)
+ monkeypatch.setattr(search_r1_agent, "encode_text", lambda _model, text: list(text.encode()))
+ retrieve = AsyncMock(return_value=retrieval_text)
+ monkeypatch.setattr(search_r1_agent, "retrieve_doc", retrieve)
+ post_reward = AsyncMock()
+ monkeypatch.setattr(search_r1_agent, "post_reward", post_reward)
+
+ agent = search_r1_agent.SearchR1Agent() if api_type == "chat" else search_r1_agent.SearchR1CompletionAgent()
+ await agent.run()
+
+ post_reward.assert_awaited_once_with("http://gateway/events", "test-key", expected_reward, expected_reason)
+ assert len(contexts) == len(responses)
+ if len(responses) > 1:
+ retrieve.assert_awaited_once()
+ feedback = f"\n\n{retrieval_text}\n\n"
+ if api_type == "chat":
+ assert contexts[1][-2:] == [
+ {"role": "assistant", "content": responses[0]},
+ {"role": "user", "content": feedback},
+ ]
+ else:
+ assert contexts[1] == contexts[0] + list(responses[0].encode()) + list(feedback.encode())
+ else:
+ retrieve.assert_not_awaited()