Clustering-Driven State Embedding for Reinforcement Learning Under Visual Distractions

Visual reinforcement learning (VRL) aims to optimize policies by utilizing information derived from visual sensory inputs. To tackle the challenges of high dimensionality and sample efficiency, numerous studies have leveraged self-supervised learning to encode observations into latent state representations. Though these algorithms have shown promise, they may be susceptible to task-irrelevant distractions, such as changes in background, color, and camera angles. Thus, we propose the clustering-driven state embedding for reinforcement learning (RL) under visual distractions (CSE-RL) to learn more robust state representations. The student’s t distribution is first utilized to model the predicted clustering assignment. In this process, the reward centers are defined by hardening the target distribution. Subsequently, the state embeddings, rewards, and their respective centers are fed into the predicted clustering assignment function for clustering. To avoid trivial solutions, the Sinkhorn-Knopp algorithm is adopted to balance observations across clusters. The proposed CSE-RL has the capacity to enhance the comprehension of inherent relationships among different states, while also filtering out irrelevant or redundant information, resulting in the creation of more efficient representations that are less prone to visual distractions. Experimental results from the distracting control suite (DCS) indicate that CSE-RL achieves comparable or even better performance compared with several state-of-the-art VRL algorithms in handling challenging visual distractions. Notably, CSE-RL exhibits remarkable improvements of 31.82% and 18.18% on DCS over baselines based on data-regularized Q (DrQ) and DrQv2 frameworks, respectively. Integrating CSE-RL into offline RL offers promising avenues for future study.