feat(algorithm): add M2PO for stale-rollout RFT - #616
Open
primorLee wants to merge 1 commit into
Open
Conversation
Implement final-paper second-moment token masking for stale rollout data, register the algorithm preset, and add tests, documentation, and an asynchronous GSM8K example.\n\nRefs agentscope-ai#615
1 task
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Related to #470 and follows the design proposal in #615.
This PR adds M2PO (Second-Moment Trust Policy Optimization) for stable reinforcement fine-tuning with stale rollout data. It implements Algorithm 1 from the final ICLR 2026 paper using Trinity-RFT's policy-loss interface.
M2PO identifies tokens in PPO's active trust-region quadrants, then removes the largest squared log importance-ratio outliers until the remaining second moment is at most
tau. Following the paper, masked losses retain the original valid-token denominator and the algorithm does not require a critic, reference model, or KL term.Changes
M2POPolicyLossFnwith the paper defaultm2_threshold=0.04;policy_loss_fn: m2poand the completealgorithm_type: m2popreset;Paper fidelity
The final paper revision specifies token masking. The authors' earlier public reference code uses adaptive clipping, so this implementation intentionally follows the final published Algorithm 1 and documents that choice.
Validation
16 passed, 73 subtests passedfor the focused policy-loss and registry suites;Config.check_and_update()validation;git diff --checkpasses.The local validation machine has no CUDA GPU, so this PR does not claim a full training-curve reproduction. The included GSM8K configuration is an integration example rather than an exact reproduction of the paper's
s=256setup.Checklist