JaxRobotarium: Training and Deploying Multi-Robot Policies in 10 Minutes

2025-05-10

AI Summary

Existing MARL platforms suffer from limited robot relevance and inadequate hardware deployment support, while lacking parallelization and hardware acceleration—hindering efficient multi-robot policy development. This paper introduces the first end-to-end, JAX-based multi-robot reinforcement learning platform tailored for the Robotarium testbed. It integrates high-fidelity dynamics simulation, GPU/TPU acceleration, vectorized parallel environment training, and seamless real-robot deployment. We propose a unified JAX-accelerated framework, design eight standardized coordination tasks—including four novel robot-centric MARL benchmarks—and implement an out-of-the-box sim-to-real evaluation pipeline. Experiments demonstrate a 20× speedup in policy training and 150× acceleration in simulation; learned policies have been successfully deployed on the physical Robotarium robot swarm. The open-source platform significantly lowers the barrier to entry for multi-robot MARL research.

Technology Category

• Intelligent Robots: Multi-Robot Systems
• Multiagent Systems: Multiagent Planning
• Machine Learning: Imitation Learning & Inverse Reinforcement Learning

Application Category

• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications
• Responsible Web: Machine-in-the-loop, human agency and autonomy
• Search and Retrieval-Augmented AI: Multilingual and cross-lingual Web search

Abstract

Multi-agent reinforcement learning (MARL) has emerged as a promising solution for learning complex and scalable coordination behaviors in multi-robot systems. However, established MARL platforms (e.g., SMAC and MPE) lack robotics relevance and hardware deployment, leaving multi-robot learning researchers to develop bespoke environments and hardware testbeds dedicated to the development and evaluation of their individual contributions. The Multi-Agent RL Benchmark and Learning Environment for the Robotarium (MARBLER) is an exciting recent step in providing a standardized robotics-relevant platform for MARL, by bridging the Robotarium testbed with existing MARL software infrastructure. However, MARBLER lacks support for parallelization and GPU/TPU execution, making the platform prohibitively slow compared to modern MARL environments and hindering adoption. We contribute JaxRobotarium, a Jax-powered end-to-end simulation, learning, deployment, and benchmarking platform for the Robotarium. JaxRobotarium enables rapid training and deployment of multi-robot reinforcement learning (MRRL) policies with realistic robot dynamics and safety constraints, supporting both parallelization and hardware acceleration. Our generalizable learning interface provides an easy-to-use integration with SOTA MARL libraries (e.g., JaxMARL). In addition, JaxRobotarium includes eight standardized coordination scenarios, including four novel scenarios that bring established MARL benchmark tasks (e.g., RWARE and Level-Based Foraging) to a realistic robotics setting. We demonstrate that JaxRobotarium retains high simulation fidelity while achieving dramatic speedups over baseline (20x in training and 150x in simulation), and provides an open-access sim-to-real evaluation pipeline through the Robotarium testbed, accelerating and democratizing access to multi-robot learning research and evaluation.

Problem

• Lack of robotics relevance in current MARL platforms
• Slow performance due to no parallelization and GPU/TPU support
• Need for standardized multi-robot learning and deployment platform

Innovation

• Jax-powered end-to-end simulation and deployment platform
• Supports parallelization and GPU/TPU hardware acceleration
• Integrates with SOTA MARL libraries for rapid training