Small improvements to the CMA-ES implementation

[Aladoro]

Thanks for maintaining this awesome repo :)

Two very small changes for CMA-ES:

1) Clip all eigenvalues of C in D to be >= 0 to avoid numerical underflow. Otherwise, this edge-case currently caused errors, as C_2 is computed by inverting D. It seems this was already partially covered for D < 0 (i.e., when numerical underflow in eigh would yield -0.0 as opposed to 0, which I believe is down to chance).

2) Convert for loops within list comprehensions for covariance computation into single einsum operations for efficiency.

[RobertTLange]

Thank you, Edoardo! The integration test doesn't seem to pass for SepCMA. Can you check what is wrong?

[Aladoro]

Oh, this error actually reveals another potential issue:

Currently, SepCMA only considers using the elites indexes when computing the weights - i.e., unintuitively, weights_truncated (shape [popsize,]) is a zero-padded copy of weights (shape [elite_popsize,])

This means that the current covariance update only considers positively-weighted elite samples (as opposed to CMA, where all samples are considered from Eq. 49)

However, this 'potential issue' is consistent with https://github.com/CyberAgentAILab/cmaes/blob/main/cmaes/_sepcma.py. Moreover, there is some ambiguity about whether this is actually warranted looking at https://hal.inria.fr/inria-00287367/document

For now, I have implemented the efficiency improvements, while keeping the old behavior exactly the same also for SepCMA. However, let me know what you think!

[RobertTLange]

Thank you, Edoardo! Looks good to me. I am merging it in and releasing it ASAP.

Re weighting in covariance estimation: Let's talk about this soon. I think you are right. I based the JAX version on the CyberAgent implementation, but using all candidates for covariance estimation might be more efficient!