> ## Documentation Index > Fetch the complete documentation index at: https://docs.together.ai/llms.txt > Use this file to discover all available pages before exploring further. # Write a reward function > How to write a reward function for GRPO, register it with your dataset, and test it before you start training. 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](/docs/post-training/quickstart-rl). * 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. ## Write a reward function A reward function takes the model's decoded completion and the row's target, and returns a float: ```python theme={null} 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](/docs/post-training/multi-turn-environments) 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](#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 `#### `. | `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`. 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: ````python theme={null} 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`: ````python theme={null} 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. Run the reward on hand-written completions before you train on it: ````python theme={null} 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 `recipes/grpo/code_output.py`: ```python theme={null} 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. ```python theme={null} 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. ```python theme={null} 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: ```python theme={null} 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](/docs/together-code-sandbox). * **Don't rely on the main thread:** Rewards can run off the main thread, for example when a [sweep](/docs/post-training/multi-lora#run-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: ```bash theme={null} 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_