fix: token-level ratio and prompt masking in GRPO strategy

- Mask prompt tokens to 0 so their logprobs excluded from ratio/KL
- Switch to token-level ratio + PPO clipping via reduction='none'
- Slice response token logprobs from full sequence output
- Replace k3 KL estimator with non-negative k1 estimator
- Fix epsilon from finifo.eps (~1e-38) to 1e-8
- Remove unused 'reduction' param from GRPOStrategy.__init__
- Clarify offline batch semantics in docstring
- Add 11 unit tests for masking, advantage, KL, sync, clipping
- Sync training.md and architecture.md docs

astrai/trainer/strategy.py

@@ -267,9 +267,14 @@ class DPOStrategy(BaseStrategy):
 class GRPOStrategy(BaseStrategy):
     """Group Relative Policy Optimization strategy.
 
-    On-policy GRPO following DeepSeek-R1: the policy model is updated while
-    a frozen ref_model stores the old-policy log-probs.  ratio = exp(logπ_θ - logπ_ref),
-    clipped PPO objective.  Call ``sync_ref_model()`` after each data-generation round.
+    Implements GRPO following DeepSeek-R1 with token-level PPO clipping.
+    Advantages are group-normalized from scalar per-response rewards and
+    broadcast across all response tokens.  The loss is computed **only on
+    response tokens** — prompt tokens are masked out.
+
+    The strategy expects offline-collected batches (``responses`` / ``rewards``
+    pre-generated by the current or a recent policy).  Call ``sync_ref_model()``
+    after each data-generation round so ``ref_model`` tracks the sampling policy.
     """
 
     def __init__(
@@ -279,7 +284,6 @@ class GRPOStrategy(BaseStrategy):
         clip_eps: float = 0.2,
         kl_coef: float = 0.01,
         group_size: int = 4,
-        reduction: str = "mean",
         sync_interval: int = 200,
         **kwargs,
     ):
@@ -290,7 +294,6 @@ class GRPOStrategy(BaseStrategy):
         self.clip_eps = clip_eps
         self.kl_coef = kl_coef
         self.group_size = group_size
-        self.reduction = reduction
         self.sync_interval = sync_interval
         self._step = 0
 
@@ -313,33 +316,54 @@ class GRPOStrategy(BaseStrategy):
         responses_flat = responses.view(-1, response_len)
         masks_flat = masks.view(-1, response_len)
         prompt_expanded = prompts.unsqueeze(1).repeat(1, group_size, 1).flatten(0, 1)
+        prompt_len = prompt_expanded.size(1)
 
         full_sequences = torch.cat([prompt_expanded, responses_flat], dim=-1)
-        full_masks = torch.cat([torch.ones_like(prompt_expanded), masks_flat], dim=-1)
-
-        log_probs_policy = get_logprobs(
-            self.model, full_sequences, full_masks, self.reduction
-        )
-        log_probs_policy = log_probs_policy.view(batch_size, group_size)
+        # Prompt tokens are masked out (0) so logprobs are computed only for
+        # response tokens.  get_logprobs shifts the mask by one position, so
+        # the first response token's logprob (predicted from the last prompt
+        # token) is correctly included.
+        full_masks = torch.cat([torch.zeros_like(prompt_expanded), masks_flat], dim=-1)
 
+        # get_logprobs returns [B*G, S-1] (S = prompt_len + response_len).
+        # Response token logprobs occupy the last ``response_len`` positions
+        # (the first response token is predicted from the last prompt token).
+        token_log_probs_policy = get_logprobs(
+            self.model, full_sequences, full_masks, "none"
+        )[:, prompt_len - 1 :]
         with torch.no_grad():
-            log_probs_ref = get_logprobs(
-                self.ref_model, full_sequences, full_masks, self.reduction
-            )
-            log_probs_ref = log_probs_ref.view(batch_size, group_size)
+            token_log_probs_ref = get_logprobs(
+                self.ref_model, full_sequences, full_masks, "none"
+            )[:, prompt_len - 1 :]
 
-        eps = torch.finfo(log_probs_policy.dtype).eps
+        # Reshape to [B, G, response_len]
+        token_log_probs_policy = token_log_probs_policy.view(batch_size, group_size, -1)
+        token_log_probs_ref = token_log_probs_ref.view(batch_size, group_size, -1)
+        token_masks = masks_flat.view(batch_size, group_size, -1).float()
+
+        # Group-normalized advantages from scalar per-response rewards.
+        eps = 1e-8
         mean = rewards.mean(dim=-1, keepdim=True)
         std = rewards.std(dim=-1, keepdim=True)
         advantages = (rewards - mean) / (std + eps)
+        # Broadcast scalar advantage to every response token: [B, G, 1]
+        advantages = advantages.unsqueeze(-1)
 
-        ratio = torch.exp(log_probs_policy - log_probs_ref)
+        # Token-level ratio and PPO clipping.
+        log_ratio = token_log_probs_policy - token_log_probs_ref
+        ratio = torch.exp(log_ratio)
 
         surr1 = ratio * advantages
         surr2 = torch.clamp(ratio, 1 - self.clip_eps, 1 + self.clip_eps) * advantages
+        per_token_policy_loss = -torch.min(surr1, surr2)
+        token_count = token_masks.sum().clamp(min=1.0)
+        policy_loss = (per_token_policy_loss * token_masks).sum() / token_count
+
+        # KL penalty with k1 estimator (non-negative): r - log(r) - 1, r=π_ref/π_θ.
+        r = torch.exp(-log_ratio)
+        kl_per_token = r - torch.log(r + eps) - 1.0
+        kl_penalty = self.kl_coef * (kl_per_token * token_masks).sum() / token_count
 
-        policy_loss = -torch.min(surr1, surr2).mean()
-        kl_penalty = self.kl_coef * (log_probs_policy - log_probs_ref).square().mean()
         total_loss = policy_loss + kl_penalty
 
         return total_loss