per_rollout_mean divides by training rows instead of rollouts, so the loss scale changes when a rollout spans several rows

Author: KunyangZhangCreated Sep 12, 2026Updated Sep 17, 2026

per_rollout_mean is the default loss mode (agentlightning/verl/config.yaml:39). In trainer.py:617-626 the advantages are rescaled once for the whole batch before the actor update:

python
normalize_advantages_by_rollout(
    batch.batch["advantages"],
    batch.batch["response_mask"],
    rollout_ids,
    num_trained_rows=len(batch),   # rows, not rollouts
)

normalize_advantages_by_rollout (per_rollout_loss.py:16-41) then divides each row by rollout_token_count * num_trained_rows.

That is exactly the per-rollout mean when a rollout occupies a single row. But a rollout can span several rows: trace_aggregator.level: transition emits one row per turn, and even the default trajectory level flushes an extra row when a turn's prompt does not extend the running context (rollout_adapter.py:542-571, reported as training/n_trace_merge_mismatch_rows). In that case the divisor counts rows, so every rollout keeps 1 / n_rows instead of 1 / n_rollouts and the batch total becomes n_rollouts / n_rows.

Reproduction

Calling the real adapter and then the same normalization trainer.py performs, on main @ 218f1f7:

[A] 2 rollouts, 1 row each   -> per-rollout mass {r1: 0.5,    r2: 0.5},   total = 1.0000
[B] 2 rollouts, 3 rows total -> per-rollout mass {r1: 0.3333, r2: 0.3333}, total = 0.6667
    (r1 hit a trace-merge mismatch; the adapter reports n_trace_merge_mismatch_rows = 1)

Case B differs from A only by a trace-merge mismatch, yet the whole batch carries 2/3 of the signal. The same effect appears in isolation:

python
normalize_advantages_by_rollout(ones(4, 3), ones(4, 3, dtype=long), ["A", "A", "A", "B"],
                                num_trained_rows=4)
# A mass 0.25, B mass 0.25, total 0.5

Question

Is the divisor the number of rollouts or the number of rows? The current test tests/verl/test_per_rollout_loss.py::test_normalize_advantages_by_rollout pins the row-based value (2 rollouts over 3 rows keep 1/3 each), so I did not want to assume it is a bug.

  • If it is meant to be the per-rollout mean: deriving the divisor from the distinct rollout_ids makes a rollout keep 1 / n_rollouts regardless of how it was split. I have that patch ready (per_rollout_loss.py, trainer.py, and the test expectation 1/31/2, plus a regression test over 1/2/3 rows per rollout). Full suite passes: 101 passed.
  • If row-based scaling is intended: the docstring ("rollout's token count and batch size") and the test would be worth stating explicitly, since switching trace_aggregator.level then silently rescales the loss.

Happy to open the PR for the first option, or rework it, once you confirm the intended semantics.

Source: microsoft/agent-lightning