Chris Lu, Oxford University

Abstract:

Recent advancements in Jax have made it possible to easily vectorise algorithms on the GPU. For one, this allows researchers to vectorise reinforcement learning environments, allowing them to quickly train reinforcement learning agents by sampling from thousands of environments simultaneously. More interestingly, however, this allows researchers to vectorise reinforcement learning algorithms themselves, allowing them to easily train thousands of agents simultaneously on a single GPU. This allows researchers to not only rapidly perform hyperparameter search and run multiple seeds simultaneously on limited hardware, but also perform meta-reinforcement learning with evolutionary strategies. I will go over the massive speedups we can achieve along with example papers and projects that have used this technique. I will then further discuss publicly available code that makes this technique broadly available to researchers.