Runner should only need to know about outcomes, not objectives vs. co…

…nstraints

Summary:
Context: In theory, a `BenchmarkRunner` should not have to know what metrics are objectives or constraints, and a test function should not have to be aware of that either. They are just generating data. That should be handled by the `OptimizationConfig`, so that various setups can be used without changing the runner and test function.

For historical reasons, runners track objectives and constraints separately and add noise to them separately, because this mimics how BoTorch test functions handle this. However, we now can and should isolate the quirks of BoTorch test functions within `BoTorchTestProblem`.

This diff:
* Removes the argument `constraint_noise_std` from `create_problem_from_botorch` and from `ParamBasedTestProblemRunner`, in favor of just using `noise_std`.
* Renames `num_objectives` to `num_outcomes` in `ParamBasedTestProblemRunner`
* Removes `evaluate_slack_true` from `ParamBasedTestProblem`, changing `evaluate_true` to return both objectives and constraints
* Updates validation

Tangentially related changes:
* For simplicity, makes `get_noise_stds` always return a dict
* Stops allowing `noise_std` to be `None` and defaults it to zero (it was eventually set to zero when it was None in the past)

Differential Revision: D64919207

ax/benchmark/benchmark_problem.py

@@ -301,8 +301,7 @@ def create_problem_from_botorch(
     *,
     test_problem_class: type[BaseTestProblem],
     test_problem_kwargs: dict[str, Any],
-    noise_std: float | list[float] | None = None,
-    constraint_noise_std: float | list[float] | None = None,
+    noise_std: float | list[float] = 0.0,
     num_trials: int,
     lower_is_better: bool = True,
     observe_noise_sd: bool = False,
@@ -321,11 +320,8 @@ def create_problem_from_botorch(
             to define the `search_space`, `optimization_config`, and `runner`.
         test_problem_kwargs: Keyword arguments used to instantiate the
             `test_problem_class`.
-        noise_std: Standard deviation of synthetic noise added to objectives. If
-            `None`, no noise is added. If a float, the same noise level is used
-            for all objectives.
-        constraint_noise_std: Standard deviation of synthetic noise added to
-            constraints.
+        noise_std: Standard deviation of synthetic noise added to outcomes. If a
+            float, the same noise level is used for all objectives.
         lower_is_better: Whether this is a minimization problem. For MOO, this
             applies to all objectives.
         num_trials: Simply the `num_trials` of the `BenchmarkProblem` created.
@@ -392,7 +388,6 @@ def create_problem_from_botorch(
                 param_names=list(search_space.parameters.keys()),
             ),
             noise_std=noise_std,
-            constraint_noise_std=constraint_noise_std,
         ),
         num_trials=num_trials,
         observe_noise_stds=observe_noise_sd,

ax/benchmark/problems/hpo/torchvision.py

@@ -118,7 +118,7 @@ def train_and_evaluate(
 @dataclass(kw_only=True)
 class PyTorchCNNTorchvisionParamBasedProblem(ParamBasedTestProblem):
     name: str  # The name of the dataset to load -- MNIST or FashionMNIST
-    num_objectives: int = 1
+    num_outcomes: int = 1
     device: torch.device = field(
         default_factory=lambda: torch.device(
             "cuda" if torch.cuda.is_available() else "cpu"

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

@@ -55,7 +55,7 @@ def jenatton_test_function(
 class Jenatton(ParamBasedTestProblem):
     """Jenatton test function for hierarchical search spaces."""
 
-    num_objectives: int = 1
+    num_outcomes: int = 1
 
     # pyre-fixme[14]: Inconsistent override
     def evaluate_true(self, params: Mapping[str, float | int | None]) -> torch.Tensor:

ax/benchmark/runners/base.py

@@ -84,7 +84,7 @@ def evaluate_oracle(self, parameters: Mapping[str, TParamValue]) -> ndarray:
         return self.get_Y_true(params=params).numpy()
 
     @abstractmethod
-    def get_noise_stds(self) -> None | float | dict[str, float]:
+    def get_noise_stds(self) -> dict[str, float]:
         """
         Return the standard errors for the synthetic noise to be applied to the
         observed values.

ax/benchmark/runners/botorch_test.py

@@ -12,8 +12,8 @@
 
 import torch
 from ax.benchmark.runners.base import BenchmarkRunner
+from ax.core.search_space import SearchSpaceDigest
 from ax.core.types import TParamValue
-from ax.exceptions.core import UnsupportedError
 from botorch.test_functions.multi_objective import MultiObjectiveTestProblem
 from botorch.test_functions.synthetic import BaseTestProblem, ConstrainedBaseTestProblem
 from botorch.utils.transforms import normalize, unnormalize
@@ -28,17 +28,17 @@ class ParamBasedTestProblem(ABC):
     (Noise - if desired - is added by the runner.)
     """
 
-    num_objectives: int
+    num_outcomes: int
 
     @abstractmethod
     def evaluate_true(self, params: Mapping[str, TParamValue]) -> Tensor:
-        """Evaluate noiselessly."""
-        ...
+        """
+        Evaluate noiselessly.
 
-    def evaluate_slack_true(self, params: Mapping[str, TParamValue]) -> Tensor:
-        raise NotImplementedError(
-            f"{self.__class__.__name__} does not support constraints."
-        )
+        Returns:
+            1d tensor of shape (num_outcomes,).
+        """
+        ...
 
 
 @dataclass(kw_only=True)
@@ -57,24 +57,32 @@ class BoTorchTestProblem(ParamBasedTestProblem):
             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.
-        num_objectives: The number of objectives.
+        num_outcomes: The number of outcomes (objectives + constraints).
     """
 
     botorch_problem: BaseTestProblem
     modified_bounds: list[tuple[float, float]] | None = None
-    num_objectives: int = 1
+    num_outcomes: 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 getattr(self.botorch_problem, "constraint_noise_std", None) is not None:
+        num_objectives = (
+            self.botorch_problem.num_objectives
+            if isinstance(self.botorch_problem, MultiObjectiveTestProblem)
+            else 1
+        )
+        num_constraints = (
+            self.botorch_problem.num_constraints
+            if isinstance(self.botorch_problem, ConstrainedBaseTestProblem)
+            else 0
+        )
+        self.num_outcomes = num_objectives + num_constraints
+
+        if (
+            self.botorch_problem.noise_std is not None
+            or getattr(self.botorch_problem, "constraint_noise_std", None) is not None
+        ):
             raise ValueError(
-                "constraint_noise_std should be set on the runner, not the test "
-                "problem."
+                "noise should be set on the `BenchmarkRunner`, not the test function."
             )
         self.botorch_problem = self.botorch_problem.to(dtype=torch.double)
 
@@ -94,22 +102,14 @@ def tensorize_params(self, params: Mapping[str, int | float]) -> torch.Tensor:
         return X
 
     # pyre-fixme [14]: inconsistent override
+    # TODO: test for MOO
     def evaluate_true(self, params: Mapping[str, float | int]) -> torch.Tensor:
         x = self.tensorize_params(params=params)
-        return self.botorch_problem(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 UnsupportedError(
-                "`evaluate_slack_true` is only supported when the BoTorch "
-                "problem is a `ConstrainedBaseTestProblem`."
-            )
-        # 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)
+        objectives = self.botorch_problem(x).view(-1)
+        if isinstance(self.botorch_problem, ConstrainedBaseTestProblem):
+            constraints = self.botorch_problem.evaluate_slack_true(x).view(-1)
+            return torch.cat([objectives, constraints], dim=-1)
+        return objectives
 
 
 @dataclass(kw_only=True)
@@ -119,7 +119,7 @@ class ParamBasedTestProblemRunner(BenchmarkRunner):
 
     Given a trial, the Runner will use its `test_problem` to evaluate the
     problem noiselessly for each arm in the trial, and then add noise as
-    specified by the `noise_std` and `constraint_noise_std`. It will return
+    specified by the `noise_std`. It will return
     metadata including the outcome names and values of metrics.
 
     Args:
@@ -132,64 +132,36 @@ class ParamBasedTestProblemRunner(BenchmarkRunner):
     """
 
     test_problem: ParamBasedTestProblem
-    noise_std: float | list[float] | None = None
-    constraint_noise_std: float | list[float] | None = None
+    noise_std: float | list[float] | dict[str, float] = 0.0
 
-    @property
-    def _is_constrained(self) -> bool:
-        return isinstance(self.test_problem, BoTorchTestProblem) and isinstance(
-            self.test_problem.botorch_problem, ConstrainedBaseTestProblem
-        )
+    def __post_init__(self, search_space_digest: SearchSpaceDigest | None) -> None:
+        super().__post_init__(search_space_digest)
+        if len(self.outcome_names) != self.test_problem.num_outcomes:
+            raise ValueError(
+                f"Number of outcomes must match number of outcomes in test problem. {self.outcome_names=}, {self.test_problem.num_outcomes=}"
+            )
 
-    def get_noise_stds(self) -> None | float | dict[str, float]:
+    def get_noise_stds(self) -> dict[str, float]:
         noise_std = self.noise_std
-        noise_std_dict: dict[str, float] = {}
-        num_obj = self.test_problem.num_objectives
-
-        # populate any noise_stds for constraints
-        if self._is_constrained:
-            constraint_noise_std = self.constraint_noise_std
-            if isinstance(constraint_noise_std, list):
-                for i, cns in enumerate(constraint_noise_std, start=num_obj):
-                    if cns is not None:
-                        noise_std_dict[self.outcome_names[i]] = cns
-            elif constraint_noise_std is not None:
-                noise_std_dict[self.outcome_names[num_obj]] = constraint_noise_std
-
-        # if none of the constraints are subject to noise, then we may return
-        # a single float or None for the noise level
-
-        if not noise_std_dict and not isinstance(noise_std, list):
-            return noise_std  # either a float or None
-
-        if isinstance(noise_std, list):
-            if not len(noise_std) == num_obj:
-                # this shouldn't be possible due to validation upon construction
-                # of the multi-objective problem, but better safe than sorry
+        if isinstance(noise_std, float):
+            return {name: noise_std for name in self.outcome_names}
+        elif isinstance(noise_std, dict):
+            if not set(noise_std.keys()) == set(self.outcome_names):
                 raise ValueError(
-                    "Noise std must have length equal to number of objectives."
+                    "Noise std must have keys equal to outcome names if given as a dict."
                 )
-        else:
-            noise_std = [noise_std for _ in range(num_obj)]
-
-        for i, noise_std_ in enumerate(noise_std):
-            if noise_std_ is not None:
-                noise_std_dict[self.outcome_names[i]] = noise_std_
-
-        return noise_std_dict
+            return noise_std
+        # list of floats
+        if not len(noise_std) == self.test_problem.num_outcomes:
+            raise ValueError(
+                f"`noise_std` must have length equal to number of outcomes. {noise_std=}, {self.test_problem.num_outcomes=}"
+            )
+        return dict(zip(self.outcome_names, noise_std))
 
     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.
+            An `m`-dim tensor of ground truth (noiseless) evaluations.
         """
-        Y_true = self.test_problem.evaluate_true(params).view(-1)
-        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(params).view(-1)],
-                dim=-1,
-            )
-        return Y_true
+        return torch.atleast_1d(self.test_problem.evaluate_true(params=params))

ax/benchmark/runners/surrogate.py

@@ -84,8 +84,11 @@ def datasets(self) -> list[SupervisedDataset]:
             self.set_surrogate_and_datasets()
         return none_throws(self._datasets)
 
-    def get_noise_stds(self) -> None | float | dict[str, float]:
-        return self.noise_stds
+    def get_noise_stds(self) -> dict[str, float]:
+        noise_std = self.noise_stds
+        if isinstance(noise_std, float):
+            return {name: noise_std for name in self.outcome_names}
+        return noise_std
 
     # pyre-fixme[14]: Inconsistent override
     def get_Y_true(self, params: Mapping[str, float | int]) -> Tensor: