Maximum Manifold Capacity Representations in State Representation Learning

TL;DR

This paper introduces a novel state representation learning method that integrates Maximum Manifold Capacity Representation (MMCR) with existing SSL techniques, improving class separability and manifold consistency while reducing computational costs.

Contribution

It presents an innovative integration of MMCR into SSL methods with a regularization strategy, and extends DIM-UA with a nuclear norm loss for robust manifold consistency.

Findings

Improved F1 score to 78% on Atari RAM data.

Enhanced DIM-UA performance with same encoding dimensions.

Significant gains with SimCLR and Barlow Twins implementations.

Abstract

The expanding research on manifold-based self-supervised learning (SSL) builds on the manifold hypothesis, which suggests that the inherent complexity of high-dimensional data can be unraveled through lower-dimensional manifold embeddings. Capitalizing on this, DeepInfomax with an unbalanced atlas (DIM-UA) has emerged as a powerful tool and yielded impressive results for state representations in reinforcement learning. Meanwhile, Maximum Manifold Capacity Representation (MMCR) presents a new frontier for SSL by optimizing class separability via manifold compression. However, MMCR demands extensive input views, resulting in significant computational costs and protracted pre-training durations. Bridging this gap, we present an innovative integration of MMCR into existing SSL methods, incorporating a discerning regularization strategy that enhances the lower bound of mutual information. We also propose a novel state representation learning method extending DIM-UA, embedding a nuclear norm loss to enforce manifold consistency robustly. On experimentation with the Atari Annotated RAM Interface, our method improves DIM-UA significantly with the same number of target encoding dimensions. The mean F1 score averaged over categories is 78% compared to 75% of DIM-UA. There are also compelling gains when implementing SimCLR and Barlow Twins. This supports our SSL innovation as a paradigm shift, enabling more nuanced high-dimensional data representations.