Reward Functions: Maths, Code Tests, Format and Length

By the end of this lesson you will be able to write a reward function in the shape TRL’s GRPOTrainer expects; extract an answer from a completion without being fooled by the rest of it; run generated code against tests in a subprocess with limits, and say precisely what those limits do and do not protect; combine correctness, format and length terms with weights that mean something; recognise four specific ways a model scores without doing the task; and measure a reward against known-good and known-bad completions before you train with it.

This is the lesson where reinforcement learning stops being a paper and becomes a Python file you are responsible for. The reward is the only statement of what you want. If it is wrong, the run will be a very efficient search for your mistake.

The shape TRL expects

A reward function takes keyword arguments and returns one float per completion. TRL passes completions, prompts, completion_ids, the trainer state, and - this is the useful part - “any additional dataset columns”. So a dataset row carrying an answer column arrives in the reward function as an answer keyword holding the list of answers for the batch.

Fragment — not complete on its own

def numeric_reward(completions, answer, **kwargs):
    """One float per completion; extra dataset columns arrive as keyword arguments."""
    return [1.0 if extract(c) == want else 0.0 for c, want in zip(completions, answer)]

Two details cost people an afternoon each. Completions arrive as strings for a standard dataset and as message lists for a conversational one, so anything that calls .split() on them breaks the moment you switch formats. And **kwargs is not optional politeness: TRL passes arguments your function does not declare, and a function without it raises a TypeError on the first step.

Several functions can run at once. GRPOConfig takes reward_weights, and the trainer logs reward as “the overall average reward after summing rewards across functions (weighted by reward_weights)”. Those weights are your statement of priority, and getting them wrong is the first reward hack in the list below.

Getting the answer out

Before anything can be scored, the answer has to be found. This is the step that looks trivial and is not.

Three markers cover most of the field. GSM8K’s own solutions end the reasoning and then give the final answer after ####, which its dataset card documents as the marker; competition maths papers use \boxed{}; and a course that controls its own prompt can simply ask for Answer: N. Take the last match, because a model that revises its answer leaves the earlier one in the text.

Then there is the fallback: if no marker is found, take the last number in the completion. It is convenient, it rescues correct answers written in prose, and it is the single most exploitable line in the file. The reward library here makes it an argument that defaults to off in training and explains why in its docstring.

✕Common mistake

Comparing answers as strings makes 15 and 15.0 different, and 1,250 and 1250 different. Parse to a number and compare with an absolute tolerance. The tests in test-rewards.py cover exactly these cases, because every one of them has silently halved somebody’s measured accuracy.

The reward library

Everything this part uses lives in one file, with its tests beside it.

RunnableAll tracks

rewards.py

"""Reward functions a program can compute, in the shape TRL's GRPOTrainer expects.

Purpose: the course's reward library. Every function here takes the keyword arguments
    TRL passes to a reward function - `completions` plus any extra columns of the
    dataset - and returns one float per completion. Nothing here asks a model for an
    opinion: each reward is arithmetic, a regular expression or a subprocess that
    either passes or fails, so the same completion always scores the same.
Platform: all (pure Python; the unit-test reward needs a POSIX system for the
    resource limits, so on Windows run it inside WSL2)
Minimum memory: 8 GB
Assumes: Python 3.10 or newer. No third-party packages. `test-rewards.py` sits beside
    this file and is the proof that the functions do what the lesson says they do.

Usage: imported by train-grpo.py and eval-pass-at-1.py:
           from rewards import build_reward_functions, extract_final_answer
       run directly to score the built-in worked examples, including the
       reward-hacking ones, and print a table:
           python3 rewards.py --demo
           python3 rewards.py --demo --show-sandbox
"""

from __future__ import annotations

import argparse
import json
import math
import os
import re
import subprocess
import sys
import tempfile
from typing import Any, Callable, Iterable, Optional, Sequence

# --------------------------------------------------------------------------------------
# Reading an answer out of a completion
# --------------------------------------------------------------------------------------

# Three ways a model marks its final answer, in the order this course prefers them.
# GSM8K's own solutions end with "#### 18", the maths literature uses \boxed{18}, and
# the course's own format reward asks for "Answer: 18".
FINAL_PATTERNS = (
    re.compile(r"####\s*(-?[0-9][0-9,]*(?:\.[0-9]+)?)"),
    re.compile(r"\\boxed\{\s*(-?[0-9][0-9,]*(?:\.[0-9]+)?)\s*\}"),
    re.compile(r"(?i)\banswer\s*[:=]\s*\$?(-?[0-9][0-9,]*(?:\.[0-9]+)?)"),
)
ANY_NUMBER = re.compile(r"-?[0-9][0-9,]*(?:\.[0-9]+)?")

# The two-block shape the format reward asks for. The working comes first inside
# <think> tags, then one line beginning "Answer:" with nothing after the number.
FORMAT_PATTERN = re.compile(
    r"^\s*<think>.*?</think>\s*Answer:\s*-?[0-9][0-9,]*(?:\.[0-9]+)?\s*$",
    re.DOTALL,
)


def text_of(completion: Any) -> str:
    """TRL hands completions as strings or as message lists; this flattens both."""
    if isinstance(completion, str):
        return completion
    if isinstance(completion, list):
        return "".join(str(m.get("content", "")) for m in completion if isinstance(m, dict))
    if isinstance(completion, dict):
        return str(completion.get("content", ""))
    return str(completion)


