Skip to content

feat(algorithm): add M2PO for stale-rollout RFT - #616

Open
primorLee wants to merge 1 commit into
agentscope-ai:mainfrom
primorLee:feat/m2po-policy-loss
Open

feat(algorithm): add M2PO for stale-rollout RFT#616
primorLee wants to merge 1 commit into
agentscope-ai:mainfrom
primorLee:feat/m2po-policy-loss

Conversation

@primorLee

Copy link
Copy Markdown

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

  • add M2POPolicyLossFn with the paper default m2_threshold=0.04;
  • register policy_loss_fn: m2po and the complete algorithm_type: m2po preset;
  • add numerical diagnostics for M2 before/after masking, masked fraction, and trust-region fraction;
  • test the vectorized mask against a literal Algorithm 1 loop, plus gradient, denominator, registration, empty-mask, and numerical edge cases;
  • add an asynchronous Qwen2.5-1.5B GSM8K example;
  • document M2PO in the English/Chinese algorithm tables and dataset-oriented example guide.

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 passed for the focused policy-loss and registry suites;
  • trainer and explorer YAML files pass full Config.check_and_update() validation;
  • targeted Black, isort, flake8, and mypy checks pass;
  • git diff --check passes.

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=256 setup.

Checklist

  • Code has passed all tests (focused tests pass; full GPU/integration CI is pending)
  • Docstrings have been added/updated in Google Style
  • Documentation has been updated
  • Code is ready for review

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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant