2025-11-22T09:37:22.788326

Foragax: An Agent-Based Modelling Framework Based on JAX

Chaturvedi, El-Gazzar, van Gerven
Foraging for resources is a ubiquitous activity conducted by living organisms in a shared environment to maintain their homeostasis. Modelling multi-agent foraging in-silico allows us to study both individual and collective emergent behaviour in a tractable manner. Agent-based modelling has proven to be effective in simulating such tasks, though scaling the simulations to accommodate large numbers of agents with complex dynamics remains challenging. In this work, we present Foragax, a general-purpose, scalable, hardware-accelerated, multi-agent foraging toolkit. Leveraging the JAX library, our toolkit can simulate thousands of agents foraging in a common environment, in an end-to-end vectorized and differentiable manner. The toolkit provides agent-based modelling tools to model various foraging tasks, including options to design custom spatial and temporal agent dynamics, control policies, sensor models, and boundary conditions. Further, the number of agents during such simulations can be increased or decreased based on custom rules. While applied to foraging, the toolkit can also be used to model and simulate a wide range of other multi-agent scenarios.
academic

Foragax: An Agent-Based Modelling Framework Based on JAX

基本信息

摘要

觅食是生物体在共享环境中为维持内稳态而进行的普遍活动。通过计算机建模多智能体觅食行为，可以以可处理的方式研究个体和集体的涌现行为。智能体建模已被证明在模拟此类任务方面有效，但将模拟扩展到容纳大量具有复杂动力学的智能体仍然具有挑战性。本文提出了Foragax，一个通用、可扩展、硬件加速的多智能体觅食工具包。利用JAX库，该工具包可以以端到端向量化和可微分的方式模拟数千个智能体在共同环境中的觅食行为。该工具包提供智能体建模工具来模拟各种觅食任务，包括设计自定义空间和时间智能体动力学、控制策略、传感器模型和边界条件的选项。此外，在模拟过程中可以根据自定义规则增加或减少智能体数量。虽然应用于觅食，该工具包也可用于建模和模拟广泛的其他多智能体场景。

研究背景与动机

问题定义

觅食是所有生物体为生存而必须解决的最古老和关键的行为优化问题之一。从最简单的细菌到高度社会化的灵长类动物，高效觅食的本能在进化的各个阶段都普遍存在。自然界中的觅食本质上是一个多智能体活动，因为相邻的生物体共享同一个觅食环境。

现有方法的局限性

  1. 可扩展性问题: 传统的智能体建模方法在处理大量智能体时面临计算瓶颈
  2. 生物真实性不足: 现有模拟中智能体数量的限制导致与自然界观察到的觅食行为存在不一致性
  3. 计算复杂性: 每个智能体的状态依赖于其他所有智能体的行为，随着智能体数量增加，模拟变得极其困难

研究动机

作者希望为觅食任务建模者提供在单一环境中同时模拟大量（数千个）智能体的能力。这种大规模模拟能力对于理解复杂的集体行为涌现至关重要，并能缩短理论研究与自然观察之间的差距。

核心贡献

  1. 提出了Foragax框架: 基于JAX的通用、可扩展、硬件加速的多智能体觅食建模工具包
  2. 实现了大规模并行模拟: 能够通过端到端向量化方法模拟数千个智能体
  3. 提供了灵活的建模接口: 支持自定义空间和时间动力学、控制策略、传感器模型和边界条件
  4. 支持动态智能体管理: 在模拟过程中可以根据自定义规则增加或减少智能体数量
  5. 实现了非情节性连续模拟: 支持不重置环境和智能体状态的长期模拟

方法详解

任务定义

Foragax专注于模拟大量智能体在共同环境中的觅食行为，同时维持内稳态。任务特点包括：

  • 输入: 智能体初始状态、环境配置、资源分布
  • 输出: 智能体行为轨迹、集体涌现行为模式
  • 约束: 资源有限性、智能体间相互作用、环境边界条件

框架架构

Foragax基于五个主要抽象层次的类设计：

1. Policy（策略）

  • 由状态变量和参数集合表示
  • 用于智能体在环境中产生行为
  • 支持多种模型：Wilson-Cowan模型、循环神经网络、多层感知器
  • 可建模为确定性或随机性质

2. Agent（智能体）

  • 由状态变量和参数集合表示
  • 支持自定义时间和空间动力学
  • 维护唯一索引和活跃状态（1表示活跃，0表示非活跃）
  • 动力学可由其策略操控

3. Agent-set（智能体集合）

  • 智能体的集合表示
  • 所有智能体共享相同的状态变量、参数和策略结构
  • 支持选择、排序、添加或删除智能体操作
  • 通过零填充保持数据结构形状恒定，维持JAX的JIT和自动向量化特性

4. Learning Algorithms（学习算法）

  • 提供调优策略和智能体参数的算法集合
  • 当前实现基于无梯度神经进化方法
  • 计划扩展到基于梯度的强化学习方法

5. Tasks（任务）

  • 允许设计具有自定义边界条件和墙体放置的觅食环境
  • 基于光线投射算法进行碰撞检测

技术创新点