def extract_final_answer(text: str, fallback_to_last_number: bool = True) -> Optional[str]:
    """The number the completion is claiming as its answer, or None.

    Marked answers win over unmarked ones. The fallback to "the last number in the
    text" is convenient and slightly dangerous: a model that ends with a page of
    working can score by accident. The lesson says to turn it off once the format
    reward has taught the model to mark its answer.
    """
    for pattern in FINAL_PATTERNS:
        matches = pattern.findall(text)
        if matches:
            return matches[-1].replace(",", "")
    if fallback_to_last_number:
        numbers = ANY_NUMBER.findall(text)
        if numbers:
            return numbers[-1].replace(",", "")
    return None


def as_number(value: Any) -> Optional[float]:
    """Parse a string or number into a float, or None if it is not one."""
    if value is None:
        return None
    if isinstance(value, (int, float)) and not isinstance(value, bool):
        return float(value)
    text = str(value).strip().replace(",", "").rstrip(".")
    if text.startswith("$"):
        text = text[1:]
    try:
        return float(text)
    except ValueError:
        return None


# --------------------------------------------------------------------------------------
# Exact match and numeric rewards
# --------------------------------------------------------------------------------------


def exact_match_reward(completions, answer, **kwargs) -> list[float]:
    """1.0 when the extracted answer string equals the reference string, else 0.0.

    Strict, cheap and unforgiving: "18.0" does not equal "18", and neither does
    "the answer is 18" if the extractor picked up something else. Use it when the
    reference answers are canonical strings and you want no tolerance at all.
    """
    out = []
    for completion, reference in zip(completions, answer):
        found = extract_final_answer(text_of(completion))
        out.append(1.0 if found is not None and found == str(reference).strip() else 0.0)
    return out


def numeric_reward(tolerance: float = 1e-6, fallback_to_last_number: bool = True):
    """1.0 when the extracted number equals the reference within a tolerance.

    This is the reward the maths lab trains against. It fixes the "18.0" problem
    above by comparing numbers rather than strings, and it is still a program: no
    judgement, no model, no ambiguity about what it will return.
    """

    def reward(completions, answer, **kwargs) -> list[float]:
        out = []
        for completion, reference in zip(completions, answer):
            want = as_number(reference)
            got = as_number(extract_final_answer(text_of(completion), fallback_to_last_number))
            if want is None or got is None:
                out.append(0.0)
            else:
                out.append(1.0 if math.isclose(got, want, rel_tol=0.0, abs_tol=tolerance) else 0.0)
        return out

    reward.__name__ = "numeric_reward"
    return reward


# --------------------------------------------------------------------------------------
# Format and length rewards
# --------------------------------------------------------------------------------------


def format_reward(pattern: re.Pattern[str] = FORMAT_PATTERN, value: float = 0.2):
    """A small reward for answering in the shape the task asked for.

    Keep it small. A format reward is a hint about where to put the answer, not a
    goal in itself; if it is worth as much as being right, the model will learn to
    produce beautifully formatted wrong answers, which is the first reward-hacking
    example in the lesson.
    """

    def reward(completions, **kwargs) -> list[float]:
        return [value if pattern.match(text_of(c)) else 0.0 for c in completions]

    reward.__name__ = "format_reward"
    return reward


def length_reward(target_tokens: int = 256, penalty: float = 0.2, tokens_per_word: float = 1.35):
    """A penalty that grows once a completion runs past a soft length budget.

    Zero up to the budget, then linear to -`penalty` at twice the budget, then flat.
    Two reasons to have it: GRPO's own objective has a documented length bias (the
    Dr. GRPO paper), and a completion that runs to the generation cap is truncated,
    so its answer never appears and it scores zero for a reason that has nothing to
    do with reasoning. The word count is an approximation of the token count; the
    ratio is a rough constant for English prose, not a measurement.
    """

    def reward(completions, **kwargs) -> list[float]:
        out = []
        for completion in completions:
            words = len(text_of(completion).split())
            approx_tokens = words * tokens_per_word
            over = max(0.0, approx_tokens - target_tokens) / max(1.0, float(target_tokens))
            out.append(-penalty * min(1.0, over))
        return out

    reward.__name__ = "length_reward"
    return reward


# --------------------------------------------------------------------------------------
# Unit tests as a reward, in a subprocess with limits
# --------------------------------------------------------------------------------------

CODE_BLOCK = re.compile(r"```(?:python)?\s*\n(.*?)```", re.DOTALL)

# The child process is started with this preamble so that a submission cannot pass by
# printing the expected output: the tests run after the candidate code, in the same
# namespace, and the harness prints one line the parent parses.
RUNNER_TEMPLATE = """\
import resource, sys
resource.setrlimit(resource.RLIMIT_CPU, ({cpu}, {cpu}))
resource.setrlimit(resource.RLIMIT_AS, ({mem}, {mem}))
resource.setrlimit(resource.RLIMIT_FSIZE, ({fsize}, {fsize}))
resource.setrlimit(resource.RLIMIT_NPROC, ({nproc}, {nproc}))
sys.stdout = open(__import__("os").devnull, "w")
namespace = {{}}
exec(compile(open({candidate!r}).read(), "candidate.py", "exec"), namespace)
exec(compile(open({tests!r}).read(), "tests.py", "exec"), namespace)
sys.stderr.write("COURSE-TESTS-PASSED\\n")
"""


def extract_code(text: str) -> str:
    """The last fenced Python block, or the whole completion if there is no fence."""
    blocks = CODE_BLOCK.findall(text)
    return blocks[-1] if blocks else text


