Foragax: An Agent-Based Modelling Framework Based on JAX

Basic Information

• Paper ID: 2409.06345
• Title: Foragax: An Agent-Based Modelling Framework Based on JAX
• Authors: Siddharth Chaturvedi, Ahmed El-Gazzar, Marcel van Gerven (Radboud University)
• Classification: cs.MA (Multi-Agent Systems)
• Publication Date: September 2024 (arXiv preprint)
• Paper Link: https://arxiv.org/abs/2409.06345
• Code Repository: https://github.com/i-m-iron-man/Foragax.git

Abstract

Foraging is a universal activity performed by organisms in shared environments to maintain homeostasis. By computationally modeling multi-agent foraging behavior, individual and collective emergent behaviors can be studied in a tractable manner. Agent-based modeling has proven effective in simulating such tasks, but extending simulations to accommodate large numbers of agents with complex dynamics remains challenging. This paper presents Foragax, a general-purpose, scalable, hardware-accelerated multi-agent foraging toolkit. Leveraging the JAX library, the toolkit can simulate thousands of agents foraging in a shared environment in an end-to-end vectorized and differentiable manner. The toolkit provides agent modeling tools to simulate various foraging tasks, including options for designing custom spatial and temporal agent dynamics, control policies, sensor models, and boundary conditions. Furthermore, agent populations can be dynamically increased or decreased during simulation according to custom rules. While applied to foraging, the toolkit can also be used for modeling and simulating a wide range of other multi-agent scenarios.

Research Background and Motivation

Problem Definition

Foraging is one of the oldest and most critical behavioral optimization problems that all organisms must solve for survival. From the simplest bacteria to highly socialized primates, the instinct for efficient foraging is ubiquitous across all stages of evolution. Foraging in nature is inherently a multi-agent activity, as adjacent organisms share the same foraging environment.

Limitations of Existing Approaches

1. Scalability Issues: Traditional agent-based modeling methods face computational bottlenecks when handling large numbers of agents
2. Insufficient Biological Realism: Limitations on agent numbers in existing simulations lead to inconsistencies with foraging behaviors observed in nature
3. Computational Complexity: Each agent's state depends on the behavior of all other agents, making simulation extremely difficult as agent numbers increase

Research Motivation

The authors aim to provide foraging task modelers with the capability to simultaneously simulate large numbers (thousands) of agents in a single environment. This large-scale simulation capability is crucial for understanding complex emergent collective behaviors and bridging the gap between theoretical research and natural observations.

Core Contributions

1. Proposed the Foragax Framework: A general-purpose, scalable, hardware-accelerated multi-agent foraging modeling toolkit based on JAX
2. Implemented Large-Scale Parallel Simulation: Capable of simulating thousands of agents through end-to-end vectorization methods
3. Provided Flexible Modeling Interface: Supports custom spatial and temporal dynamics, control policies, sensor models, and boundary conditions
4. Enabled Dynamic Agent Management: Agent populations can be increased or decreased during simulation according to custom rules
5. Implemented Non-Episodic Continuous Simulation: Supports long-term simulation without resetting environment and agent states

Methodology Details

Task Definition

Foragax focuses on simulating large numbers of agents foraging in a shared environment while maintaining homeostasis. Task characteristics include:

• Input: Initial agent states, environment configuration, resource distribution
• Output: Agent behavior trajectories, collective emergent behavior patterns
• Constraints: Resource scarcity, inter-agent interactions, environmental boundary conditions

Framework Architecture

Foragax is designed based on five main abstraction levels:

1. Policy

• Represented by state variables and parameter sets
• Used to generate agent behavior in the environment
• Supports multiple models: Wilson-Cowan model, recurrent neural networks, multilayer perceptrons
• Can be modeled as deterministic or stochastic

2. Agent

• Represented by state variables and parameter sets
• Supports custom temporal and spatial dynamics
• Maintains unique index and active status (1 for active, 0 for inactive)
• Dynamics can be controlled by its policy

3. Agent-set

• Collective representation of agents
• All agents share the same state variables, parameters, and policy structure
• Supports selection, sorting, adding, or removing agent operations
• Maintains constant data structure shape through zero-padding, preserving JAX's JIT and automatic vectorization properties

4. Learning Algorithms

• Provides a collection of algorithms for tuning policy and agent parameters
• Current implementation based on gradient-free neuroevolution methods
• Plans to extend to gradient-based reinforcement learning methods

5. Tasks

• Allows designing foraging environments with custom boundary conditions and wall placements
• Implements collision detection using ray-casting algorithms

Technical Innovations

1. JAX Backend Optimization

• Leverages JAX's just-in-time compilation (JIT) and automatic vectorization (AV)
• Enables end-to-end vectorized and differentiable computation
• Supports hardware acceleration (GPU/TPU)

2. Zero-Padding Strategy

• Pre-allocates data structures with maximum size
• Uses zero-padding to maintain constant shape
• Preserves JAX optimization properties during agent-set operations

3. Continuous Non-Episodic Simulation

