在JAX中多智能体强化学习环境和算法

这篇笔记算算是NIPS论文精读的第一步，还是得精读论文来学习，乱找别人的总结还是不能加深对方向现状的理解。这次选的是JaxMARL，一篇NIPS2024关于MARL的论文。NIPS2024总共4538篇论文，Topic明确为MARL的只有33篇，这方向，，，，回头再看看别的顶刊，顶会的情况再下判断。

论文链接： https://doi.org/10.48550/arXiv.2311.10090

代码链接： https://github.com/flairox/jaxmarl

介绍

强化学习环境通常运行在CPU上，效率比较低，需要通过硬件加速和并行化来解决这个问题，简单来说就是把环境搬到GPU上跑。要在GPU上跑环境需要很多非Python代码的支持，好在这些基础工作已经有人做了，JAX可以在GPU上实现并行化和即时编译，在此基础上又诞生了PureJaxRL，直接在GPU上并行跑训练，比CPU快了4000倍，让大规模强化学习算法的实现变得可能。JAXMARL就是在PureJaxRL、JAX等的基础上搭建的一个MARL基准，包含了环境和算法的GPU运算实现。

作者还提到一个有趣的现象：在近75篇MARL论文中，50%只用了一个评估环境，30%用了两个。所以他们实现了很多环境，希望大家能在多种场景下进行评估。其中重点实现了SMAX和STORM，SMAX是对SMAC（星际争霸2多智能体环境）的近似实现，定制化程度更高，速度快了40000倍；STORM是受 melting pot 启发开发的一个博弈环境，后文介绍STORM。关于速度提升，作者列举了很多数字，我觉得最核心的数字是12500倍，在向量化多环境下。

SMAX

SMAC有一些局限性，对复杂策略缺乏足够的随机性（可以参考"An improved benchmark for cooperative multi-agent reinforcement learning"这篇文章），开环策略——也就是仅依赖时间步、不需要观测的策略，在SMAC环境中居然表现不错。而且SMAC环境内存占用大、速度慢，所以作者提出了SMAX环境，和SMAC类似但更轻量化，有着不容易被智能体钻空子的AI，包含了SMAC的场景和SMACv2的类似场景，定制性很高。

STORM

STORM全称是Spatial-Temporal Representations of Matrix Games，中文应该是矩阵博弈时空表征环境。先说说矩阵游戏的概念，就像囚徒困境那样，表现形式是这样的：

玩家1 \ 玩家2	合作	背叛
合作	(3, 3)	(0, 5)
背叛	(5, 0)	(1, 1)

表格中的数字表示(玩家1收益, 玩家2收益)。可以看到，如果双方都合作，各得3分；如果一方背叛一方合作，背叛者得5分，合作者得0分；如果双方都背叛，各得1分。囚徒困境的博弈点在于：虽然双方合作是最优解，但每个玩家都有背叛的动机，因为无论对方选择什么，背叛总能带来更高的个人收益，最终导致双方都选择背叛，陷入次优均衡。

其中每个智能体只能选择合作或者背叛，这样的设定太抽象了，或者说太简单了。所以就有了Melting Pot的In the Matrix游戏。这个环境我还不太了解，相关论文在 https://doi.org/10.48550/arXiv.2211.13746 ，后面再读读这篇文章。总之这是一个用来衡量多智能体博弈、合作、竞争等情况的环境，具有很强的社会性，看起来挺有意思的。STORM在Melting Pot的基础上加入了随机性要素。

其他环境

• Hanabi：花火，完全合作式部分可观测卡牌游戏，详细资料可以参照 https://zh.wikipedia.org/wiki/%E8%8A%B1%E7%81%AB_(%E5%8D%A1%E7%89%8C%E9%81%8A%E6%88%B2) ，是零样本协作、心智理论与临时团队研究的基准。
• Overcooked：分手厨房，用于评估完全合作式全可观测人机协作任务的性能，玩过的都知道，详细资料可以参照 https://doi.org/10.48550/arXiv.1910.05789 ，没想到19年就有人拿分手厨房做RL环境了。
• MABrax：源于多智能体MuJoCo，是连续多智能体机器人控制任务的常用基准。
• Multi-Agent Particle Environment（MPE）：基于简单物理规则的2D世界，粒子智能体可以移动、通信并与固定地标交互。
• 硬币游戏：模拟社会困境的双玩家网格世界环境，虽然是一般和博弈的常用基准，但STORM修正了其既有缺陷。
• 开关谜题：作为调试工具的简易合作通信任务。

算法

作者使用PureJaxRL实现了IPPO和MAPPO，参照PyMARL实现了IQL、VDN和QMIX。

使用推荐

作者特别强调，这个基准环境被设计来推动CTDE、零次协调、一般和博弈、协作式连续动作这四个领域的研究，建议的评估如下：

研究领域	推荐环境
CTDE	SMAX（全部场景）、Hanabi（2–5 人）、Overcooked
零试协调	Hanabi（2 人）、Overcooked（5 个基础场景）
一般和博弈	STORM（重复囚徒困境）、STORM（匹配硬币）
协作式连续动作	MABrax