def run_tests_sandboxed(
    code: str,
    tests: str,
    timeout_s: int = 5,
    cpu_seconds: int = 5,
    memory_mb: int = 512,
    max_output_bytes: int = 1 << 20,
    max_processes: int = 64,
) -> tuple[bool, str]:
    """Run `tests` against `code` in a separate interpreter with resource limits.

    Returns (passed, reason). The child runs with -I (isolated: no user site
    directory, no PYTHON* environment variables) and -S, in a fresh temporary
    directory, with a wall-clock timeout on top of the CPU limit, so an infinite
    loop, a runaway allocation, a fork bomb and a program that writes a large file
    are all bounded.

    What this does NOT do, and the lesson says so on the page: it is not a security
    boundary. The child runs as your user, on your filesystem, with your network.
    It can read any file you can read, open a socket, and delete your work. It stops
    accidents and cheap denial of service; it does not stop a program written to do
    harm. Part 24 builds the container-and-dedicated-user sandbox that an agent
    running untrusted code needs.
    """
    if os.name != "posix":
        return False, "resource limits need a POSIX system; run this inside WSL2 on Windows"

    with tempfile.TemporaryDirectory(prefix="course-reward-") as workdir:
        candidate_path = os.path.join(workdir, "candidate.py")
        tests_path = os.path.join(workdir, "tests.py")
        runner_path = os.path.join(workdir, "runner.py")
        with open(candidate_path, "w", encoding="utf-8") as handle:
            handle.write(code)
        with open(tests_path, "w", encoding="utf-8") as handle:
            handle.write(tests)
        with open(runner_path, "w", encoding="utf-8") as handle:
            handle.write(
                RUNNER_TEMPLATE.format(
                    cpu=cpu_seconds,
                    mem=memory_mb * 1024 * 1024,
                    fsize=max_output_bytes,
                    nproc=max_processes,
                    candidate=candidate_path,
                    tests=tests_path,
                )
            )
        try:
            done = subprocess.run(
                [sys.executable, "-I", "-S", runner_path],
                cwd=workdir,
                env={"PATH": os.environ.get("PATH", ""), "HOME": workdir},
                capture_output=True,
                text=True,
                timeout=timeout_s,
                check=False,
            )
        except subprocess.TimeoutExpired:
            return False, f"timed out after {timeout_s} s"
        except OSError as exc:
            return False, f"could not start the sandbox: {exc}"

    if "COURSE-TESTS-PASSED" in (done.stderr or ""):
        return True, "passed"
    tail = (done.stderr or "").strip().splitlines()
    return False, (tail[-1][:200] if tail else f"exit code {done.returncode}")


def unit_test_reward(timeout_s: int = 5, memory_mb: int = 512, partial_credit: bool = False):
    """1.0 when the completion's code passes the task's tests, else 0.0.

    The tests come from the dataset: each row carries a `tests` column holding the
    assertions to run against the candidate. `partial_credit` splits the test source
    on blank lines and scores the fraction of blocks that pass, which gives GRPO a
    gradient on tasks where all-or-nothing rewards make every rollout in a group
    score zero and the group's advantages collapse to nothing.
    """

    def reward(completions, tests, **kwargs) -> list[float]:
        out = []
        for completion, test_source in zip(completions, tests):
            code = extract_code(text_of(completion))
            if partial_credit:
                blocks = [b for b in re.split(r"\n\s*\n", str(test_source)) if b.strip()]
                if not blocks:
                    out.append(0.0)
                    continue
                passed = sum(1 for b in blocks if run_tests_sandboxed(code, b, timeout_s, memory_mb=memory_mb)[0])
                out.append(passed / len(blocks))
            else:
                ok, _ = run_tests_sandboxed(code, str(test_source), timeout_s, memory_mb=memory_mb)
                out.append(1.0 if ok else 0.0)
        return out

    reward.__name__ = "unit_test_reward"
    return reward


# --------------------------------------------------------------------------------------
# Combining rewards, and measuring one before training with it
# --------------------------------------------------------------------------------------


def build_reward_functions(
    kind: str = "maths",
    use_format: bool = True,
    use_length: bool = True,
    target_tokens: int = 256,
    tolerance: float = 1e-6,
    fallback_to_last_number: bool = False,
) -> tuple[list[Callable[..., list[float]]], list[float]]:
    """The reward list and weights the training script passes to GRPOTrainer.

    TRL sums the functions after weighting them with `reward_weights`, so the
    weights are the statement of what matters: correctness dominates, format is a
    hint, length is a brake. Returning them together keeps the two lists the same
    length, which TRL requires.
    """
    functions: list[Callable[..., list[float]]] = []
    weights: list[float] = []

    if kind == "maths":
        functions.append(numeric_reward(tolerance, fallback_to_last_number))
        weights.append(1.0)
    elif kind == "exact":
        functions.append(exact_match_reward)
        weights.append(1.0)
    elif kind == "code":
        functions.append(unit_test_reward())
        weights.append(1.0)
    else:
        raise ValueError(f"unknown reward kind {kind!r}; expected maths, exact or code")

    if use_format:
        functions.append(format_reward())
        weights.append(1.0)
    if use_length:
        functions.append(length_reward(target_tokens))
        weights.append(1.0)
    return functions, weights


def score(functions: Sequence[Callable[..., list[float]]],
          weights: Sequence[float],
          completions: Sequence[Any],
          **columns: Any) -> list[float]:
    """Apply every reward function and return the weighted sum, one per completion."""
    totals = [0.0] * len(completions)
    for function, weight in zip(functions, weights):
        values = function(completions=list(completions), **columns)
        for i, value in enumerate(values):
            totals[i] += weight * float(value)
    return totals


