From 3c565a7f848c919ccfdbfa6b9f36ac651aab2a5a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=BB=84=E4=BA=91=E9=BE=99?=
 <76432572+nankingjing@users.noreply.github.com>
Date: Sat, 11 Jul 2026 15:06:11 +0800
Subject: [PATCH 1/3] test: add unit tests for skillopt.optimizer.scheduler

---
 tests/test_scheduler.py | 329 ++++++++++++++++++++++++++++++++++++++++
 1 file changed, 329 insertions(+)
 create mode 100644 tests/test_scheduler.py

diff --git a/tests/test_scheduler.py b/tests/test_scheduler.py
new file mode 100644
index 0000000..31f39c1
--- /dev/null
+++ b/tests/test_scheduler.py
@@ -0,0 +1,329 @@
+"""Tests for skillopt.optimizer.scheduler — edit-budget schedulers.
+
+ReflACT trainers use an edit-budget scheduler at each optimisation step to
+control how many skill edits are allowed (analogous to gradient clipping /
+learning-rate annealing in neural-network training).
+
+This module has zero LLM dependencies — all behaviour is deterministic pure
+math — making it an ideal target for precise unit tests.
+"""
+from __future__ import annotations
+
+import pytest
+
+from skillopt.optimizer.scheduler import (
+    LRScheduler,
+    ConstantScheduler,
+    LinearScheduler,
+    CosineScheduler,
+    AutonomousScheduler,
+    build_scheduler,
+)
+
+
+# ── ConstantScheduler ────────────────────────────────────────────────────────
+
+
+class TestConstantScheduler:
+    """ConstantScheduler — fixed edit budget regardless of step."""
+
+    def test_always_returns_max_lr(self) -> None:
+        s = ConstantScheduler(max_lr=8, min_lr=2, total_steps=10)
+        for _ in range(10):
+            assert s.step() == 8
+
+    def test_step_advances_internal_counter(self) -> None:
+        s = ConstantScheduler(max_lr=5, min_lr=1, total_steps=5)
+        assert s._current_step == 0
+        s.step()
+        assert s._current_step == 1
+        s.step()
+        assert s._current_step == 2
+
+    def test_get_lr_returns_max_for_arbitrary_step(self) -> None:
+        s = ConstantScheduler(max_lr=12, min_lr=1, total_steps=100)
+        assert s.get_lr(1) == 12
+        assert s.get_lr(50) == 12
+        assert s.get_lr(999) == 12
+
+    def test_state_dict_and_load_state_dict_round_trip(self) -> None:
+        s = ConstantScheduler(max_lr=8, min_lr=2, total_steps=10)
+        for _ in range(3):
+            s.step()
+        assert s._current_step == 3
+
+        state = s.state_dict()
+        s2 = ConstantScheduler(max_lr=8, min_lr=2, total_steps=10)
+        s2.load_state_dict(state)
+        assert s2._current_step == 3
+        # Step after resume lands on the correct step
+        assert s2.step() == 8
+        assert s2._current_step == 4
+
+    def test_load_state_dict_with_missing_key_defaults_to_zero(self) -> None:
+        s = ConstantScheduler(max_lr=8, min_lr=2, total_steps=10)
+        s.load_state_dict({})
+        assert s._current_step == 0
+
+
+# ── LinearScheduler ──────────────────────────────────────────────────────────
+
+
+class TestLinearScheduler:
+    """LinearScheduler — linear decay from max_lr to min_lr."""
+
+    def test_first_step_returns_max_lr(self) -> None:
+        s = LinearScheduler(max_lr=10, min_lr=2, total_steps=10)
+        assert s.step() == 10
+
+    def test_last_step_returns_min_lr(self) -> None:
+        s = LinearScheduler(max_lr=10, min_lr=2, total_steps=10)
+        for _ in range(9):
+            s.step()
+        assert s.step() == 2
+
+    def test_all_steps_return_integers(self) -> None:
+        s = LinearScheduler(max_lr=8, min_lr=2, total_steps=6)
+        for _ in range(6):
+            lr = s.step()
+            assert isinstance(lr, int)
+
+    def test_total_steps_one_returns_max_lr(self) -> None:
+        s = LinearScheduler(max_lr=10, min_lr=2, total_steps=1)
+        assert s.step() == 10
+
+    def test_total_steps_zero_returns_max_lr(self) -> None:
+        """Degenerate case: 0-step training still gets max_lr on the one call."""
+        s = LinearScheduler(max_lr=10, min_lr=2, total_steps=0)
+        assert s.step() == 10
+
+    def test_monotonically_non_increasing(self) -> None:
+        s = LinearScheduler(max_lr=20, min_lr=2, total_steps=100)
+        prev: int = 999
+        for _ in range(100):
+            lr = s.step()
+            assert lr <= prev
+            prev = lr
+
+    def test_never_below_min_lr(self) -> None:
+        s = LinearScheduler(max_lr=10, min_lr=2, total_steps=10)
+        for _ in range(20):  # overshoot
+            assert s.step() >= 2
+
+    def test_never_above_max_lr(self) -> None:
+        s = LinearScheduler(max_lr=10, min_lr=2, total_steps=10)
+        for _ in range(20):
+            assert s.step() <= 10
+
+    def test_after_total_steps_stays_at_min_lr(self) -> None:
+        s = LinearScheduler(max_lr=8, min_lr=2, total_steps=5)
+        for _ in range(5):
+            s.step()
+        # Steps beyond total_steps should plateau at min_lr
+        for _ in range(5):
+            assert s.step() == 2
+
+    def test_known_decay_sequence(self) -> None:
+        """Linear decay max_lr=10, min_lr=2, total_steps=4 (t=0.25,0.50,0.75,1.0).
+           lr = 10 + (2-10)*t = 10 - 8t
+           t=0.25: lr=8.0, t=0.50: lr=6.0, t=0.75: lr=4.0, t=1.0: lr=2.0
+        """
+        s = LinearScheduler(max_lr=10, min_lr=2, total_steps=4)
+        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:
+        """After resume, the next step value is the same as without resume."""
+        s1 = LinearScheduler(max_lr=10, min_lr=2, total_steps=5)
+        for _ in range(3):
+            s1.step()
+        resumed_lr = s1.step()  # step 4
+
+        s2 = LinearScheduler(max_lr=10, min_lr=2, total_steps=5)
+        s2.load_state_dict({"current_step": 3})
+        assert s2.step() == resumed_lr
+
+    def test_max_lr_equals_min_lr_yields_constant(self) -> None:
+        s = LinearScheduler(max_lr=5, min_lr=5, total_steps=10)
+        for _ in range(10):
+            assert s.step() == 5
+
+
+# ── CosineScheduler ──────────────────────────────────────────────────────────
+
+
+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_last_step_returns_min_lr(self) -> None:
+        s = CosineScheduler(max_lr=10, min_lr=2, total_steps=10)
+        for _ in range(9):
+            s.step()
+        assert s.step() == 2
+
+    def test_all_steps_return_integers(self) -> None:
+        s = CosineScheduler(max_lr=8, min_lr=2, total_steps=6)
+        for _ in range(6):
+            lr = s.step()
+            assert isinstance(lr, int)
+
+    def test_total_steps_one_returns_max_lr(self) -> None:
+        s = CosineScheduler(max_lr=10, min_lr=2, total_steps=1)
+        assert s.step() == 10
+
+    def test_total_steps_zero_returns_max_lr(self) -> None:
+        s = CosineScheduler(max_lr=10, min_lr=2, total_steps=0)
+        assert s.step() == 10
+
+    def test_monotonically_non_increasing(self) -> None:
+        s = CosineScheduler(max_lr=20, min_lr=2, total_steps=100)
+        prev: int = 999
+        for _ in range(100):
+            lr = s.step()
+            assert lr <= prev
+            prev = lr
+
+    def test_never_below_min_lr(self) -> None:
+        s = CosineScheduler(max_lr=10, min_lr=2, total_steps=10)
+        for _ in range(20):
+            assert s.step() >= 2
+
+    def test_never_above_max_lr(self) -> None:
+        s = CosineScheduler(max_lr=10, min_lr=2, total_steps=10)
+        for _ in range(20):
+            assert s.step() <= 10
+
+    def test_after_total_steps_stays_at_min_lr(self) -> None:
+        s = CosineScheduler(max_lr=8, min_lr=2, total_steps=5)
+        for _ in range(5):
+            s.step()
+        for _ in range(5):
+            assert s.step() == 2
+
+    def test_midpoint_close_to_mean(self) -> None:
+        """At the half-way point, cosine gives (max+min)/2 exactly."""
+        s = CosineScheduler(max_lr=20, min_lr=2, total_steps=100)
+        for _ in range(49):
+            s.step()
+        mid = s.step()  # step 50 / 100, t=0.5, cos(pi/2)=0
+        # lr = min + 0.5*(max-min)*(1+0) = (max+min)/2 = 11
+        assert mid == 11
+
+    def test_step_state_dict_resume_consistent(self) -> None:
+        s1 = CosineScheduler(max_lr=10, min_lr=2, total_steps=5)
+        for _ in range(3):
+            s1.step()
+        resumed_lr = s1.step()
+
+        s2 = CosineScheduler(max_lr=10, min_lr=2, total_steps=5)
+        s2.load_state_dict({"current_step": 3})
+        assert s2.step() == resumed_lr
+
+    def test_max_lr_equals_min_lr_yields_constant(self) -> None:
+        s = CosineScheduler(max_lr=5, min_lr=5, total_steps=10)
+        for _ in range(10):
+            assert s.step() == 5
+
+    def test_early_steps_near_max(self) -> None:
+        """Cosine annealing stays flat near max_lr early on (cos(0)≈1)."""
+        s = CosineScheduler(max_lr=100, min_lr=0, total_steps=100)
+        # step 1: t=0.01, cos(0.01*pi)≈0.9995 → lr≈99.97 → 100
+        assert s.step() == 100
+        # step 2: t=0.02, cos≈0.9980 → lr≈99.9 → 100
+        assert s.step() == 100
+
+    def test_late_steps_near_min(self) -> None:
+        """Cosine annealing flattens near min_lr at the end (cos(pi)≈-1)."""
+        s = CosineScheduler(max_lr=100, min_lr=0, total_steps=100)
+        for _ in range(99):
+            s.step()
+        # step 100: t=1.0, cos(pi)=-1 → lr=0
+        assert s.step() == 0
+
+
+# ── AutonomousScheduler ──────────────────────────────────────────────────────
+
+
+class TestAutonomousScheduler:
+    """AutonomousScheduler — no edit limit (model decides freely)."""
+
+    def test_always_returns_no_limit(self) -> None:
+        s = AutonomousScheduler(max_lr=8, min_lr=2, total_steps=10)
+        for _ in range(20):
+            assert s.step() == AutonomousScheduler.NO_LIMIT
+
+    def test_step_advances_counter(self) -> None:
+        s = AutonomousScheduler(max_lr=5, min_lr=1, total_steps=5)
+        assert s._current_step == 0
+        s.step()
+        assert s._current_step == 1
+
+    def test_get_lr_returns_no_limit(self) -> None:
+        s = AutonomousScheduler(max_lr=5, min_lr=1, total_steps=10)
+        assert s.get_lr(1) == AutonomousScheduler.NO_LIMIT
+        assert s.get_lr(50) == AutonomousScheduler.NO_LIMIT
+
+    def test_state_dict_round_trip(self) -> None:
+        s = AutonomousScheduler(max_lr=5, min_lr=1, total_steps=10)
+        for _ in range(4):
+            s.step()
+        s2 = AutonomousScheduler(max_lr=5, min_lr=1, total_steps=10)
+        s2.load_state_dict(s.state_dict())
+        assert s2._current_step == 4
+
+
+# ── build_scheduler factory ──────────────────────────────────────────────────
+
+
+class TestBuildScheduler:
+    """build_scheduler factory — creates the right scheduler from a mode name."""
+
+    def test_constant(self) -> None:
+        s = build_scheduler("constant", max_lr=8, min_lr=2, total_steps=10)
+        assert isinstance(s, ConstantScheduler)
+        assert s.max_lr == 8
+        assert s.min_lr == 2
+        assert s.total_steps == 10
+
+    def test_linear(self) -> None:
+        s = build_scheduler("linear", max_lr=12, min_lr=3, total_steps=20)
+        assert isinstance(s, LinearScheduler)
+
+    def test_cosine(self) -> None:
+        s = build_scheduler("cosine", max_lr=16, min_lr=4, total_steps=30)
+        assert isinstance(s, CosineScheduler)
+
+    def test_autonomous(self) -> None:
+        s = build_scheduler("autonomous", max_lr=8, min_lr=2, total_steps=10)
+        assert isinstance(s, AutonomousScheduler)
+
+    def test_unknown_mode_raises_value_error(self) -> None:
+        with pytest.raises(ValueError, match="Unknown scheduler mode"):
+            build_scheduler("exponential", max_lr=8, min_lr=2, total_steps=10)
+
+    def test_default_mode_is_constant(self) -> None:
+        s = build_scheduler(max_lr=8, min_lr=2, total_steps=10)
+        assert isinstance(s, ConstantScheduler)
+
+
+# ── Abstract base class ──────────────────────────────────────────────────────
+
+
+class TestLRSchedulerBase:
+    """LRScheduler — abstract base: cannot be instantiated directly."""
+
+    def test_cannot_instantiate_abstract(self) -> None:
+        with pytest.raises(TypeError):
+            LRScheduler(max_lr=8, min_lr=2, total_steps=10)  # type: ignore[abstract]
+
+    def test_concrete_subclass_instantiates_fine(self) -> None:
+        """ConstantScheduler (and others) work normally."""
+        s = ConstantScheduler(max_lr=8, min_lr=2, total_steps=10)
+        assert isinstance(s, LRScheduler)
+        assert s.step() == 8

