Compute ground-truth optimization trace on BenchmarkProblem (#2704)

Summary:
Pull Request resolved: #2704

This currently doesn't change behavior, but enables more flexibility by redefining `Problem.get_ground_truth_opt_trace` or `BenchmarkRunner.get_Y_true`. This is currently duplicative with the tracking metrics setup, which will be reaped in D61432000; these should be combined. This diff has both ways of computing the optimization trace running, along with an assertion that they give the same results.

Differential Revision: D61404056

Reviewed By: saitcakmak

ax/benchmark/benchmark.py

@@ -151,6 +151,9 @@ def benchmark_replication(
         scheduler.get_trace(optimization_config=analysis_opt_config)
     )
 
+    new_optimization_trace = problem.get_opt_trace(experiment=experiment)
+    np.testing.assert_allclose(optimization_trace, new_optimization_trace)
+
     try:
         # Catch any errors that may occur during score computation, such as errors
         # while accessing "steps" in node based generation strategies. The error

ax/benchmark/benchmark_problem.py

@@ -8,10 +8,16 @@
 from dataclasses import dataclass, field
 from typing import Any, Optional, Union
 
+import numpy as np
+import pandas as pd
+
 from ax.benchmark.metrics.base import BenchmarkMetricBase
 
 from ax.benchmark.metrics.benchmark import BenchmarkMetric
+from ax.benchmark.runners.base import BenchmarkRunner
 from ax.benchmark.runners.botorch_test import BotorchTestProblemRunner
+from ax.core.data import Data
+from ax.core.experiment import Experiment
 from ax.core.objective import MultiObjective, Objective
 from ax.core.optimization_config import (
     MultiObjectiveOptimizationConfig,
@@ -20,9 +26,9 @@
 )
 from ax.core.outcome_constraint import OutcomeConstraint
 from ax.core.parameter import ParameterType, RangeParameter
-from ax.core.runner import Runner
 from ax.core.search_space import SearchSpace
 from ax.core.types import ComparisonOp
+from ax.service.utils.best_point_mixin import BestPointMixin
 from ax.utils.common.base import Base
 from ax.utils.common.typeutils import checked_cast
 from botorch.test_functions.base import (
@@ -31,6 +37,7 @@
     MultiObjectiveTestProblem,
 )
 from botorch.test_functions.synthetic import SyntheticTestFunction
+from pyre_extensions import assert_is_instance
 
 
 def _get_name(
@@ -89,9 +96,52 @@ class BenchmarkProblem(Base):
     optimal_value: float
 
     search_space: SearchSpace = field(repr=False)
-    runner: Runner = field(repr=False)
+    runner: BenchmarkRunner = field(repr=False)
     is_noiseless: bool
 
+    def get_oracle_experiment(self, experiment: Experiment) -> Experiment:
+        records = []
+
+        new_experiment = Experiment(
+            search_space=self.search_space, optimization_config=self.optimization_config
+        )
+        for trial_index, trial in experiment.trials.items():
+            for arm in trial.arms:
+                for metric_name, metric_value in zip(
+                    self.runner.outcome_names, self.runner.evaluate_oracle(arm=arm)
+                ):
+                    records.append(
+                        {
+                            "arm_name": arm.name,
+                            "metric_name": metric_name,
+                            "mean": metric_value.item(),
+                            "sem": 0.0,
+                            "trial_index": trial_index,
+                        }
+                    )
+
+            new_experiment.attach_trial(
+                parameterizations=[arm.parameters for arm in trial.arms],
+                arm_names=[arm.name for arm in trial.arms],
+            )
+        for trial in new_experiment.trials.values():
+            trial.mark_completed()
+
+        data = Data(df=pd.DataFrame.from_records(records))
+        new_experiment.attach_data(data=data, overwrite_existing_data=True)
+        return new_experiment
+
+    def get_opt_trace(self, experiment: Experiment) -> np.ndarray:
+        """Evaluate the optimization trace of a list of Trials."""
+        oracle_experiment = self.get_oracle_experiment(experiment=experiment)
+
+        return np.array(
+            BestPointMixin._get_trace(
+                experiment=oracle_experiment,
+                optimization_config=self.optimization_config,
+            )
+        )
+
 
 # TODO: Support constrained MOO problems.
 def get_soo_config_and_outcome_names(
@@ -141,6 +191,20 @@ def get_soo_config_and_outcome_names(
     return opt_config, outcome_names
 
 
+def get_continuous_search_space(bounds: list[tuple[float, float]]) -> SearchSpace:
+    return SearchSpace(
+        parameters=[
+            RangeParameter(
+                name=f"x{i}",
+                parameter_type=ParameterType.FLOAT,
+                lower=lower,
+                upper=upper,
+            )
+            for i, (lower, upper) in enumerate(bounds)
+        ]
+    )
+
+
 def create_single_objective_problem_from_botorch(
     test_problem_class: type[SyntheticTestFunction],
     test_problem_kwargs: dict[str, Any],
@@ -169,17 +233,7 @@ def create_single_objective_problem_from_botorch(
     test_problem = test_problem_class(**test_problem_kwargs)
     is_constrained = isinstance(test_problem, ConstrainedBaseTestProblem)
 
-    search_space = SearchSpace(
-        parameters=[
-            RangeParameter(
-                name=f"x{i}",
-                parameter_type=ParameterType.FLOAT,
-                lower=lower,
-                upper=upper,
-            )
-            for i, (lower, upper) in enumerate(test_problem._bounds)
-        ]
-    )
+    search_space = get_continuous_search_space(test_problem._bounds)
 
     dim = test_problem_kwargs.get("dim", None)
     name = _get_name(
@@ -249,18 +303,11 @@ def create_multi_objective_problem_from_botorch(
     # pyre-fixme [45]: Invalid class instantiation
     test_problem = test_problem_class(**test_problem_kwargs)
 
-    problem = create_single_objective_problem_from_botorch(
-        # pyre-fixme [6]: Passing a multi-objective problem where a
-        # single-objective problem is expected.
-        test_problem_class=test_problem_class,
-        test_problem_kwargs=test_problem_kwargs,
-        lower_is_better=True,  # Seems like we always assume minimization for MOO?
-        num_trials=num_trials,
-        observe_noise_sd=observe_noise_sd,
+    dim = test_problem_kwargs.get("dim", None)
+    name = _get_name(
+        test_problem=test_problem, observe_noise_sd=observe_noise_sd, dim=dim
     )
 
-    name = problem.name
-
     n_obj = test_problem.num_objectives
     if not observe_noise_sd:
         noise_sds = [None] * n_obj
@@ -292,15 +339,25 @@ def create_multi_objective_problem_from_botorch(
             for i, metric in enumerate(metrics)
         ],
     )
+    runner = BotorchTestProblemRunner(
+        test_problem_class=test_problem_class,
+        test_problem_kwargs=test_problem_kwargs,
+        outcome_names=[
+            objective.metric.name
+            for objective in assert_is_instance(
+                optimization_config.objective, MultiObjective
+            ).objectives
+        ],
+    )
 
     return MultiObjectiveBenchmarkProblem(
         name=name,
-        search_space=problem.search_space,
+        search_space=get_continuous_search_space(test_problem._bounds),
         optimization_config=optimization_config,
-        runner=problem.runner,
+        runner=runner,
         num_trials=num_trials,
-        is_noiseless=problem.is_noiseless,
+        is_noiseless=test_problem.noise_std in (None, 0.0),
         observe_noise_stds=observe_noise_sd,
-        has_ground_truth=problem.has_ground_truth,
+        has_ground_truth=True,
         optimal_value=test_problem.max_hv,
     )

ax/benchmark/runners/botorch_test.py

@@ -17,10 +17,7 @@
 from ax.utils.common.base import Base
 from ax.utils.common.equality import equality_typechecker
 from ax.utils.common.serialization import TClassDecoderRegistry, TDecoderRegistry
-from botorch.test_functions.synthetic import (
-    ConstrainedSyntheticTestFunction,
-    SyntheticTestFunction,
-)
+from botorch.test_functions.synthetic import BaseTestProblem, ConstrainedBaseTestProblem
 from botorch.utils.transforms import normalize, unnormalize
 from pyre_extensions import assert_is_instance
 from torch import Tensor
@@ -60,30 +57,30 @@ def __eq__(self, other: Any) -> bool:
 
 class SyntheticProblemRunner(BenchmarkRunner, ABC):
     """A Runner for evaluating synthetic problems, either BoTorch
-    `SyntheticTestFunction`s or Ax benchmarking `ParamBasedTestProblem`s.
+    `BaseTestProblem`s or Ax benchmarking `ParamBasedTestProblem`s.
 
     Given a trial, the Runner will evaluate the problem noiselessly for each
     arm in the trial, as well as return some metadata about the underlying
     problem such as the noise_std.
     """
 
-    test_problem: Union[SyntheticTestFunction, ParamBasedTestProblem]
+    test_problem: Union[BaseTestProblem, ParamBasedTestProblem]
     _is_constrained: bool
-    _test_problem_class: type[Union[SyntheticTestFunction, ParamBasedTestProblem]]
+    _test_problem_class: type[Union[BaseTestProblem, ParamBasedTestProblem]]
     _test_problem_kwargs: Optional[dict[str, Any]]
 
     def __init__(
         self,
         *,
-        test_problem_class: type[Union[SyntheticTestFunction, ParamBasedTestProblem]],
+        test_problem_class: type[Union[BaseTestProblem, ParamBasedTestProblem]],
         test_problem_kwargs: dict[str, Any],
         outcome_names: list[str],
         modified_bounds: Optional[list[tuple[float, float]]] = None,
     ) -> None:
         """Initialize the test problem runner.
 
         Args:
-            test_problem_class: A BoTorch `SyntheticTestFunction` class or Ax
+            test_problem_class: A BoTorch `BaseTestProblem` class or Ax
                 `ParamBasedTestProblem` class.
             test_problem_kwargs: The keyword arguments used for initializing the
                 test problem.
@@ -105,12 +102,12 @@ def __init__(
             # abstract class with abstract method `evaluate_true`.
             test_problem_class(**test_problem_kwargs)
         )
-        if isinstance(self.test_problem, SyntheticTestFunction):
+        if isinstance(self.test_problem, BaseTestProblem):
             self.test_problem = self.test_problem.to(dtype=torch.double)
-        # A `ConstrainedSyntheticTestFunction` is a type of `SyntheticTestFunction`; a
+        # A `ConstrainedBaseTestProblem` is a type of `BaseTestProblem`; a
         # `ParamBasedTestProblem` is never constrained.
         self._is_constrained: bool = isinstance(
-            self.test_problem, ConstrainedSyntheticTestFunction
+            self.test_problem, ConstrainedBaseTestProblem
         )
         self._is_moo: bool = self.test_problem.num_objectives > 1
         self.outcome_names = outcome_names
@@ -202,10 +199,10 @@ def deserialize_init_args(
 
 class BotorchTestProblemRunner(SyntheticProblemRunner):
     """
-    A `SyntheticProblemRunner` for BoTorch `SyntheticTestFunction`s.
+    A `SyntheticProblemRunner` for BoTorch `BaseTestProblem`s.
 
     Args:
-        test_problem_class: A BoTorch `SyntheticTestFunction` class.
+        test_problem_class: A BoTorch `BaseTestProblem` class.
         test_problem_kwargs: The keyword arguments used for initializing the
             test problem.
         outcome_names: The names of the outcomes returned by the problem.
@@ -223,7 +220,7 @@ class BotorchTestProblemRunner(SyntheticProblemRunner):
     def __init__(
         self,
         *,
-        test_problem_class: type[SyntheticTestFunction],
+        test_problem_class: type[BaseTestProblem],
         test_problem_kwargs: dict[str, Any],
         outcome_names: list[str],
         modified_bounds: Optional[list[tuple[float, float]]] = None,
@@ -234,11 +231,9 @@ def __init__(
             outcome_names=outcome_names,
             modified_bounds=modified_bounds,
         )
-        self.test_problem: SyntheticTestFunction = self.test_problem.to(
-            dtype=torch.double
-        )
+        self.test_problem: BaseTestProblem = self.test_problem.to(dtype=torch.double)
         self._is_constrained: bool = isinstance(
-            self.test_problem, ConstrainedSyntheticTestFunction
+            self.test_problem, ConstrainedBaseTestProblem
         )
 
     def get_Y_true(self, arm: Arm) -> Tensor:
@@ -274,7 +269,7 @@ def get_Y_true(self, arm: Arm) -> Tensor:
             X = unnormalize(unit_X, self.test_problem.bounds)
 
         Y_true = self.test_problem.evaluate_true(X).view(-1)
-        # `SyntheticTestFunction.evaluate_true()` does not negate the outcome
+        # `BaseTestProblem.evaluate_true()` does not negate the outcome
         if self.test_problem.negate:
             Y_true = -Y_true
 

ax/benchmark/tests/problems/synthetic/hss/test_jenatton.py

@@ -14,7 +14,6 @@
     get_jenatton_benchmark_problem,
     jenatton_test_function,
 )
-from ax.benchmark.runners.base import BenchmarkRunner
 from ax.benchmark.runners.botorch_test import ParamBasedTestProblemRunner
 from ax.core.arm import Arm
 from ax.core.data import Data
@@ -96,9 +95,7 @@ def test_jenatton_test_function(self) -> None:
                 value,
             )
             self.assertAlmostEqual(
-                assert_is_instance(benchmark_problem.runner, BenchmarkRunner)
-                .get_Y_true(arm)
-                .item(),
+                benchmark_problem.runner.evaluate_oracle(arm).item(),
                 value,
                 places=6,
             )

ax/benchmark/tests/problems/test_surrogate_problems.py

@@ -34,7 +34,7 @@ def test_repr(self) -> None:
             "SOOSurrogateBenchmarkProblem(name='test', "
             "optimization_config=OptimizationConfig(objective=Objective(metric_name="
             '"branin", '
-            "minimize=False), "
+            "minimize=True), "
             "outcome_constraints=[]), num_trials=6, "
             "observe_noise_stds=True, has_ground_truth=True, "
             "tracking_metrics=[], optimal_value=0.0, is_noiseless=True)"