def evaluate_reward(functions: Sequence[Callable[..., list[float]]],
                    weights: Sequence[float],
                    cases: Iterable[dict]) -> dict:
    """Score a labelled set of completions and report whether the reward separates them.

    Each case is {"completion": ..., "should_score": "high"|"low", extra columns...}.
    A reward you have not run against known-good and known-bad answers is a reward
    you are about to train against blind: the gap between the two means is the
    signal GRPO has to work with, and any overlap is a way for the model to score
    without doing the task.
    """
    cases = list(cases)
    rows = []
    for case in cases:
        columns = {k: [v] for k, v in case.items() if k not in ("completion", "should_score", "note")}
        total = score(functions, weights, [case["completion"]], **columns)[0]
        rows.append({"note": case.get("note", ""), "expected": case["should_score"], "reward": round(total, 4)})

    highs = [r["reward"] for r in rows if r["expected"] == "high"]
    lows = [r["reward"] for r in rows if r["expected"] == "low"]
    mean = lambda xs: sum(xs) / len(xs) if xs else float("nan")  # noqa: E731 - one line, read once
    overlap = [r for r in rows if r["expected"] == "low" and highs and r["reward"] >= min(highs)]
    return {
        "rows": rows,
        "mean_high": round(mean(highs), 4),
        "mean_low": round(mean(lows), 4),
        "separation": round(mean(highs) - mean(lows), 4),
        "overlapping_low_cases": [r["note"] for r in overlap],
    }


# --------------------------------------------------------------------------------------
# Worked examples, including the ones that cheat
# --------------------------------------------------------------------------------------

DEMO_CASES = [
    {
        "note": "correct, in the requested format",
        "completion": "<think>Six boxes of four pencils is 24 pencils. Nine are used, so 15 remain.</think>\nAnswer: 15",
        "answer": "15",
        "should_score": "high",
    },
    {
        "note": "correct answer, no format",
        "completion": "Six times four is twenty-four, minus nine leaves fifteen. The answer is 15.",
        "answer": "15",
        "should_score": "high",
    },
    {
        "note": "wrong answer, perfect format (the format-reward hack)",
        "completion": "<think>Six boxes of four pencils is 24 pencils.</think>\nAnswer: 24",
        "answer": "15",
        "should_score": "low",
    },
    {
        "note": "every number in sight, hoping one is right (the last-number hack)",
        "completion": "Maybe 6, or 4, or 24, or 9, or 13, or 14, or 15",
        "answer": "15",
        "should_score": "low",
    },
    {
        "note": "right answer buried in a wall of text (the length hack)",
        "completion": "<think>" + ("Let me reconsider the problem once more. " * 120) + "</think>\nAnswer: 15",
        "answer": "15",
        "should_score": "low",
    },
    {
        "note": "refuses to answer",
        "completion": "<think>This needs more information.</think>\nAnswer: 0",
        "answer": "15",
        "should_score": "low",
    },
]

SANDBOX_CASES = [
    {
        "note": "a correct function",
        "completion": "```python\ndef add(a, b):\n    return a + b\n```",
        "tests": "assert add(2, 2) == 4\n\nassert add(-1, 1) == 0\n",
        "should_score": "high",
    },
    {
        "note": "an off-by-one function",
        "completion": "```python\ndef add(a, b):\n    return a + b + 1\n```",
        "tests": "assert add(2, 2) == 4\n\nassert add(-1, 1) == 0\n",
        "should_score": "low",
    },
    {
        "note": "an infinite loop (stopped by the timeout)",
        "completion": "```python\ndef add(a, b):\n    while True:\n        pass\n```",
        "tests": "assert add(2, 2) == 4\n",
        "should_score": "low",
    },
    {
        "note": "prints the expected output instead of computing it",
        "completion": "```python\nprint('all tests passed')\ndef add(a, b):\n    return None\n```",
        "tests": "assert add(2, 2) == 4\n",
        "should_score": "low",
    },
]


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--demo", action="store_true", help="score the built-in worked examples")
    parser.add_argument("--show-sandbox", action="store_true",
                        help="also run the unit-test reward on four candidate functions")
    parser.add_argument("--target-tokens", type=int, default=256)
    parser.add_argument("--json", action="store_true", help="print the report as JSON")
    args = parser.parse_args()

    if not args.demo:
        print(__doc__)
        return

    functions, weights = build_reward_functions(
        kind="maths", target_tokens=args.target_tokens, fallback_to_last_number=True
    )
    report = evaluate_reward(functions, weights, DEMO_CASES)

    if args.show_sandbox:
        code_functions, code_weights = build_reward_functions(kind="code", use_format=False, use_length=False)
        report["sandbox"] = evaluate_reward(code_functions, code_weights, SANDBOX_CASES)

    if args.json:
        print(json.dumps(report, indent=2))
        return

    print(f"{'reward':>8}  {'expected':<9} case")
    for row in report["rows"]:
        print(f"{row['reward']:>8.3f}  {row['expected']:<9} {row['note']}")
    print(f"\nmean reward, should score high: {report['mean_high']}")
    print(f"mean reward, should score low:  {report['mean_low']}")
    print(f"separation: {report['separation']}")
    if report["overlapping_low_cases"]:
        print("\nThese cases scored at least as well as the worst good answer:")
        for note in report["overlapping_low_cases"]:
            print(f"  - {note}")
        print("Each one is a way for the model to score without doing the task.")

    if "sandbox" in report:
        print("\nunit-test reward:")
        for row in report["sandbox"]["rows"]:
            print(f"{row['reward']:>8.3f}  {row['expected']:<9} {row['note']}")


