Optional fitness shaping within strategy

RobertTLange · Nov 19, 2022 · 2ddcdd5 · 2ddcdd5
1 parent 6c81bb0
commit 2ddcdd5
Show file tree

Hide file tree

Showing 30 changed files with 87 additions and 36 deletions.
diff --git a/evosax/strategies/ars.py b/evosax/strategies/ars.py
@@ -37,10 +37,11 @@ def __init__(
         pholder_params: Optional[Union[chex.ArrayTree, chex.Array]] = None,
         elite_ratio: float = 0.1,
         opt_name: str = "sgd",
+        **fitness_kwargs: Union[bool, int, float]
     ):
         """Augmented Random Search (Mania et al., 2018)
         Reference: https://arxiv.org/pdf/1803.07055.pdf"""
-        super().__init__(popsize, num_dims, pholder_params)
+        super().__init__(popsize, num_dims, pholder_params, **fitness_kwargs)
         assert not self.popsize & 1, "Population size must be even"
         # ARS performs antithetic sampling & allows you to select
         # "b" elite perturbation directions for the update

diff --git a/evosax/strategies/bipop_cma_es.py b/evosax/strategies/bipop_cma_es.py
@@ -25,6 +25,7 @@ def __init__(
         num_dims: Optional[int] = None,
         pholder_params: Optional[Union[chex.ArrayTree, chex.Array]] = None,
         elite_ratio: float = 0.5,
+        **fitness_kwargs: Union[bool, int, float]
     ):
         """BIPOP-CMA-ES (Hansen, 2009).
         Reference: https://hal.inria.fr/inria-00382093/document
@@ -36,6 +37,7 @@ def __init__(
             popsize=popsize,
             pholder_params=pholder_params,
             elite_ratio=elite_ratio,
+            **fitness_kwargs
         )
         from ..restarts import BIPOP_Restarter
         from ..restarts.termination import spread_criterion, cma_criterion

diff --git a/evosax/strategies/cma_es.py b/evosax/strategies/cma_es.py
@@ -86,11 +86,12 @@ def __init__(
         num_dims: Optional[int] = None,
         pholder_params: Optional[Union[chex.ArrayTree, chex.Array]] = None,
         elite_ratio: float = 0.5,
+        **fitness_kwargs: Union[bool, int, float]
     ):
         """CMA-ES (e.g. Hansen, 2016)
         Reference: https://arxiv.org/abs/1604.00772
         Inspired by: https://github.com/CyberAgentAILab/cmaes"""
-        super().__init__(popsize, num_dims, pholder_params)
+        super().__init__(popsize, num_dims, pholder_params, **fitness_kwargs)
         assert 0 <= elite_ratio <= 1
         self.elite_ratio = elite_ratio
         self.elite_popsize = max(1, int(self.popsize * self.elite_ratio))

diff --git a/evosax/strategies/de.py b/evosax/strategies/de.py
@@ -34,11 +34,12 @@ def __init__(
         popsize: int,
         num_dims: Optional[int] = None,
         pholder_params: Optional[Union[chex.ArrayTree, chex.Array]] = None,
+        **fitness_kwargs: Union[bool, int, float]
     ):
         """Differential Evolution (Storn & Price, 1997)
         Reference: https://tinyurl.com/4pje5a74"""
         assert popsize > 6
-        super().__init__(popsize, num_dims, pholder_params)
+        super().__init__(popsize, num_dims, pholder_params, **fitness_kwargs)
         self.strategy_name = "DE"
 
     @property

diff --git a/evosax/strategies/esmc.py b/evosax/strategies/esmc.py
@@ -38,11 +38,12 @@ def __init__(
         num_dims: Optional[int] = None,
         pholder_params: Optional[Union[chex.ArrayTree, chex.Array]] = None,
         opt_name: str = "adam",
+        **fitness_kwargs: Union[bool, int, float]
     ):
         """ESMC (Merchant et al., 2021)
         Reference: https://proceedings.mlr.press/v139/merchant21a.html
         """
-        super().__init__(popsize, num_dims, pholder_params)
+        super().__init__(popsize, num_dims, pholder_params, **fitness_kwargs)
         assert self.popsize & 1, "Population size must be odd"
         assert opt_name in ["sgd", "adam", "rmsprop", "clipup"]
         self.optimizer = GradientOptimizer[opt_name](self.num_dims)

