feat: using nsgaii for BO acqf optimization

bbo/algorithms/bo.py

@@ -1,3 +1,4 @@
+import logging
 from typing import List, Sequence, Optional, Union
 
 from attrs import define, field, validators
@@ -7,21 +8,26 @@
 import botorch
 from botorch import fit_gpytorch_mll
 from botorch.models import SingleTaskGP
-from botorch.acquisition import ExpectedImprovement
 import gpytorch
 from gpytorch.constraints import GreaterThan
 
 from bbo.algorithms.base import Designer
 from bbo.algorithms.random import RandomDesigner
 from bbo.utils.converters.converter import SpecType, DefaultTrialConverter
-from bbo.utils.problem_statement import ProblemStatement
+from bbo.utils.metric_config import ObjectiveMetricGoal
+from bbo.utils.problem_statement import ProblemStatement, Objective
 from bbo.utils.trial import Trial
 from bbo.algorithms.bo_utils.mean_factory import mean_factory
 from bbo.algorithms.bo_utils.kernel_factory import kernel_factory
+from bbo.algorithms.bo_utils.acqf_factory import acqf_factory
+from bbo.algorithms.evolution.nsgaii import NSGAIIDesigner
+from bbo.benchmarks.experimenters.torch_experimenter import TorchExperimenter
+
+logger = logging.getLogger(__name__)
 
 
 @define
-class BO(Designer):
+class BODesigner(Designer):
     _problem_statement: ProblemStatement = field(
         validator=validators.instance_of(ProblemStatement)
     )
@@ -30,19 +36,38 @@ class BO(Designer):
     _device: str = field(default='cpu', kw_only=True)
 
     # surrogate model configuration
-    _mean_type: Optional[str] = field(default=None, kw_only=True)
+    _mean_type: Optional[str] = field(
+        default=None, kw_only=True,
+        validator=validators.optional(validators.in_(['constant', 'mlp']))
+    )
     _mean_config: Optional[dict] = field(default=None, kw_only=True)
-    _kernel_type: Optional[str] = field(default=None, kw_only=True)
+    _kernel_type: Optional[str] = field(
+        default=None, kw_only=True,
+        validator=validators.optional(validators.in_(['matern52', 'mlp', 'kumar']))
+    )
     _kernel_config: Optional[dict] = field(default=None, kw_only=True)
 
     # surrogate model optimization configuration
-    _mll_optimizer: str = field(default='l-bfgs', kw_only=True)
+    _mll_optimizer: str = field(
+        default='l-bfgs', kw_only=True,
+        validator=validators.in_(['l-bfgs', 'adam'])
+    )
     _mll_lr: Optional[float] = field(default=None, kw_only=True)
     _mll_epochs: Optional[int] = field(default=None, kw_only=True)
 
     # acquisition function configuration
-    _acqf_type: Union[str, List[str]] = field(default='EI', kw_only=True)
-    _acqf_optimizer: str = field(default='l-bfgs', kw_only=True)
+    _acqf_type: Union[str, List[str]] = field(
+        default='qEI', kw_only=True,
+        validator=validators.or_(
+            validators.in_(['qEI', 'qUCB', 'qPI', 'qlogEI']),
+            validators.deep_iterable(validators.in_(['qEI', 'qUCB', 'qPI', 'qlogEI'])),
+        )
+    )
+    _acqf_optimizer: str = field(
+        default='l-bfgs', kw_only=True,
+        validator=validators.in_(['l-bfgs', 'nsgaii'])
+    )
+    _acqf_config: dict = field(factory=dict, kw_only=True)
 
     # internal attributes
     _trials: List[Trial] = field(factory=list, init=False)
@@ -63,6 +88,8 @@ def __attrs_post_init__(self):
     def create_model(self, train_X, train_Y):
         mean_module = mean_factory(self._mean_type, self._mean_config)
         covar_module = kernel_factory(self._kernel_type, self._kernel_config)
