JaxAHT：基于 JAX 的 Ad Hoc Teamwork 开源库

Authors:Caroline Wang, Rolando Fernandez, Zelal Su Mustafaoglu, Montek Singh Kundan, Johnny Liu, Aditya Madhan, Zhihan Wang, Lingyun Xiao, Di Yang Shi, Jiaxun Cui, Arrasy Rahman, Peter Stone

View PDF HTML (experimental)

Abstract: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.

Submission history

From: Caroline Wang [view email]
[v1] Sat, 12 Sep 2026 05:01:27 UTC (5,012 KB)
[v2] Sun, 11 Oct 2026 18:20:23 UTC (4,595 KB)