A Reliable Representation with Bidirectional Transition Model for Visual Reinforcement Learning Generalization

2023-12-04

arXiv.org

AI Summary

Weak generalization of visual representations remains a core bottleneck in visual reinforcement learning (RL). To address this, we propose the Bidirectional Transfer (BiT) model—the first to introduce bidirectional temporal consistency into visual RL representation learning—inspired by human cognitive mechanisms and grounded in the reliability principle of “predictable future and reconstructible past.” BiT jointly optimizes contrastive learning and latent-state transition modeling to simultaneously enable forward dynamics prediction and backward causal reconstruction, supporting end-to-end pixel-level training. On the DeepMind Control benchmark, BiT achieves state-of-the-art generalization performance and sample efficiency. It further demonstrates strong cross-task transferability in robotic manipulation and CARLA autonomous driving simulation. Our key contributions are: (i) establishing bidirectional temporal consistency as a novel paradigm for representation reliability; and (ii) proposing the first explicit visual RL representation framework that models bidirectional temporal structure.

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.

Problem

• Learning robust visual representations for reinforcement learning
• Bidirectional prediction of environmental transitions
• Improving generalization and sample efficiency

Innovation

• Bidirectional Transition model for reliable representations
• Forward and backward environmental transition prediction
• Competitive generalization performance and sample efficiency