🚀 Major refactor, optax support, bug fixes

[v0.2.0] - [03/2025]

This major update includes significant architectural changes, dependency updates, bug fixes, and naming standardization across the library.

Breaking Changes

• Increased minimum Python requirement to >=3.10
• Updated to JAX >=0.5.0
• Switched from setup.py to pyproject.toml for package configuration
• Restructured algorithms into two main types: PopulationBased and DistributionBased
• Removed ESLogger in favor of user-provided metrics_fn
• Removed ParameterReshape class (now handled automatically by algorithm classes)
• Removed n_devices parameter and all pmap calls to align with JAX's new sharding mechanisms
• Renamed Strategy.initialize to Strategy.init
• Renamed Fitness to Problem with .eval method replacing .rollout

Optimization and Dependencies

• Replaced custom optimizer with optax
  • Implemented ClipUp in Optax GradientTransformation
• Added ruff for code linting
• Removed chex dependency
  • Added custom types.py file with Solution, Fitness, and Metrics types
• Removed RBF (now using JAX's jax.scipy.stats.norm.pdf)

Library Structure Improvements

• Made all MA-ES strategies inherit from CMA-ES
  • Includes: cma_es, sep_cma_es, sv_cma_es, rm_es, ma_es, lm_ma_es
• For Stein Variational methods, implemented Strategies that inherit from CMA-ES/OpenES
• Refactored fitness shaping to use a direct function (fitness_shaping_fn) rather than an object
• Implemented SNES/xNES inheritance for weights
• Made all Genetic Algorithms return initial population for evaluation on first generation
• Moved Restart mechanism to core

Algorithm Fixes

• Fixed numerous bugs in evolution strategies:
  • Fixed xNES implementation
  • Fixed SimAnneal implementation
  • Fixed PBT implementation
  • Fixed GuidedES implementation
  • Fixed MR15-GA implementation
  • Fixed nan values in PGPE
  • Fixed nan values in RM-ES
  • Fixed incorrect z calculation in MA-ES
  • Fixed sep_cma_es where D was never updated
  • Fixed GESMR-GA algorithm

Naming Standardization

• Renamed EvoState to State and EvoParams to Params
• Renamed es_params/es_state to params/state
• Renamed params/x to solution or solutions
• Renamed rng to key
• Renamed gen_counter to generation_counter or counter
• Renamed gen to generation
• Renamed popsize to population_size
• Standardized naming between mean/sigma across algorithms
• Renamed get_eval_params to get_eval_solution
• Changed lrate to lr
• Renamed persistent_es.py to pes.py
• Renamed all *_strategy to _*
• Standardized all references to total_params and num_dims

Evolution Strategy Improvements

• Replaced jnp.finfo(jnp.float32).max with jnp.inf throughout codebase
• Avoided logic with None args in Strategy.__init__ and Strategy.ask
• Refactored std decay to use optax schedules
• Removed mean decay from within strategies
• Added correct implementation of c_m in MA-ES algorithms
• Fixed incorrect p notation (now properly uses p_σ and p_c)
• Updated all static_argnums to static_argnames
• Fixed GuideES to pass gradient within EvoState rather than implementing tell

Problem Modules

• Renamed "classic" to "bbo"
• Renamed "control_gym" to "gymnax"
• Added Brax support
• Improved gymnax integration
• Created base class for problems
• All problems now provide a sample_solution method
• Added noise models for BBO problems
• Added problem.eval, problem.sample_solution, dummy_solution methods

Documentation

• Updated all notebooks
• Fixed license mismatch on PyPI
• Fixed all TODOs within codebase
• Updated examples and experimental code
• Added tests for experimental features

Big shout out to @maxencefaldor 🙏

🦾 Added EvoTransformer ES & Small Fixes

[v0.1.6] - [03/2024]

Added

• Implemented Hill Climbing strategy as a simple baseline.
• Adds use_antithetic_sampling option to OpenAI-ES.
• Added EvoTransformer and EvoTF_ES strategy with example trained checkpoint.

Fixed

• Gradientless Descent best member replacement.

Changed

• SNES import DES weights directly and reuses code
• Made Sep_CMA_ES and OpenAI-ES use vector sigmas for EvoTransformer data collection.

🦎 Core Utilities Restructure & LES/LGA Ckpt Loading Fix

Added

• Adds string fitness_trafo option to FitnessShaper (e.g. z_score, etc.).
• Adds sigma_meta as kwarg to SAMR_GA and GESMR_GA.
• Adds sigma_init as kwarg to LGA and LES.
• Adds Noise-Reuse ES - NoiseReuseES - (Li et al., 2023) as a generalization of PES.
• Fix LES evolution path calculation and re-ran meta-training for checkpoint.

Fixed

• Fixed error in LGA resulting from elite_ratio=0.0 in sampling operator logit squeeze.
• Fixed range normalization in fitness transformation - range_norm_trafo - Thank you @yudonglee

Changed

• Refactored core modules and utilities. Learned evolution utilities now in subdirectory.

🐘 LGA, Warmstart strategy mean, ESLog, Fixes

[v0.1.4] - [04/2023]

Added

• Adds LGA checkpoint and optimizer class from Lange et al. (2023b).
• Adds optional init_mean to strategy.initialize to warm start strategy from e.g. pre-trained checkpoint.
• Adds n_devices option to every strategy to control reshaping for pmap in ParameterReshaper (if desired) explicitly.
• Adds mean_decay optional kwarg to LES for regularization.

Fixed

• Fix missing matplotlib requirement for BBOB Visualizer.
• Fix squeezing of sampled solutions in order to enable 1D optimization.
• Fix ESLog to work with ParameterReshaper reshaping of candidate solutions.
• Fix overflow errors in CMA-ES style ES when num_dims ** 2 is too large.

Changed

• Changed default gradient descent optimizer of ARS to Adam.

Fix PGPE bug & LES checkpoint archive 📚

[v0.1.3] - [03/2023]

• Finally solved checkpoint loading LES problem (needed MANIFEST.in)
• Fixed PGPE bug with regards to scaled noise.

FIX LES checkpoint loading 🙊

• Fix LES checkpoint loading from package data via pkgutil.

LES, Distributed ES & Fixes 🚀

Added

• Adds exponential decay of mean/weight regularization to ES that update mean (FD-ES and CMA variants). Simply provide mean_decay != 0.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.

• Adds experimental distributed ES, which sample directly on all devices (no longer only on host). Furthermore, we use pmean-like all reduce ops to construct z-scored fitness scores and gradient accumulations to update the mean estimate. So far only FD-gradient-based ES are supported. Major benefits: Scale with the number of devives and allow for larger populations/number of dimensions.

  • Supported distributed ES:
    • DistributedOpenES
  • Import via: from evosax.experimental.distributed import DistributedOpenES

• Adds RandomSearch as basic baseline.

• Adds LES (Lange et al., 2023) and a retrained trained checkpoint.

• Adds a separate example notebook for how to use the BBOBVisualizer.

Changed

• Sep_CMA_ES automatic hyperparameter calculation runs into int32 problems, when num_dims > 40k. We therefore clip the number to 40k for this calculation.

Fixed

• Fixed DES to also take flexible fitness_kwargs, temperature, sigma_init as inputs.
• Fixed PGPE exponential decay option to account for sigma update.

🦎 New strategies, API flexibility, small fixes

Added

• Adds a total_env_steps counter to both GymFitness and BraxFitness for easier sample efficiency comparability with RL algorithms.
• Support for new strategies/genetic algorithms
  • SAMR-GA (Clune et al., 2008)
  • GESMR-GA (Kumar et al., 2022)
  • SNES (Wierstra et al., 2014)
  • DES (Lange et al., 2022)
  • Guided ES (Maheswaranathan et al., 2018)
  • ASEBO (Choromanski et al., 2019)
  • CR-FM-NES (Nomura & Ono, 2022)
  • MR15-GA (Rechenberg, 1978)
• Adds full set of BBOB low-dimensional functions (BBOBFitness)
• Adds 2D visualizer animating sampled points (BBOBVisualizer)
• Adds Evosax2JAXWrapper to wrap all evosax strategies
• Adds Adan optimizer (Xie et al., 2022)

Changed

• ParameterReshaper can now be directly applied from within the strategy. You simply have to provide a pholder_params pytree at strategy instantiation (and no num_dims).
• FitnessShaper can also be directly applied from within the strategy. This makes it easier to track the best performing member across generations and addresses issue #32. Simply provide the fitness shaping settings as args to the strategy (maximize, centered_rank, ...)
• Removes Brax fitness (use EvoJAX version instead)
• Add lrate and sigma schedule to strategy instantiation

Fixed

• Fixed reward masking in GymFitness. Using jnp.sum(dones) >= 1 for cumulative return computation zeros out the final timestep, which is wrong. That's why there were problems with sparse reward gym environments (e.g. Mountain Car).
• Fixed PGPE sample indexing.
• Fixed weight decay. Falsely multiplied by -1 when maximization.

🚀 Indirect Encodings, EvoState/EvoParams, gymnax backend

[v0.0.9] - 15/06/2022

Added

• Base indirect encoding methods in experimental. Sofar support for:
  • Random projection-based decodings
  • Hypernetworks for MLP architectures
• Example notebook for infirect encodings.
• Example notebook for Brax control tasks and policy visualizations.
• Adds option to restart wrappers to copy_mean and only reset other parts of EvoState.

Changed

• Change problem wrappers to work with {"params": ...} dictionary. No longer need to define ParameterReshaper(net_params["params"]) to work without preselecting "params". Changed tests and notebooks accordingly.
• Restructured all strategies to work with flax structured dataclass and EvoState/EvoParams. Note that this will require different specification of hyperparameter settings e.g. via es_params = es_params.replace(cross_over_rate=0.9).

from flax import struct
@struct.dataclass
class EvoState:
    ...

• The core strategy API now also works without es_params being supplied in call. In this case we simply use the default settings.
• Moved all gym environment to (still private but soon to be released) gymnax.
• Updated all notebooks accordingly.

Fixed

• Makes ParameterReshaper work also with dm-haiku-style parameter dictionaries. Thanks to @vuoristo.

Another import bug fix

Fixed

• Fix gym import bug and codecov patch tolerance