Respect continuous feature bounds in LBFGSBOptimizer

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,
     )

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)