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

TL;DR

JaxRobotarium addresses the need for fast, robotics-faithful MRRL benchmarking by delivering a GPU-accelerated, end-to-end platform that bridges simulation, learning, and hardware deployment on the Robotarium. It combines a Jax-based simulator with barrier-certified collision avoidance, an interface to MARL libraries, and a unified eight-scenario benchmark, enabling rapid training and sim-to-real evaluation. The approach demonstrates up to $20\times$ faster training and up to $150\times$ faster simulation compared to baselines, while revealing nuanced sim2real gaps that can be mitigated through domain randomization. Overall, JaxRobotarium accelerates development, benchmarking, and real-world evaluation of multi-robot learning algorithms, with an open-source, globally accessible pipeline for MRRL research.

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 RL (MRRL) policies with realistic robot dynamics and safety constraints, supporting parallelization and hardware acceleration. Our generalizable learning interface integrates easily 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 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. Our code is available at https://github.com/GT-STAR-Lab/JaxRobotarium.

Figures (14)

• Figure 1: Overview of JaxRobotarium architecture. See \ref{['fig:block_diagram']} in Appendix for additional details.
• Figure 2: Jax-RPS simulates the same trajectories significantly faster than RPS. Trajectory errors are computed for Jax-RPS against the trajectory for the identical scenario initial conditions simulated in RPS.
• Figure 3: (Top) We plot the mean returns achieved in both platforms against the minimum training wall time. (Bottom) We plot the time taken to simulate timesteps across both platforms. We find that JaxRobotarium is significantly more efficient in training multi-robot policies than MARBLER. Results depict smoothed mean (solid line) with standard deviation (shaded).
• Figure 4: Scenarios (real on left and sim on right) from left to right and top to bottom as follows: Arctic Transport, Discovery, Foraging, Material Transport, Navigation, Predator Prey, RWARE, Warehouse. For detailed descriptions of each scenario, please see Appendix \ref{['appendix:scenario-details']}.
• Figure 5: JaxRobotarium architecture, colored components are novel.