diff --git a/vizier/_src/algorithms/optimizers/lbfgsb_optimizer.py b/vizier/_src/algorithms/optimizers/lbfgsb_optimizer.py index e9dc9bf43..e02ce2ee3 100644 --- a/vizier/_src/algorithms/optimizers/lbfgsb_optimizer.py +++ b/vizier/_src/algorithms/optimizers/lbfgsb_optimizer.py @@ -16,7 +16,8 @@ """L-BFGS-B Strategy optimizer.""" -from typing import Callable, Optional, Union +import itertools +from typing import Callable, Optional, Sequence, Union import attr from flax import struct @@ -51,6 +52,9 @@ class LBFGSBOptimizer: n_feature_dimensions_with_padding: types.ContinuousAndCategorical[int] = ( struct.field(pytree_node=False) ) + continuous_features_bounds: Sequence[tuple[float, float]] = struct.field( + pytree_node=False + ) # Number of parallel runs of L-BFGS-B. random_restarts: int = struct.field(pytree_node=False, default=25) # Number of iterations for each L-BFGS-B run. @@ -144,10 +148,23 @@ def _opt_score_fn(x): def setup(rng): return jax.random.uniform(rng, shape=feature_shape) - # Constraints are [0, 1]. - constraints = sp.Constraint( - bounds=(np.zeros(feature_shape), np.ones(feature_shape)) + continuous_features_min, continuous_features_max = itertools.zip_longest( + *self.continuous_features_bounds ) + bounds = [] + for bound, fill_value in ( + (continuous_features_min, 0.0), + (continuous_features_max, 1.0), + ): + # Pad the bound with fill_value to the length of the feature dimension + # with padding and account for the parallel dimension, so the shape of the + # bound is the same as `feature_shape` above. + padded_bound = bound + (fill_value,) * ( + self.n_feature_dimensions_with_padding.continuous - len(bound) + ) + bounds.append(np.stack([padded_bound] * parallel_dim, axis=0)) + + constraints = sp.Constraint(bounds=(bounds[0], bounds[1])) new_features, _ = optimize( jax.vmap(setup)(jax.random.split(init_seed, self.random_restarts)), @@ -199,9 +216,15 @@ def __call__( empty_features.continuous.shape[-1], empty_features.categorical.shape[-1], ) + continous_features_bounds = [ + (float(spec.bounds[0]), float(spec.bounds[1])) + for spec in converter.output_specs.continuous + ] + return LBFGSBOptimizer( n_feature_dimensions=n_feature_dimensions, n_feature_dimensions_with_padding=n_feature_dimensions_with_padding, + continuous_features_bounds=continous_features_bounds, random_restarts=self.random_restarts, maxiter=self.maxiter, ) diff --git a/vizier/_src/algorithms/optimizers/lbfgsb_optimizer_test.py b/vizier/_src/algorithms/optimizers/lbfgsb_optimizer_test.py index 167ea8874..924521579 100644 --- a/vizier/_src/algorithms/optimizers/lbfgsb_optimizer_test.py +++ b/vizier/_src/algorithms/optimizers/lbfgsb_optimizer_test.py @@ -42,6 +42,23 @@ def test_optimize_candidates_len(self): res = optimizer(score_fn=score_fn) self.assertLen(res.rewards, 1) + def test_singleton_constraints_are_respected(self): + problem = vz.ProblemStatement() + problem.search_space.root.add_float_param('f1', 0.0, 10.0) + problem.search_space.root.add_float_param('f2', 0.0, 10.0) + problem.search_space.root.add_float_param('f3', 5.0, 5.0) + converter = converters.TrialToModelInputConverter.from_problem(problem) + score_fn = lambda x, _: jnp.sum(x.continuous.padded_array, axis=-1) + optimizer = lo.LBFGSBOptimizerFactory(random_restarts=10, maxiter=20)( + converter + ) + res = optimizer(score_fn=score_fn) + best_candidates = vb.best_candidates_to_trials(res, converter) + best_score = score_fn(converter.to_features(best_candidates), None) + # Evaluating the score function being optimized on the best candidate should + # result in the same value as the output of the optimizer. + self.assertSequenceAlmostEqual(best_score, res.rewards) + def test_best_candidates_count_is_1(self): problem = vz.ProblemStatement() problem.search_space.root.add_float_param('f1', 0.0, 1.0)