Assistax: A Hardware-Accelerated Reinforcement Learning Benchmark for Assistive Robotics

TL;DR

Assistax introduces a fast, open-source reinforcement learning benchmark tailored for assistive robotics, utilizing hardware acceleration and multi-agent RL to evaluate robot-human interaction capabilities.

Contribution

It presents a novel, hardware-accelerated benchmark specifically designed for assistive robotics, addressing real-world embodied interaction challenges.

Findings

Assistax achieves up to 370x faster training times compared to CPU-based methods.

Provides reliable baselines for RL and MARL algorithms in assistive robotics.

Establishes a practical benchmark for advancing RL in real-world assistive scenarios.

Abstract

The development of reinforcement learning (RL) algorithms has been largely driven by ambitious challenge tasks and benchmarks. Games have dominated RL benchmarks because they present relevant challenges, are inexpensive to run and easy to understand. While games such as Go and Atari have led to many breakthroughs, they often do not directly translate to real-world embodied applications. In recognising the need to diversify RL benchmarks and addressing complexities that arise in embodied interaction scenarios, we introduce Assistax: an open-source benchmark designed to address challenges arising in assistive robotics tasks. Assistax uses JAX's hardware acceleration for significant speed-ups for learning in physics-based simulations. In terms of open-loop wall-clock time, Assistax runs up to $370\times$ faster when vectorising training runs compared to CPU-based alternatives. Assistax conceptualises the interaction between an assistive robot and an active human patient using multi-agent RL to train a population of diverse partner agents against which an embodied robotic agent's zero-shot coordination capabilities can be tested. Extensive evaluation and hyperparameter tuning for popular continuous control RL and MARL algorithms provide reliable baselines and establish Assistax as a practical benchmark for advancing RL research for assistive robotics. The code is available at: https://github.com/assistive-autonomy/assistax.