+        # logger.info('mean_module: {}'.format(mean_module))
+        # logger.info('covar_module: {}'.format(covar_module))
         model = SingleTaskGP(train_X, train_Y, covar_module=covar_module, mean_module=mean_module)
         model.likelihood.noise_covar.register_constraint('raw_noise', GreaterThan(1e-4))
         mll = gpytorch.mlls.ExactMarginalLogLikelihood(model.likelihood, model)
@@ -89,20 +116,13 @@ def optimize_model(self, mll, model, train_X, train_Y):
             raise NotImplementedError
 
     def create_acqf(self, model, train_X, train_Y):
-        def get_acqf(acqf_type, model, train_X, train_Y):
-            if acqf_type == 'EI':
-                acqf = ExpectedImprovement(model, train_Y.max())
-            else:
-                raise NotImplementedError
-            return acqf
-
         if isinstance(self._acqf_type, list):
             acqf = []
             for acqf_type in self._acqf_type:
-                acqf_tmp = get_acqf(acqf_type, model, train_X, train_Y)
+                acqf_tmp = acqf_factory(acqf_type, model, train_X, train_Y)
                 acqf.append(acqf_tmp)
         else:
-            acqf = get_acqf(self._acqf_type, model, train_X, train_Y)
+            acqf = acqf_factory(self._acqf_type, model, train_X, train_Y)
 
         return acqf
 
@@ -112,6 +132,58 @@ def optimize_acqf(self, acqf):
             next_X, _ = botorch.optim.optimize.optimize_acqf(
                 acqf, bounds=bounds, q=self._q, num_restarts=10, raw_samples=1024
             )
+        elif self._acqf_optimizer == 'nsgaii':
+            sp = self._problem_statement.search_space
+            obj = Objective()
+            if isinstance(self._acqf_type, list):
+                for name in self._acqf_type:
+                    obj.add_metric(name, ObjectiveMetricGoal.MAXIMIZE)
+            else:
+                obj.add_metric(self._acqf_type, ObjectiveMetricGoal.MAXIMIZE)
+            if obj.num_metrics() <= 1:
+                logger.warning('NSGA-II is a multi-objective optimization algorithm, but only single objective is defined')
+            nsgaii_problem_statement = ProblemStatement(sp, obj)
+            nsgaii_designer = NSGAIIDesigner(
+                nsgaii_problem_statement,
+                pop_size=self._acqf_config.get('pop_size', 20),
+                n_offsprings=self._acqf_config.get('n_offsprings', None),
+            )
+            def acqf_obj(x, acqf):
+                y = []
+                for acqf_tmp in acqf:
+                    y.append(acqf_tmp(x.unsqueeze(1)).unsqueeze(-1))
+                y = torch.hstack(y)
+                return y
+            experimenter = TorchExperimenter(lambda x: acqf_obj(x, acqf), nsgaii_problem_statement)
+            for _ in range(self._acqf_config.get('epochs', 200)):
+                trials = nsgaii_designer.suggest()
+                experimenter.evaluate(trials)
+                nsgaii_designer.update(trials)
+
+            # generate next_X for batch BO setting
+            pareto_X, _ = nsgaii_designer.result()
+            pop_X, _ = nsgaii_designer.curr_pop()
+            diff_X = [x for x in pop_X if x not in pareto_X]
+            diff_X = np.zeros((0, pareto_X.shape[-1])) if len(diff_X) == 0 else np.vstack(diff_X)
+
+            if len(pareto_X) >= self._q:
+                idx = np.random.choice(len(pareto_X), self._q, replace=False)
+                next_X = torch.from_numpy(pareto_X[idx])
+            else:
+                next_X = [pareto_X]
+                if len(diff_X) > 0:
+                    quota = min(len(diff_X), self._q-len(pareto_X))
+                    idx = np.random.choice(len(diff_X), quota, replace=False)
+                    next_X.append(diff_X[idx])
+                quota = self._q - np.vstack(next_X).shape[0]
+                if quota > 0:
+                    trials = self._init_designer.suggest(quota)
+                    features = self._converter.to_features(trials)
+                    random_X = []
+                    for name in self._converter.input_converter_dict:
+                        random_X.append(features[name])
+                    next_X.append(np.hstack(random_X))
+                next_X = torch.from_numpy(np.vstack(next_X))
         else:
             raise NotImplementedError
 

