Make BoTorch test problems into ParamBasedTestProblems

Summary:
Context:

In the benchmarks, we currently have a zoo of five runners designed to handle different ways of generating data. This will make it challenging to support other changes to runners, such as running trials asynchronously. This diff is a step towards consolidating into one Runner, abstracting away the data-generating process into a test problem.

* BoTorchTestProblemRunner has a BoTorch problem
* ParamBasedTestProblemRunner has a ParamBasedTestProblem
* SurrogateRunner has a SurrogateProblem
* `BoTorchTestProblemRunner` and `ParamBasedTestProblemRunner` subclass `BenchmarkRunner`
* `SurrogateRunner` subclasses `BenchmarkRunner`

This diff will enable going from 5 runners to 3, wrapping BoTorch problems with `ParamBasedTestProblem`. This will allow for combining `BoTorchTestProblemRunner`, `ParamBasedTestProblemRunner`, and `SyntheticProblemRunner`.

This PR:
* Introduces `BoTorchTestProblem`, which is a `ParamBasedTestProblem` and thus can be used with `ParamBasedTestProblemRunner`. It localizes tensor-related logic, including `modified_bounds`, into the BoTorch problem so it doesn't need to be handled by the runner.
* Gets rid of `SyntheticProblemRunner`, merging it with `ParamBasedTestProblemRunner`. We may want to rename this class after more consolidation.
* Makes `BoTorchTestProblemRunner` a do-nothing subclass of `ParamBasedTestProblemRunner`

Differential Revision: D63639190

ax/benchmark/benchmark_problem.py

@@ -13,7 +13,10 @@
 
 from ax.benchmark.benchmark_metric import BenchmarkMetric
 from ax.benchmark.runners.base import BenchmarkRunner
-from ax.benchmark.runners.botorch_test import BotorchTestProblemRunner
+from ax.benchmark.runners.botorch_test import (
+    BoTorchTestProblem,
+    ParamBasedTestProblemRunner,
+)
 from ax.core.data import Data
 from ax.core.experiment import Experiment
 from ax.core.objective import MultiObjective, Objective
