Return-Critic: Bridging Goal Discrepancy for Efficient Visual Reinforcement Learning

Sample inefficiency remains a challenge in pixel-based visual reinforcement learning (RL), primarily due to ineffective state representation learning. While recent advances employ auxiliary tasks to improve representation learning, their representation goals (e.g., mask reconstruction, state prediction) are misaligned with the ultimate RL goal of maximizing return, constraining further improvements in representation quality. To achieve efficient visual reinforcement learning, we propose Return-Critic (RC), an auxiliary framework that bridges goal discrepancy by return prediction. RC samples partial frames from an episode, processes them through a shared visual encoder, and employs a lightweight Transformer to predict the episode's return, forcing the encoder to learn return-relevant representation. The attention weights naturally highlight important frames, enabling a key function for prioritized learning. Extensive experiments on both online (DMControl) and offline (V-D4RL) benchmarks demonstrate that RC significantly enhances the sample efficiency, particularly achieving 68% performance boost on average across nine challenging tasks from DMControl.

When teaching artificial intelligence to perform complex tasks by processing camera feeds, a major challenge is that the learning process is often extremely slow. This happens because video feeds contain too many irrelevant details, and traditional training methods often assign extra tasks like restoring missing image parts that fail to perfectly align with the ultimate goal of achieving the highest possible score. To solve this problem, we introduce a new method called Return Critic. Instead of getting distracted by immediate or irrelevant visual details, our method trains the artificial intelligence to evaluate a sequence of frames and directly predict the total final score of the entire task. This explicitly forces the artificial intelligence to ignore background noise and focus entirely on the visual features that actually determine success or failure. Furthermore, our method automatically identifies the most critical moments in a task, such as a robot jumping or falling, and prioritizes learning from those specific experiences. When tested on various complex simulated robotic tasks, our approach significantly accelerated the learning speed and boosted the final performance simply by keeping the artificial intelligence strictly focused on the ultimate goal.