fix(jsd): make the JSD advantage follow the gradient of the JSD

tests/algorithm/advantage_fn_test.py

@@ -394,3 +394,50 @@ def test_step_wise_grpo_with_std_threshold(self):
         # Verify that all remaining experiences are from task 2
         for exp in processed_exps:
             self.assertEqual(exp.eid.task, 2)
+
+
+class TestJSDAdvantageFn(unittest.TestCase):
+    """The JSD advantage must follow the gradient of the generalized JSD."""
+
+    @staticmethod
+    def _expected_policy_gradient(advantage_fn, theta, teacher_logprobs, chunked):
+        # Exact expectation over the sampled token: E_{y~student}[A(y) * grad log student(y)].
+        student_logprobs = torch.log_softmax(theta, dim=-1)
+        detached = student_logprobs.detach()
+        if chunked:
+            cost = advantage_fn._js_divergence_per_token_chunked(
+                detached[None], teacher_logprobs[None], torch.ones(1, theta.numel())
+            )[0]
+        else:
+            cost = advantage_fn._js_divergence_per_token(detached, teacher_logprobs)
+        surrogate = (detached.exp() * -cost * student_logprobs).sum()
+        return torch.autograd.grad(surrogate, theta)[0]
+
+    def test_jsd_advantage_descends_jsd(self):
+        torch.manual_seed(0)
+        teacher_logprobs = torch.log_softmax(2 * torch.randn(8), dim=-1)
+        for beta in (0.3, 0.5, 0.9):
+            for chunked in (False, True):
+                advantage_fn = ADVANTAGE_FN.get("jsd")(lambda_coef=beta, chunk_size=3)
+                theta = torch.randn(8, requires_grad=True)
+                policy_gradient = self._expected_policy_gradient(
+                    advantage_fn, theta, teacher_logprobs, chunked
+                )
+
+                student_logprobs = torch.log_softmax(theta, dim=-1)
+                mixture = torch.logsumexp(
+                    torch.stack(
+                        [
+                            student_logprobs + np.log(1 - beta),
+                            teacher_logprobs + np.log(beta),
+                        ]
+                    ),
+                    dim=0,
+                )
+                jsd = (
+                    beta * (teacher_logprobs.exp() * (teacher_logprobs - mixture)).sum()
+                    + (1 - beta) * (student_logprobs.exp() * (student_logprobs - mixture)).sum()
+                )
+                jsd_gradient = torch.autograd.grad(jsd, theta)[0]
+
+                torch.testing.assert_close(policy_gradient, -jsd_gradient)

trinity/algorithm/advantage_fn/jsd_advantage.py

@@ -22,6 +22,11 @@ class JSDAdvantage(AdvantageFn):
     - beta=1: JSD = KL(teacher || student)  [pure KL]
     - else: JSD = beta*KL(teacher||M) + (1-beta)*KL(student||M), M = mixture in log-space.
 
+    Only the sampled token's log-probs are available and the tokens are sampled from the
+    student, so the per-token value is a REINFORCE cost whose expected gradient is the
+    gradient of the divergence. For the mixture this cost is (1-beta) * (log s - log m):
+    d/ds(y) [beta*KL(T||M) + (1-beta)*KL(S||M)] = (1-beta) * log(s(y) / m(y)).
+
     The teacher_logprobs should be stored in Experience.teacher_logprobs
     by the workflow during exploration.
     """
@@ -57,7 +62,7 @@ def _js_divergence_per_token(
         Compute JSD per token in log-space only (no exp).
         - beta=0: KL(student || teacher) = student_logprobs - teacher_logprobs
         - beta=1: KL(teacher || student) = teacher_logprobs - student_logprobs
-        - else: mixture log_m = logsumexp([s+log(1-beta), t+log(beta)]); JSD = beta*(t-log_m) + (1-beta)*(s-log_m).
+        - else: mixture log_m = logsumexp([s+log(1-beta), t+log(beta)]); cost = (1-beta)*(s-log_m).
         """
         beta = self.lambda_coef
         s = student_logprobs
@@ -81,10 +86,9 @@ def _js_divergence_per_token(
                 torch.stack([s + log_1_minus_beta, t + log_beta]),
                 dim=0,
             )
-            # KL(teacher || mixture) = t - log_m, KL(student || mixture) = s - log_m (log-space, no exp)
-            kl_teacher = t - mixture_log
-            kl_student = s - mixture_log
-            jsd = beta_t * kl_teacher + (1 - beta_t) * kl_student
+            # The tokens are sampled from the student, so t - log_m is not an estimate of
+            # KL(teacher || mixture); the gradient of the whole JSD comes from this term alone.
+            jsd = (1 - beta_t) * (s - mixture_log)
 
         if mask is not None:
             jsd = jsd * mask
@@ -128,9 +132,7 @@ def _js_divergence_per_token_chunked(
                     torch.stack([s_chunk + log_1_minus_beta, t_chunk + log_beta]),
                     dim=0,
                 )
-                kl_t = t_chunk - mixture_log
-                kl_s = s_chunk - mixture_log
-                jsd_valid[start:end] = beta_t * kl_t + (1 - beta_t) * kl_s
+                jsd_valid[start:end] = (1 - beta_t) * (s_chunk - mixture_log)
 
         out = flat_s.new_zeros(flat_s.shape)
         out[valid] = jsd_valid