• Does not reset environment and agent states
• Supports population size fluctuations
• Facilitates emergence of open-ended adaptive behaviors

Experimental Setup

Experimental Configuration

• Number of Agents: 1,000 agents
• Neural Network Scale: Rate-based neural population model with 50 neurons per agent
• Number of Resources: 300 resource patches
• Simulation Steps: 1 million steps
• Computing Platform: Cluster computer

Dynamics Models

Resource Dynamics

Resource value follows Lotka-Volterra dynamics:

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

Where:

• ϵ and α are resource growth and decay rates, respectively
• w(xn, xm) represents a function inversely proportional to the distance between the n-th resource location xn and the m-th agent location xm

Agent Movement Model

Agent position and velocity employ a double-integrator model:

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

Where:

• qm = xm, ẋmᵀ is the agent state vector
• um ∈ ℝ is the acceleration input
• Acceleration is determined by the average of linear readouts from a recurrent neural network with 50 neurons

Experimental Results

Main Results

• Computational Efficiency: Simulation of 1,000 agents, 300 resources, and 1 million steps completed in 40 minutes
• Scalability Verification: Successfully demonstrated foraging behavior of 600 agents among 600 resources
• System Stability: Spatial distribution of agents and resources remained reasonable throughout simulation

Performance Characteristics

Experimental results demonstrate that Foragax can:

1. Efficiently handle large-scale multi-agent simulations
2. Maintain real-time performance requirements
3. Generate biologically plausible foraging behavior patterns

Visualization Results

The paper provides visualizations of 600 agents (blue dots) dispersed foraging among 600 resources (red dots), with arrows indicating agent direction, demonstrating the system's spatial organization properties.

Related Work

Agent-Based Modeling Frameworks

• Agents.jl: High-performance, feature-rich agent-based modeling software
• Mesa: Agent-based modeling framework in Python
• Existing frameworks face performance bottlenecks in large-scale simulations

Foraging Theory Research

• Foraging behavior research from simple bacteria to complex primates
• Development of optimal foraging theory
• Research on the influence of internal energy reserves on reward responsiveness

Multi-Agent System Applications

• Disease transmission modeling
• Collective behavior simulation
• Scheduling optimization problems

Conclusions and Discussion

Main Conclusions

1. Foragax successfully achieves efficient simulation of large-scale multi-agent foraging behavior
2. The JAX backend significantly enhances computational performance and scalability
3. The framework provides sufficient flexibility to model various foraging scenarios
4. Continuous non-episodic simulation can reveal interesting emergent behaviors

Limitations

1. Pre-allocation Constraints: Requires prior knowledge of maximum data structure sizes
2. Memory Consumption: Zero-padding strategy may result in inefficient memory usage
3. Algorithm Scope: Currently primarily supports neuroevolution algorithms; reinforcement learning support needs improvement
4. Insufficient Validation: Lacks quantitative comparison with real biological foraging behavior

Future Directions

1. Extend to gradient-based reinforcement learning algorithms
2. Improve memory usage efficiency
3. Increase biological realism validation
4. Develop richer analytical tools

In-Depth Evaluation

Strengths

1. Strong Technical Innovation: Introduces JAX's high-performance computing capabilities to multi-agent modeling
2. High Practical Value: Addresses key computational bottlenecks in large-scale agent simulation
3. Reasonable Design: Five-layer abstract architecture is clear, easy to extend and use
4. Excellent Performance: Achieves significant performance improvements while maintaining functionality

Weaknesses

1. Limited Experimental Validation: Primarily demonstrates technical feasibility; lacks in-depth behavioral analysis
2. Insufficient Biological Validation: No quantitative comparison with real foraging behavior
3. Limited Theoretical Contribution: Primarily an engineering implementation with relatively limited theoretical innovation
4. Insufficient Documentation: Some technical details lack adequate description

Impact

1. Domain Contribution: Provides powerful computational tools for multi-agent systems research
2. Practical Value: Enables larger-scale and more complex ecological and behavioral studies
3. Reproducibility: Open-source code enhances research reproducibility and extensibility
4. Cross-Domain Applications: While focused on foraging, extensible to other multi-agent scenarios

Applicable Scenarios

1. Ecological Research: Large-scale animal foraging behavior modeling
2. Evolutionary Biology: Population evolutionary dynamics research
3. Artificial Intelligence: Multi-agent reinforcement learning algorithm testing
4. Complex Systems: Emergent behavior and self-organization phenomenon research
5. Social Science: Large-scale social dynamics modeling

References

The paper cites key literature from multiple important domains, including:

• JAX library technical documentation
• Classical works on foraging theory (Stephens & Krebs, 1986)
• Existing agent-based modeling frameworks (Agents.jl, Mesa)
• Recent advances in neuroevolution and reinforcement learning

Overall Assessment: This is a high-quality systems paper that addresses an important technical bottleneck in the field by combining modern high-performance computing techniques with multi-agent modeling. While relatively limited in theoretical innovation, its engineering value and practical utility make it an important contribution to multi-agent systems research.