ABMax: A JAX-based Agent-based Modeling Framework

TL;DR

ABMax addresses the challenge of updating a dynamic subset of agents under JAX's static-shape constraint by introducing two algorithms, $RM$ and $SCI$. The framework leverages JAX transforms (jit, vmap) and flax.struct to keep data structures compatible with transformations, enabling vectorization across multiple ABMs. It demonstrates on a canonical predation benchmark and on traffic and financial market scenarios, showing competitive performance and scalable parallel simulations. The work provides a practical, vectorization-friendly ABM toolkit that reduces boilerplate and broadens the use of large-scale ABMs on accelerators.

Abstract

Agent-based modeling (ABM) is a principal approach for studying complex systems. By decomposing a system into simpler, interacting agents, agent-based modeling (ABM) allows researchers to observe the emergence of complex phenomena. High-performance array computing libraries like JAX can help scale such computational models to a large number of agents by using automatic vectorization and just-in-time (JIT) compilation. One of the caveats of using JAX to achieve such scaling is that the shapes of arrays used in the computational model should remain immutable throughout the simulation. In the context of agent-based modeling (ABM), this can pose constraints on certain agent manipulation operations that require flexible data structures. A subset of which is represented by the ability to update a dynamically selected number of agents by applying distinct changes to them during a simulation. To this effect, we introduce ABMax, an ABM framework based on JAX that implements multiple just-in-time (JIT) compilable algorithms to provide this functionality. On the canonical predation model benchmark, ABMax achieves runtime performance comparable to state-of-the-art implementations. Further, we show that this functionality can also be vectorized, making it possible to run many similar agent-based models in parallel. We also present two examples in the form of a traffic-flow model and a financial market model to show the use case of ABMax

Figures (7)

• Figure 1: Unified Modeling Language (UML) diagram representing different classes in the ABMax framework. These classes can be divided into the helper and the core classes.
• Figure 2: Depiction of Wilensky’s discrete 2D predation model Wilensky1997. Wolves gain energy by consuming sheep, and sheep gain energy by eating grass, which regrows on each cell after a fixed delay. Both species move randomly and reproduce asexually with a given probability, spawning offspring at their current location. ABMax leverages JAX’s vmap to run multiple such simulations in parallel over identical data structures.
• Figure 3: Trend in the number of wolves and number of sheep for similar initial conditions using ten different seeds, simulated for 100 steps in a small predation environment.
• Figure 4: A cellular‐automaton traffic model: green arrows mark allowed moves, red arrows blocked actions, and purple arrows potential conflicts. Cars spawn in yellow entry cells and travel toward exit cells under green/red signals. By representing everything with ABMax data structures and applying JAX’s vmap, we can parallelize simulations across multiple roads with different car counts.
• Figure 5: The variation in the number of cars across ten different roads for 100 steps. Wherein each road has a different seed leading to different trends in the green/red signals in the exit cells, and a different number of cars are spawned at each step. Each road is seven cells long and three lanes wide, thus it can host a maximum of 21 cars.