copybara-service · Sep 20, 2023 · Sep 20, 2023
diff --git a/vizier/_src/algorithms/designers/gp/acquisitions.py b/vizier/_src/algorithms/designers/gp/acquisitions.py
@@ -25,6 +25,7 @@
 from jax import numpy as jnp
 import numpy as np
 from tensorflow_probability.substrates import jax as tfp
+from vizier._src.algorithms.designers import scalarization
 from vizier._src.jax import types
 from vizier._src.jax.models import continuous_only_kernel
 
@@ -518,6 +519,21 @@ def __call__(
     )()
 
 
+@struct.dataclass
+class ScalarizedAcquisition(AcquisitionFunction):
+  """Wrapper that scalarizes multiple objective before acquisition eval."""
+
+  acquisition_fn: AcquisitionFunction
+  scalarizer: scalarization.Scalarization
+
+  def __call__(
+      self,
+      dist: tfd.Distribution,
+      seed: Optional[jax.random.KeyArray] = None,
+  ) -> jax.Array:
+    return jnp.squeeze(self.scalarizer(self.acquisition_fn(dist, seed)))
+
+
 @struct.dataclass
 class MultiAcquisitionFunction(AcquisitionFunction):
   """Wrapper that calls multiple acquisition functions."""

diff --git a/vizier/_src/algorithms/designers/gp/acquisitions_test.py b/vizier/_src/algorithms/designers/gp/acquisitions_test.py
@@ -21,10 +21,10 @@
 from jax.config import config
 import numpy as np
 from tensorflow_probability.substrates import jax as tfp
+from vizier._src.algorithms.designers import scalarization
 from vizier._src.algorithms.designers.gp import acquisitions
 from vizier._src.jax import types
 from absl.testing import absltest
-
 tfd = tfp.distributions
 tfpk = tfp.math.psd_kernels
 tfpke = tfp.experimental.psd_kernels
@@ -70,6 +70,16 @@ def test_ei(self):
         0.34635347,
     )
 
+  def test_scalarized_ucb(self):
+    ucb = acquisitions.UCB(coefficient=2.0)
+    scalarizer = scalarization.HyperVolumeScalarization(
+        weights=jnp.array([0.1, 0.2])
+        )
+    acq = acquisitions.ScalarizedAcquisition(ucb, scalarizer)
+    self.assertAlmostEqual(
+        acq(tfd.Normal([0.1, 0.2], [1, 2])), jnp.array(21), delta=1e-5
+    )
+
   def test_pi(self):
     labels = types.PaddedArray.as_padded(jnp.array([[0.2]]))
     best_labels = acquisitions.get_best_labels(labels)

