Improving ACEPS

Author: Jan Hendrik Metzen

.gitignore

@@ -5,3 +5,5 @@ doc/_build
 doc/source/auto_examples
 doc/source/modules/generated
 doc/source/notebooks
+doc
+build

bayesian_optimization/acquisition_functions.py

@@ -7,8 +7,18 @@
 
 
 class AcquisitionFunction(object):
+    """ Abstract base class for acquisition functions."""
 
     def set_boundaries(self, boundaries):
+        """Sets boundaries of search space.
+
+        Parameters
+        ----------
+        boundaries: ndarray-like, shape=(n_params_dims, 2)
+            Box constraint on search space. boundaries[:, 0] defines the lower
+            bounds on the dimensions, boundaries[:, 1] defines the upper
+            bounds.
+        """
         pass
 
 
@@ -299,13 +309,175 @@ def set_boundaries(self, boundaries, X_candidate=None):
         self.base_entropy = entropy(p_max)
 
 
+class ContextualEntropySearch(AcquisitionFunction):
+    """
+    n_context_samples: int, default: 20
+        The number of context sampled from context space on which sample
+        policies are evaluated
+    """
+    def __init__(self, model, n_context_dims, n_context_samples,
+                 n_candidates=20, n_gp_samples=500,
+                 n_samples_y=10, n_trial_points=100):
+        self.model = model
+        self.n_context_dims = n_context_dims
+        self.n_context_samples = n_context_samples
+
+        self.n_candidates = n_candidates
+        self.n_gp_samples = n_gp_samples
+        self.n_samples_y =  n_samples_y
+        self.n_trial_points = n_trial_points
+
+    def __call__(self, x, incumbent=0, *args, **kwargs):
+        ind = list(self.nbrs.kneighbors(x[:self.n_context_dims],
+                                        return_distance=False)[0])
+        mean_entropy_reduction = \
+            np.mean([self.entropy_search_ensemble[i](x) for i in ind])
+        return mean_entropy_reduction
+
+    def set_boundaries(self, boundaries):
+        self._sample_contexts(boundaries[:self.n_context_dims])
+
+        self.entropy_search_ensemble = []
+        for i in range(self.n_context_samples):
+            boundaries_i = np.copy(boundaries)
+            boundaries_i[:self.n_context_dims] = \
+                self.context_samples[i][:, np.newaxis]
+            entropy_search_fixed_context = \
+                EntropySearch(self.model, self.n_candidates, self.n_gp_samples,
+                              self.n_samples_y, self.n_trial_points)
+            entropy_search_fixed_context.set_boundaries(boundaries_i)
+
+            self.entropy_search_ensemble.append(entropy_search_fixed_context)
+
+    def _sample_contexts(self, context_boundaries):
+        self.context_samples = \
+            np.random.uniform(context_boundaries[:, 0],
+                              context_boundaries[:, 1],
+                              (self.n_context_samples*25, self.n_context_dims))
+        from sklearn.cluster import KMeans
+        from sklearn.neighbors import NearestNeighbors
+        kmeans = KMeans(n_clusters=self.n_context_samples).fit(self.context_samples)
+        self.context_samples = kmeans.cluster_centers_
+        self.nbrs = NearestNeighbors(n_neighbors=10, algorithm='ball_tree')
+        self.nbrs.fit(self.context_samples)
+
+
+class ACEPS(ContextualEntropySearch):
+    def __init__(self, model, policy, n_context_dims,
+                 n_context_samples, n_policy_samples=20, n_gp_samples=1000,
+                 n_samples_y=10, policy_training="model-free"):
+        super(ACEPS, self).__init__(model, n_context_dims, n_context_samples,
+                                    n_policy_samples, n_gp_samples,
+                                    n_samples_y)
+        self.policy = policy
+        self.policy_training = policy_training
+
+
+    def set_boundaries(self, boundaries):
+        self._sample_contexts(boundaries[:self.n_context_dims])
+
+        selected_params = \
+            self._sample_policy_parameters(boundaries[:self.n_context_dims],
+                                           boundaries[self.n_context_dims:])
+
+        self.entropy_search_ensemble = []
+        for i in range(self.n_context_samples):
+            boundaries_i = np.copy(boundaries)
+            boundaries_i[:self.n_context_dims] = \
+                self.context_samples[i][:, np.newaxis]
+            entropy_search_fixed_context = \
+                EntropySearch(self.model, self.n_candidates, self.n_gp_samples,
+                              self.n_samples_y, self.n_trial_points)
+
+            X_candidate =  np.empty((self.n_candidates, 4)) # XXX
+            X_candidate[:, :self.n_context_dims] = self.context_samples[i]
+            X_candidate[:, self.n_context_dims:] = selected_params[:, i]
+            entropy_search_fixed_context.set_boundaries(boundaries_i,
+                                                        X_candidate)
+
+            self.entropy_search_ensemble.append(entropy_search_fixed_context)
+
+
+    def _sample_policy_parameters(self, context_boundaries, parameter_boundaries):
+        """ Sample close-to-optimal policies and let them select parameters.
+
+        We determine a set of policies which is close-to-optimal according to
+        samples drawn from the model's posterior and let these policies
+        determine parameters
+        """
+        from bolero_bayes_opt.representation.ul_policies \
+            import model_free_policy_training, model_based_policy_training_pretrained
+        # Compute policy which is close to optimal according to current model
+        contexts = self.model.gp.X_train_[:, :self.n_context_dims]
+        parameters = self.model.gp.X_train_[:, self.n_context_dims:]
+        returns = self.model.gp.y_train_
+        if self.policy_training == "model-free":
+            self.policy = model_free_policy_training(
+                        self.policy, contexts, parameters, returns,
+                        epsilon=1.0, min_eta=1e-6)
+        elif self.policy_training == "model-based":
+            self.policy.fit(contexts,
+                            [self.parameter_boundaries.mean(1)]*contexts.shape[0],
+                            weights=np.ones(contexts.shape[0]))
+            self.policy = model_based_policy_training_pretrained(
+                policy=self.policy, model=self.model.gp,
+                contexts=contexts, boundaries=self.parameter_boundaries)
+        else:
+            raise NotImplementedError()
+
+        # Draw context samples, let policy select parameters for these context
+        # (with exploration), and sample multiple possible outcomes for these
+        # (context, parameter) samples from the GP model.
+        while True:  # XXX
+            X_sample = np.empty((self.n_context_samples*5,
+                                 self.n_context_dims+parameter_boundaries.shape[0]))
+            X_sample[:, :self.n_context_dims] = \
+                np.random.uniform(context_boundaries[:, 0],
+                                  context_boundaries[:, 1],
+                                  (self.n_context_samples*5, self.n_context_dims))
+            X_sample[:, self.n_context_dims:] = \
+                [self.policy(X_sample[i, :self.n_context_dims])
+                 for i in range(self.n_context_samples*5)]
+            try:
+                y_sample = self.model.gp.sample_y(X_sample, self.n_candidates)
+                break
+            except np.linalg.LinAlgError:
+                continue
+
+        # Train for each possible outcome one policy and evaluate this policy
+        # on the context samples
+        selected_params = []  # XXX: Vectorize
+        for i in range(y_sample.shape[1]):
+            policy_sample = model_free_policy_training(
+                self.policy, X_sample[:, range(self.n_context_dims)],
+                X_sample[:, range(self.n_context_dims, X_sample.shape[1])],
+                y_sample[:, i])
+
+            params = [policy_sample(np.atleast_1d(self.context_samples[i]),
+                                    explore=False).ravel()
+                      for i in range(self.context_samples.shape[0])]
+            selected_params.append(params)
+
+        selected_params = np.array(selected_params)
+        # Enforce lower and upper bound on possible parameters
+        for i in range(selected_params.shape[2]):
+            selected_params[:, :, i][selected_params[:, :, i] < parameter_boundaries[i, 0]] = \
+                parameter_boundaries[i, 0]
+            selected_params[:, :, i][selected_params[:, :, i] > parameter_boundaries[i, 1]] = \
+                parameter_boundaries[i, 1]
+
+        return selected_params
+
+
 ACQUISITION_FUNCTIONS = {
     "PI": ProbabilityOfImprovement,
     "EI": ExpectedImprovement,
     "UCB": UpperConfidenceBound,
     "GREEDY": Greedy,
     "RANDOM": Random,
-    "EntropySearch": EntropySearch}
+    "EntropySearch": EntropySearch,
+    "ContextualEntropySearch": ContextualEntropySearch,
+    "ACEPS": ACEPS}
 
 
 def create_acquisition_function(name, model, **kwargs):

