Amortized bo implementation

src/poli_baselines/solvers/__init__.py

@@ -1,17 +1,17 @@
-# from .simple.random_mutation import RandomMutation
-# from .simple.continuous_random_mutation import ContinuousRandomMutation
-# from .simple.genetic_algorithm import FixedLengthGeneticAlgorithm
-
-# from .bayesian_optimization.vanilla_bayesian_optimization import (
-#     VanillaBayesianOptimization,
-# )
-# from .bayesian_optimization.line_bayesian_optimization import LineBO
-# from .bayesian_optimization.saas_bayesian_optimization import SAASBO
-# from .bayesian_optimization.baxus import BAxUS
-
-# from .bayesian_optimization.latent_space_bayesian_optimization import (
-#     LatentSpaceBayesianOptimization,
-# )
-
-# from .evolutionary_strategies.cma_es import CMA_ES
-# from .multi_objective.nsga_ii import DiscreteNSGAII
+# from .simple.random_mutation import RandomMutation
+# from .simple.continuous_random_mutation import ContinuousRandomMutation
+# from .simple.genetic_algorithm import FixedLengthGeneticAlgorithm
+
+# from .bayesian_optimization.vanilla_bayesian_optimization import (
+#     VanillaBayesianOptimization,
+# )
+# from .bayesian_optimization.line_bayesian_optimization import LineBO
+# from .bayesian_optimization.saas_bayesian_optimization import SAASBO
+# from .bayesian_optimization.baxus import BAxUS
+
+# from .bayesian_optimization.latent_space_bayesian_optimization import (
+#     LatentSpaceBayesianOptimization,
+# )
+
+# from .evolutionary_strategies.cma_es import CMA_ES
+# from .multi_objective.nsga_ii import DiscreteNSGAII

src/poli_baselines/solvers/bayesian_optimization/amortized/amortized_bo_wrapper.py

@@ -0,0 +1,44 @@
+import numpy as np
+from poli.core.abstract_black_box import AbstractBlackBox
+
+from poli_baselines.core.step_by_step_solver import StepByStepSolver
+from poli_baselines.solvers.bayesian_optimization.amortized.data import (
+    samples_from_arrays,
+    Population,
+)
+from poli_baselines.solvers.bayesian_optimization.amortized.deep_evolution_solver import (
+    MutationPredictorSolver,
+)
+from poli_baselines.solvers.bayesian_optimization.amortized.domains import (
+    FixedLengthDiscreteDomain,
+    Vocabulary,
+)
+
+
+class AmortizedBOWrapper(StepByStepSolver):
+    def __init__(self, black_box: AbstractBlackBox, x0: np.ndarray, y0: np.ndarray):
+        super().__init__(black_box, x0, y0)
+        self.problem_info = black_box.get_black_box_info()
+        alphabet = self.problem_info.get_alphabet()
+        if not self.problem_info.sequences_are_aligned():
+            alphabet = alphabet + [self.problem_info.get_padding_token()]
+        self.domain = FixedLengthDiscreteDomain(
+            vocab=Vocabulary(alphabet),
+            length=x0.shape[1],
+        )
+        self.solver = MutationPredictorSolver(
+            domain=self.domain,
+            initialize_dataset_fn=lambda *args, **kwargs: self.domain.encode(x0),
+        )
+
+    def next_candidate(self) -> np.ndarray:
+        samples = samples_from_arrays(
+            structures=self.domain.encode(self.x0.tolist()), rewards=self.y0.tolist()
+        )
+        x = self.solver.propose(num_samples=1, population=Population(samples))
+        s = list(self.domain.decode(x)[0])
+        if not self.problem_info.sequences_are_aligned():
+            s = s + [self.problem_info.get_padding_token()] * (
+                self.problem_info.get_max_sequence_length() - len(s)
+            )
+        return np.array([s])

src/poli_baselines/solvers/bayesian_optimization/amortized/base_solver.py

@@ -0,0 +1,82 @@
+# coding=utf-8
+# Copyright 2024 The Google Research Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Solver base class."""
+
+import abc
+
+from absl import logging
+
+from poli_baselines.solvers.bayesian_optimization.amortized import utils
+
+
+class BaseSolver(abc.ABC):
+    """Solver base class."""
+
+    def __init__(
+        self, domain, random_state=None, name=None, log_level=logging.INFO, **kwargs
+    ):
+        """Creates an instance of this class.
+
+        Args:
+          domain: An instance of a `Domain`.
+          random_state: An instance of or integer seed to build a
+            `np.random.RandomState`.
+          name: The name of the solver. If `None`, will use the class name.
+          log_level: The logging level of the solver-specific logger.
+            -2=ERROR, -1=WARN,  0=INFO, 1=DEBUG.
+          **kwargs: Named arguments stored in `self.cfg`.
+        """
+        self._domain = domain
+        self._name = name or self.__class__.__name__
+        self._random_state = utils.get_random_state(random_state)
+        self._log = utils.get_logger(self._name, level=log_level)
+
+        cfg = utils.Config(self._config())
+        cfg.update(kwargs)
+        self.cfg = cfg
+
+    def _config(self):
+        return {}
+
+    @property
+    def domain(self):
+        """Return the optimization domain."""
+        return self._domain
+
+    @property
+    def name(self):
+        """Returns the solver name."""
+        return self._name
+
+    def __str__(self):
+        return self._name
+
+    @abc.abstractmethod
+    def propose(self, num_samples, population=None, pending_samples=None, counter=0):
+        """Proposes num_samples from `self.domain`.
+
+        Args:
+          num_samples: The number of samples to return.
+          population: A `Population` of samples or None if the population is empty.
+          pending_samples: A list of structures without reward that were already
+            proposed.
+          counter: The number of times `propose` has been called with the same
+            `population`. Can be used by solver to avoid repeated computations on
+            the same `population`, e.g. updating a model.
+
+        Returns:
+          `num_samples` structures from the domain.
+        """