Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Focus On What Matters: Separated Models For Visual-Based RL Generalization

Di Zhang, Bowen Lv, Hai Zhang, Feifan Yang, Junqiao Zhao, Hang Yu, Chang Huang, Hongtu Zhou, Chen Ye, Changjun Jiang

TL;DR

This work proposes SMG (Separated Models for Generalization), a novel approach that exploits image reconstruction for generalization and incorporates two additional consistency losses to guide the agent's focus toward task-relevant areas across different scenarios, thereby achieving free from overfitting.

Abstract

A primary challenge for visual-based Reinforcement Learning (RL) is to generalize effectively across unseen environments. Although previous studies have explored different auxiliary tasks to enhance generalization, few adopt image reconstruction due to concerns about exacerbating overfitting to task-irrelevant features during training. Perceiving the pre-eminence of image reconstruction in representation learning, we propose SMG (Separated Models for Generalization), a novel approach that exploits image reconstruction for generalization. SMG introduces two model branches to extract task-relevant and task-irrelevant representations separately from visual observations via cooperatively reconstruction. Built upon this architecture, we further emphasize the importance of task-relevant features for generalization. Specifically, SMG incorporates two additional consistency losses to guide the agent's focus toward task-relevant areas across different scenarios, thereby achieving free from overfitting. Extensive experiments in DMC demonstrate the SOTA performance of SMG in generalization, particularly excelling in video-background settings. Evaluations on robotic manipulation tasks further confirm the robustness of SMG in real-world applications.

Focus On What Matters: Separated Models For Visual-Based RL Generalization

TL;DR

This work proposes SMG (Separated Models for Generalization), a novel approach that exploits image reconstruction for generalization and incorporates two additional consistency losses to guide the agent's focus toward task-relevant areas across different scenarios, thereby achieving free from overfitting.

Abstract

A primary challenge for visual-based Reinforcement Learning (RL) is to generalize effectively across unseen environments. Although previous studies have explored different auxiliary tasks to enhance generalization, few adopt image reconstruction due to concerns about exacerbating overfitting to task-irrelevant features during training. Perceiving the pre-eminence of image reconstruction in representation learning, we propose SMG (Separated Models for Generalization), a novel approach that exploits image reconstruction for generalization. SMG introduces two model branches to extract task-relevant and task-irrelevant representations separately from visual observations via cooperatively reconstruction. Built upon this architecture, we further emphasize the importance of task-relevant features for generalization. Specifically, SMG incorporates two additional consistency losses to guide the agent's focus toward task-relevant areas across different scenarios, thereby achieving free from overfitting. Extensive experiments in DMC demonstrate the SOTA performance of SMG in generalization, particularly excelling in video-background settings. Evaluations on robotic manipulation tasks further confirm the robustness of SMG in real-world applications.

Paper Structure

This paper contains 29 sections, 12 equations, 11 figures, 9 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Background
  3. Approach
  4. What Matters in a Reinforcement Learning Task?
  5. Learning Task-Relevant Representations with Separated Models
  6. Separated Models and Reconstruction
  7. Additional Loss Terms
  8. Generalize Task-Relevant Representations with Separated Models
  9. Overall Objective
  10. Experimental Results
  11. Setup
  12. DMControl Results
  13. Robotic Manipulation Results
  14. Ablation Study
  15. Related Work
  16. ...and 14 more sections

Figures (11)

  • Figure 1: A robotic manipulation task explanation for task-relevant parts in the environment.
  • Figure 2: Architecture of SMG. One-way arrows represent different types of data flows with the same input. Two-way arrows represent different types of loss.
  • Figure 3: Visualizing the reconstruction process of SMG in different tasks (from top to bottom: walker-walk, cheetah-run, peg in box).
  • Figure 4: Two types of data augmentations using in SMG.
  • Figure 5: Example of training and testing observation for DMC-GB (walker-walk). (a) is the training observation. (b-c) indicates different degrees of color change; (d-e) replaces the background with random videos, with (e) additionally removing the floor and the walker's shadow.
  • ...and 6 more figures