feat: optionally track other reward components

[Tim-Siu]

Motivation

For R1 style of RL, there are many components of reward, like format score or reward score. It is currently not trivial to track different components. It is especially inconvenient when reading validation results, when we might want to read the outcome accuracy instead of only reward score.

Design

In this PR I slightly modify verl/workers/reward_manager/naive.py to expect the compute_score function to either expect a float (final reward) or a dict (new).

The dict returned should be in the format

{"score": 1.0, "extra_info": {"format": 1.0, "answer_accuracy": 1.0}}

ray_trainer.py is also modified to handle both dict and float returned from reward_manager. It will then log the extra info from the reward function.

Example Usage

We can modify verl/verl/utils/reward_score/gsm8k.py to allow separate logging of format rewards

def compute_score(solution_str, ground_truth, method='strict', format_score=0., score=1.):
    """The scoring function for GSM8k.

    Reference: Trung, Luong, et al. "Reft: Reasoning with reinforced fine-tuning." Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers). 2024.

    Args:
        solution_str: the solution text
        ground_truth: the ground truth
        method: the method to extract the solution, choices are 'strict' and 'flexible'
        format_score: the score for the format
        score: the score for the correct answer
    """
    answer = extract_solution(solution_str=solution_str, method=method)
    if answer == ground_truth:
        return {"score": score, "extra_info": {"format": 1.0, "answer_accuracy": 1.0}}
    else:
        return {"score": format_score, "extra_info": {"format": 1.0, "answer_accuracy": 0.0}}

Breaking changes

There are no breaking changes in this PR. It is still possible to use the reward functions under utils/reward_score

Possible future work

Logging of reward and validation can be enhanced further.

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.
_{You have signed the CLA already but the status is still pending? Let us recheck it.}