Merge pull request #133 from Yif-Yang/fix/scheduler-preserve-annealing-contract

fix(scheduler): preserve existing annealing contract after #123

skillopt/optimizer/scheduler.py

@@ -63,16 +63,15 @@ class ConstantScheduler(LRScheduler):
 class LinearScheduler(LRScheduler):
     """Linear decay from ``max_lr`` to ``min_lr`` over ``total_steps``.
 
-    The scheduler guarantees that ``step()`` returns ``max_lr`` on the first
-    call and ``min_lr`` on the ``total_steps``-th call.  Intermediate steps
-    are linearly interpolated and rounded to the nearest integer.
+    For multi-step runs, the first call evaluates the first decay interval
+    (``t = 1 / total_steps``) and the ``total_steps``-th call returns
+    ``min_lr``.  Values are rounded to the nearest integer.
     """
 
     def _compute_lr(self, step: int) -> int:
         if self.total_steps <= 1:
             return self.max_lr
-        s = min(step, self.total_steps)
-        t = (s - 1) / (self.total_steps - 1)
+        t = min(step, self.total_steps) / self.total_steps
         lr = self.max_lr + (self.min_lr - self.max_lr) * t
         return max(self.min_lr, round(lr))
 
@@ -80,16 +79,15 @@ class LinearScheduler(LRScheduler):
 class CosineScheduler(LRScheduler):
     """Cosine annealing from ``max_lr`` to ``min_lr`` over ``total_steps``.
 
-    The scheduler guarantees that ``step()`` returns ``max_lr`` on the first
-    call and ``min_lr`` on the ``total_steps``-th call.  Intermediate steps
-    follow a half-cosine curve that starts and ends flat.
+    For multi-step runs, the first call evaluates the first decay interval
+    (``t = 1 / total_steps``) and the ``total_steps``-th call returns
+    ``min_lr``.  Intermediate steps follow a half-cosine curve.
     """
 
     def _compute_lr(self, step: int) -> int:
         if self.total_steps <= 1:
             return self.max_lr
-        s = min(step, self.total_steps)
-        t = (s - 1) / (self.total_steps - 1)
+        t = min(step, self.total_steps) / self.total_steps
         lr = self.min_lr + 0.5 * (self.max_lr - self.min_lr) * (1 + math.cos(math.pi * t))
         return max(self.min_lr, round(lr))
 

tests/test_scheduler.py

@@ -9,8 +9,9 @@ math — making it an ideal target for precise unit tests.
 
 Scheduler contract
 ------------------
-Decay schedulers (Linear, Cosine) guarantee:
-- First ``step()`` returns ``max_lr``.
+For multi-step runs, decay schedulers (Linear, Cosine) guarantee:
+- First ``step()`` evaluates the first decay interval at
+  ``t = 1 / total_steps`` (rounding can still produce ``max_lr``).
 - ``total_steps``-th ``step()`` returns ``min_lr``.
 - Intermediate values are monotonically non-increasing and clamped to
   ``[min_lr, max_lr]``.
@@ -81,9 +82,9 @@ class TestConstantScheduler:
 class TestLinearScheduler:
     """LinearScheduler — linear decay from max_lr to min_lr."""
 
-    def test_first_step_returns_max_lr(self) -> None:
+    def test_first_step_uses_first_decay_interval(self) -> None:
         s = LinearScheduler(max_lr=10, min_lr=2, total_steps=10)
-        assert s.step() == 10
+        assert s.step() == 9
 
     def test_last_step_returns_min_lr(self) -> None:
         s = LinearScheduler(max_lr=10, min_lr=2, total_steps=10)
@@ -134,14 +135,14 @@ class TestLinearScheduler:
 
     def test_known_decay_sequence(self) -> None:
         """Linear decay max_lr=10, min_lr=2, total_steps=4.
-           t = (step-1)/(total_steps-1) = 0, 1/3, 2/3, 1
+           t = step/total_steps = 1/4, 1/2, 3/4, 1
            lr = 10 + (2-10)*t = 10 - 8t
-           t=0: lr=10, t=1/3: lr≈7.33→7, t=2/3: lr≈4.67→5, t=1: lr=2
+           t=1/4: lr=8, t=1/2: lr=6, t=3/4: lr=4, t=1: lr=2
         """
         s = LinearScheduler(max_lr=10, min_lr=2, total_steps=4)
-        assert s.step() == 10
-        assert s.step() == 7
-        assert s.step() == 5
+        assert s.step() == 8
+        assert s.step() == 6
+        assert s.step() == 4
         assert s.step() == 2
 
     def test_step_state_dict_resume_consistent(self) -> None:
@@ -167,9 +168,9 @@ class TestLinearScheduler:
 class TestCosineScheduler:
     """CosineScheduler — cosine annealing from max_lr to min_lr."""
 
-    def test_first_step_returns_max_lr(self) -> None:
-        s = CosineScheduler(max_lr=10, min_lr=2, total_steps=10)
-        assert s.step() == 10
+    def test_first_step_uses_first_decay_interval(self) -> None:
+        s = CosineScheduler(max_lr=100, min_lr=0, total_steps=4)
+        assert s.step() == 85
 
     def test_last_step_returns_min_lr(self) -> None:
         s = CosineScheduler(max_lr=10, min_lr=2, total_steps=10)
@@ -218,8 +219,8 @@ class TestCosineScheduler:
 
     def test_midpoint_close_to_mean(self) -> None:
         """At the half-way neighbourhood, cosine is close to (max+min)/2.
-           total_steps=100, step=50 → t=49/99≈0.495.
-           cos(0.495π)≈0, lr ≈ (20+2)/2 = 11.
+           total_steps=100, step=50 → t=50/100=0.5.
+           cos(0.5π)=0, lr = (20+2)/2 = 11.
         """
         s = CosineScheduler(max_lr=20, min_lr=2, total_steps=100)
         for _ in range(49):
@@ -243,13 +244,17 @@ class TestCosineScheduler:
             assert s.step() == 5
 
     def test_early_steps_near_max(self) -> None:
-        """Cosine annealing stays near max_lr early on (cos(0)=1)."""
+        """Cosine annealing stays near max_lr early on."""
         s = CosineScheduler(max_lr=100, min_lr=0, total_steps=100)
-        # step 1: t=0, cos(0)=1 → lr=100
+        # step 1: t=0.01, cos(0.01π)≈0.9995 → lr≈99.98 → 100
         assert s.step() == 100
-        # step 2: t=1/99≈0.01, cos≈0.9995 → lr≈99.97 → 100
+        # step 2: t=0.02, cos(0.02π)≈0.9980 → lr≈99.90 → 100
         assert s.step() == 100
 
+    def test_known_decay_sequence(self) -> None:
+        s = CosineScheduler(max_lr=10, min_lr=2, total_steps=4)
+        assert [s.step() for _ in range(4)] == [9, 6, 3, 2]
+
     def test_late_steps_near_min(self) -> None:
         """Cosine annealing flattens near min_lr at the end (cos(π)=-1)."""
         s = CosineScheduler(max_lr=100, min_lr=0, total_steps=100)