feat(rewards): add a safe five-component RLCR math reward

[toffee-desuwa]

Description

As invited in #607, this PR contributes the reward implementation proposed there.

This patch adds RLCRRewardFn, a single registered reward class implementing the five-component reward structure from Damani et al., “Beyond Binary Rewards: Training LMs to Reason About Their Uncertainty”:

format + accuracy + Brier + mean confidence + endpoint-confidence auxiliary

It also adds a runnable math example and focused behavioral tests.

Why one class returns five values

Trinity associates one reward function with a task. In the standard workflow, the returned reward dictionary is reduced with sum(reward_dict.values()); it does not apply a separate per-component weight mapping there. RLCRRewardFn therefore validates a named weight mapping and returns five already weighted floats. This claim is limited to the standard reward-dictionary workflow; other scalar reward-shaping utilities in the repository have their own weighting behavior.

Response and scoring contract

The final response must end with this case-sensitive chain, followed by whitespace only:

<think>...</think><answer>...</answer><analysis>...</analysis><confidence>q</confidence>

The class uses the final complete chain, parses only the final <answer> payload, and requires finite q ∈ [0,1]. A valid response returns:

• format = 1
• accuracy = a
• brier = 1 - (a - q)^2
• mean_confidence = q
• confidence_one_or_zero = 1[q < 0.01 or q > 0.99]

Each value is multiplied by its named configured weight before it is returned.

Malformed terminal format zeros format, accuracy, and Brier. The two small confidence auxiliaries are retained only when the final confidence field is independently and safely parsed. Parser/verifier exceptions or timeouts are treated as unknown: accuracy and Brier are zero, not “wrong answer.”

Intentional differences from the official reference code

The port preserves the published five-component reward shape but deliberately hardens numerical and dispatch behavior. These deviations are part of the compatibility claim; the official reference code is reward_fns.py:

• Reject non-finite confidence (NaN, Inf, -Inf) instead of letting it enter reward arithmetic.
• Give no endpoint auxiliary reward when q is malformed or cannot be parsed safely.
• Reject out-of-range q instead of clamping it into the auxiliary reward path.
• Raise deterministic configuration errors for missing, unknown, non-numeric, or non-finite component weights.
• Verifier fixed to math_verify at construction; the source kwarg is accepted and ignored (the reference implementation dispatches the whole batch from source[0]).
• Require a strict, case-sensitive terminal <think>…</think><answer>…</answer><analysis>…</analysis><confidence>q</confidence> chain with only trailing whitespace.
• Require a balanced <think> opener. The reference format_reward (tabc) regex begins with .*?</think>, so it awards the format point to a response that has only a </think> closer and no <think> opener (for example work</think><answer>…</answer><analysis>…</analysis><confidence>q</confidence> matches). This port requires <think> to be present and balanced, and scores a missing opener as zero format, accuracy, and Brier.
• Require every reserved tag in the whole response to be globally balanced and non-nested. The reference regex matches the terminal chain against the end of the response behind a lazy .*? prefix, so extra think/answer/analysis/confidence tags that appear earlier (including nested or unbalanced ones inside the reasoning region, for example <think>x<answer>nested</answer></think> before a valid terminal chain) are absorbed by the prefix and still receive the format point. This port scans all reserved tags with a stack discipline and scores any stray, nested, or crossed reserved tag anywhere as zero format, accuracy, and Brier.

The example's weights are runnable defaults, not a frozen experimental prescription.

Files changed

• trinity/common/rewards/rlcr_reward.py
• trinity/common/rewards/__init__.py
• tests/common/rewards/test_rlcr_reward.py
• examples/rlcr_math/rlcr_math.yaml
• examples/rlcr_math/README.md

Test evidence

Local focused execution used minimal in-process stubs only for dependencies missing from the current Windows environment. Packaged commit: d350ef09fdd1fc31928fda98911817da21089e86. Exact local commands and results (paths shown relative to the repository checkout root <checkout>):

• Focused RLCR behavioral tests with isolated ray / word2number stubs — python <checkout>/run_focused_with_stubs.py → 62 passed.
• All configured pre-commit hooks on the five changed files — pre-commit run --files <checkout>/trinity/common/rewards/rlcr_reward.py <checkout>/trinity/common/rewards/__init__.py <checkout>/tests/common/rewards/test_rlcr_reward.py <checkout>/examples/rlcr_math/rlcr_math.yaml <checkout>/examples/rlcr_math/README.md → all passed, including Black 23.7.0, isort 5.12.0, flake8 6.1.0, and mypy 1.7.0.
• Extra static checks — ruff check on the changed Python files (All checks passed!), python -m compileall -q on the implementation and tests (exit 0), a PyYAML parse of examples/rlcr_math/rlcr_math.yaml (YAML_OK), and git diff --check <base>..<head> (exit 0, no output).
• Native full suite (Linux, isolated venv with [dev], [data], and [tinker] extras, CPU-only) on the patched checkout of 9c29c18b: 358 passed / 82 failed / 5 skipped / 2 errors / 176 subtests passed (38 min). The RLCR reward tests pass natively: 62/62 inside the full run and in a standalone rerun.
• Failure attribution against clean upstream: the identical failure set (82 failed + 2 errors) reproduces on the unpatched 9c29c18b baseline in the same venv (git stash swap, same paths). Zero failures are attributable to this patch. All of them require GPU training stacks (fsdp2/megatron), sglang, real model weights, or ray timing not available in a CPU-only environment; one timing-flaky scheduler test was observed on both trees. Node-by-node attribution list is retained by the submitter.
• tests/common/vllm_test.py was excluded: it calls get_model_path() at collection time and raises EnvironmentError without TRINITY_MODEL_PATH pointing at real model weights (tests/tools.py:196-201), and its vLLM engine tests target GPU inference.
• All configured pre-commit hooks also pass natively on the same box (12/12, including mypy).

AI assistance disclosure

OpenAI Codex assisted with repository inspection, the initial implementation, behavioral-test generation, and adversarial review. Anthropic Claude assisted with the native test execution, failure attribution against clean upstream, and drafting this PR description. The submitting human remains responsible for reviewing every changed line and deciding whether to submit the patch. Test evidence above reports actual command outputs, not AI-generated summaries.

Checklist

• Code has passed all tests — native full suite run on Linux: RLCR tests 62/62; zero patch-attributable failures (the 82 environmental failures reproduce identically on clean upstream; see Test evidence)
• Docstrings have been added/updated in Google Style
• Documentation has been updated
• Code is ready for review