Slice std_rewards

albertvillanova · albertvillanova · commit 4266550534ac · 2025-10-30T19:50:51.000+01:00
diff --git a/trl/experimental/grpo_with_replay_buffer/grpo_with_replay_buffer_trainer.py b/trl/experimental/grpo_with_replay_buffer/grpo_with_replay_buffer_trainer.py
@@ -270,6 +270,7 @@ def _generate_and_score_completions(
         )
         all_process_advantages = advantages.clone()  # keep the aggregated advantages for logging
         advantages = advantages[process_slice]
+        std_rewards = std_rewards[process_slice]
 
         # Calculate mean reward per function, but only for samples where the function was applied (non-NaN values)
         for i, reward_func_name in enumerate(self.reward_func_names):

Original file line number	Diff line number	Diff line change
`@@ -270,6 +270,7 @@ def _generate_and_score_completions(`
`270`	`270`	`)`
`271`	`271`	`all_process_advantages = advantages.clone() # keep the aggregated advantages for logging`
`272`	`272`	`advantages = advantages[process_slice]`
	`273`	`+ std_rewards = std_rewards[process_slice]`
`273`	`274`
`274`	`275`	`# Calculate mean reward per function, but only for samples where the function was applied (non-NaN values)`
`275`	`276`	`for i, reward_func_name in enumerate(self.reward_func_names):`