diff --git a/vizier/_src/algorithms/designers/gp/gp_models.py b/vizier/_src/algorithms/designers/gp/gp_models.py
@@ -140,41 +140,40 @@ def predict_with_aux(
 
 
 def get_vizier_gp_coroutine(
-    features: types.ModelInput,
+    data: types.ModelData,
     *,
     linear_coef: float = 0.0,
-    labels: Optional[types.PaddedArray] = None,
 ) -> sp.ModelCoroutine:
   """Gets a GP model coroutine.
 
   Args:
-    features: The features used to the train the GP model
+    data: The data used to the train the GP model
     linear_coef: If non-zero, uses a linear kernel with `linear_coef`
       hyperparameter.
-    labels: The labels used to train the GP model, default to None, which gives
-      a single-metric GP-coroutine. The last dimension should equal the number
-      of metrics.
 
   Returns:
     The model coroutine.
   """
   # Construct the multi-task GP.
-  if labels and labels.shape[-1] > 1:
+  labels_shape = data.labels.shape
+  if labels_shape[-1] > 1:
     gp_coroutine = multitask_tuned_gp_models.VizierMultitaskGaussianProcess(
         _feature_dim=types.ContinuousAndCategorical[int](
-            features.continuous.padded_array.shape[-1],
-            features.categorical.padded_array.shape[-1],
+            data.features.continuous.padded_array.shape[-1],
+            data.features.categorical.padded_array.shape[-1],
         ),
-        _num_tasks=labels.shape[-1],
+        _num_tasks=labels_shape[-1],
     )
     return sp.StochasticProcessModel(gp_coroutine).coroutine
 
   if linear_coef:
     return tuned_gp_models.VizierLinearGaussianProcess.build_model(
-        features=features, linear_coef=linear_coef
+        features=data.features, linear_coef=linear_coef
     ).coroutine
 
-  return tuned_gp_models.VizierGaussianProcess.build_model(features).coroutine
+  return tuned_gp_models.VizierGaussianProcess.build_model(
+      data.features
+  ).coroutine
 
 
 def _train_gp(spec: GPTrainingSpec, data: types.ModelData) -> GPState:

diff --git a/vizier/_src/algorithms/designers/gp/gp_models_test.py b/vizier/_src/algorithms/designers/gp/gp_models_test.py
@@ -93,7 +93,7 @@ def create_model_data(
       ard_optimizer=optimizers.default_optimizer(),
       ard_rng=jax.random.PRNGKey(0),
       coroutine=gp_models.get_vizier_gp_coroutine(
-          features=train_data.features, linear_coef=linear_coef
+          data=train_data, linear_coef=linear_coef
       ),
       ensemble_size=ensemble_size,
       ard_random_restarts=optimizers.DEFAULT_RANDOM_RESTARTS,
@@ -400,9 +400,7 @@ def test_multi_task(self):
     train_spec = gp_models.GPTrainingSpec(
         ard_optimizer=optimizers.default_optimizer(),
         ard_rng=jax.random.PRNGKey(0),
-        coroutine=gp_models.get_vizier_gp_coroutine(
-            features=model_data.features, labels=model_data.labels
-        ),
+        coroutine=gp_models.get_vizier_gp_coroutine(data=model_data),
     )
     gp = gp_models.train_gp(train_spec, model_data)
 

diff --git a/vizier/_src/algorithms/designers/gp/output_warpers.py b/vizier/_src/algorithms/designers/gp/output_warpers.py
@@ -38,9 +38,12 @@ def _validate_labels(labels_arr: types.Array) -> types.Array:
   """Checks and modifies the shape and values of the labels."""
   labels_arr = labels_arr.astype(float)
   if not (labels_arr.ndim == 2 and labels_arr.shape[-1] == 1):
-    raise ValueError('Labels need to be an array of shape (num_points, 1).')
+    raise ValueError(
+        'Labels need to be an array of shape (num_points, 1).'
+        f'Got shape: {labels_arr.shape}'
+    )
   if np.isposinf(labels_arr).any():
-    raise ValueError('Inifinity metric value is not valid.')
+    raise ValueError('Infinity metric value is not valid.')
   if np.isneginf(labels_arr).any():
     labels_arr[np.isneginf(labels_arr)] = np.nan
   return labels_arr

diff --git a/vizier/_src/algorithms/designers/gp_bandit.py b/vizier/_src/algorithms/designers/gp_bandit.py
@@ -148,8 +148,6 @@ def __attrs_post_init__(self):
     # Extra validations
     if self._problem.search_space.is_conditional:
       raise ValueError(f'{type(self)} does not support conditional search.')
-    elif len(self._problem.metric_information) != 1:
-      raise ValueError(f'{type(self)} works with exactly one metric.')
     # Extra initializations.
     # Discrete parameters are continuified to account for their actual values.
     self._converter = converters.TrialToModelInputConverter.from_problem(
@@ -175,7 +173,7 @@ def __attrs_post_init__(self):
         ),
     )
 
-    coroutine = gp_models.get_vizier_gp_coroutine(empty_data.features)
+    coroutine = gp_models.get_vizier_gp_coroutine(empty_data)
     params = sp.CoroutineWithData(coroutine, empty_data).setup(self._rng)
     model = sp.StochasticProcessWithCoroutine(coroutine, params)
     predictive = sp.UniformEnsemblePredictive(
@@ -305,7 +303,13 @@ def _generate_seed_trials(self, count: int) -> Sequence[vz.TrialSuggestion]:
   @_experimental_override_allowed
   def _warp_labels(self, labels: types.Array) -> types.Array:
     """Subclasses can override this method for experiments."""
-    return self._output_warper.warp(labels)
+    return np.concatenate(
+        [
+            self._output_warper.warp(labels[:, i : i + 1])
+            for i in range(labels.shape[1])
+        ],
+        axis=-1,
+    )
 
   @profiler.record_runtime
   def _trials_to_data(self, trials: Sequence[vz.Trial]) -> types.ModelData:
@@ -341,7 +345,7 @@ def _create_gp_spec(
         ard_optimizer=self._ard_optimizer,
         ard_rng=ard_rng,
         coroutine=gp_models.get_vizier_gp_coroutine(
-            features=data.features, linear_coef=self._linear_coef
+            data=data, linear_coef=self._linear_coef
         ),
         ensemble_size=self._ensemble_size,
        ard_random_restarts=self._ard_random_restarts,

diff --git a/vizier/_src/algorithms/designers/gp_bandit_test.py b/vizier/_src/algorithms/designers/gp_bandit_test.py
@@ -25,6 +25,7 @@
 from vizier import pyvizier as vz
 from vizier._src.algorithms.designers import gp_bandit
 from vizier._src.algorithms.designers import quasi_random
+from vizier._src.algorithms.designers import scalarization
 from vizier._src.algorithms.designers.gp import acquisitions
 from vizier._src.algorithms.optimizers import eagle_strategy as es
 from vizier._src.algorithms.optimizers import lbfgsb_optimizer as lo
@@ -393,6 +394,62 @@ def _qei_factory(data: types.ModelData) -> acquisitions.AcquisitionFunction:
         iters * n_parallel,
     )
 
+  def test_multi_metrics(self):
+    search_space = vz.SearchSpace()
+    search_space.root.add_float_param('x0', -5.0, 5.0)
+    problem = vz.ProblemStatement(
+        search_space=search_space,
+        metric_information=vz.MetricsConfig(
+            metrics=[
+                vz.MetricInformation(
+                    'obj1', goal=vz.ObjectiveMetricGoal.MAXIMIZE
+                ),
+                vz.MetricInformation(
+                    'obj2', goal=vz.ObjectiveMetricGoal.MAXIMIZE
+                ),
+            ]
+        ),
+    )
+
+    def _scalarized_ucb(
+        data: types.ModelData,
+    ) -> acquisitions.AcquisitionFunction:
+      del data
+      ucb = acquisitions.UCB(coefficient=2.0)
+      scalarizer = scalarization.HyperVolumeScalarization(
+          weights=np.array([0.1, 0.2])
+      )
+      return acquisitions.ScalarizedAcquisition(ucb, scalarizer)
+
+    scoring_fn_factory = acquisitions.bayesian_scoring_function_factory(
+        _scalarized_ucb
+    )
+
+    iters = 3
+    designer = gp_bandit.VizierGPBandit(
+        problem=problem,
+        acquisition_optimizer_factory=vectorized_optimizer_factory,
+        ard_optimizer=optimizers.JaxoptLbfgsB(
+            optimizers.LbfgsBOptions(maxiter=5, num_line_search_steps=5)
+        ),
+        scoring_function_factory=scoring_fn_factory,
+        scoring_function_is_parallel=True,
+        use_trust_region=False,
+        ensemble_size=3,
+        rng=jax.random.PRNGKey(0),
+        linear_coef=0.1,
+    )
+    self.assertLen(
+        test_runners.RandomMetricsRunner(
+            problem,
+            iters=iters,
+            verbose=1,
+            validate_parameters=True,
+            seed=1,
+        ).run_designer(designer),
+        iters,
+    )
+
 
 class GPBanditPriorsTest(parameterized.TestCase):