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

2023-06-26

AAAI Conference on Artificial Intelligence

AI Summary

To address the degradation of model-based reinforcement learning (MBRL) generalization under high-dimensional visual observations corrupted by clouds, shadows, and illumination variations, this paper proposes Dr. G—a zero-shot model-based RL framework. Our approach tackles this challenge through three key contributions: (1) a novel dual-contrastive self-supervised learning mechanism that disentangles and encodes task-relevant features from multi-view augmented data; (2) recurrent state-wise inverse dynamics modeling to enhance the world model’s temporal causal understanding; and (3) zero-shot cross-background transfer without fine-tuning. Evaluated on DeepMind Control (with complex video backgrounds) and Robosuite (with randomized environments), Dr. G achieves performance gains of 117% and 14%, respectively, over state-of-the-art methods. The implementation is publicly available.

Technology Category

• Machine Learning: Reinforcement Learning
• Intelligent Robots: Multimodal Perception & Sensor Fusion
• Computer Vision: Diffusion Models for Vision

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking

Abstract

Model-based reinforcement learning (MBRL) has been used to efficiently solve vision-based control tasks in high-dimensional 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

Problem

• Enhance MBRL robustness against visual distractions
• Generalize to unseen visual environments zero-shot
• Improve performance in real-world visual variations

Innovation

• Self-supervised zero-shot MBRL method
• Dual contrastive learning for feature extraction
• Recurrent state inverse dynamics model