Dream to Generalize: Zero-Shot Model-Based Reinforcement Learning for Unseen Visual Distractions

TL;DR

This paper introduces a zero-shot model-based reinforcement learning method called Dream to Generalize (Dr. G) that improves robustness against visual distractions in high-dimensional vision-based control tasks by using dual contrastive learning and a recurrent inverse dynamics model.

Contribution

The paper proposes a novel self-supervised approach with dual contrastive learning and a recurrent inverse dynamics model to enhance zero-shot generalization in MBRL under visual distractions.

Findings

Dr. G outperforms prior methods by 117% on simple-to-complex backgrounds.

It achieves a 14% performance gain in environments with randomized natural video backgrounds.

The approach effectively captures task-relevant features despite visual distractions.

Abstract

Model-based reinforcement learning (MBRL) has been used to efficiently solve vision-based control tasks in highdimensional image observations. Although recent MBRL algorithms perform well in trained observations, they fail when faced with visual distractions in observations. These task-irrelevant distractions (e.g., clouds, shadows, and light) may be constantly present in real-world scenarios. In this study, we propose a novel self-supervised method, Dream to Generalize (Dr. G), for zero-shot MBRL. Dr. G trains its encoder and world model with dual contrastive learning which efficiently captures task-relevant features among multi-view data augmentations. We also introduce a recurrent state inverse dynamics model that helps the world model to better understand the temporal structure. The proposed methods can enhance the robustness of the world model against visual distractions. To evaluate the generalization performance, we first train Dr. G on simple backgrounds and then test it on complex natural video backgrounds in the DeepMind Control suite, and the randomizing environments in Robosuite. Dr. G yields a performance improvement of 117% and 14% over prior works, respectively. Our code is open-sourced and available at https://github.com/JeongsooHa/DrG.git