Function reward_function

countdown_task.py:148–166

Reward function for Countdown Tasks. Total reward = 0.1 * format_reward + answer_reward

(
    response: str,
    numbers: List[int] = None,
    target: int = None,
    end_token: str = None,
)

def reward_function(
    response: str,
    numbers: List[int] = None,
    target: int = None,
    end_token: str = None,
) -> Dict[str, Any]:
    """Reward function for Countdown Tasks.

    Total reward = 0.1 * format_reward + answer_reward
    """
    format_reward = format_reward_function("<think>" + response, end_token)
    answer_reward = answer_reward_function(response, numbers, target)
    return {
        "reward": format_reward * 0.1 + answer_reward,
        "reward_info": {
            "format_reward": format_reward,
            "answer_reward": answer_reward,
        },
    }