Visual Reinforcement Learning Control With Instance-Reweighted Alignment and Instance-Dimension Uniformity.

Abstract

Visual reinforcement learning (VRL) has demonstrated remarkable capabilities in learning behaviors directly from intricate high-dimensional visual inputs. Despite these advancements, existing VRL methods still encounter obstacles such as complete collapse and dimensional collapse, resulting in representation degradation and dimensional redundancy. Contrastive learning, while helping to mitigate the complete collapse issue, is prone to the class collision dilemma. To tackle the aforementioned challenges, this article proposes a novel VRL control method with instance-reweighted alignment and instance-dimension uniformity (IAIU). In this VRL control method, the instance-reweighted alignment representation learning is introduced by minimizing the Kullback-Leibler (KL) divergence between the distributions of predicted next state representations and their weighted actual counterparts. By doing so, we aim to align state representations within the same semantic class, thereby effectively alleviating class collision. Meanwhile, an instance-dimension uniformity regularization mechanism is adopted to suppress the collapse phenomenon. This is realized by leveraging the Hilbert-Schmidt independence criterion (HSIC) and standard orthogonal constraint at the instance and dimension levels, respectively, ensuring the extraction of task-relevant state representations. In essence, IAIU's dual-strategy of alignment and uniformity not only addresses the critical issue of class collision but also guarantees uniformity with respect to both the instances and dimensions. Simulation results from the distracting control suite (DCS) benchmark demonstrate IAIU's superior performance, with substantial enhancements in both representational ability and policy efficacy. The code is available at https://github.com/anonymousforcode/IAIU.