diff --git a/evosax/strategies/full_iamalgam.py b/evosax/strategies/full_iamalgam.py
@@ -45,11 +45,12 @@ def __init__(
         num_dims: Optional[int] = None,
         pholder_params: Optional[Union[chex.ArrayTree, chex.Array]] = None,
         elite_ratio: float = 0.35,
+        **fitness_kwargs: Union[bool, int, float]
     ):
         """(Iterative) AMaLGaM (Bosman et al., 2013) - Full Covariance
         Reference: https://tinyurl.com/y9fcccx2
         """
-        super().__init__(popsize, num_dims, pholder_params)
+        super().__init__(popsize, num_dims, pholder_params, **fitness_kwargs)
         assert 0 <= elite_ratio <= 1
         self.elite_ratio = elite_ratio
         self.elite_popsize = max(1, int(self.popsize * self.elite_ratio))

diff --git a/evosax/strategies/gesmr_ga.py b/evosax/strategies/gesmr_ga.py
@@ -36,10 +36,11 @@ def __init__(
         pholder_params: Optional[Union[chex.ArrayTree, chex.Array]] = None,
         elite_ratio: float = 0.5,
         sigma_ratio: float = 0.5,
+        **fitness_kwargs: Union[bool, int, float]
     ):
         """Self-Adaptation Mutation Rate GA."""
 
-        super().__init__(popsize, num_dims, pholder_params)
+        super().__init__(popsize, num_dims, pholder_params, **fitness_kwargs)
         self.elite_ratio = elite_ratio
         self.elite_popsize = max(1, int(self.popsize * self.elite_ratio))
         self.num_sigma_groups = int(jnp.sqrt(self.popsize))

diff --git a/evosax/strategies/gld.py b/evosax/strategies/gld.py
@@ -31,10 +31,11 @@ def __init__(
         popsize: int,
         num_dims: Optional[int] = None,
         pholder_params: Optional[Union[chex.ArrayTree, chex.Array]] = None,
+        **fitness_kwargs: Union[bool, int, float]
     ):
         """Gradientless Descent (Golovin et al., 2019)
         Reference: https://arxiv.org/pdf/1911.06317.pdf"""
-        super().__init__(popsize, num_dims, pholder_params)
+        super().__init__(popsize, num_dims, pholder_params, **fitness_kwargs)
         self.strategy_name = "GLD"
 
     @property

diff --git a/evosax/strategies/indep_iamalgam.py b/evosax/strategies/indep_iamalgam.py
@@ -50,11 +50,12 @@ def __init__(
         num_dims: Optional[int] = None,
         pholder_params: Optional[Union[chex.ArrayTree, chex.Array]] = None,
         elite_ratio: float = 0.35,
+        **fitness_kwargs: Union[bool, int, float]
     ):
         """(Iterative) AMaLGaM (Bosman et al., 2013) - Diagonal Covariance
         Reference: https://tinyurl.com/y9fcccx2
         """
-        super().__init__(popsize, num_dims, pholder_params)
+        super().__init__(popsize, num_dims, pholder_params, **fitness_kwargs)
         assert 0 <= elite_ratio <= 1
         self.elite_ratio = elite_ratio
         self.elite_popsize = max(1, int(self.popsize * self.elite_ratio))

