{"type":"ScholarlyArticle","title":"Don’t Touch What Matters: Task-Aware Lipschitz Data Augmentation for Visual Reinforcement Learning","authors":[{"name":"Zhecheng Yuan","slug":"zhecheng-yuan","given":"Zhecheng","family":"Yuan","affiliation":"Tsinghua University"},{"name":"Guozheng Ma","slug":"guozheng-ma","given":"Guozheng","family":"Ma","affiliation":"Tsinghua University"},{"name":"Yao Mu","slug":"yao-mu","given":"Yao","family":"Mu","affiliation":"Hongkong University"},{"name":"Bo Xia","slug":"bo-xia","given":"Bo","family":"Xia","affiliation":"Tsinghua University"},{"name":"Bo Yuan","slug":"bo-yuan","given":"Bo","family":"Yuan","affiliation":"Tsinghua University"},{"name":"Xueqian Wang","slug":"xueqian-wang","given":"Xueqian","family":"Wang","affiliation":"Tsinghua University"},{"name":"Ping Luo","slug":"ping-luo","given":"Ping","family":"Luo","affiliation":"Hongkong University"},{"name":"Huazhe Xu","slug":"huazhe-xu","given":"Huazhe","family":"Xu","affiliation":"Stanford University"}],"abstract":"One of the key challenges in visual Reinforcement Learning (RL) is to learn policies that can generalize to unseen environments. Recently, data augmentation techniques aiming at enhancing data diversity have demonstrated proven performance in improving the generalization ability of learned policies. However, due to the sensitivity of RL training, naively applying data augmentation, which transforms each pixel in a task-agnostic manner, may suffer from instability and damage the sample efficiency, thus further exacerbating the generalization performance. At the heart of this phenomenon is the diverged action distribution and high-variance value estimation in the face of augmented images. To alleviate this issue, we propose Task-aware Lipschitz Data Augmentation (TLDA) for visual RL, which  explicitly identifies the task-correlated pixels with large Lipschitz constants, and only augments the task-irrelevant pixels for stability. We verify the effectiveness of our approach on DeepMind Control suite, CARLA and DeepMind Manipulation tasks. The extensive empirical results show that TLDA improves both sample efficiency and generalization; it outperforms previous state-of-the-art methods across 3 different visual control benchmarks.","doi":"10.24963/ijcai.2022/514","url":"https://ijcai.org/proceedings/2022/514","pdf":"https://ijcai.org/proceedings/2022/0514.pdf","markdown":"https://ijcai.org/proceedings/2022/514.md","bibtex":"https://ijcai.org/proceedings/2022/bibtex/514","volume":{"slug":"2022","title":"Proceedings of the Thirty-First International Joint Conference on Artificial Intelligence","short_title":"IJCAI 2022","url":"https://ijcai.org/proceedings/2022/","year":2022,"location":"Vienna","dates":"23-29 July 2022","editors":"Luc De Raedt","isbn":"978-1-956792-00-3","publisher":"International Joint Conferences on Artificial Intelligence Organization","doi_prefix":"10.24963/ijcai.2022"},"section":"Machine Learning","track":"Main Track","number":514,"pages":{"first":3702,"last":3708},"keywords":["Machine Learning: Deep Reinforcement Learning","Robotics: Learning in Robotics","Machine Learning: Reinforcement Learning"],"videos":[{"label":"Video #1","url":"https://slideslive.com/38986705","length_sec":89,"provider":"slideslive"},{"label":"Video #2","url":"https://slideslive.com/38984848","length_sec":352,"provider":"slideslive"}],"references":[{"position":1,"raw":"[Cobbe et al., 2019] Karl Cobbe, Oleg Klimov, Chris Hesse, Taehoon Kim, and John Schulman. Quantifying general- ization in reinforcement learning. In International Con- ference on Machine Learning, pages 1282-1289. PMLR, 2019.","title":"Quantifying generalization in reinforcement learning","authors":"Karl Cobbe, Oleg Klimov, Chris Hesse, Taehoon Kim, John Schulman","year":"2019","venue":"International Conference on Machine Learning","doi":null,"url":null,"internal_url":null},{"position":2,"raw":"[Cobbe et al., 2020] Karl Cobbe, Chris Hesse, Jacob Hilton, and John Schulman. Leveraging procedural generation to benchmark reinforcement learning. In International con- ference on machine learning, pages 2048-2056. PMLR, 2020.","title":"Leveraging procedural generation to benchmark reinforcement learning","authors":"Karl Cobbe, Chris Hesse, Jacob Hilton, John Schulman","year":"2020","venue":"International conference on machine learning","doi":null,"url":null,"internal_url":null},{"position":3,"raw":"[Dosovitskiy et al., 2017] Alexey Dosovitskiy, German Ros, Felipe Codevilla, Antonio Lopez, and Vladlen Koltun. Carla: An open urban driving simulator. In Conference on robot learning, pages 1-16. PMLR, 2017.","title":"Carla: An open urban driving simulator","authors":"Alexey Dosovitskiy, German Ros, Felipe Codevilla, Antonio Lopez, Vladlen Koltun","year":"2017","venue":"Conference on robot learning","doi":null,"url":null,"internal_url":null},{"position":4,"raw":"[Fan et al., 2021] Linxi Fan, Guanzhi Wang, De-An Huang, Zhiding Yu, Li Fei-Fei, Yuke Zhu, and Anima Anandku- mar. Secant: Self-expert cloning for zero-shot generaliza- tion of visual policies. arXiv preprint arXiv:2106.09678, 2021.","title":"Secant: Self-expert cloning for zero-shot generalization of visual policies","authors":"Linxi Fan, Guanzhi Wang, De-An Huang, Zhiding Yu, Li Fei-Fei, Yuke Zhu","year":"2021","venue":"Anima Anandkumar","doi":null,"url":null,"internal_url":null},{"position":5,"raw":"[Ferns et al., 2011] Norm Ferns, Prakash Panangaden, and Doina Precup. Bisimulation metrics for continuous markov decision processes. SIAM Journal on Computing, 40(6):1662-1714, 2011.","title":"Bisimulation metrics for continuous markov decision processes","authors":"Ferns","year":"2011","venue":"SIAM Journal on Computing","doi":null,"url":null,"internal_url":null},{"position":6,"raw":"[Finlay et al., 2018] Chris Finlay, Jeff Calder, Bilal Abbasi, and Adam Oberman. Lipschitz regularized deep neural networks generalize and are adversarially robust. arXiv preprint arXiv:1808.09540, 2018.","title":"Lipschitz regularized deep neural networks generalize and are adversarially robust","authors":"Chris Finlay, Jeff Calder, Bilal Abbasi, Adam Oberman","year":"2018","venue":null,"doi":null,"url":null,"internal_url":null},{"position":7,"raw":"[Grill et al., 2020] Jean-Bastien Grill, Florian Strub, Flo- rent Altché, Corentin Tallec, Pierre H Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhao- han Daniel Guo, Mohammad Gheshlaghi Azar, et al. Boot- strap your own latent: A new approach to self-supervised learning. arXiv preprint arXiv:2006.07733, 2020.","title":"Bootstrap your own latent: A new approach to self-supervised learning","authors":"Jean-Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Elena Pierre H Richemond, Carl Buchatskaya, Bernardo Doersch, Zhaohan Avila Pires, Mohammad Daniel Guo, Gheshlaghi Azar","year":"2020","venue":null,"doi":null,"url":null,"internal_url":null},{"position":8,"raw":"[Hansen and Wang, 2021] Nicklas Hansen and Xiaolong Wang. Generalization in reinforcement learning by soft data augmentation. In 2021 IEEE International Confer- ence on Robotics and Automation (ICRA), pages 13611- 13617. IEEE, 2021.","title":"Generalization in reinforcement learning by soft data augmentation","authors":"Wang ; Hansen, Nicklas Hansen, Xiaolong Wang","year":"2021","venue":"IEEE International Conference on Robotics and Automation (ICRA)","doi":null,"url":null,"internal_url":null},{"position":9,"raw":"[Hansen et al., 2020] Nicklas Hansen, Rishabh Jangir, Yu Sun, Guillem Alenyà, Pieter Abbeel, Alexei A Efros, Lerrel Pinto, and Xiaolong Wang. Self-supervised policy adaptation during deployment. arXiv preprint arXiv:2007.04309, 2020.","title":"Self-supervised policy adaptation during deployment","authors":"Nicklas Hansen, Rishabh Jangir, Yu Sun, Guillem Alenyà, Pieter Abbeel, Alexei A Efros, Lerrel Pinto, Xiaolong Wang","year":"2020","venue":null,"doi":null,"url":null,"internal_url":null},{"position":10,"raw":"[Hansen et al., 2021] Nicklas Hansen, Hao Su, and Xiaolong Wang. Stabilizing deep q-learning with convnets and vi- sion transformers under data augmentation. arXiv e-prints, pages arXiv-2107, 2021.","title":"Stabilizing deep q-learning with convnets and vision transformers under data augmentation","authors":"Nicklas Hansen, Hao Su, Xiaolong Wang","year":"2021","venue":null,"doi":null,"url":null,"internal_url":null},{"position":11,"raw":"[Kalashnikov et al., 2018] Dmitry Kalashnikov, Alex Irpan, Peter Pastor, Julian Ibarz, Alexander Herzog, Eric Jang, Deirdre Quillen, Ethan Holly, Mrinal Kalakrishnan, Vin- cent Vanhoucke, et al. Scalable deep reinforcement learn- ing for vision-based robotic manipulation. In Conference on Robot Learning, pages 651-673. PMLR, 2018.","title":"Scalable deep reinforcement learning for vision-based robotic manipulation","authors":"Dmitry Kalashnikov, Alex Irpan, Peter Pastor, Julian Ibarz, Alexander Herzog, Eric Jang, Deirdre Quillen, Ethan Holly, Mrinal Kalakrishnan, Vincent Vanhoucke","year":"2018","venue":"Conference on Robot Learning","doi":null,"url":null,"internal_url":null},{"position":12,"raw":"[Kostrikov et al., 2020] Ilya Kostrikov, Denis Yarats, and Rob Fergus. Image augmentation is all you need: Reg- ularizing deep reinforcement learning from pixels. arXiv preprint arXiv:2004.13649, 2020.","title":"Image augmentation is all you need: Regularizing deep reinforcement learning from pixels","authors":"Kostrikov","year":"2020","venue":null,"doi":null,"url":null,"internal_url":null},{"position":13,"raw":"[Laskin et al., 2020] Misha Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Pieter Abbeel, and Aravind Srinivas. Reinforcement learning with augmented data. Advances in Neural Information Processing Systems, 33, 2020.","title":"Reinforcement learning with augmented data","authors":"Misha Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Pieter Abbeel, Aravind Srinivas","year":"2020","venue":"Advances in Neural Information Processing Systems","doi":null,"url":null,"internal_url":null},{"position":14,"raw":"[Mnih et al., 2015] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Marc G Belle- mare, Alex Graves, Martin Riedmiller, Andreas K Fidje- land, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. nature, 518(7540):529-533, 2015.","title":"Human-level control through deep reinforcement learning","authors":"Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Marc G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski","year":"2015","venue":"nature","doi":null,"url":null,"internal_url":null},{"position":15,"raw":"[Packer et al., 2018] Charles Packer, Katelyn Gao, Jernej Kos, Philipp Krähenbühl, Vladlen Koltun, and Dawn Song. Assessing generalization in deep reinforcement learning. arXiv preprint arXiv:1810.12282, 2018.","title":"Assessing generalization in deep reinforcement learning","authors":"Charles Packer, Katelyn Gao, Jernej Kos, Philipp Krähenbühl, Vladlen Koltun, Dawn Song","year":"2018","venue":null,"doi":null,"url":null,"internal_url":null},{"position":16,"raw":"[Pauli et al., 2021] Patricia Pauli, Anne Koch, Julian Berberich, Paul Kohler, and Frank Allgower. Training robust neural networks using lipschitz bounds. IEEE Control Systems Letters, 2021.","title":"Training robust neural networks using lipschitz bounds","authors":"Patricia Pauli, Anne Koch, Julian Berberich, Paul Kohler, Frank Allgower","year":"2021","venue":"IEEE Control Systems Letters","doi":null,"url":null,"internal_url":null},{"position":17,"raw":"[Raileanu et al., 2021] Roberta Raileanu, Maxwell Gold- stein, Denis Yarats, Ilya Kostrikov, and Rob Fergus. Au- tomatic data augmentation for generalization in reinforce- ment learning. Advances in Neural Information Process- ing Systems, 34, 2021.","title":"Automatic data augmentation for generalization in reinforcement learning","authors":"Roberta Raileanu, Maxwell Goldstein, Denis Yarats, Ilya Kostrikov, Rob Fergus","year":"2021","venue":"Advances in Neural Information Processing Systems","doi":null,"url":null,"internal_url":null},{"position":18,"raw":"[Shorten and Khoshgoftaar, 2019] Connor Shorten and Taghi M Khoshgoftaar. A survey on image data augmen- tation for deep learning. Journal of Big Data, 6(1):1-48, 2019.","title":"A survey on image data augmentation for deep learning","authors":"Connor Shorten, M Taghi, Khoshgoftaar","year":"2019","venue":"Journal of Big Data","doi":null,"url":null,"internal_url":null},{"position":19,"raw":"[Tobin et al., 2017] Josh Tobin, Rachel Fong, Alex Ray, Jonas Schneider, Wojciech Zaremba, and Pieter Abbeel. Domain randomization for transferring deep neural net- works from simulation to the real world. In IROS, pages 23-30. IEEE, 2017.","title":"Domain randomization for transferring deep neural networks from simulation to the real world","authors":"Josh Tobin, Rachel Fong, Alex Ray, Jonas Schneider, Wojciech Zaremba, Pieter Abbeel","year":"2017","venue":"IROS","doi":null,"url":null,"internal_url":null},{"position":20,"raw":"[Tunyasuvunakool et al., 2020] Saran Tunyasuvunakool, Al- istair Muldal, Yotam Doron, Siqi Liu, Steven Bohez, Josh Merel, Tom Erez, Timothy Lillicrap, Nicolas Heess, and Yuval Tassa. dm control: Software and tasks for continu- ous control. Software Impacts, 6:100022, 2020.","title":"dm control: Software and tasks for continuous control","authors":"Saran Tunyasuvunakool, Alistair Muldal, Yotam Doron, Siqi Liu, Steven Bohez, Josh Merel, Tom Erez, Timothy Lillicrap, Nicolas Heess, Yuval Tassa","year":"2020","venue":"Software Impacts","doi":null,"url":null,"internal_url":null},{"position":21,"raw":"[Zhang et al., 2018] Chiyuan Zhang, Oriol Vinyals, Remi Munos, and Samy Bengio. A study on overfitting in deep reinforcement learning. arXiv preprint arXiv:1804.06893, 2018.","title":"A study on overfitting in deep reinforcement learning","authors":"Zhang","year":"2018","venue":null,"doi":null,"url":null,"internal_url":null},{"position":22,"raw":"[Zhang et al., 2020] Amy Zhang, Rowan McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learn- ing invariant representations for reinforcement learning without reconstruction. arXiv preprint arXiv:2006.10742, 2020.","title":"Learning invariant representations for reinforcement learning without reconstruction","authors":"Amy Zhang, Rowan Mcallister, Roberto Calandra, Yarin Gal, Sergey Levine","year":"2020","venue":null,"doi":null,"url":null,"internal_url":null},{"position":23,"raw":"[Zhu et al., 2017] Yuke Zhu, Roozbeh Mottaghi, Eric Kolve, Joseph J Lim, Abhinav Gupta, Li Fei-Fei, and Ali Farhadi. Target-driven visual navigation in indoor scenes using deep reinforcement learning. In ICRA, pages 3357-3364. IEEE, 2017.","title":"Target-driven visual navigation in indoor scenes using deep reinforcement learning","authors":"Yuke Zhu, Roozbeh Mottaghi, Eric Kolve, Joseph J Lim, Abhinav Gupta, Li Fei-Fei, Ali Farhadi","year":"2017","venue":"ICRA","doi":null,"url":null,"internal_url":null}],"license_note":"Copyright International Joint Conferences on Artificial Intelligence."}