1. JAX后端优化

  • 利用JAX的即时编译（JIT）和自动向量化（AV）
  • 实现端到端向量化和可微分计算
  • 支持硬件加速（GPU/TPU）

2. 零填充策略

  • 通过预先分配最大尺寸的数据结构
  • 使用零填充维持恒定形状
  • 在智能体集合操作中保持JAX优化特性

3. 连续非情节性模拟

  • 不重置环境和智能体状态
  • 支持种群规模波动
  • 促进开放式适应行为的涌现

实验设置

实验配置

  • 智能体数量: 1000个智能体
  • 神经网络规模: 每个智能体50个神经元的基于速率的神经群体模型
  • 资源数量: 300个资源补丁
  • 模拟步数: 100万步
  • 计算平台: 集群计算机

动力学模型

资源动力学

资源价值遵循Lotka-Volterra动力学：

ṡn = ϵsn - αs²n - Σ w(xn, xm)

其中：

  • ϵ和α分别为资源的增长和衰减率
  • w(xn, xm)表示与第n个资源位置xn和第m个智能体位置xm之间距离成反比的函数

智能体运动模型

智能体位置和速度采用双积分器模型：

q̇m = [0 1; 0 0]qm + [0; 1]um

其中：

  • qm = xm, ẋmᵀ为智能体状态向量
  • um ∈ ℝ为加速度输入
  • 加速度由50个神经元的循环神经网络线性读出的平均值确定

实验结果

主要结果

  • 计算效率: 1000个智能体、300个资源、100万步的模拟在40分钟内完成
  • 可扩展性验证: 成功展示了600个智能体在600个资源中的觅食行为
  • 系统稳定性: 模拟过程中智能体和资源的空间分布保持合理

性能表现

实验结果表明Foragax能够：

  1. 高效处理大规模多智能体模拟
  2. 维持实时性能要求
  3. 生成生物学上合理的觅食行为模式

可视化结果

论文提供了600个智能体（蓝色点）在600个资源（红色点）中分散觅食的可视化，箭头表示智能体方向，展示了系统的空间组织特性。

相关工作

智能体建模框架

  • Agents.jl: 高性能、功能丰富的智能体建模软件
  • Mesa: Python中的智能体建模框架
  • 现有框架在大规模模拟方面存在性能瓶颈

觅食理论研究

  • 从简单细菌到复杂灵长类的觅食行为研究
  • 最优觅食理论的发展
  • 内部能量储备对奖励响应性的影响研究

多智能体系统应用

  • 疾病传播建模
  • 群体行为模拟
  • 调度优化问题

结论与讨论

主要结论

  1. Foragax成功实现了大规模多智能体觅食行为的高效模拟
  2. JAX后端显著提升了计算性能和可扩展性
  3. 框架提供了足够的灵活性来建模各种觅食场景
  4. 连续非情节性模拟能够揭示有趣的涌现行为

局限性

  1. 预分配限制: 需要预先知道数据结构的最大尺寸
  2. 内存消耗: 零填充策略可能导致内存使用效率不高
  3. 算法范围: 当前主要支持神经进化算法，强化学习支持有待完善
  4. 验证不足: 缺乏与真实生物觅食行为的定量比较

未来方向

  1. 扩展到基于梯度的强化学习算法
  2. 改进内存使用效率
  3. 增加更多生物学真实性验证
  4. 开发更丰富的分析工具

深度评价

优点

  1. 技术创新性强: 将JAX的高性能计算能力引入多智能体建模领域
  2. 实用价值高: 解决了大规模智能体模拟的关键瓶颈问题
  3. 设计合理: 五层抽象架构清晰，易于扩展和使用
  4. 性能优异: 在保持功能性的同时实现了显著的性能提升

不足

  1. 实验验证有限: 主要展示了技术可行性，缺乏深入的行为分析
  2. 生物学验证不足: 未与真实觅食行为进行定量比较
  3. 理论贡献有限: 主要是工程实现，理论创新相对较少
  4. 文档不够详细: 某些技术细节描述不够充分

影响力

  1. 领域贡献: 为多智能体系统研究提供了强大的计算工具
  2. 实用价值: 能够支持更大规模、更复杂的生态学和行为学研究
  3. 可复现性: 开源代码增强了研究的可复现性和可扩展性
  4. 跨领域应用: 虽然专注于觅食，但可扩展到其他多智能体场景

适用场景

  1. 生态学研究: 大规模动物觅食行为建模
  2. 进化生物学: 群体进化动力学研究
  3. 人工智能: 多智能体强化学习算法测试
  4. 复杂系统: 涌现行为和自组织现象研究
  5. 社会科学: 大规模社会动力学建模

参考文献

论文引用了多个重要领域的关键文献，包括：

  • JAX库的技术文档
  • 觅食理论的经典著作（Stephens & Krebs, 1986）
  • 现有智能体建模框架（Agents.jl, Mesa）
  • 神经进化和强化学习的最新进展

总体评价: 这是一篇高质量的系统论文，通过将现代高性能计算技术与多智能体建模相结合，解决了该领域的一个重要技术瓶颈。虽然在理论创新方面相对有限，但其工程价值和实用性使其成为多智能体系统研究的重要贡献。