if __name__ == "__main__":
    main()

Download rewards.py518 lines

RunnableAll tracks

test-rewards.py

"""Tests for the reward functions, because a reward with a bug trains the bug in.

Purpose: prove that every function in rewards.py returns what the lesson says it
    returns, including on the completions that try to cheat it. A reward function is
    the only thing standing between a reinforcement-learning run and hours of
    optimising the wrong quantity, and it is ordinary code with ordinary bugs.
Platform: all (pure Python and the standard library's unittest; the sandbox tests
    are skipped automatically on a non-POSIX system)
Minimum memory: 8 GB
Assumes: Python 3.10 or newer, and rewards.py in the same directory.

Usage: python3 test-rewards.py
       python3 test-rewards.py -v        # one line per test
"""

from __future__ import annotations

import os
import unittest

import rewards


class TestAnswerExtraction(unittest.TestCase):
    def test_gsm8k_marker_wins(self):
        self.assertEqual(rewards.extract_final_answer("working 3 + 4\n#### 7"), "7")

    def test_boxed_marker(self):
        self.assertEqual(rewards.extract_final_answer("so \\boxed{42} is it"), "42")

    def test_answer_label(self):
        self.assertEqual(rewards.extract_final_answer("Answer: -12"), "-12")

    def test_thousands_separator_is_stripped(self):
        self.assertEqual(rewards.extract_final_answer("Answer: 1,250"), "1250")

    def test_currency_prefix_is_ignored(self):
        self.assertEqual(rewards.extract_final_answer("Answer: $18"), "18")

    def test_last_marked_answer_wins_over_earlier_ones(self):
        self.assertEqual(rewards.extract_final_answer("Answer: 3\nno wait\nAnswer: 5"), "5")

    def test_fallback_takes_the_last_number(self):
        self.assertEqual(rewards.extract_final_answer("6 times 4 is 24 minus 9 is 15"), "15")

    def test_fallback_can_be_disabled(self):
        self.assertIsNone(rewards.extract_final_answer("no marker here, just 15", fallback_to_last_number=False))

    def test_no_number_at_all(self):
        self.assertIsNone(rewards.extract_final_answer("I do not know"))


class TestNumericReward(unittest.TestCase):
    def setUp(self):
        self.reward = rewards.numeric_reward()

    def test_correct_answer_scores_one(self):
        self.assertEqual(self.reward(completions=["Answer: 15"], answer=["15"]), [1.0])

    def test_decimal_and_integer_forms_agree(self):
        self.assertEqual(self.reward(completions=["Answer: 15.0"], answer=["15"]), [1.0])

    def test_wrong_answer_scores_zero(self):
        self.assertEqual(self.reward(completions=["Answer: 16"], answer=["15"]), [0.0])

    def test_missing_answer_scores_zero(self):
        self.assertEqual(self.reward(completions=["I would rather not"], answer=["15"]), [0.0])

    def test_tolerance_is_absolute_and_respected(self):
        loose = rewards.numeric_reward(tolerance=0.01)
        self.assertEqual(loose(completions=["Answer: 15.005"], answer=["15"]), [1.0])
        self.assertEqual(loose(completions=["Answer: 15.5"], answer=["15"]), [0.0])

    def test_batch_is_scored_elementwise(self):
        got = self.reward(completions=["Answer: 1", "Answer: 2", "Answer: 4"], answer=["1", "3", "4"])
        self.assertEqual(got, [1.0, 0.0, 1.0])


class TestExactMatchReward(unittest.TestCase):
    def test_string_equality_is_strict(self):
        self.assertEqual(rewards.exact_match_reward(completions=["Answer: 15"], answer=["15"]), [1.0])
        self.assertEqual(rewards.exact_match_reward(completions=["Answer: 15.0"], answer=["15"]), [0.0])


class TestFormatReward(unittest.TestCase):
    def setUp(self):
        self.reward = rewards.format_reward()

    def test_the_requested_shape_scores(self):
        good = "<think>six fours are twenty-four</think>\nAnswer: 15"
        self.assertEqual(self.reward(completions=[good]), [0.2])

    def test_trailing_commentary_does_not_score(self):
        bad = "<think>working</think>\nAnswer: 15\nHope that helps!"
        self.assertEqual(self.reward(completions=[bad]), [0.0])

    def test_missing_working_block_does_not_score(self):
        self.assertEqual(self.reward(completions=["Answer: 15"]), [0.0])

    def test_a_wrong_answer_in_the_right_shape_still_scores(self):
        # This is the point of the reward-hacking example: on its own, the format
        # reward pays for presentation and knows nothing about correctness.
        self.assertEqual(self.reward(completions=["<think>x</think>\nAnswer: 999"]), [0.2])


class TestLengthReward(unittest.TestCase):
    def setUp(self):
        self.reward = rewards.length_reward(target_tokens=100, penalty=0.2)

    def test_short_completions_are_not_penalised(self):
        self.assertEqual(self.reward(completions=["a b c"]), [0.0])

    def test_the_penalty_saturates(self):
        very_long = " ".join(["word"] * 1000)
        self.assertAlmostEqual(self.reward(completions=[very_long])[0], -0.2)

    def test_the_penalty_is_monotone_in_length(self):
        short = self.reward(completions=[" ".join(["word"] * 100)])[0]
        longer = self.reward(completions=[" ".join(["word"] * 150)])[0]
        self.assertLessEqual(longer, short)


