Towards Control-Centric Representations in Reinforcement Learning from Images
This work addresses a key problem in reinforcement learning for robotics and gaming by improving representation learning in sparse reward settings, though it is incremental as it builds on bisimulation-based approaches.
The paper tackled the challenge of extracting control-centric representations in image-based reinforcement learning by introducing ReBis, which integrates reward-free control information and uses a transformer with block-wise masking and combined losses to improve performance in sparse reward environments. Empirical results on Atari and DeepMind Control Suite benchmarks showed superior performance compared to existing methods.
Image-based Reinforcement Learning is a practical yet challenging task. A major hurdle lies in extracting control-centric representations while disregarding irrelevant information. While approaches that follow the bisimulation principle exhibit the potential in learning state representations to address this issue, they still grapple with the limited expressive capacity of latent dynamics and the inadaptability to sparse reward environments. To address these limitations, we introduce ReBis, which aims to capture control-centric information by integrating reward-free control information alongside reward-specific knowledge. ReBis utilizes a transformer architecture to implicitly model the dynamics and incorporates block-wise masking to eliminate spatiotemporal redundancy. Moreover, ReBis combines bisimulation-based loss with asymmetric reconstruction loss to prevent feature collapse in environments with sparse rewards. Empirical studies on two large benchmarks, including Atari games and DeepMind Control Suit, demonstrate that ReBis has superior performance compared to existing methods, proving its effectiveness.