[BugFix] Use the same beta in KLPENPPOLoss's loss value and gradient

test/objectives/test_ppo.py

@@ -1624,6 +1624,37 @@ def test_ppo_kl_pen_decoupled_proximal_policy(self):
             out["loss_objective"], -(ratio_behav * advantage - kl.unsqueeze(-1)).mean()
         )
 
+    def test_ppo_kl_pen_beta_update_keeps_gradient(self):
+        # beta adapts in place after the penalty is computed; the gradient of
+        # loss_objective must still use the beta that produced its value
+        torch.manual_seed(self.seed)
+        batch = 16
+        td = self._create_mock_data_ppo(batch=batch)
+        td["advantage"] = torch.randn(batch, 1)
+        td["value_target"] = torch.randn(batch, 1)
+        actor = self._create_mock_actor()
+        value = self._create_mock_value()
+        betas, losses, grads = [], [], []
+        # increment=1.0 keeps beta fixed and gives the reference gradient
+        for increment in (1.0, 2.0):
+            loss_fn = KLPENPPOLoss(
+                actor,
+                value,
+                beta=1.0,
+                dtarg=1e-6,
+                increment=increment,
+                entropy_bonus=False,
+            )
+            torch.manual_seed(self.seed)
+            loss = loss_fn(td.clone())["loss_objective"]
+            grads.append(autograd.grad(loss, list(actor.parameters())))
+            losses.append(loss.detach())
+            betas.append(loss_fn.beta.item())
+        assert betas == [1.0, 2.0]
+        torch.testing.assert_close(losses[1], losses[0])
+        for grad, grad_ref in zip(grads[1], grads[0]):
+            torch.testing.assert_close(grad, grad_ref)
+
     def test_ppo_delay_actor_requires_functional(self):
         actor = self._create_mock_actor()
         value = self._create_mock_value()

torchrl/objectives/ppo.py

@@ -2038,7 +2038,9 @@ def forward(self, tensordict: TensorDictBase) -> TensorDict:
                 current_log_prob = _sum_td_features(current_log_prob)
             kl = (previous_log_prob - current_log_prob).mean(0)
         kl = kl.unsqueeze(-1)
-        neg_loss = neg_loss - self.beta * kl
+        # beta is updated in place below: penalize with a copy so that the
+        # backward pass uses the beta this loss value was computed with
+        neg_loss = neg_loss - self.beta.clone() * kl
         if kl.mean() > self.dtarg * 1.5:
             self.beta.data *= self.increment
         elif kl.mean() < self.dtarg / 1.5: