Move evolution strategy bounds to init

HISTORY.md

@@ -6,6 +6,7 @@
 
 #### API
 
+- **Backwards-incompatible:** Move evolution strategy bounds to init ({pr}`436`)
 - **Backwards-incompatible:** Use seed instead of rng in ranker ({pr}`432`)
 - **Backwards-incompatible:** Replace status and value with add_info ({pr}`430`)
 - Support custom data fields in archive, emitters, and scheduler ({pr}`421`,

ribs/emitters/_evolution_strategy_emitter.py

@@ -118,13 +118,17 @@ def __init__(
         # Check if the restart_rule is valid, discard check_restart result.
         _ = self._check_restart(0)
 
-        self._opt = _get_es(es,
-                            sigma0=sigma0,
-                            batch_size=batch_size,
-                            solution_dim=self._solution_dim,
-                            seed=opt_seed,
-                            dtype=self.archive.dtype,
-                            **(es_kwargs if es_kwargs is not None else {}))
+        self._opt = _get_es(
+            es,
+            sigma0=sigma0,
+            batch_size=batch_size,
+            solution_dim=self._solution_dim,
+            seed=opt_seed,
+            dtype=self.archive.dtype,
+            lower_bounds=self.lower_bounds,
+            upper_bounds=self.upper_bounds,
+            **(es_kwargs if es_kwargs is not None else {}),
+        )
         self._opt.reset(self._x0)
 
         self._ranker = _get_ranker(ranker, ranker_seed)
@@ -163,7 +167,7 @@ def ask(self):
             (batch_size, :attr:`solution_dim`) array -- a batch of new solutions
             to evaluate.
         """
-        return self._opt.ask(self.lower_bounds, self.upper_bounds)
+        return self._opt.ask()
 
     def _check_restart(self, num_parents):
         """Emitter-side checks for restarting the optimizer.

ribs/emitters/_gradient_arborescence_emitter.py

@@ -196,13 +196,17 @@ def __init__(self,
             lr=lr,
             **(grad_opt_kwargs if grad_opt_kwargs is not None else {}))
 
-        self._opt = _get_es(es,
-                            sigma0=sigma0,
-                            batch_size=batch_size,
-                            solution_dim=self._num_coefficients,
-                            seed=opt_seed,
-                            dtype=self.archive.dtype,
-                            **(es_kwargs if es_kwargs is not None else {}))
+        self._opt = _get_es(
+            es,
+            sigma0=sigma0,
+            batch_size=batch_size,
+            solution_dim=self._num_coefficients,
+            seed=opt_seed,
+            dtype=self.archive.dtype,
+            lower_bounds=-np.inf,  # No bounds for gradient coefficients.
+            upper_bounds=np.inf,
+            **(es_kwargs if es_kwargs is not None else {}),
+        )
 
         self._opt.reset(np.zeros(self._num_coefficients))
 
@@ -276,20 +280,7 @@ def ask(self):
             raise RuntimeError("Please call ask_dqd() and tell_dqd() "
                                "before calling ask().")
 
-        coeff_lower_bounds = np.full(
-            self._num_coefficients,
-            -np.inf,
-            dtype=self._archive.dtype,
-        )
-        coeff_upper_bounds = np.full(
-            self._num_coefficients,
-            np.inf,
-            dtype=self._archive.dtype,
-        )
-        grad_coeffs = self._opt.ask(
-            coeff_lower_bounds,
-            coeff_upper_bounds,
-        )[:, :, None]
+        grad_coeffs = self._opt.ask()[:, :, None]
         return (self._grad_opt.theta +
                 np.sum(self._jacobian_batch * grad_coeffs, axis=1))
 

ribs/emitters/opt/_cma_es.py

@@ -86,6 +86,13 @@ class CMAEvolutionStrategy(EvolutionStrategyBase):
         solution_dim (int): Size of the solution space.
         seed (int): Seed for the random number generator.
         dtype (str or data-type): Data type of solutions.
+        lower_bounds (float or np.ndarray): scalar or (solution_dim,) array
+            indicating lower bounds of the solution space. Scalars specify
+            the same bound for the entire space, while arrays specify a
+            bound for each dimension. Pass -np.inf in the array or scalar to
+            indicated unbounded space.
+        upper_bounds (float or np.ndarray): Same as above, but for upper
+            bounds (and pass np.inf instead of -np.inf).
     """
 
     def __init__(  # pylint: disable = super-init-not-called
@@ -94,12 +101,20 @@ def __init__(  # pylint: disable = super-init-not-called
             solution_dim,
             batch_size=None,
             seed=None,
-            dtype=np.float64):
+            dtype=np.float64,
+            lower_bounds=-np.inf,
+            upper_bounds=np.inf):
         self.batch_size = (4 + int(3 * np.log(solution_dim))
                            if batch_size is None else batch_size)
         self.sigma0 = sigma0
         self.solution_dim = solution_dim
         self.dtype = dtype
+
+        # Even scalars must be converted into 0-dim arrays so that they work
+        # with the bound check in numba.
+        self.lower_bounds = np.asarray(lower_bounds, dtype=self.dtype)
+        self.upper_bounds = np.asarray(upper_bounds, dtype=self.dtype)
+
         self._rng = np.random.default_rng(seed)
         self._solutions = None
 
@@ -182,17 +197,10 @@ def _transform_and_check_sol(unscaled_params, transform_mat, mean,
     # Limit OpenBLAS to single thread. This is typically faster than
     # multithreading because our data is too small.
     @threadpool_limits.wrap(limits=1, user_api="blas")
-    def ask(self, lower_bounds, upper_bounds, batch_size=None):
+    def ask(self, batch_size=None):
         """Samples new solutions from the Gaussian distribution.
 
         Args:
-            lower_bounds (float or np.ndarray): scalar or (solution_dim,) array
-                indicating lower bounds of the solution space. Scalars specify
-                the same bound for the entire space, while arrays specify a
-                bound for each dimension. Pass -np.inf in the array or scalar to
-                indicated unbounded space.
-            upper_bounds (float or np.ndarray): Same as above, but for upper
-                bounds (and pass np.inf instead of -np.inf).
             batch_size (int): batch size of the sample. Defaults to
                 ``self.batch_size``.
         """
@@ -214,8 +222,8 @@ def ask(self, lower_bounds, upper_bounds, batch_size=None):
                 (len(remaining_indices), self.solution_dim),
             ).astype(self.dtype)
             new_solutions, out_of_bounds = self._transform_and_check_sol(
-                unscaled_params, transform_mat, self.mean, lower_bounds,
-                upper_bounds)
+                unscaled_params, transform_mat, self.mean, self.lower_bounds,
+                self.upper_bounds)
             self._solutions[remaining_indices] = new_solutions
 
             # Find indices in remaining_indices that are still out of bounds

ribs/emitters/opt/_evolution_strategy_base.py

@@ -24,14 +24,23 @@ class EvolutionStrategyBase(ABC):
         batch_size (int): Number of solutions to evaluate at a time.
         seed (int): Seed for the random number generator.
         dtype (str or data-type): Data type of solutions.
+        lower_bounds (float or np.ndarray): scalar or (solution_dim,) array
+            indicating lower bounds of the solution space. Scalars specify
+            the same bound for the entire space, while arrays specify a
+            bound for each dimension. Pass -np.inf in the array or scalar to
+            indicated unbounded space.
+        upper_bounds (float or np.ndarray): Same as above, but for upper
+            bounds (and pass np.inf instead of -np.inf).
     """
 
     def __init__(self,
                  sigma0,
                  solution_dim,
                  batch_size=None,
                  seed=None,
-                 dtype=np.float64):
+                 dtype=np.float64,
+                 lower_bounds=-np.inf,
+                 upper_bounds=np.inf):
         pass
 
     @abstractmethod
@@ -55,17 +64,12 @@ def check_stop(self, ranking_values):
         """
 
     @abstractmethod
-    def ask(self, lower_bounds, upper_bounds):
+    def ask(self, batch_size=None):
         """Samples new solutions from the Gaussian distribution.
 
         Args:
-            lower_bounds (float or np.ndarray): scalar or (solution_dim,) array
-                indicating lower bounds of the solution space. Scalars specify
-                the same bound for the entire space, while arrays specify a
-                bound for each dimension. Pass -np.inf in the array or scalar to
-                indicated unbounded space.
-            upper_bounds (float or np.ndarray): Same as above, but for upper
-                bounds (and pass np.inf instead of -np.inf).
+            batch_size (int): batch size of the sample. Defaults to
+                ``self.batch_size``.
         """
 
     @abstractmethod

ribs/emitters/opt/_lm_ma_es.py

@@ -5,7 +5,6 @@
 """
 import numba as nb
 import numpy as np
-from threadpoolctl import threadpool_limits
 
 from ribs._utils import readonly
 from ribs.emitters.opt._evolution_strategy_base import EvolutionStrategyBase
@@ -23,6 +22,13 @@ class LMMAEvolutionStrategy(EvolutionStrategyBase):
         solution_dim (int): Size of the solution space.
         seed (int): Seed for the random number generator.
         dtype (str or data-type): Data type of solutions.
+        lower_bounds (float or np.ndarray): scalar or (solution_dim,) array
+            indicating lower bounds of the solution space. Scalars specify
+            the same bound for the entire space, while arrays specify a
+            bound for each dimension. Pass -np.inf in the array or scalar to
+            indicated unbounded space.
+        upper_bounds (float or np.ndarray): Same as above, but for upper
+            bounds (and pass np.inf instead of -np.inf).
         n_vectors (int): Number of vectors to use in the approximation. If None,
             this defaults to be equal to the batch size.
     """
@@ -34,12 +40,20 @@ def __init__(  # pylint: disable = super-init-not-called
             batch_size=None,
             seed=None,
             dtype=np.float64,
+            lower_bounds=-np.inf,
+            upper_bounds=np.inf,
             n_vectors=None):
         self.batch_size = (4 + int(3 * np.log(solution_dim))
                            if batch_size is None else batch_size)
         self.sigma0 = sigma0
         self.solution_dim = solution_dim
         self.dtype = dtype
+
+        # Even scalars must be converted into 0-dim arrays so that they work
+        # with the bound check in numba.
+        self.lower_bounds = np.asarray(lower_bounds, dtype=self.dtype)
+        self.upper_bounds = np.asarray(upper_bounds, dtype=self.dtype)
+
         self._rng = np.random.default_rng(seed)
         self._solutions = None
 
@@ -134,20 +148,10 @@ def _transform_and_check_sol(z, itrs, cd, m, mean, sigma, lower_bounds,
 
         return new_solutions, out_of_bounds
 
-    # Limit OpenBLAS to single thread. This is typically faster than
-    # multithreading because our data is too small.
-    @threadpool_limits.wrap(limits=1, user_api="blas")
-    def ask(self, lower_bounds, upper_bounds, batch_size=None):
+    def ask(self, batch_size=None):
         """Samples new solutions from the Gaussian distribution.
 
         Args:
-            lower_bounds (float or np.ndarray): scalar or (solution_dim,) array
-                indicating lower bounds of the solution space. Scalars specify
-                the same bound for the entire space, while arrays specify a
-                bound for each dimension. Pass -np.inf in the array or scalar to
-                indicated unbounded space.
-            upper_bounds (float or np.ndarray): Same as above, but for upper
-                bounds (and pass np.inf instead of -np.inf).
             batch_size (int): batch size of the sample. Defaults to
                 ``self.batch_size``.
         """
@@ -170,7 +174,7 @@ def ask(self, lower_bounds, upper_bounds, batch_size=None):
 
             new_solutions, out_of_bounds = self._transform_and_check_sol(
                 z, min(self.current_gens, self.n_vectors), self.cd, self.m,
-                self.mean, self.sigma, lower_bounds, upper_bounds)
+                self.mean, self.sigma, self.lower_bounds, self.upper_bounds)
             self._solutions[remaining_indices] = new_solutions
 
             # Find indices in remaining_indices that are still out of bounds
@@ -194,9 +198,6 @@ def _calc_strat_params(num_parents):
 
         return weights, mueff
 
-    # Limit OpenBLAS to single thread. This is typically faster than
-    # multithreading because our data is too small.
-    @threadpool_limits.wrap(limits=1, user_api="blas")
     def tell(self, ranking_indices, num_parents):
         """Passes the solutions back to the optimizer.
 

ribs/emitters/opt/_openai_es.py

@@ -3,7 +3,6 @@
 See here for more info: https://arxiv.org/abs/1703.03864
 """
 import numpy as np
-from threadpoolctl import threadpool_limits
 
 from ribs._utils import readonly
 from ribs.emitters.opt._adam_opt import AdamOpt
@@ -22,6 +21,13 @@ class OpenAIEvolutionStrategy(EvolutionStrategyBase):
         solution_dim (int): Size of the solution space.
         seed (int): Seed for the random number generator.
         dtype (str or data-type): Data type of solutions.
+        lower_bounds (float or np.ndarray): scalar or (solution_dim,) array
+            indicating lower bounds of the solution space. Scalars specify
+            the same bound for the entire space, while arrays specify a
+            bound for each dimension. Pass -np.inf in the array or scalar to
+            indicated unbounded space.
+        upper_bounds (float or np.ndarray): Same as above, but for upper
+            bounds (and pass np.inf instead of -np.inf).
         mirror_sampling (bool): Whether to use mirror sampling when gathering
             solutions. Defaults to True.
         adam_kwargs (dict): Keyword arguments passed to :class:`AdamOpt`.
@@ -34,13 +40,21 @@ def __init__(  # pylint: disable = super-init-not-called
             batch_size=None,
             seed=None,
             dtype=np.float64,
+            lower_bounds=-np.inf,
+            upper_bounds=np.inf,
             mirror_sampling=True,
             **adam_kwargs):
         self.batch_size = (4 + int(3 * np.log(solution_dim))
                            if batch_size is None else batch_size)
         self.sigma0 = sigma0
         self.solution_dim = solution_dim
         self.dtype = dtype
+
+        # Even scalars must be converted into 0-dim arrays so that they work
+        # with the bound check in numba.
+        self.lower_bounds = np.asarray(lower_bounds, dtype=self.dtype)
+        self.upper_bounds = np.asarray(upper_bounds, dtype=self.dtype)
+
         self._rng = np.random.default_rng(seed)
         self._solutions = None
 
@@ -96,22 +110,10 @@ def check_stop(self, ranking_values):
 
         return False
 
-    # Limit OpenBLAS to single thread. This is typically faster than
-    # multithreading because our data is too small.
-    @threadpool_limits.wrap(limits=1, user_api="blas")
-    def ask(self, lower_bounds, upper_bounds, batch_size=None):
+    def ask(self, batch_size=None):
         """Samples new solutions from the Gaussian distribution.
 
-        Note: Bounds are currently not enforced.
-
         Args:
-            lower_bounds (float or np.ndarray): scalar or (solution_dim,) array
-                indicating lower bounds of the solution space. Scalars specify
-                the same bound for the entire space, while arrays specify a
-                bound for each dimension. Pass -np.inf in the array or scalar to
-                indicated unbounded space.
-            upper_bounds (float or np.ndarray): Same as above, but for upper
-                bounds (and pass np.inf instead of -np.inf).
             batch_size (int): batch size of the sample. Defaults to
                 ``self.batch_size``.
         """
@@ -133,12 +135,11 @@ def ask(self, lower_bounds, upper_bounds, batch_size=None):
                 self.noise = self._rng.standard_normal(
                     (batch_size, self.solution_dim), dtype=self.dtype)
 
-            # TODO Numba
             new_solutions = (self.adam_opt.theta[None] +
                              self.sigma0 * self.noise)
             out_of_bounds = np.logical_or(
-                new_solutions < np.expand_dims(lower_bounds, axis=0),
-                new_solutions > np.expand_dims(upper_bounds, axis=0),
+                new_solutions < np.expand_dims(self.lower_bounds, axis=0),
+                new_solutions > np.expand_dims(self.upper_bounds, axis=0),
             )
 
             self._solutions[remaining_indices] = new_solutions
@@ -150,9 +151,6 @@ def ask(self, lower_bounds, upper_bounds, batch_size=None):
 
         return readonly(self._solutions)
 
-    # Limit OpenBLAS to single thread. This is typically faster than
-    # multithreading because our data is too small.
-    @threadpool_limits.wrap(limits=1, user_api="blas")
     def tell(
             self,
             ranking_indices,