clean code as suggested

Author: garyzhang99

tests/algorithm/policy_loss_test.py

@@ -120,7 +120,7 @@ def test_ppo_policy_loss_with_sequence_masking(self):
         policy_loss_fn_cls = POLICY_LOSS_FN.get("ppo")
         policy_loss_fn_args = policy_loss_fn_cls.default_args()
         policy_loss_fn_args["enable_sequence_masking"] = True
-        policy_loss_fn_args["delta"] = 0.1
+        policy_loss_fn_args["delta_sequence_masking"] = 0.1
         policy_loss_fn = policy_loss_fn_cls(**policy_loss_fn_args)
         loss, metrics = policy_loss_fn(log_prob=self.logprob, **self.input_data.batch)
         ppo_loss_masked = torch.tensor(0.22175675630569458)

trinity/algorithm/policy_loss_fn/ppo_policy_loss.py

@@ -21,8 +21,8 @@ def __init__(
         clip_range_high: Optional[float] = None,
         clip_ratio_c: float = 3.0,
         loss_agg_mode: Optional[str] = "token-mean",
-        enable_sequence_masking: bool = False,
-        delta: float = 0.1,
+        enable_sequence_masking: bool = False,  # introduced in DeepseekV3.2
+        delta_sequence_masking: float = 0.1,
     ) -> None:
         super().__init__(backend=backend)
         if clip_range_low is None:
@@ -39,7 +39,7 @@ def __init__(
         assert self.clip_range_high is not None, "clip_range_high must be specified."
         self.loss_agg_mode = loss_agg_mode
         self.enable_sequence_masking = enable_sequence_masking
-        self.delta = delta
+        self.delta_sequence_masking = delta_sequence_masking
 
     def __call__(  # type: ignore
         self,
@@ -55,8 +55,23 @@ def __call__(  # type: ignore
         ratio = torch.exp(negative_approx_kl)
         ppo_kl = masked_mean(-negative_approx_kl, action_mask)
 
-        # Compute sequence masking if enabled
-        sequence_mask = torch.ones_like(advantages)
+        pg_losses1 = -advantages * ratio
+        pg_losses2 = -advantages * torch.clamp(
+            ratio, 1.0 - self.clip_range_low, 1.0 + self.clip_range_high  # type: ignore
+        )
+
+        clip_pg_losses1 = torch.maximum(pg_losses1, pg_losses2)
+
+        pg_clip_frac = masked_mean(torch.gt(pg_losses2, pg_losses1).float(), action_mask)
+
+        pg_losses3 = -advantages * self.clip_ratio_c
+        clip_pg_losses2 = torch.min(pg_losses3, clip_pg_losses1)
+        pg_clipfrac_lower = masked_mean(
+            torch.gt(clip_pg_losses1, pg_losses3) * (advantages < 0).float(), action_mask
+        )
+        pg_losses = torch.where(advantages < 0, clip_pg_losses2, clip_pg_losses1)
+
+        # Apply sequence masking if enabled
         if self.enable_sequence_masking:
             # Compute sequence-level KL divergence: mean KL per sequence
             # Shape: (batch_size, seq_len) -> (batch_size,)
@@ -69,38 +84,21 @@ def __call__(  # type: ignore
             # Token-level advantage check: (batch_size, seq_len)
             has_negative_advantage = advantages < 0
             # Sequence-level KL check: (batch_size,) -> (batch_size, 1) -> (batch_size, seq_len)
-            exceeds_kl_threshold = (sequence_kl > self.delta).unsqueeze(-1).expand_as(advantages)
+            exceeds_kl_threshold = (
+                (sequence_kl > self.delta_sequence_masking).unsqueeze(-1).expand_as(advantages)
+            )
             # Mask tokens that are both negative advantage AND in high-KL sequences
             should_mask = has_negative_advantage & exceeds_kl_threshold
             sequence_mask = (~should_mask).float()
 
+            # Apply sequence mask to the losses
+            pg_losses = pg_losses * sequence_mask
+
             metrics_seq_mask = {
                 "seq_mask/masked_tokens": should_mask.float().sum().item()
                 / (action_mask.sum().item() + 1e-10),
                 "seq_mask/mean_sequence_kl": sequence_kl.mean().detach().item(),
             }
-        else:
-            metrics_seq_mask = {}
-
-        pg_losses1 = -advantages * ratio
-        pg_losses2 = -advantages * torch.clamp(
-            ratio, 1.0 - self.clip_range_low, 1.0 + self.clip_range_high  # type: ignore
-        )
-
-        clip_pg_losses1 = torch.maximum(pg_losses1, pg_losses2)
-
-        pg_clip_frac = masked_mean(torch.gt(pg_losses2, pg_losses1).float(), action_mask)
-
-        pg_losses3 = -advantages * self.clip_ratio_c
-        clip_pg_losses2 = torch.min(pg_losses3, clip_pg_losses1)
-        pg_clipfrac_lower = masked_mean(
-            torch.gt(clip_pg_losses1, pg_losses3) * (advantages < 0).float(), action_mask
-        )
-        pg_losses = torch.where(advantages < 0, clip_pg_losses2, clip_pg_losses1)
-
-        # Apply sequence mask to the losses
-        if self.enable_sequence_masking:
-            pg_losses = pg_losses * sequence_mask
 
         pg_loss = aggregate_loss(pg_losses, action_mask, loss_agg_mode=self.loss_agg_mode)
         metrics = {
@@ -109,7 +107,8 @@ def __call__(  # type: ignore
             "ppo_kl": ppo_kl.detach().item(),
             "pg_loss": pg_loss.detach().item(),
         }
-        metrics.update(metrics_seq_mask)
+        if self.enable_sequence_masking:
+            metrics.update(metrics_seq_mask)
         return pg_loss, metrics
 
     @classmethod
@@ -119,5 +118,5 @@ def default_args(cls) -> Dict:
             "clip_ratio_c": 3.0,
             "loss_agg_mode": "token-mean",
             "enable_sequence_masking": False,
-            "delta": 0.1,
+            "delta_sequence_masking": 0.1,
         }