This repository is the offcial implementation of ToM2C, "ToM2C: Target-oriented Multi-agent Communication and Cooperation with Theory of Mind (ICLR 2022)" .
To install requirements:
pip install -r requirements.txt
All the environments have been included in the code, so there is no need to install Multi-sensor Multi-target Coverage(MSMTC) or MPE(Cooperative Navigation) additionally.
To train ToM2C in MSMTC
, run this command:
python main.py --env MSMTC-v3 --model ToM2C --workers 6 --norm-reward
To train ToM2C in CN
, run this command:
python main.py --env CN --model ToM2C --workers 6 --env-steps 10 --A2C-steps 10 --norm-reward
Note that the command above will load the default environment described in the paper. If you want to change the number of agents and targets, please refer to the num-agents
and num-targets
arguments.
After running the above command, you can run the following command respectively to do Communication Reduction
mentioned in the paper:
python main.py --env MSMTC-v3 --model ToM2C --workers 6 --norm-reward --train-comm --load-model-dir [trained_model_file_path]
The above command is for cpu training. If you want to train the model on GPU, try to add --gpu-id [cuda_device_id]
in the command. Note that this implementation does NOT support multi-gpu training.
After training, you can load the trained model and render its behavior by the following command.
In CN
:
python render_test.py --env CN --model ToM2C --render --env-steps 10 --load-model-dir [trained_model_file_path]
In MSMTC
:
python render_test.py --env MSMTC-v3 --model ToM2C --render --env-steps 20 --load-model-dir [trained_model_file_path]
If you found ToM2C useful, please consider citing:
@inproceedings{
wang2021tomc,
title={ToM2C: Target-oriented Multi-agent Communication and Cooperation with Theory of Mind},
author={Yuanfei Wang and Fangwei Zhong and Jing Xu and Yizhou Wang},
booktitle={International Conference on Learning Representations},
year={2022},
url={https://openreview.net/forum?id=M3tw78MH1Bk}
}
If you have any suggestion or questions, please get in touch at yuanfei_wang@pku.edu.cn or zfw@pku.edu.cn.