bbo/algorithms/bo_utils/acqf_factory.py

@@ -0,0 +1,20 @@
+from botorch.acquisition import (
+    qExpectedImprovement,
+    qUpperConfidenceBound,
+    qProbabilityOfImprovement,
+    qLogExpectedImprovement
+)
+
+
+def acqf_factory(acqf_type, model, train_X, train_Y):
+    if acqf_type == 'qEI':
+        acqf = qExpectedImprovement(model, train_Y.max())
+    elif acqf_type == 'qUCB':
+        acqf = qUpperConfidenceBound(model, beta=0.18)
+    elif acqf_type == 'qPI':
+        acqf = qProbabilityOfImprovement(model, train_Y.max())
+    elif acqf_type == 'qlogEI':
+        acqf = qLogExpectedImprovement(model, train_Y.max())
+    else:
+        raise NotImplementedError
+    return acqf

bbo/algorithms/evolution/nsgaii.py

@@ -2,7 +2,7 @@
 
 from attrs import define, field, validators
 import numpy as np
-from pymoo.algorithms.moo.nsga2 import NSGA2
+from pymoo.algorithms.moo.nsga2 import NSGA2 as pymoo_NSGA2
 from pymoo.operators.crossover.sbx import SBX
 from pymoo.operators.mutation.pm import PM
 from pymoo.operators.sampling.rnd import FloatRandomSampling
@@ -19,20 +19,21 @@
 
 
 @define
-class NSGAII(Designer):
+class NSGAIIDesigner(Designer):
     _problem_statement: ProblemStatement = field(
         validator=validators.instance_of(ProblemStatement)
     )
     _pop_size: int = field(default=20)
+    _n_offsprings: Optional[int] = field(default=None)
 
     # internal attributes
     _last_pop = field(default=None, init=False)
 
     def __attrs_post_init__(self):
         self._converter = DefaultTrialConverter.from_problem(self._problem_statement)
