ABM-Graph-Diffusion-Network

This is the repository for the implementation of the models and the experiments described in the paper "Learning Individual Behavior in Agent-Based Models with Graph Diffusion Networks".

Graph Diffusion Network (GDN) is a framework that combines a graph neural network with a diffusion model to learn a differentiable surrogate of an ABM, from ABM-generated data. In this framework, a graph neural network captures the interactions that govern the evolution of each agent's state in response to other agents, while the diffusion model learns the distribution of possible state transitions, conditioned on these interactions.

Installation

To get started with the GDN framework, start by cloning the repository:

git clone https://github.com/fracozzi/ABM-Graph-Diffusion-Network.git
cd ABM-Graph-Diffusion-Network

Python Requirements

The following Python libraries are required to run the GDN:

• PyTorch: Used to define the neural network and train the model.
• PyTorch Geometric (pyg): Used for graph neural network implementation.
• Jupyter: To run and interact with notebooks.
• Matplotlib: For creating plots and visualizations.
• NumPy, Pickle: To manipulate arrays and save and load data.
• TQDM: For displaying progress bars when running experiments.

Our code was tested using Python version 3.12.

To install all the required Python libraries, use the requirements.txt file provided in the repository:

pip install -r requirements.txt

Running Experiments

1. Build the ABM Datasets

Before training or evaluation, you need to generate the ABM datasets for the Schelling and Predator–Prey models. From the repository root, run:

python models/abm/schelling.py
python models/abm/predatorprey.py

These scripts will produce the simulation trajectories and construct the graph‐structured datasets used by the GDN. Hyperparameters such as agent density, number of ramifications, number of timesteps...etc. can be changed in the script.

2. Train the Surrogate Model

Once the datasets are ready, you can train the Graph Diffusion Network surrogate on a chosen ABM configuration. Use the following command-line options:

python train_model.py \
  --abm_model   <predatorprey|schelling> \
  --model_type  <GDN|diffusion-only|gnn-only> \
  --parameter   <psi1|psi2|psi3|psi4> \      # for predatorprey
               <xi1|xi2|xi3> \               # for schelling
  --learning_rate  <float> \
  --T_diffusion    <int> \
  --n_epochs       <int>

Example:

Train the GDN model on the Predator–Prey model with parameter psi1, a learning rate of 1e-5, 100 diffusion steps, and 100 epochs:

python train_model.py \
  --abm_model predatorprey \
  --model_type GDN \
  --parameter psi1 \
  --learning_rate 1e-5 \
  --T_diffusion 100 \
  --n_epochs 100

Use -h or --help to display all available options and their defaults:

python train_model.py --help

3. Evaluate the Trained Surrogates

After training, assess model fidelity on held-out trajectories using the provided evaluation scripts.

python evaluate_predator_prey.py \
  --model_type <GDN|diffusion-only|gnn-only> \
  --psi        <psi1|psi2|psi3|psi4>

For example, to evaluate the GDN model trained with psi1:

python evaluate_predator_prey.py --model_type GDN --psi psi1

Notes

For further assistance, please refer to the comments in the scripts or notebooks, or feel free to contact me at: f.cozzi@uniroma1.it