agentscope-ai / agentscope-ai/Trinity-RFT

[FEATURE] Add M2PO policy loss for stale-rollout RFT

Abierto
#615 1 comentario 0 reacciones 0 asignados Ver en GitHub
Lenguaje dominante
Python
Estrellas
701
Forks
79
Merge medio
8 h 7 min
PR fusionados (30 d)
1

Descripción

### Feature Request

Add M2PO ("Prosperity before Collapse: How Far Can Off-Policy RL Reach with Stale Data on LLMs?", ICLR 2026) as a policy-loss option for stale-rollout reinforcement fine-tuning.

This proposal follows the M2PO item listed in #470. Paper: https://openreview.net/forum?id=IIgl5MWelz

### Purpose of the Feature

Asynchronous rollout/training can reuse increasingly stale samples. M2PO constrains the second moment of the token-level log importance ratio, masking only harmful trust-region tokens until the remaining mean squared log-ratio is below a threshold. This provides a model-side algorithm for stale-data stability without adding a critic or reference-model KL term.

### Expected Implementation

I propose to:

- add an `m2po` policy loss implementing Algorithm 1 from the final paper;
- constrain only active PPO trust-region quadrants: positive advantage with ratio above 1, or negative advantage with ratio below 1;
- use the paper default threshold `tau = 0.04`;
- preserve the original valid-token denominator after masking;
- register an M2PO algorithm preset without a critic, reference model, or KL penalty;
- add numerical/edge-case unit tests and an asynchronous GSM8K example.

A local prototype is ready. Its vectorized mask is tested against a literal implementation of Algorithm 1, and the relevant unit/registry/config checks pass. I have not yet run a full multi-GPU training-curve reproduction.

Two integration questions before opening the PR:

1. Since #470 explicitly requests M2PO, would maintainers prefer this to graduate directly into `trinity/algorithm/`, or first land under `trinity/plugins/`?
2. For the asynchronous example, is `bypass_old_logprobs: true` the preferred way to retain rollout-time behavior log-probabilities for M2PO?

### Additional Information

The final paper revision specifies token masking, while the authors' earlier public reference code uses adaptive clipping. The proposed implementation follows the final published Algorithm 1 and documents that choice.

### Are You Willing to Submit a PR?

- [x] Yes, I am willing to submit a PR!

Guía de contribución

Abrir la guía de contribución

Evaluación

Este issue todavía no se ha evaluado.

Recibe los nuevos issues en tu correo

Un resumen breve de issues de GitHub para principiantes.