agentscope-ai / agentscope-ai/Trinity-RFT

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

Aperta
#615 1 commento 0 reazioni 0 assegnatari Vedi su GitHub
Lingua principale
Python
Stelle
701
Fork
79
Merge medio
8h 7m
PR unite (30g)
1

Descrizione

### 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!

Guida per i contributori

Apri la guida per i contributori

Valutazione

Questa issue non è ancora stata valutata.

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.