[trainer] feat: pass non_tensor_batch to advantage estimators that can accept it

[waple0820]

Problem

compute_advantage() hands a registered advantage estimator tensors plus index. Everything a rollout attached per sample — a validity flag, a task id, benchmark metadata — lives in DataProto.non_tensor_batch and never reaches the estimator, so a custom estimator cannot use it.

The data is already there. AgentLoopOutput.extra_fields is promoted into non_tensor_batch by the agent loop (verl/experimental/agent_loop/agent_loop.py), so a rollout can attach anything it wants per sample. It just is not forwarded one function further.

GDPO hit this first and got a named special case in verl/trainer/ppo/ray_trainer.py:

# GDPO: pass raw data for per-dimension reward extraction
if adv_estimator in (AdvantageEstimator.GDPO, "gdpo"):
    adv_kwargs["non_tensor_batch"] = data.non_tensor_batch
    adv_kwargs["batch"] = data.batch

That works, but it means an estimator has to be named in ray_trainer.py to see its own data — a custom estimator registered through register_adv_est cannot.

Change

Pass both keys to any estimator whose signature can accept them (a declared parameter, or **kwargs), which subsumes the GDPO case:

for key, value in (("non_tensor_batch", data.non_tensor_batch), ("batch", data.batch)):
    if _accepts_kwarg(adv_estimator_fn, key):
        adv_kwargs[key] = value

Why this is safe for every estimator registered today

I audited all twelve @register_adv_est functions in core_algos.py:

Estimator	Signature	Effect
gdpo	declares non_tensor_batch, has **kwargs	identical to the special case it replaces
grpo_passk, reinforce_plus_plus_baseline, rloo, opo, reinforce_plus_plus, remax, gpg, optimal_token_baseline, multi_turn_optimal_token_baseline	**kwargs	receive two more keys and ignore them, as they already do for other optional keys
grpo_vectorized	fixed signature, no **kwargs	receives exactly what it did before — this is why the check exists rather than passing unconditionally
gae, grpo	never reach this branch (own branches above)	unchanged

Passing unconditionally would raise TypeError on grpo_vectorized, which is the failure mode the signature check avoids.

Motivation

A rollout that fails for infrastructure reasons — a lost environment session, an unreachable judge, an OOM-killed container — scores reward=0, which is indistinguishable from a policy that genuinely scored zero. In GRPO that zero does not merely contribute no gradient: it lowers the group baseline and inflates the advantage of whatever survived, so the batch trains away from correct behaviour.

We hit this training a browser agent: at 64-way rollout concurrency, orphaned environment sessions caused 1172 of 1280 rollouts across 20 steps (91.6%) to fail before the policy ever acted, and mean reward fell from 15.45% to 0.39%. Excluding those samples from the group statistics needs a per-sample flag inside the estimator, which is what this change makes possible.

Upstream context, for what produces such a flag: NeMo Gym is adding mask_sample to its verify contract so an environment can report an infrastructure failure rather than an implicit zero (NVIDIA-NeMo/Gym#2608, NVIDIA-NeMo/Gym#2611). With this change a recipe can consume it with a plain @register_adv_est estimator; without it, the only options are monkey-patching compute_advantage or vendoring core_algos.py.

Tests

tests/trainer/ppo/test_adv_estimator_kwargs_on_cpu.py:

• _accepts_kwarg recognises a declared parameter and a **kwargs catch-all, and rejects a fixed signature.
• A registered estimator receives non_tensor_batch and can read a per-sample field from it.
• An estimator with a fixed signature is still called successfully.

Note on the V1 trainer

verl/trainer/ppo/v1/trainer_base.py::_compute_advantage selects a fixed field list from the TransferQueue (uid, response_mask, rm_scores, rollout_log_probs, old_log_probs, ref_log_prob, values) before building the DataProto, so under V1 a custom non-tensor field still does not reach the estimator even with this change. I left that out to keep this PR to one idea, but I am happy to follow up with a way to opt extra fields into that selection if you think that is the right direction.

[CLAassistant]

All committers have signed the CLA.