Add Diversifier QD Meta Algorithm - JAX backend

[dietmarwo]

This PR adds a new JAX-based QD meta algorithm called Diversifier. It is a generalization of CMA-ME.

It uses a MAP-Elites archive not for solution candidate generation, but only to modify the fitness values told (via tell) to the wrapped algorithm. This modification changes the fitness ranking of the population to favor exploration over exploitation. Tested with CR-FM-NES and CMA-ES, but other wrapped algorithms may work as well. Based on fcmaes diversifier.py (see MapElites.adoc).

The generalization over CMA-ME is necessary in the EvoJAX context, because CMA-ES struggles with a very high number of decision variables. Therefore CR-FM-NES-ME is superior here - as possibly are other not yet tested alternatives.

https://doi.org/10.1145/2739480.2754664 proposes the QD score (sum of fitness values of all elites in the map) as metric for comparison.

For Brax-Ant CR-FM-NES-ME (Diversifier applied to CR-FM-NES), compared with MAP-Elites, reaches a higher QD-score for high iteration numbers (see details below). So MAP-Elites should only be preferred for a low evaluation budget or if you want to maximize the number of occupied niches.

On a NVIDIA 3090 + AMD 5950, Linux Mint with optimized configurations we measured:

• MAP-Elites has a the same optimizer overhead (evaluation/sec rate for the same popsize).
• MAP-Elites has a higher number of occupied niches.

but

• CR-FM-NES-ME has a much higher QD score and found a better global optimum for a high evaluation budget.

Detailed measurements for the Brax-Ant example (NVIDIA 3090 + AMD 5950, Linux Mint):

After 20 minutes MAP-Elites is in the lead, but slows down from there. CR_FM_NES-ME continues to improve until 500 minutes / 8 million evaluations. CR_FM_NES-ME can even produce a good global optimum - 4107 - thereby still occupying 6138 niches with a mean score of 1208. After 500 minutes MAP-Elites continues to improve where CR_FM_NES-ME does not, but at that time CR_FM_NES-ME has a >70% lead in score.

CR_FM_NES-ME with init-std = 0.159, popsize = 512, fitness_weight 0.0

20 min QD score: 1692282 occupied: 4936 max score: 558 mean score: 342 evaluations: 263680 50 min QD score: 2724260 occupied: 5628 max score: 918 mean score: 484 evaluations: 704512 100 min QD score: 4289807 occupied: 6087 max score: 1442 mean score: 704 evaluations: 1496576 200 min QD score: 5928753 occupied: 6138 max score: 2363 mean score: 965 evaluations: 3072000 300 min QD score: 6524518 occupied: 6138 max score: 2862 mean score: 1063 evaluations: 4710400 400 min QD score: 7353257 occupied: 6138 max score: 3889 mean score: 1198 evaluations: 6348800 500 min QD score: 7418018 occupied: 6138 max score: 4107 mean score: 1208 evaluations: 7884800 600 min QD score: 7444092 occupied: 6138 max score: 4211 mean score: 1212 evaluations: 9523200

MAP-Elites iso-sigma = 0.05, line-sigma = 0.2, popsize = 1024: (line-sigma = 0.3 is worse)

20 min QD score: 2509773 occupied: 5621 max score: 643 mean score: 446 evaluations: 346112 50 min QD score: 3022521 occupied: 6375 max score: 724 mean score: 474 evaluations: 915456 100 min QD score: 3383041 occupied: 6786 max score: 769 mean score: 498 evaluations: 1941504 200 min QD score: 3713977 occupied: 7107 max score: 825 mean score: 522 evaluations: 3936256 300 min QD score: 3915492 occupied: 7265 max score: 927 mean score: 538 evaluations: 5922816 400 min QD score: 4065677 occupied: 7400 max score: 927 mean score: 549 evaluations: 7941120 500 min QD score: 4179020 occupied: 7498 max score: 927 mean score: 557 evaluations: 9958400 600 min QD score: 4272665 occupied: 7566 max score: 927 mean score: 564 evaluations: 12083200 700 min QD score: 4351397 occupied: 7632 max score: 941 mean score: 570 evaluations: 14094336 800 min QD score: 4415351 occupied: 7675 max score: 1003 mean score: 575 evaluations: 16040960

These results indicate that it should be possible to apply MAP-Elites to the resulting CR_FM_NES-ME archive to further improve occupancy and score. As algorithm wrapped by Diversifier,py CRFMNES can be replaced by FCRFMC (same algorithm but implemented in C++). We got the same results, but this may reduce the GPU load for smaller GPUs/TPUs and is definively advantageous for CPU alone executions. On the Nvidia 3090 CRFMNES is slightly faster.

Note that 'fitness_weight' is a concept neither used in CMA-ME nor in fcmaes fcmaes diversifier. All these use implicitely fitness_weight=0. For fcmaes the reason is that there are other means to improve the elites of a given map, so the focus is on exploration here. We use as default fitness_weight=0, because for Brax Ant the final QD score is higher - but the final global optimum found is lower.

fcmaes even supports sequences of wrapped algorithms, something probably not relevant for EvoJAX.

Increasing popsize to 1024 closes the evaluations / sec gap to MAP-Elites, the rate is 34% higher than with popsize = 512. But popsize = 1024 seems to produce lower occupancy - which is quite suprising:

