skip_whiten_if_response_mask_sum_eq_0

Author: ToasterSC

roll/utils/functionals.py

@@ -683,9 +683,13 @@ def compute_advantage(
 ):
     if response_mask is None:
         response_mask = data.batch["response_mask"][:, 1:]
+    if response_mask.sum() == 0:
+        skip_whiten = True
+    else:
+        skip_whiten = False
 
     token_level_rewards = data.batch["token_level_rewards"].float()
-    if whiten_rewards:
+    if whiten_rewards and not skip_whiten:
         token_level_rewards = masked_whiten(values=token_level_rewards, mask=response_mask)
     # data.batch['token_level_rewards'] = token_level_rewards
     token_level_rewards = token_level_rewards * response_mask
@@ -708,7 +712,7 @@ def compute_advantage(
         raise NotImplementedError
 
     data.batch["raw_advantages"] = advantages
-    if whiten_advantages:
+    if whiten_advantages and not skip_whiten:
         # TODO whiten过程中是否要考虑response的长度？
         advantages = masked_whiten(values=advantages, mask=response_mask)
     advantages = advantages * response_mask