group rewards

compute_group_rewards is called once per prompt-group. it receives every completed rollout in the group (a sequence of BaseRollout), letting you compute rewards that depend on cross-rollout context. it runs alongside compute_reward, not instead of it; the trainer merges both reward dicts. return a mapping keyed by rollout id:

from collections.abc import Mapping, Sequence

from benchmax.envs import BaseRollout, RewardMap

async def compute_group_rewards(
    self, rollouts: Sequence[BaseRollout]
) -> Mapping[str, RewardMap] | None:
    ...
    return {r.rollout_id: {"component": score(r)} for r in rollouts}

each RewardMap value is just a dict of component name to float score. as with compute_reward, the transcript is rollout.messages and per-example data is rollout.example_args; all rollouts in a group share one example. per-rollout and group reward component names must not overlap.

keep every component non-negative, the same as compute_reward. see best practices.

pairwise win-rate

judges are usually better at “A beats B” than at producing calibrated absolute scores. fan out one judge call per pair, then aggregate into a win-rate per rollout.

import asyncio

from benchmax.rewards import extract_completion_text

async def compute_group_rewards(self, rollouts):
    texts = {r.rollout_id: extract_completion_text(r.messages) for r in rollouts}
    ids = [r.rollout_id for r in rollouts]
    wins = {rid: 0 for rid in ids}
    ground_truth = rollouts[0].example_args.get("ground_truth", "")

    # fan out all pairwise comparisons concurrently
    pairs = [(a, b) for i, a in enumerate(ids) for b in ids[i+1:]]
    results = await asyncio.gather(*[
        self._judge_pair(texts[a], texts[b], ground_truth) for a, b in pairs
    ])

    for (a, b), winner in zip(pairs, results):
        if winner == "a":
            wins[a] += 1
        elif winner == "b":
            wins[b] += 1

    max_wins = len(ids) - 1
    return {rid: {"win_rate": wins[rid] / max_wins} for rid in ids}

the async signature exists for exactly this fan-out pattern: all n*(n-1)/2 judge calls run in parallel.

rank transform

convert raw scores to in-group rank before they become advantages. this is robust to reward-model scale drift and outliers: a score of 9.8 when everyone else scored 9.7 matters; a score of 9.8 when everyone else scored 2.0 matters differently. ranking normalizes both cases.

async def compute_group_rewards(self, rollouts):
    raw = [
        self._score(
            extract_completion_text(r.messages),
            r.example_args.get("ground_truth", ""),
        )
        for r in rollouts
    ]

    # rank within group: 0.0 = worst, 1.0 = best
    order = sorted(range(len(raw)), key=lambda i: raw[i])
    ranks = [0.0] * len(raw)
    for position, idx in enumerate(order):
        ranks[idx] = position / (len(raw) - 1) if len(raw) > 1 else 1.0

    return {r.rollout_id: {"rank": rank} for r, rank in zip(rollouts, ranks)}

self-consistency (label-free)

when the example carries no ground truth, or it’s unreliable, majority consensus across the group acts as a pseudo-label. rollouts that agree with the majority get a positive signal; outliers get none.

from collections import Counter

async def compute_group_rewards(self, rollouts):
    answers = [
        self._extract_answer(extract_completion_text(r.messages)) for r in rollouts
    ]

    # find the most common answer
    counts = Counter(answers)
    majority, majority_count = counts.most_common(1)[0]

    # reward agreement with majority, scaled by how strong the consensus is.
    # outliers score 0, not a negative value (keep components non-negative).
    consensus_strength = majority_count / len(answers)
    return {
        r.rollout_id: {"consistency": consensus_strength if a == majority else 0.0}
        for r, a in zip(rollouts, answers)
    }

useful for tasks where you have prompts but no labels, or where ground truth is expensive to produce. works best when the group size is large enough (≥8) that the majority answer is meaningful.

test coverage / complementarity

reward passing the cases others missed. this pushes the group toward pass@k coverage (one rollout specializes in the edge cases another handles easily) rather than having all rollouts pile onto the same easy tests.

async def compute_group_rewards(self, rollouts):
    # passed[i] is a set of test-case ids that rollout i passed
    passed = [self._run_tests(r) for r in rollouts]
    pass_counts = Counter(tc for p in passed for tc in p)

    rewards = {}
    for r, p in zip(rollouts, passed):
        # tests passed by fewer rollouts are worth more
        complementarity = sum(1.0 / pass_counts[tc] for tc in p)
        # normalize by group size so scale is stable
        rewards[r.rollout_id] = {"complementarity": complementarity / len(rollouts)}

    return rewards

a test passed by only one rollout contributes 1/1 = 1.0; one passed by all contributes 1/n ≈ 0. the result incentivizes finding unique solutions rather than rediscovering the majority path.

diversity scaling

without diversity pressure, the model can collapse to a single strategy per prompt. scale_by_diversity groups the rollouts into clusters of similar responses, then divides each rollout’s reward by the size of its cluster. a strategy that five rollouts converge on has its reward split five ways; a unique strategy keeps full reward. this rewards exploring different approaches instead of piling onto one.

from benchmax.rewards import NgramDiversityConfig, extract_completion_text, scale_by_diversity

async def compute_group_rewards(self, rollouts):
    texts = [extract_completion_text(r.messages) for r in rollouts]
    raw = [
        {"diverse_quality": self._score(text, r.example_args)}
        for text, r in zip(texts, rollouts)
    ]
    scaled, cluster_info = await scale_by_diversity(
        rewards=raw,
        texts=texts,
        config=NgramDiversityConfig(),
    )
    return {r.rollout_id: reward for r, reward in zip(rollouts, scaled)}

note the scaled component gets its own name (diverse_quality): individual and group reward keys must stay disjoint, so don’t reuse a key that compute_reward already returns.

what counts as “similar” depends on the method:

method	compares	how	tradeoff
ngram	surface form	character n-gram overlap (Jaccard similarity), so responses that share wording cluster together	fast, offline, no API calls. good when distinct strategies also read differently
llm	meaning	a judge model groups responses by underlying tactic, so two differently-worded responses with the same approach still cluster	catches paraphrases, but costs judge calls and needs a Judge

the config value selects the method and its parameters:

from benchmax.envs import InjectedAuth
from benchmax.rewards import Judge, LLMDiversityConfig, NgramDiversityConfig

# fast offline clustering
config = NgramDiversityConfig(similarity_threshold=0.5)

# semantic clustering via an LLM judge
config = LLMDiversityConfig(
    judge=Judge(
        model="gpt-5.4-mini",
        base_url="https://api.openai.com/v1",
        auth=InjectedAuth("judge"),
    ),
)

ngram clustering is offline and deterministic. a failed llm clustering call raises JudgeError, which the runtime records as an operational failure with zeroed reward keys for the group; it never silently mis-scores.

the telestich example trains with group rewards end to end, combining deterministic checks with an llm judge.