Reinforcement learning with verifiable rewards, on one node

Most of what a business needs from a small model is not creativity. It is compliance: produce this format, call this tool with these arguments, say "I do not know" when the document does not contain the answer. Those are things you can check with code. And if you can check it with code, you can train for it with reinforcement learning, without a human labelling anything.

This is the recipe we use for post-training small models on verifiable tasks, on a single machine. It is not the only way. It is the one that has kept working.

Why verifiable rewards#

Classic RLHF needs a reward model trained on human preferences, which needs a lot of human preferences, which needs a lot of budget. Verifiable rewards skip that. The reward is a function: does the output parse, does it match the schema, does the tool call name a real tool with valid arguments, does the extracted total equal the sum of the lines, does the test suite pass.

The signal is exact, cheap and infinitely repeatable. The limitation is the flip side: it only works for tasks where "correct" can be computed. That covers a surprising share of production work.

GRPO in one paragraph#

Group Relative Policy Optimisation samples a group of completions for each prompt, scores every one with the reward function, and uses each completion's score relative to the group's mean as its advantage. Completions better than their siblings are reinforced; worse ones are pushed down. There is no separate value model to train, which is most of why it fits on one node. A KL penalty against the starting model keeps the policy from wandering off into reward-hacking territory too quickly.

Before you start: can the base model do it at all?#

RL sharpens behaviour the model can already produce sometimes. If the base model never emits valid JSON for your schema, the group will be all zeros, the advantages will be all zeros, and nothing will learn. Run a hundred samples first. If the success rate is under a few per cent, do a short supervised fine-tune on a few hundred examples to get into range, then switch to RL.

The reward function#

This is the part to spend your time on, because it is the specification. Everything the model becomes, it becomes to satisfy this function. Ours are boring on purpose: a weighted sum of checks that are each easy to read.

import json
from jsonschema import validate, ValidationError
 
ALLOWED_TOOLS = {"lookup_order", "reschedule", "handoff"}
 
def parse(text: str):
    try:
        return json.loads(text), None
    except json.JSONDecodeError as e:
        return None, str(e)
 
def reward(prompt: dict, completion: str) -> float:
    score = 0.0
    obj, err = parse(completion)
    if obj is None:
        return -1.0                                 # unparseable: strongly negative
    score += 0.2                                    # valid JSON at all
 
    try:
        validate(obj, prompt["schema"])
        score += 0.3                                # matches the schema
    except ValidationError:
        return score - 0.5
 
    tool = obj.get("tool")
    if tool in ALLOWED_TOOLS:
        score += 0.2
    elif tool is not None:
        score -= 0.5                                # invented a tool
 
    # Task checks: the part that stops "empty but valid" from winning.
    expected = prompt["expected"]
    if tool == expected["tool"]:
        score += 0.2
    if obj.get("args", {}).get("order_id") == expected.get("order_id"):
        score += 0.1
    if expected["tool"] == "handoff" and tool == "handoff":
        score += 0.2                                # knowing when to stop is rewarded
 
    # Length: discourage padding without punishing legitimate content.
    if len(completion) > 800:
        score -= 0.1
    return score

The highlighted block is the one people forget. Without task checks, the model learns that an empty object which passes the schema scores well, and it will produce exactly that, forever. Every reward function needs at least one check that requires the model to have actually done the job.

Data: prompts and verifiers, not answers#

The training set is a list of prompts, each with whatever the verifier needs. For a tool-use task that is the schema, the allowed tools and the expected call. For extraction it is the document and the ground-truth fields. You do not need model-written answers at all, which is the second reason this is cheap.

{"prompt": "Caller says: my order 48213 has not arrived, can you check?", "schema": {"type": "object", "required": ["tool", "args"]}, "expected": {"tool": "lookup_order", "order_id": "48213"}}

A few thousand prompts is plenty for a narrow task. Variety matters more than volume: phrasings, edge cases, and a healthy share of prompts where the right answer is to hand off.

The training run#

TRL's GRPOTrainer does the sampling, scoring and updates. With LoRA adapters and vLLM handling generation, a model in the one-to-eight-billion range trains comfortably on a single GPU with 80 GB of memory, and faster on a node with several. The configuration below is the shape; the exact argument names move with versions.

from datasets import load_dataset
from trl import GRPOConfig, GRPOTrainer
from peft import LoraConfig
from reward import reward
 
ds = load_dataset("json", data_files="train.jsonl")["train"]
 
def reward_fn(completions, **kw):
    return [reward(p, c) for p, c in zip(kw["prompt_meta"], completions)]
 
cfg = GRPOConfig(
    output_dir="out/tool-use-grpo",
    num_generations=8,              # group size per prompt
    max_completion_length=256,
    per_device_train_batch_size=8,
    gradient_accumulation_steps=4,
    learning_rate=1e-6,
    beta=0.04,                      # KL penalty against the reference model
    temperature=0.9,                # diversity inside the group
    use_vllm=True,
    bf16=True,
    logging_steps=10,
    save_steps=200,
)
 
trainer = GRPOTrainer(
    model="/models/base-3b",
    reward_funcs=reward_fn,
    args=cfg,
    train_dataset=ds,
    peft_config=LoraConfig(r=32, lora_alpha=64, target_modules="all-linear"),
)
trainer.train()

Watch three curves: mean reward (should rise), reward standard deviation inside groups (should stay above zero, otherwise nothing is being learned), and KL to the reference (should rise slowly, not jump).

How it goes wrong#

• Reward hacking. The model finds the cheapest output that scores. Empty objects, the same tool every time, a hand-off for everything. The fix is always in the reward function, never in the trainer.
• Length collapse. Completions shrink to the minimum that passes. Fine for tool calls, bad for anything with content. Add a floor to the length term if it matters.
• KL drift. Reward climbs, then general capability quietly degrades. Raise beta, lower the learning rate, and run the general evals alongside the task eval.
• All-zero groups. The task is too hard for the base model. Go back to the supervised step.

Gate it with evals#

A rising reward curve is not evidence the model is better. It is evidence the model is better at the reward function. Before anything ships, the checkpoint goes through the same evaluation harness as every other change: the task set, the general set, the safety set, and a replay of production samples. If the task score rose and anything else fell, it does not ship.

That is the whole recipe. A reward function you can read, prompts with verifiers, one node, and an eval gate that does not care how good the training curve looked.