Write a reward function

In group relative policy optimization (GRPO), the reward function is the only learning signal. For each prompt, the model generates several completions, your function scores each one, and the update favors completions that scored above their group’s mean. This page shows how to use the cookbook’s built-in datasets, register your own dataset and reward, and choose a reward shape. It also shows how to test a reward and check a baseline score before you spend GPU time on training.

Requirements

A cookbook checkout. To set one up, see the RL quickstart.
A dataset where each row has a prompt and a target to score completions against, either on Hugging Face or in a form you can load into the same shape.

A reward function takes the model’s decoded completion and the row’s target, and returns a float:

def exact_match_reward(prediction: str, target: str) -> float:
    return 1.0 if prediction.strip() == target.strip() else 0.0

A good reward function:

Returns values in [0, 1]: This keeps advantages on a predictable scale.
Varies across completions: A reward that scores most completions the same leaves GRPO nothing to learn from. A group where every completion scores the same has zero-variance advantages and is skipped. Before training, score a sample of the base model’s outputs and check that the scores spread.
Is deterministic: The same prediction and target always score the same.
Is fast: It runs at least batch_size * group_size times per step, and more often with dynamic sampling or validation, or when it doubles as the metric. The cookbook scores single-turn completions one at a time, so a reward that waits on an API call, a model, or a subprocess adds that wait to every step. To score with something slow, such as a judge model, write a one-turn environment instead, because the cookbook runs every episode’s async step concurrently.

Built-in datasets

The cookbook registers four datasets, each paired with a reward. To train on one, pass its name as dataset on GRPOConfig. Their entries in post_training_cookbook/datasets/registry.py also serve as working examples of the fields described in Register a dataset and reward.

Name	Hugging Face dataset	Task	Reward
`gsm8k`	`openai/gsm8k`	Grade-school math word problems. The prompt asks for a final line of `#### <answer>`.	`1.0` when the number after the first `####`, or inside `\boxed{}` when there’s no `####`, matches the reference.
`deepmath`	`zwhe99/DeepMath-103K`	Competition math. The prompt asks for the answer inside `\boxed{}`.	`1.0` when the boxed answer is mathematically equivalent to the reference.
`deepmath-lv5`	`zwhe99/DeepMath-103K`	The `deepmath` problems at difficulty 5 and below.	Same as `deepmath`.
`game24`	`nlile/24-game`	Combine four numbers with arithmetic to make 24, ending with a boxed expression. Only solvable rows are used.	Partial credit for a valid expression that uses each number once and for its closeness to 24, plus a bonus for hitting 24 exactly. Validation reports strict 0-or-1 accuracy instead.

Register a dataset and reward

The cookbook’s dataset registry maps each dataset name to how its rows become prompts, what its target is, and how its completions are scored. The steps below register a dataset of Python coding problems, where each row has a prompt and an expected_output.

Write the reward function

Add your reward function to post_training_cookbook/datasets/rewards.py. This example runs the generated code and compares its standard output with the expected output:

import functools
import subprocess
import sys
import tempfile


@functools.lru_cache(maxsize=1024)
def code_output_reward(prediction: str, target: str) -> float:
    """Score 1.0 if the code's stdout matches the expected output."""
    code = _extract_code_block(prediction)
    if code is None:
        return 0.0
    try:
        with tempfile.NamedTemporaryFile(
            mode="w", suffix=".py", delete=True
        ) as f:
            f.write(code)
            f.flush()
            result = subprocess.run(
                [sys.executable, f.name],
                capture_output=True,
                text=True,
                timeout=5,
            )
        return 1.0 if result.stdout.strip() == target.strip() else 0.0
    except Exception:
        # A timeout, a crash, or unrunnable code all score 0.
        return 0.0


def _extract_code_block(text: str) -> str | None:
    """Extract the last ```python ... ``` block from text."""
    blocks = text.split("```python")
    if len(blocks) < 2:
        return None
    last_block = blocks[-1]
    end = last_block.find("```")
    return last_block.strip() if end == -1 else last_block[:end].strip()

