JAX implementation of CMAES

[moskomule]

Hi, I'm really amazed by this library.

Currently, CMAES is just a wrapper. I implemented a JAX CMAES based on https://github.com/CyberAgentAILab/cmaes/.

[hardmaru]

Dear @moskomule,

Thanks for the contribution!

As discussed in the README file, can you please test your implementation on the 4 tasks, and report the results?

[moskomule]

Hi, thank you.

As you may know, CMA-ES includes eigendecomposition of covariance matrices, which makes it too slow for MNIST's CNN policy (at least in my environment). So, can I skip MNIST results?

[lerrytang]

Hi, thank you.

As you may know, CMA-ES includes eigendecomposition of covariance matrices, which makes it too slow for MNIST's CNN policy (at least in my environment). So, can I skip MNIST results?

Do you mind sharing some log to show how slow it is?
If it is due to some legacy hardware, I can run the tests.

[lerrytang]

There can be multiple CMA-ES implementations in EvoJAX, we are considering to use the naming rule: {algo_name}_{src_name / contributor}.
For example, we will change ours to something like CMAES_OriginalCPU, and yours may be CMAES_CyberAgent.

[moskomule]

OK, I will update so.

[moskomule]

Hi, thank you.
As you may know, CMA-ES includes eigendecomposition of covariance matrices, which makes it too slow for MNIST's CNN policy (at least in my environment). So, can I skip MNIST results?

Do you mind sharing some log to show how slow it is? If it is due to some legacy hardware, I can run the tests.

As for MNIST, I previously run a reduced CNN, but for the original CNN policy, even a single A6000 GPU throws OOM. So it's appreciated if you could run it.

One question: can I modify the examples to select solvers as an option, or do you prefer to keep the examples as is?

[lerrytang]

Thanks for the update.
I'd like to keep the examples as-is and avoid the situation where every new PR pertaining to algos has to change the examples.

[RobertTLange]

@moskomule I put together a small repository for running/logging the tests/benchmarks. Maybe this can be of help to: https://github.com/RobertTLange/evojax-benchmarks.

I have also been thinking of porting the evosax version of CMA-ES to evojax and have previously evolved small All-CNN-C-like networks with CMA-ES. Most of the large dimensionality results from the dense layer after the last convolutional layer. All-CNN architectures circumvent FC layers and perform the readout with a 1x1 convolution with appropriate window size.

Let me know if you would like to work on something like that together.

[lerrytang]

@RobertTLange Thanks!

@moskomule
I'm getting an OOM error on CartPole, does this happen to you too?
By the way, can you also fix the missing CLA error below?