[RFC] Upstream Environment-Regularized Policy Optimization (ERPO) into ROLL
- Dominant language
- Python
- Stars
- 3.4k
- Forks
- 312
- Avg merge
- 1h 2m
- Merged PRs (30d)
- 2
Description
## Background
We would like to propose upstream support for **Environment-Regularized Policy Optimization (ERPO)** in ROLL.
ERPO addresses the stability–exploration dilemma in LLM policy optimization by moving regularization from the response side to the query side. It introduces:
1. **Query-KL (QKL):** regularizes the query distribution induced by the current policy against the pre-RL reference model while leaving the response score function unconstrained.
2. **Reference-derived per-query weighting:** assigns a dataset-static weight to each query according to its likelihood under the reference model, reducing estimator variance and improving robustness.
ERPO is estimator-agnostic and can be integrated into GRPO-, PPO-, and REINFORCE-style training pipelines without additional model forward passes.
On six mathematical reasoning benchmarks with Qwen2.5-Math-7B, ERPO improves the average Avg@32 score from 0.274 to 0.336 while maintaining stable long-horizon training. Detailed results, configurations, and training dynamics are available in the ERPO repository and paper.
Paper: https://arxiv.org/abs/2608.23311
Repository: https://github.com/AlibabaResearch/ERPO
## Existing Implementation
We have released an incremental ERPO implementation based on ROLL. The current implementation includes:
* ERPO computation in `roll/pipeline/rlvr/actor_worker.py`
* The `prepare_backward_batch` integration hook
* Dynamic reference-derived per-query loss weighting
* Query-side KL through `kl_loss_mask_mode: prompt`
* The `dynamic_prompt_logp_loss_weight` configuration option
* ERPO and GRPO controlled-comparison recipes under `examples/erpo/`
* Training and evaluation data with checksum manifests
The current repository provides an incremental patch and still requires users to merge the files manually into a full ROLL installation.
## Goals
We propose integrating ERPO into the upstream ROLL RLVR pipeline with the following goals:
* Provide ERPO as an optional, configuration-controlled algorithm.
* Reuse the existing reference-model forward computation without introducing additional forward passes.
* Support distributed query-weight computation through the existing ROLL parallel infrastructure.
* Provide a reproducible ERPO example configuration.
* Add documentation and correctness tests for Query-KL and per-query weighting.
## Non-Goals
This proposal does not aim to:
* Change the default RLVR algorithm or configuration.
* Introduce ERPO-specific dependencies.
* Include experiment checkpoints or training logs in the ROLL repository.
* Modify unrelated checkpointing, evaluation-frequency, or asynchronous-upload behavior.
* Require users who do not enable ERPO to change their existing recipes.
## Proposed Integration
The proposed upstream changes include:
1. Add the ERPO-related configuration field:
```yaml
dynamic_prompt_logp_loss_weight: true
```
2. Support query-side KL through:
```yaml
use_kl_loss: true
kl_loss_coef: 1.0e-2
kl_loss_mask_mode: prompt
```
3. Integrate the reference-derived per-query weight into the actor backward batch preparation.
4. Keep all ERPO behavior disabled by default so that existing training configurations remain unchanged.
5. Add an ERPO example recipe and a matched GRPO baseline recipe.
6. Add unit tests for Query-KL, per-query weighting, and backward compatibility.
## Request for Feedback
Would the ROLL be open to upstreaming ERPO into the RLVR pipeline? We would also appreciate guidance on whether ERPO should be exposed as a named algorithm or as a combination of reusable configuration options.
Once the preferred integration design is confirmed, we can prepare a focused PR against the latest ROLL main branch.
Contributor guide
No contributing guide indexed for this repository
Research direction
Start by reviewing the existing ERPO computation and prepare_backward_batch integration in roll/pipeline/rlvr/actor_worker.py, then inspect the ERPO and GRPO recipes under examples/erpo/. Confirm how the configuration options preserve default behavior and reuse the existing reference-model forward pass. Done means an optional upstream integration with reproducible recipes, documentation, and correctness and backward-compatibility tests.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- machine-learning
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 45/100