fix: pad agg_logp with NaN on zero-variance skip to keep is_s alignment

The zero-variance bail at train.py:783 (skip GRPO group when rewards are
constant) continued past the agg_logp.extend at line 821. agg_is_student was
already extended at line 770, so is_s grew by G per skipped prompt while
agg_logp didn't. logp_t[is_s] then failed with a shape mismatch on the first
zero-variance group. Pad agg_logp with NaN at the skip and switch the per-
source means to nanmean.

Caught by #52 vanilla matched-control crashing at step 0.

src/projected_grpo/train.py

@@ -781,6 +781,10 @@ def main(cfg: Config) -> int:
             # dominant pathology with our binary-ish reward shape on a weak 2B
             # substrate (every group can clip to 0.25 = format_only).
             if (rewards.max() - rewards.min()).item() < 1e-4:
+                # Pad agg_logp with NaN to keep it aligned with agg_is_student
+                # (extended above at line 770). Skipping the gen_logp forward
+                # here is the whole point of the zero-variance bail.
+                agg_logp.extend([float("nan")] * len(rs))
                 continue
             centered = rewards - rewards.mean()
             adv = centered if cfg.unbiased else centered / (rewards.std() + 1e-4)
@@ -943,9 +947,11 @@ def main(cfg: Config) -> int:
         gt_s_n   = int((g_t & is_s).sum())
         gt_t_n   = int((g_t & ~is_s).sum())
         rew_s_mean = rewards_t[is_s].mean().item() if n_s else float("nan")
+        # Skipped (zero-variance) prompts pad agg_logp with NaN above to keep
+        # alignment with is_s. nanmean drops them from the per-source means.
         logp_t = torch.tensor(agg_logp, dtype=torch.float32) if agg_logp else torch.zeros(0)
-        lp_s_mean = logp_t[is_s].mean().item()  if n_s else float("nan")
-        lp_t_mean = logp_t[~is_s].mean().item() if n_t else float("nan")
+        lp_s_mean = logp_t[is_s].nanmean().item()  if n_s else float("nan")
+        lp_t_mean = logp_t[~is_s].nanmean().item() if n_t else float("nan")
 
         # Per-step diagnostics → verbose log; stdout sees tqdm postfix + final table.
         n_fin = sum(agg_finished)