class TestCombinedReward(unittest.TestCase):
    def setUp(self):
        self.functions, self.weights = rewards.build_reward_functions(
            kind="maths", target_tokens=100, fallback_to_last_number=True
        )

    def test_correct_and_well_formatted_beats_correct_alone(self):
        formatted = rewards.score(self.functions, self.weights,
                                  ["<think>w</think>\nAnswer: 15"], answer=["15"])[0]
        plain = rewards.score(self.functions, self.weights, ["the answer is 15"], answer=["15"])[0]
        self.assertGreater(formatted, plain)

    def test_correctness_outweighs_format(self):
        right_ugly = rewards.score(self.functions, self.weights, ["the answer is 15"], answer=["15"])[0]
        wrong_pretty = rewards.score(self.functions, self.weights,
                                     ["<think>w</think>\nAnswer: 24"], answer=["15"])[0]
        self.assertGreater(right_ugly, wrong_pretty)

    def test_the_demo_set_rewards_good_answers_more_on_average(self):
        report = rewards.evaluate_reward(self.functions, self.weights, rewards.DEMO_CASES)
        self.assertGreater(report["mean_high"], report["mean_low"])

    def test_the_demo_set_still_lets_two_hacks_through(self):
        # The whole point of the worked examples: this reward is better on average
        # and it is not safe. Two bad completions score as well as a good one, and
        # the test says so rather than the page hoping the reader noticed.
        report = rewards.evaluate_reward(self.functions, self.weights, rewards.DEMO_CASES)
        self.assertEqual(len(report["overlapping_low_cases"]), 2)

    def test_turning_off_the_last_number_fallback_closes_one_of_them(self):
        functions, weights = rewards.build_reward_functions(
            kind="maths", target_tokens=100, fallback_to_last_number=False
        )
        spray = "Maybe 6, or 4, or 24, or 9, or 13, or 14, or 15"
        self.assertEqual(rewards.score(functions, weights, [spray], answer=["15"])[0], 0.0)

    def test_the_same_change_also_zeroes_a_correct_but_unmarked_answer(self):
        # The price of closing that hole, and the reason the format reward exists:
        # with no fallback, an answer the extractor cannot find is an answer that
        # did not happen, however right it was.
        functions, weights = rewards.build_reward_functions(
            kind="maths", target_tokens=100, fallback_to_last_number=False
        )
        marked = "Six times four is twenty-four, minus nine. Answer: 15"
        self.assertEqual(rewards.score(functions, weights, [marked], answer=["15"])[0], 1.0)
        unmarked = "Six times four is twenty-four, minus nine leaves 15."
        self.assertEqual(rewards.score(functions, weights, [unmarked], answer=["15"])[0], 0.0)


@unittest.skipUnless(os.name == "posix", "the sandbox uses POSIX resource limits")
class TestUnitTestReward(unittest.TestCase):
    def setUp(self):
        self.reward = rewards.unit_test_reward(timeout_s=10)

    def test_correct_code_passes(self):
        completion = "```python\ndef add(a, b):\n    return a + b\n```"
        self.assertEqual(self.reward(completions=[completion], tests=["assert add(2, 2) == 4\n"]), [1.0])

    def test_incorrect_code_fails(self):
        completion = "```python\ndef add(a, b):\n    return a - b\n```"
        self.assertEqual(self.reward(completions=[completion], tests=["assert add(2, 2) == 4\n"]), [0.0])

    def test_an_infinite_loop_is_stopped(self):
        completion = "```python\ndef add(a, b):\n    while True:\n        pass\n```"
        self.assertEqual(self.reward(completions=[completion], tests=["assert add(2, 2) == 4\n"]), [0.0])

    def test_printing_the_answer_does_not_pass(self):
        completion = "```python\nprint('ok')\ndef add(a, b):\n    return None\n```"
        self.assertEqual(self.reward(completions=[completion], tests=["assert add(2, 2) == 4\n"]), [0.0])

    def test_code_outside_a_fence_is_still_run(self):
        self.assertEqual(self.reward(completions=["def add(a, b):\n    return a + b\n"],
                                     tests=["assert add(2, 2) == 4\n"]), [1.0])

    def test_partial_credit_scores_the_fraction_of_blocks(self):
        partial = rewards.unit_test_reward(timeout_s=10, partial_credit=True)
        completion = "```python\ndef add(a, b):\n    return abs(a) + abs(b)\n```"
        tests = ["assert add(2, 2) == 4\n\nassert add(-1, 1) == 0\n"]
        self.assertAlmostEqual(partial(completions=[completion], tests=tests)[0], 0.5)


class TestMessageShapedCompletions(unittest.TestCase):
    """TRL hands conversational datasets back as message lists, not strings."""

    def test_a_message_list_is_flattened(self):
        completion = [{"role": "assistant", "content": "Answer: 15"}]
        self.assertEqual(rewards.numeric_reward()(completions=[completion], answer=["15"]), [1.0])


if __name__ == "__main__":
    unittest.main()

Download test-rewards.py213 lines

RunnableAll tracks

run the tests before you trust the rewards

python3 test-rewards.py

Output — what you should see

....................................
----------------------------------------------------------------------
Ran 36 tests in 5.1s

OK

A reward function is ordinary code with ordinary bugs, and it is code whose bugs are invisible: a reward that returns 1.0 slightly too often produces a training curve that looks like success. Tests are the cheapest defence there is, and thirty-six of them run in about five seconds.

Unit tests as a reward, and what the sandbox is