-        self._impl = NSGA2(
+        self._impl = pymoo_NSGA2(
             pop_size=self._pop_size,
-            n_offsprings=self._pop_size,
+            n_offsprings=self._n_offsprings or self._pop_size,
             sampling=FloatRandomSampling(),
             crossover=SBX(eta=15, prob=0.9),
             mutation=PM(eta=20),
@@ -69,4 +70,12 @@ def update(self, completed: Sequence[Trial]) -> None:
         F = np.concatenate(F, axis=-1)
         static = StaticProblem(self._nsga_problem, F=F)
         Evaluator().eval(static, self._last_pop)
-        self._impl.tell(infills=self._last_pop)
+        self._impl.tell(infills=self._last_pop)
+
+    def result(self):
+        res = self._impl.result()
+        return np.atleast_2d(res.X), np.atleast_2d(res.F)
+
+    def curr_pop(self):
+        pop = self._impl.pop
+        return pop.get('X'), pop.get('F')

bbo/benchmarks/experimenters/numpy_experimenter.py

@@ -22,7 +22,7 @@ def __init__(
             problem_statement, scale=False, onehot_embed=False,
         )
 
-    def evaluate(self, suggestions: List[Trial], *, inplace=True):
+    def evaluate(self, suggestions: List[Trial]):
         features = self._converter.to_features(suggestions)
         m = self._impl(features)
         metrics = self._converter.to_metrics(m)

bbo/benchmarks/experimenters/torch_experimenter.py

@@ -0,0 +1,34 @@
+from typing import List, Callable
+
+from torch import Tensor
+
+from bbo.benchmarks.experimenters.base import BaseExperimenter
+from bbo.utils.problem_statement import ProblemStatement
+from bbo.utils.converters.torch_converter import TorchTrialConverter
+from bbo.utils.trial import Trial
+
+
+class TorchExperimenter(BaseExperimenter):
+    def __init__(
+        self,
+        impl: Callable[[Tensor], Tensor],
+        problem_statement: ProblemStatement,
+    ):
+        self._dim = problem_statement.search_space.num_parameters()
+        self._impl = impl
+        self._problem_statement = problem_statement
+
+        self._converter = TorchTrialConverter.from_problem(
+            problem_statement, scale=False, onehot_embed=False,
+        )
+
+    def evaluate(self, suggestions: List[Trial]):
+        features = self._converter.to_features(suggestions)
+        m = self._impl(features)
+        metrics = self._converter.to_metrics(m)
+        for suggestion, m in zip(suggestions, metrics):
+            suggestion.complete(m)
+        return suggestions
+
+    def problem_statement(self) -> ProblemStatement:
+        return self._problem_statement

bbo/utils/converters/torch_converter.py

@@ -0,0 +1,66 @@
+from typing import Dict, Sequence, List, Tuple
+
+import torch
+from torch import Tensor
+
+from bbo.utils.converters.base import (
+    BaseInputConverter,
+    BaseOutputConverter,
+    BaseTrialConverter,
+)
+from bbo.utils.converters.converter import ArrayTrialConverter, NumpyArraySpec
+from bbo.utils.problem_statement import ProblemStatement
+from bbo.utils.metric_config import MetricInformation
+from bbo.utils.trial import ParameterDict, MetricDict, Trial
+
+
+class TorchTrialConverter(BaseTrialConverter):
+    def __init__(
+        self,
+        input_converters: Sequence[BaseInputConverter],
+        output_converters: Sequence[BaseOutputConverter],
+    ):
+        self._impl = ArrayTrialConverter(input_converters, output_converters)
+
+    @classmethod
+    def from_problem(
+        cls,
+        problem: ProblemStatement,
+        *,
+        scale: bool = True,
+        onehot_embed: bool = False,
+    ):
+        converter = cls([], [])
+        converter._impl = ArrayTrialConverter.from_problem(
+            problem, scale=scale, onehot_embed=onehot_embed
+        )
+        return converter
+
+    def convert(self, trials: Sequence[Trial]) -> Tuple[Tensor, Tensor]:
+        return self.to_features(trials), self.to_labels(trials)
+
+    def to_features(self, trials: Sequence[Trial]) -> Tensor:
+        return torch.from_numpy(self._impl.to_features(trials))
+
+    def to_labels(self, trials: Sequence[Trial]) -> Tensor:
+        return torch.from_numpy(self._impl.to_labels(trials))
+
+    def to_trials(self, features: Tensor, labels: Tensor=None) -> Sequence[Trial]:
+        features = features.detach().numpy()
+        if labels is not None:
+            labels = labels.detach().numpy()
+        return self._impl.to_trials(features, labels)
+
+    def to_parameters(self, features: Tensor) -> List[ParameterDict]:
+        return self._impl.to_parameters(features.detach().numpy())
+
+    def to_metrics(self, labels: Tensor) -> List[MetricDict]:
+        return self._impl.to_metrics(labels.detach().numpy())
+
+    @property
+    def output_spec(self) -> Dict[str, NumpyArraySpec]:
+        return {k: v.output_spec for k, v in self._impl.input_converter_dict.items()}
+
+    @property
+    def metric_spec(self) -> Dict[str, MetricInformation]:
+        return {k: v.metric_information for k, v in self._impl.output_converter_dict.items()}