From 5f0ff2e06fd362cdcbc6fbe83ba8184295bd03f2 Mon Sep 17 00:00:00 2001
From: David-Wu1119 <133224895+David-Wu1119@users.noreply.github.com>
Date: Wed, 11 Nov 2026 00:49:36 -0500
Subject: [PATCH] fix(kl): make the adaptive KL coefficient update work

`KLFn.update_kl_coef` clipped the proportional error with `torch.clip` on a
Python float, which raises `TypeError` (torch.clip only accepts tensors). It
is called from `apply_kl_penalty_to_reward` on every trainer step, so any run
with a reward KL penalty and `kl_penalty_fn_args.adaptive: true` crashed at
the first step instead of adapting the coefficient.

Clip with `min`/`max`, matching verl's AdaptiveKLController
(`np.clip(current_kl / target - 1, -0.2, 0.2)`).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 tests/algorithm/kl_fn_test.py    | 25 +++++++++++++++++++++++++
 trinity/algorithm/kl_fn/kl_fn.py |  4 ++--
 2 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/tests/algorithm/kl_fn_test.py b/tests/algorithm/kl_fn_test.py
index 4f8bc3393dd..f1742ee0ad2 100644
--- a/tests/algorithm/kl_fn_test.py
+++ b/tests/algorithm/kl_fn_test.py
@@ -129,3 +129,28 @@ def test_kl_loss_aggregation_modes(self):
             old_logprob=self.old_logprob,
         )
         self.assertGreater(kl_loss_sum.item(), kl_loss_mean.item())
+
+    def test_adaptive_kl_penalty_updates_coef(self):
+        from verl import DataProto
+
+        batch_size, prompt_len, resp_len = 4, 3, 5
+        response_mask = torch.ones(batch_size, resp_len)
+        batch = DataProto.from_dict(
+            {
+                "responses": torch.zeros(batch_size, resp_len, dtype=torch.long),
+                "attention_mask": torch.ones(batch_size, prompt_len + resp_len),
+                "response_mask": response_mask,
+                "token_level_scores": torch.zeros(batch_size, resp_len),
+                "old_log_probs": torch.zeros(batch_size, resp_len),
+                "ref_log_prob": torch.full((batch_size, resp_len), -0.5),
+            }
+        )
+        # k1 KL is 0.5 per token, so the proportional error is clipped to +0.2 / -0.2
+        for target_kl, multiplier in [(0.1, 1.2), (1.0, 0.8)]:
+            kl_fn = KL_FN.get("k1")(adaptive=True, kl_coef=0.1, target_kl=target_kl, horizon=10)
+            _, metrics = kl_fn.apply_kl_penalty_to_reward(batch)
+            self.assertAlmostEqual(metrics["kl"], 0.5, places=6)
+            self.assertAlmostEqual(metrics["kl_coef"], 0.1)
+            # verl AdaptiveKLController: coef *= 1 + clip(kl / target - 1, -0.2, 0.2) * n / horizon
+            expected = 0.1 * (1 + (multiplier - 1) * batch_size / 10)
+            self.assertAlmostEqual(kl_fn.kl_coef, expected, places=6)
diff --git a/trinity/algorithm/kl_fn/kl_fn.py b/trinity/algorithm/kl_fn/kl_fn.py
index 312f5303293..3ac46b23623 100644
--- a/trinity/algorithm/kl_fn/kl_fn.py
+++ b/trinity/algorithm/kl_fn/kl_fn.py
@@ -37,8 +37,8 @@ def update_kl_coef(self, current_kl: float, batch_size: int) -> None:
         """Update kl coefficient."""
         if self.adaptive:
             target_kl = self.target_kl
-            proportional_error = torch.clip(current_kl / target_kl - 1, -0.2, 0.2).item()  # type: ignore
-            multiplier = 1 + proportional_error * batch_size / self.horizon
+            proportional_error = min(max(current_kl / target_kl - 1, -0.2), 0.2)  # type: ignore
+            multiplier = 1 + proportional_error * batch_size / self.horizon  # type: ignore
             self.kl_coef *= multiplier
 
     def apply_kl_penalty_to_reward(self, experiences: Any) -> Tuple[Any, Dict]: