rl: true per-token entropy bonus (maximize r + beta*H, explicit diversity knob, no KL)

scripts/rl.py

@@ -174,7 +174,7 @@ def update(actor, opt, submodule, ids, attn, p_len, marker, old_lp, adv, dirs_re
     gen_mask = attn[:, p_len:].bool()
     total_tok = max(int(gen_mask.sum()), 1)
     lo, hi = 1 - a.clip_eps, 1 + a.clip_eps
-    loss_sum, clipped_tok = 0.0, 0
+    loss_sum, clipped_tok, ent_sum = 0.0, 0, 0.0
     opt.zero_grad(set_to_none=True)
     for s in range(0, n, a.micro_batch):
         e = min(s + a.micro_batch, n)
@@ -183,12 +183,20 @@ def update(actor, opt, submodule, ids, attn, p_len, marker, old_lp, adv, dirs_re
                                 STEER_COEFF, device, torch.bfloat16)
         with hooked(submodule, hook):
             logits = actor(input_ids=b_ids, attention_mask=b_attn).logits[:, p_len - 1 : -1]
-        new_lp = torch.log_softmax(logits.float(), -1).gather(-1, b_ids[:, p_len:, None]).squeeze(-1)
+        logp_full = torch.log_softmax(logits.float(), -1)
         del logits
+        new_lp = logp_full.gather(-1, b_ids[:, p_len:, None]).squeeze(-1)
         m = gen_mask[s:e].to(device)
         ratio = torch.exp(new_lp - old_lp[s:e].to(device)).clamp(max=a.tis_cap)  # TIS, upper only
         A = adv[s:e, None].to(device)
         loss = (-torch.minimum(ratio * A, ratio.clamp(lo, hi) * A) * m).sum() / total_tok
+        if a.entropy_coef > 0:
+            # true per-token entropy (unbiased, logits are already here) — maximize r + β·H(π):
+            # keeps the policy stochastic for Bo-N without KL's behavior-anchoring side effect
+            ent = -(logp_full.exp() * logp_full).sum(-1)
+            ent_sum += float((ent.detach() * m).sum())
+            loss = loss - a.entropy_coef * (ent * m).sum() / total_tok
+        del logp_full
         loss.backward()  # micro-losses share the global normalizer → grads sum correctly
         loss_sum += loss.item()
         clipped_tok += int((((ratio < lo) | (ratio > hi)) & m).sum())
@@ -199,7 +207,8 @@ def update(actor, opt, submodule, ids, attn, p_len, marker, old_lp, adv, dirs_re
     else:  # stepping Adam on nan/inf grads corrupts moments AND weights
         opt.zero_grad(set_to_none=True)
         print(f"[update] non-finite grad norm ({gn}) — skipping step", flush=True)
-    return {"loss": loss_sum, "grad_norm": gn, "clipfrac": clipped_tok / total_tok}
+    return {"loss": loss_sum, "grad_norm": gn, "clipfrac": clipped_tok / total_tok,
+            "entropy": ent_sum / total_tok}
 
 
 def main():
@@ -226,6 +235,8 @@ def main():
     ap.add_argument("--len-penalty-start", type=int, default=cfg.len_penalty_start)
     ap.add_argument("--len-penalty-per-tok", type=float, default=cfg.len_penalty_per_tok)
     ap.add_argument("--no-gates", action="store_true", help="disable fluency/distinct/len shaping")
+    ap.add_argument("--entropy-coef", type=float, default=cfg.entropy_coef,
+                    help="β for maximize r + β·H(π): direct diversity pressure (Bo-N depends on it)")
     ap.add_argument("--tp", type=int, default=int(os.environ.get("WORLD_SIZE", "1")))
     ap.add_argument("--vllm-gpu-mem", type=float, default=0.35)
     ap.add_argument("--vllm-max-len", type=int, default=1024)
@@ -337,6 +348,7 @@ def main():
         log = {"reward/mean": raw_r.mean().item(), "reward/std": raw_r.std().item(),
                "reward/max": raw_r.max().item(), "reward/shaped_mean": r.mean().item(),
                "reward/gate_frac": gate_frac, "ratio/clipfrac": stats["clipfrac"],
+               "policy/entropy": stats["entropy"],
                "loss": stats["loss"], "grad_norm": stats["grad_norm"],
                "rollout/mean_logp": torch.cat(old_lps).mean().item(),
                "rollout/len_mean": n_gen / (B * G), "tokens_per_sec": n_gen / secs,

src/mxf/config.py

@@ -68,6 +68,7 @@ class RLConfig:
     lr: float = 1e-6
     clip_eps: float = 0.2
     tis_cap: float = 2.0              # TIS upper ratio cap — absorbs residual vLLM/HF kernel mismatch
+    entropy_coef: float = 0.0         # β in maximize r + β·H(π); explicit diversity knob (no KL)
     max_new_tokens: int = 96
     min_new_tokens: int = 16
     temperature: float = 1.0