agentscope-ai / agentscope-ai/Trinity-RFT
[FEATURE] Add M2PO policy loss for stale-rollout RFT
- 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.