Added reset method to StoppingCriterion (#2927)

Summary:
Pull Request resolved: #2927

A custom StoppingCriterion currently does not work as intended when multiple restarts are used in fit_gpytorch_mll_(torch).

Due to its preservation of state, the previous fitting attempt's data will be saved, and generally, the criterion will immediately be satisfied on the new attempt, thus stopping fitting on the first iteration.

StoppingCriterion is now a Protocol with methods `__call__` and `reset`.

Reviewed By: Balandat

Differential Revision: D78343624

fbshipit-source-id: 6897a9c3f50f47fa7db4e9617e51438afd56d7b8

Author: Carl Hvarfner

botorch/generation/gen.py

@@ -601,7 +601,7 @@ def assign_grad():
             return loss
 
         _optimizer.step(assign_grad)
-        stop = stopping_criterion.evaluate(fvals=loss.detach())
+        stop = stopping_criterion(fvals=loss.detach())
         if timeout_sec is not None:
             runtime = time.monotonic() - start_time
             if runtime > timeout_sec:

botorch/optim/core.py

@@ -20,6 +20,7 @@
 import numpy.typing as npt
 
 from botorch.optim.closures import NdarrayOptimizationClosure
+from botorch.optim.stopping import StoppingCriterion
 from botorch.optim.utils.numpy_utils import get_bounds_as_ndarray
 from botorch.optim.utils.timeout import minimize_with_timeout
 from numpy import asarray, float64 as np_float64
@@ -153,7 +154,7 @@ def torch_minimize(
     scheduler: LRScheduler | Callable[[Optimizer], LRScheduler] | None = None,
     step_limit: int | None = None,
     timeout_sec: float | None = None,
-    stopping_criterion: Callable[[Tensor], bool] | None = None,
+    stopping_criterion: StoppingCriterion | None = None,
 ) -> OptimizationResult:
     r"""Generic torch.optim-based optimization routine.
 
@@ -190,6 +191,10 @@ def torch_minimize(
     if not (scheduler is None or isinstance(scheduler, LRScheduler)):
         scheduler = scheduler(optimizer)
 
+    if stopping_criterion is not None:
+        # Reset stopping criterion to ensure clean state for new optimization run
+        stopping_criterion.reset()
+
     _bounds = (
         {}
         if bounds is None

botorch/optim/fit.py

@@ -22,7 +22,7 @@
     scipy_minimize,
     torch_minimize,
 )
-from botorch.optim.stopping import ExpMAStoppingCriterion
+from botorch.optim.stopping import ExpMAStoppingCriterion, StoppingCriterion
 from botorch.optim.utils import get_parameters_and_bounds, TorchAttr
 from botorch.utils.types import DEFAULT
 from gpytorch.mlls.marginal_log_likelihood import MarginalLogLikelihood
@@ -118,7 +118,7 @@ def fit_gpytorch_mll_torch(
     closure: Callable[[], tuple[Tensor, Sequence[Tensor | None]]] | None = None,
     closure_kwargs: dict[str, Any] | None = None,
     step_limit: int | None = None,
-    stopping_criterion: Callable[[Tensor], bool] | None = DEFAULT,  # pyre-ignore [9]
+    stopping_criterion: StoppingCriterion | None = DEFAULT,
     optimizer: Optimizer | Callable[..., Optimizer] = Adam,
     scheduler: _LRScheduler | Callable[..., _LRScheduler] | None = None,
     callback: Callable[[dict[str, Tensor], OptimizationResult], None] | None = None,

botorch/optim/optimize.py

@@ -826,7 +826,7 @@ def optimize_acqf_cyclic(
     if q > 1:
         cyclic_options = cyclic_options or {}
         stopping_criterion = ExpMAStoppingCriterion(**cyclic_options)
-        stop = stopping_criterion.evaluate(fvals=acq_vals)
+        stop = stopping_criterion(fvals=acq_vals)
         base_X_pending = acq_function.X_pending
         idxr = torch.ones(q, dtype=torch.bool, device=opt_inputs.bounds.device)
         while not stop:
@@ -847,7 +847,7 @@ def optimize_acqf_cyclic(
                 candidates[i] = candidate_i
                 acq_vals[i] = acq_val_i
                 idxr[i] = 1
-            stop = stopping_criterion.evaluate(fvals=acq_vals)
+            stop = stopping_criterion(fvals=acq_vals)
         acq_function.set_X_pending(base_X_pending)
     return candidates, acq_vals
 

botorch/optim/stopping.py

@@ -6,21 +6,20 @@
 
 from __future__ import annotations
 
-from abc import ABC, abstractmethod
+from typing import Protocol
 
 import torch
 from torch import Tensor
 
 
-class StoppingCriterion(ABC):
-    r"""Base class for evaluating optimization convergence.
+class StoppingCriterion(Protocol):
+    r"""Protocol for evaluating optimization convergence.
 
-    Stopping criteria are implemented as a objects rather than a function, so that they
+    Stopping criteria are implemented as objects rather than functions, so that they
     can keep track of past function values between optimization steps.
     """
 
-    @abstractmethod
-    def evaluate(self, fvals: Tensor) -> bool:
+    def __call__(self, fvals: Tensor) -> bool:
         r"""Evaluate the stopping criterion.
 
         Args:
@@ -30,15 +29,20 @@ def evaluate(self, fvals: Tensor) -> bool:
                 true for all elements.
 
         Returns:
-            Stopping indicator (if True, stop the optimziation).
+            Stopping indicator (if True, stop the optimization).
         """
-        pass  # pragma: no cover
+        ...  # pragma: no cover
 
-    def __call__(self, fvals: Tensor) -> bool:
-        return self.evaluate(fvals)
+    def reset(self) -> None:
+        r"""Reset the stopping criterion to its initial state.
+
+        This method should be called before starting a new optimization run
+        to ensure that any internal state from previous runs is cleared.
+        """
+        ...  # pragma: no cover
 
 
-class ExpMAStoppingCriterion(StoppingCriterion):
+class ExpMAStoppingCriterion:
     r"""Exponential moving average stopping criterion.
 
     Computes an exponentially weighted moving average over window length `n_window`
@@ -80,7 +84,7 @@ def __init__(
         self.weights = weights / weights.sum()
         self._prev_fvals = None
 
-    def evaluate(self, fvals: Tensor) -> bool:
+    def __call__(self, fvals: Tensor) -> bool:
         r"""Evaluate the stopping criterion.
 
         Args:
@@ -125,3 +129,11 @@ def evaluate(self, fvals: Tensor) -> bool:
             return True
 
         return False
+
+    def reset(self) -> None:
+        r"""Reset the stopping criterion to its initial state.
+
+        Resets the iteration counter and clears any stored function values.
+        """
+        self.iter = 0
+        self._prev_fvals = None

test/optim/test_core.py

@@ -17,6 +17,7 @@
     scipy_minimize,
     torch_minimize,
 )
