Fix rewards_accuracy to use per-sample comparison instead of scalar

Previously compared mean chosen vs mean rejected rewards (always 0 or 1).
Now computes per-sample accuracy and averages, matching dpo_tune_cache.py.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: finbarrtimbers

open_instruct/olmo_core_train_modules.py

@@ -79,6 +79,7 @@ def __init__(
         self._total_rejected_logps = torch.tensor(0.0, device=device)
         self._total_chosen_rewards = torch.tensor(0.0, device=device)
         self._total_rejected_rewards = torch.tensor(0.0, device=device)
+        self._total_rewards_accuracy = torch.tensor(0.0, device=device)
         self._total_aux_loss = torch.tensor(0.0, device=device) if args.load_balancing_loss else None
 
         if args.packing:
@@ -161,13 +162,15 @@ def train_batch(self, batch: dict[str, Any], dry_run: bool = False) -> None:
         self._total_rejected_logps.zero_()
         self._total_chosen_rewards.zero_()
         self._total_rejected_rewards.zero_()
+        self._total_rewards_accuracy.zero_()
         if self._total_aux_loss is not None:
             self._total_aux_loss.zero_()
         total_loss = self._total_loss
         total_chosen_logps = self._total_chosen_logps
         total_rejected_logps = self._total_rejected_logps
         total_chosen_rewards = self._total_chosen_rewards
         total_rejected_rewards = self._total_rejected_rewards
+        total_rewards_accuracy = self._total_rewards_accuracy
         total_aux_loss = self._total_aux_loss
 
         for micro_batch_idx, micro_batch in enumerate(micro_batches):
@@ -209,6 +212,9 @@ def train_batch(self, batch: dict[str, Any], dry_run: bool = False) -> None:
                 if self.args.loss_type.computes_reward_metrics:
                     total_chosen_rewards += chosen_rewards.mean().detach() / num_micro_batches
                     total_rejected_rewards += rejected_rewards.mean().detach() / num_micro_batches
+                    total_rewards_accuracy += (
+                        chosen_rewards > rejected_rewards
+                    ).float().mean().detach() / num_micro_batches
                 if total_aux_loss is not None and aux_loss is not None:
                     total_aux_loss += aux_loss.detach() / num_micro_batches
 
@@ -222,14 +228,13 @@ def train_batch(self, batch: dict[str, Any], dry_run: bool = False) -> None:
             self.record_metric("train/logps_rejected", total_rejected_logps, ReduceType.mean)
 
             if self.args.loss_type.computes_reward_metrics:
-                accuracy = (total_chosen_rewards > total_rejected_rewards).float()
                 margin = total_chosen_rewards - total_rejected_rewards
                 self.record_metric("train/rewards_chosen", total_chosen_rewards, ReduceType.mean)
                 self.record_metric("train/rewards_rejected", total_rejected_rewards, ReduceType.mean)
                 self.record_metric(
                     "train/rewards_average", (total_chosen_rewards + total_rejected_rewards) / 2, ReduceType.mean
                 )
-                self.record_metric("train/rewards_accuracy", accuracy, ReduceType.mean)
+                self.record_metric("train/rewards_accuracy", total_rewards_accuracy, ReduceType.mean)
                 self.record_metric("train/rewards_margin", margin, ReduceType.mean)
 
             chosen_lengths = (batch["chosen_labels"] != -100).sum()