The cookbook scores every completion twice, once with compute_reward and once with compute_metric, which falls back to the reward when you don’t set one. The lru_cache decorator returns the cached score for the second call, so each completion’s code runs only once.

Add a DatasetConfig entry to post_training_cookbook/datasets/registry.py:

from post_training_cookbook.datasets.rewards import code_output_reward

DATASETS["code_output"] = DatasetConfig(
    hf_name="my-org/python-coding-problems",
    hf_config=None,
    hf_split="train",
    format_question=lambda row: (
        f"{row['prompt']}\n\n"
        "Write Python code that solves this problem. "
        "Put your solution in a ```python code block."
    ),
    extract_target=lambda row: row["expected_output"],
    compute_reward=code_output_reward,
    few_shot_messages=[
        {
            "role": "user",
            "content": (
                "Write a function that returns the sum of two numbers, "
                "then print the result of calling it with 3 and 4.\n\n"
                "Write Python code that solves this problem. "
                "Put your solution in a ```python code block."
            ),
        },
        {
            "role": "assistant",
            "content": (
                "```python\n"
                "def add(a, b):\n"
                "    return a + b\n\n"
                "print(add(3, 4))\n"
                "```"
            ),
        },
    ],
)

Field	Description
`hf_name`, `hf_config`, `hf_split`	The Hugging Face dataset repo or a local folder, its config (`None` if the dataset has none), and the split.
`format_question`	Takes a row and returns the prompt text.
`extract_target`	Takes a row and returns the ground truth the reward scores against.
`compute_reward`	Takes the decoded completion and the target, and returns the reward.
`compute_metric`	Optional. A separate accuracy function for validation. Falls back to `compute_reward` when `None`.
`filter_fn`	Optional. Takes a row and returns `True` to keep it.
`few_shot_messages`	Optional. Examples prepended to every prompt as earlier user and assistant turns. Write them in the output format the reward expects.

To train on data that isn’t on Hugging Face, save it as a JSON Lines file named train.jsonl, one row per line, in a folder of its own. Then set hf_name to the folder’s path, such as "./my_problems", hf_config to None, and hf_split to "train". The cookbook loads the folder with the Hugging Face datasets library, which reads train.jsonl into the train split. A relative path resolves from the directory you run the recipe in.

Test the reward

Run the reward on hand-written completions before you train on it:

from post_training_cookbook.datasets.rewards import code_output_reward

assert code_output_reward('```python\nprint("hello")\n```', "hello") == 1.0
assert code_output_reward('```python\nprint("world")\n```', "hello") == 0.0
assert code_output_reward("The answer is hello", "hello") == 0.0

Then add the same cases to tests/datasets/test_rewards.py and run uv run pytest tests/datasets. Include one edge case per failure mode: an empty completion, a very long completion, output in an unexpected format, and, for this reward, a timeout.

Create and run the recipe

Create recipes/grpo/code_output.py:

from post_training_cookbook.grpo.config import GRPOConfig
from post_training_cookbook.grpo.train import train_single_turn

config = GRPOConfig(
    model="Qwen/Qwen3.6-35B-A3B",
    dataset="code_output",
    num_examples=5000,
    num_val_examples=200,
    batch_size=8,
    group_size=8,
    num_epochs=3,
    learning_rate=1e-5,
    max_model_length=2048,
    max_prompt_length=512,
    max_sample_tokens=1024,
    enable_thinking=False,
    temperature=0.8,
    wandb_project="grpo-code-output",
)

if __name__ == "__main__":
    train_single_turn(config)

A code solution takes more tokens than a math answer, so this recipe sets max_sample_tokens higher than the GSM8K recipe does, and temperature=0.8 gives the completions in each group more room to diverge. Run it with uv run python recipes/grpo/code_output.py.

Reward shapes

Binary: 1.0 for correct, 0.0 for wrong, like exact_match_reward above. Use it for tasks with a single checkable answer, such as math and factual questions. Tiered: Partial credit at several levels. Use it when partial progress deserves credit, such as in multi-step problems or in puzzles where getting the format right is a large part of the task.

