support archer

Author: huangjunyi.0

verl/trainer/ppo/core_algos.py

@@ -975,18 +975,19 @@ def compute_policy_loss_vanilla(
     return pg_loss, pg_clipfrac, ppo_kl, pg_clipfrac_lower
 
 
-@register_policy_loss("adc")
-def compute_policy_loss_adc(
+@register_policy_loss("archer")
+def compute_policy_loss_archer(
     old_log_prob: torch.Tensor,
     log_prob: torch.Tensor,
     advantages: torch.Tensor,
     response_mask: torch.Tensor,
     loss_agg_mode: str = "token-mean",
     config: Optional[DictConfig | AlgoConfig] = None,
     rollout_log_probs: torch.Tensor | None = None,
+    entropy: torch.Tensor | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
     """
-    ADC policy loss (Asymmetric Dual-Clipping):
+    ancher policy loss (Asymmetric Dual-Clipping):
     - For advantages > 0: invert importance ratio (use 1/ratio).
     - Extend dual-clip to advantages > 0 as well.
     - Use Soft Clip in dual-clip region to limit weight while preserving gradients.
@@ -995,71 +996,43 @@ def compute_policy_loss_adc(
 
     assert config is not None
     assert not isinstance(config, AlgoConfig)
-    clip_ratio = config.clip_ratio
-    clip_ratio_low = config.clip_ratio_low if config.clip_ratio_low is not None else clip_ratio
-    clip_ratio_high = config.clip_ratio_high if config.clip_ratio_high is not None else clip_ratio
+    clip_ratio_low = config.clip_ratio_low
+    clip_ratio_high = config.clip_ratio_hig
     clip_ratio_c = config.get("clip_ratio_c", 3.0)
+    token_entropy_quantile = config.get("token_entropy_quantile", 0.8)
+    masked_entropy = torch.where(response_mask.bool(), entropy.detach(), torch.nan)  # (bsz, response_length)
+    q80 = torch.nanquantile(masked_entropy, q=token_entropy_quantile, dim=-1, keepdim=True)  # (bsz, 1)
+    high_entropy_mask = (masked_entropy <= q80) & response_mask # only low entropy token is True
 
-    cliprange = clip_ratio
-    cliprange_low = clip_ratio_low
-    cliprange_high = clip_ratio_high
-
-    assert clip_ratio_c > 1.0, (
-        "The lower bound of the clip_ratio_c for dual-clip PPO should be greater than 1.0," + f" but get the value: {clip_ratio_c}."
-    )
-
-    negative_approx_kl = log_prob - old_log_prob
-    # Clamp for stability
-    negative_approx_kl = torch.clamp(negative_approx_kl, min=-20.0, max=20.0)
-
-    # Standard ratio
-    ratio = torch.exp(negative_approx_kl)
-    ppo_kl = verl_F.masked_mean(-negative_approx_kl, response_mask)
-
-    # For A>0: invert IS (use 1 / ratio). For A<=0: use standard ratio.
-    ratio_adc = torch.where(advantages > 0, 1.0 / ratio, ratio)
-
-    # Standard PPO loss (base branch for gradient direction)
-    pg_losses1 = -advantages * ratio_adc
-
-    if cliprange_low is None:
-        cliprange_low = cliprange
-    if cliprange_high is None:
-        cliprange_high = cliprange
-
-    pos_mask = advantages > 0
-    neg_mask = advantages < 0
+    ratio = torch.exp(torch.clamp(log_prob - old_log_prob, min=-20.0, max=20.0))
 
-    # Standard PPO clip (Hard Clip) on the ratio driving gradients
-    pg_losses2 = -advantages * torch.clamp(ratio_adc, 1 - cliprange_low, 1 + cliprange_high)
-    clip_pg_losses_base = torch.maximum(pg_losses1, pg_losses2)
+    negative_clip_ratio = torch.where(high_entropy_mask, torch.clamp(ratio, min=1-clip_ratio_low, max=None), torch.clamp(ratio, min=1-clip_ratio_high, max=None))
+    positive_clip_ratio = torch.where(high_entropy_mask, torch.clamp(ratio, min=None, max=1+clip_ratio_low), torch.clamp(ratio, min=None, max=1+clip_ratio_high))
 
-    pg_losses_dual = -advantages * clip_ratio_c
+    clip_ratio = torch.where(advantages < 0, negative_clip_ratio, positive_clip_ratio)
 
-    # Apply asymmetric dual-clip selection:
-    # - adv > 0: cap magnitude from below via min(base, -A * clip_ratio_c)
-    # - adv < 0: cap magnitude from above via max(base, -A * clip_ratio_c)
-    pg_losses_pos = torch.maximum(clip_pg_losses_base, pg_losses_dual)
-    pg_losses_neg = torch.minimum(clip_pg_losses_base, pg_losses_dual)
-    pg_losses = torch.where(pos_mask, pg_losses_pos, torch.where(neg_mask, pg_losses_neg, clip_pg_losses_base))
+    pg_clipfrac_upper = verl_F.masked_mean(torch.gt(ratio, clip_ratio).float(), response_mask)
+    pg_clipfrac_lower = verl_F.masked_mean(torch.lt(ratio, clip_ratio).float(), response_mask)
 
-    # Metrics
-    pg_clipfrac = verl_F.masked_mean(torch.gt(pg_losses2, pg_losses1).float(), response_mask)
+    negative_pg_losses_clip = -advantages * negative_clip_ratio
+    positive_pg_losses_clip = -advantages * (positive_clip_ratio / positive_clip_ratio.detach()) / positive_clip_ratio.detach()
 
-    # Dual-clip trigger fraction (both sides), measured by ratio exceeding clip_ratio_c
-    # For adv>0 we monitor the inverted ratio (ratio_adc); for adv<0 we monitor the standard ratio.
-    lower_clip_pos = pos_mask & (ratio_adc > clip_ratio_c)
-    lower_clip_neg = neg_mask & (ratio > clip_ratio_c)
-    pg_clipfrac_lower = verl_F.masked_mean((lower_clip_pos | lower_clip_neg).float(), response_mask)
+    negative_dual_clip_ratio = torch.clamp(negative_clip_ratio, min=None, max=clip_ratio_c)
+    negative_clipped_mask = torch.gt(negative_clip_ratio, negative_dual_clip_ratio)
+    negative_pg_clipfrac_dual = verl_F.masked_mean(negative_clipped_mask.float(), response_mask & (advantages < 0))
+    negative_pg_losses_dual = -advantages * negative_dual_clip_ratio.detach() * log_prob
+    negative_pg_losses = torch.where(negative_clipped_mask, negative_pg_losses_dual, negative_pg_losses_clip)
 
-    if config.tis_imp_ratio_cap > 0 and rollout_log_probs is not None:
-        tis_imp_ratio = torch.exp(old_log_prob - rollout_log_probs)
-        tis_imp_ratio = torch.clamp(tis_imp_ratio, max=config.tis_imp_ratio_cap)
-        pg_losses = pg_losses * tis_imp_ratio
+    positive_dual_clip_ratio = torch.clamp(1/positive_clip_ratio, min=None, max=clip_ratio_c)
+    positive_clipped_mask = torch.gt(1/positive_clip_ratio, positive_dual_clip_ratio)
+    positive_pg_clipfrac_dual = verl_F.masked_mean(positive_clipped_mask.float(), response_mask & (advantages > 0))
+    positive_pg_losses_dual = -advantages * positive_dual_clip_ratio.detach() * log_prob
+    positive_pg_losses = torch.where(positive_clipped_mask, positive_pg_losses_dual, positive_pg_losses_clip)
 
+    pg_losses = torch.where(advantages < 0, negative_pg_losses, positive_pg_losses)
     pg_loss = agg_loss(loss_mat=pg_losses, loss_mask=response_mask, loss_agg_mode=loss_agg_mode)
 
-    return pg_loss, pg_clipfrac, ppo_kl, pg_clipfrac_lower
+    return pg_loss, pg_clipfrac_upper, pg_clipfrac_lower, negative_pg_clipfrac_dual, positive_pg_clipfrac_dual
 
 @register_policy_loss("gspo")
 def compute_policy_loss_gspo(

verl/workers/actor/dp_actor.py

@@ -27,7 +27,7 @@
 
 import verl.utils.torch_functional as verl_F
 from verl import DataProto
-from verl.trainer.ppo.core_algos import agg_loss, get_policy_loss_fn, kl_penalty
+from verl.trainer.ppo.core_algos import agg_loss, get_policy_loss_fn, kl_penalty, compute_policy_loss_archer
 from verl.utils.attention_utils import index_first_axis, pad_input, rearrange, unpad_input
 from verl.utils.device import get_device_id, get_device_name
 from verl.utils.fsdp_utils import FSDPModule, fsdp2_clip_grad_norm_
@@ -427,7 +427,9 @@ def update_policy(self, data: DataProto):
                         loss_scale_factor = 1 / self.gradient_accumulation
 
                     # all return: (bsz, response_length)
-                    calculate_entropy = False
+                    loss_mode = self.config.policy_loss.get("loss_mode", "vanilla")
+                    is_archer = (loss_mode == "archer")
+                    calculate_entropy = is_archer
                     if entropy_coeff != 0:
                         calculate_entropy = True
                     entropy, log_prob, aux_loss = self._forward_micro_batch(
@@ -439,20 +441,31 @@ def update_policy(self, data: DataProto):
                     else:
                         old_log_prob = model_inputs["old_log_probs"]
 
-                    loss_mode = self.config.policy_loss.get("loss_mode", "vanilla")
                     # vanilla -> verl.trainer.ppo.core_algos.compute_policy_loss_vanilla
                     # gpg -> verl.trainer.ppo.core_algos.compute_policy_loss_gpg
                     # clip_cov -> verl.trainer.ppo.core_algos.compute_policy_loss_clip_cov
-                    policy_loss_fn = get_policy_loss_fn(loss_mode)
-                    pg_loss, pg_clipfrac, ppo_kl, pg_clipfrac_lower = policy_loss_fn(
-                        old_log_prob=old_log_prob,
-                        log_prob=log_prob,
-                        advantages=advantages,
-                        response_mask=response_mask,
-                        loss_agg_mode=loss_agg_mode,
-                        config=self.config,
-                        rollout_log_probs=rollout_log_probs,
-                    )
+                    if is_archer:
+                        pg_loss, pg_clipfrac_upper, pg_clipfrac_lower, negative_pg_clipfrac_dual, positive_pg_clipfrac_dual = compute_policy_loss_archer(
+                            old_log_prob=old_log_prob,
+                            log_prob=log_prob,
+                            advantages=advantages,
+                            response_mask=response_mask,
+                            loss_agg_mode=loss_agg_mode,
+                            config=self.config,
+                            rollout_log_probs=rollout_log_probs,
+                            entropy=entropy,
+                        )
+                    else:    
+                        policy_loss_fn = get_policy_loss_fn(loss_mode)
+                        pg_loss, pg_clipfrac, ppo_kl, pg_clipfrac_lower = policy_loss_fn(
+                            old_log_prob=old_log_prob,
+                            log_prob=log_prob,
+                            advantages=advantages,
+                            response_mask=response_mask,
+                            loss_agg_mode=loss_agg_mode,
+                            config=self.config,
+                            rollout_log_probs=rollout_log_probs,
+                        )
 
                     if entropy_coeff != 0:
                         entropy_loss = agg_loss(loss_mat=entropy, loss_mask=response_mask, loss_agg_mode=loss_agg_mode)
@@ -491,15 +504,26 @@ def update_policy(self, data: DataProto):
                     else:
                         loss = policy_loss * loss_scale_factor
                     loss.backward()
-
-                    micro_batch_metrics.update(
-                        {
-                            "actor/pg_loss": pg_loss.detach().item() * loss_scale_factor,
-                            "actor/pg_clipfrac": pg_clipfrac.detach().item(),
-                            "actor/ppo_kl": ppo_kl.detach().item(),
-                            "actor/pg_clipfrac_lower": pg_clipfrac_lower.detach().item(),
-                        }
-                    )
+                    
+                    if is_archer:
+                        micro_batch_metrics.update(
+                            {
+                                "actor/pg_loss": pg_loss.detach().item() * loss_scale_factor,
+                                "actor/pg_clipfrac_upper": pg_clipfrac_upper.detach().item(),
+                                "actor/pg_clipfrac_lower": pg_clipfrac_lower.detach().item(),
+                                "actor/negative_pg_clipfrac_dual": negative_pg_clipfrac_dual.detach().item(),
+                                "actor/positive_pg_clipfrac_dual": positive_pg_clipfrac_dual.detach().item(),
+                            }
+                        )
+                    else:
+                        micro_batch_metrics.update(
+                            {
+                                "actor/pg_loss": pg_loss.detach().item() * loss_scale_factor,
+                                "actor/pg_clipfrac": pg_clipfrac.detach().item(),
+                                "actor/ppo_kl": ppo_kl.detach().item(),
+                                "actor/pg_clipfrac_lower": pg_clipfrac_lower.detach().item(),
+                            }
+                        )
                     append_to_dict(metrics, micro_batch_metrics)
 
                 grad_norm = self._optimizer_step()