Authors: Michał Bortkiewicz, Władek Pałucki, Vivek Myers, Tadeusz Dziarmaga, Tomasz Arczewski, Łukasz Kuciński, Benjamin Eysenbach
Abstract: Self-supervision has the potential to transform reinforcement learning (RL),
paralleling the breakthroughs it has enabled in other areas of machine
learning. While self-supervised learning in other domains aims to find patterns
in a fixed dataset, self-supervised goal-conditioned reinforcement learning
(GCRL) agents discover new behaviors by learning from the goals achieved during
unstructured interaction with the environment. However, these methods have
failed to see similar success, both due to a lack of data from slow
environments as well as a lack of stable algorithms. We take a step toward
addressing both of these issues by releasing a high-performance codebase and
benchmark JaxGCRL for self-supervised GCRL, enabling researchers to train
agents for millions of environment steps in minutes on a single GPU. The key to
this performance is a combination of GPU-accelerated environments and a stable,
batched version of the contrastive reinforcement learning algorithm, based on
an infoNCE objective, that effectively makes use of this increased data
throughput. With this approach, we provide a foundation for future research in
self-supervised GCRL, enabling researchers to quickly iterate on new ideas and
evaluate them in a diverse set of challenging environments. Website + Code:
https://github.com/MichalBortkiewicz/JaxGCRL
Source: http://arxiv.org/abs/2408.11052v1
