`nemo_rl.utils.length_penalty`#

Per-prompt-group length bonuses/penalties for rollout rewards.

Rewards conciseness among high-quality generations by applying:

A flat bonus to the shortest generation among top scorers in each prompt group.
Optional flat penalties on the longest reasoning / longest answer among top-percentile scorers (with at least two eligible rollouts to compare).
Independent zero-centered penalties for reasoning and answer length.
Optional Iglewicz–Hoaglin modified-Z (MAD) high-side outliers among positive scorers. Config keys: reasoning_zmad_threshold, reasoning_zmad_penalty, answer_zmad_threshold, answer_zmad_penalty. If reasoning_zmad_threshold or answer_zmad_threshold is ≤ 0, that channel is off and its penalty is ignored (no flagging). If threshold > 0 but the matching penalty is 0, nothing is subtracted.

Supports per-agent filtering and parameter overrides via config.

Module Contents#

Classes#

`LengthPenaltyParams`	Length-adjustment parameters for one agent.
`ProfileBandChannel`	One `{a, b, f}` band: multiplier 1.0 up to `a`, linear to `f` at `b`.
`ProfileBandDefaults`
`ProfileBandConfig`	Global `{a, b, f}` bands for datasets without per-row `profile_band`.
`LengthPenaltyConfig`	`grpo.length_penalty`: per-prompt-group length penalties for NeMo-Gym rollouts.

Functions#

`_is_binary_reward`
`_extract_reasoning_and_answer_text`	Extract reasoning and answer text from the Response API output items.
`apply_group_length_penalties`	Apply per-prompt-group length bonuses/penalties.
`_resolve_global_profile_band`	Parse `length_penalty.profile_band` into per-channel {a, b, f} blocks.
`_merged_profile_band`	Merge per-row profile_band over global defaults (row channel wins).
`_apply_profile_band_multipliers`	Apply per-channel profile_band multipliers to correct rollouts.
`_band_multiplier`	Per-channel profile_band reward multiplier.
`_group_length_profile_gate_info`	Prompt-level gate for group-relative length penalties.
`_profile_band_numeric_value`
`_resolve_agent_params`	Resolve length bonus parameters for a given agent.
`_zmad_local_outliers`	Indices into `lengths` with Iglewicz–Hoaglin modified Z (MAD) > `z_thresh`.
`_apply_length_penaltyes_and_penalties`	Apply length-based bonuses/penalties to a single prompt group.
`_compute_length_weights`	Compute zero-centered weights where shorter = higher weight.

Data#

`logger`
`_ZMAD_MIN_MAD_REL`
`ProfileBandChannelName`
`_PROFILE_BAND_CHANNELS`
`_NON_BINARY_WARNED_AGENTS`
`_BINARY_REWARD_TOL`

API#

nemo_rl.utils.length_penalty.logger#: ‘getLogger(…)’

nemo_rl.utils.length_penalty._ZMAD_MIN_MAD_REL#: 0.015

nemo_rl.utils.length_penalty.ProfileBandChannelName#: None

nemo_rl.utils.length_penalty._PROFILE_BAND_CHANNELS: tuple[nemo_rl.utils.length_penalty.ProfileBandChannelName, ...]#: (‘total’, ‘reasoning’, ‘answer’)

class nemo_rl.utils.length_penalty.LengthPenaltyParams#

Bases: pydantic.BaseModel

Length-adjustment parameters for one agent.

Used for grpo.length_penalty.default and for each entry under grpo.length_penalty.agent_overrides. Only the keys a user sets in an override replace the corresponding default value. Unknown keys are forbidden because a misspelled parameter would silently leave a penalty off. See docs/guides/length-penalty.md for the per-key semantics.

enabled: bool#: True

length_type: Literal[tokens, chars]#: ‘tokens’

top_percentile: float#: 0.5

reasoning_bonus: float#: 0.0

answer_bonus: float#: 0.0

total_bonus: float#: 0.0

longest_reasoning_penalty: float#: 0.0

longest_answer_penalty: float#: 0.0

longest_total_penalty: float#: 0.0

group_reasoning_length_penalty_coeff: float#: 0.0

group_answer_length_penalty_coeff: float#: 0.0

group_total_length_penalty_coeff: float#: 0.0

reasoning_zmad_threshold: float#: 0.0

reasoning_zmad_penalty: float#: 0.0

answer_zmad_threshold: float#: 0.0

answer_zmad_penalty: float#: 0.0

total_zmad_threshold: float#: 0.0

total_zmad_penalty: float#: 0.0

profiled_length_penalty: float#: 0.0

profiled_length_n_std: float#: 1.0