CR_FM_NES-ME with init-std = 0.159, popsize = 1024, fitness_weight 0.0

20 min QD score: 1864258 occupied: 4856 max score: 538 mean score: 383 evaluations: 350208 50 min QD score: 2702674 occupied: 5402 max score: 848 mean score: 500 evaluations: 905216 100 min QD score: 3853807 occupied: 5873 max score: 1288 mean score: 656 evaluations: 1879040 200 min QD score: 5005292 occupied: 5947 max score: 1781 mean score: 841 evaluations: 3891200 300 min QD score: 6425120 occupied: 5963 max score: 2936 mean score: 1077 evaluations: 5963776 400 min QD score: 7103424 occupied: 5976 max score: 3783 mean score: 1188 evaluations: 8192000 500 min QD score: 7282457 occupied: 5980 max score: 4111 mean score: 1217 evaluations: 10240000 600 min QD score: 7371868 occupied: 5982 max score: 4227 mean score: 1232 evaluations: 12288000 700 min QD score: 7405531 occupied: 5982 max score: 4276 mean score: 1237 evaluations: 14336000 800 min QD score: 7464371 occupied: 5983 max score: 4307 mean score: 1247 evaluations: 16384000 900 min QD score: 7509290 occupied: 5989 max score: 4325 mean score: 1253 evaluations: 18432000 1000min QD score: 7514184 occupied: 5989 max score: 4342 mean score: 1254 evaluations: 20480000

But why can't we have our cake and eat it too?

This is not part of the PR but discusses what could be done in the future:

Both Diversifier and MAP-Elites share the same archive management. They differ only in population generation. In the future both could be unified into a single MD solver - still called MAP-Elites. This new implementation could randomly chose the way "ask" works. We define a probability, a wrapped solver is used instead of the standard mechanism. If this probability is 0, we have the old MAP-Elites. If it is 1.0, we have Diversifier. The interesting question is: What happens for values in between? Lets try 0.5. This can easily be implemented as:

    def ask(self) -> jnp.ndarray:
        self.key, key = jax.random.split(self.key)
        if jax.random.uniform(key) > 0.5: # a parameter to play with
            self.population = self.solver.ask() # population from wrapped solver
            self.solver_asked = True
        else: # population from MA-Elites generator
            self.key, mutate_key, parents = self._sample_parents(
                                key=self.key,
                                occupancy=self.occupancy_lattice,
                                params=self.params_lattice)      
            self.population = self._gen_pop(parents, mutate_key)
            self.solver_asked = False
        return self.population

    def tell(self, fitness: Union[np.ndarray, jnp.ndarray]) -> None:      
        if self.solver_asked:   
            lattice_fitness = self.fitness_lattice[self.bin_idx]
            to_tell = self._get_to_tell(fitness, lattice_fitness, self.fitness_weight)
            self.solver.tell(to_tell)
        # update lattice 

MAP-Elites + CR_FM_NES-ME with iso-sigma = 0.05, line-sigma = 0.2, init-std = 0.159, popsize = 1024, fitness_weight 0.0

20 min QD score: 2201387 occupied: 5266 max score: 672 mean score: 418 evaluations: 344064 50 min QD score: 2738691 occupied: 6105 max score: 672 mean score: 448 evaluations: 892928 100 min QD score: 3348423 occupied: 6656 max score: 857 mean score: 503 evaluations: 1859584 200 min QD score: 4233393 occupied: 7135 max score: 1103 mean score: 593 evaluations: 3851264 300 min QD score: 5139277 occupied: 7334 max score: 1586 mean score: 700 evaluations: 5893120 400 min QD score: 5776929 occupied: 7457 max score: 1884 mean score: 774 evaluations: 7943168 500 min QD score: 6098261 occupied: 7537 max score: 2104 mean score: 809 evaluations: 9947136 600 min QD score: 7240351 occupied: 7603 max score: 2890 mean score: 952 evaluations: 11999232 700 min QD score: 7800357 occupied: 7660 max score: 3421 mean score: 1018 evaluations: 14023680 800 min QD score: 8004109 occupied: 7699 max score: 3735 mean score: 1039 evaluations: 16000000

This is a 81% QD score increase compared to MAP-Elites alone thereby also improving occupancy.

900 min QD score: 8115904 occupied: 7744 max score: 3917 mean score: 1048 evaluations: 18140160 1000min QD score: 8195842 occupied: 7772 max score: 4019 mean score: 1054 evaluations: 20133888 1100min QD score: 8259249 occupied: 7799 max score: 4090 mean score: 1059 evaluations: 22155264 1200min QD score: 8319027 occupied: 7826 max score: 4130 mean score: 1062 evaluations: 24177664 1300min QD score: 8362034 occupied: 7847 max score: 4156 mean score: 1065 evaluations: 26213376

QD-score 8362034 probably is a challenge for each algorithm independent from the evaluation budget.

[dietmarwo]

Added a small fix for CRFMNES - applied jax.jit to generate_population. On my hardware it is only a small improvement. Tried that also for tell, but got no improvement there, so I left it as it is. I also did some experiments with MAP-Elites (applying simulated binary crossover or adding some random solutions to sample_parents) inspired by fcmaes mapelites, but Iso+LineDD proved to be optimal for Brax Ant. Only a small occupation improvement could be achieved, not worth to complicate things there.