agentscope-ai / agentscope-ai/Trinity-RFT

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

Đang mở
#615 1 bình luận 0 reaction 0 người được giao Xem trên GitHub
Ngôn ngữ chính
Python
Star
701
Fork
79
Merge trung bình
8 giờ 7 phút
Pull request đã merge (30 ngày)
1

Mô tả

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

Hướng dẫn đóng góp

Mở hướng dẫn đóng góp

Đánh giá

Issue này chưa được đánh giá.

Nhận issue mới trong hộp thư của bạn

Bản tóm tắt ngắn những issue GitHub phù hợp với người mới.