refactor(algorithms): standardize advantage-estimator results by tianyi-zhang-02 · Pull Request #3512 · NVIDIA-NeMo/RL

tianyi-zhang-02 · 2026-08-06T05:00:49Z

What does this PR do?

Standardizes all six advantage estimators on one explicit result type:

AdvantageResult(advantages, returns=None, metrics={})

GRPO, synchronous GRPO, PPO, and SingleController now consume named fields instead of branching on tensors versus tuples. OPD metrics travel in the result instead of through the last_metrics side channel.

The earlier AdvantageEstimator Protocol and its fake None defaults are removed. GDPOAdvantageEstimator.repeated_batch and GeneralizedAdvantageEstimator.values remain honestly required, while the typed estimator configs merged in #3773 stay unchanged.

Addresses #2909.

Validation

Final SHA: 6eb502a06122c27eb4763645a40a5175467f4427, fast-forwarded from the existing PR head and merged with current main (ccbcd4cc5).

Environment: Runpod Secure Cloud, nvcr.io/nvidia/nemo-rl:v0.7.0, Python 3.13.14, PyTorch 2.11.0+cu130; tests ran CPU-only because this refactor changes result/caller contracts rather than GPU kernels.

14 targeted estimator/caller tests passed.
296 related tests across estimator, GRPO, PPO, and SingleController files passed; 7 GPU-only tests skipped.
Ruff check and format-check passed on all 10 changed Python files.

copy-pr-bot · 2026-08-06T05:00:54Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

tianyi-zhang-02 · 2026-08-07T06:41:21Z

Rebased onto current main and pushed a follow-up commit. Removing the isinstance(result, tuple) branch in ppo.py means every compute_advantage test double has to return an AdvantageResult too, and five were missed — four of them in tests that landed on main after this branch was cut:

tests/unit/algorithms/test_ppo.py — DummyAdvantageEstimator returned a bare tuple, so test_ppo_train_excludes_overlong_samples_from_advantage and test_ppo_train_wires_logprob_mask_to_advantage_training_and_metrics failed with AttributeError: 'tuple' object has no attribute 'advantages'.
tests/unit/algorithms/test_grpo.py — three MagicMock return_values set to bare tensors.
tests/unit/single_controller/test_train_pump.py — _FakeAdvEstimator returned a bare tensor; this branch had not touched that file at all.

Verified by running test_ppo.py, test_grpo.py and test_advantage_estimator.py against origin/main and against this branch and diffing the results: identical, with no failures introduced. Before the fix this branch had two failures that main did not.

Worth flagging one thing I noticed while doing this: the AdvantageEstimator Protocol added here has no reference anywhere else in the tree, so nothing is actually checked against it today. It is documentation rather than enforcement unless a call site is annotated with it — happy to either annotate one or drop the Protocol and keep just the dataclass, whichever you prefer.

tianyi-zhang-02 · 2026-08-14T22:36:20Z

Fixed something I should have caught before opening this: the protocol was defined and then referenced nowhere, and two estimators didn't satisfy it — GDPOAdvantageEstimator required repeated_batch and GeneralizedAdvantageEstimator required values, both positional and outside the contract. A protocol nothing is checked against is just a docstring that can drift.

Latest push makes it real:

The contract's parameters are keyword-only, which is how both loops call it. Each passes the union of what all estimators need and leans on **kwargs for the rest, so "may accept extra arguments, may not require them" is the actual invariant.
repeated_batch and values get keyword defaults with explicit errors. Both are still required in practice — the defaults only stop the signatures from demanding more than the shared contract. One side effect worth having: ppo.py adds values to the kwargs only if "values" in train_data, so GAE without a critic used to fail with a bare TypeError about a missing positional argument. It now says to enable the critic.
Both _create_advantage_estimator functions are annotated with the protocol.

