PaperlayerView exact revision on arXiv

Source-linked AI summary

JaxAHT: A JAX-Based Library for Ad Hoc Teamwork

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

arXiv:2609.13716v1cs.AI

TL;DR

AHT requires agents to coordinate with unfamiliar partners, while research faces computational, standardization, and evaluation-teammate gaps. JaxAHT unifies the research lifecycle in a JAX-based library and supports large-scale benchmarking; results show no algorithm consistently performs best, and agent modeling primarily benefits role-based tasks.

  • Problem

    AHT requires coordination with unknown, unprecoordinated partners, while computational barriers and limited standardization and teammate diversity constrain its research lifecycle.

    Source: §1 · pp. 1–2

  • Method

    JaxAHT combines teammate generation, ego-agent training, and evaluation against held-out teammates in a JAX-based framework with standardized algorithms and diverse benchmark teammates.

    Source: §1 · p. 2, §3 · p. 3

  • Results

    No existing algorithm performs best consistently; agent-modeling methods benefit primarily in role-based tasks and sometimes harm learning versus a non-agent-modeling baseline.

    Source: §6 · p. 10

  • Takeaways & Limitations

    The unified pipeline, diverse teammates, and reduced compute barriers make rigorous large-scale AHT studies tractable at academic scale.

    Source: §1 · p. 2

    The work targets existing AHT benchmarks rather than new environments, and excludes full Hanabi from benchmarking because of prohibitive compute requirements.

    Source: §6 · p. 10

Abstract

from arXiv · show

Ad Hoc Teamwork (AHT) addresses the challenge of designing agents capable of coordinating with novel partners without prior coordination. However, progress in the field is hindered by the prohibitive computational cost of the AHT research lifecycle, the lack of standardized benchmark implementations, and the absence of a diverse, validated evaluation teammate suite. In this work, we introduce 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 PyTorch counterparts. Alongside the library, we contribute a diverse suite of evaluation teammates across the domains of Level-Based Foraging, Overcooked, and Hanabi. To illustrate the value of the framework, we use it to conduct a large-scale, compute-controlled benchmark study comparing teammate generation and AHT agent learning methods, finding that no algorithm consistently performs best, and that agent modeling primarily offers benefits in role-based scenarios with diverse teammates.

1 INTRODUCTION

AHT research is constrained by computational cost, inconsistent evaluation teammates, and the complexity of assembling benchmark implementations. JaxAHT unifies and accelerates the research lifecycle while providing evaluation teammates and enabling large-scale benchmarking.

1 INTRODUCTION

  • AHT research faces costly iteration, limited teammate diversity validation, incomparable results, and time-intensive orchestration across algorithm implementations.These bottlenecks raise the barrier to entry and limit investigation of methods addressing different components of AHT.Source: §1 · p. 1
  • JaxAHT provides a unified framework for teammate generation, ego-agent training, and cooperative MARL, supporting the full AHT research lifecycle.It is the first open-source, JAX-based library for that lifecycle.Source: §1 · pp. 1–2
  • JaxAHT provides a diverse suite of heuristic, RL-based, and human-gameplay-informed evaluation teammates for LBF, Overcooked, and Hanabi.The work also introduces the first open-source human gameplay dataset for LBF.Source: §1 · p. 2
  • JaxAHT’s algorithm implementations achieve approximately 95x wall-clock speedup over PyTorch.The paper also uses the framework for a large-scale benchmark comparing AHT algorithms.Source: §1 · p. 2
  • JaxAHT’s unified framework and evaluation suite support rigorous, large-scale AHT studies that were previously intractable.Code, pregenerated teammates, and human gameplay data are open-sourced.Source: §1 · p. 2

2 RELATED WORK

Related work spans JAX-based reinforcement-learning libraries and AHT benchmarks, but few libraries jointly support multiple AHT algorithms and environments. JaxAHT is positioned as an end-to-end JAX framework for the full lifecycle across environments and algorithm classes.

JAX-based RL Libraries

  • JAX libraries use compilation and auto-vectorization for accelerated RL; PureJaxRL reported over 1000× speedups in highly parallel settings, while JaxMARL brought JAX to MARL.Jumanji supplies JAX environments, and NiceWebRL extends them to web-based human-subject studies.Source: §2 · p. 2

