Foragax: An Agent-Based Modelling Framework Based on JAX

TL;DR

Foragax is a scalable, hardware-accelerated agent-based modelling toolkit built on JAX, enabling efficient simulation of thousands of agents in complex foraging tasks and adaptable to various multi-agent scenarios.

Contribution

The paper introduces Foragax, a novel, flexible, and high-performance framework for large-scale multi-agent simulations using JAX, with customizable dynamics and control policies.

Findings

Supports thousands of agents in a common environment

End-to-end vectorized and differentiable simulations

Flexible for various multi-agent scenarios

Abstract

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.