bolero_bayes_opt/optimizer/aceps.py

@@ -7,6 +7,8 @@
 import numpy as np
 from scipy.stats import entropy, norm
 
+from bayesian_optimization import create_acquisition_function
+from bayesian_optimization.acquisition_functions import AcquisitionFunction
 from bayesian_optimization.model import ParametricModelApproximation
 
 from .bocps import BOCPSOptimizer
@@ -117,40 +119,22 @@ def get_next_parameters(self, params, explore=True):
         # the context in get_desired_context()
         params[:] = self.parameters
 
-    def _determine_contextparams(self, optimizer, context_samples=None):
+    def _determine_contextparams(self, optimizer):
         """Select context and params jointly using ACEPS."""
-        if context_samples is None:
-            # Select self.n_query_points possible contexts for the rollout
-            # randomly
-            context_samples = \
-                self.rng.uniform(self.context_boundaries[:, 0],
-                                 self.context_boundaries[:, 1],
-                                 (self.n_query_points, self.context_dims))
-        # Let the base-acquisition function select parameters for these
-        # contexts
-        selected_params = np.empty((context_samples.shape[0], self.dimension))
-        for i in range(context_samples.shape[0]):
-            selected_params[i] = \
-                self._determine_next_query_point(context_samples[i], optimizer)
-
-        # We cannot evaluate query points without having seen at least two
-        # datapoints. We thus select a random query point
-        if self.model.last_training_size < 2:
-            rand_ind = np.random.choice(context_samples.shape[0])
-            return context_samples[rand_ind], selected_params[rand_ind]
-
-        # Determine for every of the possible context-parameter pairs their
-        # ACEPS score and select the one with maximal score
-        cx_boundaries = np.vstack((self.context_boundaries, self.boundaries))
-        aceps = ACEPS(self.model, self.policy, cx_boundaries,
-                      n_context_dims=self.context_dims, **self.aceps_params)
-        aceps_performance = np.empty(context_samples.shape[0])
-        for i in range(context_samples.shape[0]):
-            aceps_performance[i] = \
-                aceps(np.hstack((context_samples[i], selected_params[i])))
-        opt_ind = np.argmax(aceps_performance)
-
-        return context_samples[opt_ind], selected_params[opt_ind]
+        cx_boundaries = np.empty((self.context_dims + self.dimension, 2))
+        cx_boundaries[:self.context_dims] = self.context_boundaries
+        cx_boundaries[self.context_dims:] = self.boundaries
+
+        # Determine optimal parameters for fixed context
+        cx = optimizer.select_query_point(cx_boundaries)
+        return cx[:self.context_dims], cx[self.context_dims:]
+
+    def _create_acquisition_function(self, name, model, **kwargs):
+        if not name in ["ContextualEntropySearch", "ACEPS"]:
+            raise ValueError("%s acquisition function not supported."
+                             % name)
+
+        return create_acquisition_function(name, model, **kwargs)
 
 
 class ACEPS(object):