Move multiobjective to simplify from_problem factory

PiperOrigin-RevId: 650835016

Author: xingyousong

vizier/_src/algorithms/designers/gp_bandit.py

@@ -119,7 +119,9 @@ class VizierGPBandit(vza.Designer, vza.Predictor):
       default=optimizers.DEFAULT_RANDOM_RESTARTS, kw_only=True
   )
   _num_seed_trials: int = attr.field(default=1, kw_only=True)
-  _linear_coef: float = attr.field(default=0.0, kw_only=True)
+  # Linear coef is set to 1.0 as prior and uses VizierLinearGaussianProcess
+  # which uses a sum of Matern and linear but ARD still tunes its amplitude.
+  _linear_coef: float = attr.field(default=1.0, kw_only=True)
   _scoring_function_factory: acq_lib.ScoringFunctionFactory = attr.field(
       factory=lambda: default_scoring_function_factory,
       kw_only=True,
@@ -142,6 +144,11 @@ class VizierGPBandit(vza.Designer, vza.Predictor):
       factory=output_warpers.create_default_warper, kw_only=True
   )
 
+  # Multi-objective settings.
+  _num_scalars: int = attr.field(default=1000, kw_only=True)
+  _ref_scaling: float = attr.field(default=0.01, kw_only=True)
+  _num_samples: Optional[int] = attr.field(default=None, kw_only=True)
+
   # ------------------------------------------------------------------
   # Internal attributes which should not be set by callers.
   # ------------------------------------------------------------------
@@ -182,6 +189,57 @@ def __attrs_post_init__(self):
         seed=int(jax.random.randint(self._rng, [], 0, 2**16)),
     )
 
+    m_info = self._problem.metric_information
+    if not m_info.is_single_objective:
+      num_obj = len(m_info.of_type(vz.MetricType.OBJECTIVE))
+      self._rng, weights_rng = jax.random.split(self._rng)
+      weights = jnp.abs(
+          jax.random.normal(weights_rng, shape=(self._num_scalars, num_obj))
+      )
+      weights = weights / jnp.linalg.norm(weights, axis=-1, keepdims=True)
+
+      if self._num_samples is None:
+
+        def _scalarized_ucb(
+            data: types.ModelData,
+        ) -> acq_lib.AcquisitionFunction:
+          scalarizer = scalarization.HyperVolumeScalarization(
+              weights,
+              acq_lib.get_reference_point(data.labels, self._ref_scaling),
+          )
+          return acq_lib.ScalarizedAcquisition(
+              acq_lib.UCB(),
+              scalarizer,
+              reduction_fn=lambda x: jnp.mean(x, axis=0),
+          )
+
+        acq_fn_factory = _scalarized_ucb
+
+      else:
+
+        def _scalarized_sample_ehvi(
+            data: types.ModelData,
+        ) -> acq_lib.AcquisitionFunction:
+          scalarizer = scalarization.HyperVolumeScalarization(
+              weights,
+              acq_lib.get_reference_point(data.labels, self._ref_scaling),
+          )
+          return acq_lib.ScalarizedAcquisition(
+              acq_lib.Sample(self._num_samples),
+              scalarizer,
+              # We need to reduce across the scalarization and sample axes.
+              reduction_fn=lambda x: jnp.mean(jax.nn.relu(x), axis=[0, 1]),
+          )
+
+        acq_fn_factory = _scalarized_sample_ehvi
+
+      # Multi-objective overrides.
+      self._scoring_function_factory = (
+          acq_lib.bayesian_scoring_function_factory(acq_fn_factory)
+      )
+      self._scoring_function_is_parallel = True
+      self._use_trust_region = False
+
     self._acquisition_optimizer = self._acquisition_optimizer_factory(
         self._converter
     )
@@ -578,67 +636,7 @@ def from_problem(
       cls,
       problem: vz.ProblemStatement,
       seed: Optional[int] = None,
-      num_scalarizations: int = 1000,
-      reference_scaling: float = 0.01,
-      num_samples: int | None = None,
       **kwargs,
   ) -> 'VizierGPBandit':
     rng = jax.random.PRNGKey(seed or 0)
-    # Linear coef is set to 1.0 as prior and uses VizierLinearGaussianProcess
-    # which uses a sum of Matern and linear but ARD still tunes its amplitude.
-    if problem.is_single_objective:
-      return cls(problem, linear_coef=1.0, rng=rng, **kwargs)
-    else:
-      num_obj = len(problem.metric_information.of_type(vz.MetricType.OBJECTIVE))
-      rng, weights_rng = jax.random.split(rng)
-      weights = jnp.abs(
-          jax.random.normal(weights_rng, shape=(num_scalarizations, num_obj))
-      )
-      weights = weights / jnp.linalg.norm(weights, axis=-1, keepdims=True)
-
-      if num_samples is None:
-
-        def _scalarized_ucb(
-            data: types.ModelData,
-        ) -> acq_lib.AcquisitionFunction:
-          scalarizer = scalarization.HyperVolumeScalarization(
-              weights,
-              acq_lib.get_reference_point(data.labels, reference_scaling),
-          )
-          return acq_lib.ScalarizedAcquisition(
-              acq_lib.UCB(),
-              scalarizer,
-              reduction_fn=lambda x: jnp.mean(x, axis=0),
-          )
-
-        acq_fn_factory = _scalarized_ucb
-      else:
-
-        def _scalarized_sample_ehvi(
-            data: types.ModelData,
-        ) -> acq_lib.AcquisitionFunction:
-          scalarizer = scalarization.HyperVolumeScalarization(
-              weights,
-              acq_lib.get_reference_point(data.labels, reference_scaling),
-          )
-          return acq_lib.ScalarizedAcquisition(
-              acq_lib.Sample(num_samples),
-              scalarizer,
-              # We need to reduce across the scalarization and sample axes.
-              reduction_fn=lambda x: jnp.mean(jax.nn.relu(x), axis=[0, 1]),
-          )
-
-        acq_fn_factory = _scalarized_sample_ehvi
-
-      scoring_function_factory = acq_lib.bayesian_scoring_function_factory(
-          acq_fn_factory
-      )
-      return cls(
-          problem,
-          linear_coef=1.0,
-          scoring_function_factory=scoring_function_factory,
-          scoring_function_is_parallel=True,
-          use_trust_region=False,
-          rng=rng,
-          **kwargs,
-      )
+    return cls(problem, rng=rng, **kwargs)