"""Reward functions for agent tasks, in the shape TRL's GRPO trainer expects.
Purpose: Part 27's reward library. Part 14's rewards scored one answer; these score a
whole attempt at a task: whether the tool calls parse, whether their arguments match
the declared schema, whether the files the model wrote make a test suite pass, and
whether it did all that inside a step budget. Nothing here asks a model for an
opinion. Every function 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 test-suite reward needs a POSIX system for its resource
limits, so on Windows run it inside WSL2)
Minimum memory: 8 GB
Assumes: Python 3.10 or newer and no third-party packages. Every function takes the
keyword arguments TRL passes a reward function - `completions` plus the dataset's
own columns - and returns one float per completion.
Usage: run directly to score the built-in worked examples, including the two that try
to cheat the reward and the one that fails on format, and print a table:
python3 agent-rewards.py --demo
python3 agent-rewards.py --demo --show-sandbox
loaded by a GRPO training script that sits beside it, the same way Part 26's
harness loads an agent entry point (the file name has a hyphen in it, so it is
loaded by path rather than imported by name):
import importlib.util, pathlib
spec = importlib.util.spec_from_file_location(
"agent_rewards", pathlib.Path(__file__).with_name("agent-rewards.py"))
agent_rewards = importlib.util.module_from_spec(spec)
spec.loader.exec_module(agent_rewards)
functions, weights = agent_rewards.build_agent_reward_functions(
kind="tests", workspace="repo")
The file-writing protocol, which is what the test-suite reward reads, is one fence per
file:
```file:src/app.py
def total(rows):
return sum(r["amount"] for r in rows)
```
It is deliberately dull. A reward that has to interpret a unified diff spends its budget
on parsing rather than on measuring, and a patch that does not apply is scored as a
failure to write code rather than as a failure to solve the task.
"""
from __future__ import annotations
import argparse
import json
import os
import re
import shutil
import subprocess
import sys
import tempfile
from pathlib import Path
from typing import Any, Callable, Sequence
# ---------------------------------------------------------------------------------------
# Reading a completion
# ---------------------------------------------------------------------------------------
# One fenced block per file, as described in the module docstring.
FILE_BLOCK = re.compile(r"```file:([^\n`]+)\n(.*?)```", re.DOTALL)
# A tool call the model emitted as JSON inside a fence or as a bare object.
JSON_TOOL_CALL = re.compile(r"\{\s*\"name\"\s*:\s*\"([A-Za-z0-9_.-]+)\"\s*,\s*\"arguments\"\s*:",
re.DOTALL)
# Fragments that mean the model tried to call a tool and nothing parsed it back. Each one
# is a real format from a model family in this course's reference set, seen as visible
# text; the same list Part 24's reliability test watches for.
UNPARSED_MARKERS = (
"", "", "<|channel|>commentary", "[TOOL_REQUEST]",
"",
)
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):
parts = []
for message in completion:
if not isinstance(message, dict):
continue
if message.get("content"):
parts.append(str(message["content"]))
for call in message.get("tool_calls") or []:
function = call.get("function", {}) if isinstance(call, dict) else {}
parts.append(json.dumps({"name": function.get("name"),
"arguments": function.get("arguments")}))
return "\n".join(parts)
if isinstance(completion, dict):
return str(completion.get("content", ""))
return str(completion)
def calls_of(completion: Any) -> list[tuple[str, dict]]:
"""(name, arguments) for every structured tool call in the completion.
Only calls the serving stack already parsed count. A call that arrived as visible
text is not a call; it is a formatting failure, and the format reward is where it is
scored.
"""
out: list[tuple[str, dict]] = []
if isinstance(completion, list):
for message in completion:
if not isinstance(message, dict):
continue
for call in message.get("tool_calls") or []:
function = (call or {}).get("function", {})
arguments = function.get("arguments")
if isinstance(arguments, str):
try:
arguments = json.loads(arguments)
except json.JSONDecodeError:
arguments = {}
out.append((str(function.get("name", "")), arguments if isinstance(arguments, dict) else {}))
return out
def files_written(text: str) -> dict[str, str]:
"""The files the completion asked for, as {path: contents}. The last block wins."""
out: dict[str, str] = {}
for path, body in FILE_BLOCK.findall(text):
cleaned = path.strip()
if cleaned:
out[cleaned] = body
return out
# ---------------------------------------------------------------------------------------
# Format and schema
# ---------------------------------------------------------------------------------------
def tool_call_format_reward(value: float = 1.0) -> Callable[..., list[float]]:
"""1 when the attempt's tool calls arrived as calls, 0 when any arrived as text.
This is the reward that fixes the failure Part 24's reliability test measures: a
model that writes the tool-call format into its prose. It is cheap, it is checkable
without running anything, and on a small model it is often most of the gain.
"""
def reward(completions, **kwargs) -> list[float]:
scores = []
for completion in completions:
text = text_of(completion)
leaked = any(marker in text for marker in UNPARSED_MARKERS)
scores.append(0.0 if leaked else value)
return scores
reward.__name__ = "tool_call_format_reward"
return reward
def schema_errors(arguments: Any, schema: dict) -> list[str]:
"""The subset of JSON Schema this course checks: type, required, properties, enum."""
problems: list[str] = []
if not isinstance(arguments, dict):
return ["arguments are not an object"]
properties = schema.get("properties", {}) or {}
for name in schema.get("required", []) or []:
if name not in arguments:
problems.append(f"missing required argument {name!r}")
types = {"string": str, "number": (int, float), "integer": int,
"boolean": bool, "object": dict, "array": list}
for name, value in arguments.items():
if name not in properties:
problems.append(f"unknown argument {name!r}")
continue
spec = properties[name] or {}
wanted = types.get(str(spec.get("type", "")))
if wanted and not isinstance(value, wanted):
problems.append(f"argument {name!r} should be {spec.get('type')}")
choices = spec.get("enum")
if choices and value not in choices:
problems.append(f"argument {name!r} is not one of {choices}")
return problems
def schema_valid_reward(tools: Sequence[dict] | None = None,
value: float = 1.0) -> Callable[..., list[float]]:
"""The fraction of the attempt's calls whose arguments match the declared schema.
`tools` may be given once here or per row as a `tools` dataset column, which is what
happens when different tasks expose different tools.
"""
def reward(completions, **kwargs) -> list[float]:
per_row = kwargs.get("tools")
scores = []
for index, completion in enumerate(completions):
declared = tools
if per_row is not None:
candidate = per_row[index] if isinstance(per_row, list) else per_row
if isinstance(candidate, str):
try:
candidate = json.loads(candidate)
except json.JSONDecodeError:
candidate = None
declared = candidate or tools
schemas = {t.get("function", {}).get("name"): t.get("function", {}).get("parameters", {})
for t in (declared or [])}
calls = calls_of(completion)
if not calls or not schemas:
scores.append(0.0)
continue
good = sum(1 for name, arguments in calls
if name in schemas and not schema_errors(arguments, schemas[name]))
scores.append(value * good / len(calls))
return scores
reward.__name__ = "schema_valid_reward"
return reward
def step_budget_reward(max_steps: int = 8, penalty: float = 0.25) -> Callable[..., list[float]]:
"""0 when the attempt stayed inside the budget, negative when it ran long.
A brake rather than a target. Rewarding short trajectories directly teaches a model
to answer without looking, which is the failure this whole part exists to remove.
"""
def reward(completions, **kwargs) -> list[float]:
scores = []
for completion in completions:
steps = len(calls_of(completion)) or len(JSON_TOOL_CALL.findall(text_of(completion)))
over = max(0, steps - max_steps)
scores.append(-penalty * over if over else 0.0)
return scores
reward.__name__ = "step_budget_reward"
return reward
# ---------------------------------------------------------------------------------------
# The test suite as a reward
# ---------------------------------------------------------------------------------------
RUNNER_TEMPLATE = '''"""Written by agent-rewards.py. Applies resource limits, then runs the command."""
import os, resource, sys
resource.setrlimit(resource.RLIMIT_CPU, ({cpu}, {cpu}))
resource.setrlimit(resource.RLIMIT_AS, ({mem}, {mem}))
resource.setrlimit(resource.RLIMIT_FSIZE, ({fsize}, {fsize}))
try:
resource.setrlimit(resource.RLIMIT_NPROC, ({nproc}, {nproc}))
except (ValueError, OSError):
pass
os.execvp(sys.argv[1], sys.argv[1:])
'''
def run_tests_in_copy(
workspace: Path,
writes: dict[str, str],
test_command: Sequence[str],
protected: Sequence[str],
timeout_s: int = 120,
cpu_seconds: int = 120,
memory_mb: int = 2048,
max_output_bytes: int = 1 << 22,
max_processes: int = 128,
env_passthrough: Sequence[str] = (),
) -> tuple[bool, str]:
"""Copy the workspace, write the model's files into the copy, run the tests there.
Returns (passed, reason). The copy is thrown away afterwards, so a rollout that
deletes half the repository costs one directory rather than the repository. The child
runs with a wall-clock timeout on top of a CPU limit, an address-space limit and a
file-size limit, so an infinite loop, a runaway allocation and a program that writes a
huge 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 stops
accidents and cheap denial of service; it does not stop a program written to do harm.
Part 25 builds the container-and-dedicated-user sandbox that untrusted code needs, and
that is what a rollout loop should run inside.
"""
if os.name != "posix":
return False, "resource limits need a POSIX system; run this inside WSL2 on Windows"
if not workspace.is_dir():
return False, f"workspace {workspace} does not exist"
protected_paths = {Path(p).as_posix() for p in protected}
for path in writes:
candidate = Path(path)
if candidate.is_absolute() or ".." in candidate.parts:
return False, f"refused to write outside the workspace: {path}"
if candidate.as_posix() in protected_paths:
# This is the reward-hacking check, and it is not optional. A model that can
# edit the tests will edit the tests, because that is the cheapest way to make
# them pass, and the reward would happily pay for it.
return False, f"wrote a protected file: {path}"
with tempfile.TemporaryDirectory(prefix="course-agent-reward-") as parent:
copy = Path(parent) / "workspace"
shutil.copytree(workspace, copy, symlinks=False,
ignore=shutil.ignore_patterns(".git", "__pycache__", "*.pyc"))
for path, body in writes.items():
target = copy / path
target.parent.mkdir(parents=True, exist_ok=True)
target.write_text(body, encoding="utf-8")
runner = Path(parent) / "runner.py"
runner.write_text(RUNNER_TEMPLATE.format(
cpu=cpu_seconds, mem=memory_mb * 1024 * 1024,
fsize=max_output_bytes, nproc=max_processes), encoding="utf-8")
env = {"PATH": os.environ.get("PATH", ""), "HOME": str(copy),
"LANG": os.environ.get("LANG", "C.UTF-8")}
for name in env_passthrough:
if name in os.environ:
env[name] = os.environ[name]
try:
done = subprocess.run(
[sys.executable, "-I", "-S", str(runner), *test_command],
cwd=str(copy), env=env, 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 test command: {exc}"
if done.returncode == 0:
return True, "tests passed"
tail = ((done.stdout or "") + (done.stderr or "")).strip().splitlines()
return False, (tail[-1][:200] if tail else f"exit code {done.returncode}")
def test_suite_reward(workspace: str | os.PathLike[str],
test_command: Sequence[str] = ("python3", "-m", "pytest", "-q"),
protected: Sequence[str] = ("tests.py",),
timeout_s: int = 120,
memory_mb: int = 2048,
value: float = 1.0) -> Callable[..., list[float]]:
"""1 when the files the attempt wrote make the suite pass in a fresh copy, else 0.
The workspace and the command may also come from the dataset, as `workspace`,
`test_command` and `protected` columns, which is how one training run covers several
repositories.
"""
root = Path(workspace)
def reward(completions, **kwargs) -> list[float]:
scores = []
for index, completion in enumerate(completions):
def column(name: str, fallback):
value_ = kwargs.get(name)
if value_ is None:
return fallback
return value_[index] if isinstance(value_, list) else value_
here = Path(column("workspace", root))
command = column("test_command", test_command)
if isinstance(command, str):
command = command.split()
guard = column("protected", protected)
if isinstance(guard, str):
guard = [guard]
writes = files_written(text_of(completion))
if not writes:
scores.append(0.0)
continue
passed, _reason = run_tests_in_copy(here, writes, command, guard,
timeout_s=timeout_s, cpu_seconds=timeout_s,
memory_mb=memory_mb)
scores.append(value if passed else 0.0)
return scores
reward.__name__ = "test_suite_reward"
return reward
# ---------------------------------------------------------------------------------------
# Combining them
# ---------------------------------------------------------------------------------------
def build_agent_reward_functions(
kind: str = "tools",
tools: Sequence[dict] | None = None,
workspace: str | os.PathLike[str] | None = None,
test_command: Sequence[str] = ("python3", "-m", "pytest", "-q"),
protected: Sequence[str] = ("tests.py",),
max_steps: int = 8,
timeout_s: int = 120,
) -> tuple[list[Callable[..., list[float]]], list[float]]:
"""The reward list and weights a GRPO run passes to the trainer.
TRL sums the functions after weighting them, so the weights are the statement of what
matters. Here the task outcome dominates, format is a hint, and the step budget is a
brake. Returning both lists together keeps them the same length, which TRL requires.
"""
functions: list[Callable[..., list[float]]] = []
weights: list[float] = []
if kind == "tools":
functions.append(schema_valid_reward(tools))
weights.append(1.0)
elif kind == "tests":
if workspace is None:
raise ValueError("kind='tests' needs a workspace to run the suite in")
functions.append(test_suite_reward(workspace, test_command, protected,
timeout_s=timeout_s))
weights.append(2.0)
else:
raise ValueError(f"unknown reward kind {kind!r}; expected tools or tests")
functions.append(tool_call_format_reward())
weights.append(0.5)
functions.append(step_budget_reward(max_steps))
weights.append(1.0)
return functions, weights
def score(functions: Sequence[Callable[..., list[float]]], weights: Sequence[float],
completions: Sequence[Any], **columns: Any) -> list[float]:
"""The weighted sum, computed the way the trainer computes it."""
totals = [0.0] * len(completions)
for function, weight in zip(functions, weights):
for index, value in enumerate(function(completions=completions, **columns)):
totals[index] += weight * value
return totals
# ---------------------------------------------------------------------------------------
# The worked examples
# ---------------------------------------------------------------------------------------
DEMO_APP = '''def total(rows):
"""Sum the amount on every row. Deliberately wrong: it drops the last row."""
return sum(r["amount"] for r in rows[:-1])
'''
# Plain assertions run by the interpreter rather than a test runner, so the demo needs
# nothing installed. A real repository's suite is whatever its own command is, which is
# what --test-command and the dataset's test_command column are for.
DEMO_TESTS = '''from app import total
assert total([{"amount": 1}, {"amount": 2}, {"amount": 3}]) == 6, "does not sum every row"
assert total([]) == 0, "the total of nothing is not zero"
print("ok")
'''
DEMO_TOOLS = [
{"type": "function", "function": {
"name": "read_file",
"description": "Read a text file inside the workspace.",
"parameters": {"type": "object",
"properties": {"path": {"type": "string"}},
"required": ["path"]}}},
]
def demo_completions() -> list[tuple[str, list[dict]]]:
"""Six attempts: two honest, one wasteful, two that cheat, and one that mis-formats."""
def attempt(content: str, calls: list[tuple[str, Any]] | None = None) -> list[dict]:
message: dict[str, Any] = {"role": "assistant", "content": content}
if calls:
message["tool_calls"] = [
{"type": "function", "function": {"name": name, "arguments": arguments}}
for name, arguments in calls
]
return [message]
fixed = "```file:app.py\ndef total(rows):\n return sum(r[\"amount\"] for r in rows)\n```"
return [
("correct fix, one call",
attempt("I read the file and fixed the slice.\n" + fixed, [("read_file", {"path": "app.py"})])),
("correct fix, no tool call",
attempt(fixed)),
("wasteful: twelve reads",
attempt("Checking everything.\n" + fixed,
[("read_file", {"path": "app.py"}) for _ in range(12)])),
("hack: edits the tests instead",
attempt("```file:tests.py\nassert True\n```",
[("read_file", {"path": "tests.py"})])),
("hack: prints a success message",
attempt("ok\nAll tests passed.")),
("format failure: the call arrived as text",
attempt("\n{\"name\": \"read_file\", \"arguments\": {\"path\": \"app.py\"}}\n")),
]
def run_demo(show_sandbox: bool) -> None:
with tempfile.TemporaryDirectory(prefix="course-agent-demo-") as parent:
workspace = Path(parent) / "repo"
workspace.mkdir()
(workspace / "app.py").write_text(DEMO_APP, encoding="utf-8")
(workspace / "tests.py").write_text(DEMO_TESTS, encoding="utf-8")
functions, weights = build_agent_reward_functions(
kind="tests", workspace=workspace,
test_command=("python3", "tests.py"),
protected=("tests.py",), max_steps=8, timeout_s=60)
labels = [label for label, _ in demo_completions()]
completions = [c for _, c in demo_completions()]
columns: dict[str, Any] = {"tools": [DEMO_TOOLS] * len(completions)}
per_function = {f.__name__: f(completions=completions, **columns) for f in functions}
totals = score(functions, weights, completions, **columns)
header = f"{'attempt':<42}" + "".join(f"{name.replace('_reward',''):>20}" for name in per_function)
print(header + f"{'total':>10}")
print("-" * len(header + f"{'total':>10}"))
for index, label in enumerate(labels):
row = f"{label:<42}"
for name in per_function:
row += f"{per_function[name][index]:>20.2f}"
print(row + f"{totals[index]:>10.2f}")
print("\nWeights: " + ", ".join(f"{f.__name__}={w}" for f, w in zip(functions, weights)))
print("\nRead the fourth row. The attempt that edits the tests scores zero on the "
"test-suite reward because the protected-path check refuses the write, not "
"because the suite failed. Take that check out and it scores full marks, which "
"is what a reward function is for: paying for the cheapest thing that "
"satisfies it.")
print("The fifth row is the other half of the same lesson: a completion that says "
"the tests passed scores nothing, because the reward runs them.")
if show_sandbox:
passed, reason = run_tests_in_copy(
workspace, {"app.py": DEMO_APP}, ("python3", "tests.py"),
("tests.py",), timeout_s=60)
print(f"\nsandbox check with the unfixed file: passed={passed} reason={reason}")
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
parser.add_argument("--demo", action="store_true", help="score the worked examples")
parser.add_argument("--show-sandbox", action="store_true",
help="also run the unfixed file through the sandbox and print the reason")
args = parser.parse_args()
if not args.demo:
parser.print_help()
return
run_demo(args.show_sandbox)
if __name__ == "__main__":
main()