Benchmarks and Evaluation for AHT

  • Overcooked and Hanabi are widely used AHT environments, yet few libraries jointly support multiple AHT algorithms and environments.The AHT lifecycle includes developing training and evaluation teammate sets, training agents, and testing them against unseen teammates.Source: §2 · p. 2, Figure 1 · p. 2
  • ZSC-Eval focuses on teammate generation and partner-selection and generalization measures, but its PyTorch implementation remains computationally expensive.AH2AC2 focuses on Hanabi and offers human proxies trained on about 100k games behind an API; SocialJax targets social dilemmas rather than AHT.Source: §2 · pp. 2–3
  • JaxAHT is presented as the first end-to-end JAX framework supporting the complete AHT lifecycle across multiple environments and algorithm classes.It complements related work by jointly supporting teammate generation, ego-agent training, and evaluation.Source: §2 · p. 3

3 PROBLEM SETTING AND DESIGN PHILOSOPHY

The paper formalizes AHT as cooperation with teammates whose behavior is unknown and evaluates agents against held-out partners. JaxAHT maps the resulting research stages and algorithm types onto shared abstractions for training and evaluation.

3 PROBLEM SETTING AND DESIGN PHILOSOPHY

  • The cooperative setting is modeled as a Dec-POMDP with agent-specific action and observation spaces, shared reward, and state-transition and observation functions.Each agent acts from private observations, yielding decentralized policies.Source: §3 · p. 3
  • AHT asks an ego agent to coordinate with partners it has not previously known or had a chance to pre-coordinate with, then evaluates generalization on held-out teammates.Training partners come from Ttrain, while evaluation partners are drawn from Teval.Source: §3 · p. 3
  • The three stages are constructing training teammates, learning the ego policy with them, and evaluating it on a standardized held-out teammate set.JaxAHT reports returns relative to each teammate’s best response because raw returns are not comparable across teammates of differing competence.Source: §3 · p. 3
  • Teammate-generation, AHT-agent-training, and cooperative-MARL algorithms can support any lifecycle stage, so JaxAHT offers a unified interface for individual components or combinations.Research on one algorithm type may require others, such as an AHT-agent-training method to evaluate teammate generation.Source: §3 · p. 3
  • Algorithms use single-file implementations for understandability and extension, while a shared agent and population interface lets agents trained by one method be used by another.The library also supports Hydra configuration management and WandB logging.Source: §3 · pp. 3–4
  • JaxAHT’s common agent and population abstractions support MARL, teammate generation, ego-agent learning, and unified methods that jointly construct training teammates and ego policies.The unified methods use library building blocks and MARL procedures to train best-response agents.Source: Figure 2 · p. 4, §3 · p. 4

4 JAXAHT LIBRARY

JaxAHT packages algorithms, supported environments, and standardized held-out teammate evaluation into a JAX-accelerated library for the AHT research lifecycle.

Algorithms

  • JaxAHT provides thorough hyperparameter sweeps for benchmarked tasks to supply performant baseline configurations for future research.The algorithm summary is presented in Table 2.Source: §4 · p. 4, Table 2 · p. 4
  • JaxAHT implements teammate-generation, ego-training, MARL, and unified-training algorithms, including methods that jointly train ego agents and teammate populations.Supported methods include FCP, BRDiv, LBRDiv, CoMeDi, PPO, LIAM, MeLIBA, IPPO, TrajeDi, COLE, and ROTATE. Their approaches include diversity objectives, latent interaction-history representations, cooperative-compatibility sampling, and regret-based teammate generation.Source: §4 · pp. 4–5, Table 2 · p. 4

Environments

  • The library currently supports LBF, Overcooked, and Hanabi, selected for their prevalence in AHT research and ability to support diverse cooperative conventions.The supported environments build on JAX-based re-implementations of commonly used AHT environments.Source: §4 · p. 5, Table 3 · p. 5

Evaluation Teammates

  • Held-out evaluation teammates span hand-designed heuristics, algorithmically generated policies, and policies learned from human data across supported environments.Empirical best-response references support per-teammate upper bounds and BRProx-normalized returns; aggregate metrics include 95% bootstrapped confidence intervals. Heuristics impose distinct coordination demands, while human-data policies are derived differently across Overcooked and LBF; Hanabi human proxies are accessed through an evaluation API.Source: §4 · p. 5

