feat: implemented reward redistribution and implicit reward redistribution#6
Open
maharajamihir wants to merge 3 commits into
Open
feat: implemented reward redistribution and implicit reward redistribution#6maharajamihir wants to merge 3 commits into
maharajamihir wants to merge 3 commits into
Conversation
There was a problem hiding this comment.
Pull Request Overview
This PR introduces reward redistribution logic (both explicit and implicit) into different components of the training pipeline.
- In the reward manager, explicit reward redistribution is applied by splitting the reward among tokens.
- In the PPO trainer and core algorithms, implicit redistribution is implemented via scaling advantages and configuration updates.
- The configuration file now includes new flags to toggle these reward redistribution approaches.
Reviewed Changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated no comments.
| File | Description |
|---|---|
| verl/workers/reward_manager/naive.py | Applies explicit reward redistribution to token scores. |
| verl/trainer/ppo/ray_trainer.py | Incorporates implicit reward redistribution and adds an assert to enforce mutual exclusivity. |
| verl/trainer/ppo/core_algos.py | Implements advantage scaling when implicit redistribution is enabled. |
| verl/trainer/config/ppo_trainer.yaml | Introduces new configuration flags for both redistribution methods. |
Comments suppressed due to low confidence (2)
verl/trainer/ppo/ray_trainer.py:898
- Using an assert for configuration validation may be bypassed in optimized runs; consider replacing it with an explicit check that raises an appropriate exception to ensure reliable error handling in production.
assert not ( self.config.algorithm.redistribute_reward and self.config.algorithm.redistribute_reward_implicit ), ( "Only one of `redistribute_reward` and `redistribute_reward_implicit` can be True" )
verl/trainer/ppo/ray_trainer.py:570
- [nitpick] The assignment of redistribution flags to meta_info is repeated in the test loop; consider abstracting this logic into a helper function to reduce duplication and improve maintainability.
test_batch.meta_info["redistribute_reward"] = ( self.config.algorithm.redistribute_reward )
emergenz
reviewed
May 14, 2025
emergenz
reviewed
May 14, 2025
emergenz
reviewed
May 14, 2025
| reward = score | ||
|
|
||
| reward_tensor[i, valid_response_length - 1] = reward | ||
| if data_item.meta_info["redistribute_reward"]: |
Collaborator
There was a problem hiding this comment.
Do we exclusively use the naive reward manager?
Collaborator
There was a problem hiding this comment.
Also implement RR for the batch reward manager
emergenz
requested changes
May 14, 2025
emergenz
left a comment
Collaborator
There was a problem hiding this comment.
Implicit reward redistribution currently contains a bug
emergenz
reviewed
May 16, 2025
| # Scale the advantages by (T-t)/T | ||
| if redistribute_reward_implicit: | ||
| # Compute per‐example length T from the response mask | ||
| lengths = response_mask.sum(dim=1, keepdim=True).to(advantages) |
emergenz
reviewed
May 16, 2025
| # Scale the advantages by (T-t)/T | ||
| if redistribute_reward_implicit: | ||
| # Compute per‐example length T from the response mask | ||
| lengths = response_mask.sum(dim=1, keepdim=True).to(advantages) |
Collaborator
There was a problem hiding this comment.
Suggested change
| lengths = response_mask.sum(dim=1, keepdim=True).to(advantages) | |
| lengths = response_mask.sum(dim=1, keepdim=True) |
emergenz
reviewed
May 16, 2025
| # Compute per‐example length T from the response mask | ||
| lengths = response_mask.sum(dim=1, keepdim=True).to(advantages) | ||
| # Build a [1, T_max] tensor of time indices | ||
| time_idx = torch.arange(gen_len, device=advantages.device, dtype=advantages.dtype).unsqueeze(0) |
Collaborator
There was a problem hiding this comment.
Suggested change
| time_idx = torch.arange(gen_len, device=advantages.device, dtype=advantages.dtype).unsqueeze(0) | |
| time_idx = torch.arange(gen_len).unsqueeze(0) |
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.
No description provided.