Bandit dispatcher class

baybe/surrogates/__init__.py

@@ -1,6 +1,9 @@
 """BayBE surrogates."""
 
-from baybe.surrogates.bandit import BetaBernoulliMultiArmedBanditSurrogate
+from baybe.surrogates.bandits.base import MultiArmedBanditSurrogate
+from baybe.surrogates.bandits.beta_bernoulli import (
+    BetaBernoulliMultiArmedBanditSurrogate,
+)
 from baybe.surrogates.custom import CustomONNXSurrogate, register_custom_architecture
 from baybe.surrogates.gaussian_process.core import GaussianProcessSurrogate
 from baybe.surrogates.linear import BayesianLinearSurrogate
@@ -15,6 +18,7 @@
     "CustomONNXSurrogate",
     "GaussianProcessSurrogate",
     "MeanPredictionSurrogate",
+    "MultiArmedBanditSurrogate",
     "NGBoostSurrogate",
     "RandomForestSurrogate",
 ]

baybe/surrogates/bandits/__init__.py

@@ -0,0 +1,11 @@
+"""Bandit surrogates."""
+
+from baybe.surrogates.bandits.base import MultiArmedBanditSurrogate
+from baybe.surrogates.bandits.beta_bernoulli import (
+    BetaBernoulliMultiArmedBanditSurrogate,
+)
+
+__all__ = [
+    "BetaBernoulliMultiArmedBanditSurrogate",
+    "MultiArmedBanditSurrogate",
+]

baybe/surrogates/bandits/base.py

@@ -0,0 +1,57 @@
+"""Multi-armed bandit surrogates."""
+
+from __future__ import annotations
+
+from attrs import define, field
+from attrs.validators import instance_of
+from pandas import DataFrame
+
+from baybe.objectives.base import Objective
+from baybe.objectives.single import SingleTargetObjective
+from baybe.priors.base import Prior
+from baybe.searchspace.core import SearchSpace
+from baybe.surrogates.bandits.beta_bernoulli import (
+    BetaBernoulliMultiArmedBanditSurrogate,
+)
+from baybe.surrogates.base import Surrogate
+from baybe.targets.binary import BinaryTarget
+
+
+def _get_bandit_class(
+    searchspace: SearchSpace, objective: Objective
+) -> type[BetaBernoulliMultiArmedBanditSurrogate]:
+    """Retrieve the appropriate bandit class for the given modelling context."""
+    match searchspace, objective:
+        case _, SingleTargetObjective(_target=BinaryTarget()):
+            return BetaBernoulliMultiArmedBanditSurrogate
+        case _:
+            raise NotImplementedError(
+                f"Currently, only a single target of type '{BinaryTarget.__name__}' "
+                f"is supported."
+            )
+
+
+@define
+class MultiArmedBanditSurrogate:
+    """A bandit surrogate class dispatching class.
+
+    Follows the strategy design pattern to dispatch to the appropriate bandit model.
+    """
+
+    prior: Prior = field(validator=instance_of(Prior))
+    """The prior distribution assumed for each arm of the bandit."""
+
+    _bandit_model: Surrogate | None = field(init=False, default=None, eq=False)
+    """The specific bandit model to which is being dispatched."""
+
+    def fit(
+        self, searchspace: SearchSpace, objective: Objective, measurements: DataFrame
+    ) -> None:
+        """Instantiate an appropriate bandit model and fit it to the data."""
+        cls = _get_bandit_class(searchspace, objective)
+        self._bandit_model = cls(self.prior)
+        self._bandit_model.fit(searchspace, objective, measurements)
+
+    def __getattr__(self, name):
+        # If the attribute is not found, try to get it from the bandit object
+        return getattr(self._bandit_model, name)

baybe/surrogates/bandit.py → baybe/surrogates/bandits/beta_bernoulli.py

@@ -1,10 +1,11 @@
-"""Multi-armed bandit surrogates."""
+"""A Beta-Bernoulli multi-armed bandit surrogate."""
 
 from __future__ import annotations
 
 from typing import TYPE_CHECKING, Any, ClassVar
 
 from attrs import define, field
+from attrs.validators import instance_of
 
 from baybe.exceptions import IncompatibleSearchSpaceError, ModelNotTrainedError
 from baybe.parameters.categorical import CategoricalParameter
@@ -28,7 +29,9 @@ class BetaBernoulliMultiArmedBanditSurrogate(Surrogate):
     supports_transfer_learning: ClassVar[bool] = False
     # See base class.
 
-    prior: BetaPrior = field(factory=lambda: BetaPrior(1, 1))
+    prior: BetaPrior = field(
+        factory=lambda: BetaPrior(1, 1), validator=instance_of(BetaPrior)
+    )
     """The beta prior for the win rates of the bandit arms. Uniform by default."""
 
     # TODO: type should be `torch.Tensor | None` but is currently

examples/Multi_Armed_Bandit/bernoulli_multi_armed_bandit.py

@@ -17,12 +17,13 @@
 from baybe.acquisition import PosteriorStandardDeviation, qThompsonSampling
 from baybe.acquisition.base import AcquisitionFunction
 from baybe.parameters import CategoricalParameter
+from baybe.priors.basic import BetaPrior
 from baybe.recommenders import (
     BotorchRecommender,
     RandomRecommender,
     TwoPhaseMetaRecommender,
 )
-from baybe.surrogates import BetaBernoulliMultiArmedBanditSurrogate
+from baybe.surrogates import MultiArmedBanditSurrogate
 from baybe.targets import BinaryTarget
 
 ### An Imaginary Use Case
@@ -97,7 +98,7 @@ def draw_arm(self, arm_index: int) -> bool:
     # For bandits, one-hot parameter encoding is required:
     encoding="OHE",
 )
-surrogate = BetaBernoulliMultiArmedBanditSurrogate()
+surrogate = MultiArmedBanditSurrogate(BetaPrior(1, 1))
 
 
 # For each simulation, we report the trajectory of earned rewards and the estimated win rates of the bandit arms:

tests/test_iterations.py

@@ -37,7 +37,7 @@
 )
 from baybe.recommenders.pure.nonpredictive.base import NonPredictiveRecommender
 from baybe.searchspace import SearchSpaceType
-from baybe.surrogates.bandit import BetaBernoulliMultiArmedBanditSurrogate
+from baybe.surrogates.bandits import BetaBernoulliMultiArmedBanditSurrogate
 from baybe.surrogates.base import IndependentGaussianSurrogate, Surrogate
 from baybe.surrogates.custom import CustomONNXSurrogate
 from baybe.surrogates.gaussian_process.presets import (