For code tasks the reward is “did it pass the tests”, which means running text a model wrote. The implementation in rewards.py starts a fresh interpreter with -I and -S, in a temporary directory, with four limits set before the candidate code is reached and a wall-clock timeout on top.

The limits come from Python’s resource module, whose documented availability is “Unix, not WASI”:

• RLIMIT_CPU, “the maximum amount of processor time (in seconds) that a process can use”, which stops a busy loop;
• RLIMIT_AS, “the maximum area (in bytes) of address space which may be taken by the process”, which stops a runaway allocation;
• RLIMIT_FSIZE, “the maximum size of a file which the process may create”, which stops the disk filling;
• RLIMIT_NPROC, “the maximum number of processes the current process may create”, which stops a fork bomb.

On top of those, subprocess.run takes a timeout, and the documentation is precise about what happens: “if the timeout expires, the child process will be killed and waited for”, with the exception “re-raised after the child process has terminated”. The CPU limit handles compute; the wall-clock timeout handles a process that is asleep rather than busy.

⛨Security

This is a safety net, not a security boundary. The child runs as your user, on your filesystem, with your network. It can read any file you can read, open a socket, and delete your work. What the limits stop is accidental damage and cheap denial of service: infinite loops, memory exhaustion, disk filling, fork bombs. What they do not stop is code written to do harm.

That is acceptable here because the code comes from a small model you are training on arithmetic on your own machine. It is not acceptable for anything else. Part 24 builds the sandbox an agent needs

• a container or a dedicated user, with the filesystem and the network as the boundary rather than a resource limit - and that is the pattern to reach for the moment untrusted input is involved.

One more design decision in that function is worth naming. The tests run after the candidate code, in the same namespace, and the harness signals success on standard error rather than standard output, with the child’s standard output redirected to the null device. Without that, a model that prints “all tests passed” scores full marks. There is a test for it.

Format, length and combination

A format reward is a regular expression: does the completion have the shape the prompt asked for? Its job is to make the answer findable, which is what lets you turn off the last-number fallback. Keep it small, because it is paid on every completion and says nothing at all about correctness.

A length reward is a penalty that starts at a soft budget and grows. Two reasons to have one. The GRPO objective has its own documented length bias, described in the previous lesson, and a completion that hits the generation cap is truncated, so its answer never appears and it scores zero for a reason unrelated to reasoning.

Combining them is a statement of priorities, and the weights are the statement:

Term	Weight in the lab	What it says
Numeric correctness	1.0, and the reward is 1.0 or 0.0	Being right is the point
Format	1.0, and the reward is 0.2 or 0.0	Put the answer where it can be found
Length	1.0, and the reward is 0.0 down to -0.2	Do not pad, and do not run into the cap

The correctness term is worth five times the format term, and that ratio is the design. Make them equal and the cheapest way to score becomes formatting, since format is easy and arithmetic is not.

Four worked hacks

Run the reward library’s demonstration and read the table it prints. It scores six completions against a reward that a reasonable person would write on a first attempt.

RunnableAll tracks

score the worked examples, including the ones that cheat

python3 rewards.py --demo --show-sandbox

Output — what you should see

  reward  expected  case
   1.200  high      correct, in the requested format
   1.000  high      correct answer, no format
   0.200  low       wrong answer, perfect format (the format-reward hack)
   1.000  low       every number in sight, hoping one is right (the last-number hack)
   1.000  low       right answer buried in a wall of text (the length hack)
   0.200  low       refuses to answer

mean reward, should score high: 1.1
mean reward, should score low:  0.6
separation: 0.5

These cases scored at least as well as the worst good answer:
  - every number in sight, hoping one is right (the last-number hack)
  - right answer buried in a wall of text (the length hack)
Each one is a way for the model to score without doing the task.

Read that output carefully, because it is the lesson. The reward is better on average and it is not safe. Two bad completions score as well as a good one, and an optimiser given ten thousand attempts will find both.

Answer spraying. “Maybe 6, or 4, or 24, or 9, or 13, or 14, or 15” scores full marks because the last number happens to be right. Turning the fallback off closes it, at the price of scoring zero for a correct answer that was not marked. That price is why the format reward exists.

Format farming. A wrong answer in a perfect shape still collects the format term. Harmless at a weight of 0.2 against a correctness term of 1.0; a catastrophe if you raise it to 1.0 because the model kept forgetting the format.

Length inflation. The right answer after 120 repetitions of “let me reconsider” collects the correctness and format terms, and the length penalty saturates at -0.2. If you want length to be a real constraint, it has to be able to outweigh being right, and then you have to accept that the model will sometimes stop early rather than finish.

Printing the answer. In the sandbox demonstration, the candidate that prints “all tests passed” and defines a function returning None scores zero, because the harness discards the child’s standard output and reports through a channel the candidate cannot forge.

?Reality check

The claim to resist is that a verifiable reward cannot be hacked. It cannot be flattered, which is different. Skalse and colleagues define reward hacking as what happens when “optimizing an imperfect proxy reward function leads to poor performance according to the true reward function”, and a checker is a proxy like any other. What a program-computed reward buys you is that its failures are findable: you can enumerate them, write them down as test cases, and score them.

Measure the reward before you train with it

The cheapest hour in this whole part is the one spent scoring a labelled set of completions before starting a run.

Evaluating a reward before spending a GPU-hour on it

