From 3eb73f88c7be71367e9cd73e65c811993ddbe734 Mon Sep 17 00:00:00 2001 From: Lior Belenki Date: Wed, 20 Mar 2024 18:07:22 -0700 Subject: [PATCH] Update scheduled params API PiperOrigin-RevId: 617678719 --- .../designers/scheduled_designer.py | 4 ++ .../designers/scheduled_designer_test.py | 26 ++++---- .../designers/scheduled_gp_bandit.py | 17 +++--- .../designers/scheduled_gp_ucb_pe.py | 59 +++++++++++++++---- 4 files changed, 76 insertions(+), 30 deletions(-) diff --git a/vizier/_src/algorithms/designers/scheduled_designer.py b/vizier/_src/algorithms/designers/scheduled_designer.py index 2e86d911d..ed776c3e2 100644 --- a/vizier/_src/algorithms/designers/scheduled_designer.py +++ b/vizier/_src/algorithms/designers/scheduled_designer.py @@ -138,6 +138,10 @@ def __attrs_post_init__(self): self._designer = self._designer_factory(self._problem) self._update_designer_state() + @property + def scheduled_params(self) -> dict[str, ScheduledParam]: + return self._scheduled_params + def update( self, completed: vza.CompletedTrials, all_active: vza.ActiveTrials ) -> None: diff --git a/vizier/_src/algorithms/designers/scheduled_designer_test.py b/vizier/_src/algorithms/designers/scheduled_designer_test.py index 1bd445955..667db2772 100644 --- a/vizier/_src/algorithms/designers/scheduled_designer_test.py +++ b/vizier/_src/algorithms/designers/scheduled_designer_test.py @@ -125,22 +125,22 @@ def _gp_bandit_factory(problem): scheduled_desinger = scheduled_gp_bandit.ScheduledGPBanditFactory( gp_bandit_factory=_gp_bandit_factory, - expected_total_num_trials=5, - init_ucb_coef=4.0, - final_ucb_coef=1.0, - ucb_coef_decay_rate=1.2, + expected_total_num_trials=2, + init_ucb_coefficient=4.0, + final_ucb_coefficient=1.0, + decay_ucb_coefficient=1.2, )(problem) self.assertLen( test_runners.RandomMetricsRunner( problem, - iters=5, + iters=2, batch_size=1, verbose=1, validate_parameters=True, seed=1, ).run_designer(scheduled_desinger), - 5, + 2, ) @@ -163,10 +163,16 @@ def _gp_ucb_pe_factory( scheduled_desinger = scheduled_gp_ucb_pe.ScheduledGPUCBPEFactory( gp_ucb_pe_factory=_gp_ucb_pe_factory, - expected_total_num_trials=5, - init_ucb_coef=4.0, - final_ucb_coef=1.0, - ucb_coef_decay_rate=1.2, + expected_total_num_trials=3, + init_ucb_coefficient=4.0, + final_ucb_coefficient=1.0, + decay_ucb_coefficient=1.2, + init_explore_region_ucb_coefficient=1.0, + final_explore_region_ucb_coefficient=0.5, + decay_explore_region_ucb_coefficient=1.2, + init_ucb_overwrite_probability=0.25, + final_ucb_overwrite_probability=0.0, + decay_ucb_overwrite_probability=1.0, )(problem) self.assertLen( diff --git a/vizier/_src/algorithms/designers/scheduled_gp_bandit.py b/vizier/_src/algorithms/designers/scheduled_gp_bandit.py index 03f9a273e..b45503de9 100644 --- a/vizier/_src/algorithms/designers/scheduled_gp_bandit.py +++ b/vizier/_src/algorithms/designers/scheduled_gp_bandit.py @@ -30,11 +30,10 @@ class ScheduledGPBanditFactory: """Scheduled GP-Bandit factory.""" _gp_bandit_factory: Callable[[vz.ProblemStatement], gp_bandit.VizierGPBandit] - _init_ucb_coef: float - _final_ucb_coef: float - _ucb_coef_decay_rate: float + _init_ucb_coefficient: float + _final_ucb_coefficient: float + _decay_ucb_coefficient: float _expected_total_num_trials: int - _ucb_coef_param_name: str = attrs.field(default='ucb_coef', init=False) def __call__( self, @@ -45,20 +44,20 @@ def __call__( def _gp_bandit_state_updater(designer, params): designer._scoring_function_factory = ( # pylint: disable=protected-access acquisitions.bayesian_scoring_function_factory( - lambda _: acquisitions.UCB(params[self._ucb_coef_param_name]) + lambda _: acquisitions.UCB(params['ucb_coefficient']) ) ) ucb_coef_param = scheduled_designer.ExponentialScheduledParam( - init_value=self._init_ucb_coef, - final_value=self._final_ucb_coef, - rate=self._ucb_coef_decay_rate, + init_value=self._init_ucb_coefficient, + final_value=self._final_ucb_coefficient, + rate=self._decay_ucb_coefficient, ) return scheduled_designer.ScheduledDesigner( problem, designer_factory=self._gp_bandit_factory, designer_state_updater=_gp_bandit_state_updater, - scheduled_params={self._ucb_coef_param_name: ucb_coef_param}, + scheduled_params={'ucb_coefficient': ucb_coef_param}, expected_total_num_trials=self._expected_total_num_trials, ) diff --git a/vizier/_src/algorithms/designers/scheduled_gp_ucb_pe.py b/vizier/_src/algorithms/designers/scheduled_gp_ucb_pe.py index 9bdf3dac5..59871dcd3 100644 --- a/vizier/_src/algorithms/designers/scheduled_gp_ucb_pe.py +++ b/vizier/_src/algorithms/designers/scheduled_gp_ucb_pe.py @@ -25,7 +25,6 @@ from vizier._src.algorithms.designers import scheduled_designer -# TODO: Add scheduling to more params. @attrs.define(kw_only=True) class ScheduledGPUCBPEFactory: """Scheduled GP-UCB-PE factory.""" @@ -33,11 +32,21 @@ class ScheduledGPUCBPEFactory: _gp_ucb_pe_factory: Callable[ [vz.ProblemStatement], gp_ucb_pe.VizierGPUCBPEBandit ] - _init_ucb_coef: float - _final_ucb_coef: float - _ucb_coef_decay_rate: float _expected_total_num_trials: int - _ucb_coef_param_name: str = attrs.field(default='ucb_coef', init=False) + # ucb_coefficient + _init_ucb_coefficient: float + _final_ucb_coefficient: float + _decay_ucb_coefficient: float + + # explore_region_ucb_coefficient + _init_explore_region_ucb_coefficient: float + _final_explore_region_ucb_coefficient: float + _decay_explore_region_ucb_coefficient: float + + # ucb_overwrite_probability + _init_ucb_overwrite_probability: float + _final_ucb_overwrite_probability: float + _decay_ucb_overwrite_probability: float def __call__( self, @@ -48,22 +57,50 @@ def __call__( def _gp_ucb_pe_state_updater(designer, params): # Create a new copy of the config and replace the ucb coefficient value. updated_config = dc.replace( - designer._config, ucb_coefficient=params[self._ucb_coef_param_name] # pylint: disable=protected-access + designer._config, # pylint: disable=protected-access + ucb_coefficient=params['ucb_coefficient'], + explore_region_ucb_coefficient=params[ + 'explore_region_ucb_coefficient' + ], + ucb_overwrite_probability=params['ucb_overwrite_probability'], ) # Update the designer config. As the GP_UCB_PE state is recomputed every # suggest() call based on the config, this effectively updates the state. + # TODO: Instead of changing internal state, use public method. designer._config = updated_config # pylint: disable=protected-access - ucb_coef_param = scheduled_designer.ExponentialScheduledParam( - init_value=self._init_ucb_coef, - final_value=self._final_ucb_coef, - rate=self._ucb_coef_decay_rate, + ucb_coefficient_param = scheduled_designer.ExponentialScheduledParam( + init_value=self._init_ucb_coefficient, + final_value=self._final_ucb_coefficient, + rate=self._decay_ucb_coefficient, + ) + + explore_region_ucb_coefficient_param = ( + scheduled_designer.ExponentialScheduledParam( + init_value=self._init_explore_region_ucb_coefficient, + final_value=self._final_explore_region_ucb_coefficient, + rate=self._decay_explore_region_ucb_coefficient, + ) ) + ucb_overwrite_probability_param = ( + scheduled_designer.ExponentialScheduledParam( + init_value=self._init_ucb_overwrite_probability, + final_value=self._final_ucb_overwrite_probability, + rate=self._decay_ucb_overwrite_probability, + ) + ) + + scheduled_params = { + 'ucb_coefficient': ucb_coefficient_param, + 'explore_region_ucb_coefficient': explore_region_ucb_coefficient_param, + 'ucb_overwrite_probability': ucb_overwrite_probability_param, + } + return scheduled_designer.ScheduledDesigner( problem, designer_factory=self._gp_ucb_pe_factory, designer_state_updater=_gp_ucb_pe_state_updater, - scheduled_params={self._ucb_coef_param_name: ucb_coef_param}, + scheduled_params=scheduled_params, expected_total_num_trials=self._expected_total_num_trials, )