[Feature] Reasoning reward utilities: group-relative length-scaled accuracy (GRPO-LEAD) + graduated format reward
- Dominant language
- Python
- Stars
- 19.3k
- Forks
- 3k
- Avg merge
- 1d 20h
- Merged PRs (30d)
- 194
Description
### Feature request
Add two reward functions to `trl.rewards`, both built on existing primitives:
1. **`get_length_scaled_accuracy_reward(alpha=0.5, incorrect_reward=-1.0, reasoning_delimiters=None)`** : a reward-function factory that wraps `reasoning_accuracy_reward` and scales the reward of *correct* completions by their length **relative to the other correct completions of the same prompt** (a z-score within the GRPO group). Implements the length-aware shaping from **GRPO-LEAD** (https://huggingface.co/papers/2504.09696).
2. **`graduated_format_reward(completions, **kwargs)`** : a graduated variant of `think_format_reward` that gives partial credit (`1.0 / 0.5 / 0.25 / 0.0`) for partial `...` structure, providing a denser signal during cold-start training when strict-format matches are rare.
### Motivation
**Length-scaled accuracy : how it differs from the existing `get_cosine_scaled_reward`.**
TRL already has [`get_cosine_scaled_reward`](https://huggingface.co/docs/trl/main/rewards#get_cosine_scaled_reward) (#6066), which also rewards shorter correct completions. The two are complementary, not redundant:
| | `get_cosine_scaled_reward` (existing) | `get_length_scaled_accuracy_reward` (proposed) |
|---|---|---|
| Length reference | **absolute** token count vs a fixed `max_len` budget | **group-relative** z-score among the prompt's correct completions |
| Needs a length budget? | yes (`max_len`) | no — adapts to each prompt's own length distribution |
| Shaping curve | cosine interpolation between fixed bounds | `exp(-alpha * z)`; a group's correct rewards have geometric mean `1.0` |
| Paper | Demystifying Long CoT | GRPO-LEAD |
The group-relative formulation is useful when there is no good global length budget (mixed-difficulty data): on an easy prompt where every completion is short, and on a hard prompt where every completion is long, GRPO-LEAD still applies the same relative pressure toward the shorter correct answers, because it standardizes within the group. A fixed `max_len` schedule cannot do that.
**Graduated format reward.** `think_format_reward` is binary: a completion that emits only `` (or both tags in the wrong order) scores exactly `0.0`, identical to a completion with no tags at all. Early in training that signal is very sparse. A graduated version gives the policy a gradient to climb toward the target structure. Nothing equivalent exists in `trl.rewards` today.
Both functions follow existing conventions: `get_length_scaled_accuracy_reward` mirrors the picklable callable-class factory pattern of `get_cosine_scaled_reward` / `get_soft_overlong_punishment` (so it works under the async GRPO rollout worker), and `graduated_format_reward` sits alongside `think_format_reward`.
### Your contribution
I have an implementation ready (functions, docs, a `paper_index.md` entry for GRPO-LEAD that cross-references the cosine reward, and unit tests including a picklability test) and will open a PR referencing this issue.
Contributor guide
Research direction
Start in trl.rewards by reading get_cosine_scaled_reward, get_soft_overlong_punishment, reasoning_accuracy_reward, and think_format_reward to follow the existing conventions. Review the mentioned unit tests and picklability test, then update the documentation and paper_index.md entry. Done means both reward utilities, their tests, and the GRPO-LEAD documentation are included.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- machine-learning
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 35/100