Kheperax: a Lightweight JAX-based Robot Control Environment for Benchmarking Quality-Diversity Algorithms

Abstract

This work introduces a new lightweight and massively parallelizable implementation of a Quality-Diversity (QD) task: the libfastsim maze. This QD task involves finding a collection of neural network controllers navigating a robot to diverse positions in a maze. The proposed implementation, called Kheperax, can be used as a benchmark task for standard QD algorithms, but also for Model-based, Unsupervised and Uncertain QD algorithms. It can automatically run and parallelize parameter evaluations on hardware accelerators, such as Graphical Processing Units (GPUs). When evaluating large batches of parameters, Kheperax is at least 4 times faster than the original libfastsim implementation. The source code is available online1.