fix(metrics): correct four RL trainer diagnostics

Two of the four are padding problems and two are not. No loss or gradient path
changes in any of them.

Online DPO summed raw `(B, T)` logprobs for `objective/kl`, `objective/entropy`
and, through `non_score_reward`, `objective/rlhf_reward`, two lines below a loss
that already multiplies by `~cr_padding_mask`. PPO masked `pg_loss` and
`pg_clipfrac` but left `policy/approxkl_avg`, `policy/entropy_avg` and
`policy/ratio_avg` on a bare `.mean()`. Both now use the mask already in scope.

The padded positions hold `log p(pad | context)`: `_forward` returns
`selective_log_softmax` over the full padded width and nothing zeroes it before
the stat sites. That value rises toward zero as the policy learns to emit pad
after EOS, so the contamination shrinks during training instead of sitting at a
fixed offset. Holding the true masked entropy fixed and moving only `p(pad|ctx)`
from 0.01 to 0.99 slides the reported value from 48.99 to 37.51, a downward
trend that nothing real is producing.

XPO logged `logps/chosen` and `logps/rejected` as policy plus reference logps.
Online DPO logs the policy term alone and Nash-MD logs the two separately, so
the sum matches neither and has no interpretation. Now the policy term.

GRPO computed `policy_loss` on every loss branch but appended it only inside the
entropy-bonus block, so the metric vanished for standard runs. The append moves
out. Its value is unchanged: it is captured after the KL term is folded in, so
with `beta != 0` it carries that too, and excludes only the entropy bonus and
the MoE auxiliary loss.

Resolves #6809

Author: behroozazarkhalili

trl/experimental/online_dpo/online_dpo_trainer.py

@@ -1249,7 +1249,9 @@ def training_step(
         self.stats["logps/chosen"].append(self.accelerator.gather_for_metrics(chosen_logprobs_sum).mean().item())
         self.stats["logps/rejected"].append(self.accelerator.gather_for_metrics(rejected_logprobs_sum).mean().item())
 
-        kl = logprobs - ref_logprobs
+        # Exclude padding positions, as the loss path above does: `logprobs` and `ref_logprobs` span the padded
+        # completion width, so summing them raw adds post-EOS garbage to every objective below.
+        kl = (logprobs - ref_logprobs) * ~padding_mask
         mean_kl = kl.sum(1).mean()
         self.stats["objective/kl"].append(self.accelerator.gather_for_metrics(mean_kl).mean().item())
         non_score_reward = (-self.beta * kl).sum(1)
@@ -1262,7 +1264,7 @@ def training_step(
             rlhf_reward = rewards + non_score_reward
             self.stats["objective/rlhf_reward"].append(self.accelerator.gather_for_metrics(rlhf_reward).mean().item())
 
-        mean_entropy = -logprobs.sum(1).mean()
+        mean_entropy = -(logprobs * ~padding_mask).sum(1).mean()
         self.stats["objective/entropy"].append(self.accelerator.gather_for_metrics(mean_entropy).mean().item())
         chosen_rewards = self.beta * (chosen_logprobs_sum - chosen_ref_logprobs_sum)
         gathered_chosen_rewards = self.accelerator.gather_for_metrics(chosen_rewards)

trl/experimental/ppo/ppo_trainer.py

@@ -854,7 +854,7 @@ def repeat_generator():
                                 )
                                 prob_dist = torch.nn.functional.softmax(logits, dim=-1)
                                 entropy = torch.logsumexp(logits, dim=-1) - torch.sum(prob_dist * logits, dim=-1)
-                                approxkl = 0.5 * (logprobs_diff**2).mean()
+                                approxkl = 0.5 * masked_mean(logprobs_diff**2, ~padding_mask[micro_batch_inds])
                                 approxkl_stats[ppo_epoch_idx, minibatch_idx, gradient_accumulation_idx] = approxkl
                                 pg_clipfrac_stats[ppo_epoch_idx, minibatch_idx, gradient_accumulation_idx] = (
                                     pg_clipfrac
@@ -864,8 +864,12 @@ def repeat_generator():
                                 vf_clipfrac_stats[ppo_epoch_idx, minibatch_idx, gradient_accumulation_idx] = (
                                     vf_clipfrac
                                 )
-                                entropy_stats[ppo_epoch_idx, minibatch_idx, gradient_accumulation_idx] = entropy.mean()
-                                ratio_stats[ppo_epoch_idx, minibatch_idx, gradient_accumulation_idx] = ratio.mean()
+                                entropy_stats[ppo_epoch_idx, minibatch_idx, gradient_accumulation_idx] = masked_mean(
+                                    entropy, ~padding_mask[micro_batch_inds]
+                                )
+                                ratio_stats[ppo_epoch_idx, minibatch_idx, gradient_accumulation_idx] = masked_mean(
+                                    ratio, ~padding_mask[micro_batch_inds]
+                                )
                         gradient_accumulation_idx += 1
                     minibatch_idx += 1
                     # del everything and empty cache

trl/experimental/xpo/xpo_trainer.py

@@ -384,8 +384,8 @@ def gather_mean(tensor):
         rejected_ref_logprobs = torch.where(~chosen_mask, ref_logprobs_model_data_sum, ref_logprobs_ref_data_sum)
         rejected_log_ratios = rejected_model_logprobs - rejected_ref_logprobs
 
-        self.stats["logps/chosen"].append(gather_mean(chosen_model_logprobs.mean() + chosen_ref_logprobs.mean()))
-        self.stats["logps/rejected"].append(gather_mean(rejected_model_logprobs.mean() + rejected_ref_logprobs.mean()))
+        self.stats["logps/chosen"].append(gather_mean(chosen_model_logprobs.mean()))
+        self.stats["logps/rejected"].append(gather_mean(rejected_model_logprobs.mean()))
 
         # Log rewards
         # Compute various statistics

trl/trainer/grpo_trainer.py

@@ -3272,6 +3272,11 @@ def _compute_loss(self, model, inputs):
         else:
             raise ValueError(f"Unknown loss type: {self.loss_type}")
 
+        # Log for every run, not only the ones with an entropy bonus enabled. `policy_loss` is captured after the
+        # KL term is folded in above, so with `beta != 0` it carries that too; it excludes only the entropy bonus
+        # and the MoE auxiliary loss, both added below.
+        self._metrics[mode]["policy_loss"].append(self.accelerator.gather(policy_loss).nanmean().item())
+
         # Entropy bonus: add entropy regularization to encourage exploration. _entropy_bonus_enabled is set
         # whenever a non-zero static coef is set OR adaptive mode is enabled (adaptive stays enabled even when
         # entropy_coef has been decremented to entropy_coef_min so it can recover once entropy drops again).
@@ -3299,8 +3304,6 @@ def _compute_loss(self, model, inputs):
 
             loss = loss - apply_coef * entropy_loss
 
-            self._metrics[mode]["policy_loss"].append(self.accelerator.gather(policy_loss).nanmean().item())
-
             # Adaptive update. Gated on train mode so evaluation cannot mutate the entropy controller state.
             if self.use_adaptive_entropy and mode == "train":
                 # Accumulate the entropy sum and active-token count of every micro-batch into a running window