From 9f9f87c0644eadcd88bf19452037472682df6f3e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=BB=84=E4=BA=91=E9=BE=99?=
 <76432572+nankingjing@users.noreply.github.com>
Date: Tue, 14 Jul 2026 00:45:12 +0800
Subject: [PATCH 2/3] fix(scheduler): use consistent (step-1)/(total_steps-1)
 endpoint formula

LinearScheduler and CosineScheduler now guarantee step() returns max_lr
on the first call and min_lr on the total_steps-th call, using the
standard endpoint parameterisation t = (step-1)/(total_steps-1).

This resolves the inconsistent first-step semantics reported in review.
---
 skillopt/optimizer/scheduler.py | 20 ++++++++++++++++----
 1 file changed, 16 insertions(+), 4 deletions(-)

diff --git a/skillopt/optimizer/scheduler.py b/skillopt/optimizer/scheduler.py
index 63b944e..416d3bb 100644
--- a/skillopt/optimizer/scheduler.py
+++ b/skillopt/optimizer/scheduler.py
@@ -61,23 +61,35 @@ class ConstantScheduler(LRScheduler):
 
 
 class LinearScheduler(LRScheduler):
-    """Linear decay from ``max_lr`` to ``min_lr`` over ``total_steps``."""
+    """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.
+    """
 
     def _compute_lr(self, step: int) -> int:
         if self.total_steps <= 1:
             return self.max_lr
