rewards

the reward function defines what the model optimizes for during training. every design choice here directly shapes the model’s behavior.

defining rewards

compute_reward runs after each rollout. it receives the completed rollout as a BaseRollout: rollout.messages is the full conversation transcript (a list of Message dicts) and rollout.example_args is the per-example data (every field from your example payload other than prompt_messages, e.g. ground_truth). return a dict mapping reward component names to float scores. component names are inferred from the returned dictionaries; there is no separate reward-shape declaration.

extract_completion_text is a helper that pulls the assistant’s text out of a messages list. use it at the top of your reward function to get the text you’ll score.

import re

from benchmax.envs import BaseEnv, BaseRollout
from benchmax.rewards import extract_completion_text

class MyEnv(BaseEnv):
    async def compute_reward(self, rollout: BaseRollout) -> dict[str, float]:
        text = extract_completion_text(rollout.messages)
        match = re.search(r"<answer>(.*?)</answer>", text, re.DOTALL)
        answer = match.group(1).strip() if match else ""

        ground_truth = rollout.example_args.get("ground_truth", "")
        return {"correct": 1.0 if answer == ground_truth.strip() else 0.0}

multiple reward components

return multiple keys to score different dimensions independently. each component is tracked separately on the platform, so you can monitor and debug them individually in the rollout inspector.

class MyEnv(BaseEnv):
    async def compute_reward(self, rollout: BaseRollout) -> dict[str, float]:
        text = extract_completion_text(rollout.messages)
        ground_truth = rollout.example_args.get("ground_truth", "")

        correctness = await self._judge_correctness(text, ground_truth)
        format_ok = 1.0 if self._is_valid_format(text) else 0.0

        return {"correctness": correctness, "format": format_ok}

to weight one dimension over another, scale the value inside compute_reward before returning it (or make a component conditional on another, like gating conciseness on correctness). the platform shows each component’s score over time as a separate chart, making it easy to see which dimension is improving and which is lagging.

reward dictionaries may be conditional. when a rollout omits a component that another rollout in the same training batch emitted, the trainer fills the missing value with 0.0. if every rollout in a batch returns an empty dictionary, the batch has an all-zero aggregate reward and no named components.

going further

two patterns extend the basic per-rollout reward:

best practices

• keep reward components non-negative. each component should stay in [0, 1]: return 0 for failure, not a negative score. negative values make the reward breakdown harder to read. to express a penalty, gate the component (return 0 when a condition fails) or reward the absence of the bad quality instead. (an llm-judge negative rubric still returns [0, 1]; invert it in compute_reward.)
• gate secondary rewards on primary ones. for example, if you have a conciseness reward, gate it on correctness, so the model doesn’t learn to produce short, wrong answers.
• start simple. design with a strict few reward components first. add more only when you see specific behaviors to encourage or discourage.
• watch for reward hacking. if the model finds a shortcut that scores well but produces bad output, your reward function has a gap. inspect high-reward completions regularly.
• use deterministic gates before expensive judges. check cheap conditions first (is there an answer tag? does it parse?) and return 0 immediately for obvious failures. saves LLM judge cost.
• log intermediate values. use logger.info() inside compute_reward to trace how scores were computed. see logging.
• test locally first. run your reward function on hand-crafted completions before launching. see testing.