5 EXPERIMENTAL STUDIES

The studies report JaxAHT’s runtime gains, evaluate the quality and diversity of its teammates, and benchmark AHT methods under controlled compute budgets. Results vary by task: no method dominates, agent modeling helps in role-based Overcooked settings, and benchmark agents can exceed PPO-derived best-response estimates.

5.1 WALL CLOCK TIME COMPARISON

  • 95x speedup is achieved by JaxAHT’s LIAM implementation across 300 seeds, while single-seed speedups are ∼5x for LIAM and ∼25x for BRDiv.The comparison used one Nvidia A6000 GPU, 8 environments, and 3 million timesteps.Source: Table 4 · p. 5, §5.1 · p. 6

5.2 EVALUATION TEAMMATE QUALITY AND DIVERSITY

  • Evaluation teammates generally have stronger diagonal than off-diagonal cross-play returns, indicating that higher returns require distinct response strategies.Some high-return off-diagonal blocks suggest partial behavioral overlap or transfer, especially within agent types.Source: §5.2 · p. 6
  • In Overcooked Coordination Ring, teammate–best-response pairs occupy largely non-overlapping latent regions, supporting behavioral variety in the evaluation set.The analysis uses a classifier on interaction data and t-SNE embeddings.Source: §5.2 · p. 7, Figure 4 · p. 7

5.3 BENCHMARKING AHT ALGORITHM PERFORMANCE

  • The benchmark selects hyperparameters on validation teammates disjoint from both training and evaluation teammates, with training timesteps controlled for fair comparisons.The study covers five representative tasks across LBF, Overcooked, and Mini Hanabi.Source: §5.3 · p. 7
  • Different hyperparameter configurations produce a wide range of performances, motivating compute-controlled tuning for algorithm comparisons.JaxAHT makes such tuning feasible under its JAX-based framework and releases the optimized hyperparameters.Source: §5.3 · pp. 7–8
  • Agent modeling outperforms PPO in role-based Overcooked settings, while LIAM and MeLIBA offer little benefit on LBF and perform worse on Mini Hanabi.The authors attribute the split to Overcooked’s complementary roles and hypothesize that latent teammate representations help model role-dependent best responses.Source: §5.3 · pp. 8–9, Figure 6 · p. 9
  • ROTATE, FCP, and CoMeDi are strongest overall, but no AHT method performs best across all benchmark scenarios.Each performs best on at least one task.Source: §5.3 · p. 9, Figure 6 · p. 9

5.4 CHALLENGES OF BEST RESPONSE RETURN ESTIMATION

  • Benchmark AHT agents sometimes outperform PPO-derived best-response estimates, which the authors attribute to training against diverse populations rather than one teammate.The estimate takes the maximum return across PPO best-response policies.Source: §5.4 · p. 9

6 DISCUSSION AND CONCLUSION

JaxAHT combines a JAX-based full-lifecycle AHT library and validated evaluation teammates with a controlled benchmark, which finds no consistently best algorithm and role-dependent benefits from agent modeling. The work also identifies benchmark-scope limitations and directions for future research.

6 DISCUSSION AND CONCLUSION

  • No single existing algorithm performs best; agent-modeling-based methods benefit primarily role-based tasks and sometimes harm learning versus a non-agent-modeling baseline.The benchmark also finds that standard single-agent PPO training may systematically underestimate best-response returns.Source: §6 · p. 10
  • The unified framework, released hyperparameters, and evaluation teammates support joint study of teammate generation and ego-agent learning and fairer future comparisons.Population-pretrained agents are suggested as a promising, underexplored direction for more capable best-response estimation.Source: §6 · p. 10
  • The study covers existing AHT benchmark tasks rather than developing new tasks, and excludes full Hanabi because of its prohibitive compute requirements.Richer state spaces, continuous control, and more complex role structure are identified as complementary task-development directions; sample-efficient Hanabi methods remain an open challenge.Source: §6 · p. 10

A CODE AND DATA

The appendices provide links to the library and released datasets, explain cross-play-based teammate evaluation and empirical upper-bound refinement, and detail behavioral-diversity visualizations.

