# MADiff: Offline Multi-agent Learning with Diffusion Models

![Python 3.8](https://img.shields.io/badge/Python-3.8-blue)
![Code style](https://img.shields.io/badge/code%20style-black-000000.svg)
![MIT](https://img.shields.io/badge/license-MIT-blue)
[![arXiv](https://img.shields.io/badge/arXiv-Paper-<COLOR>.svg)](https://arxiv.org/abs/2305.17330)

This is the official implementation of "MADiff: Offline Multi-agent Learning with Diffusion Models".

![MADiff](/assets/images/madiff.png)

## Setup

### Installation

```bash
sudo apt-get update
sudo apt-get install libssl-dev libcurl4-openssl-dev swig
conda create -n madiff python=3.8
conda activate madiff
pip install torch==1.12.1+cu113 --extra-index-url https://download.pytorch.org/whl/cu113
pip install -r requirements.txt
```

### Setup MPE

We use the MPE dataset from [OMAR](https://github.com/ling-pan/OMAR). The dataset download link and instructions can be found in OMAR's [repo](https://github.com/ling-pan/OMAR). Since their BaiduPan download links might be inconvenient for non-Chinese users, we maintain a anonymous mirror [repo](https://osf.io/jxawh/?view_only=dd3264a695af4c03bffde0350b8e8c4a) in OSF for acquiring the dataset.

The downloaded dataset should be placed under `diffuser/datasets/data/mpe`.

Install MPE environment:

```bash
pip install -e third_party/multiagent-particle-envs
pip install -e third_party/ddpg-agent
```

### Setup MA-Mujoco

1. Install MA-Mujoco:

    ```bash
    pip install -e third_party/multiagent_mujoco
    ```

2. We use the MA-Mujoco dataset from [off-the-grid MARL](https://sites.google.com/view/og-marl). We preprocess the dataset to concatenate trajectories to full episodes and save them as `.npy` files for easier loading. The original dataset can be downloaded from links below.

+ [2halfcheetah](https://s3.kao.instadeep.io/offline-marl-dataset/2halfcheetah.zip)

+ [2ant](https://s3.kao.instadeep.io/offline-marl-dataset/2ant.zip)

+ [4ant](https://s3.kao.instadeep.io/offline-marl-dataset/4ant.zip)

    The downloaded dataset should be placed under `diffuser/datasets/data/mamujoco`.

3. Install off-the-grid MARL and transform the original dataset.

    ```bash
    pip install -r ./third_party/og_marl/requirements.txt
    pip install -e ./third_party/og_marl
    python scripts/transform_og_marl_dataset.py --env_name mamujoco --map_name <map> --quality <dataset>
    ```

### Setup SMAC

1. Run `scripts/smac.sh` to install *StarCraftII*.

2. Install SMAC:

    ```bash
    pip install git+https://github.com/oxwhirl/smac.git
    ```

3. We use the SMAC dataset from [off-the-grid MARL](https://sites.google.com/view/og-marl). We preprocess the dataset to concatenate trajectories to full episodes and save them as `.npy` files for easier loading. The original dataset can be downloaded from links below.

+ [3m](https://s3.kao.instadeep.io/offline-marl-dataset/3m.zip)

+ [2s3z](https://s3.kao.instadeep.io/offline-marl-dataset/2s3z.zip)

+ [5m_vs_6m](https://s3.kao.instadeep.io/offline-marl-dataset/5m_vs_6m.zip)

+ [8m](https://s3.kao.instadeep.io/offline-marl-dataset/8m.zip)

    The downloaded dataset should be placed under `diffuser/datasets/data/smac`.

4. Install off-the-grid MARL and transform the original dataset.

    ```bash
    pip install -r ./third_party/og_marl/requirements.txt
    pip install -e ./third_party/og_marl
    python scripts/transform_og_marl_dataset.py --env_name smac --map_name <map> --quality <dataset>
    ```

## Training and Evaluation
To start training, run the following commands

```bash
# multi-agent particle environment
python run_experiment.py -e exp_specs/mpe/<task>/mad_mpe_<task>_attn_<dataset>.yaml  # CTCE
python run_experiment.py -e exp_specs/mpe/<task>/mad_mpe_<task>_ctde_<dataset>.yaml  # CTDE
# ma-mujoco
python run_experiment.py -e exp_specs/mamujoco/<task>/mad_mamujoco_<task>_attn_<dataset>_history.yaml  # CTCE
python run_experiment.py -e exp_specs/mamujoco/<task>/mad_mamujoco_<task>_ctde_<dataset>_history.yaml  # CTDE
# smac
python run_experiment.py -e exp_specs/smac/<map>/mad_smac_<map>_attn_<dataset>_history.yaml  # CTCE
python run_experiment.py -e exp_specs/smac/<map>/mad_smac_<map>_ctde_<dataset>_history.yaml  # CTDE
```

To evaluate the trained model, first replace the `log_dir` with those need to be evaluated in `exp_specs/eval_inv.yaml` and run
```bash
python run_experiment.py -e exp_specs/eval_inv.yaml
```

## Acknowledgements

The codebase is built upon [decision-diffuser repo](https://github.com/anuragajay/decision-diffuser) and [ILSwiss](https://github.com/Ericonaldo/ILSwiss).