def tiered_reward(prediction: str, target: str) -> float:
    score = 0.0
    answer = extract_answer(prediction)
    if answer is not None:
        score += 0.2  # followed the format
    if answer and is_close(answer, target):
        score += 0.4  # close to the answer
    if answer == target:
        score += 0.4  # exact match
    return score

Continuous: A distance or similarity score. Use it for tasks where approximate answers have value, such as translation or summarization.

from difflib import SequenceMatcher


def similarity_reward(prediction: str, target: str) -> float:
    return SequenceMatcher(None, prediction.strip(), target.strip()).ratio()

Separate reward and metric: The reward you train on and the metric you report can differ. A smooth, tiered reward gives GRPO a gradient to follow, while a strict binary metric tells you how often the model is actually right:

DATASETS["my_task"] = DatasetConfig(
    ...,
    compute_reward=tiered_reward,
    compute_metric=exact_match_accuracy,
)

compute_reward drives the advantages. compute_metric is what validation reports as val/accuracy.

Design guidance

Reward the format you want: If you want \boxed{} answers, give partial credit for the format even when the content is wrong. The model has to learn to emit the format before it can get the content right.
Penalize degenerate outputs: Return 0.0 for empty responses, repetition, or outputs that don’t attempt the task.
Guard against reward hacking: Watch val/accuracy alongside training reward. If training reward rises while validation accuracy falls, or validation response length creeps up, the model has found a way to score well on your reward without solving the task.
Sandbox code execution: A reward that executes model output, like the code example above, needs a timeout and should run isolated from your machine, for example in a container or a code sandbox.
Don’t rely on the main thread: Rewards can run off the main thread, for example when a sweep trains each arm in its own thread. Timeouts built on signal.alarm fail there, so use a subprocess timeout like the one in the code example.

Check a baseline before training

To see how a model performs on your task before you spend GPU time training it, run one of the cookbook’s recipes/serverless_eval/ scripts. Each script scores a hosted chat model on a dataset’s validation split, without a training session or model resource:

uv run python recipes/serverless_eval/gsm8k.py

The scripts cover gsm8k, deepmath, game24, and the multi-turn tool-use environment. To score your own registered dataset, copy a script and change its dataset. Each run writes its episodes to runs/serverless_eval_<label>_<model>_<timestamp>/episodes.json. To score Qwen3.8 27B before training, deploy Qwen/Qwen3.8-27B-FP8 on dedicated model inference. The scripts take no arguments, so edit the script’s model to the endpoint string. The same flow works for a model you trained. For an adapter attached to a multi-LoRA deployment, use the adapter’s model string, such as your-project-slug/my-sweep:ml_abc123. Then run the script with TOGETHER_BASE_URL set to the dedicated inference host:

export TOGETHER_BASE_URL=https://api-inference.together.ai/v1
uv run python recipes/serverless_eval/gsm8k.py

Stop the endpoint when you’re done, because it bills while it runs. The baseline isn’t directly comparable to training-time validation, because it applies no codec, token budget, or prompt-length filter. Watch the eval/length_capped metric in the output, which reports the fraction of episodes the endpoint cut off at max_tokens. If it’s high, the score measures truncation rather than ability, so raise max_tokens, especially for models with thinking on.

Multi-turn rewards

In a multi-turn environment, the reward comes from the environment’s step result instead of a (prediction, target) function. It’s assigned per episode, not per completion. For details, see Multi-turn environments.

Next steps

Reinforcement learning with GRPO

The settings that shape how your reward turns into updates.

Multi-turn environments

Reward whole episodes of tool use instead of single completions.

Monitoring

The metrics that show whether your reward is giving the model a signal.

RL quickstart

Run the GSM8K recipe your dataset will slot into.

GET STARTED

SERVERLESS

INFERENCE APIS

DEDICATED MODEL INFERENCE

DEDICATED CONTAINER INFERENCE

POST-TRAINING

FINE-TUNING

GPU CLUSTERS

CODE EXECUTION

ADMINISTRATION

Write a reward function

Requirements