diff --git a/evosax/strategies/ipop_cma_es.py b/evosax/strategies/ipop_cma_es.py
@@ -25,6 +25,7 @@ def __init__(
         num_dims: Optional[int] = None,
         pholder_params: Optional[Union[chex.ArrayTree, chex.Array]] = None,
         elite_ratio: float = 0.5,
+        **fitness_kwargs: Union[bool, int, float]
     ):
         """IPOP-CMA-ES (Auer & Hansen, 2005).
         Reference: http://www.cmap.polytechnique.fr/~nikolaus.hansen/cec2005ipopcmaes.pdf
@@ -36,6 +37,7 @@ def __init__(
             num_dims=num_dims,
             pholder_params=pholder_params,
             elite_ratio=elite_ratio,
+            **fitness_kwargs
         )
         from ..restarts import IPOP_Restarter
         from ..restarts.termination import cma_criterion, spread_criterion

diff --git a/evosax/strategies/lm_ma_es.py b/evosax/strategies/lm_ma_es.py
@@ -46,11 +46,12 @@ def __init__(
         pholder_params: Optional[Union[chex.ArrayTree, chex.Array]] = None,
         elite_ratio: float = 0.5,
         memory_size: int = 10,
+        **fitness_kwargs: Union[bool, int, float]
     ):
         """Limited Memory MA-ES (Loshchilov et al., 2017)
         Reference: https://arxiv.org/pdf/1705.06693.pdf
         """
-        super().__init__(popsize, num_dims, pholder_params)
+        super().__init__(popsize, num_dims, pholder_params, **fitness_kwargs)
         assert 0 <= elite_ratio <= 1
         self.elite_ratio = elite_ratio
         self.elite_popsize = max(1, int(self.popsize * self.elite_ratio))

diff --git a/evosax/strategies/ma_es.py b/evosax/strategies/ma_es.py
@@ -42,11 +42,12 @@ def __init__(
         num_dims: Optional[int] = None,
         pholder_params: Optional[Union[chex.ArrayTree, chex.Array]] = None,
         elite_ratio: float = 0.5,
+        **fitness_kwargs: Union[bool, int, float]
     ):
         """MA-ES (Bayer & Sendhoff, 2017)
         Reference: https://www.honda-ri.de/pubs/pdf/3376.pdf
         """
-        super().__init__(popsize, num_dims, pholder_params)
+        super().__init__(popsize, num_dims, pholder_params, **fitness_kwargs)
         assert 0 <= elite_ratio <= 1
         self.elite_ratio = elite_ratio
         self.elite_popsize = max(1, int(self.popsize * self.elite_ratio))

diff --git a/evosax/strategies/open_es.py b/evosax/strategies/open_es.py
@@ -36,11 +36,12 @@ def __init__(
         num_dims: Optional[int] = None,
         pholder_params: Optional[Union[chex.ArrayTree, chex.Array]] = None,
         opt_name: str = "adam",
+        **fitness_kwargs: Union[bool, int, float]
     ):
         """OpenAI-ES (Salimans et al. (2017)
         Reference: https://arxiv.org/pdf/1703.03864.pdf
         Inspired by: https://github.com/hardmaru/estool/blob/master/es.py"""
-        super().__init__(popsize, num_dims, pholder_params)
+        super().__init__(popsize, num_dims, pholder_params, **fitness_kwargs)
         assert not self.popsize & 1, "Population size must be even"
         assert opt_name in ["sgd", "adam", "rmsprop", "clipup"]
         self.optimizer = GradientOptimizer[opt_name](self.num_dims)

diff --git a/evosax/strategies/pbt.py b/evosax/strategies/pbt.py
@@ -32,10 +32,11 @@ def __init__(
         popsize: int,
         num_dims: Optional[int] = None,
         pholder_params: Optional[Union[chex.ArrayTree, chex.Array]] = None,
+        **fitness_kwargs: Union[bool, int, float]
     ):
         """Synchronous Population-Based Training (Jaderberg et al., 2017)
         Reference: https://arxiv.org/abs/1711.09846"""
-        super().__init__(popsize, num_dims, pholder_params)
+        super().__init__(popsize, num_dims, pholder_params, **fitness_kwargs)
         self.strategy_name = "PBT"
 
     @property

diff --git a/evosax/strategies/persistent_es.py b/evosax/strategies/persistent_es.py
@@ -40,12 +40,13 @@ def __init__(
         num_dims: Optional[int] = None,
         pholder_params: Optional[Union[chex.ArrayTree, chex.Array]] = None,
         opt_name: str = "adam",
+        **fitness_kwargs: Union[bool, int, float]
     ):
         """Persistent ES (Vicol et al., 2021).
         Reference: http://proceedings.mlr.press/v139/vicol21a.html
         Inspired by: http://proceedings.mlr.press/v139/vicol21a/vicol21a-supp.pdf
         """
-        super().__init__(popsize, num_dims, pholder_params)
+        super().__init__(popsize, num_dims, pholder_params, **fitness_kwargs)
         assert not self.popsize & 1, "Population size must be even"
         assert opt_name in ["sgd", "adam", "rmsprop", "clipup"]
         self.optimizer = GradientOptimizer[opt_name](self.num_dims)

diff --git a/evosax/strategies/pgpe.py b/evosax/strategies/pgpe.py
@@ -39,11 +39,12 @@ def __init__(
         pholder_params: Optional[Union[chex.ArrayTree, chex.Array]] = None,
         elite_ratio: float = 1.0,
         opt_name: str = "adam",
+        **fitness_kwargs: Union[bool, int, float]
     ):
         """PGPE (e.g. Sehnke et al., 2010)
         Reference: https://tinyurl.com/2p8bn956
         Inspired by: https://github.com/hardmaru/estool/blob/master/es.py"""
-        super().__init__(popsize, num_dims, pholder_params)
+        super().__init__(popsize, num_dims, pholder_params, **fitness_kwargs)
         assert 0 <= elite_ratio <= 1
         self.elite_ratio = elite_ratio
         self.elite_popsize = max(1, int(self.popsize / 2 * self.elite_ratio))

diff --git a/evosax/strategies/pso.py b/evosax/strategies/pso.py
@@ -36,10 +36,11 @@ def __init__(
         popsize: int,
         num_dims: Optional[int] = None,
         pholder_params: Optional[Union[chex.ArrayTree, chex.Array]] = None,
+        **fitness_kwargs: Union[bool, int, float]
     ):
         """Particle Swarm Optimization (Kennedy & Eberhart, 1995)
         Reference: https://ieeexplore.ieee.org/document/488968"""
-        super().__init__(popsize, num_dims, pholder_params)
+        super().__init__(popsize, num_dims, pholder_params, **fitness_kwargs)
         self.strategy_name = "PSO"
 
     @property

diff --git a/evosax/strategies/rm_es.py b/evosax/strategies/rm_es.py
@@ -67,11 +67,12 @@ def __init__(
         pholder_params: Optional[Union[chex.ArrayTree, chex.Array]] = None,
         elite_ratio: float = 0.5,
         memory_size: int = 10,
+        **fitness_kwargs: Union[bool, int, float]
     ):
         """Rank-m ES (Li & Zhang, 2017)
         Reference: https://ieeexplore.ieee.org/document/8080257
         """
-        super().__init__(popsize, num_dims, pholder_params)
+        super().__init__(popsize, num_dims, pholder_params, **fitness_kwargs)
         assert 0 <= elite_ratio <= 1
         self.elite_ratio = elite_ratio
         self.elite_popsize = max(1, int(self.popsize * self.elite_ratio))

diff --git a/evosax/strategies/samr_ga.py b/evosax/strategies/samr_ga.py
@@ -35,10 +35,11 @@ def __init__(
         num_dims: Optional[int] = None,
         pholder_params: Optional[Union[chex.ArrayTree, chex.Array]] = None,
         elite_ratio: float = 0.0,
+        **fitness_kwargs: Union[bool, int, float]
     ):
         """Self-Adaptation Mutation Rate GA."""
 
-        super().__init__(popsize, num_dims, pholder_params)
+        super().__init__(popsize, num_dims, pholder_params, **fitness_kwargs)
         self.elite_ratio = elite_ratio
         self.elite_popsize = max(1, int(self.popsize * self.elite_ratio))
         self.strategy_name = "SAMR_GA"

diff --git a/evosax/strategies/sep_cma_es.py b/evosax/strategies/sep_cma_es.py
@@ -63,12 +63,13 @@ def __init__(
         num_dims: Optional[int] = None,
         pholder_params: Optional[Union[chex.ArrayTree, chex.Array]] = None,
         elite_ratio: float = 0.5,
+        **fitness_kwargs: Union[bool, int, float]
     ):
         """Separable CMA-ES (e.g. Ros & Hansen, 2008)
         Reference: https://hal.inria.fr/inria-00287367/document
         Inspired by: github.com/CyberAgentAILab/cmaes/blob/main/cmaes/_sepcma.py
         """
-        super().__init__(popsize, num_dims, pholder_params)
+        super().__init__(popsize, num_dims, pholder_params, **fitness_kwargs)
         assert 0 <= elite_ratio <= 1
         self.elite_ratio = elite_ratio
         self.elite_popsize = max(1, int(self.popsize * self.elite_ratio))

diff --git a/evosax/strategies/sim_anneal.py b/evosax/strategies/sim_anneal.py
@@ -38,11 +38,12 @@ def __init__(
         popsize: int,
         num_dims: Optional[int] = None,
         pholder_params: Optional[Union[chex.ArrayTree, chex.Array]] = None,
+        **fitness_kwargs: Union[bool, int, float]
     ):
         """Simulated Annealing (Rasdi Rere et al., 2015)
         Reference: https://www.sciencedirect.com/science/article/pii/S1877050915035759
         """
-        super().__init__(popsize, num_dims, pholder_params)
+        super().__init__(popsize, num_dims, pholder_params, **fitness_kwargs)
         self.strategy_name = "SimAnneal"
 
     @property

diff --git a/evosax/strategies/simple_es.py b/evosax/strategies/simple_es.py
@@ -34,11 +34,12 @@ def __init__(
         num_dims: Optional[int] = None,
         pholder_params: Optional[Union[chex.ArrayTree, chex.Array]] = None,
         elite_ratio: float = 0.5,
+        **fitness_kwargs: Union[bool, int, float]
     ):
         """Simple Gaussian Evolution Strategy (Rechenberg, 1975)
         Reference: https://onlinelibrary.wiley.com/doi/abs/10.1002/fedr.19750860506
         Inspired by: https://github.com/hardmaru/estool/blob/master/es.py"""
-        super().__init__(popsize, num_dims, pholder_params)
+        super().__init__(popsize, num_dims, pholder_params, **fitness_kwargs)
         self.elite_ratio = elite_ratio
         self.elite_popsize = max(1, int(self.popsize * self.elite_ratio))
         self.strategy_name = "SimpleES"

diff --git a/evosax/strategies/simple_ga.py b/evosax/strategies/simple_ga.py
@@ -36,12 +36,13 @@ def __init__(
         num_dims: Optional[int] = None,
         pholder_params: Optional[Union[chex.ArrayTree, chex.Array]] = None,
         elite_ratio: float = 0.5,
+        **fitness_kwargs: Union[bool, int, float]
     ):
         """Simple Genetic Algorithm (Such et al., 2017)
         Reference: https://arxiv.org/abs/1712.06567
         Inspired by: https://github.com/hardmaru/estool/blob/master/es.py"""
 
-        super().__init__(popsize, num_dims, pholder_params)
+        super().__init__(popsize, num_dims, pholder_params, **fitness_kwargs)
         self.elite_ratio = elite_ratio
         self.elite_popsize = max(1, int(self.popsize * self.elite_ratio))
         self.strategy_name = "SimpleGA"

diff --git a/evosax/strategies/snes.py b/evosax/strategies/snes.py
@@ -44,11 +44,12 @@ def __init__(
         popsize: int,
         num_dims: Optional[int] = None,
         pholder_params: Optional[Union[chex.ArrayTree, chex.Array]] = None,
+        **fitness_kwargs: Union[bool, int, float]
     ):
         """Separable Exponential Natural ES (Wierstra et al., 2014)
         Reference: https://www.jmlr.org/papers/volume15/wierstra14a/wierstra14a.pdf
         """
-        super().__init__(popsize, num_dims, pholder_params)
+        super().__init__(popsize, num_dims, pholder_params, **fitness_kwargs)
         self.strategy_name = "SNES"
 
     @property

diff --git a/evosax/strategies/xnes.py b/evosax/strategies/xnes.py
@@ -41,11 +41,12 @@ def __init__(
         popsize: int,
         num_dims: Optional[int] = None,
         pholder_params: Optional[Union[chex.ArrayTree, chex.Array]] = None,
+        **fitness_kwargs: Union[bool, int, float]
     ):
         """Exponential Natural ES (Wierstra et al., 2014)
         Reference: https://www.jmlr.org/papers/volume15/wierstra14a/wierstra14a.pdf
         Inspired by: https://github.com/chanshing/xnes"""
-        super().__init__(popsize, num_dims, pholder_params)
+        super().__init__(popsize, num_dims, pholder_params, **fitness_kwargs)
         self.strategy_name = "xNES"
 
     @property

diff --git a/evosax/strategy.py b/evosax/strategy.py
@@ -4,7 +4,7 @@
 from typing import Tuple, Optional, Union
 from functools import partial
 from flax import struct
-from .utils import get_best_fitness_member, ParameterReshaper
+from .utils import get_best_fitness_member, ParameterReshaper, FitnessShaper
 
 
 @struct.dataclass
@@ -33,6 +33,7 @@ def __init__(
         popsize: int,
         num_dims: Optional[int] = None,
         pholder_params: Optional[Union[chex.ArrayTree, chex.Array]] = None,
+        **fitness_kwargs: Union[bool, int, float]
     ):
         """Base Class for an Evolution Strategy."""
         self.popsize = popsize
@@ -48,6 +49,9 @@ def __init__(
             self.num_dims is not None
         ), "Provide either num_dims or pholder_params to strategy."
 
+        # Setup optional fitness shaper
+        self.fitness_shaper = FitnessShaper(**fitness_kwargs)
+
     @property
     def default_params(self) -> EvoParams:
         """Return default parameters of evolution strategy."""
@@ -108,8 +112,11 @@ def tell(
         if self.use_param_reshaper:
             x = self.param_reshaper.flatten(x)
 
+        # Perform fitness reshaping inside of strategy tell call (if desired)
+        fitness_re = self.fitness_shaper.apply(x, fitness)
+
         # Update the search state based on strategy-specific update
-        state = self.tell_strategy(x, fitness, state, params)
+        state = self.tell_strategy(x, fitness_re, state, params)
 
         # Check if there is a new best member & update trackers
         best_member, best_fitness = get_best_fitness_member(x, fitness, state)