trinity.common.rewards.dapo_reward module

trinity.common.rewards.dapo_reward module#

Reward Function with Overlong Reward Shaping described in DAPO (https://arxiv.org/pdf/2503.14476)

class trinity.common.rewards.dapo_reward.MathDAPORewardFn(enable_overlong_penalty: bool | None = None, penalty_factor: float | None = None, max_response_length: int | None = None, cache_length: int | None = None)[source]#

Bases: RewardFn

A reward function that follows the definition in DAPO for math task.

__init__(enable_overlong_penalty: bool | None = None, penalty_factor: float | None = None, max_response_length: int | None = None, cache_length: int | None = None) → None[source]#

Initialize DAPO math reward settings.

Parameters:

  • enable_overlong_penalty – Whether to apply overlong response shaping.

  • penalty_factor – Magnitude for overlong penalties.

  • max_response_length – Maximum allowed response length in tokens.

  • cache_length – Soft-penalty transition window in tokens.

compute_overlong_penalty(response_token)[source]#

Compute soft/hard penalty for long responses.

Parameters:

response_token – Response token ids.

Returns:

Length-based shaping value, where negative values penalize overlong outputs.

Return type:

float