Add mean_decay_coeff kwarg to relevant ES

RobertTLange · Dec 20, 2022 · ffaafb7 · ffaafb7
1 parent 3c7b751
commit ffaafb7
Show file tree

Hide file tree

Showing 26 changed files with 317 additions and 65 deletions.
diff --git a/.gitignore b/.gitignore
@@ -1,3 +1,4 @@
+anims/
 examples/experimental
 # bbob.py
 # Standard ROB excludes

diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -11,7 +11,19 @@
     - [ ] Wavelet Based Encoding (van Steenkiste, 2016)
     - [ ] CNN Hypernetwork (Ha - start with simple MLP)
 
-### [v0.1.0] - [TBD]
+### [v0.1.1] - [TBD]
+
+##### Added
+
+- Add exponential decay of mean/weight regularization to ES that update mean (FD-ES and CMA variants). Simply provide `mean_decay_coeff` != 1.0 argument at strategy instantiation to strategy. Note that covariance estimates may be a bit off, but this circumvents constant increase of mean norm due to stochastic process nature.
+
+##### Changed
+
+##### Fixed
+
+- Fixed DES to also take flexible `fitness_kwargs`, `temperature`, `sigma_init` as inputs
+
+### [v0.1.0] - [12/2022]
 
 ##### Added
 

diff --git a/evosax/strategies/ars.py b/evosax/strategies/ars.py
@@ -43,11 +43,18 @@ def __init__(
         sigma_init: float = 0.03,
         sigma_decay: float = 1.0,
         sigma_limit: float = 0.01,
+        mean_decay_coeff: float = 1.0,
         **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, **fitness_kwargs)