On the test: runtime_checkable only checks that compute_advantage exists, which every estimator passed even with an incompatible signature — so an isinstance test would have been close to vacuous. It inspects the signature instead and asserts no estimator requires anything outside the contract. Removing either default turns it red.

yuki-97

hi @tianyi-zhang-02 , sorry for late reply and thanks for so many contributions! will assign reviewers.

@bg51717 to review

tianyi-zhang-02 · 2026-08-17T03:28:10Z

Small follow-up on my last push. The note explaining why values carries a default went into _compute_gae's Args block instead of compute_advantage's — and _compute_gae has neither the default nor the check the note describes (values: torch.Tensor is still required there, and the guard sits in compute_advantage above it, so "checked explicitly below" pointed the wrong way too).

Moved to compute_advantage, which turned out to have no Args block at all — which is how it ended up in the wrong place: the string I matched on was unique in the file, so it looked right.

Nothing about the behaviour changes; the docstring in a PR about documenting the estimator contract just shouldn't be the part that's wrong.

bg51717

Thanks for the refactor. Unifying the estimator return type and removing the last_metrics side channel makes the call sites clearer. I left a few small comments.

bg51717 · 2026-08-20T07:55:18Z

+            raise ValueError(
+                "GAE needs per-token value estimates; enable the critic so that "
+                "'values' is present in the training batch, or switch "
+                "grpo.adv_estimator.name to an estimator that needs no value model."


Should this be ppo.adv_estimator.name? GAE is constructed from PPO's adv_estimator config.

You are right — GAE is built in ppo.py::_create_advantage_estimator from master_config.ppo["adv_estimator"], so grpo.adv_estimator.name sends the reader to a key that does not control it. Fixed to ppo.adv_estimator.name.

bg51717 · 2026-08-20T08:00:35Z

+def _create_advantage_estimator(master_config: MasterConfig) -> AdvantageEstimator:
    """Create and return an advantage estimator based on configuration.

    PPO's training loop consumes a `(advantages, returns)` pair from a


This still describes the old tuple return contract. Could we update it ?

Fixed. It now says that both supported estimators return an AdvantageResult, that gae populates returns while raw_reward leaves it None, and that the loop writes train_data["returns"] only when it is present.

I got this wrong once on the way: my first rewrite claimed both populate returns. RawRewardAdvantageEstimator returns AdvantageResult(adv) with returns left None, which is exactly why the loop still guards with if result.returns is not None.

bg51717 · 2026-08-20T08:10:25Z

        rewards,
        mask,
-        repeated_batch,
+        repeated_batch=None,


@yuki-97, could you take a quick look at the API design here? The PR uses an AdvantageResult dataclass for the common return value and an AdvantageEstimator Protocol instead of a shared base class. To conform to this Protocol, GDPOAdvantageEstimator, which requires repeated_batch, now gives that argument a default value of None and checks at runtime whether a valid value was actually provided. I'm not sure whether this is a reasonable trade-off.

Fair question — the fake default is the part I like least too.

Here is why it is there. Both loops call compute_advantage with the union of what every estimator might need, and lean on **kwargs for the rest. So the real invariant is "may accept extra arguments, may not require any beyond prompt_ids/rewards/mask". repeated_batch and values were the only two breaking it. A protocol that two of six estimators fail is not enforcing much.

It did buy one concrete thing. ppo.py adds values only if "values" in train_data. So GAE without a critic was reachable, and it failed with a bare TypeError about a missing positional argument. Now it tells you to enable the critic. That is strictly better than the failure it replaced.

Three ways to go:

Leave it. Uniform signature, explicit error, docstring says it is required.

Drop the Protocol, keep AdvantageResult. Signatures stay honest, and the dataclass is the part carrying its weight anyway. This is what the PR was two pushes ago.

Pass one typed context object instead of loose kwargs. Then every estimator takes one argument honestly. I would want this long-term, but it touches every estimator and both loops, so it should be its own PR.

If the fake default is the sticking point, I would rather do 2 here and propose 3 separately. Your call :)

tianyi-zhang-02 · 2026-08-20T17:59:46Z

Also merged current main in, which surfaced a real break worth flagging: the async PPO loop added since this branch was cut still had the tuple branching this PR removes elsewhere —

if isinstance(result, tuple):
    advantages, returns = result
else:
    advantages, returns = result, None

An AdvantageResult is not a tuple, so it took the else branch and handed the dataclass itself to torch.masked_select. test_ppo_train_critic_warmup_reuses_generation_until_policy_update fails on the merge and passes on main, so this was mine to fix, not a pre-existing failure. Now matches the sync path directly above it.

That is the second time a new compute_advantage consumer has landed while this sat — worth keeping in mind for whenever it merges. tests/unit/algorithms/, tests/unit/single_controller/: 859 passed, 54 skipped.

tianyi-zhang-02 · 2026-08-24T05:01:39Z

@bg51717 @yuki-97 — new information that bears on the API question at advantage_estimator.py:165, which I left as three options and which is still the thing holding this up.

#3773 (feat(sc): support PPO in single controller) rewrites the same file from the other side. It replaces estimator_config: dict with a pydantic GAEConfig(BaseModel) and turns estimator_config["gae_lambda"] into estimator_config.gae_lambda, and it narrows AdvantageEstimatorConfig.name to a Literal.

That is option 3 from my earlier reply — "pass one typed object instead of loose lookups" — applied to the config side. So the direction question looks answered: typed objects, not dicts. Which makes the return side of this PR (AdvantageResult instead of a bare tensor or a (advantages, returns) tuple) the matching half rather than a competing style.

It also means the two will conflict. Both rewrite the estimator __init__s and the compute_advantage signatures. #3773 is the feature and this has been open since 2026-08-06, so I would rather not make you resolve that: happy to rebase this onto #3773 once it settles, or to close this and hand you the AdvantageResult change to fold into #3773 directly if that is less churn. Whichever you prefer.

On the repeated_batch=None default that prompted the question: if #3773 lands the typed-config direction, option 3 stops being a big separate change and the fake default can just go away — every estimator would take its own typed config plus the shared tensors, and nothing would need a runtime "was this actually passed" check. I am happy to do that as the follow-up rather than defend the default.

The estimators shared no declared interface: compute_advantage returned a bare tensor from GRPO/GDPO/Reinforce++/OPD but a (advantages, returns) tuple from RawReward/GAE, so ppo.py had to branch on isinstance(result, tuple), and OPD published its logging metrics by assigning self.last_metrics, which grpo.py read back through a hasattr() probe that only existed because the other estimators lack the attribute. Introduce AdvantageResult(advantages, returns=None, metrics={}) as the single return type and an AdvantageEstimator Protocol describing the contract. Migrate all six estimators and the five call sites; the tuple branch in ppo.py and the hasattr probe in grpo.py both disappear. Declared as a Protocol rather than a base class so estimators defined outside nemo_rl satisfy it structurally, matching how dataset classes are already pluggable by dotted path. Behavior is unchanged: only the wrapper moves. GAE's private _compute_gae still returns its tuple, since three call sites inside the estimator destructure it. Addresses NVIDIA-NeMo#2909 Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>

Removing the isinstance(result, tuple) branch in ppo.py means every compute_advantage double has to return an AdvantageResult too. Five were missed, four of them in tests added to main after this branch was cut: - tests/unit/algorithms/test_ppo.py DummyAdvantageEstimator returned a bare tuple, so test_ppo_train_excludes_overlong_samples_from_advantage and test_ppo_train_wires_logprob_mask_to_advantage_training_and_metrics failed with AttributeError: tuple object has no attribute advantages. - tests/unit/algorithms/test_grpo.py had three MagicMock return_values set to bare tensors. - tests/unit/single_controller/test_train_pump.py _FakeAdvEstimator returned a bare tensor; that file was not touched by this branch at all. Verified by running test_ppo.py, test_grpo.py and test_advantage_estimator.py against origin/main and against this branch: identical results, no failures introduced. Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>

Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>

The AdvantageResult docstring said 'returns' is populated by GAE and the raw-reward estimator. Only GAE populates it -- RawRewardAdvantageEstimator returns AdvantageResult(adv) with returns left None, three lines away in the same file. The six public compute_advantage docstrings also still described the old return types (a bare tensor, or a tuple). Ruff's pydocstyle rules check that a Returns section exists, never that it is accurate, so none of this was caught. Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>

The protocol was defined but referenced nowhere, and two estimators did not satisfy it: GDPO required ``repeated_batch`` and GAE required ``values``, both as positional parameters outside the contract. A protocol nothing is checked against is documentation that can go stale silently. - Declare the contract's parameters keyword-only, which is how both loops actually call it: each passes the union of what all estimators need and relies on ``**kwargs`` to absorb the rest. - Give GDPO's ``repeated_batch`` and GAE's ``values`` keyword defaults with explicit errors. Both are still required in practice; the defaults only stop the signatures from demanding more than the shared contract. GAE previously raised a bare TypeError when ``values`` was absent from the batch (ppo.py adds it conditionally), which now says to enable the critic instead. - Annotate both ``_create_advantage_estimator`` functions with the protocol. ``runtime_checkable`` only checks that ``compute_advantage`` exists, which every estimator passed even with an incompatible signature, so the test inspects the signature: no estimator may require anything outside the contract. It fails if either default is removed. Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>

The note explaining why ``values`` carries a default landed in ``_compute_gae``'s Args block, which has neither the default nor the check it describes -- ``values: torch.Tensor`` there is still required, and the guard lives in ``compute_advantage`` above it, so "checked explicitly below" pointed the wrong way as well as at the wrong function. Moves it to ``compute_advantage``, which had no Args block at all, and matches the wording to the equivalent note on GDPO's ``repeated_batch``. Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>

…rn contract Both from review. The GAE guard told the reader to change grpo.adv_estimator.name, but GAE is built by ppo.py's _create_advantage_estimator from master_config.ppo['adv_estimator']. That function's docstring also still described the pre-AdvantageResult tuple return. Restates it as it now works: gae populates returns, raw_reward leaves it None, and the loop writes train_data['returns'] only when present. Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>

The async PPO loop that landed on main since this branch was cut still had the tuple branching this PR removes elsewhere: if isinstance(result, tuple): advantages, returns = result else: advantages, returns = result, None An AdvantageResult is not a tuple, so it took the else branch and passed the dataclass itself to torch.masked_select -- test_ppo_train_critic_warmup_reuses_generation_until_policy_update fails with "argument 'input' must be Tensor, not AdvantageResult". Matches the sync path above it. Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>

…kens guard The rebase onto main left the new AdvantageResult call sitting *after* main's guarded block instead of replacing it, so the estimator ran twice and ran even when the sequence mask had removed the whole chunk -- which is exactly what test_advantage_stage_skips_estimator_when_seq_mask_removes_whole_chunk asserts must not happen. Also migrates the four estimator doubles main added since this PR opened (NVIDIA-NeMo#3768's OPD ones plus the GAE-like tuple) onto AdvantageResult, which is the point of the contract: a double that still returns a bare tensor or a tuple is modelling an interface that no longer exists. Signed-off-by: Tianyi Zhang <zhangtianyi975@gmail.com> Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>

Two gaps that only show once this branch is merged with anything else: - nine controller stubs predate _teacher_logprobs_required, which main's NVIDIA-NeMo#3768 made _advantage_stage read unconditionally; - the estimator double returns a bare tensor, which NVIDIA-NeMo#3512 replaces with AdvantageResult. The double now resolves AdvantageResult through the module rather than importing it, so it returns whichever shape the checkout actually has and this file does not depend on a name that exists only on the other branch. Signed-off-by: Tianyi Zhang <zhangtianyi975@gmail.com> Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>

…t has NVIDIA-NeMo#3512 replaces the bare-tensor return with AdvantageResult. The double now resolves that class through the module instead of importing it, so it returns a tensor on this branch and an AdvantageResult once NVIDIA-NeMo#3512 lands, and this file does not depend on a name that exists only there. An earlier version of this commit also added _teacher_logprobs_required to nine controller stubs. main already sets it in seven places and the suite is green without the other two, so that was duplicate work that collided with NVIDIA-NeMo#3512 on the same lines -- dropped. Signed-off-by: Tianyi Zhang <zhangtianyi975@gmail.com> Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>

Two gaps in this PR's own parity claim. grpo.py logs advantages at :3471 and clips at :3487; grpo_sync.py at :903 then :915. Both report pre-clip. This PR clipped first, so SC's advantages/mean|max|min would have meant something no other driver's does. Clip after the log; the data-plane write is unchanged. The clip guard's 'not self._is_ppo' half was untested, because the PPO test handed the ppo block a GRPO-shaped stub carrying advantage_clip_low -- a field real PPOConfig does not declare. Deleting the clause left every test green. PPOConfig is extra="allow", so the clause is exactly what stops a user-set ppo.advantage_clip_low from running GRPO clipping on a PPO run; there is a test for that now, and the PPO stub matches the real config. Also drop the importlib lookup for AdvantageResult: it does not exist on main, NVIDIA-NeMo#3512 is unmerged, and if that branch ever returned the wrapper the production code at single_controller.py:2028 would break anyway. The guide said setup rejects reward_scaling. It does not any more. Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>

Signed-off-by: Tianyi Zhang <zhangtianyi975@gmail.com> # Conflicts: # tests/unit/single_controller/test_single_controller_actor.py

Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>

Retain the unified AdvantageResult while removing the Protocol that forced heterogeneous estimators into an artificial three-argument signature. Restore GDPO's repeated_batch and GAE's values as honest required inputs, drop the now-vacuous signature tests, and update the data-plane example for the result wrapper. Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>

tianyi-zhang-02 · 2026-08-30T19:55:23Z

Follow-up on the API concern from your review, on final SHA 6eb502a06122c27eb4763645a40a5175467f4427:

took option 2 from the thread: removed the unused Protocol and kept AdvantageResult as the common return contract;
removed the fake None defaults, so GDPO's repeated_batch and GAE's values are required again;
left feat(sc): support PPO in single controller #3773's typed estimator configs intact and migrated the GRPO/PPO/SingleController callers, including OPD metrics;
Runpod validation: 14 targeted tests and 296 related CPU tests pass, 7 GPU-only tests skip, and Ruff check/format are clean.

cc @bg51717 — this directly resolves the trade-off you flagged without keeping a runtime-only signature contract.

tianyi-zhang-02 requested review from a team as code owners August 6, 2026 05:00

nemo-automation-bot Bot added the community-request label Aug 6, 2026

tianyi-zhang-02 mentioned this pull request Aug 6, 2026

Formalize the advantage-estimator interface (shared Protocol + explicit metrics, unified return type) #2909

Open

tianyi-zhang-02 force-pushed the refactor/advantage-estimator-protocol branch from 7cae64c to ab220c1 Compare August 7, 2026 06:40

tianyi-zhang-02 requested a review from a team as a code owner August 8, 2026 11:38

svcnvidia-nemo-ci added the waiting-on-maintainers Waiting on maintainers to respond label Aug 9, 2026

svcnvidia-nemo-ci added waiting-on-maintainers Waiting on maintainers to respond and removed waiting-on-maintainers Waiting on maintainers to respond labels Aug 15, 2026

yuki-97 reviewed Aug 17, 2026

View reviewed changes

yuki-97 requested a review from bg51717 August 17, 2026 03:03

svcnvidia-nemo-ci added waiting-on-maintainers Waiting on maintainers to respond and removed waiting-on-maintainers Waiting on maintainers to respond labels Aug 17, 2026

tianyi-zhang-02 force-pushed the refactor/advantage-estimator-protocol branch from 3442920 to 78b8224 Compare August 19, 2026 19:51

bg51717 reviewed Aug 20, 2026

View reviewed changes

svcnvidia-nemo-ci added waiting-on-customer Waiting on the original author to respond and removed waiting-on-maintainers Waiting on maintainers to respond labels Aug 20, 2026

svcnvidia-nemo-ci added waiting-on-maintainers Waiting on maintainers to respond and removed waiting-on-customer Waiting on the original author to respond labels Aug 20, 2026

tianyi-zhang-02 mentioned this pull request Aug 24, 2026

feat(sc): support PPO in single controller #3773

Merged

4 tasks

svcnvidia-nemo-ci added waiting-on-maintainers Waiting on maintainers to respond and removed waiting-on-maintainers Waiting on maintainers to respond labels Aug 24, 2026

tianyi-zhang-02 mentioned this pull request Aug 26, 2026

fix(algorithms): honour the configured KL clamps in the reward-side KL #3853

Open

tianyi-zhang-02 added 8 commits August 26, 2026 22:23

chore(tests): ignore generated unit result files

e5d4d40

Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>

tianyi-zhang-02 force-pushed the refactor/advantage-estimator-protocol branch from 6fa1956 to dd2ebe9 Compare August 27, 2026 02:31

tianyi-zhang-02 and others added 3 commits August 29, 2026 02:18

Merge remote-tracking branch 'upstream/main' into fix3512

c01c733

Signed-off-by: Tianyi Zhang <zhangtianyi975@gmail.com> # Conflicts: # tests/unit/single_controller/test_single_controller_actor.py

Merge upstream/main into w3512-advresult-clean

c1ac535

Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>

tianyi-zhang-02 changed the title ~~refactor(algorithms): formalize the advantage-estimator contract~~ refactor(algorithms): standardize advantage-estimator results Aug 30, 2026

svcnvidia-nemo-ci added waiting-on-maintainers Waiting on maintainers to respond and removed waiting-on-maintainers Waiting on maintainers to respond labels Aug 30, 2026

Conversation

tianyi-zhang-02 commented Aug 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Validation

Uh oh!

copy-pr-bot Bot commented Aug 6, 2026

Uh oh!

tianyi-zhang-02 commented Aug 7, 2026

Uh oh!

tianyi-zhang-02 commented Aug 14, 2026

Uh oh!

yuki-97 left a comment

Choose a reason for hiding this comment

Uh oh!

tianyi-zhang-02 commented Aug 17, 2026

Uh oh!

bg51717 left a comment

Choose a reason for hiding this comment

Uh oh!

bg51717 Aug 20, 2026

Choose a reason for hiding this comment

Uh oh!

tianyi-zhang-02 Aug 20, 2026

Choose a reason for hiding this comment

Uh oh!

bg51717 Aug 20, 2026

Choose a reason for hiding this comment

Uh oh!

tianyi-zhang-02 Aug 20, 2026

Choose a reason for hiding this comment

Uh oh!

bg51717 Aug 20, 2026

Choose a reason for hiding this comment

Uh oh!

tianyi-zhang-02 Aug 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

tianyi-zhang-02 commented Aug 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

tianyi-zhang-02 commented Aug 24, 2026

Uh oh!

tianyi-zhang-02 commented Aug 30, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

tianyi-zhang-02 commented Aug 6, 2026 •

edited

Loading

tianyi-zhang-02 Aug 20, 2026 •

edited

Loading

tianyi-zhang-02 commented Aug 20, 2026 •

edited

Loading