profiled_length_min_samples: int#: 2

pass_rate_length_penalty_weight: float#: 0.0

profile_band_total: bool#: False

profile_band_reasoning: bool#: False

profile_band_answer: bool#: False

group_length_penalty_profile_gate: bool#: False

group_length_penalty_profile_gate_channel: nemo_rl.utils.length_penalty.ProfileBandChannelName#: ‘total’

group_length_penalty_profile_gate_field: Literal[a, b, f]#: ‘a’

group_length_penalty_profile_gate_positive_only: bool#: True

class nemo_rl.utils.length_penalty.ProfileBandChannel#

Bases: pydantic.BaseModel

One {a, b, f} band: multiplier 1.0 up to a, linear to f at b.

f is bounded to [0, 1]: a value outside it would make rewards negative or larger for longer rollouts.

a: float#: None

b: float#: None

f: float#: ‘Field(…)’

validate_b_gt_a() → nemo_rl.utils.length_penalty.ProfileBandChannel#

class nemo_rl.utils.length_penalty.ProfileBandDefaults#

Bases: pydantic.BaseModel

total: Optional[nemo_rl.utils.length_penalty.ProfileBandChannel]#: None

reasoning: Optional[nemo_rl.utils.length_penalty.ProfileBandChannel]#: None

answer: Optional[nemo_rl.utils.length_penalty.ProfileBandChannel]#: None

class nemo_rl.utils.length_penalty.ProfileBandConfig#

Bases: pydantic.BaseModel

Global {a, b, f} bands for datasets without per-row profile_band.

enabled: bool#: False

defaults: nemo_rl.utils.length_penalty.ProfileBandDefaults#: ‘Field(…)’

class nemo_rl.utils.length_penalty.LengthPenaltyConfig#

Bases: pydantic.BaseModel

grpo.length_penalty: per-prompt-group length penalties for NeMo-Gym rollouts.

agent_overrides.<agent> may be null to keep default for that agent. Channels listed under profile_band.defaults are implicitly enabled unless default sets the matching profile_band_* flag.

verbose: bool#: False

default: nemo_rl.utils.length_penalty.LengthPenaltyParams#: ‘Field(…)’

agent_overrides: Optional[dict[str, Optional[nemo_rl.utils.length_penalty.LengthPenaltyParams]]]#: None

profile_band: Optional[nemo_rl.utils.length_penalty.ProfileBandConfig]#: None

nemo_rl.utils.length_penalty._NON_BINARY_WARNED_AGENTS: set[str]#: ‘set(…)’

nemo_rl.utils.length_penalty._BINARY_REWARD_TOL#: 1e-06

nemo_rl.utils.length_penalty._is_binary_reward(value: float) → bool#

nemo_rl.utils.length_penalty._extract_reasoning_and_answer_text( result: dict[str, Any], ) → tuple[str, str]#: Extract reasoning and answer text from the Response API output items.

nemo_rl.utils.length_penalty.apply_group_length_penalties( results: list[dict[str, Any]], length_penalty_config: nemo_rl.utils.length_penalty.LengthPenaltyConfig | dict[str, Any], group_size: int, tokenizer: Any = None, ) → dict[str, float]#

Apply per-prompt-group length bonuses/penalties.

Mutates full_result["reward"] in place. No-ops when no length-adjustment feature is enabled.

Parameters:

results – List of per-generation result dicts, group_size contiguous rows per prompt group.
length_penalty_config – The grpo.length_penalty block (a dict is validated into :class:LengthPenaltyConfig).
group_size – Number of contiguous rows per prompt group.
tokenizer – Tokenizer for computing reasoning/answer token counts.

Returns:

