Evosax

- Fix LES checkpoint loading from package data via `pkgutil`.

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.

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 solution 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 lr 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.

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 `State`.

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 `State`/`Params`.

python
from flax import struct
struct.dataclass
class State:
...


- The core strategy API now also works without `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.

Fixed

- Fix gym import bug and codecov patch tolerance.

Fixed

- Bug due to `subpops` import in `experimental`.