@@ -384,9 +387,8 @@ def create_problem_from_botorch(
         name=name,
         search_space=search_space,
         optimization_config=optimization_config,
-        runner=BotorchTestProblemRunner(
-            # pyre-ignore[45]: Can't instantiate abstract class
-            test_problem=test_problem_class(**test_problem_kwargs),
+        runner=ParamBasedTestProblemRunner(
+            test_problem=BoTorchTestProblem(botorch_problem=test_problem),
             outcome_names=outcome_names,
             search_space_digest=extract_search_space_digest(
                 search_space=search_space,

ax/benchmark/problems/synthetic/discretized/mixed_integer.py

@@ -21,7 +21,10 @@
 from ax.benchmark.benchmark_metric import BenchmarkMetric
 
 from ax.benchmark.benchmark_problem import BenchmarkProblem
-from ax.benchmark.runners.botorch_test import BotorchTestProblemRunner
+from ax.benchmark.runners.botorch_test import (
+    BoTorchTestProblem,
+    ParamBasedTestProblemRunner,
+)
 from ax.core.objective import Objective
 from ax.core.optimization_config import OptimizationConfig
 from ax.core.parameter import ParameterType, RangeParameter
@@ -101,8 +104,11 @@ def _get_problem_from_common_inputs(
         test_problem = test_problem_class(dim=dim)
     else:
         test_problem = test_problem_class(dim=dim, bounds=test_problem_bounds)
-    runner = BotorchTestProblemRunner(
-        test_problem=test_problem, outcome_names=[metric_name], modified_bounds=bounds
+    runner = ParamBasedTestProblemRunner(
+        test_problem=BoTorchTestProblem(
+            botorch_problem=test_problem, modified_bounds=bounds
+        ),
+        outcome_names=[metric_name],
     )
     return BenchmarkProblem(
         name=benchmark_name + ("_observed_noise" if observe_noise_sd else ""),

ax/benchmark/problems/synthetic/hss/jenatton.py

@@ -63,7 +63,7 @@ def evaluate_true(self, params: Mapping[str, float | int | None]) -> torch.Tenso
         # `jenatton_test_function`, for 1st positional argument, expected
         # `Optional[float]` but got `Union[None, bool, float, int, str]`.
         value = jenatton_test_function(**params)
-        return torch.tensor(value)
+        return torch.tensor(value, dtype=torch.double)
 
 
 def get_jenatton_benchmark_problem(

ax/benchmark/runners/botorch_test.py

@@ -7,14 +7,12 @@
 
 from abc import ABC, abstractmethod
 from collections.abc import Mapping
-from dataclasses import asdict, dataclass
+from dataclasses import dataclass
 
 import torch
 from ax.benchmark.runners.base import BenchmarkRunner
-from ax.core.search_space import SearchSpaceDigest
 from ax.core.types import TParamValue
-from ax.utils.common.base import Base
-from ax.utils.common.equality import equality_typechecker
+from botorch.test_functions.multi_objective import MultiObjectiveTestProblem
 from botorch.test_functions.synthetic import BaseTestProblem, ConstrainedBaseTestProblem
 from botorch.utils.transforms import normalize, unnormalize
 from torch import Tensor
@@ -46,7 +44,65 @@ def evaluate_slack_true(self, params: Mapping[str, TParamValue]) -> Tensor:
 
 
 @dataclass(kw_only=True)
-class SyntheticProblemRunner(BenchmarkRunner, ABC):
+class BoTorchTestProblem(ParamBasedTestProblem):
+    botorch_problem: BaseTestProblem
+    modified_bounds: list[tuple[float, float]] | None = None
+    num_objectives: int = 1
+
+    def __post_init__(self) -> None:
+        if isinstance(self.botorch_problem, MultiObjectiveTestProblem):
+            self.num_objectives = self.botorch_problem.num_objectives
+        if self.botorch_problem.noise_std is not None:
+            raise ValueError(
+                "noise_std should be set on the runner, not the test problem."
+            )
+        if (
+            hasattr(self.botorch_problem, "constraint_noise_std")
+            and self.botorch_problem.constraint_noise_std is not None
+        ):
+            raise ValueError(
+                "constraint_noise_std should be set on the runner, not the test "
+                "problem."
+            )
+        if self.botorch_problem.negate:
+            raise ValueError(
+                "negate should be set on the runner, not the test problem."
+            )
+        self.botorch_problem = self.botorch_problem.to(dtype=torch.double)
+
+    def tensorize_params(self, params: Mapping[str, int | float]) -> torch.Tensor:
+        X = torch.tensor(
+            [value for _key, value in [*params.items()][: self.botorch_problem.dim]],
+            dtype=torch.double,
+        )
+
+        if self.modified_bounds is not None:
+            # Normalize from modified bounds to unit cube.
+            unit_X = normalize(
+                X, torch.tensor(self.modified_bounds, dtype=torch.double).T
+            )
+            # Unnormalize from unit cube to original problem bounds.
+            X = unnormalize(unit_X, self.botorch_problem.bounds)
+        return X
+
+    # pyre-fixme [14]: inconsistent override
+    def evaluate_true(self, params: Mapping[str, float | int]) -> torch.Tensor:
+        x = self.tensorize_params(params=params)
+        return self.botorch_problem.evaluate_true(x)
+
+    # pyre-fixme [14]: inconsistent override
+    def evaluate_slack_true(self, params: Mapping[str, float | int]) -> torch.Tensor:
+        if not isinstance(self.botorch_problem, ConstrainedBaseTestProblem):
+            raise NotImplementedError  # TODO unsupportedError
+        # todo: could return x so as to not recompute
+        # or could do both methods together, track indices of outcomes,
+        # and only negate the non-constraints
+        x = self.tensorize_params(params=params)
+        return self.botorch_problem.evaluate_slack_true(x)
+
+
+@dataclass(kw_only=True)
+class ParamBasedTestProblemRunner(BenchmarkRunner):
     """A Runner for evaluating synthetic problems, either BoTorch
     `BaseTestProblem`s or Ax benchmarking `ParamBasedTestProblem`s.
 
@@ -69,15 +125,17 @@ class SyntheticProblemRunner(BenchmarkRunner, ABC):
         search_space_digest: Used to extract target fidelity and task.
     """
 
-    test_problem: BaseTestProblem | ParamBasedTestProblem
+    test_problem: ParamBasedTestProblem
     modified_bounds: list[tuple[float, float]] | None = None
     constraint_noise_std: float | list[float] | None = None
     noise_std: float | list[float] | None = None
     negate: bool = False
 
     @property
     def _is_constrained(self) -> bool:
-        return isinstance(self.test_problem, ConstrainedBaseTestProblem)
+        return isinstance(self.test_problem, BoTorchTestProblem) and isinstance(
+            self.test_problem.botorch_problem, ConstrainedBaseTestProblem
+        )
 
     def get_noise_stds(self) -> None | float | dict[str, float]:
         noise_std = self.noise_std
@@ -116,134 +174,24 @@ def get_noise_stds(self) -> None | float | dict[str, float]:
 
         return noise_std_dict
 
-
-@dataclass(kw_only=True)
-class BotorchTestProblemRunner(SyntheticProblemRunner):
-    """
-    A `SyntheticProblemRunner` for BoTorch `BaseTestProblem`s.
-
-    Args:
-        test_problem: A BoTorch `BaseTestProblem`.
-        outcome_names: The names of the outcomes returned by the problem.
-        modified_bounds: The bounds that are used by the Ax search space
-            while optimizing the problem. If different from the bounds of the
-            test problem, we project the parameters into the test problem
-            bounds before evaluating the test problem.
-            For example, if the test problem is defined on [0, 1] but the Ax
-            search space is integers in [0, 10], an Ax parameter value of
-            5 will correspond to 0.5 while evaluating the test problem.
-            If modified bounds are not provided, the test problem will be
-            evaluated using the raw parameter values.
-        search_space_digest: Used to extract target fidelity and task.
-    """
-
-    test_problem: BaseTestProblem
-
-    def __post_init__(self, search_space_digest: SearchSpaceDigest | None) -> None:
-        super().__post_init__(search_space_digest=search_space_digest)
-        if self.test_problem.noise_std is not None:
-            raise ValueError(
-                "noise_std should be set on the runner, not the test problem."
-            )
-        if (
-            hasattr(self.test_problem, "constraint_noise_std")
-            and self.test_problem.constraint_noise_std is not None
-        ):
-            raise ValueError(
-                "constraint_noise_std should be set on the runner, not the test "
-                "problem."
-            )
-        if self.test_problem.negate:
-            raise ValueError(
-                "negate should be set on the runner, not the test problem."
-            )
-        self.test_problem = self.test_problem.to(dtype=torch.double)
-
     def get_Y_true(self, params: Mapping[str, TParamValue]) -> Tensor:
-        """
-        Convert the arm to a tensor and evaluate it on the base test problem.
-
-        Convert the tensor to original bounds -- only if modified bounds were
-        provided -- and evaluates the test problem. See the docstring for
-        `modified_bounds` in `BotorchTestProblemRunner.__init__` for details.
-
-        Args:
-            params: Parameterization to evaluate. It will be converted to a
-                `batch_shape x d`-dim tensor of point(s) at which to evaluate the
-                test problem.
+        """Evaluates the test problem.
 
         Returns:
             A `batch_shape x m`-dim tensor of ground truth (noiseless) evaluations.
         """
-        X = torch.tensor(
-            [value for _key, value in [*params.items()][: self.test_problem.dim]],
-            dtype=torch.double,
-        )
-
-        if self.modified_bounds is not None:
-            # Normalize from modified bounds to unit cube.
-            unit_X = normalize(
-                X, torch.tensor(self.modified_bounds, dtype=torch.double).T
-            )
-            # Unnormalize from unit cube to original problem bounds.
-            X = unnormalize(unit_X, self.test_problem.bounds)
-
-        Y_true = self.test_problem.evaluate_true(X).view(-1)
-        # `BaseTestProblem.evaluate_true()` does not negate the outcome
+        Y_true = self.test_problem.evaluate_true(params).view(-1)
+        # `ParamBasedTestProblem.evaluate_true()` does not negate the outcome
         if self.negate:
             Y_true = -Y_true
-
         if self._is_constrained:
             # Convention: Concatenate objective and black box constraints. `view()`
             # makes the inputs 1d, so the resulting `Y_true` are also 1d.
             Y_true = torch.cat(
-                [Y_true, self.test_problem.evaluate_slack_true(X).view(-1)],
+                [Y_true, self.test_problem.evaluate_slack_true(params).view(-1)],
                 dim=-1,
             )
-
         return Y_true
 
-    @equality_typechecker
-    def __eq__(self, other: Base) -> bool:
-        """
-        Compare equality by comparing dicts, except for `test_problem`.
-
-        Dataclasses are compared by comparing the results of calling asdict on
-        them. However, equality checks don't work as needed with BoTorch test
-        problems, e.g. Branin() == Branin() is False. To get around that, the
-        test problem is stripped from the dictionary. This doesn't make the
-        check less sensitive, as long as the problem has not been modified,
-        because the test problem class and keyword arguments will still be
-        compared.
-        """
-        if not isinstance(other, type(self)):
-            return False
-        self_as_dict = asdict(self)
-        other_as_dict = asdict(other)
-        self_as_dict.pop("test_problem")
-        other_as_dict.pop("test_problem")
-        return (self_as_dict == other_as_dict) and (
-            type(self.test_problem) is type(other.test_problem)
-        )
-
-
-@dataclass(kw_only=True)
-class ParamBasedTestProblemRunner(SyntheticProblemRunner):
-    """
-    A `SyntheticProblemRunner` for `ParamBasedTestProblem`s. See
-    `SyntheticProblemRunner` for more information.
-    """
-
-    test_problem: ParamBasedTestProblem
 
-    def get_Y_true(self, params: Mapping[str, TParamValue]) -> Tensor:
-        """Evaluates the test problem.
-
-        Returns:
-            A `batch_shape x m`-dim tensor of ground truth (noiseless) evaluations.
-        """
-        Y_true = self.test_problem.evaluate_true(params).view(-1)
-        # `ParamBasedTestProblem.evaluate_true()` does not negate the outcome
-        if self.negate:
-            Y_true = -Y_true
-        return Y_true
+BotorchTestProblemRunner = ParamBasedTestProblemRunner

ax/benchmark/tests/problems/test_mixed_integer_problems.py

@@ -15,13 +15,16 @@
     get_discrete_hartmann,
     get_discrete_rosenbrock,
 )
-from ax.benchmark.runners.botorch_test import BotorchTestProblemRunner
+from ax.benchmark.runners.botorch_test import (
+    BoTorchTestProblem,
+    ParamBasedTestProblemRunner,
+)
 from ax.core.arm import Arm
 from ax.core.parameter import ParameterType
 from ax.core.trial import Trial
 from ax.utils.common.testutils import TestCase
-from ax.utils.common.typeutils import checked_cast
 from botorch.test_functions.synthetic import Ackley, Hartmann, Rosenbrock
+from pyre_extensions import assert_is_instance
 
 
 class MixedIntegerProblemsTest(TestCase):
@@ -34,10 +37,10 @@ def test_problems(self) -> None:
             name = problem_cls.__name__
             problem = constructor()
             self.assertEqual(f"Discrete {name}", problem.name)
-            self.assertIsInstance(
-                checked_cast(BotorchTestProblemRunner, problem.runner).test_problem,
-                problem_cls,
-            )
+            runner = assert_is_instance(problem.runner, ParamBasedTestProblemRunner)
+            test_problem = assert_is_instance(runner.test_problem, BoTorchTestProblem)
+            botorch_problem = test_problem.botorch_problem
+            self.assertIsInstance(botorch_problem, problem_cls)
             self.assertEqual(len(problem.search_space.parameters), dim)
             self.assertEqual(
                 sum(
@@ -51,12 +54,7 @@ def test_problems(self) -> None:
                 expected_bounds = [(-5.0, 10.0) for _ in range(dim)]
             else:
                 expected_bounds = [(0.0, 1.0) for _ in range(dim)]
-            self.assertEqual(
-                checked_cast(
-                    BotorchTestProblemRunner, problem.runner
-                ).test_problem._bounds,
-                expected_bounds,
-            )
+            self.assertEqual(botorch_problem._bounds, expected_bounds)
             self.assertGreaterEqual(problem.optimal_value, problem_cls().optimal_value)
 
         # Test that they match correctly to the original problems.
@@ -101,17 +99,18 @@ def test_problems(self) -> None:
         ]
 
         for problem, params, expected_arg in cases:
-            runner = checked_cast(BotorchTestProblemRunner, problem.runner)
+            runner = assert_is_instance(problem.runner, ParamBasedTestProblemRunner)
+            test_problem = assert_is_instance(runner.test_problem, BoTorchTestProblem)
             trial = Trial(experiment=MagicMock())
             # pyre-fixme: Incompatible parameter type [6]: In call
             # `Arm.__init__`, for argument `parameters`, expected `Dict[str,
             # Union[None, bool, float, int, str]]` but got `dict[str, float]`.
             arm = Arm(parameters=params, name="--")
             trial.add_arm(arm)
             with patch.object(
-                runner.test_problem,
+                test_problem.botorch_problem,
                 attribute="evaluate_true",
-                wraps=runner.test_problem.evaluate_true,
+                wraps=test_problem.botorch_problem.evaluate_true,
             ) as mock_call:
                 runner.run(trial)
             actual = mock_call.call_args[0][0]