Merge pull request AI-Hypercomputer#4140 from AI-Hypercomputer:pr/fix-eval-reward-question-v2

PiperOrigin-RevId: 930602952

Author: Google-ML-Automation

src/maxtext/trainers/post_train/rl/evaluate_rl.py

@@ -183,7 +183,7 @@ def score_responses(tmvp_config, question, responses, answers):
   raise ValueError(f"Unknown eval_mode: {eval_mode!r}")
 
 
-def _compute_row_reward(reward_fns, prompt, responses, answer, row_idx):
+def _compute_row_reward(reward_fns, prompt, question, responses, answer, row_idx):
   """Sum the per-function reward scores across all sampled responses for one prompt.
 
   Honors the sampling strategy `evaluate()` ran with: when `num_passes > 1`
@@ -193,6 +193,12 @@ def _compute_row_reward(reward_fns, prompt, responses, answer, row_idx):
   total by the number of (prompt, response) pairs to get the per-sample
   mean reward, mirroring tunix's GRPO per-rollout reward aggregation.
 
+  `question` is forwarded as a reward-fn kwarg because the built-in
+  `check_numbers` reward (and any user reward that wants to inspect the
+  raw question) reads it from `kwargs`. Training-time tunix passes the
+  whole dataset row to each reward fn; we mirror the kwargs the trainer
+  passes so eval-time and train-time reward calls are interchangeable.
+
   Returns a tuple `(score_sum, n_responses)`. On any exception the
   failure is logged and `(0.0, 0)` is returned so the caller's running
   mean is not corrupted.
@@ -203,7 +209,7 @@ def _compute_row_reward(reward_fns, prompt, responses, answer, row_idx):
     score_sum = 0.0
     for resp in responses:
       for fn in reward_fns:
-        scores = fn(prompts=[prompt], completions=[resp], answer=[answer])
+        scores = fn(prompts=[prompt], completions=[resp], answer=[answer], question=question)
         if scores:
           score_sum += float(scores[0])
     return score_sum, len(responses)
@@ -283,7 +289,7 @@ def evaluate(
       # the actual per-(prompt, response) count at the end. See
       # `_compute_row_reward` for details.
       if use_reward:
-        row_sum, row_count = _compute_row_reward(reward_fns, prompt, responses, answer, total)
+        row_sum, row_count = _compute_row_reward(reward_fns, prompt, question, responses, answer, total)
         reward_sum += row_sum
         reward_count += row_count
 

tests/post_training/unit/evaluate_rl_test.py

@@ -216,25 +216,27 @@ class TestComputeRowReward(unittest.TestCase):
   def _two_fns(self):
     """Return two reward functions whose per-response scores can be summed."""
 
-    # Each fn must accept prompts, completions, answer as keyword args and
-    # return a list of per-completion scores. The helper calls fn once per
-    # response with single-element lists, so the returned list has length 1.
-    def fn1(prompts, completions, answer):  # pylint: disable=unused-argument
+    # Each fn must accept prompts, completions, answer, question as keyword
+    # args and return a list of per-completion scores. The helper calls fn
+    # once per response with single-element lists, so the returned list has
+    # length 1.
+    def fn1(prompts, completions, answer, question):  # pylint: disable=unused-argument
       return [1.0 for _ in completions]
 
-    def fn2(prompts, completions, answer):  # pylint: disable=unused-argument
+    def fn2(prompts, completions, answer, question):  # pylint: disable=unused-argument
       return [float(len(c)) for c in completions]
 
     return [fn1, fn2]
 
   @pytest.mark.cpu_only
   def test_single_response_single_fn(self):
-    def fn(prompts, completions, answer):  # pylint: disable=unused-argument
+    def fn(prompts, completions, answer, question):  # pylint: disable=unused-argument
       return [2.5 for _ in completions]
 
     score_sum, count = evaluate_rl._compute_row_reward(
         reward_fns=[fn],
         prompt="p",
+        question="q",
         responses=["abc"],
         answer="gold",
         row_idx=0,
@@ -247,6 +249,7 @@ def test_sums_across_reward_fns_for_single_response(self):
     score_sum, count = evaluate_rl._compute_row_reward(
         reward_fns=self._two_fns(),
         prompt="p",
+        question="q",
         responses=["abcd"],
         answer="gold",
         row_idx=0,
@@ -261,6 +264,7 @@ def test_sums_across_passes_for_multiple_responses(self):
     score_sum, count = evaluate_rl._compute_row_reward(
         reward_fns=self._two_fns(),
         prompt="p",
+        question="q",
         responses=["a", "bcd", "ef"],
         answer="gold",
         row_idx=0,
@@ -276,6 +280,7 @@ def test_empty_responses_returns_zero_and_zero_count(self):
     score_sum, count = evaluate_rl._compute_row_reward(
         reward_fns=self._two_fns(),
         prompt="p",
+        question="q",
         responses=[],
         answer="gold",
         row_idx=0,
@@ -293,13 +298,65 @@ def _boom(**kwargs):  # pylint: disable=unused-argument
     score_sum, count = evaluate_rl._compute_row_reward(
         reward_fns=[_boom],
         prompt="p",
+        question="q",
         responses=["abc"],
         answer="gold",
         row_idx=0,
     )
     self.assertEqual(score_sum, 0.0)
     self.assertEqual(count, 0)  # zero count so the caller's mean isn't biased
 
+  @pytest.mark.cpu_only
+  def test_question_is_forwarded_to_reward_fn(self):
+    """Regression: helper must pass `question` through to reward fns.
+
+    The built-in `check_numbers` reward reads `kwargs["question"]`; if the
+    helper omits it, every eval row produces `KeyError('question')` and
+    `mean_reward` collapses to 0.0.
+    """
+    received = {}
+
+    def fn(prompts, completions, answer, question):  # pylint: disable=unused-argument
+      received["question"] = question
+      return [1.0 for _ in completions]
+
+    evaluate_rl._compute_row_reward(
+        reward_fns=[fn],
+        prompt="p",
+        question="What is 2+2?",
+        responses=["abc"],
+        answer="gold",
+        row_idx=0,
+    )
+    self.assertEqual(received["question"], "What is 2+2?")
+
+  @pytest.mark.cpu_only
+  def test_integrates_with_real_check_numbers_reward(self):
+    """End-to-end: real `check_numbers` from utils_rl must not raise on the
+    eval-time kwargs the helper passes (regression for the original
+    `KeyError('question')` failure mode in production)."""
+    from maxtext.trainers.post_train.rl import utils_rl  # pylint: disable=import-outside-toplevel
+
+    config = _make_config(eval_mode="pass")
+
+    # `check_numbers` takes tmvp_config positionally via partial; mirror what
+    # train_rl.py's make_reward_fn does.
+    def wrapped_check_numbers(**kwargs):
+      return utils_rl.check_numbers(tmvp_config=config, **kwargs)
+
+    # A correct-answer response should yield a non-zero score (proves the
+    # kwargs all reached the inside of check_numbers).
+    score_sum, count = evaluate_rl._compute_row_reward(
+        reward_fns=[wrapped_check_numbers],
+        prompt="solve: 2+2",
+        question="What is 2+2?",
+        responses=["<reasoning>2+2=4</reasoning><answer>4</answer>"],
+        answer='["4"]',  # json-encoded list of acceptable answers
+        row_idx=0,
+    )
+    self.assertEqual(count, 1)
+    self.assertGreater(score_sum, 0.0)
+
 
 if __name__ == "__main__":
   unittest.main()