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

2025-07-29

AI Summary

Existing RL benchmarks are predominantly game-based and poorly suited to real-world embodied applications—particularly assistive robotics—and lack systematic evaluation of zero-shot human-robot collaboration. To address this, we introduce Assistax: the first open-source RL benchmark specifically designed for assistive robotics tasks. Implemented in JAX, it enables hardware-accelerated, vectorized physics simulation and multi-agent training, supporting efficient parallel learning for continuous-control and multi-agent RL algorithms. Its key contributions are: (1) clinically inspired patient-robot interaction simulation environments; (2) a suite of diverse partner agents to quantitatively assess zero-shot collaborative generalization; and (3) up to 370× speedup over CPU-based simulation, alongside optimized baselines for major RL algorithms. Extensive experiments validate Assistax’s effectiveness and scalability as a practical benchmark for embodied AI research.

Technology Category

• Intelligent Robots: Embodied AI
• Humans and AI: Human-AI Collaboration / Human-AI Teaming
• Multiagent Systems: Adversarial Agents

Application Category

• Search and Retrieval-Augmented AI: Assisted, interactive, and conversational search
• User Modeling, Personalization and Recommendation: User modeling and simulation for interactive and conversational systems
• Economics, Online Markets and Human Computation: Research challenges in human and human-AI computation

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 imes$ 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.

Problem

• Develop hardware-accelerated RL benchmark for assistive robotics
• Address real-world challenges in human-robot interaction
• Enable fast training via physics-based simulations

Innovation

• Hardware-accelerated RL using JAX for speed
• Multi-agent RL for human-robot interaction
• Vectorised training for faster simulation runs