Add Step-wise GRPO Advantage #153

pan-x-c · 2025-07-30T12:13:10Z

Description

As the title says

Checklist

Please check the following items before code is ready to be reviewed.

Code has passed all tests
Docstrings have been added/updated in Google Style
Documentation has been updated
Code is ready for review

gemini-code-assist

Summary of Changes

Hello @pan-x-c, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces a novel approach to advantage calculation and propagation within the trinity framework. It provides a new strategy that derives a single advantage score from the final step of a trajectory and then broadcasts this score across all steps of that trajectory, aiming to improve how reward signals are utilized for policy updates, particularly in grouped reinforcement learning settings.

Highlights

New Advantage Strategy: I've introduced a new StepWiseGRPOStrategy class, which is a specialized AddStrategy designed to calculate and apply advantages in a step-wise manner, inspired by the rLLM framework.
Group Advantage Calculation: The strategy includes a calculate_group_advantage method that computes a normalized score for each run based on the reward of its last experience relative to the group's mean and standard deviation. This score represents the overall advantage for that run.
Advantage Broadcasting: A broadcast_advantages method is implemented to take the single group advantage score calculated for a run and apply it to all experiences within that run's trajectory. It also includes an optional enable_step_norm feature to normalize advantages by the trajectory length.
Registry Integration: The new StepWiseGRPOStrategy is registered under the name 'step_wise_grpo' within the ADD_STRATEGY registry, making it discoverable and configurable for use in the system.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	`/gemini review`	Performs a code review for the current pull request in its current state.
Pull Request Summary	`/gemini summary`	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	`/gemini help`	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments or fill out our survey to provide feedback.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

gemini-code-assist

Code Review

This pull request introduces a new StepWiseGRPOStrategy for calculating advantages by broadcasting scores from the last step to all previous steps. The implementation is a good starting point, but I've found a few issues that need attention. There's a critical issue with potential None rewards that could cause a crash, a type mismatch when storing scores, and some leftover debugging print statements. I've also suggested minor improvements to code clarity. Addressing these points will make the new strategy more robust and maintainable.

trinity/algorithm/add_strategy/step_wise_add_strategy.py

trinity/algorithm/add_strategy/__init__.py

pan-x-c · 2025-07-31T12:48:21Z

/unittest-module-algorithm

github-actions · 2025-07-31T12:49:04Z

Summary

Tests 📝	Passed ✅	Failed ❌	Skipped ⏭️	Other ❓	Flaky 🍂	Duration ⏱️
8	8	0	0	0	0	3ms

Tests

Test Name	Status	Duration
tests/algorithm/add_strategy_test.py::TestAddStrategy::test_grpo_args	✅	1ms
tests/algorithm/add_strategy_test.py::TestAddStrategy::test_reward_variance_strategy	✅	1ms
tests/algorithm/add_strategy_test.py::TestAddStrategy::test_step_wise_grpo_strategy	✅	1ms
tests/algorithm/policy_loss_test.py::VerlPolicyLossTest::test_dpo_policy_loss	✅	1ms
tests/algorithm/policy_loss_test.py::VerlPolicyLossTest::test_mix_policy_loss	✅	1ms
tests/algorithm/policy_loss_test.py::VerlPolicyLossTest::test_opmd_policy_loss	✅	1ms
tests/algorithm/policy_loss_test.py::VerlPolicyLossTest::test_ppo_policy_loss	✅	1ms
tests/algorithm/policy_loss_test.py::VerlPolicyLossTest::test_sft_policy_loss	✅	1ms

Github Test Reporter by CTRF 💚

pan-x-c · 2025-07-31T12:50:22Z

/unittest-module-explorer

pan-x-c · 2025-07-31T12:58:29Z

/unittest-module-explorer

github-actions · 2025-07-31T13:04:32Z

Summary

Tests 📝	Passed ✅	Failed ❌	Skipped ⏭️	Other ❓	Flaky 🍂	Duration ⏱️
19	19	0	0	0	0	324ms

