JaxAHT: 一个基于JAX的临时团队协作库

# JaxAHT：一个基于JAX的临时团队协作研究库
来源：https://arxiv.org/html/2609.13716
Caroline Wang††谢谢：通信作者：caroline\.l\.wang@utexas\.edu
Rolando Fernandez隶属：德克萨斯大学奥斯汀分校计算机科学系
Zelal Su Mustafaoglu隶属：德克萨斯大学奥斯汀分校计算机科学系
Montek Kundan隶属：德克萨斯大学奥斯汀分校计算机科学系
Jiaxun Cui隶属：德克萨斯大学奥斯汀分校计算机科学系隶属：Google
Lingyun Xiao隶属：德克萨斯大学奥斯汀分校计算机科学系
Zhihan Wang隶属：德克萨斯大学奥斯汀分校计算机科学系
Di Yang Shi隶属：德克萨斯大学奥斯汀分校计算机科学系
Aditya Madhan隶属：德克萨斯大学奥斯汀分校计算机科学系
Johnny Liu隶属：德克萨斯大学奥斯汀分校计算机科学系
Arrasy Rahman隶属：德克萨斯大学奥斯汀分校计算机科学系
Peter Stone隶属：德克萨斯大学奥斯汀分校计算机科学系隶属：Sony AI

###### 摘要
临时团队协作旨在解决设计能够与未知伙伴协调工作的智能体的挑战。然而，该领域的进展受到AHT研究生命周期高昂计算成本、缺乏标准化基准实现以及缺少多样化、经过验证的评估队友套件的阻碍。在本工作中，我们推出了JaxAHT——首个基于JAX的开源库，旨在加速并标准化AHT研究生命周期。利用JAX的硬件加速和大规模并行化能力，JaxAHT为队友生成、自我智能体训练以及针对未知队友的评估提供了一个统一的框架，相比PyTorch对应实现实现了约95倍的墙钟时间加速。除了该库，我们还贡献了一套跨越基于等级的觅食、Overcooked和Hanabi领域的多样化评估队友。为展示该框架的价值，我们使用它进行了大规模、计算受控的基准研究，比较了队友生成和AHT智能体学习方法，发现没有算法能始终保持最佳性能，并且智能体建模主要在基于角色、队友多样化的场景中提供收益。

## 1引言
在现实世界中真正自主运行的智能体必须具备社会智能，以便与其他智能体在共享环境中协作。临时团队协作用形式化表达了这一挑战，要求智能体能够与行为未知的队友有效协调。与传统的多智能体强化学习（MARL）中智能体联合训练不同，AHT智能体必须实时适应未知伙伴——这一能力对于在搜救行动和人机协调等领域进行现实世界部署至关重要。
尽管其重要性，AHT研究面临着三个关键瓶颈。首先，计算成本严重限制了研究迭代速度和严谨性。AHT研究生命周期需要生成多样化的队友群体，训练自我智能体以对抗生成的队友，并在一组未知的队友上进行评估。在标准PyTorch框架下，这些阶段的计算成本高昂，限制了迭代速度，并阻碍了超参数搜索等理想实验实践的执行。其次，缺乏标准的评估队友导致研究人员自行设计队友，但多样性验证有限，这引发了评估彻底性的疑问，并使得不同发表的研究结果无法比较。第三，缺乏标准化的基准实现，加上AHT研究生命周期的复杂性，迫使研究人员执行耗时的算法实现编排工作。这不仅提高了入门门槛，也限制了对解决AHT问题不同组成部分的方法的研究。

本工作推出了JaxAHT，一个基于JAX的库，旨在缓解所有这三个瓶颈。通过利用JAX的硬件加速和大规模并行化能力，JaxAHT提供了统一且加速的AHT训练和评估流水线，以及一套跨越各种标准基准任务的多样化评估队友。本文的贡献如下：