-        t = min(step, self.total_steps) / self.total_steps
+        s = min(step, self.total_steps)
+        t = (s - 1) / (self.total_steps - 1)
         lr = self.max_lr + (self.min_lr - self.max_lr) * t
         return max(self.min_lr, round(lr))
 
 
 class CosineScheduler(LRScheduler):
-    """Cosine annealing from ``max_lr`` to ``min_lr`` over ``total_steps``."""
+    """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.
+    """
 
     def _compute_lr(self, step: int) -> int:
         if self.total_steps <= 1:
             return self.max_lr
-        t = min(step, self.total_steps) / self.total_steps
+        s = min(step, self.total_steps)
+        t = (s - 1) / (self.total_steps - 1)
         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))
 

From daf303542989b236584390da9cdc89ce6154162d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=BB=84=E4=BA=91=E9=BE=99?=
 <76432572+nankingjing@users.noreply.github.com>
Date: Tue, 14 Jul 2026 00:45:22 +0800
Subject: [PATCH 3/3] test(scheduler): update sequence expectation to match
 consistent endpoint formula

The LinearScheduler.known_decay_sequence test now reflects the new
t = (step-1)/(total_steps-1) contract where step 1 returns max_lr.
Updated docstrings for midpoint and early-step tests to match.

All 43 tests now pass against the updated formula.
---
 tests/test_scheduler.py | 38 +++++++++++++++++++++++++-------------
 1 file changed, 25 insertions(+), 13 deletions(-)

