JaxAHT: A JAX-Based Library for Ad Hoc Teamwork

Caroline Wang, Rolando Fernandez, Jiaxun Cui, Johnny Liu, Aditya Madhan, Zhihan Wang, Lingyun Xiao, Di Yang Shi, Arrasy Rahman, Peter Stone

Keywords: JAX, benchmark, ad hoc teamwork, multi-agent systems

TL;DR: Proposes a JAX based benchmark for ad hoc teamwork

Abstract:

Ad Hoc Teamwork (AHT) addresses the challenge of designing agents capable of coordinating with novel partners without prior coordination. However, progress in the field is currently hindered by the lack of a systematic evaluation framework and the prohibitive computational cost of generating diverse populations of training and evaluation partners. In this work, we introduce **JaxAHT**, the first open-source, JAX-based library designed to accelerate and standardize the AHT research lifecycle. Leveraging the hardware acceleration and massive parallelization capabilities of JAX, the library provides a unified pipeline for teammate generation, AHT agent training, and evaluation against unseen teammates. JaxAHT provides native integration with standard AHT research environments. Preliminary experiments demonstrate that our implementations achieve significant wall-clock time speedups compared to PyTorch counterparts while successfully reproducing established performance hierarchies on held-out evaluation teammates. The codebase
is available at https://github.com/LARG/jax-aht.

Submission Number: 54