fix(search-r1): score the selected answer action

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:

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,21 +249,20 @@ 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)
                 break
 
             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,23 +332,22 @@ 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))
                 break
 
             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)
 

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"),
+    [
+        (["<answer>London</answer><answer>Paris</answer>"], 2, "", 0.0, "em_miss"),
+        (["<answer>Paris</answer><answer>London</answer>"], 2, "", 1.0, "em_match"),
+        (["<answer><answer>Paris</answer></answer>"], 2, "", 0.0, "em_miss"),
+        (["<search>capital</search>", "I do not know."], 1, "<answer>Paris</answer>", 0.0, "em_miss"),
+        (["<search>capital</search>", "<answer>Paris</answer>"], 2, "Paris is the capital.", 1.0, "em_match"),
+        (["<search>capital</search>", "<answer>Paris</answer>"], 1, "Paris is the capital.", 1.0, "em_match"),
+        (["No action tags."], 2, "", 0.0, "invalid_action"),
+        (
+            ["<search>capital</search>", "<search>again</search><answer>Paris</answer>"],
+            1,
+            "Paris is the capital.",
+            0.0,
+            "em_miss",
+        ),
+        (
+            ["<search>capital</search>", "<answer>London</answer><answer>Paris</answer>"],
+            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<information>{retrieval_text}</information>\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()