JaxAHT: A JAX-Based Library for Ad Hoc Teamwork

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.