- •统一AHT框架：JaxAHT是首个支持完整AHT研究生命周期的框架，它在一个统一框架内提供了队友生成、自我智能体训练和协作MARL算法。
- •基于JAX的AHT库：JaxAHT是首个用于AHT研究的JAX库，提供了常用AHT算法的标准化实现。我们证明，JaxAHT算法实现相比PyTorch实现了约95倍的墙钟时间加速。
- •多样化评估队友套件：我们在标准AHT基准任务——基于等级的觅食（LBF）、Overcooked和Hanabi上，提供了一套多样化的评估队友，包括启发式队友、基于RL的队友以及基于人类游戏数据生成的队友。我们还首次公开了LBF的人类游戏数据集。
- •基准研究：我们在JaxAHT评估队友上进行了基准研究，严格比较了所有算法，展示了JaxAHT可能进行的大规模、统一的经验分析，并为当前AHT研究现状提供了见解。

通过降低计算障碍、提供多样化的评估队友套件，并为AHT生命周期的方法提供统一框架，JaxAHT使得以前难以进行的严格、大规模AHT研究成为可能。预生成队友和人类游戏的代码和数据集已开源；链接见附录A。

参考图注图1：AHT生命周期。AHT研究需要开发独立的训练和评估队友集，训练AHT智能体，并评估这些智能体对抗未知队友的能力，以评估协作泛化能力。

## 2相关工作
| 库 | JAX | 多环境 | 自我算法 | 队友生成 | 评估队友 | 人类数据 |
|---|---|---|---|---|---|---|
| JaxAHT | ✓ | ✓ | ✓ | ✓ | ✓ | ✓ |
| ZSC‑Eval（Wang等，2024b） | ✓ | ✓ | | ✓ | ✓ | ✓ |
| AH2AC2（Dizdarević等，2025） | ✓ | | ✓ | | ✓ | |
| SocialJax（Guo等，2026） | ✓ | ✓ | ✓ | | | |

表1：库比较。比较了密切相关库包含的功能。

我们简要概述了基于JAX的强化学习库，以及专门针对AHT的基准和评估的相关工作。

##### 基于JAX的强化学习库。
JAX通过JIT编译和自动向量化，支持GPU加速的数值计算和大规模机器学习。PureJaxRL范式表明，当并行运行多个环境和种子时，环境、策略和训练循环都在GPU上执行的端到端JAX实现，与传统的基于PyTorch的方法相比，可实现超过1000倍的加速。PureJaxRL提供了核心强化学习算法的最小化、单文件实现，这些实现完全JIT编译且可向量化。NiceWebRL将JAX环境扩展到基于网络的人类受试者实验，支持人机交互研究。JaxMARL将JAX的优势带入了MARL，实现了标准环境和算法。Jumanji提供了基于JAX的环境，能够在降低复杂性的同时加速迭代和大规模实验。

##### AHT的基准与评估。
诸如Overcooked和Hanabi等环境被广泛用于评估AHT方法，但很少有库同时支持多种AHT算法和环境。最密切相关的工作是ZSC‑Eval、AH2AC2和SocialJax。ZSC‑Eval专注于队友生成算法，并提出了BRDiv用于伙伴选择和BRProx用于衡量泛化能力，但其PyTorch实现使其计算成本高昂。AH2AC2仅专注于计算密集型的Hanabi基准，提供了在大约10万场人类游戏上训练的黄金标准人类代理。然而，访问这些代理受限于一个API。SocialJax为社会困境（而非AHT）提供了环境和算法的JAX实现。结构化比较见表1。JaxAHT补充了先前的工作，是第一个支持完整AHT研究生命周期、跨多个环境和算法类别的端到端JAX框架。

