# Towards Control-Centric Representations in Reinforcement Learning from Images

arXiv: https://arxiv.org/abs/2310.16655 · HTML: https://mlsanity.com/paper/2310.16655/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 43 of its 62 references that are in the map.
- First posted on arXiv 25 Oct 2023.

## Builds on (most cited first, 30 of 43)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Representation Learning with Contrastive Predictive Coding](https://mlsanity.com/paper/1807.03748/) (2018)
- [Masked Autoencoders Are Scalable Vision Learners](https://mlsanity.com/paper/2111.06377/) (2021)
- [Grad-CAM: Visual Explanations from Deep Networks via Gradient-based Localization](https://mlsanity.com/paper/1610.02391/) (2016)
- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018)
- [BEiT: BERT Pre-Training of Image Transformers](https://mlsanity.com/paper/2106.08254/) (2021)
- [The Kinetics Human Action Video Dataset](https://mlsanity.com/paper/1705.06950/) (2017)
- [Dream to Control: Learning Behaviors by Latent Imagination](https://mlsanity.com/paper/1912.01603/) (2019)
- [Learning Latent Dynamics for Planning from Pixels](https://mlsanity.com/paper/1811.04551/) (2018)
- [Mastering Diverse Domains through World Models](https://mlsanity.com/paper/2301.04104/) (2023)
- [Mastering Atari with Discrete World Models](https://mlsanity.com/paper/2010.02193/) (2020)
- [DeepMind Control Suite](https://mlsanity.com/paper/1801.00690/) (2018)
- [Deep Reinforcement Learning at the Edge of the Statistical Precipice](https://mlsanity.com/paper/2108.13264/) (2021)
- [Rainbow: Combining Improvements in Deep Reinforcement Learning](https://mlsanity.com/paper/1710.02298/) (2017)
- [CURL: Contrastive Unsupervised Representations for Reinforcement Learning](https://mlsanity.com/paper/2004.04136/) (2020)
- [Image Augmentation Is All You Need: Regularizing Deep Reinforcement Learning from Pixels](https://mlsanity.com/paper/2004.13649/) (2020)
- [Masked Feature Prediction for Self-Supervised Visual Pre-Training](https://mlsanity.com/paper/2112.09133/) (2021)
- [Learning Invariant Representations for Reinforcement Learning without Reconstruction](https://mlsanity.com/paper/2006.10742/) (2020)
- [Masked Autoencoders As Spatiotemporal Learners](https://mlsanity.com/paper/2205.09113/) (2022)
- [Reinforcement Learning with Unsupervised Auxiliary Tasks](https://mlsanity.com/paper/1611.05397/) (2016)
- [Transformers are Sample-Efficient World Models](https://mlsanity.com/paper/2209.00588/) (2022)
- [Data-Efficient Reinforcement Learning with Self-Predictive Representations](https://mlsanity.com/paper/2007.05929/) (2020)
- [Improving Sample Efficiency in Model-Free Reinforcement Learning from Images](https://mlsanity.com/paper/1910.01741/) (2019)
- [Decoupling Representation Learning from Reinforcement Learning](https://mlsanity.com/paper/2009.08319/) (2020)
- [Understanding self-supervised Learning Dynamics without Contrastive Pairs](https://mlsanity.com/paper/2102.06810/) (2021)
- [Stochastic Latent Actor-Critic: Deep Reinforcement Learning with a Latent Variable Model](https://mlsanity.com/paper/1907.00953/) (2019)
- [Reinforcement Learning with Prototypical Representations](https://mlsanity.com/paper/2102.11271/) (2021)
- [Masked World Models for Visual Control](https://mlsanity.com/paper/2206.14244/) (2022)
- [TransDreamer: Reinforcement Learning with Transformer World Models](https://mlsanity.com/paper/2202.09481/) (2022)
- [Scalable methods for computing state similarity in deterministic Markov Decision Processes](https://mlsanity.com/paper/1911.09291/) (2019)

## Abstract (arXiv)

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.