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
Organizations: Department of Computer Science, The University of Texas at Austin · Google · Sony AI
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.
Figures & tables
Library
JAX
Multi-Env
Ego Algs.
Team Gen.
Eval Teammates
Human Data
JaxAHT
✓
✓
✓
✓
✓
✓
ZSCEval ( Wang et al., 2024b )
✓
✓
✓
✓
AH2AC2 ( Dizdarević et al., 2025 )
✓
✓
ICRL4AHT ( Jing et al., 2026 )
✓
✓
✓
✓
Assistax ( Hinckeldey et al., 2026 )
✓
✓
✓
Table 1: Library comparison. JaxAHT is the only library that covers the full AHT research lifecycle.
Figure 1: JaxAHT Design Philosophy. JaxAHT unifies the full AHT research lifecycle within a shared codebase. During training, common agent and population abstractions support MARL, teammate-generation, ego-agent learning, and unified AHT methods that jointly construct training teammates and ego policies. During evaluation, learned agents are tested against a standardized suite of heuristic, algorithmic, and human-data-informed teammates, together with BR references, enabling consistent and scalable benchmarking across AHT methods.
Category
Algorithm
Ego Agent Training
PPO ( Schulman et al., 2017 )
LIAM ( Papoudakis et al., 2021 )
MeLIBA ( Zintgraf et al., 2021 )
Teammate Generation
FCP ( Strouse et al., 2021 )
BRDiv ( Rahman et al., 2023 )
LBRDiv ( Rahman et al., 2024 )
Table 2: JaxAHT Algorithms.
Environment
Source
Eval Teammates
Variants
Level-Based Foraging
Jumanji
✓
7x7 grid, constant levels
12x12 grid
Overcooked-v1
JaxMARL
✓
Asymmetric Advantages
Coordination Ring
Counter Circuit
Cramped Room
Table 3: JaxAHT Environments. Overview of environments supported by JaxAHT.
Algorithm
Seeds
PyTorch
JAX
LIAM
1
40m 27s
08m 22s
BRDiv
1
44m 18s
01m 46s
LIAM
3
121m 21s*
09m 10s
LIAM
10
404m 30s*
12m 17s
LIAM
300
202h 15m*
126m
Table 4: Wall-clock time comparison on LBF 7x7. JaxAHT achieves ∼5 x (LIAM) and ∼25 x (BRDiv) single-seed speedups over PyTorch, and 95x for LIAM across 300 seeds. Asterisks indicate runtimes for sequentially running PyTorch implementations.
Figure 2: BR return estimates for JaxAHT evaluation teammates. Each bar is the estimated BR return for an evaluation teammate, colored by teammate type. Estimates are the maximum return obtained against that teammate by any trained BR policy in the cross-play matrix, or the known optimal return for certain LBF heuristics. Note that the maximum return on LBF is 0.5, on Hanabi is 25, and on Mini Hanabi is 8. For Overcooked-v1 tasks, agents are awarded 20 points per delivered dish, with no analytic maximum.
Figure 3: Representations of Behavior Trajectories for Evaluation Teammates. Each point reflects a single embedded behavior trajectory of a teammate-BR pair. Teammate-BR pairs are meaningfully distinct.
Figure 4: Hyperparameter sweep return distributions on Cramped Room. Each point represents the performance of a distinct hyperparameter configuration on a validation set of teammates. Best performing configurations are selected for the benchmark study.
Figure 5: AHT Algorithm Benchmark Results. (Left) Ego training methods with agent modeling (LIAM, MeLIBA) outperform PPO only in Overcooked environments, which have role-based conventions, and in the presence of diverse teammate populations (CoMeDi). (Right) No single AHT method dominates on all tasks; the best overall are ROTATE, FCP, and CoMeDi.
Appendix figures & tables11 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 6: LBF cross-play matrices using normalized return. Rows are evaluation teammates and columns are best-response policies.
Figure 7: Hanabi cross-play matrices using normalized return. Rows are evaluation teammates and columns are best-response policies.
Figure 8: Overcooked cross-play matrices using normalized return. Rows are evaluation teammates and columns are best-response policies. Diagonal entries reflect matched best-response performance; off-diagonal structure indicates the degree of transferability across teammate types.
Figure 9: Representations of Behavior Trajectories for RL and Heuristic Teammates on LBF 7x7.
Figure 10: Hyperparameter sweep return distributions for ego methods on Cramped Room.
Figure 11: Hyperparameter sweep return distributions on LBF 12x12.
Figure 12: BR Return Bounds Gap. AHT agents can sometimes improve the empirical BR estimates derived by RL.
Figure 13: Teammate generation and unified methods performance, broken down by RL agent type.
Figure 14: Training statistics for all methods on LBF 12x12.
Env
Agent
Behavior
Hanabi
Cautious
Derived from human gameplay; avoids plays that can lose lives.
IGGI
Extends Cautious with conservative play and predictable discards.
Piers
Risk-taking extension of IGGI, including Hail Mary plays when the deck is empty.
Internal ( Osawa, 2015 )
Tracks hints received about its own hand, without maintaining the full history of hints given to its partner.
Outer ( Osawa, 2015 )
Tracks hints about its own hand and hints given to its partner, using the latter to avoid redundant tells.
Van Den Bergh
Rule-based agent whose rule ordering was found by a genetic algorithm ( van den Bergh et al., 2017 ) .
Appendix
Table 5: Heuristic teammate suite. The hand-designed teammates span convention-based Hanabi policies, LBF food-targeting strategies, and Overcooked-v1 role-specialist policies; the listed behaviors define the coordination demands used in held-out evaluation.
C2DL, Institute of Automation, Chinese Academy of Sciences · School of Artificial Intelligence, University of Chinese Academy of Sciences · University of Science and Technology of China +5