# Learning Latent State Spaces for Planning through Reward Prediction

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

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 9 Dec 2019.

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

- [Latent State Design for World Models under Sufficiency Constraints](https://mlsanity.com/paper/2605.01694/) (2026)
- [Robot Motion Planning as Video Prediction: A Spatio-Temporal Neural Network-based Motion Planner](https://mlsanity.com/paper/2208.11287/) (2022)
- [Learning Tactile Insertion in the Real World](https://mlsanity.com/paper/2405.00383/) (2024)
- [Situationally-Aware Dynamics Learning](https://mlsanity.com/paper/2505.19574/) (2025)

## Abstract (arXiv)

Model-based reinforcement learning methods typically learn models for high-dimensional state spaces by aiming to reconstruct and predict the original observations. However, drawing inspiration from model-free reinforcement learning, we propose learning a latent dynamics model directly from rewards. In this work, we introduce a model-based planning framework which learns a latent reward prediction model and then plans in the latent state-space. The latent representation is learned exclusively from multi-step reward prediction which we show to be the only necessary information for successful planning. With this framework, we are able to benefit from the concise model-free representation, while still enjoying the data-efficiency of model-based algorithms. We demonstrate our framework in multi-pendulum and multi-cheetah environments where several pendulums or cheetahs are shown to the agent but only one of which produces rewards. In these environments, it is important for the agent to construct a concise latent representation to filter out irrelevant observations. We find that our method can successfully learn an accurate latent reward prediction model in the presence of the irrelevant information while existing model-based methods fail. Planning in the learned latent state-space shows strong performance and high sample efficiency over model-free and model-based baselines.