Formalize the advantage-estimator interface (shared Protocol + explicit metrics, unified return type)
- 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
Assessment
This issue has not been assessed yet.