A CODE AND DATA

  • The JaxAHT library is open-sourced on GitHub, and evaluation-teammate, best-response, and LBF human-gameplay datasets are released on Hugging Face.The listed assets include evaluation teammate policies, best-response policies, and LBF human gameplay data.Source: A CODE AND DATA · p. 15

B BEST-RESPONSE POLICIES AND UPPER BOUND ESTIMATION

  • Cross-play matrices pair evaluation teammates with best-response policies, and empirical normalization bounds are updated when a cross-play score exceeds the recorded value.The bounds depend on the particular evaluated best-response policies; Figures 7 and 8 show LBF and Hanabi matrices with teammates as rows and best-response policies as columns.Source: B BEST-RESPONSE POLICIES AND UPPER… · p. 15, Figure 7 · p. 15, Figure 8 · p. 16

C T-SNE BEHAVIORAL DIVERSITY DETAILS

  • Behavioral-diversity embeddings are learned from joint interaction data to predict teammate–best-response identity, then projected into two dimensions with t-SNE.The classifier is an LSTM encoder followed by a linear layer; training and testing use fixed per-pair trajectory counts, with different counts for LBF 7x7 and Coordination Ring.Source: C T-SNE BEHAVIORAL DIVERSITY DETAILS · p. 16, Figure 10 · p. 16
  • Figure 9’s Overcooked cross-play matrix uses normalized returns: diagonal entries show matched best-response performance, while off-diagonal structure indicates transferability across teammate types.Some off-diagonal blocks retain high returns, suggesting partial behavioral overlap or transferability, especially within agent types such as PPO.Source: Figure 9 · p. 17

D BENCHMARK STUDY DETAILS

This appendix classifies benchmark tasks by convergence difficulty and details the hyperparameter sweep, including its search procedures, validation selection, and return-distribution plots.

D BENCHMARK STUDY DETAILS

  • LBF tasks, Cramped Room, and Mini Hanabi are categorized as easy, while Coordination Ring is categorized as hard based on the general timesteps required for MARL convergence.Source: D BENCHMARK STUDY DETAILS · p. 18

D.1 HYPERPARAMETER SWEEP

  • For each algorithm and task, 140 configurations were swept, except FCP’s 90-configuration grid search and CoR’s Bayesian sweep; the highest-mean-return configuration on an independent validation set was selected.The network architecture was fixed; algorithm-specific hyperparameters and base learning rate, entropy coefficient, and clip parameter varied, while other base RL hyperparameters were held constant.Source: §D.1 · p. 18
  • Ego methods received 7M timesteps on easy tasks and 15M on hard tasks, while teammate-generation and unified methods used separate task-dependent budgets.Ego agents trained on generated teammates for 30M steps on easy tasks and 60M on hard tasks; unified methods received 130M and 260M timesteps, respectively.Source: §D.1 · p. 18
  • Figure 11 plots final normalized-return distributions across algorithm–task combinations on Cramped Room.Source: Figure 11 · p. 18, §D.1 · p. 18

D.2 BENCHMARK STUDY

The benchmark appendix details compute-controlled comparisons, supplemental results, and the methods included in the study. Its supplemental figures include improved best-response return bounds attained by AHT agents, while the algorithm overview spans standard ego baselines, teammate-generation methods, and unified approaches.

D.2 BENCHMARK STUDY

  • The benchmark uses larger timestep budgets than the hyperparameter sweep, with matched PPO architectures and learning configurations for teammate-generation comparisons.Teammate generation receives 165M/330M timesteps on easy/hard tasks, followed by ego training; unified methods receive 195M/390M total timesteps, respectively.Source: §D.2 · p. 19

D.2.1 SUPPLEMENTAL FIGURES

  • AHT agents improve best-response return bounds, and may improve empirical best-response estimates derived by RL.Supplemental figures also show return breakdowns by RL agent type and training statistics on LBF 12x12.Source: §D.2.1 · p. 20, Figure 13 · p. 20, Figure 14 · p. 20

