JaxAHT: A JAX-Based Library for Ad Hoc Teamwork

As autonomous agents are increasingly deployed, they must be able to coordinate with previously unseen agents, such as humans. This challenge has been formalized as Ad Hoc Teamwork (AHT), and has attracted a substantial body of research. However, progress has been hindered by the prohibitive computational cost of the AHT research lifecycle, lack of standardized benchmark implementations, and the absence of a diverse, validated evaluation teammate suite. This work introduces **JaxAHT**, the first open-source, JAX-based library designed to accelerate and standardize the AHT research lifecycle. Leveraging JAX's hardware acceleration and massive parallelization capabilities, JaxAHT provides a unified framework for teammate generation, ego agent training, and evaluation against unseen teammates, achieving approximately 95x wall-clock speedup over existing PyTorch counterparts. Along with the library, we contribute a diverse suite of evaluation teammates across Level-Based Foraging, Overcooked, and Hanabi. Using JaxAHT, we conduct a large-scale, compute-controlled benchmark study of AHT algorithms, finding that no algorithm consistently performs best, that agent modeling primarily offers benefits in role-based scenarios, and that best response returns are underestimated by reinforcement learning from scratch.