Reward Functions & GRPO

Score an agent's outputs with a Python function and improve it with GRPO, reward modelling, or CPO

Grade an agent’s answers with a plain Python function, then let GRPO train the model to score higher — no preference pairs required.

praisonai-train llm adds three RL / preference objectives on top of SFT: grpo (group-relative policy optimisation, scored by your reward functions), reward (train a reward model from preference pairs), and cpo (contrastive preference optimisation, no reference model). Only GRPO needs reward functions — the other two learn from chosen / rejected pairs.

Quick Start

Start with an Agent whose behaviour you want to score, write one function that grades its outputs, then run GRPO on it.

1. Start from an Agent

   Pick the behaviour you want to reward — here, concise answers.

   from praisonaiagents import Agent
   
   agent = Agent(instructions="Answer briefly and correctly.")
   agent.start("What is the capital of France?")

2. Write a reward function

   TRL calls each reward function as (prompts, completions, **kwargs) and expects one float per completion. Put it in an importable module — e.g. myproject/rewards.py.

   # myproject/rewards.py
   def length_penalty(prompts, completions, **kwargs) -> list[float]:
       """Reward short answers: 1.0 for <= 20 words, decreasing after."""
       return [max(0.0, 1.0 - len(c.split()) / 20) for c in completions]

3. Point reward_funcs at it and run GRPO

   reward_funcs names each callable by its module:function import path. GRPO needs only a prompt column — it generates its own completions.

   # config.yaml
   method: grpo
   model_name: "unsloth/gemma-2-2b-it-bnb-4bit"
   max_seq_length: 2048
   reward_funcs:
     - myproject.rewards:length_penalty
   num_generations: 8
   dataset:
     - name: "your-org/your-prompts"
       split: "train"

   pip install "praisonai-train[llm]"
   praisonai-train llm config.yaml

Which method should I use?

Dataset shapes

Each method fails fast before training if the required columns are missing — the error names both the missing columns and the columns your dataset actually has.

Reward functions

A reward function is a Python callable that scores generated completions. GRPO calls it once per generation batch and uses the returned floats to push the model toward higher-scoring answers.

The module:function convention

A YAML config can only carry a string, so reward_funcs names each callable by a dotted import path — the same module:function form console_scripts and gunicorn use.

method: grpo
reward_funcs:
  - myproject.rewards:length_penalty
  - myproject.rewards:contains_answer
num_generations: 8

A single spec doesn’t have to be a list — reward_funcs: myproject.rewards:length_penalty works too. You can also pass callables directly from the Python API.

The signature TRL expects

def my_reward(prompts, completions, **kwargs) -> list[float]:
    ...

One float per completion. Extra columns from your dataset arrive in **kwargs, keyed by column name.

Resolved before the model loads

Every reward path is imported and validated before the multi-gigabyte model loads — a typo costs five seconds, not a lost session.

Why an import path, not a registry?

Two alternatives were considered and deliberately not offered:

Signature check is a warning, not a failure

If a reward function’s signature doesn’t accept completions, the trainer prints a warning and continues — TRL may still call it correctly, and refusing would block working code:

WARNING: my_reward does not take 'completions'; TRL calls reward functions as (prompts, completions, **kwargs)

Callables whose signature can’t be read (builtins, C callables) are left to TRL rather than guessed at.

Python API — praisonai_train.rewards

Resolve import paths yourself, or pass callables straight into a config.

from praisonai_train.rewards import resolve, resolve_all, require, RewardError

fn = resolve("myproject.rewards:length_penalty")   # -> callable
fns = resolve_all([                                  # collects all failures at once
    "myproject.rewards:length_penalty",
    "myproject.rewards:contains_answer",
])
require("grpo", fns)                                 # raises RewardError if empty

GRPO config keys

Three keys apply to GRPO, alongside the shared training keys.

Fast rollouts (fast_inference)

GRPO always ran the slow HF-generate path for its rollouts. Flip fast_inference: true to swap it for vLLM.

method: grpo
model_name: "unsloth/gemma-2-2b-it-bnb-4bit"
fast_inference: true
gpu_memory_utilization: 0.6
max_lora_rank: 32
num_generations: 8
reward_funcs:
  - myproject.rewards:length_penalty

The two companion knobs are ignored unless fast_inference is on. See Unsloth Surface for the full walkthrough.

Errors you might see

Each error runs before the model loads and tells the failing part apart, so you know exactly what to fix.

method 'grpo' needs reward_funcs: a list of 'module:function' import paths
that score completions. Example:
  reward_funcs:
    - myproject.rewards:length_penalty

cannot import 'nomodule' for reward function 'nomodule:f': No module named
'nomodule'. Is it on PYTHONPATH?

'mymod' has no 'no_such'. Callables there: contains_answer, length_penalty

reward function 'mymod.func' is not a 'module:function' path.
Example: myproject.rewards:length_penalty

'mymod:some_int' is a int, not a function

Best Practices

Related