# Control-Aware Representations for Model-based Reinforcement Learning

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

## Key facts

- Cited by 8 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 24 Jun 2020.

## Built on by (most cited first, 8 of 8)

- [Temporal Predictive Coding For Model-Based Planning In Latent Space](https://mlsanity.com/paper/2106.07156/) (2021)
- [Value Gradient weighted Model-Based Reinforcement Learning](https://mlsanity.com/paper/2204.01464/) (2022)
- [Towards Deeper Deep Reinforcement Learning with Spectral Normalization](https://mlsanity.com/paper/2106.01151/) (2021)
- [Deciding What to Model: Value-Equivalent Sampling for Reinforcement Learning](https://mlsanity.com/paper/2206.02072/) (2022)
- [Inference via Interpolation: Contrastive Representations Provably Enable Planning and Inference](https://mlsanity.com/paper/2403.04082/) (2024)
- [On Rate-Distortion Theory in Capacity-Limited Cognition & Reinforcement Learning](https://mlsanity.com/paper/2210.16877/) (2022)
- [Bayesian Reinforcement Learning with Limited Cognitive Load](https://mlsanity.com/paper/2305.03263/) (2023)
- [Robust Robotic Control from Pixels using Contrastive Recurrent State-Space Models](https://mlsanity.com/paper/2112.01163/) (2021)

## Abstract (arXiv)

A major challenge in modern reinforcement learning (RL) is efficient control of dynamical systems from high-dimensional sensory observations. Learning controllable embedding (LCE) is a promising approach that addresses this challenge by embedding the observations into a lower-dimensional latent space, estimating the latent dynamics, and utilizing it to perform control in the latent space. Two important questions in this area are how to learn a representation that is amenable to the control problem at hand, and how to achieve an end-to-end framework for representation learning and control. In this paper, we take a few steps towards addressing these questions. We first formulate a LCE model to learn representations that are suitable to be used by a policy iteration style algorithm in the latent space. We call this model control-aware representation learning (CARL). We derive a loss function for CARL that has close connection to the prediction, consistency, and curvature (PCC) principle for representation learning. We derive three implementations of CARL. In the offline implementation, we replace the locally-linear control algorithm (e.g.,~iLQR) used by the existing LCE methods with a RL algorithm, namely model-based soft actor-critic, and show that it results in significant improvement. In online CARL, we interleave representation learning and control, and demonstrate further gain in performance. Finally, we propose value-guided CARL, a variation in which we optimize a weighted version of the CARL loss function, where the weights depend on the TD-error of the current policy. We evaluate the proposed algorithms by extensive experiments on benchmark tasks and compare them with several LCE baselines.