Allow treating discrete/integer parameters with many feasible values as continuous for trust regions.

vizier/_src/algorithms/designers/gp_bandit.py

@@ -50,6 +50,13 @@
 from vizier.pyvizier.converters import padding
 from vizier.utils import profiler
 
+
+# The maximum number of feasible values to use for the trust region. If a
+# discrete or integer parameter has more than this many feasible values, it is
+# considered continuous and always included in the trust region computation.
+_MAX_NUM_FEASIBLE_VALUES_FOR_TRUST_REGION = 1000
+
+
 default_acquisition_optimizer_factory = vb.VectorizedOptimizerFactory(
     strategy_factory=es.VectorizedEagleStrategyFactory(
         eagle_config=es.EagleStrategyConfig()
@@ -246,7 +253,9 @@ def acq_fn_factory(data: types.ModelData) -> acq_lib.AcquisitionFunction:
     scoring_fn = self._scoring_function_factory(
         empty_data,
         predictive,
-        self._converter.continuous_feasible_values,
+        self._converter.continuous_feasible_values(
+            max_num_feasible_values=_MAX_NUM_FEASIBLE_VALUES_FOR_TRUST_REGION
+        ),
         self._use_trust_region,
     )
     if (
@@ -532,7 +541,9 @@ def suggest(self, count: int = 1) -> Sequence[vz.TrialSuggestion]:
     scoring_fn = self._scoring_function_factory(
         data,
         gp,
-        self._converter.continuous_feasible_values,
+        self._converter.continuous_feasible_values(
+            max_num_feasible_values=_MAX_NUM_FEASIBLE_VALUES_FOR_TRUST_REGION
+        ),
         self._use_trust_region,
     )
     logging.info('Optimizing acquisition: %s', scoring_fn)

vizier/_src/algorithms/designers/gp_ucb_pe.py

@@ -54,6 +54,12 @@
 tfd = tfp.distributions
 
 
+# The maximum number of feasible values to use for the trust region. If a
+# discrete or integer parameter has more than this many feasible values, it is
+# considered continuous and always included in the trust region computation.
+_MAX_NUM_FEASIBLE_VALUES_FOR_TRUST_REGION = 1000
+
+
 class MultimetricPromisingRegionPenaltyType(enum.Enum):
   """The type of penalty to apply to the points outside the promising region.
 
@@ -1390,7 +1396,9 @@ def suggest(
     )
     tr = acquisitions.TrustRegion(
         trusted=tr_features,
-        continuous_feasible_values=self._converter.continuous_feasible_values,
+        continuous_feasible_values=self._converter.continuous_feasible_values(
+            max_num_feasible_values=_MAX_NUM_FEASIBLE_VALUES_FOR_TRUST_REGION
+        ),
     )
 
     acquisition_problem = copy.deepcopy(self._problem)

vizier/_src/algorithms/designers/gp_ucb_pe_test.py

@@ -615,6 +615,10 @@ def test_discrete_parameters_are_explored(
     root.add_discrete_param('discrete_double_1', [-3.0, 7.0])
     root.add_discrete_param('discrete_double_single_feasible_value', [0.37])
     root.add_float_param('double', min_value=-5.0, max_value=5.0)
+    root.add_int_param('integer_with_many_feasible_values', -1, 6 * 10**7)
+    root.add_discrete_param(
+        'discrete_double_many_feasible_values', np.linspace(-10.0, 10.0, 1000)
+    )
     problem = vz.ProblemStatement(space)
     problem.metric_information.append(
         vz.MetricInformation(

vizier/pyvizier/converters/jnp_converters.py

@@ -246,29 +246,52 @@ def output_specs(self):  # TODO: Add back pytype
   def metric_specs(self) -> Sequence[vz.MetricInformation]:
     return self._impl.metric_specs
 
-  @property
-  def continuous_feasible_values(self) -> list[jax.Array]:
+  def continuous_feasible_values(
+      self, max_num_feasible_values: int | None = None
+  ) -> list[jax.Array]:
     """Returns a list of feasible values for each continuified parameter.
 
     The list is ordered the same as the parameters in the search space. An empty
     array means that the parameter is continuous and all values within its
     range are feasible. The returned feasible values are in the scaled space.
-    Categorical parameters are ignored.
+    Categorical parameters are ignored. If `max_num_feasible_values` is
+    specified, discrete or integer parameters with more than
+    `max_num_feasible_values` feasible values are considered continuous and
+    empty arrays are returned for them.
+
+    Args:
+      max_num_feasible_values: If specified, discrete or integer parameters with
+        more than this many feasible values are considered continuous and empty
+        arrays are returned for them.
+
+    Returns:
+      A list of feasible values for each continuified parameter. The list is
+      ordered the same as the parameters in the search space.
     """
 
     continuous_feasible_values = []
     for param in self._problem.search_space.parameters:
       if param.type in [vz.ParameterType.DISCRETE, vz.ParameterType.INTEGER]:
-        converted_feasible_values = self.to_features([
-            vz.TrialSuggestion(parameters={param.name: value})
-            for value in param.feasible_values
-        ]).continuous.unpad()
-        continuous_feasible_values.append(
-            # `self.to_features` returns NaNs for missing parameters (we call it
-            # with one specific parameter instead of all parameters in the
-            # search space), so we remove them here.
-            converted_feasible_values[~jnp.isnan(converted_feasible_values)]
-        )
+        if param.type == vz.ParameterType.INTEGER:
+          num_feasible_values = param.bounds[1] - param.bounds[0] + 1
+        else:
+          num_feasible_values = len(param.feasible_values)
+        if (
+            max_num_feasible_values is not None
+            and num_feasible_values > max_num_feasible_values
+        ):
+          continuous_feasible_values.append(jnp.asarray([]))
+        else:
+          converted_feasible_values = self.to_features([
+              vz.TrialSuggestion(parameters={param.name: value})
+              for value in param.feasible_values
+          ]).continuous.unpad()
+          continuous_feasible_values.append(
+              # `self.to_features` returns NaNs for missing parameters (we call
+              # it with one specific parameter instead of all parameters in the
+              # search space), so we remove them here.
+              converted_feasible_values[~jnp.isnan(converted_feasible_values)]
+          )
       elif param.type == vz.ParameterType.DOUBLE:
         continuous_feasible_values.append(jnp.asarray([]))
       elif param.type == vz.ParameterType.CUSTOM:

vizier/pyvizier/converters/jnp_converters_test.py

@@ -110,7 +110,7 @@ def setUp(self):
     self.maxDiff = None
 
   def test_continuous_feasible_values(self):
-    """Tests the `continuous_feasible_values` property.
+    """Tests the `continuous_feasible_values` method.
 
     The returned feasible values are expected to be in the scaled space, and
     categorical parameters are ignored.
@@ -125,6 +125,7 @@ def test_continuous_feasible_values(self):
     )
     root.add_int_param('integer', -2, 2)
     root.add_float_param('float', -3, 4)
+    root.add_int_param('integer_with_many_feasible_values', -1, 6 * 10**7)
     converter = jnpc.TrialToModelInputConverter.from_problem(
         vz.ProblemStatement(
             search_space=space,
@@ -135,19 +136,22 @@ def test_continuous_feasible_values(self):
             ],
         )
     )
-    self.assertLen(converter.continuous_feasible_values, 5)
+    continuous_feasible_values = converter.continuous_feasible_values(
+        max_num_feasible_values=1000
+    )
+    self.assertLen(continuous_feasible_values, 6)
     self.assertSequenceAlmostEqual(
-        converter.continuous_feasible_values[0],
+        continuous_feasible_values[0],
         # The original feasible values are linearly scaled to [0, 1].
         (np.asarray([-0.4, -0.3, 0.4, 1.1, 1.15]) + 0.4) / 1.55,
     )
     self.assertSequenceAlmostEqual(
         # The original feasible values [5, 9] are linearly scaled to [0, 1].
-        converter.continuous_feasible_values[1],
+        continuous_feasible_values[1],
         np.asarray([0.0, 1.0]),
     )
     self.assertSequenceAlmostEqual(
-        converter.continuous_feasible_values[2],
+        continuous_feasible_values[2],
         # The original feasible values [1e-5, 1e-2, 1e-1] are log-then-linearly
         # scaled to [0, 1].
         (
@@ -156,13 +160,15 @@ def test_continuous_feasible_values(self):
         ),
     )
     self.assertSequenceAlmostEqual(
-        converter.continuous_feasible_values[3],
+        continuous_feasible_values[3],
         # The integer parameter with range [-2, 2] are internally scaled to
         # evenly spaced values in [0, 1].
         np.asarray([0.0, 0.25, 0.5, 0.75, 1.0]),
     )
     # Feasible values for float parameters are empty.
-    self.assertEmpty(converter.continuous_feasible_values[4])
+    self.assertEmpty(continuous_feasible_values[4])
+    # Feasible values for int parameters with many feasible values are empty.
+    self.assertEmpty(continuous_feasible_values[5])
 
   def test_back_to_back_conversion(self):
     converter = jnpc.TrialToContinuousAndCategoricalConverter.from_study_config(