Tests

Test Name	Status	Duration
tests/explorer/explorer_test.py::BaseExplorerCase::test_explorer	✅	2ms
tests/explorer/explorer_test.py::TestExplorerCountdownEval::test_explorer	✅	94ms
tests/explorer/explorer_test.py::TestExplorerCountdownNoEval::test_explorer	✅	84ms
tests/explorer/explorer_test.py::TestExplorerWithAddStrategy::test_explorer	✅	59ms
tests/explorer/scheduler_test.py::SchedulerTest::test_concurrent_operations	✅	4ms
tests/explorer/scheduler_test.py::SchedulerTest::test_get_results	✅	19ms
tests/explorer/scheduler_test.py::SchedulerTest::test_multi_step_execution	✅	4ms
tests/explorer/scheduler_test.py::SchedulerTest::test_scheduler_all_methods	✅	14ms
tests/explorer/scheduler_test.py::SchedulerTest::test_scheduler_restart_after_stop	✅	8ms
tests/explorer/scheduler_test.py::SchedulerTest::test_split_tasks	✅	7ms
tests/explorer/scheduler_test.py::SchedulerTest::test_wait_all	✅	7ms
tests/explorer/scheduler_test.py::SchedulerTest::test_wait_all_timeout_with_multi_batch	✅	13ms
tests/explorer/workflow_test.py::WorkflowTest::test_gsm8k_workflow	✅	1ms
tests/explorer/workflow_test.py::WorkflowTest::test_math_boxed_workflow	✅	1ms
tests/explorer/workflow_test.py::WorkflowTest::test_math_complex_workflow	✅	1ms
tests/explorer/workflow_test.py::WorkflowTest::test_math_fraction_workflow	✅	1ms
tests/explorer/workflow_test.py::WorkflowTest::test_math_workflow	✅	1ms
tests/explorer/workflow_test.py::WorkflowTest::test_rm_gallery_workflow	✅	1ms
tests/explorer/workflow_test.py::WorkflowTest::test_workflow_resettable	✅	1ms

Github Test Reporter by CTRF 💚

pan-x-c added 2 commits July 30, 2025 14:30

fix test reporter

6def24b

add step_wise_grpo

014ac8d

gemini-code-assist bot reviewed Jul 30, 2025

View reviewed changes

hiyuchang reviewed Jul 31, 2025

View reviewed changes

trinity/algorithm/add_strategy/__init__.py Show resolved Hide resolved

pan-x-c added 5 commits July 31, 2025 13:51

fix comments

3d3198e

fix torch dtype

c8cd1c5

add tests for add_strategy

fe7d212

add more tests

290ac3e

fix pre-commit

c09f92c

Merge branch 'main' into feature/adv_broadcast

afbcf40

hiyuchang approved these changes Jul 31, 2025

View reviewed changes

pan-x-c merged commit e10b5d2 into modelscope:main Jul 31, 2025
1 check passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Add Step-wise GRPO Advantage #153

Add Step-wise GRPO Advantage #153

Uh oh!

pan-x-c commented Jul 30, 2025 •

edited

Loading

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

pan-x-c commented Jul 31, 2025

Uh oh!

github-actions bot commented Jul 31, 2025

Uh oh!

pan-x-c commented Jul 31, 2025

Uh oh!

pan-x-c commented Jul 31, 2025

Uh oh!

github-actions bot commented Jul 31, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Add Step-wise GRPO Advantage #153

Add Step-wise GRPO Advantage #153

Uh oh!

Conversation

pan-x-c commented Jul 30, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Checklist

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

pan-x-c commented Jul 31, 2025

Uh oh!

github-actions bot commented Jul 31, 2025

Summary

Tests

Uh oh!

pan-x-c commented Jul 31, 2025

Uh oh!

pan-x-c commented Jul 31, 2025

Uh oh!

github-actions bot commented Jul 31, 2025

Summary

Tests

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

pan-x-c commented Jul 30, 2025 •

edited

Loading