Clean up sampler usage in pick_best_out_of_sample_point_acqf_class (facebook#5095)

Summary:
Pull Request resolved: facebook#5095

This was generating a SobolQMCSampler regardless of the model it is used with. SobolQMCSampler does not support PosteriorList, which was leading to issues with LILO integration. This diff cleans up the generated options in `pick_best_out_of_sample_point_acqf_class`, since BoTorch (in `MCSamplerMixin.get_posterior_samples`) will dispatch to proper defaults without them.
This led to the `options` input becoming unused, so I removed those rather than having them get ignored silently.

Reviewed By: hvarfner

Differential Revision: D97956036

fbshipit-source-id: e363c433d4769488f9225fc2ecca6395442638c2

Author: saitcakmak

ax/generators/torch/botorch_modular/surrogate.py

@@ -50,12 +50,8 @@
 from ax.generators.types import TConfig
 from ax.generators.utils import best_in_sample_point
 from ax.utils.common.base import Base
-from ax.utils.common.constants import Keys
 from ax.utils.common.logger import get_logger
-from ax.utils.common.typeutils import (
-    _argparse_type_encoder,
-    assert_is_instance_optional,
-)
+from ax.utils.common.typeutils import _argparse_type_encoder
 from ax.utils.stats.model_fit_stats import (
     DIAGNOSTIC_FN_DIRECTIONS,
     DIAGNOSTIC_FNS,
@@ -1023,7 +1019,6 @@ def best_out_of_sample_point(
         self,
         search_space_digest: SearchSpaceDigest,
         torch_opt_config: TorchOptConfig,
-        options: TConfig | None = None,
     ) -> tuple[Tensor, Tensor]:
         """Finds the best predicted point and the corresponding value of the
         appropriate best point acquisition function.
@@ -1032,9 +1027,6 @@ def best_out_of_sample_point(
             search_space_digest: A `SearchSpaceDigest`.
             torch_opt_config: A `TorchOptConfig`; none-None `fixed_features` is
                 not supported.
-            options: Optional. If present, `seed_inner` (default None) and `qmc`
-                (default True) will be parsed from `options`; any other keys
-                will be ignored.
 
         Returns:
             A two-tuple (`candidate`, `acqf_value`), where `candidate` is a 1d
@@ -1048,18 +1040,8 @@ def best_out_of_sample_point(
             # TODO (ref: https://fburl.com/diff/uneqb3n9)
             raise NotImplementedError("Fixed features not yet supported.")
 
-        options = options or {}
-        botorch_acqf_class, botorch_acqf_options = (
-            pick_best_out_of_sample_point_acqf_class(
-                outcome_constraints=torch_opt_config.outcome_constraints,
-                seed_inner=assert_is_instance_optional(
-                    options.get(Keys.SEED_INNER, None), int
-                ),
-                qmc=assert_is_instance(
-                    options.get(Keys.QMC, True),
-                    bool,
-                ),
-            )
+        botorch_acqf_class = pick_best_out_of_sample_point_acqf_class(
+            outcome_constraints=torch_opt_config.outcome_constraints,
         )
 
         # Avoiding circular import between `Surrogate` and `Acquisition`.
@@ -1070,7 +1052,6 @@ def best_out_of_sample_point(
             botorch_acqf_class=botorch_acqf_class,
             search_space_digest=search_space_digest,
             torch_opt_config=torch_opt_config,
-            botorch_acqf_options=botorch_acqf_options,
         )
         candidates, acqf_value, _ = acqf.optimize(
             n=1,

ax/generators/torch/tests/test_surrogate.py

@@ -1652,7 +1652,6 @@ def test_best_out_of_sample_point(self) -> None:
             candidate, acqf_value = surrogate.best_out_of_sample_point(
                 search_space_digest=self.search_space_digest,
                 torch_opt_config=torch_opt_config,
-                options=self.options,
             )
             candidate_in_bounds = all(
                 ((x >= b[0]) & (x <= b[1]) for x, b in zip(candidate, self.bounds))

ax/generators/torch/tests/test_utils.py

@@ -1471,17 +1471,14 @@ def test_choose_botorch_acqf_class_with_map_key_fixing(self) -> None:
         self.assertEqual(acqf_class, qLogNoisyExpectedImprovement)
 
     def test_pick_best_out_of_sample_point_acqf_class(self) -> None:
-        # Unconstrained: PosteriorMean with no options.
-        acqf_class, acqf_options = pick_best_out_of_sample_point_acqf_class(
+        # Unconstrained: PosteriorMean.
+        acqf_class = pick_best_out_of_sample_point_acqf_class(
             outcome_constraints=None,
         )
         self.assertEqual(acqf_class, PosteriorMean)
-        self.assertEqual(acqf_options, {})
 
-        # Constrained: qSimpleRegret with no explicit sampler (get_sampler
-        # auto-dispatches, which handles PosteriorList correctly).
-        acqf_class, acqf_options = pick_best_out_of_sample_point_acqf_class(
+        # Constrained: qSimpleRegret.
+        acqf_class = pick_best_out_of_sample_point_acqf_class(
             outcome_constraints=(torch.tensor([[1.0, 0.0]]), torch.tensor([[0.5]])),
         )
         self.assertEqual(acqf_class, qSimpleRegret)
-        self.assertEqual(acqf_options, {})

ax/generators/torch/utils.py

@@ -9,7 +9,6 @@
 import logging
 from collections.abc import Callable, Sequence
 from dataclasses import dataclass
-from typing import Any, cast
 
 import numpy.typing as npt
 import torch
@@ -404,18 +403,10 @@ def objective(samples: Tensor, X: Tensor | None = None) -> Tensor:
 
 def pick_best_out_of_sample_point_acqf_class(
     outcome_constraints: tuple[Tensor, Tensor] | None = None,
-    mc_samples: int = 512,
-    qmc: bool = True,
-    seed_inner: int | None = None,
-) -> tuple[type[AcquisitionFunction], dict[str, Any]]:
+) -> type[AcquisitionFunction]:
     if outcome_constraints is None:
-        acqf_class = PosteriorMean
-        acqf_options = {}
-    else:
-        acqf_class = qSimpleRegret
-        acqf_options = {}
-
-    return cast(type[AcquisitionFunction], acqf_class), acqf_options
+        return PosteriorMean
+    return qSimpleRegret
 
 
 def predict_from_model(