Use weighted list reward functions

separate the calculation of advantage from the calculation of rewards

src/retool_trainer.py

@@ -7,7 +7,7 @@ import datasets
 import torch
 import torch.utils.data
 import transformers
-from accelerate.utils import broadcast_object_list, gather, gather_object, is_peft_model, set_seed
+#from accelerate.utils import broadcast_object_list, gather, gather_object, is_peft_model, set_seed
 from datasets import Dataset, IterableDataset
 from packaging import version
 from torch import nn
@@ -163,19 +163,96 @@ class ReToolTrainer(Trainer):  # Change this line
             return self._check_equivalence(predicted, ground_truth)
         return False
 
-    def _compute_rewards_and_advantages(self, completions_text, ground_truths, device):
-        """Simplified reward and advantage computation for ReTool."""
+    def _compute_rewards(self, inputs, prompts, completions, completion_ids_list=None):
+        """Calculate rewards for completions and combine them according to weights."""
+        device = self.device  # Your device might be set differently
+        rewards_per_func = torch.zeros(len(prompts), len(self.reward_funcs), device=device)
+
+        # Extract additional arguments from inputs if needed
+        reward_kwargs = {}
+        if isinstance(inputs, list) and len(inputs) > 0 and isinstance(inputs[0], dict):
+            keys = [key for key in inputs[0] if key not in ["prompt", "completion", "completion_ids"]]
+            reward_kwargs = {key: [example[key] for example in inputs] for key in keys}
         
-        # Compute binary rewards
-        rewards = []
-        for completion_text, ground_truth in zip(completions_text, ground_truths):
-            if self._is_correct_answer(completion_text, ground_truth):
-                rewards.append(1.0)
-            else:
-                rewards.append(-1.0)
+        # Add correct_answers to kwargs if present (common in math reasoning tasks)
+        if "correct_answers" in reward_kwargs:
+            reward_kwargs["solution"] = reward_kwargs["correct_answers"]  # Alias for compatibility
+
+        # Calculate rewards for each function with non-zero weight
+        for i, (reward_func, func_name) in enumerate(zip(self.reward_funcs, self.reward_func_names)):
+            # Skip computation if weight is zero
+            if abs(self.reward_weights[i].item()) < 1e-6:
+                rewards_per_func[:, i] = float('nan')
+                if self.verbose:
+                    print(f"Skipping reward '{func_name}' (zero weight)")
+                continue
+            
+            # Calculate reward
+            try:
+                # Call the reward function with appropriate arguments
+                rewards = reward_func(
+                    prompts=prompts, 
+                    completions=completions,
+                    completion_ids=completion_ids_list if completion_ids_list is not None else None,
+                    **reward_kwargs
+                )
+                
+                # Convert None values to NaN and ensure it's a tensor
+                rewards = [r if r is not None else float('nan') for r in rewards]
+                rewards_per_func[:, i] = torch.tensor(rewards, dtype=torch.float32, device=device)
+                
+                # Log reward statistics if verbose
+                if self.verbose:
+                    valid_rewards = [r for r in rewards if not (r is None or (isinstance(r, float) and math.isnan(r)))]
+                    if valid_rewards:
+                        print(f"Reward '{func_name}': min={min(valid_rewards):.4f}, max={max(valid_rewards):.4f}, "
+                            f"mean={sum(valid_rewards)/len(valid_rewards):.4f}")
+            except Exception as e:
+                print(f"Error in reward function '{func_name}': {e}")
+                rewards_per_func[:, i] = float('nan')
+        
+        # Combine rewards using weights
+        rewards = (rewards_per_func * self.reward_weights.to(device).unsqueeze(0)).nansum(dim=1)
         
-        # For now: advantages = rewards (skip group normalization)
-        advantages = torch.tensor(rewards, dtype=torch.float32, device=device)
+        # Convert to list for easier handling
+        final_rewards = rewards.cpu().tolist()
+        
+        return final_rewards
+    
+
+    def compute_rewards_and_advantages(self, inputs, prompts, completions, completion_ids_list=None):
+        """Calculate rewards and compute advantages based on those rewards."""
+        # First calculate rewards
+        rewards = self.compute_rewards(inputs, prompts, completions, completion_ids_list)
+        
+        # Convert to tensor if not already
+        if not isinstance(rewards, torch.Tensor):
+            rewards = torch.tensor(rewards, dtype=torch.float32, device=self.device)
+        
+        # For now, simple advantage calculation
+        advantages = rewards.clone()  # Simple case: advantages = rewards
+        
+        # If later I want to implement GRPO-style advantage calculation:
+        if self.use_grouped_advantages:
+            # Reshape rewards into groups (assuming self.num_generations is set)
+            grouped_rewards = rewards.view(-1, self.num_generations)
+            
+            # Calculate statistics per group
+            mean_grouped_rewards = grouped_rewards.mean(dim=1)
+            std_grouped_rewards = grouped_rewards.std(dim=1)
+            
+            # Expand means and stds to match original shape
+            mean_expanded = mean_grouped_rewards.repeat_interleave(self.num_generations, dim=0)
+            std_expanded = std_grouped_rewards.repeat_interleave(self.num_generations, dim=0)
+            
+            # Compute advantages: reward - baseline
+            advantages = rewards - mean_expanded
+            
+            # Optionally normalize advantages
+            if self.normalize_advantages:
+                # Avoid division by zero
+                std_expanded = torch.clamp(std_expanded, min=1e-8)
+                advantages = advantages / std_expanded
         
         return advantages