Remove inferring reference point in Scheduler for GSS (#2383)

Summary:
Pull Request resolved: #2383

We should be inferring the reference point in the GSS when needed. This moves the call to infer the reference point to inside of the GSS.

Reviewed By: Balandat

Differential Revision: D56227773

fbshipit-source-id: 6a63844fd3d083b12b4f693a6f9a66b9fd89b00b

Author: sdaulton

ax/global_stopping/strategies/improvement.py

@@ -22,8 +22,10 @@
 from ax.core.types import ComparisonOp
 from ax.global_stopping.strategies.base import BaseGlobalStoppingStrategy
 from ax.modelbridge.modelbridge_utils import observed_hypervolume
-from ax.plot.pareto_utils import get_tensor_converter_model
-from ax.service.utils.best_point import fill_missing_thresholds_from_nadir
+from ax.plot.pareto_utils import (
+    get_tensor_converter_model,
+    infer_reference_point_from_experiment,
+)
 from ax.utils.common.logger import get_logger
 from ax.utils.common.typeutils import checked_cast, not_none
 
@@ -78,6 +80,7 @@ def __init__(
         self.window_size = window_size
         self.improvement_bar = improvement_bar
         self.hv_by_trial: Dict[int, float] = {}
+        self._inferred_objective_thresholds: Optional[List[ObjectiveThreshold]] = None
 
     def _should_stop_optimization(
         self,
@@ -104,6 +107,9 @@ def _should_stop_optimization(
                 when computing hv of the pareto front against. This is used only in the
                 MOO setting. If not specified, the objective thresholds on the
                 experiment's optimization config will be used for the purpose.
+                If no thresholds are provided, they are automatically inferred. They are
+                only inferred once for each instance of the strategy (i.e. inferred
+                thresholds don't update with additional data).
             kwargs: Unused.
 
         Returns:
@@ -138,10 +144,25 @@ def _should_stop_optimization(
             return stop, message
 
         if isinstance(experiment.optimization_config, MultiObjectiveOptimizationConfig):
+            if objective_thresholds is None:
+                # self._inferred_objective_thresholds is cached and only computed once.
+                if self._inferred_objective_thresholds is None:
+                    # only infer reference point if there is data on the experiment.
+                    data = experiment.fetch_data()
+                    if not data.df.empty:
+                        # We infer the nadir reference point to be used by the GSS.
+                        self._inferred_objective_thresholds = (
+                            infer_reference_point_from_experiment(
+                                experiment=experiment, data=data
+                            )
+                        )
+                # TODO: move this out into a separate infer_objective_thresholds
+                # instance method or property that handles the caching.
+                objective_thresholds = self._inferred_objective_thresholds
             return self._should_stop_moo(
                 experiment=experiment,
                 trial_to_check=trial_to_check,
-                objective_thresholds=objective_thresholds,
+                objective_thresholds=not_none(objective_thresholds),
             )
         else:
             return self._should_stop_single_objective(
@@ -152,7 +173,7 @@ def _should_stop_moo(
         self,
         experiment: Experiment,
         trial_to_check: int,
-        objective_thresholds: Optional[List[ObjectiveThreshold]] = None,
+        objective_thresholds: List[ObjectiveThreshold],
     ) -> Tuple[bool, str]:
         """
         This is the "should_stop_optimization" method of this class, specialized
@@ -186,13 +207,6 @@ def _should_stop_moo(
         data_df_reference = data_df[data_df["trial_index"] <= reference_trial_index]
         data_df = data_df[data_df["trial_index"] <= trial_to_check]
 
-        optimization_config = checked_cast(
-            MultiObjectiveOptimizationConfig, experiment.optimization_config
-        ).clone_with_args(objective_thresholds=objective_thresholds)
-        objective_thresholds = fill_missing_thresholds_from_nadir(
-            experiment=experiment, optimization_config=optimization_config
-        )
-
         # Computing or retrieving HV at "window_size" iteration before
         if reference_trial_index in self.hv_by_trial:
             hv_reference = self.hv_by_trial[reference_trial_index]

ax/global_stopping/tests/test_strategies.py

@@ -269,7 +269,7 @@ def test_multi_objective(self) -> None:
         ]
         exp = self._create_multi_objective_experiment(metric_values=metric_values)
         gss = ImprovementGlobalStoppingStrategy(
-            min_trials=3, window_size=3, improvement_bar=0.1
+            min_trials=3, window_size=3, improvement_bar=0.3
         )
         stop, message = gss.should_stop_optimization(experiment=exp, trial_to_check=4)
         self.assertFalse(stop)
@@ -280,9 +280,9 @@ def test_multi_objective(self) -> None:
 
         self.assertEqual(
             message,
-            "The improvement in hypervolume in the past 3 trials (=0.000) is "
-            "less than improvement_bar (=0.1) times the hypervolume at the "
-            "start of the window (=0.055).",
+            "The improvement in hypervolume in the past 3 trials (=0.289) is "
+            "less than improvement_bar (=0.3) times the hypervolume at the "
+            "start of the window (=0.104).",
         )
 
         # Now we select a very far custom reference point against which the pareto front

ax/plot/pareto_utils.py

@@ -6,24 +6,20 @@
 
 # pyre-strict
 
-import copy
 from copy import deepcopy
 from itertools import combinations
 from logging import Logger
-from typing import cast, Dict, List, NamedTuple, Optional, Tuple, Union
+from typing import Dict, List, NamedTuple, Optional, Tuple, Union
 
 import numpy as np
 import torch
 from ax.core.batch_trial import BatchTrial
 from ax.core.data import Data
 from ax.core.experiment import Experiment
 from ax.core.metric import Metric
-from ax.core.objective import ScalarizedObjective
+from ax.core.objective import MultiObjective, ScalarizedObjective
 from ax.core.observation import ObservationFeatures
-from ax.core.optimization_config import (
-    MultiObjectiveOptimizationConfig,
-    OptimizationConfig,
-)
+from ax.core.optimization_config import MultiObjectiveOptimizationConfig
 from ax.core.outcome_constraint import (
     ComparisonOp,
     ObjectiveThreshold,
@@ -43,6 +39,7 @@
 from ax.models.torch.posterior_mean import get_PosteriorMean
 from ax.models.torch_base import TorchModel
 from ax.utils.common.logger import get_logger
+from ax.utils.common.typeutils import checked_cast
 from ax.utils.stats.statstools import relativize
 from botorch.utils.multi_objective import is_non_dominated
 from botorch.utils.multi_objective.hypervolume import infer_reference_point
@@ -615,11 +612,24 @@ def infer_reference_point_from_experiment(
     # when calculating the Pareto front. Also, defining a multiplier to turn all
     # the objectives to be maximized. Note that the multiplier at this point
     # contains 0 for outcome_constraint metrics, but this will be dropped later.
-    dummy_rp = copy.deepcopy(
-        experiment.optimization_config.objective_thresholds  # pyre-ignore
+    opt_config = checked_cast(
+        MultiObjectiveOptimizationConfig, experiment.optimization_config
     )
+    inferred_rp = _get_objective_thresholds(optimization_config=opt_config)
     multiplier = [0] * len(objective_orders)
-    for ot in dummy_rp:
+    if len(opt_config.objective_thresholds) > 0:
+        inferred_rp = deepcopy(opt_config.objective_thresholds)
+    else:
+        inferred_rp = []
+        for objective in checked_cast(MultiObjective, opt_config.objective).objectives:
+            ot = ObjectiveThreshold(
+                metric=objective.metric,
+                bound=0.0,  # dummy value
+                op=ComparisonOp.LEQ if objective.minimize else ComparisonOp.GEQ,
+                relative=False,
+            )
+            inferred_rp.append(ot)
+    for ot in inferred_rp:
         # In the following, we find the index of the objective in
         # `objective_orders`. If there is an objective that does not exist
         # in `obs_data`, a ValueError is raised.
@@ -640,12 +650,10 @@ def infer_reference_point_from_experiment(
         modelbridge=mb_reference,
         observation_features=obs_feats,
         observation_data=obs_data,
-        objective_thresholds=dummy_rp,
+        objective_thresholds=inferred_rp,
         use_model_predictions=False,
     )
-
     if len(frontier_observations) == 0:
-        opt_config = cast(OptimizationConfig, mb_reference._optimization_config)
         outcome_constraints = opt_config._outcome_constraints
         if len(outcome_constraints) == 0:
             raise RuntimeError(
@@ -665,10 +673,11 @@ def infer_reference_point_from_experiment(
             modelbridge=mb_reference,
             observation_features=obs_feats,
             observation_data=obs_data,
-            objective_thresholds=dummy_rp,
+            objective_thresholds=inferred_rp,
             use_model_predictions=False,
         )
-        opt_config._outcome_constraints = outcome_constraints  # restoring constraints
+        # restoring constraints
+        opt_config._outcome_constraints = outcome_constraints
 
     # Need to reshuffle columns of `f` and `obj_w` to be consistent
     # with objective_orders.
@@ -698,15 +707,38 @@ def infer_reference_point_from_experiment(
         x for (i, x) in enumerate(objective_orders) if multiplier[i] != 0
     ]
 
-    # Constructing the objective thresholds.
-    # NOTE: This assumes that objective_thresholds is already initialized.
-    nadir_objective_thresholds = copy.deepcopy(
-        experiment.optimization_config.objective_thresholds
-    )
-
-    for obj_threshold in nadir_objective_thresholds:
+    for obj_threshold in inferred_rp:
         obj_threshold.bound = rp[
             objective_orders_reduced.index(obj_threshold.metric.name)
         ].item()
+    return inferred_rp
 
-    return nadir_objective_thresholds
+
+def _get_objective_thresholds(
+    optimization_config: MultiObjectiveOptimizationConfig,
+) -> List[ObjectiveThreshold]:
+    """Get objective thresholds for an optimization config.
+
+    This will return objective thresholds with dummy values if there are
+    no objective thresholds on the optimization config.
+
+    Args:
+        optimization_config: Optimization config.
+
+    Returns:
+        List of objective thresholds.
+    """
+    if optimization_config.objective_thresholds is not None:
+        return deepcopy(optimization_config.objective_thresholds)
+    objective_thresholds = []
+    for objective in checked_cast(
+        MultiObjective, optimization_config.objective
+    ).objectives:
+        ot = ObjectiveThreshold(
+            metric=objective.metric,
+            bound=0.0,  # dummy value
+            op=ComparisonOp.LEQ if objective.minimize else ComparisonOp.GEQ,
+            relative=False,
+        )
+        objective_thresholds.append(ot)
+    return objective_thresholds

ax/service/scheduler.py

@@ -41,7 +41,6 @@
     MultiObjectiveOptimizationConfig,
     OptimizationConfig,
 )
-from ax.core.outcome_constraint import ObjectiveThreshold
 from ax.core.runner import Runner
 from ax.core.types import TModelPredictArm, TParameterization
 
@@ -59,7 +58,6 @@
 )
 from ax.modelbridge.base import ModelBridge
 from ax.modelbridge.generation_strategy import GenerationStrategy
-from ax.plot.pareto_utils import infer_reference_point_from_experiment
 from ax.service.utils.best_point_mixin import BestPointMixin
 from ax.service.utils.scheduler_options import SchedulerOptions, TrialType
 from ax.service.utils.with_db_settings_base import DBSettings, WithDBSettingsBase
@@ -202,9 +200,6 @@ class Scheduler(WithDBSettingsBase, BestPointMixin):
     # applications where the user wants to run the optimization loop to exhaust
     # the declared number of trials.
     __ignore_global_stopping_strategy: bool = False
-    # In MOO cases, the following will be populated by an inferred reference point
-    # for pareto front after a certain number of completed trials.
-    __inferred_reference_point: Optional[List[ObjectiveThreshold]] = None
     # Default kwargs passed when fetching data if not overridden on `SchedulerOptions`
     DEFAULT_FETCH_KWARGS = {
         "overwrite_existing_data": True,
@@ -444,26 +439,8 @@ def completion_criterion(self) -> Tuple[bool, str]:
             and self.options.global_stopping_strategy is not None
         ):
             gss = not_none(self.options.global_stopping_strategy)
-            if (
-                self.experiment.is_moo_problem
-                and self.__inferred_reference_point is None
-                and len(self.experiment.trials_by_status[TrialStatus.COMPLETED])
-                >= gss.min_trials
-            ):
-                # only infer reference point if there is data on the experiment.
-                data = self.experiment.fetch_data()
-                if not data.df.empty:
-                    # We infer the nadir reference point to be used by the GSS.
-                    self.__inferred_reference_point = (
-                        infer_reference_point_from_experiment(
-                            self.experiment,
-                            data=data,
-                        )
-                    )
-
             stop_optimization, global_stopping_msg = gss.should_stop_optimization(
-                experiment=self.experiment,
-                objective_thresholds=self.__inferred_reference_point,
+                experiment=self.experiment
             )
             if stop_optimization:
                 return True, global_stopping_msg

ax/service/tests/scheduler_test_utils.py

@@ -793,67 +793,6 @@ def test_run_preattached_trials_only(self) -> None:
             all(t.completed_successfully for t in scheduler.experiment.trials.values())
         )
 
-    def test_inferring_reference_point(self) -> None:
-        init_test_engine_and_session_factory(force_init=True)
-        experiment = get_branin_experiment_with_multi_objective()
-        experiment.runner = self.runner
-        gs = self._get_generation_strategy_strategy_for_test(
-            experiment=experiment,
-            generation_strategy=self.sobol_GS_no_parallelism,
-        )
-
-        scheduler = Scheduler(
-            experiment=experiment,
-            generation_strategy=gs,
-            options=SchedulerOptions(
-                # Stops the optimization after 5 trials.
-                global_stopping_strategy=DummyGlobalStoppingStrategy(
-                    min_trials=2, trial_to_stop=5
-                ),
-            ),
-            db_settings=self.db_settings,
-        )
-
-        with patch(
-            "ax.service.scheduler.infer_reference_point_from_experiment"
-        ) as mock_infer_rp:
-            scheduler.run_n_trials(max_trials=10)
-            mock_infer_rp.assert_called_once()
-
-    def test_inferring_reference_point_no_data(self) -> None:
-        init_test_engine_and_session_factory(force_init=True)
-        experiment = get_branin_experiment_with_multi_objective()
-        experiment.runner = self.runner
-        gs = self._get_generation_strategy_strategy_for_test(
-            experiment=experiment,
-            generation_strategy=self.sobol_GS_no_parallelism,
-        )
-
-        scheduler = Scheduler(
-            experiment=experiment,
-            generation_strategy=gs,
-            options=SchedulerOptions(
-                # Stops the optimization after 5 trials.
-                global_stopping_strategy=DummyGlobalStoppingStrategy(
-                    min_trials=0,
-                    trial_to_stop=5,
-                ),
-            ),
-            db_settings=self.db_settings,
-        )
-        empty_data = Data(
-            df=pd.DataFrame(
-                columns=["metric_name", "arm_name", "trial_index", "mean", "sem"]
-            )
-        )
-        with patch(
-            "ax.service.scheduler.infer_reference_point_from_experiment"
-        ) as mock_infer_rp, patch.object(
-            scheduler.experiment, "fetch_data", return_value=empty_data
-        ):
-            scheduler.run_n_trials(max_trials=1)
-            mock_infer_rp.assert_not_called()
-
     def test_global_stopping(self) -> None:
         gs = self._get_generation_strategy_strategy_for_test(
             experiment=self.branin_experiment,