Learning Robust Representations via Bidirectional Transition for Visual Reinforcement Learning

TL;DR

This paper introduces a Bidirectional Transition (BiT) model that enhances visual reinforcement learning by predicting environmental transitions both forward and backward, leading to more reliable and generalizable representations across various control tasks.

Contribution

The paper proposes a novel BiT model that leverages bidirectional transition prediction to improve representation reliability and generalization in visual reinforcement learning.

Findings

Competitive performance on DeepMind Control suite

Effective in robotic manipulation tasks

Demonstrates wide applicability in simulation environments

Abstract

Visual reinforcement learning has proven effective in solving control tasks with high-dimensional observations. However, extracting reliable and generalizable representations from vision-based observations remains a central challenge. Inspired by the human thought process, when the representation extracted from the observation can predict the future and trace history, the representation is reliable and accurate in comprehending the environment. Based on this concept, we introduce a Bidirectional Transition (BiT) model, which leverages the ability to bidirectionally predict environmental transitions both forward and backward to extract reliable representations. Our model demonstrates competitive generalization performance and sample efficiency on two settings of the DeepMind Control suite. Additionally, we utilize robotic manipulation and CARLA simulators to demonstrate the wide applicability of our method.