diff --git a/examples/nemo_gym/nemotron-3-super/super_launch.sh b/examples/nemo_gym/nemotron-3-super/super_launch.sh index c058cf2fa38..ea71205cabd 100755 --- a/examples/nemo_gym/nemotron-3-super/super_launch.sh +++ b/examples/nemo_gym/nemotron-3-super/super_launch.sh @@ -70,8 +70,26 @@ fi # ---- Derived paths ---- CODE_DIR=$(realpath "$PWD") WANDB_NAME="${EXP_NAME}" -CHECKPOINT_DIR="results/${EXP_NAME}" -LOG_DIR="logs/${EXP_NAME}" +# Optional stable results layout: set RESULTS_DIR to get +# $RESULTS_DIR/checkpoints (stable -> singleton auto-resume) +# $RESULTS_DIR/runs//{logs,slurm} (per-submission; runs/latest symlink) +# $RESULTS_DIR/ray_logs/-logs (ray.sub infra logs via BASE_LOG_DIR) +# Unset: legacy snapshot-relative results/ and logs/ dirs. +RESULTS_DIR="${RESULTS_DIR:-}" +if [[ -n "${RESULTS_DIR}" ]]; then + CHECKPOINT_DIR="${CHECKPOINT_DIR:-${RESULTS_DIR}/checkpoints}" + RUN_DIR="${RESULTS_DIR}/runs/$(date +%Y%m%d-%H%M)" + LOG_DIR="${RUN_DIR}/logs" + SLURM_LOG_DIR="${RUN_DIR}/slurm" + mkdir -p "${CHECKPOINT_DIR}" "${LOG_DIR}" "${SLURM_LOG_DIR}" + ln -sfn "$(realpath "${RUN_DIR}")" "${RESULTS_DIR}/runs/latest" + export BASE_LOG_DIR="${BASE_LOG_DIR:-${RESULTS_DIR}/ray_logs}" + mkdir -p "${BASE_LOG_DIR}" +else + CHECKPOINT_DIR="results/${EXP_NAME}" + LOG_DIR="logs/${EXP_NAME}" + SLURM_LOG_DIR="" +fi VLLM_CACHE_DIR="${PERSISTENT_CACHE}/vllm_compile_cache" FLASHINFER_CUBIN_CACHE="${PERSISTENT_CACHE}/flashinfer_cubins" @@ -134,7 +152,7 @@ export LISTEN_PORT=6000 export NGINX_PORT=6000 export NEMO_SKILLS_SANDBOX_PORT=6000 export SANDBOX_CONTAINER -export SANDBOX_COMMAND="/start-with-nginx.sh" +export SANDBOX_COMMAND="${SANDBOX_COMMAND:-/start-with-nginx.sh}" export SANDBOX_ENV_VARS="NEMO_SKILLS_SANDBOX_PORT=${NEMO_SKILLS_SANDBOX_PORT}" # ---- Build the run command ---- @@ -160,8 +178,8 @@ export COMMAND="export HF_MODULES_CACHE=${HF_MODULES_CACHE_DIR} ; \ PYTHONPATH=${SNAPSHOT_DIR}:\${PYTHONPATH:-} \ python ./examples/nemo_gym/run_grpo_nemo_gym.py \ --config ${CONFIG_PATH} \ - env.nemo_gym.uv_venv_dir=${GYM_VENV_DIR} \ - env.nemo_gym.skip_venv_if_present=true \ + ++env.nemo_gym.uv_venv_dir=${GYM_VENV_DIR} \ + ++env.nemo_gym.skip_venv_if_present=true \ policy.model_name=${MODEL_PATH} \ checkpointing.checkpoint_dir=${CHECKPOINT_DIR} \ logger.log_dir=${LOG_DIR} \ @@ -216,11 +234,27 @@ SBATCH_CMD=( --job-name="${WANDB_NAME}" --partition="${SLURM_PARTITION}" --time="${SLURM_TIME_LIMIT}" - --gres=gpu:8 + --gres=gpu:"${GPUS_PER_NODE:-8}" --exclusive - --dependency=singleton + --dependency=singleton${SLURM_EXTRA_DEPENDENCY:+,${SLURM_EXTRA_DEPENDENCY}} ray.sub ) +if [[ -n "${SLURM_QOS:-}" ]]; then + SBATCH_CMD=("${SBATCH_CMD[@]:0:1}" --qos="${SLURM_QOS}" "${SBATCH_CMD[@]:1}") +fi +if [[ -n "${SLURM_COMMENT:-}" ]]; then + SBATCH_CMD=("${SBATCH_CMD[@]:0:1}" --comment="${SLURM_COMMENT}" "${SBATCH_CMD[@]:1}") +fi +if [[ -n "${SLURM_LOG_DIR}" ]]; then + SBATCH_CMD=("${SBATCH_CMD[@]:0:1}" --output="${SLURM_LOG_DIR}/%j.out" --error="${SLURM_LOG_DIR}/%j.err" "${SBATCH_CMD[@]:1}") +fi +if [[ -n "${SLURM_SEGMENT:-}" ]]; then + if (( SBATCH_NUM_NODES % SLURM_SEGMENT != 0 )); then + echo "Error: SBATCH_NUM_NODES=${SBATCH_NUM_NODES} not divisible by SLURM_SEGMENT=${SLURM_SEGMENT}" >&2 + exit 1 + fi + SBATCH_CMD=("${SBATCH_CMD[@]:0:1}" --segment="${SLURM_SEGMENT}" "${SBATCH_CMD[@]:1}") +fi if [[ "$DRY_RUN" == true ]]; then echo "" diff --git a/nemo_rl/algorithms/async_utils/trajectory_collector.py b/nemo_rl/algorithms/async_utils/trajectory_collector.py index 2136a593557..dd0ff5bcd94 100644 --- a/nemo_rl/algorithms/async_utils/trajectory_collector.py +++ b/nemo_rl/algorithms/async_utils/trajectory_collector.py @@ -1291,6 +1291,7 @@ async def _iter_rollout_groups( max_rollout_turns=None, greedy=False, reward_penalty_config=self.master_config.reward_penalties, + length_penalty_config=self.master_config.grpo.model_dump(), thinking_tags=get_nemo_gym_thinking_tags(self.master_config.env), mask_env_flagged_samples=should_mask_flagged_samples( self.master_config.env diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index 9639c884293..bff55e9c385 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -3015,6 +3015,7 @@ def grpo_train( greedy=False, effort_config=_get_effort_config(master_config), reward_penalty_config=master_config.reward_penalties, + length_penalty_config=master_config.grpo.model_dump(), thinking_tags=get_nemo_gym_thinking_tags(master_config.env), mask_env_flagged_samples=should_mask_flagged_samples( master_config.env @@ -3938,6 +3939,11 @@ def validate( greedy=False, effort_config=_get_effort_config(master_config), reward_penalty_config=master_config.reward_penalties, + # No length_penalty_config here: validation metrics + # (accuracy/pass_k) must reflect the raw env reward, and the + # adjustment code groups by the TRAINING stride + # (num_generations_per_prompt), which does not match + # val_num_generations_per_prompt. thinking_tags=get_nemo_gym_thinking_tags(master_config.env), mask_env_flagged_samples=should_mask_flagged_samples( master_config.env diff --git a/nemo_rl/experience/rollouts.py b/nemo_rl/experience/rollouts.py index 4b03741f199..b920aebc1b1 100644 --- a/nemo_rl/experience/rollouts.py +++ b/nemo_rl/experience/rollouts.py @@ -68,6 +68,7 @@ GenerationOutputSpec, GenerationSamplingParams, ) +from nemo_rl.utils.length_penalty import apply_group_length_penalties from nemo_rl.utils.multimodal_payload_metrics import ( collect_multimodal_payload_metrics, print_multimodal_payload_metrics, @@ -2271,6 +2272,7 @@ async def run_async_nemo_gym_rollout( greedy: bool = False, effort_config: Optional[EffortLevelsConfig] = None, reward_penalty_config: dict[str, Any] | BaseModel | None = None, + length_penalty_config: dict[str, Any] | BaseModel | None = None, thinking_tags: list[str] | tuple[str, ...] | None = None, mask_env_flagged_samples: bool = True, returns_entire_batch: bool = False, @@ -2302,6 +2304,7 @@ async def run_async_nemo_gym_rollout( greedy: Must be ``False`` because this path does not support greedy mode. effort_config: Optional configuration for effort-based reward shaping. reward_penalty_config: Optional reward-penalty configuration. + length_penalty_config: Optional GRPO config block for length adjustments. thinking_tags: Optional opening and closing tags used by thinking penalties. mask_env_flagged_samples: Whether to carry env-driven ``mask_sample`` flags in the rollout batch for loss masking. @@ -2477,6 +2480,7 @@ async def run_async_nemo_gym_rollout( log_full_result_tables=log_full_result_tables, effort_config=effort_config, reward_penalty_config=reward_penalty_config, + length_penalty_config=length_penalty_config, thinking_tags=thinking_tags, mask_env_flagged_samples=mask_env_flagged_samples, ) @@ -2514,6 +2518,7 @@ def run_nemo_gym_rollout_sync( greedy: bool = False, effort_config: Optional[EffortLevelsConfig] = None, reward_penalty_config: dict[str, Any] | BaseModel | None = None, + length_penalty_config: dict[str, Any] | BaseModel | None = None, thinking_tags: list[str] | tuple[str, ...] | None = None, sampling_params: Optional[GenerationSamplingParams] = None, mask_env_flagged_samples: bool = True, @@ -2578,6 +2583,7 @@ async def _consume_rollout() -> NemoGymRolloutResult: greedy=greedy, effort_config=effort_config, reward_penalty_config=reward_penalty_config, + length_penalty_config=length_penalty_config, thinking_tags=thinking_tags, mask_env_flagged_samples=mask_env_flagged_samples, returns_entire_batch=True, @@ -2604,6 +2610,7 @@ def _postprocess_single_nemo_gym_group( log_full_result_tables: bool, effort_config: Optional[EffortLevelsConfig] = None, reward_penalty_config: dict[str, Any] | BaseModel | None = None, + length_penalty_config: dict[str, Any] | BaseModel | None = None, thinking_tags: list[str] | tuple[str, ...] | None = None, mask_env_flagged_samples: bool = True, ) -> NemoGymRolloutResult: @@ -2620,6 +2627,27 @@ def _postprocess_single_nemo_gym_group( ) penalty_counts = apply_reward_penalties(results, resolved_reward_penalty_config) + if length_penalty_config is not None: + grpo_config = ( + length_penalty_config.model_dump() + if isinstance(length_penalty_config, BaseModel) + else dict(length_penalty_config) + ) + # Callers pass the whole grpo config block; runs without a + # grpo.length_penalty section are untouched by this block. + if grpo_config.get("length_penalty"): + # Copy the per-row fields the length adjustments consume. + for nemo_gym_row, result in zip(nemo_gym_rows, results): + result["agent_ref"] = nemo_gym_row["agent_ref"] + result["profiled_rewards"] = nemo_gym_row.get("profiled_rewards") + result["profiled_output_lengths"] = nemo_gym_row.get( + "profiled_output_lengths" + ) + result["profile_band"] = nemo_gym_row.get("profile_band") + apply_group_length_penalties( + results, {"grpo": grpo_config}, tokenizer=tokenizer + ) + # Prepare for the rollout metrics calculation below. Not strictly necessary here, but good to have parity with `run_async_multi_turn_rollout` with timer.time(f"{timer_prefix}/prepare_for_metrics_calculation"): batch_size = len(nemo_gym_rows) diff --git a/nemo_rl/utils/length-penalty.md b/nemo_rl/utils/length-penalty.md new file mode 100644 index 00000000000..60ffdc469b3 --- /dev/null +++ b/nemo_rl/utils/length-penalty.md @@ -0,0 +1,597 @@ +# Length Penalty Algorithms + +This file documents the length-penalty and length-bonus algorithms implemented in: + +`nemo_rl/utils/length_penalty.py` + +Configure `grpo.length_penalty`. The length adjustments mutate `full_result["reward"]` in +place during rollout postprocessing. + +All algorithms are resolved per prompt group. Unless otherwise stated, only rollouts with +`reward > 0` participate in length comparisons and receive length-based adjustments. + +## Binary Rewards Requirement + +Length adjustments are defined for **binary (0/1) environment rewards** only. Every algorithm +and the final clamp assume it. At rollout time each prompt group's rewards are checked: a group +containing any graded or negative reward is skipped entirely — no adjustment, no clamp, rewards +pass through untouched — and a warning is logged once per agent. To silence the warning for a +deliberately graded agent (e.g. a genrm judge), disable it explicitly under `agent_overrides` +with `enabled: false`. + +Consequences of binariness worth knowing: + +- The adjusted reward is clamped at 0, so penalties can wipe a correct rollout's reward out but + never flip its sign — and an all-wrong group stays all-zero (no within-group variance, hence + no GRPO gradient, matching vanilla behavior on groups with no correctness signal). +- `top_percentile` is effectively inert: all correct rollouts tie at the top score, so every + positive rollout is always a "top scorer". + +## Common Config + +```yaml +grpo: + length_penalty: + verbose: true + default: + enabled: true + length_type: tokens + top_percentile: 0.5 + reasoning_bonus: 0.0 + answer_bonus: 0.0 + total_bonus: 0.0 + longest_reasoning_penalty: 0.0 + longest_answer_penalty: 0.0 + longest_total_penalty: 0.0 + group_reasoning_length_penalty_coeff: 0.0 + group_answer_length_penalty_coeff: 0.0 + group_total_length_penalty_coeff: 0.0 + reasoning_zmad_threshold: 0.0 + reasoning_zmad_penalty: 0.0 + answer_zmad_threshold: 0.0 + answer_zmad_penalty: 0.0 + total_zmad_threshold: 0.0 + total_zmad_penalty: 0.0 + profiled_length_penalty: 0.0 + profiled_length_n_std: 1.0 + profiled_length_min_samples: 2 + profile_band_total: false + profile_band_reasoning: false + profile_band_answer: false + agent_overrides: + math_with_judge_simple_agent: + enabled: true + group_total_length_penalty_coeff: 0.1 + instruction_following_simple_agent: + enabled: false + genrm_simple_agent: + enabled: false + code_gen_simple_agent: + enabled: true + total_zmad_threshold: 2.0 + total_zmad_penalty: 0.1 +``` + +`length_type` may be: + +- `tokens`: lengths are tokenizer token counts. +- anything else: lengths fall back to character counts. + +`agent_overrides` can override any supported parameter per agent. If an agent is missing from +`agent_overrides`, the implementation falls back to `default`. To disable length adjustments for +a specific environment or agent while keeping the default enabled, set `enabled: false` for that +agent: + +```yaml +grpo: + length_penalty: + default: + enabled: true + group_total_length_penalty_coeff: 0.1 + agent_overrides: + instruction_following_simple_agent: + enabled: false + genrm_simple_agent: + enabled: false +``` + +## Flag Reference + +| Flag | Short description | +| --- | --- | +| `verbose` | Prints per-group length adjustment details during rollout processing. | +| `default` | Default length-adjustment config used for agents without an override. | +| `agent_overrides` | Per-agent config overrides keyed by agent name. | +| `enabled` | Enables length adjustment for this config block. | +| `length_type` | Selects length unit: `tokens` uses tokenizer counts; other values use character counts. | +| `top_percentile` | Fraction of positive scorers treated as top scorers for longest-penalty selection. | +| `reasoning_bonus` | Flat bonus for the shortest positive/top-scoring reasoning trace in a prompt group. | +| `answer_bonus` | Flat bonus for the shortest positive/top-scoring answer in a prompt group. | +| `total_bonus` | Flat bonus for the shortest positive/top-scoring reasoning + answer total length. | +| `longest_reasoning_penalty` | Flat penalty for the longest reasoning trace among top scorers. | +| `longest_answer_penalty` | Flat penalty for the longest answer among top scorers. | +| `longest_total_penalty` | Flat penalty for the longest reasoning + answer total length among top scorers. | +| `group_reasoning_length_penalty_coeff` | Dense group-relative coefficient for reasoning length; shorter positive rollouts get higher adjustment. | +| `group_answer_length_penalty_coeff` | Dense group-relative coefficient for answer length. | +| `group_total_length_penalty_coeff` | Dense group-relative coefficient for total reasoning + answer length. | +| `reasoning_zmad_threshold` | Modified-Z threshold for flagging long reasoning outliers. | +| `reasoning_zmad_penalty` | Flat penalty applied to reasoning lengths above the zMAD threshold. | +| `answer_zmad_threshold` | Modified-Z threshold for flagging long answer outliers. | +| `answer_zmad_penalty` | Flat penalty applied to answer lengths above the zMAD threshold. | +| `total_zmad_threshold` | Modified-Z threshold for flagging long total-length outliers. | +| `total_zmad_penalty` | Flat penalty applied to total lengths above the zMAD threshold. | +| `profiled_length_penalty` | Flat penalty for rollouts longer than a per-prompt profiled-length threshold. | +| `profiled_length_n_std` | Number of standard deviations used in `mean + n_std * std` for profiled-length thresholding. | +| `profiled_length_min_samples` | Minimum PASSING profiled rollouts required; below this no profiled penalty is applied. | +| `pass_rate_length_penalty_weight` | Weight of the pass-rate-scaled (MAI-style) dense length penalty; 0 disables. | +| `profile_band_total` | Enables per-prompt `{a,b,f}` multiplier on total length for correct rollouts. | +| `profile_band_reasoning` | Enables per-prompt `{a,b,f}` multiplier on reasoning length for correct rollouts. | +| `profile_band_answer` | Enables per-prompt `{a,b,f}` multiplier on answer length for correct rollouts. | +| `group_length_penalty_profile_gate` | Gates group-relative length coefficients using a per-prompt `profile_band` threshold. | +| `group_length_penalty_profile_gate_channel` | Selects which profile-band channel to gate on: `reasoning`, `answer`, or `total`. | +| `group_length_penalty_profile_gate_field` | Selects which field from the chosen profile-band channel to use as the gate threshold, usually `a`. | +| `group_length_penalty_profile_gate_positive_only` | If true, computes the gate mean using only `reward > 0` rollouts; if false, uses all rollouts. | + +## Per-Prompt Data Format + +Some algorithms depend on metadata stored on each training-data row. The rollout code copies +these fields from `extra_env_info` into each rollout result before applying length adjustments. + +At minimum, a row still looks like a normal NeMo-Gym training example. The length-related fields +are extra keys: + +```json +{ + "problem": "Solve ...", + "expected_answer": "42", + "agent_name": "math_with_judge_simple_agent", + "extra_env_info": { + "profiled_rewards": [1, 1, 0, 1, 0, 1, 1, 1], + "profiled_output_lengths": [18342, 17110, 32768, 19004, 28991, 16820, 17455, 18101], + "profile_band": { + "total": {"a": 18138.6667, "b": 23756.0123, "f": 0.9}, + "reasoning": {"a": 17686.3333, "b": 23111.0123, "f": 0.9}, + "answer": {"a": 452.3333, "b": 1097.3333, "f": 0.9} + } + } +} +``` + +Some data files store these fields at top level instead of inside `extra_env_info`; the important +part is that by rollout time the result has: + +```json +{ + "profiled_rewards": [1, 1, 0, 1, 0, 1, 1, 1], + "profiled_output_lengths": [18342, 17110, 32768, 19004, 28991, 16820, 17455, 18101], + "profile_band": { + "total": {"a": 18138.6667, "b": 23756.0123, "f": 0.9}, + "reasoning": {"a": 17686.3333, "b": 23111.0123, "f": 0.9}, + "answer": {"a": 452.3333, "b": 1097.3333, "f": 0.9} + } +} +``` + +Field usage: + +- `profiled_rewards`: used by `profiled_length_penalty` to identify passing profiled rollouts. +- `profiled_output_lengths`: used by `profiled_length_penalty` to compute + `mean + n_std * std`. +- `profile_band.total`: used by `profile_band_total` and by profile-gated group-relative + penalties when `group_length_penalty_profile_gate_channel: total`. +- `profile_band.reasoning`: used by `profile_band_reasoning` and by profile-gated + group-relative penalties when the gate channel is `reasoning`. +- `profile_band.answer`: used by `profile_band_answer` and by profile-gated group-relative + penalties when the gate channel is `answer`. + +The profile-band values mean: + +```text +a: full reward / no penalty up to this length +b: multiplier reaches f at this length +f: multiplier at b +``` + +For profile-gated group-relative penalties, `a`, `b`, or `f` can be selected as the gate field, +though `a` is the normal choice: + +```yaml +group_length_penalty_profile_gate: true +group_length_penalty_profile_gate_channel: total +group_length_penalty_profile_gate_field: a +``` + +## Implemented Algorithms + +### 1. Shortest Rollout Bonus + +Config keys: + +- `reasoning_bonus` +- `answer_bonus` +- `total_bonus` + +For each prompt group, the code finds the shortest positive rollout for the selected channel +and adds a flat bonus if that rollout is also a top scorer. + +Channels: + +- `reasoning_bonus`: shortest non-empty reasoning length. +- `answer_bonus`: shortest non-empty answer length. +- `total_bonus`: shortest non-empty reasoning + answer length. + +This is a sparse adjustment: usually only one rollout per group gets the bonus for each enabled +channel. + +Example: + +```yaml +grpo: + length_penalty: + default: + enabled: true + total_bonus: 0.1 +``` + +### 2. Longest Top-Scorer Penalty + +Config keys: + +- `longest_reasoning_penalty` +- `longest_answer_penalty` +- `longest_total_penalty` +- `top_percentile` + +For each prompt group, the code first selects positive rollouts in the top score percentile. +Among those, it subtracts a flat penalty from the longest rollout for the selected channel. + +Channels: + +- `longest_reasoning_penalty` +- `longest_answer_penalty` +- `longest_total_penalty` + +The implementation requires at least two eligible top-scorer rollouts to compare. + +Example: + +```yaml +grpo: + length_penalty: + default: + enabled: true + top_percentile: 0.5 + longest_total_penalty: 0.1 +``` + +### 3. Group Relative-Length Scaling + +Config keys: + +- `group_reasoning_length_penalty_coeff` +- `group_answer_length_penalty_coeff` +- `group_total_length_penalty_coeff` + +This is a dense group-relative adjustment over positive rollouts. + +For each enabled channel: + +1. Find the shortest and longest positive rollout lengths in the group. +2. Convert each length to a raw weight where shorter is larger: + + ```text + raw_weight = 1 - (length - min_length) / (max_length - min_length) + ``` + +3. Zero-center the weights by subtracting the mean raw weight. +4. Multiply by the configured coefficient. + +Shorter positive rollouts receive positive adjustment; longer positive rollouts receive negative +adjustment. If all lengths are equal, the adjustment is zero. + +Example: + +```yaml +grpo: + length_penalty: + default: + enabled: true + group_total_length_penalty_coeff: 0.1 +``` + +### 4. zMAD Long-Outlier Penalty + +Config keys: + +- `reasoning_zmad_threshold` +- `reasoning_zmad_penalty` +- `answer_zmad_threshold` +- `answer_zmad_penalty` +- `total_zmad_threshold` +- `total_zmad_penalty` + +This flags high-side length outliers among positive rollouts using the Iglewicz-Hoaglin modified +Z score: + +```text +modified_z = 0.6745 * (length - median_length) / MAD +``` + +If `modified_z > threshold`, the corresponding flat penalty is subtracted. + +Only long-side outliers are penalized. Short outliers are not penalized. + +The implementation also has a fixed MAD floor: + +```text +MAD / median >= 0.015 +``` + +If the MAD is too small, no zMAD outliers are flagged. + +Example: + +```yaml +grpo: + length_penalty: + default: + enabled: true + total_zmad_threshold: 2.5 + total_zmad_penalty: 0.1 +``` + +### 5. Profiled Length Threshold Penalty + +Config keys: + +- `profiled_length_penalty` +- `profiled_length_n_std` +- `profiled_length_min_samples` + +This uses per-prompt profiling metadata: + +- `profiled_rewards` +- `profiled_output_lengths` + +For each prompt group: + +1. Select profiled lengths from passing rollouts (`profiled_rewards > 0`) only. +2. If there are fewer than `profiled_length_min_samples` passing rollouts, apply **no penalty** + for this prompt. A profile that mostly failed carries no signal about the right thinking + budget — the problem may simply need more than the profiled model had. +3. Otherwise compute: + + ```text + threshold = mean(profiled_lengths) + profiled_length_n_std * std(profiled_lengths) + ``` + +4. Penalize rollouts whose total generated length is greater than or equal to the threshold. + +Example: + +```yaml +grpo: + length_penalty: + default: + enabled: true + profiled_length_penalty: 0.1 + profiled_length_n_std: 1.0 + profiled_length_min_samples: 2 +``` + +### 6. Profile-Band Multiplier + +Config keys: + +- `profile_band_total` +- `profile_band_reasoning` +- `profile_band_answer` + +This uses per-row `profile_band` metadata with channel-specific `{a, b, f}` values: + +```json +{ + "profile_band": { + "total": {"a": 10000, "b": 15000, "f": 0.9}, + "reasoning": {"a": 9000, "b": 14000, "f": 0.9}, + "answer": {"a": 500, "b": 1000, "f": 0.9} + } +} +``` + +For an enabled channel, the multiplier is: + +```text +length <= a: multiplier = 1 +a < length < b: multiplier = 1 - (length - a) / (b - a) * (1 - f) +length >= b: multiplier = f +``` + +So the multiplier interpolates linearly from `1` at `a` down to `f` at `b`, then stays at `f` for +all lengths past `b`. + +Profile-band multipliers are applied only to rollouts whose original environment reward is +positive. + +Example: + +```yaml +grpo: + length_penalty: + default: + enabled: true + profile_band_total: true +``` + +#### Global Defaults (dataset without per-prompt bands) + +When the dataset has no per-prompt `profile_band` metadata, global `{a, b, f}` values can be +set directly in the config under `length_penalty.profile_band`. Only the channels listed under +`defaults` are activated: + +```yaml +grpo: + length_penalty: + profile_band: + enabled: true + defaults: + total: {a: 10000, b: 20000, f: 0.5} +``` + +```yaml +grpo: + length_penalty: + profile_band: + enabled: true + defaults: + reasoning: {a: 9000, b: 14000, f: 0.9} +``` + +```yaml +grpo: + length_penalty: + profile_band: + enabled: true + defaults: + answer: {a: 500, b: 1000, f: 0.9} +``` + +The first config applies the multiplier on total length only, the second on reasoning length +only, and the last on answer length only. Multiple channels may be listed together. + +Semantics: + +- Channels under `defaults` are implicitly enabled — no need to also set + `profile_band_total/reasoning/answer: true` under `length_penalty.default`. Per-agent + `agent_overrides` can still disable a channel (e.g. `profile_band_total: false`). +- Per-prompt `profile_band` metadata, when present on a row, takes precedence over the global + defaults on a per-channel basis (a row that only provides `total` still falls back to the + global `reasoning`/`answer` blocks if those are configured). +- The global band also feeds profile-gated group-relative penalties + (`group_length_penalty_profile_gate`) when rows lack metadata. +- A malformed channel block (missing `a`/`b`/`f`, or `b <= a`) is ignored with a warning. + +### 7. Profile-Gated Group Relative-Length Scaling + +Config keys: + +- `group_length_penalty_profile_gate` +- `group_length_penalty_profile_gate_channel` +- `group_length_penalty_profile_gate_field` +- `group_length_penalty_profile_gate_positive_only` +- plus one or more group-relative coefficients: + - `group_reasoning_length_penalty_coeff` + - `group_answer_length_penalty_coeff` + - `group_total_length_penalty_coeff` + +This is a gate on group-relative length scaling. It does not define a separate penalty by itself. + +Defaults: `group_length_penalty_profile_gate_channel: total`, +`group_length_penalty_profile_gate_field: a`, +`group_length_penalty_profile_gate_positive_only: true`. + +For each prompt group: + +1. Read a threshold from `profile_band[channel][field]`, for example `profile_band["total"]["a"]`. + The band is the row's `profile_band` metadata merged over the global + `length_penalty.profile_band.defaults` (row channels win), so the gate also works on datasets + without per-prompt bands when global defaults are configured. +2. Compute the mean rollout length for the selected channel. +3. If `group_length_penalty_profile_gate_positive_only` is true, use only positive rollouts in + that mean. +4. Enable group-relative length scaling only if: + + ```text + mean_length > profile_band[channel][field] + ``` + +If the gate is closed, all group-relative coefficients are set to zero for that prompt group. + +The gate fails closed: when the threshold cannot be resolved — no `profile_band` available for +the channel (neither per-row nor global), an unknown channel name, or no eligible rollouts to +average (e.g. `positive_only: true` and the whole group scored 0) — group-relative scaling is +disabled for that prompt group. The verbose per-group logs record the reason +(`missing_profile_limit`, `unknown_channel`, `no_lengths`, `mean_le_limit`, `mean_gt_limit`). + +`group_length_penalty_profile_gate_positive_only` only affects the gate decision. It does not +change the rollouts that receive the group-relative adjustment after the gate opens. In the +current implementation, group-relative length scaling itself still applies only to positive +rollouts. + +Example: + +```yaml +grpo: + length_penalty: + default: + enabled: true + group_total_length_penalty_coeff: 0.1 + group_length_penalty_profile_gate: true + group_length_penalty_profile_gate_channel: total + group_length_penalty_profile_gate_field: a + group_length_penalty_profile_gate_positive_only: true +``` + +### 8. Pass-Rate-Scaled Length Penalty (MAI-style) + +Config keys: + +- `pass_rate_length_penalty_weight` + +Implements the MAI-paper length penalty `- w_len * R_len(y_i)` with + +```text +R_len(y_i) = rho_q * |y_i| / l_max +``` + +where `rho_q` is the prompt group's pass rate (fraction of rollouts with `reward > 0`), +`|y_i|` is the rollout's total generated length (reasoning + answer, in `length_type` units), +and `l_max` is the longest total length among the group's CORRECT rollouts (not a hyperparameter — +the penalty is self-normalizing over the set it applies to; long wrong rollouts do not dilute it). + +Properties: + +- Dense pressure scaled by difficulty: easy prompts (high pass rate) get strong shortening + pressure; hard prompts (low pass rate) get little. +- All-wrong groups (`rho_q = 0`) receive exactly zero penalty by construction, preserving the + no-gradient-without-correctness-signal invariant. +- Applied to correct rollouts only (consistent with the rest of this module; under binary + rewards plus the 0-clamp this matches the paper's behavior for wrong rollouts). +- The longest correct rollout in a group loses exactly `w_len * rho_q`; shorter ones lose + proportionally less. + +Example: + +```yaml +grpo: + length_penalty: + default: + enabled: true + pass_rate_length_penalty_weight: 0.2 +``` + +## Practical Notes + +- Most length algorithms act only on positive rollouts (`reward > 0`). +- Additive penalties can wipe a rollout's reward out but never flip its sign: the adjusted + reward is clamped at 0 (stacked flat penalties exceeding the reward produce 0, not a negative + value). Non-binary rewards never reach the clamp — their groups are skipped wholesale (see + Binary Rewards Requirement). +- `profile_band_*` multipliers also apply only to originally correct rollouts. +- Group-relative scaling can reduce average length aggressively because it gives dense per-group + pressure. +- zMAD is more selective: it only hits high-side outliers. + +## Final Recommendations + +1. For domains where longer reasoning does not strongly correlate with higher accuracy, apply + group-relative length scaling. This gives steady pressure toward shorter correct rollouts. + +2. For domains where longer reasoning does correlate with higher accuracy, leave length + unpenalized. Penalizing length in those domains can remove useful reasoning and hurt task + performance. + +3. If you have a target length in mind, for example when this model should be less verbose on a + domain than another reference model, use profile-gated group-relative length scaling. The + profile gate lets the penalty activate only when the prompt group's rollout lengths exceed a + per-prompt target. diff --git a/nemo_rl/utils/length_penalty.py b/nemo_rl/utils/length_penalty.py new file mode 100644 index 00000000000..5d9b64bef24 --- /dev/null +++ b/nemo_rl/utils/length_penalty.py @@ -0,0 +1,1161 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Per-prompt-group length bonuses/penalties for rollout rewards. + +Rewards conciseness among high-quality generations by applying: +1. A flat bonus to the shortest generation among top scorers in each prompt group. +2. Optional flat penalties on the longest reasoning / longest answer among + top-percentile scorers (with at least two eligible rollouts to compare). +3. Independent zero-centered penalties for reasoning and answer length. +4. 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. +""" + +from __future__ import annotations + +import logging +import statistics +from typing import Any + +logger = logging.getLogger(__name__) + +# MAD/median floor for zMAD (fixed; matches ``flag_reasoning_length_outliers`` default). +_ZMAD_MIN_MAD_REL = 0.015 + +_PARAM_KEYS = ( + "enabled", + "reasoning_bonus", + "answer_bonus", + "total_bonus", + "longest_reasoning_penalty", + "longest_answer_penalty", + "longest_total_penalty", + "top_percentile", + "group_reasoning_length_penalty_coeff", + "group_answer_length_penalty_coeff", + "group_total_length_penalty_coeff", + "length_type", + "reasoning_zmad_threshold", + "reasoning_zmad_penalty", + "answer_zmad_threshold", + "answer_zmad_penalty", + "total_zmad_threshold", + "total_zmad_penalty", + "profiled_length_penalty", + "profiled_length_n_std", + "profiled_length_min_samples", + "pass_rate_length_penalty_weight", + "profile_band_total", + "profile_band_reasoning", + "profile_band_answer", + "group_length_penalty_profile_gate", + "group_length_penalty_profile_gate_channel", + "group_length_penalty_profile_gate_field", + "group_length_penalty_profile_gate_positive_only", +) + +# Param keys that should be merged as bools rather than floats. +_BOOL_PARAM_KEYS = frozenset( + { + "enabled", + "profile_band_total", + "profile_band_reasoning", + "profile_band_answer", + "group_length_penalty_profile_gate", + "group_length_penalty_profile_gate_positive_only", + } +) + +_STR_PARAM_KEYS = frozenset( + { + "length_type", + "group_length_penalty_profile_gate_channel", + "group_length_penalty_profile_gate_field", + } +) + +# Length adjustments are defined for binary (0/1) env rewards only. Agents +# already warned about non-binary rewards (warn once per agent, then skip +# their prompt groups). +_NON_BINARY_WARNED_AGENTS: set[str] = set() +_BINARY_REWARD_TOL = 1e-6 + + +def _is_binary_reward(value: float) -> bool: + v = float(value) + return abs(v) <= _BINARY_REWARD_TOL or abs(v - 1.0) <= _BINARY_REWARD_TOL + + +def _extract_reasoning_and_answer_text(result: dict[str, Any]) -> tuple[str, str]: + """Extract reasoning and answer text from the Response API output items.""" + fr = result.get("full_result", {}) + response_obj = fr.get("response", {}) + output_items = ( + response_obj.get("output", []) + if isinstance(response_obj, dict) + else getattr(response_obj, "output", []) + ) + + reasoning_text = "" + answer_text = "" + for item in output_items: + item_type = ( + item.get("type", "") + if isinstance(item, dict) + else getattr(item, "type", "") + ) + if item_type == "reasoning": + summaries = ( + item.get("summary", []) + if isinstance(item, dict) + else getattr(item, "summary", []) + ) + for s in summaries: + t = s.get("text", "") if isinstance(s, dict) else getattr(s, "text", "") + reasoning_text += t + elif item_type == "message": + content = ( + item.get("content", []) + if isinstance(item, dict) + else getattr(item, "content", []) + ) + if isinstance(content, list): + for c in content: + t = ( + c.get("text", "") + if isinstance(c, dict) + else getattr(c, "text", "") + ) + answer_text += t + elif isinstance(content, str): + answer_text += content + + return reasoning_text, answer_text + + +_TOP_LEVEL_LENGTH_PENALTY_KEYS = frozenset( + {"verbose", "default", "agent_overrides", "profile_band"} +) +_PROFILE_BAND_BLOCK_KEYS = frozenset({"enabled", "defaults"}) +_PROFILE_BAND_CHANNELS = frozenset({"total", "reasoning", "answer"}) + + +def _reject_unknown_length_penalty_keys(length_cfg: dict[str, Any]) -> None: + """Raise on unknown config keys instead of silently ignoring them.""" + + def _check(block: Any, allowed: frozenset, where: str) -> None: + if not isinstance(block, dict): + return + unknown = sorted(set(block) - allowed) + if unknown: + raise ValueError( + f"Unknown key(s) {unknown} in {where}; allowed: {sorted(allowed)}" + ) + + _check(length_cfg, _TOP_LEVEL_LENGTH_PENALTY_KEYS, "grpo.length_penalty") + _check( + length_cfg.get("default"), frozenset(_PARAM_KEYS), "grpo.length_penalty.default" + ) + agents_cfg = length_cfg.get("agent_overrides") + if isinstance(agents_cfg, dict): + for agent_name, overrides in agents_cfg.items(): + _check( + overrides, + frozenset(_PARAM_KEYS), + f"grpo.length_penalty.agent_overrides.{agent_name}", + ) + pb_cfg = length_cfg.get("profile_band") + if isinstance(pb_cfg, dict): + _check(pb_cfg, _PROFILE_BAND_BLOCK_KEYS, "grpo.length_penalty.profile_band") + _check( + pb_cfg.get("defaults"), + _PROFILE_BAND_CHANNELS, + "grpo.length_penalty.profile_band.defaults", + ) + + +def apply_group_length_penalties( + results: list[dict[str, Any]], + master_config: dict[str, Any], + tokenizer: Any = None, +) -> None: + """Apply per-prompt-group length bonuses/penalties. + + Reads ``grpo.length_penalty`` for configuration and mutates + ``full_result["reward"]`` in place. No-ops when no length-adjustment + feature is enabled. + + Args: + results: List of per-generation result dicts. + master_config: Full training config dict. + tokenizer: Tokenizer for computing reasoning/answer token counts. + """ + grpo_config = master_config.get("grpo", {}) + length_cfg = dict(grpo_config.get("length_penalty", {}) or {}) + if not length_cfg: + return + + _reject_unknown_length_penalty_keys(length_cfg) + default_cfg = length_cfg.get("default", {}) + agents_cfg = length_cfg.get("agent_overrides") + global_band = _resolve_global_profile_band(length_cfg.get("profile_band")) + verbose = bool(length_cfg.get("verbose", False)) + # `enabled` defaults True here AND in the per-group param resolution: a + # configured `default:` block is intent-to-enable; omitting `enabled` must + # not silently no-op (and must not depend on unrelated keys being present). + if not default_cfg.get("enabled", True) and not agents_cfg and not global_band: + return + + num_gens = master_config["grpo"]["num_generations_per_prompt"] + defaults: dict[str, Any] = {} + for k in _PARAM_KEYS: + if k == "length_type": + defaults[k] = default_cfg.get(k, "tokens") + elif k == "group_length_penalty_profile_gate_channel": + defaults[k] = default_cfg.get(k, "total") + elif k == "group_length_penalty_profile_gate_field": + defaults[k] = default_cfg.get(k, "a") + elif k == "enabled": + defaults[k] = default_cfg.get(k, True) + elif k == "group_length_penalty_profile_gate_positive_only": + defaults[k] = default_cfg.get(k, True) + elif k in _BOOL_PARAM_KEYS: + defaults[k] = default_cfg.get(k, False) + elif k == "profiled_length_min_samples": + defaults[k] = default_cfg.get(k, 2) + elif k == "profiled_length_n_std": + defaults[k] = default_cfg.get(k, 1.0) + elif k == "top_percentile": + defaults[k] = default_cfg.get(k, 0.5) + else: + defaults[k] = default_cfg.get(k, 0.0) + # Channels listed under length_penalty.profile_band.defaults are implicitly + # enabled — unless the user explicitly configured the channel flag, which + # always wins (e.g. profile_band_total: false stays false). + for _ch in global_band: + if f"profile_band_{_ch}" not in default_cfg: + defaults[f"profile_band_{_ch}"] = True + + n = len(results) + original_rewards = [r["full_result"]["reward"] for r in results] + agent_names = [r["agent_ref"]["name"] for r in results] + + # Extract text once; lengths computed per-group based on resolved length_type + texts: list[tuple[str, str]] = [] + for r in results: + texts.append(_extract_reasoning_and_answer_text(r)) + + # Phase 1: calculate all adjustments per-group + all_adjustments = [0.0] * n + all_reasoning_adj = [0.0] * n + all_answer_adj = [0.0] * n + all_total_adj = [0.0] * n + all_reasoning_bonus = [0.0] * n + all_answer_bonus = [0.0] * n + all_total_bonus = [0.0] * n + all_reasoning_longest_pen = [0.0] * n + all_answer_longest_pen = [0.0] * n + all_total_longest_pen = [0.0] * n + all_zmad_reasoning_adj = [0.0] * n + all_zmad_answer_adj = [0.0] * n + all_zmad_total_adj = [0.0] * n + all_profiled_length_adj = [0.0] * n + all_pass_rate_len_adj = [0.0] * n + reasoning_lengths = [0] * n + answer_lengths = [0] * n + total_lengths = [0] * n + groups_adjusted = 0 + group_gate_infos: dict[int, dict[str, Any]] = {} + # Rows whose group passed the binary-rewards check; all other rows are + # left completely untouched (no adjustment, no clamp, no reward writeback). + binary_ok = [False] * n + + for g in range(0, n, num_gens): + agent_name = agent_names[g] + group_size = min(num_gens, n - g) + # Length adjustments are defined for binary (0/1) env rewards only: + # every algorithm and the phase-3 clamp assume it. Skip (and warn once + # per agent) on graded or negative rewards. + if any( + not _is_binary_reward(original_rewards[g + k]) for k in range(group_size) + ): + if agent_name not in _NON_BINARY_WARNED_AGENTS: + _NON_BINARY_WARNED_AGENTS.add(agent_name) + logger.warning( + f"length penalties require binary (0/1) env rewards; agent " + f"{agent_name} produced non-binary rewards — skipping length " + f"penalties for its prompt groups" + ) + continue + for k in range(group_size): + binary_ok[g + k] = True + if any(results[g + k].get("low_effort_applied") for k in range(group_size)): + continue + params = _resolve_agent_params(agent_name, agents_cfg, defaults) + if params is None: + continue + if not params.pop("enabled", True): + continue + + group_lt = params.pop("length_type", "tokens") + use_tokens = group_lt == "tokens" + + for k in range(group_size): + idx = g + k + r_text, a_text = texts[idx] + if use_tokens and tokenizer is not None: + reasoning_lengths[idx] = ( + len(tokenizer.encode(r_text, add_special_tokens=False)) + if r_text + else 0 + ) + answer_lengths[idx] = ( + len(tokenizer.encode(a_text, add_special_tokens=False)) + if a_text + else 0 + ) + else: + reasoning_lengths[idx] = len(r_text) + answer_lengths[idx] = len(a_text) + + group_reasoning = reasoning_lengths[g : g + num_gens] + group_answer = answer_lengths[g : g + num_gens] + group_total = [r + a for r, a in zip(group_reasoning, group_answer)] + total_lengths[g : g + group_size] = group_total[:group_size] + group_rewards = original_rewards[g : g + num_gens] + gate_info = _group_length_profile_gate_info( + band=_merged_profile_band(results[g].get("profile_band"), global_band), + params=params, + rewards=group_rewards[:group_size], + reasoning_lengths=group_reasoning[:group_size], + answer_lengths=group_answer[:group_size], + total_lengths=group_total[:group_size], + ) + group_gate_infos[g] = gate_info + if gate_info["enabled"] and not gate_info["open"]: + params["group_reasoning_length_penalty_coeff"] = 0.0 + params["group_answer_length_penalty_coeff"] = 0.0 + params["group_total_length_penalty_coeff"] = 0.0 + ( + _, + adjustments, + reasoning_adjs, + answer_adjs, + total_adjs, + r_bonus, + a_bonus, + t_bonus, + r_lpen, + a_lpen, + t_lpen, + zmad_r_adj, + zmad_a_adj, + zmad_t_adj, + ) = _apply_length_penaltyes_and_penalties( + group_rewards, group_reasoning, group_answer, group_total, **params + ) + + for k in range(len(adjustments)): + all_adjustments[g + k] = adjustments[k] + all_reasoning_adj[g + k] = reasoning_adjs[k] + all_answer_adj[g + k] = answer_adjs[k] + all_total_adj[g + k] = total_adjs[k] + all_reasoning_bonus[g + k] = r_bonus[k] + all_answer_bonus[g + k] = a_bonus[k] + all_total_bonus[g + k] = t_bonus[k] + all_reasoning_longest_pen[g + k] = r_lpen[k] + all_answer_longest_pen[g + k] = a_lpen[k] + all_total_longest_pen[g + k] = t_lpen[k] + all_zmad_reasoning_adj[g + k] = zmad_r_adj[k] + all_zmad_answer_adj[g + k] = zmad_a_adj[k] + all_zmad_total_adj[g + k] = zmad_t_adj[k] + + if any(a != 0.0 for a in adjustments): + groups_adjusted += 1 + + # Profiled length penalty: penalize rollouts longer than mean + n_std of + # passing profiled lengths for this prompt. If fewer than min_samples + # profiled rollouts passed, the profiled model found the problem hard + # and its lengths carry no budget signal — apply no penalty at all. + plp = params.get("profiled_length_penalty", 0.0) + if plp > 0.0: + p_rewards = results[g].get("profiled_rewards") + p_lengths = results[g].get("profiled_output_lengths") + if p_rewards is not None and p_lengths is not None: + min_samples = int(params.get("profiled_length_min_samples", 2)) + passing = [l for r, l in zip(p_rewards, p_lengths) if r > 0] + if len(passing) >= min_samples: + mean_l = statistics.mean(passing) + std_l = statistics.stdev(passing) if len(passing) >= 2 else 0.0 + n_std = float(params.get("profiled_length_n_std", 1.0)) + threshold = mean_l + n_std * std_l + for k in range(group_size): + idx = g + k + if total_lengths[idx] >= threshold: + all_profiled_length_adj[idx] = -plp + + # Pass-rate-scaled length penalty (MAI): -w * rho_q * |y_i| / l_max on + # correct rollouts, where rho_q is the group's pass rate. Easy prompts + # (high pass rate) get strong shortening pressure; hard prompts get + # little; all-wrong groups (rho_q = 0) get exactly none. + prlp_w = params.get("pass_rate_length_penalty_weight", 0.0) + if prlp_w > 0.0: + pass_rate = sum( + 1 for k in range(group_size) if original_rewards[g + k] > 0 + ) / float(group_size) + if pass_rate > 0.0: + # l_max is the longest CORRECT rollout, so the penalty is + # self-normalizing over the set it applies to: the longest + # correct rollout loses exactly w * rho_q, shorter ones + # proportionally less — decoupled from wrong-rollout lengths + # (long wrong rambles must not dilute the pressure). + l_max = float( + max( + total_lengths[g + k] + for k in range(group_size) + if original_rewards[g + k] > 0 + ) + ) + if l_max > 0.0: + for k in range(group_size): + idx = g + k + if original_rewards[idx] <= 0: + continue + all_pass_rate_len_adj[idx] = ( + -prlp_w * pass_rate * total_lengths[idx] / l_max + ) + + # Phase 2: debug print (only when verbose flag is set) + if verbose: + num_groups = n // num_gens if num_gens > 0 else 0 + print(f"\n{'=' * 70}", flush=True) + print( + f"[Rollout] {n} samples, {num_groups} groups, {groups_adjusted} adjusted" + f" default longest_reasoning_penalty={defaults['longest_reasoning_penalty']}" + f" longest_answer_penalty={defaults['longest_answer_penalty']}", + flush=True, + ) + + for g in range(0, n, num_gens): + agent_name = agent_names[g] + group_size = min(num_gens, n - g) + low_effort = any( + results[g + k].get("low_effort_applied") for k in range(group_size) + ) + params = _resolve_agent_params(agent_name, agents_cfg, defaults) + skipped = params is None + disabled = params is not None and not params.get("enabled", True) + + if low_effort: + print( + f"\n group {g // num_gens} agent={agent_name} [low_effort — skipped]", + flush=True, + ) + elif skipped: + print( + f"\n group {g // num_gens} agent={agent_name} [skipped]" + f" (default longest_reasoning_penalty={defaults['longest_reasoning_penalty']}" + f" longest_answer_penalty={defaults['longest_answer_penalty']})", + flush=True, + ) + elif disabled: + print( + f"\n group {g // num_gens} agent={agent_name} [disabled]" + f" longest_reasoning_penalty={params['longest_reasoning_penalty']}" + f" longest_answer_penalty={params['longest_answer_penalty']}", + flush=True, + ) + else: + lt = params.get("length_type", "tokens") + unit = "tok" if lt == "tokens" else "chr" + print( + f"\n group {g // num_gens} agent={agent_name}" + f" length_type={unit}" + f" reasoning_bonus={params['reasoning_bonus']} answer_bonus={params['answer_bonus']}" + f" total_bonus={params['total_bonus']}" + f" longest_reasoning_penalty={params['longest_reasoning_penalty']}" + f" longest_answer_penalty={params['longest_answer_penalty']}" + f" longest_total_penalty={params['longest_total_penalty']}" + f" top_pct={params['top_percentile']}" + f" reasoning_coeff={params['group_reasoning_length_penalty_coeff']}" + f" answer_coeff={params['group_answer_length_penalty_coeff']}" + f" total_coeff={params['group_total_length_penalty_coeff']}" + f" reasoning_zmad_threshold={params['reasoning_zmad_threshold']}" + f" reasoning_zmad_penalty={params['reasoning_zmad_penalty']}" + f" answer_zmad_threshold={params['answer_zmad_threshold']}" + f" answer_zmad_penalty={params['answer_zmad_penalty']}" + f" total_zmad_threshold={params['total_zmad_threshold']}" + f" total_zmad_penalty={params['total_zmad_penalty']}" + f" profiled_length_penalty={params['profiled_length_penalty']}" + f" profiled_length_n_std={params['profiled_length_n_std']}" + f" profiled_length_min_samples={params['profiled_length_min_samples']}", + flush=True, + ) + gate = group_gate_infos.get(g) + if gate and gate["enabled"]: + print( + f" profile_gate channel={gate['channel']} field={gate['field']}" + f" positive_only={gate['positive_only']}" + f" mean={gate['mean']}" + f" limit={gate['limit']}" + f" open={gate['open']}" + f" reason={gate['reason']}", + flush=True, + ) + for k in range(group_size): + idx = g + k + orig = original_rewards[idx] + profiled_adj = ( + all_profiled_length_adj[idx] if all_adjustments[idx] >= 0 else 0.0 + ) + final = max( + 0.0, + orig + + all_adjustments[idx] + + profiled_adj + + all_pass_rate_len_adj[idx], + ) + print( + f" [{k}] reward={orig:.4f}" + f" reasoning_len={reasoning_lengths[idx]}" + f" reasoning_adj={all_reasoning_adj[idx]:+.4f}" + f" reasoning_bonus={all_reasoning_bonus[idx]:+.4f}" + f" longest_reasoning_penalty_adj={all_reasoning_longest_pen[idx]:+.4f}" + f" answer_len={answer_lengths[idx]}" + f" answer_adj={all_answer_adj[idx]:+.4f}" + f" answer_bonus={all_answer_bonus[idx]:+.4f}" + f" longest_answer_penalty_adj={all_answer_longest_pen[idx]:+.4f}" + f" total_len={total_lengths[idx]}" + f" total_adj={all_total_adj[idx]:+.4f}" + f" total_bonus={all_total_bonus[idx]:+.4f}" + f" longest_total_penalty_adj={all_total_longest_pen[idx]:+.4f}" + f" zmad_r={all_zmad_reasoning_adj[idx]:+.4f}" + f" zmad_a={all_zmad_answer_adj[idx]:+.4f}" + f" zmad_t={all_zmad_total_adj[idx]:+.4f}" + f" profiled_len_adj={all_profiled_length_adj[idx]:+.4f}" + f" pass_rate_len_adj={all_pass_rate_len_adj[idx]:+.4f}" + f" final_reward={final:.4f}", + flush=True, + ) + + print(f"{'=' * 70}\n", flush=True) + + # Phase 3: apply additive adjustments (binary-verified rows only) + additive_base_rewards = [0.0] * n + for i, r in enumerate(results): + if not binary_ok[i]: + continue + # The profiled-length penalty stacks only on rollouts whose group + # adjustments are non-negative: a rollout already penalized by the + # group-relative channels should not be double-penalized for the same + # excess length. + profiled_adj = all_profiled_length_adj[i] if all_adjustments[i] >= 0 else 0.0 + additive_delta = all_adjustments[i] + profiled_adj + all_pass_rate_len_adj[i] + additive_base_rewards[i] = original_rewards[i] + additive_delta + # Rewards here are binary, so any negative value is penalty-created. + # Length penalties may wipe a reward out but never flip its sign; this + # also keeps all-wrong groups variance-free (no gradient from a group + # with no correctness signal). + if additive_base_rewards[i] < 0: + additive_base_rewards[i] = 0.0 + r["full_result"]["reward"] = additive_base_rewards[i] + + # Phase 4: apply per-prompt profile_band multipliers (correct rollouts only). + _apply_profile_band_multipliers( + results=results, + original_rewards=original_rewards, + base_rewards=additive_base_rewards, + total_lengths=total_lengths, + reasoning_lengths=reasoning_lengths, + answer_lengths=answer_lengths, + agent_names=agent_names, + agents_cfg=agents_cfg, + defaults=defaults, + num_gens=num_gens, + global_band=global_band, + binary_ok=binary_ok, + ) + + +def _resolve_global_profile_band(pb_cfg: Any) -> dict[str, dict[str, Any]]: + """Parse ``length_penalty.profile_band`` into per-channel {a, b, f} blocks. + + Returns only channels ("total", "reasoning", "answer") present under + ``defaults`` with a complete, well-formed block. Empty dict when the + section is absent or disabled. + """ + if not isinstance(pb_cfg, dict) or not pb_cfg.get("enabled", False): + return {} + pb_defaults = pb_cfg.get("defaults") + if not isinstance(pb_defaults, dict): + return {} + band: dict[str, dict[str, Any]] = {} + for ch in ("total", "reasoning", "answer"): + ch_cfg = pb_defaults.get(ch) + if not isinstance(ch_cfg, dict): + continue + a, b, f = ch_cfg.get("a"), ch_cfg.get("b"), ch_cfg.get("f") + if a is None or b is None or f is None or b <= a: + logger.warning( + f"length_penalty.profile_band.defaults.{ch} is malformed " + f"(a={a}, b={b}, f={f}); ignoring this channel" + ) + continue + band[ch] = {"a": a, "b": b, "f": f} + return band + + +def _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).""" + if not global_band: + return row_band + if not row_band: + return dict(global_band) + merged: dict[str, Any] = dict(global_band) + merged.update(row_band) + return merged + + +def _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, Any] | 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``). + """ + n = len(results) + global_band = global_band or {} + for g in range(0, n, num_gens): + agent_name = agent_names[g] + group_size = min(num_gens, n - g) + if any(results[g + k].get("low_effort_applied") for k in range(group_size)): + continue + params = _resolve_agent_params(agent_name, agents_cfg, defaults) + if params is None: + continue + use_total = bool(params.get("profile_band_total", False)) + use_rsn = bool(params.get("profile_band_reasoning", False)) + use_ans = bool(params.get("profile_band_answer", False)) + if not (use_total or use_rsn or use_ans): + continue + band = _merged_profile_band(results[g].get("profile_band"), global_band) + if not band: + continue + ch_total = band.get("total") if use_total else None + ch_rsn = band.get("reasoning") if use_rsn else None + ch_ans = band.get("answer") if use_ans else None + for k in range(group_size): + idx = g + k + # Only rows whose group passed the binary-rewards check. + if binary_ok is not None and not binary_ok[idx]: + continue + # Gate on the env reward (correct rollouts only). + if original_rewards[idx] <= 0: + continue + # The phase-3 clamp floors the base at 0; multiplying a negative + # base by m < 1 would RAISE the reward for longer rollouts, so + # scale only strictly-positive bases. + if base_rewards[idx] <= 0: + continue + current_reward = base_rewards[idx] + + total_m = _band_multiplier(total_lengths[idx], ch_total) + total_delta = current_reward * total_m - current_reward + current_reward += total_delta + + reasoning_m = _band_multiplier(reasoning_lengths[idx], ch_rsn) + reasoning_delta = current_reward * reasoning_m - current_reward + current_reward += reasoning_delta + + answer_m = _band_multiplier(answer_lengths[idx], ch_ans) + answer_delta = current_reward * answer_m - current_reward + current_reward += answer_delta + + results[idx]["full_result"]["reward"] = current_reward + + +def _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 + """ + if not ch: + return 1.0 + a = ch.get("a") + b = ch.get("b") + f = ch.get("f") + if a is None or b is None or f is None or b <= a: + return 1.0 + if rl <= a: + return 1.0 + if rl >= b: + return float(f) + return 1.0 - (rl - a) / (b - a) * (1.0 - float(f)) + + +def _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``. + """ + enabled = bool(params.get("group_length_penalty_profile_gate", False)) + channel = str(params.get("group_length_penalty_profile_gate_channel", "total")) + field = str(params.get("group_length_penalty_profile_gate_field", "a")) + positive_only = bool( + params.get("group_length_penalty_profile_gate_positive_only", True) + ) + info = { + "enabled": enabled, + "open": True, + "channel": channel, + "field": field, + "positive_only": positive_only, + "mean": None, + "limit": None, + "reason": "disabled", + } + if not enabled: + return info + + limit = _profile_band_numeric_value(band, channel, field) + info["limit"] = limit + if limit is None: + info["open"] = False + info["reason"] = "missing_profile_limit" + return info + + length_by_channel = { + "reasoning": reasoning_lengths, + "answer": answer_lengths, + "total": total_lengths, + } + candidate_lengths = length_by_channel.get(channel) + if candidate_lengths is None: + info["open"] = False + info["reason"] = "unknown_channel" + return info + + if positive_only: + lengths = [l for r, l in zip(rewards, candidate_lengths) if r > 0] + else: + lengths = list(candidate_lengths) + if not lengths: + info["open"] = False + info["reason"] = "no_lengths" + return info + + mean_length = float(statistics.mean(lengths)) + info["mean"] = mean_length + info["open"] = mean_length > limit + info["reason"] = "mean_gt_limit" if info["open"] else "mean_le_limit" + return info + + +def _profile_band_numeric_value( + band: dict[str, Any] | None, channel: str, field: str +) -> float | None: + if not isinstance(band, dict): + return None + channel_block = band.get(channel) + if not isinstance(channel_block, dict): + return None + value = channel_block.get(field) + if isinstance(value, bool): + return None + if isinstance(value, (int, float)): + return float(value) + return None + + +def _resolve_agent_params( + agent_name: str, + agents_cfg: dict[str, Any] | None, + defaults: dict[str, Any], +) -> dict[str, Any] | None: + """Resolve length bonus parameters for a given agent.""" + if agents_cfg is None: + return dict(defaults) + + if agent_name not in agents_cfg: + print( + f"[length_penalty] WARNING: agent '{agent_name}' not found in " + f"agent_overrides, falling back to defaults", + flush=True, + ) + return dict(defaults) + + overrides = agents_cfg[agent_name] + if overrides is None: + return dict(defaults) + + merged = dict(defaults) + for key in _PARAM_KEYS: + if key in overrides: + if key in _STR_PARAM_KEYS: + merged[key] = overrides[key] + elif key in _BOOL_PARAM_KEYS: + merged[key] = bool(overrides[key]) + else: + merged[key] = float(overrides[key]) + return merged + + +def _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``.""" + if len(lengths) < 2: + return set() + med = statistics.median(lengths) + devs = [abs(x - med) for x in lengths] + mad = statistics.median(devs) + if mad == 0: + return set() + if min_mad_rel > 0 and mad / max(med, 1e-9) < min_mad_rel: + return set() + out: set[int] = set() + for k, x in enumerate(lengths): + mz = 0.6745 * (x - med) / mad + if mz > z_thresh: + out.add(k) + return out + + +def _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. + + 1. Reasoning bonus: shortest non-empty reasoning among positive scorers; awarded only if that + sample satisfies ``reward >= top_threshold``. + 2. Answer bonus: same pattern for shortest non-empty answer. + 3. Total bonus: same pattern for shortest combined (reasoning + answer) length. + 4. Longest penalties: subtract from longest non-empty reasoning / answer / total among + top-percentile scorers; needs at least two eligible rollouts to compare. + 5. Independent zero-centered penalties for reasoning, answer, and total lengths. + 6. 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. + """ + n = len(rewards) + zeros = [0.0] * n + if n < 2: + return ( + list(rewards), + list(zeros), + list(zeros), + list(zeros), + list(zeros), + list(zeros), + list(zeros), + list(zeros), + list(zeros), + list(zeros), + list(zeros), + list(zeros), + list(zeros), + list(zeros), + ) + + positive_indices = [i for i in range(n) if rewards[i] > 0] + if len(positive_indices) < 2: + return ( + list(rewards), + list(zeros), + list(zeros), + list(zeros), + list(zeros), + list(zeros), + list(zeros), + list(zeros), + list(zeros), + list(zeros), + list(zeros), + list(zeros), + list(zeros), + list(zeros), + ) + + adjusted = list(rewards) + adjustments = [0.0] * n + reasoning_adjs = [0.0] * n + answer_adjs = [0.0] * n + total_adjs = [0.0] * n + r_bonus_per = [0.0] * n + a_bonus_per = [0.0] * n + t_bonus_per = [0.0] * n + r_longest_pen_per = [0.0] * n + a_longest_pen_per = [0.0] * n + t_longest_pen_per = [0.0] * n + zmad_reasoning_adj = [0.0] * n + zmad_answer_adj = [0.0] * n + zmad_total_adj = [0.0] * n + + pos_reasoning = [reasoning_lengths[i] for i in positive_indices] + pos_answer = [answer_lengths[i] for i in positive_indices] + pos_total = [total_lengths[i] for i in positive_indices] + pos_rewards = [rewards[i] for i in positive_indices] + + sorted_scores = sorted(pos_rewards, reverse=True) + threshold_idx = max(0, int(len(pos_rewards) * top_percentile) - 1) + top_threshold = sorted_scores[threshold_idx] + top_scorer_indices = [i for i in positive_indices if rewards[i] >= top_threshold] + + # Reasoning bonus: shortest non-empty reasoning among top scorers + if reasoning_bonus > 0: + valid = [ + (pi, pos_reasoning[k]) + for k, pi in enumerate(positive_indices) + if pos_reasoning[k] > 0 + ] + if valid: + shortest_pi, _ = min(valid, key=lambda x: x[1]) + if adjusted[shortest_pi] >= top_threshold: + adjusted[shortest_pi] += reasoning_bonus + adjustments[shortest_pi] += reasoning_bonus + r_bonus_per[shortest_pi] = reasoning_bonus + + # Answer bonus: shortest non-empty answer among top scorers + if answer_bonus > 0: + valid = [ + (pi, pos_answer[k]) + for k, pi in enumerate(positive_indices) + if pos_answer[k] > 0 + ] + if valid: + shortest_pi, _ = min(valid, key=lambda x: x[1]) + if adjusted[shortest_pi] >= top_threshold: + adjusted[shortest_pi] += answer_bonus + adjustments[shortest_pi] += answer_bonus + a_bonus_per[shortest_pi] = answer_bonus + + # Total bonus: shortest combined (reasoning + answer) among top scorers + if total_bonus > 0: + valid = [ + (pi, pos_total[k]) + for k, pi in enumerate(positive_indices) + if pos_total[k] > 0 + ] + if valid: + shortest_pi, _ = min(valid, key=lambda x: x[1]) + if adjusted[shortest_pi] >= top_threshold: + adjusted[shortest_pi] += total_bonus + adjustments[shortest_pi] += total_bonus + t_bonus_per[shortest_pi] = total_bonus + + # Longest reasoning penalty: longest among top-percentile scorers only + if longest_reasoning_penalty > 0: + valid = [ + (pi, reasoning_lengths[pi]) + for pi in top_scorer_indices + if reasoning_lengths[pi] > 0 + ] + if len(valid) >= 2: + longest_pi, _ = max(valid, key=lambda x: x[1]) + pen = -longest_reasoning_penalty + adjusted[longest_pi] += pen + adjustments[longest_pi] += pen + r_longest_pen_per[longest_pi] = pen + + # Longest answer penalty: longest among top-percentile scorers only + if longest_answer_penalty > 0: + valid = [ + (pi, answer_lengths[pi]) + for pi in top_scorer_indices + if answer_lengths[pi] > 0 + ] + if len(valid) >= 2: + longest_pi, _ = max(valid, key=lambda x: x[1]) + pen = -longest_answer_penalty + adjusted[longest_pi] += pen + adjustments[longest_pi] += pen + a_longest_pen_per[longest_pi] = pen + + # Longest total penalty: longest combined length among top-percentile scorers only + if longest_total_penalty > 0: + valid = [ + (pi, total_lengths[pi]) + for pi in top_scorer_indices + if total_lengths[pi] > 0 + ] + if len(valid) >= 2: + longest_pi, _ = max(valid, key=lambda x: x[1]) + pen = -longest_total_penalty + adjusted[longest_pi] += pen + adjustments[longest_pi] += pen + t_longest_pen_per[longest_pi] = pen + + # Independent reasoning, answer, and total length penalties (zero-centered) + if ( + group_reasoning_length_penalty_coeff > 0 + or group_answer_length_penalty_coeff > 0 + or group_total_length_penalty_coeff > 0 + ): + reasoning_weights = _compute_length_weights(pos_reasoning) + answer_weights = _compute_length_weights(pos_answer) + total_weights = _compute_length_weights(pos_total) + + for k, i in enumerate(positive_indices): + r_adj = reasoning_weights[k] * group_reasoning_length_penalty_coeff + a_adj = answer_weights[k] * group_answer_length_penalty_coeff + t_adj = total_weights[k] * group_total_length_penalty_coeff + combined_adj = r_adj + a_adj + t_adj + reasoning_adjs[i] = r_adj + answer_adjs[i] = a_adj + total_adjs[i] = t_adj + if combined_adj != 0: + adjusted[i] += combined_adj + adjustments[i] += combined_adj + + zm = _ZMAD_MIN_MAD_REL + ztr = float(reasoning_zmad_threshold) + zpr = float(reasoning_zmad_penalty) + zta = float(answer_zmad_threshold) + zpa = float(answer_zmad_penalty) + ztt = float(total_zmad_threshold) + zpt = float(total_zmad_penalty) + + if len(positive_indices) >= 2: + if ztr > 0.0: + if zpr != 0.0: + for local_k in _zmad_local_outliers(pos_reasoning, ztr, zm): + gi = positive_indices[local_k] + adjusted[gi] -= zpr + adjustments[gi] -= zpr + zmad_reasoning_adj[gi] -= zpr + if zta > 0.0: + if zpa != 0.0: + for local_k in _zmad_local_outliers(pos_answer, zta, zm): + gi = positive_indices[local_k] + adjusted[gi] -= zpa + adjustments[gi] -= zpa + zmad_answer_adj[gi] -= zpa + if ztt > 0.0: + if zpt != 0.0: + for local_k in _zmad_local_outliers(pos_total, ztt, zm): + gi = positive_indices[local_k] + adjusted[gi] -= zpt + adjustments[gi] -= zpt + zmad_total_adj[gi] -= zpt + + return ( + adjusted, + adjustments, + reasoning_adjs, + answer_adjs, + total_adjs, + r_bonus_per, + a_bonus_per, + t_bonus_per, + r_longest_pen_per, + a_longest_pen_per, + t_longest_pen_per, + zmad_reasoning_adj, + zmad_answer_adj, + zmad_total_adj, + ) + + +def _compute_length_weights(lengths: list[int]) -> list[float]: + """Compute zero-centered weights where shorter = higher weight. + + Returns all zeros if all lengths are equal. + """ + max_len = max(lengths) + min_len = min(lengths) + + if max_len == min_len: + return [0.0] * len(lengths) + + span = max_len - min_len + raw_weights = [1.0 - ((length - min_len) / span) for length in lengths] + mean_weight = sum(raw_weights) / len(raw_weights) + return [w - mean_weight for w in raw_weights] diff --git a/scripts/build_profile_band_dataset.py b/scripts/build_profile_band_dataset.py new file mode 100644 index 00000000000..afd4e94bdf1 --- /dev/null +++ b/scripts/build_profile_band_dataset.py @@ -0,0 +1,225 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Attach a per-prompt ``profile_band`` block to a profiled training JSONL. + +Reads a JSONL produced by ``profile_run`` (rows must carry +``profiled_rewards``, ``profiled_output_lengths``, +``profiled_reasoning_lengths``, ``profiled_answer_lengths``, and +``pass_rate``) and writes a new JSONL where each row has an additional +``profile_band`` field consumed by Mechanism 6 in +``nemo_rl/utils/length_penalty.py``:: + + profile_band: + total: {a, b, f} # only present if data is non-degenerate + reasoning: {a, b, f} + answer: {a, b, f} + +For each channel: + - reference set = passing profiled rollouts (reward > 0). If too few + (< ``min_passing``), fall back to all profiled rollouts. + - a = mean(reference) + - b = mean(reference) + n_std * std(reference) + - f = looked up in ``f_table`` by row's ``pass_rate``. + - channel block is OMITTED if std == 0 (degenerate / cap-clamped), + if reference set has fewer than 2 samples, or if pass_rate is not + in the f_table. + +Config yaml shape:: + + n_std: 2.0 + min_passing: 2 + channels: [total, reasoning, answer] + f_table: # pass_rate -> f; omitted pass_rates skip the row entirely + - {pass_rate: 1.000, f: 0.6} + - {pass_rate: 0.875, f: 0.7} + - {pass_rate: 0.750, f: 0.8} + - {pass_rate: 0.625, f: 0.9} + +Usage:: + + python scripts/build_profile_band_dataset.py \ + --input /path/to/dapo17k_profiled_boxed_nanov3.jsonl \ + --config /path/to/profile_band.yaml \ + --output /path/to/dapo17k_profiled_band_boxed_nanov3.jsonl +""" + +from __future__ import annotations + +import argparse +import json +import statistics +import sys +from collections import Counter +from pathlib import Path +from typing import Any + +# yaml is in stdlib via pyyaml on the cluster; fall back to a tiny parser if absent. +try: + import yaml # type: ignore +except ImportError: + yaml = None # type: ignore + +CHANNEL_TO_LENGTHS_KEY = { + "total": "profiled_output_lengths", + "reasoning": "profiled_reasoning_lengths", + "answer": "profiled_answer_lengths", +} + + +def parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + p.add_argument("--input", required=True, help="Path to profiled JSONL (input).") + p.add_argument("--config", required=True, help="Path to profile_band yaml config.") + p.add_argument("--output", required=True, help="Path to write augmented JSONL.") + p.add_argument( + "--quiet", action="store_true", help="Suppress per-row diagnostics summary." + ) + return p.parse_args() + + +def load_config(path: str) -> dict[str, Any]: + if yaml is None: + sys.exit("PyYAML not available; install pyyaml to use this script.") + with open(path) as f: + cfg = yaml.safe_load(f) + if not isinstance(cfg, dict): + sys.exit(f"config at {path} did not parse as a dict") + cfg.setdefault("n_std", 2.0) + cfg.setdefault("min_passing", 2) + cfg.setdefault("channels", ["total", "reasoning", "answer"]) + cfg.setdefault("f_table", []) + # Normalize f_table to a {rounded_pass_rate: f} dict. + f_table: dict[float, float] = {} + for entry in cfg["f_table"]: + pr = round(float(entry["pass_rate"]), 4) + f_table[pr] = float(entry["f"]) + cfg["_f_table"] = f_table + bad_channels = [c for c in cfg["channels"] if c not in CHANNEL_TO_LENGTHS_KEY] + if bad_channels: + sys.exit( + f"unknown channels in config: {bad_channels}; valid: {sorted(CHANNEL_TO_LENGTHS_KEY)}" + ) + return cfg + + +def lookup_f(pass_rate: float, f_table: dict[float, float]) -> float | None: + return f_table.get(round(float(pass_rate), 4)) + + +def channel_block( + lengths: list[int], + rewards: list[float], + n_std: float, + min_passing: int, + f_value: float, +) -> dict[str, float] | None: + """Compute {a, b, f} for one channel; return None if degenerate.""" + passing = [l for l, r in zip(lengths, rewards) if r is not None and r > 0] + ref = passing if len(passing) >= min_passing else list(lengths) + if len(ref) < 2: + return None + mean_l = statistics.mean(ref) + std_l = statistics.stdev(ref) + if std_l <= 0: + return None + a = mean_l + b = mean_l + n_std * std_l + return {"a": float(a), "b": float(b), "f": float(f_value)} + + +def build_band(row: dict[str, Any], cfg: dict[str, Any]) -> dict[str, Any] | None: + """Construct the profile_band dict for this row, or None to skip.""" + f = lookup_f(row.get("pass_rate", -1.0), cfg["_f_table"]) + if f is None: + return None + rewards = row.get("profiled_rewards") or [] + band: dict[str, Any] = {} + for ch_name in cfg["channels"]: + lengths_key = CHANNEL_TO_LENGTHS_KEY[ch_name] + lengths = row.get(lengths_key) + if lengths is None: + continue + block = channel_block( + lengths=lengths, + rewards=rewards, + n_std=cfg["n_std"], + min_passing=cfg["min_passing"], + f_value=f, + ) + if block is not None: + band[ch_name] = block + return band or None + + +def main() -> None: + args = parse_args() + cfg = load_config(args.config) + + in_path = Path(args.input) + out_path = Path(args.output) + out_path.parent.mkdir(parents=True, exist_ok=True) + + n_total = 0 + n_with_band = 0 + n_pr_skipped = 0 + n_pr_skipped_by_passrate: Counter[float] = Counter() + n_channels_emitted: Counter[str] = Counter() + + with in_path.open() as fin, out_path.open("w") as fout: + for line in fin: + line = line.rstrip("\n") + if not line: + continue + row = json.loads(line) + n_total += 1 + band = build_band(row, cfg) + if band is None: + pr = round(float(row.get("pass_rate", -1.0)), 4) + if lookup_f(pr, cfg["_f_table"]) is None: + n_pr_skipped += 1 + n_pr_skipped_by_passrate[pr] += 1 + else: + row["profile_band"] = band + n_with_band += 1 + for ch in band: + n_channels_emitted[ch] += 1 + fout.write(json.dumps(row) + "\n") + + if not args.quiet: + print(f"input : {in_path}") + print(f"output : {out_path}") + print(f"rows in: {n_total}") + print( + f"rows w/ profile_band: {n_with_band} ({100 * n_with_band / n_total:.1f}%)" + ) + print(f"rows skipped (pass_rate not in f_table): {n_pr_skipped}") + if n_pr_skipped_by_passrate: + print(" by pass_rate:") + for pr in sorted(n_pr_skipped_by_passrate): + print(f" {pr:.3f}: {n_pr_skipped_by_passrate[pr]}") + print("channels emitted (per row, summed):") + for ch in cfg["channels"]: + c = n_channels_emitted.get(ch, 0) + print(f" {ch:>10}: {c} ({100 * c / n_total:.1f}% of rows)") + print("f_table (rounded pass_rate -> f):") + for pr in sorted(cfg["_f_table"]): + print(f" {pr:.3f} -> {cfg['_f_table'][pr]}") + print(f"n_std={cfg['n_std']} min_passing={cfg['min_passing']}") + + +if __name__ == "__main__": + main() diff --git a/tests/unit/utils/test_length_penalty.py b/tests/unit/utils/test_length_penalty.py new file mode 100644 index 00000000000..3bde11fd545 --- /dev/null +++ b/tests/unit/utils/test_length_penalty.py @@ -0,0 +1,595 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for the profile_band multiplier and group relative-length +scaling algorithms in nemo_rl/utils/length_penalty.py.""" + +import pytest + +from nemo_rl.utils.length_penalty import ( + _band_multiplier, + apply_group_length_penalties, +) + +AGENT = "math_with_judge_simple_agent" + + +def make_result(reasoning: str, answer: str, reward: float, band=None): + """Build a minimal rollout result dict in the shape the module consumes.""" + result = { + "full_result": { + "reward": reward, + "response": { + "output": [ + {"type": "reasoning", "summary": [{"text": reasoning}]}, + {"type": "message", "content": [{"text": answer}]}, + ] + }, + }, + "agent_ref": {"name": AGENT}, + } + if band is not None: + result["profile_band"] = band + return result + + +def make_config(default=None, profile_band=None, num_gens=2): + length_penalty = {} + if default is not None: + length_penalty["default"] = {"length_type": "chars", **default} + if profile_band is not None: + length_penalty["profile_band"] = profile_band + return { + "grpo": { + "num_generations_per_prompt": num_gens, + "length_penalty": length_penalty, + } + } + + +def rewards_of(results): + return [r["full_result"]["reward"] for r in results] + + +class TestBandMultiplier: + """Direct tests of the {a, b, f} multiplier shape.""" + + CH = {"a": 10, "b": 20, "f": 0.5} + + def test_at_or_below_a_is_one(self): + assert _band_multiplier(5, self.CH) == 1.0 + assert _band_multiplier(10, self.CH) == 1.0 + + def test_linear_interpolation_between_a_and_b(self): + assert _band_multiplier(15, self.CH) == pytest.approx(0.75) + + def test_exactly_b_is_f(self): + assert _band_multiplier(20, self.CH) == pytest.approx(0.5) + + def test_clamps_at_f_past_b(self): + # Past b the multiplier stays at f; it must NOT keep decaying to 0. + assert _band_multiplier(25, self.CH) == pytest.approx(0.5) + assert _band_multiplier(30, self.CH) == pytest.approx(0.5) + assert _band_multiplier(10_000, self.CH) == pytest.approx(0.5) + + def test_missing_or_malformed_channel_is_noop(self): + assert _band_multiplier(100, None) == 1.0 + assert _band_multiplier(100, {}) == 1.0 + assert _band_multiplier(100, {"a": 10, "b": 20}) == 1.0 # missing f + assert _band_multiplier(100, {"a": 20, "b": 10, "f": 0.5}) == 1.0 # b <= a + + +class TestProfileBandPerRow: + """profile_band multipliers driven by per-row dataset metadata.""" + + def test_total_channel_scales_correct_rollouts(self): + band = {"total": {"a": 10, "b": 20, "f": 0.5}} + results = [ + make_result("12345", "12345", 1.0, band=band), # total 10 -> x1.0 + make_result("1234567890", "1234567890", 1.0, band=band), # 20 -> x0.5 + ] + cfg = make_config(default={"enabled": True, "profile_band_total": True}) + apply_group_length_penalties(results, cfg) + assert rewards_of(results) == pytest.approx([1.0, 0.5]) + + def test_zero_reward_rollouts_untouched(self): + band = {"total": {"a": 10, "b": 20, "f": 0.5}} + results = [ + make_result("1234567890", "1234567890", 0.0, band=band), + make_result("1234567890", "1234567890", 1.0, band=band), + ] + cfg = make_config(default={"enabled": True, "profile_band_total": True}) + apply_group_length_penalties(results, cfg) + assert rewards_of(results) == pytest.approx([0.0, 0.5]) + + def test_reasoning_channel_ignores_answer_length(self): + band = {"reasoning": {"a": 10, "b": 20, "f": 0.5}} + long_answer = "x" * 100 # must not affect the reasoning channel + results = [ + make_result("12345", long_answer, 1.0, band=band), # reasoning 5 -> x1.0 + make_result("123456789012345", long_answer, 1.0, band=band), # 15 -> x0.75 + ] + cfg = make_config(default={"enabled": True, "profile_band_reasoning": True}) + apply_group_length_penalties(results, cfg) + assert rewards_of(results) == pytest.approx([1.0, 0.75]) + + def test_missing_row_band_is_noop(self): + results = [ + make_result("12345", "12345", 1.0), + make_result("1234567890123456789012345", "12345", 1.0), + ] + cfg = make_config(default={"enabled": True, "profile_band_total": True}) + apply_group_length_penalties(results, cfg) + assert rewards_of(results) == pytest.approx([1.0, 1.0]) + + def test_channel_not_enabled_in_config_is_noop(self): + band = {"total": {"a": 10, "b": 20, "f": 0.5}} + results = [ + make_result("1234567890", "1234567890", 1.0, band=band), + make_result("12345", "12345", 1.0, band=band), + ] + cfg = make_config(default={"enabled": True}) # no profile_band_* flag + apply_group_length_penalties(results, cfg) + assert rewards_of(results) == pytest.approx([1.0, 1.0]) + + +class TestProfileBandGlobalDefaults: + """profile_band driven by config-level length_penalty.profile_band defaults.""" + + def test_global_total_only(self): + cfg = make_config( + default={"enabled": True}, + profile_band={ + "enabled": True, + "defaults": {"total": {"a": 10, "b": 20, "f": 0.5}}, + }, + ) + results = [ + make_result("12345", "12345", 1.0), # total 10 -> x1.0 + make_result("1234567890123456789012345", "12345", 1.0), # 30 -> x0.5 + ] + apply_group_length_penalties(results, cfg) + assert rewards_of(results) == pytest.approx([1.0, 0.5]) + + def test_global_works_without_default_block(self): + # Channels under defaults are implicitly enabled; no length_penalty.default + # `enabled` or profile_band_* booleans required. + cfg = make_config( + default={}, # only length_type + profile_band={ + "enabled": True, + "defaults": {"reasoning": {"a": 10, "b": 20, "f": 0.5}}, + }, + ) + results = [ + make_result("123456789012345", "xx", 1.0), # reasoning 15 -> x0.75 + make_result("12345", "xx", 1.0), # 5 -> x1.0 + ] + apply_group_length_penalties(results, cfg) + assert rewards_of(results) == pytest.approx([0.75, 1.0]) + + def test_row_band_wins_over_global(self): + cfg = make_config( + default={"enabled": True}, + profile_band={ + "enabled": True, + "defaults": {"total": {"a": 10, "b": 20, "f": 0.5}}, + }, + ) + generous = {"total": {"a": 100, "b": 200, "f": 0.5}} + results = [ + make_result("12345", "12345", 1.0, band=generous), + make_result("1234567890123456789012345", "12345", 1.0, band=generous), + ] + apply_group_length_penalties(results, cfg) + assert rewards_of(results) == pytest.approx([1.0, 1.0]) + + def test_disabled_block_is_noop(self): + cfg = make_config( + default={"enabled": True}, + profile_band={ + "enabled": False, + "defaults": {"total": {"a": 10, "b": 20, "f": 0.5}}, + }, + ) + results = [ + make_result("1234567890123456789012345", "12345", 1.0), + make_result("12345", "12345", 1.0), + ] + apply_group_length_penalties(results, cfg) + assert rewards_of(results) == pytest.approx([1.0, 1.0]) + + def test_malformed_global_channel_ignored(self): + cfg = make_config( + default={"enabled": True}, + profile_band={ + "enabled": True, + "defaults": {"total": {"a": 20, "b": 10, "f": 0.5}}, # b <= a + }, + ) + results = [ + make_result("1234567890123456789012345", "12345", 1.0), + make_result("12345", "12345", 1.0), + ] + apply_group_length_penalties(results, cfg) + assert rewards_of(results) == pytest.approx([1.0, 1.0]) + + +class TestGroupRelativeLengthScaling: + """Dense zero-centered group relative-length penalty.""" + + def test_two_rollouts_symmetric_adjustment(self): + # lengths 10 and 30: raw weights 1 and 0, centered +0.5/-0.5, coeff 0.1. + cfg = make_config( + default={"enabled": True, "group_total_length_penalty_coeff": 0.1} + ) + results = [ + make_result("12345", "12345", 1.0), # total 10 -> +0.05 + make_result("1234567890123456789012345", "12345", 1.0), # 30 -> -0.05 + ] + apply_group_length_penalties(results, cfg) + assert rewards_of(results) == pytest.approx([1.05, 0.95]) + + def test_three_rollouts_zero_centered(self): + # lengths 10/20/30 -> raw weights 1/0.5/0 -> centered +0.5/0/-0.5. + cfg = make_config( + default={"enabled": True, "group_total_length_penalty_coeff": 0.1}, + num_gens=3, + ) + results = [ + make_result("12345", "12345", 1.0), # 10 + make_result("1234567890", "1234567890", 1.0), # 20 + make_result("123456789012345", "123456789012345", 1.0), # 30 + ] + apply_group_length_penalties(results, cfg) + assert rewards_of(results) == pytest.approx([1.05, 1.0, 0.95]) + # Zero-centered: the group's mean reward is unchanged by the adjustment. + assert sum(rewards_of(results)) == pytest.approx(3.0) + + def test_equal_lengths_no_adjustment(self): + cfg = make_config( + default={"enabled": True, "group_total_length_penalty_coeff": 0.1} + ) + results = [ + make_result("12345", "12345", 1.0), + make_result("12345", "12345", 1.0), + ] + apply_group_length_penalties(results, cfg) + assert rewards_of(results) == pytest.approx([1.0, 1.0]) + + def test_only_positive_rollouts_participate(self): + # The zero-reward rollout is neither adjusted nor part of min/max, so + # the two positives (10 and 30) still get the symmetric +/-0.05. + cfg = make_config( + default={"enabled": True, "group_total_length_penalty_coeff": 0.1}, + num_gens=3, + ) + results = [ + make_result("12345", "12345", 1.0), # 10 -> +0.05 + make_result("1" * 1000, "1" * 1000, 0.0), # untouched, excluded + make_result("1234567890123456789012345", "12345", 1.0), # 30 -> -0.05 + ] + apply_group_length_penalties(results, cfg) + assert rewards_of(results) == pytest.approx([1.05, 0.0, 0.95]) + + def test_reasoning_channel_uses_reasoning_length_only(self): + # Same total lengths, different reasoning/answer split: only the + # reasoning coefficient is on, so the shorter-reasoning rollout wins. + cfg = make_config( + default={"enabled": True, "group_reasoning_length_penalty_coeff": 0.1} + ) + results = [ + make_result("12345", "123456789012345", 1.0), # reasoning 5 -> +0.05 + make_result("123456789012345", "12345", 1.0), # reasoning 15 -> -0.05 + ] + apply_group_length_penalties(results, cfg) + assert rewards_of(results) == pytest.approx([1.05, 0.95]) + + def test_zero_coefficient_is_noop(self): + cfg = make_config( + default={"enabled": True, "group_total_length_penalty_coeff": 0.0} + ) + results = [ + make_result("12345", "12345", 1.0), + make_result("1234567890123456789012345", "12345", 1.0), + ] + apply_group_length_penalties(results, cfg) + assert rewards_of(results) == pytest.approx([1.0, 1.0]) + + def test_agent_override_disables_for_agent(self): + cfg = make_config( + default={"enabled": True, "group_total_length_penalty_coeff": 0.1} + ) + cfg["grpo"]["length_penalty"]["agent_overrides"] = {AGENT: {"enabled": False}} + results = [ + make_result("12345", "12345", 1.0), + make_result("1234567890123456789012345", "12345", 1.0), + ] + apply_group_length_penalties(results, cfg) + assert rewards_of(results) == pytest.approx([1.0, 1.0]) + + +class TestProfiledLengthPenalty: + """Profiled-length threshold penalty and its passing-samples requirement.""" + + @staticmethod + def make_profiled(reasoning, answer, reward, p_rewards, p_lengths): + r = make_result(reasoning, answer, reward) + r["profiled_rewards"] = p_rewards + r["profiled_output_lengths"] = p_lengths + return r + + def cfg(self, min_samples=2): + return make_config( + default={ + "enabled": True, + "profiled_length_penalty": 0.3, + "profiled_length_n_std": 1.0, + "profiled_length_min_samples": min_samples, + } + ) + + def test_enough_passes_penalizes_over_threshold(self): + # Passing profiled lengths 10 and 14: threshold = 12 + 1*std(=~2.83) ~ 14.83. + p_rewards, p_lengths = [1, 1, 0], [10, 14, 100] + results = [ + self.make_profiled("12345", "12345", 1.0, p_rewards, p_lengths), # 10 < thr + self.make_profiled( + "1234567890", "1234567890", 1.0, p_rewards, p_lengths + ), # 20 >= thr + ] + apply_group_length_penalties(results, self.cfg()) + assert rewards_of(results) == pytest.approx([1.0, 0.7]) + # The failing profiled length (100) must not have entered the threshold: + # with it, mean+std would exceed 20 and nothing would be penalized. + + def test_one_pass_below_min_samples_no_penalty(self): + # Only 1 passing profiled rollout with min_samples=2: no penalty for + # anyone — no fallback to failing profiled lengths. + p_rewards, p_lengths = [1, 0, 0], [10, 100, 120] + results = [ + self.make_profiled("1234567890", "1234567890", 1.0, p_rewards, p_lengths), + self.make_profiled("1" * 50, "1" * 50, 1.0, p_rewards, p_lengths), + ] + apply_group_length_penalties(results, self.cfg()) + assert rewards_of(results) == pytest.approx([1.0, 1.0]) + + def test_zero_passes_no_penalty(self): + p_rewards, p_lengths = [0, 0, 0], [10, 12, 14] + results = [ + self.make_profiled("1234567890", "1234567890", 1.0, p_rewards, p_lengths), + self.make_profiled("1" * 50, "1" * 50, 1.0, p_rewards, p_lengths), + ] + apply_group_length_penalties(results, self.cfg()) + assert rewards_of(results) == pytest.approx([1.0, 1.0]) + + def test_one_pass_allowed_when_min_samples_is_one(self): + # min_samples=1 opts in to single-pass thresholds: threshold = 10 + 0. + p_rewards, p_lengths = [1, 0], [10, 100] + results = [ + self.make_profiled("1234", "1234", 1.0, p_rewards, p_lengths), # 8 < 10 + self.make_profiled( + "1234567890", "1234567890", 1.0, p_rewards, p_lengths + ), # 20 >= 10 + ] + apply_group_length_penalties(results, self.cfg(min_samples=1)) + assert rewards_of(results) == pytest.approx([1.0, 0.7]) + + +class TestPassRateLengthPenalty: + """MAI-style penalty: -w * rho_q * |y_i| / l_max on correct rollouts.""" + + def cfg(self, w=0.2, num_gens=2): + return make_config( + default={"enabled": True, "pass_rate_length_penalty_weight": w}, + num_gens=num_gens, + ) + + def test_all_correct_group_scales_by_relative_length(self): + # rho_q = 1.0, l_max = 20: penalties are w*10/20 and w*20/20. + results = [ + make_result("12345", "12345", 1.0), # total 10 -> -0.1 + make_result("1234567890", "1234567890", 1.0), # total 20 -> -0.2 + ] + apply_group_length_penalties(results, self.cfg(w=0.2)) + assert rewards_of(results) == pytest.approx([0.9, 0.8]) + + def test_pass_rate_scales_penalty(self): + # 1 of 2 correct -> rho_q = 0.5; only the correct rollout is penalized: + # 1.0 - 0.2 * 0.5 * 20/20 = 0.9. The wrong rollout stays at 0. + results = [ + make_result("1234567890", "1234567890", 1.0), # l_max contributor + make_result("12345", "12345", 0.0), + ] + apply_group_length_penalties(results, self.cfg(w=0.2)) + assert rewards_of(results) == pytest.approx([0.9, 0.0]) + + def test_all_wrong_group_gets_zero_penalty(self): + # rho_q = 0 -> no penalty at all, group stays variance-free. + results = [ + make_result("12345", "12345", 0.0), + make_result("1234567890", "1234567890", 0.0), + ] + apply_group_length_penalties(results, self.cfg(w=0.2)) + assert rewards_of(results) == pytest.approx([0.0, 0.0]) + + def test_zero_weight_is_noop(self): + results = [ + make_result("12345", "12345", 1.0), + make_result("1234567890", "1234567890", 1.0), + ] + apply_group_length_penalties(results, self.cfg(w=0.0)) + assert rewards_of(results) == pytest.approx([1.0, 1.0]) + + def test_wrong_rollout_length_does_not_set_normalizer(self): + # l_max comes from CORRECT rollouts only: the wrong rollout's total of + # 40 is ignored, the correct rollout (total 20) is its own max and + # loses exactly w * rho = 0.2 * 0.5. A long wrong ramble must not + # dilute the pressure on correct rollouts. + results = [ + make_result("1234567890", "1234567890", 1.0), + make_result("1" * 20, "1" * 20, 0.0), + ] + apply_group_length_penalties(results, self.cfg(w=0.2)) + assert rewards_of(results) == pytest.approx([1.0 - 0.2 * 0.5, 0.0]) + + +class TestReviewFixes: + """Regression tests for review findings on PR #3852.""" + + def test_flat_penalty_exceeding_reward_clamps_at_zero(self): + # Stacked/oversized flat penalties may wipe a correct rollout's reward + # out but never flip its sign: 1.0 - 1.5 clamps to 0.0, not -0.5. + results = [ + make_result("1234", "1234", 1.0), # len 8 < threshold 10: untouched + make_result("1234567890", "1234567890", 1.0), # len 20 >= 10: clamped + ] + for r in results: + r["profiled_rewards"] = [1, 1] + r["profiled_output_lengths"] = [10, 10] + cfg = make_config(default={"enabled": True, "profiled_length_penalty": 1.5}) + apply_group_length_penalties(results, cfg) + assert rewards_of(results) == pytest.approx([1.0, 0.0]) + + def test_negative_env_reward_group_skipped(self): + # Length adjustments require binary rewards: a group containing a + # negative env reward is skipped wholesale — no adjustment, no clamp. + cfg = make_config( + default={"enabled": True, "group_total_length_penalty_coeff": 0.1} + ) + results = [ + make_result("12345", "12345", -1.0), + make_result("1234567890", "1234567890", 1.0), + ] + apply_group_length_penalties(results, cfg) + assert rewards_of(results) == pytest.approx([-1.0, 1.0]) + + def test_graded_rewards_group_skipped(self): + # Graded (non-binary) rewards also skip the group untouched. + cfg = make_config( + default={"enabled": True, "group_total_length_penalty_coeff": 0.1} + ) + results = [ + make_result("12345", "12345", 0.5), + make_result("1234567890", "1234567890", 1.0), + ] + apply_group_length_penalties(results, cfg) + assert rewards_of(results) == pytest.approx([0.5, 1.0]) + + def test_all_wrong_group_stays_variance_free(self): + # profiled_length_penalty has no positive-reward gate, but the binary + # clamp floors penalty-created negatives at 0: an all-wrong group must + # stay all-zero (no within-group variance -> no GRPO gradient). + results = [ + make_result("1234", "1234", 0.0), # len 8, under threshold + make_result("1234567890", "1234567890", 0.0), # len 20, over + ] + for r in results: + r["profiled_rewards"] = [1, 1] + r["profiled_output_lengths"] = [10, 10] + cfg = make_config(default={"enabled": True, "profiled_length_penalty": 0.3}) + apply_group_length_penalties(results, cfg) + assert rewards_of(results) == pytest.approx([0.0, 0.0]) + + def test_band_multiplier_never_rewards_length_on_penalized_base(self): + # Flat penalty exceeds the reward, so bases clamp to 0; the band phase + # must leave them at 0 (base * m - base > 0 for base < 0 would have + # made longer rollouts score HIGHER pre-clamp). + band = {"total": {"a": 1, "b": 2, "f": 0.1}} + results = [ + make_result("12345", "12345", 1.0, band=band), + make_result("1234567890", "1234567890", 1.0, band=band), + ] + for r in results: + r["profiled_rewards"] = [1, 1] + r["profiled_output_lengths"] = [5, 5] + cfg = make_config( + default={ + "enabled": True, + "profiled_length_penalty": 1.5, # threshold 5: both clamped to 0 + "profile_band_total": True, + } + ) + apply_group_length_penalties(results, cfg) + r_short, r_long = rewards_of(results) + assert r_short == pytest.approx(0.0) + assert r_long == pytest.approx(0.0) + assert r_long <= r_short # longer must never beat shorter + + def test_explicit_false_channel_not_overridden_by_global_defaults(self): + cfg = make_config( + default={"enabled": True, "profile_band_total": False}, + profile_band={ + "enabled": True, + "defaults": {"total": {"a": 10, "b": 20, "f": 0.5}}, + }, + ) + results = [ + make_result("1234567890123456789012345", "12345", 1.0), # total 30 + make_result("12345", "12345", 1.0), + ] + apply_group_length_penalties(results, cfg) + assert rewards_of(results) == pytest.approx([1.0, 1.0]) + + def test_default_block_without_enabled_is_active(self): + # `enabled` defaults True consistently: a default: block that omits it + # applies its penalties (previously the early return used default False + # and silently no-oped). + cfg = make_config(default={"group_total_length_penalty_coeff": 0.2}) + results = [ + make_result("12345", "12345", 1.0), + make_result("1234567890123456789012345", "12345", 1.0), + ] + apply_group_length_penalties(results, cfg) + assert rewards_of(results) == pytest.approx([1.1, 0.9]) + + def test_behavior_independent_of_empty_agent_overrides(self): + # The same default: block must behave identically with or without an + # unrelated empty agent_overrides entry (previously it flipped the + # early return and changed behavior). + base = {"group_total_length_penalty_coeff": 0.2} + outs = [] + for extra_overrides in (None, {"some_agent": {}}): + cfg = make_config(default=dict(base)) + if extra_overrides is not None: + cfg["grpo"]["length_penalty"]["agent_overrides"] = extra_overrides + results = [ + make_result("12345", "12345", 1.0), + make_result("1234567890123456789012345", "12345", 1.0), + ] + apply_group_length_penalties(results, cfg) + outs.append(rewards_of(results)) + assert outs[0] == pytest.approx(outs[1]) + + def test_unknown_key_raises(self): + cfg = make_config(default={"enabled": True, "group_total_length_coeff": 0.1}) + results = [make_result("12345", "12345", 1.0), make_result("123", "123", 1.0)] + try: + apply_group_length_penalties(results, cfg) + except ValueError as e: + assert "group_total_length_coeff" in str(e) + else: + raise AssertionError("expected ValueError for unknown config key") + + def test_longest_penalty_under_binary_rewards(self): + # Under binary rewards all correct rollouts tie at the top score, so + # top_percentile is inert and both rollouts are eligible top scorers; + # the longest one takes the flat penalty (default top_percentile 0.5). + results = [ + make_result("12345", "12345", 1.0), + make_result("1234567890", "1234567890", 1.0), + ] + cfg = make_config(default={"enabled": True, "longest_total_penalty": 0.2}) + apply_group_length_penalties(results, cfg) + assert rewards_of(results) == pytest.approx([1.0, 0.8])