fix: detect rollouts truncated by the per-turn generation cap

nemo_rl/experience/rollouts.py

@@ -2313,6 +2313,63 @@ async def _consume_rollout() -> NemoGymRolloutResult:
     return asyncio.run(_consume_rollout())
 
 
+def _nemo_gym_sample_metrics(
+    result: dict,
+    *,
+    max_total_tokens: int,
+    max_new_tokens: Optional[int],
+) -> dict[str, Any]:
+    """Per-sample rollout metrics, including whether generation was cut short.
+
+    ``hit_max_tokens`` becomes ``final_batch["truncated"]``, which
+    ``grpo.overlong_filtering`` uses to drop a sample from the loss.
+
+    Prefer the agent's own ``is_truncated``. It is the same ground truth the native
+    generation path uses -- the model server's ``finish_reason == "length"`` -- carried
+    through the agent, so it distinguishes "the budget stopped this" from "the model
+    emitted EOS" without knowing which budget applied. NeMo-Gym's ``verifiers_agent``
+    reports it; agents that do not simply omit the field.
+
+    Fall back to comparing lengths when it is absent. A rollout can be stopped by either
+    of two budgets, so both are checked: the conversation filling the context
+    (``max_total_tokens``), or a single turn reaching the per-turn generation cap
+    (``max_new_tokens``). Checking only the first misses every rollout that stopped below
+    the context, which is not a corner case -- ``max_new_tokens`` applies per turn, so a
+    run with a small ``max_new_tokens`` never fills the context at all and its truncated
+    samples keep the near-zero reward that overlong filtering exists to discard.
+
+    The fallback compares with ``>=`` rather than ``==``. The message log is
+    post-processed before it reaches here -- reasoning content is re-wrapped in thinking
+    tags, for one -- so an exact match can be thrown off by a single token. It also
+    cannot tell a turn that ran out of budget from one that happened to emit EOS on the
+    final allowed token, which is why the reported flag wins when there is one.
+    """
+    message_log = result["message_log"]
+    assistant_lengths = [
+        len(m["token_ids"]) for m in message_log if m["role"] == "assistant"
+    ]
+    total_tokens = sum(len(m["token_ids"]) for m in message_log)
+    max_gen_tokens_per_turn = max(assistant_lengths, default=0)
+
+    reported_truncated = (result.get("full_result") or {}).get("is_truncated")
+    if reported_truncated is not None:
+        hit_max_tokens = bool(reported_truncated)
+    else:
+        hit_max_tokens = total_tokens >= max_total_tokens
+        if max_new_tokens is not None:
+            hit_max_tokens = hit_max_tokens or max_gen_tokens_per_turn >= max_new_tokens
+
+    return {
+        "total_reward": result["full_result"]["reward"],
+        "assistant_tokens": sum(assistant_lengths),
+        "total_tokens": total_tokens,
+        "turn_count": sum(1 for m in message_log if m["role"] == "user"),
+        "hit_max_tokens": hit_max_tokens,
+        # max_gen_tokens_per_turn: Diagnostic for long single generations
+        "max_gen_tokens_per_turn": max_gen_tokens_per_turn,
+    }
+
+
 def _postprocess_single_nemo_gym_group(
     nemo_gym_rows: list[dict],
     results: list[dict],
@@ -2359,28 +2416,15 @@ def _postprocess_single_nemo_gym_group(
             max_total_tokens_per_sample = policy_generation.cfg[
                 "max_total_sequence_length"
             ]
+        # Absent for backends whose generation config omits it; the context check then
+        # applies on its own, matching the previous behaviour.
+        max_new_tokens = policy_generation.cfg.get("max_new_tokens")
         all_sample_metrics = [
-            {
-                "total_reward": r["full_result"]["reward"],
-                "assistant_tokens": sum(
-                    len(m["token_ids"])
-                    for m in r["message_log"]
-                    if m["role"] == "assistant"
-                ),
-                "total_tokens": sum(len(m["token_ids"]) for m in r["message_log"]),
-                "turn_count": sum(1 for m in r["message_log"] if m["role"] == "user"),
-                "hit_max_tokens": sum(len(m["token_ids"]) for m in r["message_log"])
-                == max_total_tokens_per_sample,
-                # max_gen_tokens_per_turn: Diagnostic for long single generations
-                "max_gen_tokens_per_turn": max(
-                    (
-                        len(m["token_ids"])
-                        for m in r["message_log"]
-                        if m["role"] == "assistant"
-                    ),
-                    default=0,
-                ),
-            }
+            _nemo_gym_sample_metrics(
+                r,
+                max_total_tokens=max_total_tokens_per_sample,
+                max_new_tokens=max_new_tokens,
+            )
             for r in results
         ]
 

tests/unit/experience/test_rollouts.py

@@ -1299,6 +1299,124 @@ def test_postprocess_nemo_gym_group_returns_task_index(log_full_result_tables):
     ) is log_full_result_tables
 
 
+def _gym_result(prompt_tokens, assistant_turns, reward=1.0, is_truncated=None):
+    message_log = [
+        {
+            "role": "user",
+            "content": "prompt",
+            "token_ids": torch.zeros(prompt_tokens, dtype=torch.long),
+        }
+    ]
+    for turn_tokens in assistant_turns:
+        message_log.append(
+            {
+                "role": "assistant",
+                "content": "answer",
+                "token_ids": torch.zeros(turn_tokens, dtype=torch.long),
+                "generation_logprobs": torch.zeros(turn_tokens),
+            }
+        )
+    full_result = {"reward": reward}
+    if is_truncated is not None:
+        full_result["is_truncated"] = is_truncated
+    return {
+        "input_message_log": message_log[:1],
+        "message_log": message_log,
+        "full_result": full_result,
+    }
+
+
+@pytest.mark.parametrize(
+    "prompt_tokens,assistant_turns,max_total_tokens,max_new_tokens,expected",
+    [
+        # The conversation filled the context. Detected before this check also
+        # considered the per-turn cap, and still detected now.
+        (10, [90], 100, None, True),
+        # A turn reached the generation cap well inside the context. This is the case
+        # that used to go unflagged, so overlong_filtering silently kept the sample.
+        (10, [64], 1024, 64, True),
+        # Neither budget reached: no single turn hits the cap, context has room left.
+        (10, [20, 20], 1024, 64, False),
+        # `>=`, not `==`: the message log is post-processed, so an exact match can be
+        # thrown off by a token or two.
+        (10, [95], 100, None, True),
+        # No generation cap configured: context check alone, matching prior behaviour.
+        (10, [64], 1024, None, False),
+    ],
+)
+def test_nemo_gym_sample_metrics_detects_both_truncation_budgets(
+    prompt_tokens, assistant_turns, max_total_tokens, max_new_tokens, expected
+):
+    metrics = rollouts_mod._nemo_gym_sample_metrics(
+        _gym_result(prompt_tokens, assistant_turns),
+        max_total_tokens=max_total_tokens,
+        max_new_tokens=max_new_tokens,
+    )
+
+    assert metrics["hit_max_tokens"] is expected
+
+
+@pytest.mark.parametrize(
+    "reported,assistant_turns,max_new_tokens,expected",
+    [
+        # Reported flag wins over the length fallback in BOTH directions. The second case
+        # is the one lengths cannot get right: a turn that emitted EOS on its final
+        # allowed token looks identical to one that ran out of budget.
+        (True, [4], 64, True),
+        (False, [64], 64, False),
+    ],
+)
+def test_nemo_gym_sample_metrics_prefers_the_reported_flag(
+    reported, assistant_turns, max_new_tokens, expected
+):
+    metrics = rollouts_mod._nemo_gym_sample_metrics(
+        _gym_result(10, assistant_turns, is_truncated=reported),
+        max_total_tokens=1024,
+        max_new_tokens=max_new_tokens,
+    )
+
+    assert metrics["hit_max_tokens"] is expected
+
+
+def test_nemo_gym_sample_metrics_reports_lengths():
+    metrics = rollouts_mod._nemo_gym_sample_metrics(
+        _gym_result(10, [30, 20], reward=2.5),
+        max_total_tokens=1024,
+        max_new_tokens=None,
+    )
+
+    assert metrics["total_reward"] == 2.5
+    assert metrics["assistant_tokens"] == 50
+    assert metrics["total_tokens"] == 60
+    assert metrics["turn_count"] == 1
+    assert metrics["max_gen_tokens_per_turn"] == 30
+
+
+def test_postprocess_nemo_gym_group_flags_generation_cap_truncation():
+    """final_batch["truncated"] is what grpo.overlong_filtering masks on."""
+    rows = [{"agent_ref": {"name": "agent"}} for _ in range(2)]
+    # First rollout stops at the generation cap; second finishes early. Neither comes
+    # close to max_model_len, so the context check alone would flag neither.
+    results = [_gym_result(10, [16]), _gym_result(10, [4])]
+
+    rollout_result = rollouts_mod._postprocess_single_nemo_gym_group(
+        nemo_gym_rows=rows,
+        results=results,
+        timer=rollouts_mod.Timer(),
+        timer_prefix="timing/rollout",
+        policy_generation=type(
+            "_PolicyGeneration",
+            (),
+            {"cfg": {"vllm_cfg": {"max_model_len": 1024}, "max_new_tokens": 16}},
+        )(),
+        input_batch=BatchedDataDict({"loss_multiplier": torch.ones(2)}),
+        tokenizer=type("_Tokenizer", (), {"pad_token_id": 0})(),
+        log_full_result_tables=False,
+    )
+
+    assert rollout_result.final_batch["truncated"].tolist() == [True, False]
+
+
 def test_run_nemo_gym_rollout_sync_drains_entire_batch(monkeypatch):
     input_batch = BatchedDataDict({"loss_multiplier": torch.ones(3)})
     expected = rollouts_mod.NemoGymRolloutResult(