"""Chapter 14 lab: where rewards come from. You will build three pieces of real reward plumbing: Part A a math-answer verifier, the reward function of RLVR (Tasks 1-3) Part B a rubric aggregator, for tasks with no answer key (Task 4) Part C order-swapping, to cancel an LLM judge's position bias (Task 5) Do the tasks in order and run the tests for one task at a time: python -m pytest -q -k extract_boxed # Task 1 python -m pytest -q -k parse_number # Task 2 python -m pytest -q -k math_reward # Task 3 python -m pytest -q -k rubric # Task 4 python -m pytest -q -k "debiased or verdict or swapping" # Task 5 Then run `python run.py` to watch a tiny policy exploit a naive verifier but not yours, and to see what order-swapping can and cannot fix. Only the standard library is needed here (re, fractions). Do not use sympy or eval(): the point is to see exactly what a verifier accepts. """ from __future__ import annotations import re # noqa: F401 (you will probably want it) from fractions import Fraction from typing import Callable, Sequence from helpers import normalize_text # noqa: F401 (used in Task 3) # --------------------------------------------------------------------------- # Part A: a robust math-answer verifier # --------------------------------------------------------------------------- def extract_boxed(text: str) -> list[str]: r"""Task 1. Return the contents of every \boxed{...} in `text`, in order. Models are asked to put their final answer in \boxed{...}. The contents can contain braces themselves, e.g. \boxed{\frac{1}{2}}, so a regex like r"\\boxed\{(.*?)\}" is WRONG: it stops at the first "}" and returns "\frac{1". Instead, walk the string and count brace depth. Rules: * Allow whitespace between \boxed and "{" (e.g. "\boxed {7}"). * A backslash escapes the next character: "\{" and "\}" are not braces for counting purposes (so \boxed{\{1, 2\}} returns "\{1, 2\}"). * Strip leading/trailing whitespace from each extracted string. * If a \boxed{ is never closed, ignore it (and anything after it). * \boxed not followed by "{" is skipped. Examples: extract_boxed(r"so \boxed{\frac{1}{2}}.") -> [r"\frac{1}{2}"] extract_boxed(r"\boxed{3} or \boxed{ 4 }") -> ["3", "4"] extract_boxed("no box here") -> [] extract_boxed(r"\boxed{5") -> [] """ # TODO: find each "\boxed", skip spaces, check for "{", then scan forward # counting depth (+1 on "{", -1 on "}", skip the char after a backslash) # until depth returns to 0. raise NotImplementedError("Task 1: extract_boxed") def parse_number(s: str) -> Fraction | None: r"""Task 2. Parse a numeric answer into an exact Fraction, or return None. The same number can be written many ways. All of these must give Fraction(1, 2): "0.5", ".5", "0.50", "1/2", "2/4", r"\frac{1}{2}", r"\dfrac{1}{2}", r"\tfrac12", r"\frac12", "$0.5$", "x = 1/2", r"\left(\frac{1}{2}\right)". Apply these steps in order (each is one or two lines): 1. Strip whitespace and "$". Remove \left, \right, \displaystyle and the LaTeX spacing commands \! \, \; "\ ". Replace the unicode minus "−" with "-", and the LaTeX thousands separator "{,}" with ",". 2. Remove a trailing unit in \text{...}, \mbox{...} or \mathrm{...} (e.g. "5\text{ cm}"), degree signs "^\circ" / "^{\circ}", and percent signs "\%" / "%". (So "50\%" parses as 50: we compare numbers, and the gold answer is written the same way.) 3. If there is an "=", keep only what follows the LAST "=" ("x = 5" -> "5"). 4. Strip again, and drop one trailing period ("5." -> "5"). 5. Thousands separators: if the string looks like 1,234 or 12,345,678.5 (groups of exactly three digits), delete the commas. Leave "1,2" alone, so it fails to parse: it is probably a list, not a number. 6. If ONE matching pair of parentheses wraps the whole string, remove it: "(3)" -> "3", but "(1)/(2)" stays as it is (its first "(" closes early). 7. Turn \dfrac and \tfrac into \frac, then \frac{a}{b} into (a)/(b), and the shorthand \frac12 (two single digits) into (1)/(2). Remove all remaining whitespace. 8. What remains must be: an optional sign, then a number (integer or decimal, maybe in parentheses), optionally "/" and a second number (maybe in parentheses). Anything else -> None. Division by zero -> None. Fraction("0.125") gives the exact value 1/8, so use Fraction, not float. Returns None for non-numbers: "", "abc", r"\sqrt{2}", "x+1", "1,2", "1/0", "1.2.3". """ # TODO: steps 1-8. Regexes you may find useful: # r"\\frac\{([^{}]*)\}\{([^{}]*)\}" braced fraction # r"^[+-]?\d{1,3}(,\d{3})+(\.\d+)?$" a number with thousands separators # number: r"[+-]?(?:\d+\.?\d*|\.\d+)" raise NotImplementedError("Task 2: parse_number") def math_reward(response: str, gold: str, max_answers: int = 1) -> float: r"""Task 3. The verifier: 1.0 if `response` gives the gold answer, else 0.0. Algorithm: 1. candidates = extract_boxed(response). 2. If there are none, fall back to the group captured by the LAST match of r"answer\s*(?:is\s*:?|:)\s*([^\n]+?)\s*(?:\.(?:\s|$)|$)" with flags re.IGNORECASE | re.MULTILINE. It catches "Answer: 42", "The answer is 42." and "Final answer is: 42", and stops at the end of the sentence, so "The answer is 85. I'm 100% sure." gives "85" (the lazy +? and the "period followed by a space or the end" do that). If nothing matches either, return 0.0. 3. Anti-hedging: count the DISTINCT candidates, where two candidates are the same if they are the same answer (see `same` below): "\boxed{0.5} ... so \boxed{1/2}" is one distinct answer, not two. If there are more than `max_answers` distinct candidates, return 0.0. (Without this rule a policy learns to box every number it can think of.) 4. Return 1.0 if the LAST candidate is the same answer as `gold`, else 0.0. Two strings x, y are "the same answer" if: * both parse as numbers (parse_number) and the Fractions are equal, or * otherwise, normalize_text(x) == normalize_text(y) (for answers like "x^2+1"). Examples (gold = "1/2"): r"... \boxed{\frac{1}{2}}" -> 1.0 "So the answer is 0.5." -> 1.0 r"\boxed{0.5} or maybe \boxed{0.25}" -> 0.0 (hedging) r"\boxed{1/3}" -> 0.0 Example (gold = "57"): r"\boxed{157}" -> 0.0 (no substring matching!) """ # TODO: write a small helper same(x, y), then steps 1-4. # Tip: stop counting distinct answers as soon as you exceed max_answers # (a "spam" response can box hundreds of numbers). raise NotImplementedError("Task 3: math_reward") # --------------------------------------------------------------------------- # Part B: rubric rewards for tasks without an answer key # --------------------------------------------------------------------------- # Default weights per category, as in the "Rubrics as Rewards" paper (Gunjal et # al., 2025). Pitfall criteria are phrased positively ("avoids X"), so passing # one ADDS reward, and failing one withholds it. DEFAULT_WEIGHTS = {"essential": 1.0, "important": 0.7, "optional": 0.3, "pitfall": 0.9} def rubric_reward(rubric: Sequence[dict], checks: Sequence[float], gate_essential: bool = False) -> float: r"""Task 4. Combine per-criterion judgments into one scalar reward in [0, 1]. Args: rubric: a list of criteria, each a dict like {"criterion": "States the diagnosis", "category": "essential"} with category in DEFAULT_WEIGHTS. An item may instead carry its own "weight" (a positive number), which overrides the category default. checks: one number per criterion in [0, 1]: did the response meet it? (1 = yes, 0 = no; a judge may also give partial credit.) gate_essential: if True, return 0.0 whenever any ESSENTIAL criterion has a check below 1. (A design choice of ours, not from the paper: it stops a policy from collecting many small criteria while missing the point.) Returns the weighted average r = sum_j w_j c_j / sum_j w_j (the chapter's "explicit aggregation" formula). Raises ValueError if len(rubric) != len(checks) or the weights sum to <= 0. """ # TODO raise NotImplementedError("Task 4: rubric_reward") # --------------------------------------------------------------------------- # Part C: cancelling position bias by swapping the order # --------------------------------------------------------------------------- def debiased_preference(judge: Callable, prompt: str, a, b, n_samples: int = 1) -> float: r"""Task 5a. P(a is better than b), averaged over both presentation orders. `judge(prompt, first, second)` returns the judge's probability (or a sampled 0/1 verdict) that the answer shown FIRST is better. Call it n_samples times with a shown first and n_samples times with b shown first, average each, and combine: P_debiased(a > b) = 1/2 * [ P(a wins | a first) + (1 - P(b wins | b first)) ] This is the "balanced position calibration" of Wang et al. (2023) and what RLAIF (Lee et al., 2023) does with its two inferences. Any bonus the judge gives to the first slot is given once to a and once to b, so it cancels in the decision (the chapter proves this). """ # TODO raise NotImplementedError("Task 5a: debiased_preference") def consistent_verdict(judge: Callable, prompt: str, a, b, n_samples: int = 1) -> str: r"""Task 5b. MT-Bench's conservative rule: a win only if it wins in BOTH orders. Average n_samples judgments in each order, exactly as in Task 5a. Then: * return "a" if a is preferred (prob > 0.5) when shown first AND when shown second, * return "b" if b is preferred in both orders, * return "tie" otherwise (the verdict flipped with the order, or a prob was exactly 0.5). """ # TODO raise NotImplementedError("Task 5b: consistent_verdict")