+from botorch.optim.stopping import ExpMAStoppingCriterion
 from botorch.utils.testing import BotorchTestCase
 from numpy import allclose
 from scipy.optimize import OptimizeResult
@@ -254,11 +255,11 @@ def _callback(parameters, result, out) -> None:
             self.assertEqual(result.step, len(step_results))
 
             # Test `stopping_criterion`
-            stopping_decisions = iter((False, False, True, False))
+            max3_stopping_criterion = ExpMAStoppingCriterion(maxiter=3, n_window=5)
             result = torch_minimize(
                 closure=closure,
                 parameters=closure.parameters,
-                stopping_criterion=lambda fval: next(stopping_decisions),
+                stopping_criterion=max3_stopping_criterion,
             )
             self.assertEqual(result.step, 3)
             self.assertEqual(result.status, OptimizationStatus.STOPPED)

test/optim/test_stopping.py

@@ -7,15 +7,11 @@
 from __future__ import annotations
 
 import torch
-from botorch.optim.stopping import ExpMAStoppingCriterion, StoppingCriterion
+from botorch.optim.stopping import ExpMAStoppingCriterion
 from botorch.utils.testing import BotorchTestCase
 
 
 class TestStoppingCriterion(BotorchTestCase):
-    def test_abstract_raises(self):
-        with self.assertRaises(TypeError):
-            StoppingCriterion()
-
     def test_exponential_moving_average(self):
         for dtype in (torch.float, torch.double):
             tkwargs = {"device": self.device, "dtype": dtype}
@@ -25,8 +21,8 @@ def test_exponential_moving_average(self):
             self.assertEqual(sc.maxiter, 2)
             self.assertEqual(sc.n_window, 10)
             self.assertEqual(sc.rel_tol, 1e-5)
-            self.assertFalse(sc.evaluate(fvals=torch.ones(1, **tkwargs)))
-            self.assertTrue(sc.evaluate(fvals=torch.zeros(1, **tkwargs)))
+            self.assertFalse(sc(fvals=torch.ones(1, **tkwargs)))
+            self.assertTrue(sc(fvals=torch.zeros(1, **tkwargs)))
 
             # test convergence
             n_window = 4
@@ -43,7 +39,7 @@ def test_exponential_moving_average(self):
                 if not minimize:
                     f_vals = -f_vals
                 for i, fval in enumerate(f_vals):
-                    if sc.evaluate(fval):
+                    if sc(fval):
                         self.assertEqual(i, 10)
                         break
                 # test multiple components
@@ -55,6 +51,16 @@ def test_exponential_moving_average(self):
                     df = -df
                 f_vals = torch.stack([f_vals, f_vals + df], dim=-1)
                 for i, fval in enumerate(f_vals):
-                    if sc.evaluate(fval):
+                    if sc(fval):
                         self.assertEqual(i, 10)
                         break
+
+                # Test reset functionality - verify state after use, reset, and reuse
+                self.assertGreater(sc.iter, 0)
+                self.assertIsNotNone(sc._prev_fvals)
+                sc.reset()
+                self.assertEqual(sc.iter, 0)
+                self.assertIsNone(sc._prev_fvals)
+                # Verify criterion works after reset
+                self.assertFalse(sc(f_vals[0]))
+                self.assertEqual(sc.iter, 1)