+        super().__init__(
+            popsize,
+            num_dims,
+            pholder_params,
+            mean_decay_coeff,
+            **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/asebo.py b/evosax/strategies/asebo.py
@@ -48,6 +48,7 @@ def __init__(
         sigma_init: float = 0.03,
         sigma_decay: float = 1.0,
         sigma_limit: float = 0.01,
+        mean_decay_coeff: float = 1.0,
         **fitness_kwargs: Union[bool, int, float],
     ):
         """ASEBO (Choromanski et al., 2019)
@@ -56,7 +57,13 @@ def __init__(
         1. We always sample a fixed population size per generation
         2. We keep a fixed archive of gradients to estimate the subspace
         """
-        super().__init__(popsize, num_dims, pholder_params, **fitness_kwargs)
+        super().__init__(
+            popsize,
+            num_dims,
+            pholder_params,
+            mean_decay_coeff,
+            **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/bipop_cma_es.py b/evosax/strategies/bipop_cma_es.py
@@ -26,6 +26,7 @@ def __init__(
         pholder_params: Optional[Union[chex.ArrayTree, chex.Array]] = None,
         elite_ratio: float = 0.5,
         sigma_init: float = 1.0,
+        mean_decay_coeff: float = 1.0,
         **fitness_kwargs: Union[bool, int, float]
     ):
         """BIPOP-CMA-ES (Hansen, 2009).
@@ -39,6 +40,7 @@ def __init__(
             pholder_params=pholder_params,
             elite_ratio=elite_ratio,
             sigma_init=sigma_init,
+            mean_decay_coeff=mean_decay_coeff,
             **fitness_kwargs,
         )
         from ..restarts import BIPOP_Restarter
@@ -47,7 +49,10 @@ def __init__(
         self.wrapped_strategy = BIPOP_Restarter(
             self.strategy,
             stop_criteria=[spread_criterion, cma_criterion],
-            strategy_kwargs={"elite_ratio": elite_ratio},
+            strategy_kwargs={
+                "elite_ratio": elite_ratio,
+                "mean_decay_coeff": mean_decay_coeff,
+            },
         )
 
     @property

diff --git a/evosax/strategies/cma_es.py b/evosax/strategies/cma_es.py
@@ -87,12 +87,19 @@ def __init__(
         pholder_params: Optional[Union[chex.ArrayTree, chex.Array]] = None,
         elite_ratio: float = 0.5,
         sigma_init: float = 1.0,
+        mean_decay_coeff: float = 1.0,
         **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, **fitness_kwargs)
+        super().__init__(
+            popsize,
+            num_dims,
+            pholder_params,
+            mean_decay_coeff,
+            **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/cr_fm_nes.py b/evosax/strategies/cr_fm_nes.py
@@ -80,12 +80,19 @@ def __init__(
         num_dims: Optional[int] = None,
         pholder_params: Optional[Union[chex.ArrayTree, chex.Array]] = None,
         sigma_init: float = 1.0,
+        mean_decay_coeff: float = 1.0,
         **fitness_kwargs: Union[bool, int, float]
     ):
         """Cost-Reduced Fast-Moving Natural ES (Nomura & Ono, 2022)
         Reference: https://arxiv.org/abs/2201.11422
         """
-        super().__init__(popsize, num_dims, pholder_params, **fitness_kwargs)
+        super().__init__(
+            popsize,
+            num_dims,
+            pholder_params,
+            mean_decay_coeff,
+            **fitness_kwargs
+        )
         assert not self.popsize & 1, "Population size must be even"
         self.strategy_name = "CR_FM_NES"
 

diff --git a/evosax/strategies/des.py b/evosax/strategies/des.py
@@ -22,7 +22,7 @@ class EvoParams:
     temperature: float = 12.5  # Temperature for softmax weights
     lrate_sigma: float = 0.1  # Learning rate for population std
     lrate_mean: float = 1.0  # Learning rate for population mean
-    sigma_init: float = 1.0  # Standard deviation
+    sigma_init: float = 0.1  # Standard deviation
     init_min: float = 0.0
     init_max: float = 0.0
     clip_min: float = -jnp.finfo(jnp.float32).max
@@ -45,15 +45,29 @@ def __init__(
         popsize: int,
         num_dims: Optional[int] = None,
         pholder_params: Optional[Union[chex.ArrayTree, chex.Array]] = None,
+        temperature: float = 12.5,
+        sigma_init: float = 0.1,
+        mean_decay_coeff: float = 1.0,
+        **fitness_kwargs: Union[bool, int, float]
     ):
         """Discovered Evolution Strategy (Lange et al., 2022)"""
-        super().__init__(popsize, num_dims, pholder_params)
+        super().__init__(
+            popsize,
+            num_dims,
+            pholder_params,
+            mean_decay_coeff,
+            **fitness_kwargs
+        )
         self.strategy_name = "DES"
+        self.temperature = temperature
+        self.sigma_init = sigma_init
 
     @property
     def params_strategy(self) -> EvoParams:
         """Return default parameters of evolution strategy."""
-        return EvoParams()
+        return EvoParams(
+            temperature=self.temperature, sigma_init=self.sigma_init
+        )
 
     def initialize_strategy(
         self, rng: chex.PRNGKey, params: EvoParams

diff --git a/evosax/strategies/esmc.py b/evosax/strategies/esmc.py
@@ -44,12 +44,19 @@ def __init__(
         sigma_init: float = 0.03,
         sigma_decay: float = 1.0,
         sigma_limit: float = 0.01,
+        mean_decay_coeff: float = 1.0,
         **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, **fitness_kwargs)
+        super().__init__(
+            popsize,
+            num_dims,
+            pholder_params,
+            mean_decay_coeff,
+            **fitness_kwargs
+        )
         assert self.popsize & 1, "Population size must be odd"
         assert opt_name in ["sgd", "adam", "rmsprop", "clipup", "adan"]
         self.optimizer = GradientOptimizer[opt_name](self.num_dims)

diff --git a/evosax/strategies/full_iamalgam.py b/evosax/strategies/full_iamalgam.py
@@ -49,12 +49,19 @@ def __init__(
         sigma_init: float = 0.0,
         sigma_decay: float = 0.99,
         sigma_limit: float = 0.0,
+        mean_decay_coeff: float = 1.0,
         **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, **fitness_kwargs)
+        super().__init__(
+            popsize,
+            num_dims,
+            pholder_params,
+            mean_decay_coeff,
+            **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/gld.py b/evosax/strategies/gld.py
@@ -31,11 +31,18 @@ def __init__(
         popsize: int,
         num_dims: Optional[int] = None,
         pholder_params: Optional[Union[chex.ArrayTree, chex.Array]] = None,
+        mean_decay_coeff: float = 1.0,
         **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, **fitness_kwargs)
+        super().__init__(
+            popsize,
+            num_dims,
+            pholder_params,
+            mean_decay_coeff,
+            **fitness_kwargs
+        )
         self.strategy_name = "GLD"
 
     @property

diff --git a/evosax/strategies/guided_es.py b/evosax/strategies/guided_es.py
@@ -48,13 +48,20 @@ def __init__(
         sigma_init: float = 0.03,
         sigma_decay: float = 1.0,
         sigma_limit: float = 0.01,
+        mean_decay_coeff: float = 1.0,
         **fitness_kwargs: Union[bool, int, float],
     ):
         """Guided ES (Maheswaranathan et al., 2018)
         Reference: https://arxiv.org/abs/1806.10230
         Note that there are a couple of JAX-based adaptations:
         """
-        super().__init__(popsize, num_dims, pholder_params, **fitness_kwargs)
+        super().__init__(
+            popsize,
+            num_dims,
+            pholder_params,
+            mean_decay_coeff,
+            **fitness_kwargs,
+        )
         assert not self.popsize & 1, "Population size must be even"
         assert opt_name in ["sgd", "adam", "rmsprop", "clipup", "adan"]
         assert (

diff --git a/evosax/strategies/indep_iamalgam.py b/evosax/strategies/indep_iamalgam.py
@@ -54,12 +54,19 @@ def __init__(
         sigma_init: float = 0.0,
         sigma_decay: float = 0.99,
         sigma_limit: float = 0.0,
+        mean_decay_coeff: float = 1.0,
         **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, **fitness_kwargs)
+        super().__init__(
+            popsize,
+            num_dims,
+            pholder_params,
+            mean_decay_coeff,
+            **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
@@ -26,6 +26,7 @@ def __init__(
         pholder_params: Optional[Union[chex.ArrayTree, chex.Array]] = None,
         elite_ratio: float = 0.5,
         sigma_init: float = 1.0,
+        mean_decay_coeff: float = 1.0,
         **fitness_kwargs: Union[bool, int, float]
     ):
         """IPOP-CMA-ES (Auer & Hansen, 2005).
@@ -39,6 +40,7 @@ def __init__(
             pholder_params=pholder_params,
             elite_ratio=elite_ratio,
             sigma_init=sigma_init,
+            mean_decay_coeff=mean_decay_coeff,
             **fitness_kwargs
         )
         from ..restarts import IPOP_Restarter
@@ -47,7 +49,10 @@ def __init__(
         self.wrapped_strategy = IPOP_Restarter(
             self.strategy,
             stop_criteria=[spread_criterion, cma_criterion],
-            strategy_kwargs={"elite_ratio": elite_ratio},
+            strategy_kwargs={
+                "elite_ratio": elite_ratio,
+                "mean_decay_coeff": mean_decay_coeff,
+            },
         )
 
     @property

diff --git a/evosax/strategies/lm_ma_es.py b/evosax/strategies/lm_ma_es.py
@@ -47,12 +47,19 @@ def __init__(
         elite_ratio: float = 0.5,
         memory_size: int = 10,
         sigma_init: float = 1.0,
+        mean_decay_coeff: float = 1.0,
         **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, **fitness_kwargs)
+        super().__init__(
+            popsize,
+            num_dims,
+            pholder_params,
+            mean_decay_coeff,
+            **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
@@ -43,12 +43,19 @@ def __init__(
         pholder_params: Optional[Union[chex.ArrayTree, chex.Array]] = None,
         elite_ratio: float = 0.5,
         sigma_init: float = 1.0,
+        mean_decay_coeff: float = 1.0,
         **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, **fitness_kwargs)
+        super().__init__(
+            popsize,
+            num_dims,
+            pholder_params,
+            mean_decay_coeff,
+            **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
@@ -42,12 +42,19 @@ def __init__(
         sigma_init: float = 0.03,
         sigma_decay: float = 1.0,
         sigma_limit: float = 0.01,
+        mean_decay_coeff: float = 1.0,
         **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, **fitness_kwargs)
+        super().__init__(
+            popsize,
+            num_dims,
+            pholder_params,
+            mean_decay_coeff,
+            **fitness_kwargs
+        )
         assert not self.popsize & 1, "Population size must be even"
         assert opt_name in ["sgd", "adam", "rmsprop", "clipup", "adan"]
         self.optimizer = GradientOptimizer[opt_name](self.num_dims)

diff --git a/evosax/strategies/persistent_es.py b/evosax/strategies/persistent_es.py
@@ -46,13 +46,20 @@ def __init__(
         sigma_init: float = 0.03,
         sigma_decay: float = 1.0,
         sigma_limit: float = 0.01,
+        mean_decay_coeff: float = 1.0,
         **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, **fitness_kwargs)
+        super().__init__(
+            popsize,
+            num_dims,
+            pholder_params,
+            mean_decay_coeff,
+            **fitness_kwargs
+        )
         assert not self.popsize & 1, "Population size must be even"
         assert opt_name in ["sgd", "adam", "rmsprop", "clipup", "adan"]
         self.optimizer = GradientOptimizer[opt_name](self.num_dims)