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 --num-agents 7 --num-targets 7 --workers 12 --env-steps 10 --A2C-steps 10 --norm-reward --gpu-id 0
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.