1. 1Write down what should score highCorrect answers in several styles: marked, unmarked, terse, worked through. If your reward only likes one of them, it is a format reward wearing a correctness label.
2. 2Write down what should score lowWrong answers, refusals, and one deliberate attempt at each hack you can think of. This is the part people skip.
3. 3Score both setsThe gap between the means is the signal GRPO has to work with. A small gap is a slow run; a gap of zero is no run at all.
4. 4Look for overlapAny low case scoring as well as the worst high case is a way to score without doing the task, and the optimiser will find it before you do.
5. 5Fix, and add a testEvery hack you close becomes a test in test-rewards.py, so the next change to the reward cannot quietly reopen it.

The demonstration in rewards.py is this loop, already written, with six maths cases and four code cases. Replace them with cases from your own task and the same four steps apply.

There is a second measurement worth making once the run starts: score the model’s actual rollouts, sort them by reward, and read the top five and the bottom five. The top five tell you what the model is being paid for, which is not always what you meant to pay for.

=Expected result

A reward ready to train with has three properties. Every correct completion in your labelled set scores higher than every incorrect one. The separation between the two means is large relative to the format and length terms. And every hack you could think of is a test case that passes. If any of the three is missing, fix the reward: it is minutes of work, against hours of optimising the wrong quantity.

Test the verifier as adversarial input-handling code

A verifier is part of the optimisation target, so the policy has repeated opportunities to exploit it. Test correct answers, wrong answers, malformed outputs, multiple answers, extra text and deliberately misleading strings before training. For arithmetic, parsing “the answer is 7, or perhaps 9” as correct merely because seven appears somewhere rewards ambiguity.

Define a strict answer extraction contract and separate parsing success from correctness. For executable tasks, run generated code in the lab’s isolated environment with resource limits and tests the policy cannot modify. A timeout is an outcome that needs explicit scoring, not a reason to drop the sample from the dataset.

Keep development tests for the verifier separate from held-out task evaluation. After training, inspect high-reward failures manually; these are especially informative because they show where the measured objective differs from the intended one. Do not improve the reported score by weakening the verifier after seeing failures. Version the verifier alongside the dataset and model so a reward curve retains its meaning across experiments.

Recap

A TRL reward function takes completions plus the dataset’s extra columns as keyword arguments and returns one float per completion, and it must accept **kwargs and handle message-list completions. Answer extraction is where most of the errors live: prefer an explicit marker, take the last match, compare as numbers with a tolerance, and treat the last-number fallback as a known hole. Unit tests as a reward mean running generated code, which belongs in a separate interpreter with CPU, address space, file size and process limits and a wall-clock timeout, and that combination is a safety net against accidents rather than a security boundary. Format and length terms shape where the answer goes and how long it may be, and their weights relative to correctness are the design. Four hacks are worth knowing by name: answer spraying, format farming, length inflation and printing the expected output. And the way to find the fifth is to score a labelled set before training and look for overlap.

Check your understanding

Question 1. A reward function raises TypeError on the first training step. What is the most likely cause?

• The dataset has the wrong columns
• The function does not accept **kwargs, and TRL passes arguments it does not declare
• The model is too small
• The reward weights do not sum to one

Show the answer and why

Answer: The function does not accept **kwargs, and TRL passes arguments it does not declare

TRL passes prompts, completion_ids, trainer state, logging callables and every extra dataset column. A function that declares only what it uses has to swallow the rest, which is what **kwargs is for.

Question 2. The reward extracts the last number in the completion when no answer marker is found. Why is that dangerous, and what makes it safe to switch off?

• It is slow; a faster regular expression fixes it
• A completion that lists many numbers can score by accident; switching it off is only safe once a format reward has taught the model to mark its answer
• It only works for integers
• It is not dangerous at all

Show the answer and why

Answer: A completion that lists many numbers can score by accident; switching it off is only safe once a format reward has taught the model to mark its answer

The two are a pair. Without the fallback a correct unmarked answer scores zero; without the format reward the fallback is a way to score without solving anything. The lab uses both and weights the format term low.

Question 3. What do the resource limits and the timeout around generated code actually protect against? Select all that apply.

• An infinite loop consuming the machine
• A runaway memory allocation
• A program that fills the disk
• A program written to read your files and send them somewhere

Show the answer and why

Answer: An infinite loop consuming the machine, A runaway memory allocation, A program that fills the disk

The first three are accidents and cheap denial of service, and the limits bound all of them. The fourth is a different problem: the child runs as your user with your network. That needs a container or a dedicated user, which is what Part 24 builds.

Question 4. You raise the format reward from 0.2 to 1.0 because the model keeps forgetting the format. What is the likely outcome?

• The model produces the format and stays as accurate as before
• Formatting becomes as valuable as being right, and since formatting is much easier, the model can score well by producing immaculate wrong answers
• Training becomes faster
• The KL divergence falls

Show the answer and why

Answer: Formatting becomes as valuable as being right, and since formatting is much easier, the model can score well by producing immaculate wrong answers

The weights are the specification. When the cheap term is worth as much as the expensive one, the optimiser takes the cheap one, and the reward curve will look excellent while accuracy does not move.

Question 5. Why does the unit-test harness report success on standard error with the child's standard output redirected away?

• To make the logs shorter
• So that a candidate which simply prints the expected output cannot forge a pass; the signal travels on a channel it does not control
• Because standard output is slower
• To support Windows

Show the answer and why

Answer: So that a candidate which simply prints the expected output cannot forge a pass; the signal travels on a channel it does not control

It is a small design decision that closes a whole class of cheating, and there is a test asserting it. Anywhere a model can write to the channel a grader reads, it will eventually write what the grader wants to see.

Check my answersClear