[NeurIPS 2024] MADiff: Offline Multi-agent Learning with Diffusion Models

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

We omit the standard deviation of the results for brevity. The full results can be found in our paper.

The peformances on MPE datasets released in OMAR paper. The results are averaged over 5 random seeds.

The peformances on MA-Mujoco datasets released in off-the-grid MARL benchmark. The results are averaged over 5 random seeds.

The peformances on SMAC datasets released in off-the-grid MARL benchmark. The results are averaged over 5 random seeds.

* MADiff-C is not meant to be a fair comparison with baseline methods but to show if MADiff-D fills the gap for coordination without global information.

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

We use the MPE dataset from OMAR. The dataset download link and instructions can be found in OMAR's repo. Since their BaiduPan download links might be inconvenient for non-Chinese users, we maintain an anonymous mirror repo in OSF for acquiring the dataset.

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

Install MPE environment:

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

1. Install MA-Mujoco:

   pip install -e third_party/multiagent_mujoco

2. We use the MA-Mujoco dataset from off-the-grid 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 the Huggingface repo.

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

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

   pip install -r ./third_party/og-marl/install_environments/requirements/mamujoco.txt
   pip install -e ./third_party/og-marl
   python scripts/transform_og_marl_dataset.py --env_name mamujoco --map_name <map> --quality <dataset>

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

2. Install SMAC:

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

3. We use the SMAC dataset from off-the-grid 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 the Huggingface repo.

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

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

   pip install -r ./third_party/og-marl/install_environments/requirements/smacv1.txt
   pip install -e ./third_party/og-marl
   python scripts/transform_og_marl_dataset.py --env_name smac --map_name <map> --quality <dataset>

Current OG-MARL datasets use the Flashbax Vault format and require Python 3.10. Do not replace third_party/og-marl or install Flashbax into the Python 3.8 MADiff environment. Use the separate environment-vault.yml environment and the Vault converter described in docs/vault_dataset_conversion.md.

This checkout reads datasets from /home/lotus/lotus/lhh/My_DoF/DoF_Trajectory/diffuser/datasets/data by default. Set MADIFF_DATA_ROOT to override that root on another machine. Training artifacts, checkpoints, TensorBoard events, and evaluation results are written below this project's logs/ directory.

Training progress bars are enabled by default and show the dataset, seed, epoch, global step, and latest logged loss. Add show_progress: false under an experiment's constants section to disable them.

To start training, run the following commands

# 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

python run_experiment.py -e exp_specs/eval_inv.yaml

@article{zhu2023madiff,
  title={MADiff: Offline Multi-agent Learning with Diffusion Models},
  author={Zhu, Zhengbang and Liu, Minghuan and Mao, Liyuan and Kang, Bingyi and Xu, Minkai and Yu, Yong and Ermon, Stefano and Zhang, Weinan},
  journal={arXiv preprint arXiv:2305.17330},
  year={2023}
}

The codebase is built upon decision-diffuser repo and ILSwiss.