diff --git a/tests/test_scheduler.py b/tests/test_scheduler.py
index 31f39c1..40bf3ec 100644
--- a/tests/test_scheduler.py
+++ b/tests/test_scheduler.py
@@ -6,6 +6,15 @@ learning-rate annealing in neural-network training).
 
 This module has zero LLM dependencies — all behaviour is deterministic pure
 math — making it an ideal target for precise unit tests.
+
+Scheduler contract
+------------------
+Decay schedulers (Linear, Cosine) guarantee:
+- First ``step()`` returns ``max_lr``.
+- ``total_steps``-th ``step()`` returns ``min_lr``.
+- Intermediate values are monotonically non-increasing and clamped to
+  ``[min_lr, max_lr]``.
+- Beyond ``total_steps`` the value plateaus at ``min_lr``.
 """
 from __future__ import annotations
 
@@ -124,14 +133,15 @@ class TestLinearScheduler:
             assert s.step() == 2
 
     def test_known_decay_sequence(self) -> None:
-        """Linear decay max_lr=10, min_lr=2, total_steps=4 (t=0.25,0.50,0.75,1.0).
+        """Linear decay max_lr=10, min_lr=2, total_steps=4.
+           t = (step-1)/(total_steps-1) = 0, 1/3, 2/3, 1
            lr = 10 + (2-10)*t = 10 - 8t
-           t=0.25: lr=8.0, t=0.50: lr=6.0, t=0.75: lr=4.0, t=1.0: lr=2.0
+           t=0: lr=10, t=1/3: lr≈7.33→7, t=2/3: lr≈4.67→5, t=1: lr=2
         """
         s = LinearScheduler(max_lr=10, min_lr=2, total_steps=4)
