# Heterogeneity Distance Measurement and HetDPS

Official implementation of the NeurIPS 2025 submission ***Heterogeneity in Multi-Agent Reinforcement Learning***.

## Introduction

In this repository, we provide the implementation of the *heterogeneity distance* measurement and *HetDPS* algorithm. All implementations are integrated into the MARL training process. In the MARL training process, the environment-related heterogeneity distance is quantified at regular intervals, and the policy distance is also quantified. Then, the HetDPS algorithm is run. The code also provides visualization functions, which can visualize the heterogeneity distance matrix and policy distance matrix, and support saving the experimental results to the local host.

## Requirements

To run this repository you need:

i) Install the code's requirements. A virtual environment based on Conda is recommended. Install with:
```setup
conda create --name madps --file requirements.txt
```
Please note that some libraries may have order dependencies, so when the automatic installation fails, you can try to install them manually.

ii) Install the supported MARL-environments, for example:
- Multi-agent Spread of [Petting-Zoo Version](https://github.com/semitable/PettingZoo)




## Running

You can then simply run `main_Het.py` using:
```python
python main_Het.py with algorithm_mode='hetdps' env_name=\'pettingzoo:pz-mpe-large-spread-v1\' time_limit=50 parallel_envs=32 experiment_label='NeurIPS' 

```





## Structure of this repository
This repository is structured as follows:
### 1. Actor-Critic Training and Execution Framework (`main_Het.py`)

main_Het.py includes:
- Multi-agent Environment Construction and Maintenance.
- Environment Sampling and Sample Pool Construction.
- Neural Network Evaluating and Training.

More details can be found in the comments of main function in main_Het.py.

### 2. Multi-agent Heterogeneity Distance Computing and HetDPS (`utils_Het.py`)

utils_Het.py includes:

- `compute_fusions` (including policy distance):
    1. Train Conditional VAE (Learning the conditional representations of agents' decisions).
    2. Use VAE and agents' models to calculate dij（Compute the multi-agent policy distance matrix using the learned conditional representations).

- `compute_implicit_het` (including meta-het distance):
  1. Train Conditional VAE (Learning the conditional representations of agents' decisions).
  2. Use VAE and agents' models to calculate dij（Compute the multi-agent policy distance matrix using the learned conditional representations).

- `compute_dynamic_parameter_sharing` (including dynamic parameter sharing distance):
  

- `calculate_N_Gaussians_BD`: This function is used for parallel calculation of the Bhattacharyya distance between multiple Gaussian distributions using PyTorch.
- `calculate_N_Gaussians_Hellinger_through_BD`: This function is used for parallel calculation of the Hellinger distance between multiple Gaussian distributions using PyTorch. It requires results from the `calculate_N_Gaussians_BD` function.
- `calculate_N_Gaussians_WD`: This function is used for the parallel calculation of the Wasserstein distance between multiple Gaussian distributions using PyTorch.

### 3. Multi-agent Nerual Network Models (model_NF.py)

model_NF.py includes:
- `MADPSNet`: Multi-agent neural network models that support dynamic adjustment and hierarchical adjustment of parameter sharing.
- `MultiAgentFCNetwork`: Multi-agent neural network models that support adjustment of parameter sharing (Replication of [SePS](https://proceedings.mlr.press/v139/christianos21a/christianos21a.pdf) algorithm).
- `Policy`: Multi-agent policy models.
- `ConditionalVAE`: Conditional VAE model.


