feat(length-reward): support global reasoning profile bands

nemo_rl/utils/length_adjustments.md

@@ -49,6 +49,8 @@ grpo:
       profile_band_total: false
       profile_band_reasoning: false
       profile_band_answer: false
+      reasoning_end_token_id: null
+      profile_band: null
     agent_overrides:
       math_with_judge_simple_agent:
         enabled: true
@@ -117,6 +119,8 @@ grpo:
 | `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. |
+| `reasoning_end_token_id` | Optional model-specific reasoning-end token used to split newly generated assistant token IDs exactly. |
+| `profile_band` | Optional global `{total,reasoning,answer}` band used when a dataset row has no `profile_band` metadata. |
 | `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`. |
@@ -415,7 +419,8 @@ Config keys:
 - `profile_band_reasoning`
 - `profile_band_answer`
 
-This uses per-row `profile_band` metadata with channel-specific `{a, b, f}` values:
+This uses channel-specific `{a, b, f}` values. Per-row `profile_band` metadata
+takes precedence when present:
 
 ```json
 {
@@ -427,15 +432,34 @@ This uses per-row `profile_band` metadata with channel-specific `{a, b, f}` valu
 }
 ```
 
+When rows do not carry that metadata, configure a global fallback under
+`grpo.length_bonus.default`. Agent overrides may replace the fallback:
+
+```yaml
+grpo:
+  length_bonus:
+    default:
+      enabled: true
+      profile_band_reasoning: true
+      reasoning_end_token_id: 13
+      profile_band:
+        reasoning: {a: 1024, b: 4096, f: 0.95}
+```
+
+When `reasoning_end_token_id` is set, reasoning and answer lengths are computed
+from newly generated assistant token IDs only. The original prompt prefix is
+excluded, and the first matching end token in each generated assistant turn is
+the boundary. This supports models such as Nemotron Omni whose prompt opens the
+reasoning block and whose generation emits only the closing token.
+
 For an enabled channel, the multiplier is:
 
 ```text
 length <= a: multiplier = 1
