diff --git a/examples/rlcr_math/README.md b/examples/rlcr_math/README.md new file mode 100644 index 00000000000..36ddb6bbfbb --- /dev/null +++ b/examples/rlcr_math/README.md @@ -0,0 +1,29 @@ +# RLCR reward on math tasks + +This example wires the five-component RLCR reward into Trinity's standard +`math_workflow`. The dataset fields are explicit: `prompt` supplies the user +prompt and `solution` supplies the ground-truth response passed to the reward as +`truth`. The reward extracts only the final `` payload before parsing and +verification. + +Set the model and checkpoint paths as needed, then launch the config with the +usual Trinity entry point: + +```bash +TRINITY_MODEL_PATH=/path/to/model \ +trinity run --config examples/rlcr_math/rlcr_math.yaml +``` + +The system prompt requires this case-sensitive terminal sequence: + +```text +.........q +``` + +Only whitespace may follow ``, and `q` must be finite and in +`[0, 1]`. `rlcr_reward` returns five already weighted floats because the standard +workflow sums reward-dictionary values without applying per-component weights. + +The weights in `rlcr_math.yaml` are runnable configuration defaults, not frozen +research parameters. Any study using this example should preregister its own +weights and other training choices before launch. diff --git a/examples/rlcr_math/rlcr_math.yaml b/examples/rlcr_math/rlcr_math.yaml new file mode 100644 index 00000000000..7e06796f1f3 --- /dev/null +++ b/examples/rlcr_math/rlcr_math.yaml @@ -0,0 +1,69 @@ +project: Trinity-RFT-example +name: rlcr-math +checkpoint_root_dir: ${oc.env:TRINITY_CHECKPOINT_ROOT_DIR,./checkpoints} +model: + model_path: ${oc.env:TRINITY_MODEL_PATH,Qwen/Qwen2.5-7B-Instruct} + max_response_tokens: 3072 + max_model_len: 4096 +algorithm: + algorithm_type: grpo + repeat_times: 8 + optimizer: + lr: 5e-7 +cluster: + node_num: 1 + gpu_per_node: 8 +buffer: + total_epochs: 1 + batch_size: 96 + explorer_input: + taskset: + name: rlcr-math + storage_type: file + path: open-r1/DAPO-Math-17k-Processed + subset_name: all + split: train + format: + prompt_key: prompt + response_key: solution + system_prompt: |- + Solve the math problem and finish with exactly this four-tag sequence: + your reasoningyour final answercheck whether the final answer is correctq + Use the tags exactly as written and put nothing but whitespace after . + Set q to a finite number from 0 to 1 representing the probability that your final answer is correct. + rollout_args: + temperature: 1.0 + logprobs: 0 + default_workflow_type: math_workflow + default_reward_fn_type: rlcr_reward + reward_fn_args: + weights: + format: 0.5 + accuracy: 0.5 + brier: 0.5 + mean_confidence: 1.0e-5 + confidence_one_or_zero: 1.0e-5 + trainer_input: + experience_buffer: + name: rlcr_math_buffer + storage_type: queue +explorer: + eval_interval: 10 + runner_per_model: 8 + rollout_model: + engine_num: 2 + tensor_parallel_size: 1 + enable_prefix_caching: false + enforce_eager: true + dtype: bfloat16 + seed: 42 +synchronizer: + sync_method: nccl + sync_interval: 1 + sync_timeout: 1200 +trainer: + save_interval: 100 + grad_clip: 1.0 + use_dynamic_bsz: true + max_token_len_per_gpu: 16384 + ulysses_sequence_parallel_size: 1 diff --git a/tests/common/rewards/test_rlcr_reward.py b/tests/common/rewards/test_rlcr_reward.py new file mode 100644 index 00000000000..c80bd83cbd2 --- /dev/null +++ b/tests/common/rewards/test_rlcr_reward.py @@ -0,0 +1,451 @@ +# -*- coding: utf-8 -*- +"""Behavioral contract tests for the RLCR reward function.""" + +import math + +import pytest + +from trinity.common.rewards import REWARD_FUNCTIONS +from trinity.common.rewards.reward_fn import RewardFn +from trinity.common.rewards.rlcr_reward import RLCRRewardFn, _parse_terminal + + +def _identity(value: str) -> str: + return value.strip() + + +def _equal(answer: str, truth: str) -> bool: + return answer == truth + + +def _reward(**kwargs) -> RLCRRewardFn: + return RLCRRewardFn(answer_parser=_identity, verifier=_equal, **kwargs) + + +def _response( + answer: str = "42", + confidence: str = "0.8", + *, + prefix: str = "", + separator: str = "", +) -> str: + return ( + f"{prefix}work{separator}" + f"{answer}{separator}" + f"check{separator}" + f"{confidence}" + ) + + +def _assert_all_zero(scores: dict[str, float]) -> None: + assert scores == { + "format": 0.0, + "accuracy": 0.0, + "brier": 0.0, + "mean_confidence": 0.0, + "confidence_one_or_zero": 0.0, + } + + +def test_registry_resolves_rlcr_lazily() -> None: + reward_cls = REWARD_FUNCTIONS.get("rlcr_reward") + + assert reward_cls is RLCRRewardFn + assert issubclass(reward_cls, RewardFn) + + +def test_correct_answer_returns_five_weighted_components() -> None: + scores = _reward()(response=_response(), truth="42") + + assert scores == pytest.approx( + { + "format": 0.5, + "accuracy": 0.5, + "brier": 0.5 * (1.0 - (1.0 - 0.8) ** 2), + "mean_confidence": 1e-5 * 0.8, + "confidence_one_or_zero": 0.0, + } + ) + assert all(type(value) is float for value in scores.values()) + assert sum(scores.values()) == pytest.approx(1.480008) + + +def test_wrong_answer_uses_the_same_brier_formula() -> None: + scores = _reward()(response=_response(answer="41", confidence="0.2"), truth="42") + + assert scores == pytest.approx( + { + "format": 0.5, + "accuracy": 0.0, + "brier": 0.5 * (1.0 - (0.0 - 0.2) ** 2), + "mean_confidence": 1e-5 * 0.2, + "confidence_one_or_zero": 0.0, + } + ) + + +def test_default_parser_and_timeout_aware_verifier_path() -> None: + scores = RLCRRewardFn()(response=_response(answer="42", confidence="1"), truth="42") + + assert scores["accuracy"] == pytest.approx(0.5) + assert scores["brier"] == pytest.approx(0.5) + + +def test_parser_receives_extracted_answer_payload_and_truth_separately() -> None: + parser_calls: list[str] = [] + verifier_calls: list[tuple[str, str]] = [] + + def parser(value: str) -> str: + parser_calls.append(value) + return value.strip() + + def verifier(answer: str, truth: str) -> bool: + verifier_calls.append((answer, truth)) + return True + + response = _response( + answer="final", + confidence="0.7", + prefix="decoy0.1", + ) + scores = RLCRRewardFn(answer_parser=parser, verifier=verifier)( + response=response, + truth="gold", + source=["code", "math"], + ) + + assert parser_calls == ["final", "gold"] + assert verifier_calls == [("final", "gold")] + assert scores["accuracy"] == pytest.approx(0.5) + + +def test_complete_earlier_reserved_tags_are_allowed_and_terminal_values_win() -> None: + prefix = ( + "draft old workwrong" + "old check0.1 revised: " + ) + + scores = _reward()(response=_response(prefix=prefix, separator=" \n\t"), truth="42") + + assert scores["format"] == pytest.approx(0.5) + assert scores["accuracy"] == pytest.approx(0.5) + assert scores["mean_confidence"] == pytest.approx(8e-6) + + +@pytest.mark.parametrize( + "response", + [ + "work42x0.8", + "xx420.8", + "x42x0.8xnested" + _response(), + "x42x0.8", + "unclosed " + _response(), + "x42nested" + "x0.8", + "x42x0.8", + ], + ids=[ + "missing-think-opener", + "wrong-order", + "truncated-confidence", + "trailing-junk", + "nested-prefix", + "crossed-tags", + "unbalanced-prefix", + "nested-terminal", + "case-sensitive", + ], +) +def test_malformed_terminal_chain_fails_the_three_main_components(response: str) -> None: + scores = _reward()(response=response, truth="42") + + assert scores["format"] == 0.0 + assert scores["accuracy"] == 0.0 + assert scores["brier"] == 0.0 + + +def test_only_trailing_whitespace_is_allowed_after_terminal_chain() -> None: + scores = _reward()(response=_response() + " \r\n\t", truth="42") + + assert scores["format"] == pytest.approx(0.5) + + +@pytest.mark.parametrize( + ("response", "reason"), + [ + ( + "x42x", + "missing_tag", + ), + ( + "x42x" "0.8", + "unbalanced_tags", + ), + ( + "x42x" + "0.8", + "crossed_tags", + ), + ( + "xnested" + _response(), + "nested_tags", + ), + ( + "xx42" + "0.8", + "wrong_order", + ), + (_response() + " trailing junk", "trailing_junk"), + (_response(confidence=""), "confidence_empty"), + (_response(confidence="not-a-number"), "confidence_non_numeric"), + (_response(confidence="NaN"), "confidence_non_finite"), + (_response(confidence="1.01"), "confidence_out_of_range"), + ], + ids=[ + "missing", + "unbalanced", + "crossed", + "nested", + "wrong-order", + "trailing-junk", + "empty", + "non-numeric", + "non-finite", + "out-of-range", + ], +) +def test_terminal_parser_exposes_stable_structured_failure_reasons( + response: str, reason: str +) -> None: + result = _parse_terminal(response) + + assert result.ok is False + assert result.answer is None + assert result.confidence is None + assert result.reason == reason + + +def test_terminal_parser_success_has_no_failure_reason() -> None: + result = _parse_terminal(_response(answer="42", confidence="0.8")) + + assert result.ok is True + assert result.answer == "42" + assert result.confidence == pytest.approx(0.8) + assert result.reason is None + + +@pytest.mark.parametrize( + "confidence", + ["", "not-a-number", "NaN", "Inf", "+inf", "-inf", "-0.01", "1.01"], +) +def test_invalid_confidence_zeros_every_component(confidence: str) -> None: + scores = _reward()(response=_response(confidence=confidence), truth="42") + + _assert_all_zero(scores) + + +def test_safe_confidence_still_drives_micros_when_terminal_format_is_invalid() -> None: + response = "42x1" + + scores = _reward()(response=response, truth="42") + + assert scores == pytest.approx( + { + "format": 0.0, + "accuracy": 0.0, + "brier": 0.0, + "mean_confidence": 1e-5, + "confidence_one_or_zero": 1e-5, + } + ) + + +def test_a_later_truncated_confidence_makes_the_micro_confidence_unsafe() -> None: + response = _response(confidence="0.8") + "0.1" + + scores = _reward()(response=response, truth="42") + + _assert_all_zero(scores) + + +@pytest.mark.parametrize( + "response", + [ + "42x" "01", + "42x" "1", + "42x" "1", + "42x" "01", + ], + ids=["nested", "orphan-closer", "leading-orphan-closer", "truncated-last"], +) +def test_malformed_confidence_structure_never_earns_micro_rewards(response: str) -> None: + scores = _reward()(response=response, truth="42") + + _assert_all_zero(scores) + + +def test_multiple_complete_confidence_pairs_keep_the_last_safe_q_without_think() -> None: + response = ( + "0.2draft" + "42x1" + ) + + scores = _reward()(response=response, truth="42") + + assert scores == pytest.approx( + { + "format": 0.0, + "accuracy": 0.0, + "brier": 0.0, + "mean_confidence": 1e-5, + "confidence_one_or_zero": 1e-5, + } + ) + + +@pytest.mark.parametrize( + ("confidence", "expected"), + [("0", 1e-5), ("0.009", 1e-5), ("0.01", 0.0), ("0.99", 0.0), ("0.991", 1e-5), ("1", 1e-5)], +) +def test_extreme_confidence_uses_strict_open_thresholds(confidence: str, expected: float) -> None: + scores = _reward()(response=_response(confidence=confidence), truth="42") + + assert scores["confidence_one_or_zero"] == pytest.approx(expected) + + +@pytest.mark.parametrize("stage", ["answer", "truth", "verify"]) +def test_verification_errors_zero_accuracy_and_brier_without_low_q_reward(stage: str) -> None: + def parser(value: str) -> str: + if stage == "answer" and value == "42": + raise TimeoutError("answer parse timed out") + if stage == "truth" and value == "gold": + raise ValueError("truth parse failed") + return value + + def verifier(answer: str, truth: str) -> bool: + if stage == "verify": + raise TimeoutError("verify timed out") + return answer == truth + + scores = RLCRRewardFn(answer_parser=parser, verifier=verifier)( + response=_response(answer="42", confidence="0"), + truth="gold", + ) + + assert scores == pytest.approx( + { + "format": 0.5, + "accuracy": 0.0, + "brier": 0.0, + "mean_confidence": 0.0, + "confidence_one_or_zero": 1e-5, + } + ) + + +def test_missing_truth_is_a_verification_error_not_an_incorrect_answer() -> None: + scores = _reward()(response=_response(confidence="0"), truth=None) + + assert scores["format"] == pytest.approx(0.5) + assert scores["accuracy"] == 0.0 + assert scores["brier"] == 0.0 + + +def test_custom_weights_are_matched_by_name_not_position() -> None: + weights = { + "confidence_one_or_zero": 5.0, + "mean_confidence": 4.0, + "brier": 3.0, + "accuracy": 2.0, + "format": 1.0, + } + + scores = _reward(weights=weights)(response=_response(confidence="1"), truth="42") + + assert scores == pytest.approx( + { + "format": 1.0, + "accuracy": 2.0, + "brier": 3.0, + "mean_confidence": 4.0, + "confidence_one_or_zero": 5.0, + } + ) + assert sum(scores.values()) == pytest.approx(15.0) + + +@pytest.mark.parametrize( + ("weights", "message"), + [ + ( + {"format": 1, "accuracy": 1, "brier": 1, "mean_confidence": 1}, + "missing", + ), + ( + { + "format": 1, + "accuracy": 1, + "brier": 1, + "mean_confidence": 1, + "confidence_one_or_zero": 1, + "extra": 1, + }, + "unknown", + ), + ( + { + "format": "0.5", + "accuracy": 1, + "brier": 1, + "mean_confidence": 1, + "confidence_one_or_zero": 1, + }, + "numeric", + ), + ( + { + "format": math.nan, + "accuracy": 1, + "brier": 1, + "mean_confidence": 1, + "confidence_one_or_zero": 1, + }, + "finite", + ), + ( + { + "format": math.inf, + "accuracy": 1, + "brier": 1, + "mean_confidence": 1, + "confidence_one_or_zero": 1, + }, + "finite", + ), + ( + { + "format": True, + "accuracy": 1, + "brier": 1, + "mean_confidence": 1, + "confidence_one_or_zero": 1, + }, + "numeric", + ), + ], +) +def test_invalid_named_weight_mappings_fail_deterministically( + weights: dict[str, object], message: str +) -> None: + with pytest.raises((TypeError, ValueError), match=message): + _reward(weights=weights) # type: ignore[arg-type] + + +@pytest.mark.parametrize("bad_callable", [None, 3, "callable"]) +def test_injected_parser_and_verifier_must_be_callable(bad_callable: object) -> None: + with pytest.raises(TypeError, match="answer_parser"): + RLCRRewardFn(answer_parser=bad_callable, verifier=_equal) # type: ignore[arg-type] + with pytest.raises(TypeError, match="verifier"): + RLCRRewardFn(answer_parser=_identity, verifier=bad_callable) # type: ignore[arg-type] diff --git a/trinity/common/rewards/__init__.py b/trinity/common/rewards/__init__.py index b20bc7a6d57..e14a748cbe3 100644 --- a/trinity/common/rewards/__init__.py +++ b/trinity/common/rewards/__init__.py @@ -13,6 +13,7 @@ "countdown_reward": "trinity.common.rewards.countdown_reward.CountDownRewardFn", "accuracy_reward": "trinity.common.rewards.accuracy_reward.AccuracyReward", "math_dapo_reward": "trinity.common.rewards.dapo_reward.MathDAPORewardFn", + "rlcr_reward": "trinity.common.rewards.rlcr_reward.RLCRRewardFn", }, ) diff --git a/trinity/common/rewards/rlcr_reward.py b/trinity/common/rewards/rlcr_reward.py new file mode 100644 index 00000000000..7f9139b460d --- /dev/null +++ b/trinity/common/rewards/rlcr_reward.py @@ -0,0 +1,245 @@ +# -*- coding: utf-8 -*- +"""A safe port of the five-component RLCR reward for math tasks.""" + +import math as _math +import re as _re +from collections.abc import Callable as _Callable +from collections.abc import Mapping as _Mapping +from dataclasses import dataclass as _dataclass +from numbers import Real as _Real +from typing import Any as _Any + +from trinity.common.rewards.eval_utils import ( + simple_answer_parser as _simple_answer_parser, +) +from trinity.common.rewards.eval_utils import ( + verify_with_timeout as _verify_with_timeout, +) +from trinity.common.rewards.reward_fn import RewardFn as _RewardFn + +__all__ = ["RLCRRewardFn"] + +_COMPONENTS = ( + "format", + "accuracy", + "brier", + "mean_confidence", + "confidence_one_or_zero", +) +_DEFAULT_WEIGHTS = { + "format": 0.5, + "accuracy": 0.5, + "brier": 0.5, + "mean_confidence": 1e-5, + "confidence_one_or_zero": 1e-5, +} +_RESERVED_TAG_RE = _re.compile(r"") +_CONFIDENCE_TAG_RE = _re.compile(r"") +_TERMINAL_TAGS = ( + "", + "", + "", + "", + "", + "", + "", + "", +) + + +@_dataclass(frozen=True) +class _ConfidenceParseResult: + confidence: float | None + reason: str | None + + +@_dataclass(frozen=True) +class _TerminalParseResult: + ok: bool + answer: str | None + confidence: float | None + reason: str | None + + +def _terminal_failure(reason: str) -> _TerminalParseResult: + return _TerminalParseResult(False, None, None, reason) + + +def _parse_finite_confidence(payload: str) -> _ConfidenceParseResult: + payload = payload.strip() + if not payload: + return _ConfidenceParseResult(None, "confidence_empty") + try: + confidence = float(payload) + except (TypeError, ValueError, OverflowError): + return _ConfidenceParseResult(None, "confidence_non_numeric") + if not _math.isfinite(confidence): + return _ConfidenceParseResult(None, "confidence_non_finite") + if not 0.0 <= confidence <= 1.0: + return _ConfidenceParseResult(None, "confidence_out_of_range") + return _ConfidenceParseResult(confidence, None) + + +def _scan_confidence(response: str) -> _ConfidenceParseResult: + """Return the last q only when every confidence tag is balanced and unnested.""" + if not isinstance(response, str): + return _ConfidenceParseResult(None, "response_not_string") + + open_tag: _re.Match[str] | None = None + last_payload: str | None = None + for tag in _CONFIDENCE_TAG_RE.finditer(response): + if tag.group(0) == "": + if open_tag is not None: + return _ConfidenceParseResult(None, "nested_tags") + open_tag = tag + else: + if open_tag is None: + return _ConfidenceParseResult(None, "unbalanced_tags") + last_payload = response[open_tag.end() : tag.start()] + open_tag = None + + if open_tag is not None: + return _ConfidenceParseResult(None, "unbalanced_tags") + if last_payload is None: + return _ConfidenceParseResult(None, "missing_tag") + return _parse_finite_confidence(last_payload) + + +def _parse_safe_confidence(response: str) -> float | None: + """Parse the final closed confidence even if the full terminal chain is invalid.""" + return _scan_confidence(response).confidence + + +def _parse_terminal(response: str) -> _TerminalParseResult: + """Validate the strict terminal tag chain and return its answer and confidence.""" + if not isinstance(response, str): + return _terminal_failure("response_not_string") + + tags = list(_RESERVED_TAG_RE.finditer(response)) + open_name: str | None = None + for tag in tags: + token = tag.group(0) + is_close = token.startswith(" dict[str, float]: + if weights is None: + return dict(_DEFAULT_WEIGHTS) + if not isinstance(weights, _Mapping): + raise TypeError("weights must be a named mapping") + + provided = set(weights) + expected = set(_COMPONENTS) + missing = sorted(expected - provided) + unknown = sorted(provided - expected) + if missing or unknown: + details = [] + if missing: + details.append(f"missing keys: {missing}") + if unknown: + details.append(f"unknown keys: {unknown}") + raise ValueError("invalid weights mapping; " + "; ".join(details)) + + validated: dict[str, float] = {} + for name in _COMPONENTS: + value = weights[name] + if isinstance(value, bool) or not isinstance(value, _Real): + raise TypeError(f"weight '{name}' must be numeric") + numeric = float(value) + if not _math.isfinite(numeric): + raise ValueError(f"weight '{name}' must be finite") + validated[name] = numeric + return validated + + +class RLCRRewardFn(_RewardFn): + """Return already weighted RLCR format, accuracy, and calibration components.""" + + def __init__( + self, + weights: _Mapping[str, _Any] | None = None, + answer_parser: _Callable[[str], _Any] = _simple_answer_parser, + verifier: _Callable[[_Any, _Any], bool] = _verify_with_timeout, + ) -> None: + if not callable(answer_parser): + raise TypeError("answer_parser must be callable") + if not callable(verifier): + raise TypeError("verifier must be callable") + self._weights = _validate_weights(weights) + self._answer_parser = answer_parser + self._verifier = verifier + + def __call__( # type: ignore[override] + self, + response: str, + prompt: str | None = None, + truth: str | None = None, + **kwargs: _Any, + ) -> dict[str, float]: + del prompt, kwargs + safe_confidence = _parse_safe_confidence(response) + mean_confidence = safe_confidence if safe_confidence is not None else 0.0 + confidence_one_or_zero = float( + safe_confidence is not None and (safe_confidence < 0.01 or safe_confidence > 0.99) + ) + + terminal = _parse_terminal(response) + format_score = float(terminal.ok) + accuracy_score = 0.0 + brier_score = 0.0 + + if terminal.ok: + try: + if terminal.answer is None or terminal.confidence is None or truth is None: + raise ValueError("answer, confidence, and truth are required for verification") + parsed_answer = self._answer_parser(terminal.answer) + parsed_truth = self._answer_parser(str(truth)) + accuracy_score = float(bool(self._verifier(parsed_answer, parsed_truth))) + brier_score = 1.0 - (accuracy_score - terminal.confidence) ** 2 + except Exception: + # An infrastructure failure is not evidence that the answer is wrong. + accuracy_score = 0.0 + brier_score = 0.0 + + return { + "format": float(self._weights["format"] * format_score), + "accuracy": float(self._weights["accuracy"] * accuracy_score), + "brier": float(self._weights["brier"] * brier_score), + "mean_confidence": float(self._weights["mean_confidence"] * mean_confidence), + "confidence_one_or_zero": float( + self._weights["confidence_one_or_zero"] * confidence_one_or_zero + ), + }