Remember the first PID gradient for the next derivative

tests/test_pid_grad_buffer.py

@@ -0,0 +1,28 @@
+"""The first PID step must remember its gradient."""
+import torch
+import torch_optimizer as optim
+
+
+def test_second_step_uses_the_previous_gradient():
+    param = torch.nn.Parameter(torch.tensor([0.0]))
+    opt = optim.PID(
+        [param],
+        lr=1.0,
+        momentum=0.5,
+        dampening=0.0,
+        weight_decay=0.0,
+        integral=0.0,
+        derivative=10.0,
+    )
+    param.grad = torch.tensor([1.0])
+    opt.step()
+    assert torch.equal(
+        opt.state[param]["grad_buffer"], torch.tensor([1.0])
+    )
+    # first derivative is 0: there is no earlier gradient
+    assert torch.equal(param.detach(), torch.tensor([-1.0]))
+
+    param.grad = torch.tensor([3.0])
+    opt.step()
+    # D = (1 - momentum) * (3 - 1) = 1, so the step is 3 + 10 * 1
+    assert torch.allclose(param.detach(), torch.tensor([-14.0]))

torch_optimizer/pid.py

@@ -104,7 +104,10 @@ def step(self, closure: OptLossClosure = None) -> OptFloat:
                         g_buf = param_state["grad_buffer"] = torch.zeros_like(
                             p, memory_format=torch.preserve_format
                         )
-                        g_buf = d_p
+                        # copy into the buffer. Rebinding g_buf = d_p left
+                        # the stored gradient at 0, so the next step's
+                        # derivative was taken against a zero previous grad.
+                        g_buf.copy_(d_p)
 
                         d_buf = param_state["d_buffer"] = torch.zeros_like(
                             p, memory_format=torch.preserve_format