length_penalty/* rollout metrics, all per-row means over results (pre-adjustment env reward, reward delta, rollouts whose correct reward was wiped to 0, rollouts adjusted, rollouts in groups skipped for non-binary or multi-component rewards). Per-row so the async path, which calls this once per prompt group and averages the per-group values, reports the same numbers as the batched paths. Empty when the block enables nothing.

nemo_rl.utils.length_penalty._resolve_global_profile_band( pb_cfg: nemo_rl.utils.length_penalty.ProfileBandConfig | None, ) → dict[str, dict[str, Any]]#

Parse length_penalty.profile_band into per-channel {a, b, f} blocks.

Returns only the channels (“total”, “reasoning”, “answer”) present under defaults. Empty dict when the section is absent or disabled.

nemo_rl.utils.length_penalty._merged_profile_band( row_band: dict[str, Any] | None, global_band: dict[str, dict[str, Any]], ) → dict[str, Any] | None#: Merge per-row profile_band over global defaults (row channel wins).

nemo_rl.utils.length_penalty._apply_profile_band_multipliers( results: list[dict[str, Any]], original_rewards: list[float], base_rewards: list[float], total_lengths: list[int], reasoning_lengths: list[int], answer_lengths: list[int], agent_names: list[str], agents_cfg: dict[str, nemo_rl.utils.length_penalty.LengthPenaltyParams | None] | None, defaults: dict[str, Any], num_gens: int, global_band: dict[str, dict[str, Any]] | None = None, binary_ok: list[bool] | None = None, ) → None#

Apply per-channel profile_band multipliers to correct rollouts.

Each enabled channel contributes a multiplier in [0.0, 1.0] derived from the per-row {a, b, f} block, falling back to length_penalty.profile_band.defaults for channels the row does not provide. Mutates scalar rewards in place.

Skips any group where the low-effort bypass already replaced the reward (parity with Phase 1 of apply_group_length_penalties).

nemo_rl.utils.length_penalty._band_multiplier(rl: int, ch: dict[str, Any] | None) → float#

Per-channel profile_band reward multiplier.

Returns 1.0 if the channel block is missing or malformed (no-op). Otherwise: rl <= a -> 1.0 a < rl < b -> linear interpolation from 1.0 down to f rl >= b -> f

nemo_rl.utils.length_penalty._group_length_profile_gate_info( *, band: dict[str, Any] | None, params: dict[str, Any], rewards: list[float], reasoning_lengths: list[int], answer_lengths: list[int], total_lengths: list[int], ) → dict[str, Any]#

Prompt-level gate for group-relative length penalties.

When enabled, group-relative coefficients are applied only if the mean rollout length exceeds a prompt-specific threshold from profile_band.

nemo_rl.utils.length_penalty._profile_band_numeric_value( band: dict[str, Any] | None, channel: str, field: str, ) → float | None#

nemo_rl.utils.length_penalty._resolve_agent_params( agent_name: str, agents_cfg: dict[str, nemo_rl.utils.length_penalty.LengthPenaltyParams | None] | None, defaults: dict[str, Any], ) → dict[str, Any] | None#: Resolve length bonus parameters for a given agent.

nemo_rl.utils.length_penalty._zmad_local_outliers( lengths: list[int], z_thresh: float, min_mad_rel: float, ) → set[int]#: Indices into lengths with Iglewicz–Hoaglin modified Z (MAD) > z_thresh.

nemo_rl.utils.length_penalty._apply_length_penaltyes_and_penalties(

rewards: list[float],

reasoning_lengths: list[int],

answer_lengths: list[int],

total_lengths: list[int],

reasoning_bonus: float,

answer_bonus: float,

total_bonus: float,

longest_reasoning_penalty: float,

longest_answer_penalty: float,

longest_total_penalty: float,

top_percentile: float,

group_reasoning_length_penalty_coeff: float,

group_answer_length_penalty_coeff: float,

group_total_length_penalty_coeff: float,

reasoning_zmad_threshold: float = 0.0,

reasoning_zmad_penalty: float = 0.0,

answer_zmad_threshold: float = 0.0,

answer_zmad_penalty: float = 0.0,

total_zmad_threshold: float = 0.0,

total_zmad_penalty: float = 0.0,

**_kwargs,

) → tuple[list[float], list[float], list[float], list[float], list[float], list[float], list[float], list[float], list[float], list[float], list[float], list[float], list[float], list[float]]#

Apply length-based bonuses/penalties to a single prompt group.

Only samples with reward > 0 participate. Samples with reward <= 0 are left untouched and excluded from weight computation.

Reasoning bonus: shortest non-empty reasoning among positive scorers; awarded only if that sample satisfies reward >= top_threshold.
Answer bonus: same pattern for shortest non-empty answer.
Total bonus: same pattern for shortest combined (reasoning + answer) length.
Longest penalties: subtract from longest non-empty reasoning / answer / total among top-percentile scorers; needs at least two eligible rollouts to compare.
Independent zero-centered penalties for reasoning, answer, and total lengths.
Optional MAD modified-Z outliers among positives for reasoning, answer, and total lengths. Each channel runs only if its threshold is > 0; otherwise that channel is disabled and its penalty is ignored.

nemo_rl.utils.length_penalty._compute_length_weights(lengths: list[int]) → list[float]#

Compute zero-centered weights where shorter = higher weight.

Returns all zeros if all lengths are equal.

nemo_rl.utils.length_penalty#

Module Contents#

Classes#

Functions#

Data#

API#

`nemo_rl.utils.length_penalty`#