-length > a:  multiplier = max(0, 1 - (length - a) / (b - a) * (1 - f))
+a < length < b: multiplier = 1 - (length - a) / (b - a) * (1 - f)
+length >= b: multiplier = f
 ```
 
-So the multiplier is `f` at `b`, then the same slope continues past `b` until clamped at `0`.
-
 Profile-band multipliers are applied only to rollouts whose original environment reward is
 positive.
 

nemo_rl/utils/length_adjustments.py

@@ -33,8 +33,11 @@
 
 import logging
 import statistics
+from numbers import Real
 from typing import Any
 
+import torch
+
 logger = logging.getLogger(__name__)
 
 # MAD/median floor for zMAD (fixed; matches ``flag_reasoning_length_outliers`` default).
@@ -65,6 +68,8 @@
     "profile_band_total",
     "profile_band_reasoning",
     "profile_band_answer",
+    "profile_band",
+    "reasoning_end_token_id",
     "group_length_penalty_profile_gate",
     "group_length_penalty_profile_gate_channel",
     "group_length_penalty_profile_gate_field",
@@ -87,6 +92,10 @@
     "group_length_penalty_profile_gate_field",
 })
 
+_MAPPING_PARAM_KEYS = frozenset({"profile_band"})
+
+_OPTIONAL_INT_PARAM_KEYS = frozenset({"reasoning_end_token_id"})
+
 _GDPO_LENGTH_FEATURE_PARAM_KEYS = frozenset({
     "reasoning_bonus",
     "answer_bonus",
@@ -176,6 +185,62 @@ def _extract_reasoning_and_answer_text(result: dict[str, Any]) -> tuple[str, str
     return reasoning_text, answer_text
 
 
+def _generated_assistant_token_lengths(
+    result: dict[str, Any], reasoning_end_token_id: int
+) -> tuple[int, int]:
+    """Return reasoning and answer token lengths for newly generated turns.
+
+    NeMo Gym returns the full conversation in ``message_log`` and the original
+    prompt as its prefix in ``input_message_log``. Only messages after that
+    prefix are generated by the current rollout. Each generated assistant turn
+    is split at its first reasoning-end token; an unfinished turn is counted
+    entirely as reasoning.
+    """
+    input_message_count = len(result.get("input_message_log", []))
+    generated_messages = result.get("message_log", [])[input_message_count:]
+
+    reasoning_length = 0
+    answer_length = 0
+    for message in generated_messages:
+        if message.get("role") != "assistant":
+            continue
+        token_ids = message.get("token_ids")
+        if token_ids is None:
+            raise ValueError(
+                "Generated assistant messages must include token_ids when "
+                "reasoning_end_token_id is configured"
+            )
+        if isinstance(token_ids, torch.Tensor):
+            if token_ids.ndim != 1:
+                raise ValueError(
+                    "Generated assistant token_ids must be one-dimensional when "
+                    "reasoning_end_token_id is configured, got "
+                    f"shape={tuple(token_ids.shape)}"
+                )
+            token_id_list = token_ids.tolist()
+        elif isinstance(token_ids, list) and all(
+            isinstance(token_id, int) and not isinstance(token_id, bool)
+            for token_id in token_ids
+        ):
+            token_id_list = token_ids
+        else:
+            raise TypeError(
+                "Generated assistant token_ids must be a one-dimensional integer "
+                "sequence when "
+                "reasoning_end_token_id is configured"
+            )
+
+        try:
+            split_idx = token_id_list.index(reasoning_end_token_id)
+        except ValueError:
+            reasoning_length += len(token_id_list)
+        else:
+            reasoning_length += split_idx
+            answer_length += len(token_id_list) - split_idx - 1
+
+    return reasoning_length, answer_length
+
+
 def _extract_gdpo_length_feature_params(feature_cfg: Any) -> dict[str, Any]:
     if not isinstance(feature_cfg, dict):
         return {}
@@ -297,6 +362,19 @@ def apply_group_length_adjustments(
             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 _MAPPING_PARAM_KEYS:
+            defaults[k] = default_cfg.get(k)
+        elif k in _OPTIONAL_INT_PARAM_KEYS:
+            value = default_cfg.get(k)
+            if value is not None and (
+                isinstance(value, bool)
+                or not isinstance(value, int)
+                or value < 0
+            ):
+                raise ValueError(
+                    f"{k} must be a non-negative integer or null, got {value!r}"
+                )
+            defaults[k] = value
         elif k in _BOOL_PARAM_KEYS:
             defaults[k] = default_cfg.get(k, False)
         elif k == "profiled_length_min_samples":
@@ -350,11 +428,19 @@ def apply_group_length_adjustments(
 
         group_lt = params.pop("length_type", "tokens")
         use_tokens = group_lt == "tokens"
+        reasoning_end_token_id = params.pop("reasoning_end_token_id", None)
+        configured_profile_band = params.pop("profile_band", None)
 
         for k in range(group_size):
             idx = g + k
             r_text, a_text = texts[idx]
-            if use_tokens and tokenizer is not None:
+            if use_tokens and reasoning_end_token_id is not None:
+                reasoning_lengths[idx], answer_lengths[idx] = (
+                    _generated_assistant_token_lengths(
+                        results[idx], reasoning_end_token_id
+                    )
+                )
+            elif 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:
@@ -366,8 +452,9 @@ def apply_group_length_adjustments(
         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]
+        profile_band = results[g].get("profile_band") or configured_profile_band
         gate_info = _group_length_profile_gate_info(
-            band=results[g].get("profile_band"),
+            band=profile_band,
             params=params,
             rewards=group_rewards[:group_size],
             reasoning_lengths=group_reasoning[:group_size],
@@ -701,7 +788,7 @@ def _apply_profile_band_multipliers(
         use_ans = bool(params.get("profile_band_answer", False))
         if not (use_total or use_rsn or use_ans):
             continue
-        band = results[g].get("profile_band")
+        band = results[g].get("profile_band") or params.get("profile_band")
         if not band:
             continue
         ch_total = band.get("total") if use_total else None
@@ -771,18 +858,25 @@ def _band_multiplier(rl: int, ch: dict[str, Any] | None) -> float:
     Otherwise:
         rl <= a    -> 1.0
         rl == b    -> f
-        rl > a     -> same linear slope continues past b, floored at 0.0
+        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:
+    if any(
+        isinstance(value, bool) or not isinstance(value, Real)
+        for value in (a, b, f)
+    ):
+        return 1.0
+    if b <= a or not 0.0 <= float(f) <= 1.0:
         return 1.0
     if rl <= a:
         return 1.0
-    return max(0.0, 1.0 - (rl - a) / (b - a) * (1.0 - float(f)))
+    if rl >= b:
+        return float(f)
+    return 1.0 - (rl - a) / (b - a) * (1.0 - float(f))
 
 
 def _group_length_profile_gate_info(
@@ -894,6 +988,20 @@ def _resolve_agent_params(
                 merged[key] = overrides[key]
             elif key in _BOOL_PARAM_KEYS:
                 merged[key] = bool(overrides[key])
+            elif key in _MAPPING_PARAM_KEYS:
+                value = overrides[key]
+                merged[key] = None if value is None else dict(value)
+            elif key in _OPTIONAL_INT_PARAM_KEYS:
+                value = overrides[key]
+                if value is not None and (
+                    isinstance(value, bool)
+                    or not isinstance(value, int)
+                    or value < 0
+                ):
+                    raise ValueError(
+                        f"{key} must be a non-negative integer or null, got {value!r}"
+                    )
+                merged[key] = value
             else:
                 merged[key] = float(overrides[key])
     return merged

tests/unit/algorithms/test_async_utils.py

@@ -1234,6 +1234,59 @@ async def test_drain_payload_metrics_returns_collector_interval(self, monkeypatc
         assert metrics["payload_bytes/nemo_gym_return/logical_media"] == 150
         assert metrics["payload_ratio/nemo_gym_return/physical_to_logical"] == 0.2
 
+    def test_nemo_gym_forwards_length_adjustment_config(self, monkeypatch):
+        collector = self.create_local_collector()
+        collector.master_config.grpo.length_bonus = {
+            "default": {
+                "enabled": True,
+                "profile_band_reasoning": True,
+            }
+        }
+        collector.master_config.policy["generation"] = {
+            "stop_token_ids": [1],
+            "stop_strings": ["stop"],
+        }
+        repeated_batch = BatchedDataDict(
+            {
+                "extra_env_info": [
+                    {"_ng_task_index": 7},
+                    {"_ng_task_index": 7},
+                ],
+                "loss_multiplier": torch.ones(2),
+            }
+        )
+        captured_kwargs = {}
+
+        async def fake_rollouts(**kwargs):
+            captured_kwargs.update(kwargs)
+            yield SimpleNamespace(
+                task_index=7,
+                final_batch=BatchedDataDict({"loss_multiplier": torch.ones(2)}),
+                rollout_metrics={},
+            )
+
+        import nemo_rl.experience.rollouts as rollouts_mod
+
+        monkeypatch.setattr(rollouts_mod, "run_async_nemo_gym_rollout", fake_rollouts)
+
+        async def collect_groups():
+            return [
+                group
+                async for group in collector._iter_rollout_groups(
+                    repeated_batch=repeated_batch,
+                    num_generations=2,
+                    use_nemo_gym=True,
+                    task_index_to_group_index={7: 0},
+                )
+            ]
+
+        groups = asyncio.run(collect_groups())
+
+        assert len(groups) == 1
+        assert captured_kwargs["length_adjustment_config"] == (
+            collector.master_config.grpo.model_dump()
+        )
+
     def test_collection_loop_marks_errored_on_crash(self):
         """A crash sets errored (not data_exhausted) so driver guards fail fast."""
         collector = self.create_local_collector()