Robust Representation learning by Clustering with Bisimulation Metrics for Visual Reinforcement Learning with Distractions
This is the code of paper Robust Representation learning by Clustering with Bisimulation Metrics for Visual Reinforcement Learning with Distractions. Qiyuan Liu, Qi Zhou, Rui Yang, Jie Wang*. AAAI 2023.
Since CBM is tested under Distracting Control Suite, you need to download the DAVIS 2017 dataset. Make sure to select the 2017 TrainVal - Images and Annotations (480p). The training images will be used as distracting backgrounds. Other dependencies is recorded in the file requirements.txt
.
pip install -r requirements.txt
For example, test the performance of CBM combined with DrQ-v2 on Cartpole Swingup under Distracting Control Suite "easy" distractions, run
python scripts/run.py configs/cbm_drqv2_easy/cartpole.json
If you find this code useful, please consider citing the following paper.
@article{liu2023robust,
title={Robust Representation Learning by Clustering with Bisimulation Metrics for Visual Reinforcement Learning with Distractions},
author={Liu, Qiyuan and Zhou, Qi and Yang, Rui and Wang, Jie},
journal={arXiv preprint arXiv:2302.12003},
year={2023}
}
We use the environment wrapper from Distracting Control Suite and DrQ-v2. Thanks for their contributions.