alibaba / alibaba/ROLL

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

Open
#490 0 comments 0 reactions 0 assignees View on GitHub
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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.