D.2 BENCHMARK STUDY

  • The study includes standard ego baselines, two-stage teammate generation followed by ego training, and unified methods that train teammates and ego agents simultaneously.Figure 12 shows hyperparameter-sweep return distributions on LBF 12x12; the appendix points readers to the open-source benchmark branch for configurations and reproduction.Source: §D.2 · p. 18, Figure 12 · p. 19, E ALGORITHMS OVERVIEW · p. 22

E ALGORITHMS OVERVIEW

  • PPO ego trains directly against the available teammate distribution, while MeLIBA adapts from interaction history and LIAM infers latent partner information under partial observability.Source: E ALGORITHMS OVERVIEW · p. 22
  • FCP trains an independent-policy population before ego training, whereas BRDiv creates teammate–best-response pairs and LBRDiv seeks a compact teammate set with sufficient coverage.Source: E ALGORITHMS OVERVIEW · p. 22
  • CoMeDi sequentially generates conventions using self-play, compatibility, and mixed-play objectives, while ROTATE alternates teammate generation with ego training against accumulated teammates.ROTATE generates teammates in response to the current ego agent’s cooperative-regret weaknesses.Source: E ALGORITHMS OVERVIEW · pp. 22–23
  • TrajeDi regularizes a policy population for diverse trajectory distributions, while COLE adds strategies to improve cooperative compatibility across the population.Source: E ALGORITHMS OVERVIEW · p. 23

F EVALUATION TEAMMATE DESCRIPTIONS

JaxAHT provides heuristic evaluation teammates across Hanabi, Level-Based Foraging, and Overcooked-v1, with distinct rule-based, strategic, and role-specialist policies.

F EVALUATION TEAMMATE DESCRIPTIONS

  • Hanabi includes eight rule-based agents, while Level-Based Foraging spans 12 strategies and Overcooked-v1 includes three role-based agents.The Hanabi suite draws seven agents from Walton-Rivers et al. (2017) and includes SmartBot; Overcooked-v1 includes role specialists such as Onion and Plate.Source: F EVALUATION TEAMMATE DESCRIPTIONS · p. 23
  • Table 5 summarizes the human-designed heuristic teammates used across the benchmark environments.Its caption describes convention-based Hanabi policies, LBF food-targeting strategies, and Overcooked-v1 role specialists as defining held-out coordination demands.Source: F EVALUATION TEAMMATE DESCRIPTIONS · p. 23

G LBF HUMAN DATA COLLECTION

The appendix describes an IRB-approved LBF human-gameplay study, its game procedure and variants, participant recruitment and filtering, and software and data sources used in the work.

G LBF HUMAN DATA COLLECTION

  • The LBF human-gameplay data collection was approved by the University of Texas at Austin IRB, and participation was voluntary.Participants received contact information for study questions in accordance with IRB requirements.Source: G LBF HUMAN DATA COLLECTION · p. 23

G.1 PROCEDURE

  • Participants played ten two-player LBF games with a uniformly sampled automated bot, with a 50-step maximum per game.The first two games were randomly selected warm-ups and excluded from the dataset; the remaining eight games supplied data.Source: §G.1 · pp. 23–24

G.2 VARIANTS

  • The four game variants crossed 7x7 or 12x12 grids with whether fruits had different levels, using applicable heuristic agents and one IPPO teammate.Table 5 describes human-designed teammates spanning convention-based Hanabi, food-targeting LBF, and role-specialist Overcooked policies.Source: §G.2 · p. 24, Table 5 · p. 24

G.3 PROLIFIC

  • 180 Prolific participants provided 1,440 episodes; filtering for unsuccessful or idle actions and repeated zero scores left 1,299 episodes.Participants were paid between 8 and 9 USD per hour, and the study interface is shown in Figure 16.Source: §G.3 · p. 25, Figure 16 · p. 25

H ATTRIBUTIONS

  • The implementation builds IPPO and uses Overcooked and Hanabi environments from JaxMARL, LBF from Jumanji, and released human data and OBL weights.The cited sources include AH2AC2 human gameplay, Overcooked human data, and pretrained Off-Belief Learning weights.Source: H ATTRIBUTIONS · p. 25

I COMPUTE INFRASTRUCTURE

Experiments ran on several server configurations, pairing different CPU systems with NVIDIA accelerators.

I COMPUTE INFRASTRUCTURE

Loading 2609.13716v1…