# JaxAHT: A JAX-Based Library for Ad Hoc Teamwork

arXiv: https://arxiv.org/abs/2609.13716 · HTML: https://mlsanity.com/paper/2609.13716/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 35 of its 50 references that are in the map.
- First posted on arXiv 12 Sep 2026.

## Builds on (most cited first, 30 of 35)

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Deep Reinforcement Learning at the Edge of the Statistical Precipice](https://mlsanity.com/paper/2108.13264/) (2021)
- [Conflict-Averse Gradient Descent for Multi-task Learning](https://mlsanity.com/paper/2110.14048/) (2021)
- [On the Utility of Learning about Humans for Human-AI Coordination](https://mlsanity.com/paper/1910.05789/) (2019)
- [Is Independent Learning All You Need in the StarCraft Multi-Agent Challenge?](https://mlsanity.com/paper/2011.09533/) (2020)
- [The Complexity of Decentralized Control of Markov Decision Processes](https://mlsanity.com/paper/1301.3836/) (2013)
- [The Hanabi Challenge: A New Frontier for AI Research](https://mlsanity.com/paper/1902.00506/) (2019)
- [Collaborating with Humans without Human Data](https://mlsanity.com/paper/2110.08176/) (2021)
- [Structured State Space Models for In-Context Reinforcement Learning](https://mlsanity.com/paper/2303.03982/) (2023)
- [Discovered Policy Optimisation](https://mlsanity.com/paper/2210.05639/) (2022)
- [Off-Belief Learning](https://mlsanity.com/paper/2103.04000/) (2021)
- [JaxMARL: Multi-Agent RL Environments and Algorithms in JAX](https://mlsanity.com/paper/2311.10090/) (2023)
- [Jumanji: a Diverse Suite of Scalable Reinforcement Learning Environments in JAX](https://mlsanity.com/paper/2306.09884/) (2023)
- [Agent Modelling under Partial Observability for Deep Reinforcement Learning](https://mlsanity.com/paper/2006.09447/) (2020)
- [Modeling Strong and Human-Like Gameplay with KL-Regularized Search](https://mlsanity.com/paper/2112.07544/) (2021)
- [A Game-Theoretic Model and Best-Response Learning Method for Ad Hoc Coordination in Multiagent Systems](https://mlsanity.com/paper/1506.01170/) (2015)
- [A Survey of Ad Hoc Teamwork Research](https://mlsanity.com/paper/2202.10450/) (2022)
- [Melting Pot 2.0](https://mlsanity.com/paper/2211.13746/) (2022)
- [A Survey of Progress on Cooperative Multi-agent Reinforcement Learning in Open Environment](https://mlsanity.com/paper/2312.01058/) (2023)
- [Cooperative Open-ended Learning Framework for Zero-shot Coordination](https://mlsanity.com/paper/2302.04831/) (2023)
- [Deep Interactive Bayesian Reinforcement Learning via Meta-Learning](https://mlsanity.com/paper/2101.03864/) (2021)
- [ZSC-Eval: An Evaluation Toolkit and Benchmark for Multi-agent Zero-shot Coordination](https://mlsanity.com/paper/2310.05208/) (2023)
- [Generating Teammates for Training Robust Ad Hoc Teamwork Agents via Best-Response Diversity](https://mlsanity.com/paper/2207.14138/) (2022)
- [Human-compatible driving partners through data-regularized self-play reinforcement learning](https://mlsanity.com/paper/2403.19648/) (2024)
- [Diverse Conventions for Human-AI Collaboration](https://mlsanity.com/paper/2310.15414/) (2023)
- [N-Agent Ad Hoc Teamwork](https://mlsanity.com/paper/2404.10740/) (2024)
- [Minimum Coverage Sets for Training Robust Ad Hoc Teamwork Agents](https://mlsanity.com/paper/2308.09595/) (2023)
- [Towards Few-shot Coordination: Revisiting Ad-hoc Teamplay Challenge In the Game of Hanabi](https://mlsanity.com/paper/2308.10284/) (2023)
- [Evaluating and Modelling Hanabi-Playing Agents](https://mlsanity.com/paper/1704.07069/) (2017)

## Built on by (most cited first, 1 of 1)

- [Grounded Joint-Attention Other-Play for Zero-Shot Coordination](https://mlsanity.com/paper/2610.06025/) (2026)

## Abstract (arXiv)

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.