`nemo_rl.algorithms.reward_functions`#

Module Contents#

Classes#

RewardShapingConfig

Configuration for reward function processing.

Functions#

apply_reward_shaping

Process rewards by applying penalties for responses exceeding max_response_length. Currently, this function only supports DAPO reward shaping as illustrated in the DAPO paper : https://arxiv.org/pdf/2503.14476.

Data#

Tensor

API#

nemo_rl.algorithms.reward_functions.Tensor#: ‘TypeVar(…)’

class nemo_rl.algorithms.reward_functions.RewardShapingConfig#

Bases: pydantic.BaseModel

Configuration for reward function processing.

This configuration enables custom reward shaping, currently supporting DAPO-style penalties for responses that exceed the maximum response length threshold.

enabled: bool#: False

overlong_buffer_length: int | None#: None

overlong_buffer_penalty: float | None#: None

max_response_length: int | None#: None

stop_properly_penalty_coef: float | None#: None

nemo_rl.algorithms.reward_functions.apply_reward_shaping( batch: nemo_rl.distributed.batched_data_dict.BatchedDataDict, cfg: nemo_rl.algorithms.reward_functions.RewardShapingConfig, ) → nemo_rl.distributed.batched_data_dict.BatchedDataDict#