NVIDIA-NeMo / NVIDIA-NeMo/RL

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

Open
#2,909 1 comment 0 reactions 1 assignee Claimed by @terrykong View on GitHub
enhancement Feature
Dominant language
Python
Stars
2k
Forks
561
Avg merge
4d 5h
Merged PRs (30d)
145

Description

**Is your feature request related to a problem? Please describe.**

The advantage estimators in `nemo_rl/algorithms/advantage_estimator.py` share no base class or `Protocol`, so the contract is entirely informal and fragile:

1. **Inconsistent return types.** `compute_advantage` returns a bare tensor in `GRPOAdvantageEstimator`, `GDPOAdvantageEstimator`, `ReinforcePlusPlusAdvantageEstimator`, and `OPDAdvantageEstimator`, but a `(advantages, returns)` tuple in `RawRewardAdvantageEstimator` and `GeneralizedAdvantageEstimator`. Callers must "know" which is which.
2. **Metrics via side-channel.** `OPDAdvantageEstimator` emits logging metrics as a side effect into `self.last_metrics` (`advantage_estimator.py:537,586`) rather than returning them, and the caller peeks at the attribute out-of-band, guarded by `hasattr` only because the other estimators don't define it (`grpo.py:3852-3856`):
```python
if hasattr(adv_estimator, "last_metrics") and adv_estimator.last_metrics:
rollout_metrics.update(adv_estimator.last_metrics)
```

Nothing enforces the attribute or return shapes, so adding an estimator means matching undocumented conventions.

**Describe the solution you'd like**

- A shared `AdvantageEstimator` base class / `Protocol` with an explicit metrics hook, e.g. `get_metrics() -> dict[str, float]` defaulting to `{}`. The call site becomes `rollout_metrics.update(adv_estimator.get_metrics())`.
- A unified return type, e.g. an `AdvantageResult(advantages, returns, metrics)` dataclass (`returns=None` where unused), removing the bare-tensor-vs-tuple split.

Touches all estimators in `advantage_estimator.py` plus the `compute_advantage` call sites in `grpo.py`.

**Additional context**

Follow-up from the MOPD PR #2780 review.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.