JaxAHT: A JAX-Based Library for Ad Hoc Teamwork
Not on the map. Why?
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.
Key facts. Cited by 1 paper 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
35 of 50 references.
References · 35
- 2019PyTorch: An Imperative Style, High-Performance Deep Learning Librarycited by 20,305
- 2017Proximal Policy Optimization Algorithmscited by 15,215
- 2021Deep Reinforcement Learning at the Edge of the Statistical Precipicecited by 807
- 2021Conflict-Averse Gradient Descent for Multi-task Learningcited by 459
- 2019On the Utility of Learning about Humans for Human-AI Coordinationcited by 349
- 2020Is Independent Learning All You Need in the StarCraft Multi-Agent Challenge?cited by 294
- 2013The Complexity of Decentralized Control of Markov Decision Processescited by 231
- 2019The Hanabi Challenge: A New Frontier for AI Researchcited by 197
- 2021Collaborating with Humans without Human Datacited by 142
- 2023Structured State Space Models for In-Context Reinforcement Learningcited by 111
- 2022Discovered Policy Optimisationcited by 104
- 2021Off-Belief Learningcited by 75
23 more
- 2023JaxMARL: Multi-Agent RL Environments and Algorithms in JAXcited by 63
- 2023Jumanji: a Diverse Suite of Scalable Reinforcement Learning Environments in JAXcited by 58
- 2020Agent Modelling under Partial Observability for Deep Reinforcement Learningcited by 57
- 2021Modeling Strong and Human-Like Gameplay with KL-Regularized Searchcited by 50
- 2015A Game-Theoretic Model and Best-Response Learning Method for Ad Hoc Coordination in Multiagent Systemscited by 44
- 2022A Survey of Ad Hoc Teamwork Researchcited by 43
- 2022Melting Pot 2.0cited by 40
- 2023A Survey of Progress on Cooperative Multi-agent Reinforcement Learning in Open Environmentcited by 39
- 2023Cooperative Open-ended Learning Framework for Zero-shot Coordinationcited by 30
- 2021Deep Interactive Bayesian Reinforcement Learning via Meta-Learningcited by 29
- 2023ZSC-Eval: An Evaluation Toolkit and Benchmark for Multi-agent Zero-shot Coordinationcited by 24
- 2022Generating Teammates for Training Robust Ad Hoc Teamwork Agents via Best-Response Diversitycited by 17
- 2024Human-compatible driving partners through data-regularized self-play reinforcement learningcited by 17
- 2023Diverse Conventions for Human-AI Collaborationcited by 15
- 2024N-Agent Ad Hoc Teamworkcited by 15
- 2023Minimum Coverage Sets for Training Robust Ad Hoc Teamwork Agentscited by 13
- 2023Towards Few-shot Coordination: Revisiting Ad-hoc Teamplay Challenge In the Game of Hanabicited by 10
- 2017Evaluating and Modelling Hanabi-Playing Agentscited by 9
- 2025SocialJax: An Evaluation Suite for Multi-agent Reinforcement Learning in Sequential Social Dilemmascited by 8
- 2025Ad-Hoc Human-AI Coordination Challengecited by 8
- 2025ROTATE: Regret-driven Open-ended Training for Ad Hoc Teamworkcited by 6
- 2025Predicting Multi-Agent Specialization via Task Parallelizabilitycited by 5
- 2025Partner Modelling Emerges in Recurrent Agents (But Only When It Matters)cited by 4
Built on by
Citing papers per year: 2026: 1.
1 paper.
Similar, not cited
Close in wording, no citation either way
- 2025PADiff: Predictive and Adaptive Diffusion Policies for Ad Hoc Teamwork
- 2022Knowledge-based and Data-driven Reasoning and Learning for Ad Hoc Teamwork
- 2023Cooperation on the Fly: Exploring Language Agents for Ad Hoc Teamwork in the Avalon Game
- 2023Knowledge-based Reasoning and Learning under Partial Observability in Ad Hoc Teamwork
- 2026Benchmarking the Limits of In-Context Reinforcement Learning for Ad-Hoc Teamwork
- 2025ToMCAT: Theory-of-Mind for Cooperative Agents in Teams via Multiagent Diffusion Policies
- 2025ReCollab: Retrieval-Augmented LLMs for Cooperative Ad-hoc Teammate Modeling
- 2023MetaAgents: Large Language Model Based Agents for Decision-Making on Teaming