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, et al.
cs.AI
Quick verdict
First JAX library for the full AHT lifecycle; compute-controlled benchmark is convincing, but the headline 95x speedup rests on extrapolation.
Summary
The paper introduces JaxAHT, an open-source JAX library covering the full ad hoc teamwork (AHT) research lifecycle: teammate generation (FCP, BRDiv, LBRDiv, CoMeDi), ego-agent training (PPO, LIAM, MeLIBA), cooperative MARL (IPPO), and unified methods (TrajeDi, COLE, ROTATE), evaluated on Level-Based Foraging, Overcooked-v1, and Hanabi. On one A6000 GPU, single-seed JAX training runs 5x (LIAM) to 25x (BRDiv) faster than PyTorch counterparts, and 300 parallel LIAM seeds complete in 126 minutes — about 95x faster than an extrapolated PyTorch cost. The authors contribute a standardized evaluation teammate suite (heuristic, RL-based, and human-data-derived policies) with best-response-normalized returns, cross-play diversity matrices, and t-SNE behavioral embeddings, plus the first open-source LBF human gameplay dataset (1,299 filtered episodes from 180 participants). A compute-controlled benchmark — 140 hyperparameter configurations per algorithm-task, validation teammates disjoint from evaluation teammates, five seeds — finds no consistently best algorithm; ROTATE, FCP, and CoMeDi are strongest overall, COLE weakest, and agent modeling (LIAM, MeLIBA) helps only in role-based Overcooked while hurting on Mini Hanabi. The paper also shows PPO-based best-response estimates can be systematically underestimated. Limitations: the 95x figure extrapolates an unmeasured PyTorch baseline, benchmark rankings rest on five seeds without significance tests, full Hanabi and several Overcooked layouts are excluded, and Hanabi's gold human proxies remain API-gated.
Review complete
Assessment at a glance
5
risk flags identified
● high 0
● medium 3
● low 2
6 strengths · 6 weaknesses · 5 author questions
AI text signal = Binoculars zero-shot detector score over 20 sampled passages (Fast-DetectGPT co-reported) — absence does not prove human authorship, and a flag does not imply misconduct.
Levels reflect verification complexity — not a quality score. Evolve never rates your paper.
Key claims under review
- JaxAHT is the first open-source JAX library supporting the complete AHT lifecycle (teammate generation, ego-agent training, cooperative MARL, evaluation), with measured single-seed speedups of 5x (LIAM) and 25x (BRDiv) and an extrapolated ~95x at 300 parallel LIAM seeds.
- The released evaluation teammate suite spans heuristic, RL-based, and human-data-derived policies across LBF, Overcooked-v1, and Hanabi, validated via BR returns, cross-play matrices, and t-SNE behavioral embeddings, and includes the first open-source LBF human gameplay dataset (1,299 filtered episodes).
- A compute-controlled benchmark with 140 hyperparameter configurations per algorithm-task and disjoint validation teammates finds no consistently best algorithm; ROTATE, FCP, and CoMeDi are strongest overall and COLE weakest.
- Explicit agent modeling (LIAM, MeLIBA) helps only in role-based Overcooked and markedly hurts on Mini Hanabi relative to plain PPO.
- PPO-derived best-response return estimates can systematically underestimate achievable returns, implying that BR-dependent AHT algorithms need explicit sensitivity analysis to BR estimation error.
Strengths
Benchmark methodology is itself a contribution
Sec. 5.3 + App. D: 140 random-search hyperparameter configs per algorithm-task (grid of 90 for FCP, Bayesian for CoR), validation teammates disjoint from both training and evaluation teammates, strictly controlled budgets (195M/390M for generation+unified vs 30M/60M ego), identical S5 architecture across ego methods, five seeds. This directly addresses the hyperparameter-fairness gap the paper identifies in prior AHT comparisons.
Genuinely unified lifecycle abstractions
Sec. 3-4 and Table 2: shared agent/population interfaces let agents trained by one algorithm feed another (e.g., Fig. 6a crosses ego methods LIAM/MeLIBA/PPO with populations from FCP and CoMeDi), an orchestration the paper argues was previously infeasible; 11 algorithms in single-file JAX modules over Jumanji/JaxMARL environments (Table 3).
Two-pronged teammate validation, with ground-truth anchors
Candid, useful negative findings
Exemplary artifact release and human-study hygiene
Wall-clock results are strong even before extrapolation
Weaknesses
Headline 95x speedup is extrapolated, not measured
Rankings rest on five seeds with no significance testing
Normalization bounds are known-biased, no sensitivity analysis
Ego-method conclusions depend on only two training populations
Benchmark covers a subset of the offered suite
Author overlap with compared methods
Full Review Report
$1.99 · one-timeUnlock the complete preliminary analysis Evolve produced for this paper.
Report preview
Scope & Claims
JaxAHT (arXiv:2609.13716, cs.AI) is a library-and-benchmark paper from UT Austin's LARG group. Sec. 1 claims four contributions: (1) the first framework supporting the full AHT lifecycle — teammate generation, ego-agent training, cooperative MARL — in one codebase; (2) the first JAX-based AHT library, with ~95x wall-clock speedup over PyTorch; (3) a diverse, standardized evaluation teammate suite for LBF, Overcooked-v1, and Hanabi, plus the first open-source LBF human gameplay dataset; (4) a compute-controlled benchmark yielding three field-level findings: no algorithm consistently wins, agent modeling helps mainly in role-based scenarios, and PPO-based BR return estimation can be systematically underestimated. Claims are mostly supported within scope; the 95x figure is the weakest link.
The remaining sections cover Methodology Assessment, Statistical Evidence, Results & Reproducibility, Limitations & Risks, and Verdict & Recommendations.
- All 6 strengths & 6 weaknesses, analyzed in detail
- 5 risk flags with severity, section & reviewer questions
- Methods & statistical evidence assessment
- 5 questions a peer reviewer would ask
- Full six-dimension evaluation report & recommendations
One-time purchase · this paper only
One-time purchase, this paper only. AI-generated preliminary analysis — not peer review.
Frequently asked questions
What is JaxAHT?
JaxAHT is an open-source JAX library for ad hoc teamwork (AHT) research — building agents that coordinate with previously unseen partners. It unifies teammate generation (FCP, BRDiv, LBRDiv, CoMeDi), ego-agent training (PPO, LIAM, MeLIBA), cooperative MARL (IPPO), and unified methods (TrajeDi, COLE, ROTATE) over the LBF, Overcooked-v1, and Hanabi environments, and ships a standardized evaluation teammate suite. Code is on GitHub (LARG/jax-aht); teammate policies, best responses, and an LBF human gameplay dataset are on Hugging Face.
How much faster is JaxAHT than PyTorch?
On a single Nvidia A6000 (8 parallel environments, 3M timesteps, LBF 7x7), measured single-seed training was about 5x faster for LIAM (8m22s vs 40m27s) and about 25x faster for BRDiv (1m46s vs 44m18s). The headline ~95x figure refers to training 300 LIAM seeds in parallel in 126 minutes and is extrapolated from the single-seed PyTorch time rather than directly measured.
Which ad hoc teamwork algorithm performed best in the JaxAHT benchmark?
None dominated across all tasks. Under controlled compute budgets and a 140-configuration hyperparameter sweep per algorithm-task, ROTATE, FCP, and CoMeDi were strongest overall (each best on at least one task), COLE was frequently weakest, and agent-modeling methods (LIAM, MeLIBA) helped only in role-based Overcooked layouts while performing markedly worse than plain PPO on Mini Hanabi.
What is ad hoc teamwork (AHT)?
Ad hoc teamwork is the problem of designing an agent that can cooperate effectively with teammates whose behavior is unknown in advance and with whom it could not pre-coordinate — formalized as training one 'ego' agent against a set of training teammates and evaluating its expected team return against held-out, unseen teammates.
About this review
This is an AI-generated preliminary review by Evolve.Pub, produced from the openly available text of the paper. It is not peer review and does not represent the authors' views. The canonical record of the paper is at the publisher link above.
- Figures (cross-play matrices, t-SNE plots, Fig. 6 bars) appear in the extraction only as flattened axis-label and value text, so figure-level values could not be visually inspected; every number cited in this review comes from tables or body text (Tables 1-5, Figs. 5/11/12 medians, Sec. 5 timings).
- The PDF header lists 12 authors (including Aditya Madhan, Johnny Liu, Arrasy Rahman, Peter Stone) while the supplied metadata lists 8; this review follows the full text.
Go deeper on this paper
Wrote a paper? Stress-test it before reviewers do.
Upload your manuscript for the same structured scrutiny — reasoning, statistics, evidence, figures, and reproducibility.