Jaxaht speeds ad hoc teamwork agent research by 95 times
JaxAHT: A JAX-Based Library for Ad Hoc Teamwork
Artificial Intelligence
Summary
Ad hoc teamwork means creating agents that can work well with new partners without planning together beforehand. Progress in this area has been slow because testing and training these agents takes a lot of computing power, and there are few standard tools or teammates to test with. The authors introduce JaxAHT, a new open-source software tool that uses fast computing techniques to speed up this research by about 95 times compared to older tools. They also provide a variety of test partner agents from different games to help researchers evaluate teamwork. Their studies show that no single method is best at all times, and understanding teammates helps mainly when roles differ widely.
What this means in practice
- •For multi-agent system developers: Train and test agents that collaborate with unknown partners more efficiently using JaxAHT’s fast hardware-accelerated framework.
- •For game ai programmers: Use JaxAHT’s diverse set of evaluation teammates to benchmark teamwork strategies in cooperative games like Overcooked and Hanabi.
Authors
Caroline Wang, Rolando Fernandez, Zelal Su Mustafaoglu, Montek Kundan, Jiaxun Cui, Lingyun Xiao, Zhihan Wang, Di Yang Shi, Aditya Madhan, Johnny Liu, Arrasy Rahman, Peter Stone
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 hindered by the prohibitive computational cost of the AHT research lifecycle, the lack of standardized benchmark implementations, and the absence of a diverse, validated evaluation teammate suite. In this work, we introduce 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 PyTorch counterparts. Alongside the library, we contribute a diverse suite of evaluation teammates across the domains of Level-Based Foraging, Overcooked, and Hanabi. To illustrate the value of the framework, we use it to conduct a large-scale, compute-controlled benchmark study comparing teammate generation and AHT agent learning methods, finding that no algorithm consistently performs best, and that agent modeling primarily offers benefits in role-based scenarios with diverse teammates.