JaxAHT: A Unified Jax-based Benchmark for Ad Hoc Teamwork

A unified toolkit for fast iteration on AHT algorithms and evaluation.

AHT benchmark workflow diagram showing teammate generation, RL training, evaluation, and analysis — The Jax AHT benchmark aims to provide a Jax-accelerated and unified framework for the advancement of ad hoc teamwork research. The figure demonstrates three key components of ad hoc teamwork research, all of which are included in the benchmark suite: 1) Teammate generation algorithms, where the goal is to generate diverse policies to be used as partners for ego agent training. 2) Ego agent training algorithms, where the goal is to create a generalizable ego agent that performs well with whichever teammate it is matched up with. 3) Evaluation, where the ego agent is partnered with evaluation teammate policies that are unavailable during the ego agent training phase.

The development and evaluation of ad hoc teamwork (AHT) agents are often hindered by the lack of standardized implementations and reproducible evaluation pipelines, making it difficult to compare methods or build upon prior work. Experimentation is also often slow due to the lengthy reinforcement learning training periods required for many parts of the AHT training workflow, such as teammate generation or ego agent training.

The JaxAHT library addresses these challenges by introducing:

A unified benchmark with transparent evaluation protocols and strong baselines for fair, replicable research.
A Jax-based implementation of various AHT methods, which is significantly faster than previous PyTorch/Tensorflow implementations.
A shared agent interface that allows agents trained by one algorithm to be seamlessly reused in other parts of the AHT workflow (i.e., teammate generation to ego agent training).
A single-file implementation style inspired by JaxMARL and CleanRL, which makes it easier for users to understand and build upon the implemented AHT methods.

These design choices lower the barrier to entry, enable faster iterations through the AHT agent development process (i.e., code implementation, training, and evaluation), and provide a reliable foundation for advancing AHT research.

Citation

If you use this benchmark in your research, please cite:

@misc{jaxaht2025,
  author = {Learning Agents Research Group},
  title = {JaxAHT},
  year = {2025},
  month = {September},
  note = {Version 1.0.0},
  url = {https://github.com/LARG/jax-aht},
}