[RFC] Upstream Environment-Regularized Policy Optimization (ERPO) into ROLL

[adtureven]

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:

dynamic_prompt_logp_loss_weight: true

2. Support query-side KL through:

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.