Source code for ax.benchmark.runners.base

# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

# pyre-strict

from abc import ABC, abstractmethod
from math import sqrt
from typing import Any, Dict, List, Optional, Tuple, Union

import torch
from ax.core.arm import Arm
from ax.core.base_trial import BaseTrial
from ax.core.batch_trial import BatchTrial
from ax.core.runner import Runner
from ax.core.trial import Trial

from ax.utils.common.typeutils import checked_cast
from torch import Tensor


[docs]class BenchmarkRunner(Runner, ABC): @property @abstractmethod def outcome_names(self) -> List[str]: """The names of the outcomes of the problem (in the order of the outcomes).""" pass # pragma: no cover
[docs] def get_Y_true(self, arm: Arm) -> Tensor: """Function returning the ground truth values for a given arm. The synthetic noise is added as part of the Runner's `run()` method. For problems that do not have a ground truth, the Runner must implement the `get_Y_Ystd()` method instead.""" raise NotImplementedError( "Must implement method `get_Y_true()` for Runner " f"{self.__class__.__name__} as it does not implement a " "`get_Y_Ystd()` method." )
[docs] def get_noise_stds(self) -> Union[None, float, Dict[str, float]]: """Function returning the standard errors for the synthetic noise to be applied to the observed values. For problems that do not have a ground truth, the Runner must implement the `get_Y_Ystd()` method instead.""" raise NotImplementedError( "Must implement method `get_Y_Ystd()` for Runner " f"{self.__class__.__name__} as it does not implement a " "`get_noise_stds()` method." )
[docs] def get_Y_Ystd(self, arm: Arm) -> Tuple[Tensor, Optional[Tensor]]: """Function returning the observed values and their standard errors for a given arm. This function is unused for problems that have a ground truth (in this case `get_Y_true()` is used), and is required for problems that do not have a ground truth.""" raise NotImplementedError( "Must implement method `get_Y_Ystd()` for Runner " f"{self.__class__.__name__} as it does not implement a " "`get_Y_true()` method." )
[docs] def run(self, trial: BaseTrial) -> Dict[str, Any]: """Run the trial by evaluating its parameterization(s). Args: trial: The trial to evaluate. Returns: A dictionary with the following keys: - Ys: A dict mapping arm names to lists of corresponding outcomes, where the order of the outcomes is the same as in `outcome_names`. - Ystds: A dict mapping arm names to lists of corresponding outcome noise standard deviations (possibly nan if the noise level is unobserved), where the order of the outcomes is the same as in `outcome_names`. - Ys_true: A dict mapping arm names to lists of corresponding ground truth outcomes, where the order of the outcomes is the same as in `outcome_names`. If the benchmark problem does not provide a ground truth, this key will not be present in the dict returned by this function. - "outcome_names": A list of metric names. """ Ys, Ys_true, Ystds = {}, {}, {} noise_stds = self.get_noise_stds() if noise_stds is not None: # extract arm weights to adjust noise levels accordingly if isinstance(trial, BatchTrial): # normalize arm weights (we assume that the noise level is defined) # w.r.t. to a single arm allocated all of the sample budget nlzd_arm_weights = { arm: weight / sum(trial.arm_weights.values()) for arm, weight in trial.arm_weights.items() } else: nlzd_arm_weights = {checked_cast(Trial, trial).arm: 1.0} # generate a tensor of noise levels that we'll reuse below if isinstance(noise_stds, float): noise_stds_tsr = torch.full( (len(self.outcome_names),), noise_stds, dtype=torch.double, ) else: noise_stds_tsr = torch.tensor( [noise_stds[metric_name] for metric_name in self.outcome_names], dtype=torch.double, ) for arm in trial.arms: try: # Case where we do have a ground truth Y_true = self.get_Y_true(arm) Ys_true[arm.name] = Y_true.tolist() if noise_stds is None: # No noise, so just return the true outcome. Ystds[arm.name] = [0.0] * len(Y_true) Ys[arm.name] = Y_true.tolist() else: # We can scale the noise std by the inverse of the relative sample # budget allocation to each arm. This works b/c (i) we assume that # observations per unit sample budget are i.i.d. and (ii) the # normalized weights sum to one. std = noise_stds_tsr.to(Y_true) / sqrt(nlzd_arm_weights[arm]) Ystds[arm.name] = std.tolist() Ys[arm.name] = (Y_true + std * torch.randn_like(Y_true)).tolist() except NotImplementedError: # Case where we don't have a ground truth. Y, Ystd = self.get_Y_Ystd(arm) Ys[arm.name] = Y.tolist() Ystds[arm.name] = Ystd.tolist() if Ystd is not None else None run_metadata = { "Ys": Ys, "Ystds": Ystds, "outcome_names": self.outcome_names, } if Ys_true: # only add key if we actually have a ground truth run_metadata["Ys_true"] = Ys_true return run_metadata