Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add a Python/JAX port of CR-FM-NES #46

Merged
merged 5 commits into from Oct 4, 2022

Conversation

dietmarwo
Copy link
Contributor

@dietmarwo dietmarwo commented Sep 29, 2022

This PR adds a Python/JAX port of Fast Moving Natural Evolution Strategy for High-Dimensional Problems (CR-FM-NES), see https://arxiv.org/abs/2201.11422 . Derived from https://github.com/nomuramasahir0/crfmnes.

This variant is slightly faster than FCRFMC (the C++ port) on fast GPUs/TPUs, but slower on CPUs and for smaller dimensions.
It uses 32 bit accuracy (FCRFMC uses 64 bit) which mostly doesn't harm the convergence (with Waterworld MA being the exception for very high iteration numbers).

Wall time and convergence is mostly comparable with PGPE (as FCRFMC) for the benchmarks. Slower in the beginning, but improving at higher iterations.

Since there are no for-loops I found no beneficial applications of 'jax.jit', just converted most 'np.arrays' into 'jnp.arrays' deployed on the GPUs/TPUs.

def sort_indices_by(evals: np.ndarray, z: jnp.ndarray) -> jnp.ndarray:

uses not evals: jnp.ndarray because this slowed things down on my NVIDIA 3090.

Since this is Python code, no missing shared libraries on Ubuntu 18 this time.

Added test results for CRFMNES (this Python implementation) at EvoJax.adoc.

@lerrytang lerrytang self-assigned this Oct 2, 2022
@lerrytang lerrytang merged commit 1ab663f into google:main Oct 4, 2022
@lerrytang
Copy link
Contributor

Thanks for the contribution!
Together with FCRFMC, will you send a pr to update this page?

@dietmarwo
Copy link
Contributor Author

On vacation until end of next week. Will update the page when I am back.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

None yet

2 participants