-        assert s.step() == 8
-        assert s.step() == 6
-        assert s.step() == 4
+        assert s.step() == 10
+        assert s.step() == 7
+        assert s.step() == 5
         assert s.step() == 2
 
     def test_step_state_dict_resume_consistent(self) -> None:
@@ -207,12 +217,14 @@ class TestCosineScheduler:
             assert s.step() == 2
 
     def test_midpoint_close_to_mean(self) -> None:
-        """At the half-way point, cosine gives (max+min)/2 exactly."""
+        """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.
+        """
         s = CosineScheduler(max_lr=20, min_lr=2, total_steps=100)
         for _ in range(49):
             s.step()
-        mid = s.step()  # step 50 / 100, t=0.5, cos(pi/2)=0
-        # lr = min + 0.5*(max-min)*(1+0) = (max+min)/2 = 11
+        mid = s.step()  # step 50
         assert mid == 11
 
     def test_step_state_dict_resume_consistent(self) -> None:
@@ -231,19 +243,19 @@ class TestCosineScheduler:
             assert s.step() == 5
 
     def test_early_steps_near_max(self) -> None:
-        """Cosine annealing stays flat near max_lr early on (cos(0)≈1)."""
+        """Cosine annealing stays near max_lr early on (cos(0)=1)."""
         s = CosineScheduler(max_lr=100, min_lr=0, total_steps=100)
-        # step 1: t=0.01, cos(0.01*pi)≈0.9995 → lr≈99.97 → 100
+        # step 1: t=0, cos(0)=1 → lr=100
         assert s.step() == 100
-        # step 2: t=0.02, cos≈0.9980 → lr≈99.9 → 100
+        # step 2: t=1/99≈0.01, cos≈0.9995 → lr≈99.97 → 100
         assert s.step() == 100
 
     def test_late_steps_near_min(self) -> None:
-        """Cosine annealing flattens near min_lr at the end (cos(pi)≈-1)."""
+        """Cosine annealing flattens near min_lr at the end (cos(π)=-1)."""
         s = CosineScheduler(max_lr=100, min_lr=0, total_steps=100)
         for _ in range(99):
             s.step()
-        # step 100: t=1.0, cos(pi)=-1 → lr=0
+        # step 100: t=1, cos(π)=-1 → lr=0
         assert s.step() == 0