## 3问题设定与设计哲学
我们考虑部分可观测性下的协作多智能体决策，将其建模为N智能体分散部分可观测马尔可夫决策过程（Dec‑POMDP）。Dec‑POMDP定义为$\mathcal{M}=\langle\mathcal{S},\{\mathcal{A}_{i}\}_{i=1}^{N},P,\{\mathcal{O}_{i}\}_{i=1}^{N},\Omega,r,\rho,\gamma\rangle$，其中$\mathcal{S}$是状态空间，$\mathcal{A}_{i}$和$\mathcal{O}_{i}$分别是智能体$i$的动作和观测空间，$P(s'\mid s,\mathbf{a})$是转移函数，$\Omega(\mathbf{o}\mid s',\mathbf{a})$是观测函数，$r(s,\mathbf{a})$是共享的团队奖励，$\rho$是初始状态分布，$\gamma\in[0,1)$是折扣因子。每个智能体$i$根据其局部动作‑观测历史$\tau_{i}^{t}=(o_{i}^{0},a_{i}^{0},\dots,o_{i}^{t})$选择一个动作$a_{i}^{t}$，该历史以智能体$i$在时间$t$收到的私有观测$o_{i}^{t}$结尾，从而产生一个分散策略$\pi_{i}(a_{i}^{t}\mid\tau_{i}^{t})$。联合策略记为$\bm{\pi}=(\pi_{1},\dots,\pi_{N})$。AHT研究的是这样一种情境：自我智能体必须与行为未知且未曾有机会预先协调的队友有效协调。令$e\in\{1,\dots,N\}$表示自我智能体，$\bm{\pi}_{-e}$表示与自我智能体配对的$N-1$个队友的联合策略，这些队友在训练期间来自$\mathcal{T}_{\mathrm{train}}$，在评估时来自$\mathcal{T}_{\mathrm{eval}}$。遵循标准的AHT评估实践，自我智能体被训练以进行广泛合作，然后在从未见过的、来自保留集$\mathcal{T}_{\mathrm{eval}}$的队友上进行评估：
$J_{\mathrm{AHT}}(\pi_{e})=\frac{1}{|\mathcal{T}_{\mathrm{eval}}|}\sum_{\bm{\pi}_{-e}\in\mathcal{T}_{\mathrm{eval}}}J(\pi_{e},\bm{\pi}_{-e})$，
其中$J(\pi_{e},\bm{\pi}_{-e})$是期望的折扣团队回报。在实践中解决此问题涉及三个阶段，每个阶段对应JaxAHT的一个组件。首先，必须构建一组训练队友$\mathcal{T}_{\mathrm{train}}$，可以通过手工设计或依赖队友生成或协作MARL算法。其次，通过与从$\mathcal{T}_{\mathrm{train}}$采样的队友交互，使用自我智能体训练或协作MARL算法来学习自我策略$\pi_{e}$。第三，生成的自我策略必须在一组保留的队友上进行评估，这需要一个标准化的评估队友集。由于原始回报在不同能力的队友间不可比，JaxAHT报告每个回报相对于该队友的*最佳响应*的比例。最佳响应$\mathrm{BR}(\bm{\pi}_{-e})\in\arg\max_{\pi}J(\pi,\bm{\pi}_{-e})$，是能与该队友取得最高回报的策略，可以通过自我智能体训练或MARL算法训练来近似。最佳响应推理是先前AHT工作中构建队友群体的核心。111给定一个智能体，可以将MARL算法应用于训练针对该智能体的最佳响应，将其视为环境的一部分。同样的思想允许协作MARL算法训练自我策略。

为了支持这些阶段，AHT研究依赖于三种主要算法类型：队友生成算法、AHT智能体训练算法和协作MARL算法。前两种是AHT算法的类型，而第三种传统上被作为独立的问题领域进行研究，尽管这三种算法都可用于支持AHT研究的每个阶段。对一种算法类型的研究通常需要其他算法进行训练或评估；例如，评估一种队友生成方法需要一种AHT智能体训练方法。JaxAHT为这些过程提供了统一接口，支持对单个组件或其组合的研究。

图2展示了这些组件如何协同工作。

参考图注图2：JaxAHT设计哲学。JaxAHT在统一的代码库中整合了完整的AHT研究生命周期。在训练期间，通用的智能体和群体抽象支持MARL、队友生成、自我智能体学习以及联合构建训练队友和自我策略的统一AHT方法。在评估期间，学习到的智能体在...