[Bug fix] bugs in LM_MA_ES

Hi, thanks for such a nice library. I found a bug in LMMAES and will send a pullrequest.

Background Recently I'm interested in LM-MA-ES and use both implementation in numpy and one in jax, this library has.

Here this is my first test code and unfortunately only evosax's LM_MA_ES sampled nan-valued solution vectors and its step-size grew to infinity after several generations;

from evosax.problems import BBOBFitness
from evosax.strategies import LM_MA_ES
import jax.numpy as jnp
import jax

#Initialization
seed = 0
rng = jax.random.PRNGKey(seed)
dim = 2000 # high dimension, but same bug occured in lower dimension too.
sphere = BBOBFitness("Sphere", num_dims=dim, seed_id=2)
default_lamnda_and_m = int(4+jnp.floor(3*jnp.log(dim)))
mean = 3
sigma = 2
max_evals = 3*10**5*default_lamnda_and_m
criterion = 10**-10
strategy = LM_MA_ES(popsize=default_lamnda_and_m,
                    num_dims=dim, memory_size=default_lamnda_and_m                    )
es_params = strategy.default_params.replace(init_min=mean, init_max=mean, sigma_init=sigma)
state = strategy.initialize(rng, es_params)

best = float('inf')
t = 0
while best > criterion:
    rng, rng_gen, rng_eval = jax.random.split(rng, 3)
    x, state = strategy.ask(rng_gen, state, es_params)
    if jnp.isnan(x).any():
        print("nan exist in x.")

    fitness = sphere.rollout(rng_eval, x)
    state = strategy.tell(x, fitness, state, es_params)
    best = min(best, fitness.min())
    if (t + 1) % 50 == 0:
        print(f"# Gen: {t+1}|Fitness: {best:.12f}")
    if jnp.isnan(fitness).any():
        print('nan exists in fitness.')
        break
    if jnp.isnan(state.mean).any():
        print('nan exists in mean.')
    if jnp.isinf(state.sigma):
        print('sigma is inf.')
    t += 1

print(f"Gens:{t},evals:{t*default_lamnda_and_m},lamb:{default_lamnda_and_m}")

So I started bug fixing and finally I did. I will send my pull requests.

RobertTLange / evosax

[Bug fix] bugs in LM_MA_ES #55