作者提出了一套评估指标计算方法：先对算法在相关环境类别上的表现进行标准化，计算每个随机种子的平均性能。然后，在每个环境中，基于这些汇总的统计数据计算四分位均值。这种方法可以对环境类别进行公平比较，不会因为某些类别包含更多独立场景而使其权重过高。他们还写了算指标的代码，可以直接用。

性能评估

具体的对比情况还是直接看原文比较方便，这里给出作者的结论：在单个GPU上训练，对于MPE环境下的QMIX算法，只需要198.4秒就能完成1024次独立训练任务，而基于PyMARL的单一训练需要1小时10分钟，相当于每个智能体实现了21,500倍的加速比。附录中对IPPO算法的重复实验显示加速比为12,500倍。图5c展示了采用JaxMARL的IPPO实现在SMAX环境对比PyMARL在SMAC环境训练所获得的加速效果：在不同数量的环境推演线程下，JaxMARL最高可以实现31倍的加速比。

他们还对每个环境的实现效果进行了评估，这里就不详细说了。按照他们的说法，自然是和环境的经典实现效果上差别不大，部分有改进和简化在后文有说明，可以放心使用。相关工作和结论部分就不看了。

附录

附录里写了前面提到的每个环境的介绍，还有算法、API、环境的详细评估、默认超参等等，很细致。这里只聊聊他们实现的环境，其他部分还不如直接跑跑代码，看看文档。

SMAX

说说一些重点差异。首先单位被建模成二维世界的一个圆，支持智能体自我对抗，还能自定义单位类型。观测方面：每个智能体可以感知视野范围内所有友方与敌方单位的生命值、上一步动作、位置坐标、武器冷却状态及单位类型，采用《星际争霸II》设定的动态视野与攻击范围，而不是SMAC的固定数值。奖励方面：SMAC的伤害输出在奖励中的占比会随智能体的数量进行线性缩放，SMAX把奖励均分为伤害输出和对战胜负，解决了在很多智能体的情况下伤害输出占比过大的问题，比如在27m_vs_30m场景里，随机策略甚至能得10/20，大部分奖励都来自于伤害输出。

AI这块，SMAX的AI是去中心化的，而SMAC和SMACv2的AI是中心化的，据说可以在去中心化的前提下达到和中心化AI一样的强度，不过这样从原理上保证任何低于50%的胜率都代表明确的学习失败，因为算法可以直接复制AI的策略。

执行上，SMAX部分动作是同步执行的，可能导致平局发生，此时双方奖励均为0。每个环境步包括8个小时间步，动作是离散的，由四向移动、停止及针对每个敌人的射击构成。

设定上，移除了异虫的回血机制（因为这对训练的影响很微弱，只有0.38/s），星灵的护盾机制（脱战10s开始回护盾，在微操任务中基本没用，以生命值补偿替代），删除了一些单位，比如医疗运输机、巨像、爆虫等，进行环境的精简。可惜了，不能看到智能体枪兵甩毒爆，运输机闪光头了。

初始位置生成机制和SMACv2一致，不过在"surrounded"这种初始位置设定下，盟友和敌人的初始位置是对称生成的，保证了50%胜率的理论可达性。

STORM

这是一个8×8的网格环境，智能体只有四个动作：前进、左转、右转、交互。观测范围有限，碰撞时静止优先，如果都在运动则随机判定。智能体要去收集地图上的合作币和背叛币，收集到后颜色会转变，表示可以交互了。交互时向前方发射一道交互光束（挺有趣的设定），如果碰到一个也准备好的智能体就会发生交互，根据博弈矩阵给出奖励收益。交互完成后，两个智能体在五步内会陷入冻结状态，向周围的智能体公布它获取的硬币，解除冻结后在新位置重生，硬币归零，重新开始获取硬币。每个episode结束后，硬币的位置会随机改变。作者强调了STORM和melting pot的区别，不过这里就不介绍了，后面研究完melting pot再做介绍。他还特别强调了这个硬币随机改变机制的重要性。

硬币博弈

这是一个双智能体的网格环境游戏，模拟了迭代囚徒困境等社会困境。地图是环形的，上下联通，左右联通，就像贪吃蛇的地图。红蓝两个智能体在其中移动，地图上有红蓝两种硬币，任何智能体捡到硬币就加一分，不过红色智能体收集到蓝色硬币会给蓝色智能体减2分，反之亦然。一个硬币被消耗后，同样的硬币会在地图随机位置刷新。如果两个智能体同时吃一个硬币，它们都可以吃到这个硬币。每个训练回合的step是固定的。博弈点在于，如果两个智能体都选择看到硬币就吃，那么总奖励期望是0。

Hanabi

花火，完全合作式、部分可观测的多玩家卡牌游戏。玩家可以观察其他玩家的手牌却看不到自己的牌。玩家团队必须通过有限的线索交流，按特定顺序打出卡牌才能获胜。是零样本协作与临时团队研究领域的常用基准。作者的实现支持自定义游戏配置（如花色/等级数量、玩家数量、提示令牌数量），比经典实现更加灵活和可读。

Overcooked

分手厨房，用来评估完全合作式、完全可观察的人机协作任务性能，看人玩玩就知道了。