RobertTLange / evosax

Evolution Strategies in JAX 🦎

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

Interest in CR-FM-NES Strategy?

nhansendev opened this issue · comments

After some experiments with xNES demonstrated impressive performance I started searching for the latest "successor", if any existed. This led me to Fast-Moving NES (FM-NES) and its extension Cost-Reduced Fast Moving NES (CR-FM-NES).

When compared to xNES and CMA-ES, FM-NES demonstrated similar, or better performance per their paper:
" The results show that the performance of FM-NES is better than that of the other methods on the Ellipsoid problem, the Rosenbrock problem, the Cigar problem, the IC-Rosenbrock problem, and the IC-Cigar problem. In particular, for the Rosenbrock problem, the Cigar problem, the IC-Rosenbrock problem, and the IC-Cigar problem, which have ridge structure, FM-NES has succeeded in finding the optimal solutions about 1.71, 1.78, 1.67, and 1.42 times faster than DX-NES-IC, respectively. Furthermore, FM-NES showed better performance than CMA-ES on the unconstrained problems with ridge structure, the Rosenbrock problem and the Cigar problem, on which DX-NES-IC showed worse performance than CMAES. On the other hand, on the Sphere problem, the IC-Sphere problem, and the IC-Ellipsoid problem, FM-NES shows almost the same performance as DX-NES-IC and better performance than xNES, CMA-ES, xNES/R, and CMA-ES/R. Overall, the result suggests the effectiveness of FM-NES. "

When compared to FM-NES, Sep-CMA, and VD-CMA, CR-FM-NES demonstrated similar, or better performance per their paper:
" Our experimental results suggest that compared to FM-NES, CR-FM-NES has achieved significant speedup on some problems or at least achieved competitive performance. Furthermore, CR-FM-NES showed favorable performance over VD-CMA [35] and SepCMA [38].
The main limitation of the proposed method lies on the restricted representation of the covariance matrix. As noted in [35], CR-FM-NES will fail optimization when the covariance matrix cannot approximate the inverse Hessian of the objective function. "

performance

The authors provided a publicly available, numpy-based implementation of CR-FM-NES on github. I thought it could be a good fit for evosax, though it would be challenging for me to implement myself as I am still new to the evosax framework and jax in general.

Is this something you are interested in?

Hi @Obliman, Thank you for bringing this up! Yes, I am very much interested in adding more scalable ES to evosax. E.g. a VkD-CMA version has been on my bucket list for a while. After checking the numpy implementation it looks fairly easy to rewrite it in JAX. Do you want to give it a try, open a PR and I help out when you get stuck? Otherwise, this would go on my ToDo-list for after the end of my summer internship. Cheers, Rob

I've started on it in a fork of the repository here.
So far it's just been converted from numpy to jax.numpy, but it's a start.

Thank you for giving it a first shot! If you want to you can try to make your implementation match the evosax ask/tell API and separate out the es_state and es_params. Otherwise, I would continue with your first port attempt at a later point.

FYI: After some experimentation I've found that the control over VRAM usage that I get with Numpy+Pytorch is better for my application than jax. Due to this I do not intend to work on the evosax variant of the algorithm for the foreseeable future.

By the way, thanks for building this library. It's been very informative and great for experimenting with ES methods.

Added in PR #34 and released in v0.10.0