QDax: A Library for Quality-Diversity and Population-based Algorithms with Hardware Acceleration

Felix Chalumeau | Bryan Lim 1 | Raphael Boige | Maxime Allard 1 | Luca Grillotti 1 | Manon Flageat 1 | Valentin Mace | Guillaume Richard | Arthur Flajolet | Thomas Pierrot | Antoine Cully 1

1 Imperial College London

Published

ABSTRACT

QDax is an open-source library with a streamlined and modular API for Quality-Diversity (QD) optimization algorithms in jax. The library serves as a versatile tool for optimization purposes, ranging from black-box optimization to continuous control. qdax offers implementations of popular QD, Neuroevolution, and Reinforcement Learning (RL) algorithms, supported by various examples. All the implementations can be just-in-time compiled with Jax, facilitating efficient execution across multiple accelerators, including GPUs and TPUs. These implementations effectively demonstrate the framework’s flexibility and user-friendliness, easing experimentation for research purposes. Furthermore, the library is thoroughly documented and tested with 95% coverage.