# Simplifying Model-based RL: Learning Representations, Latent-space Models, and Policies with One Objective

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

## Key facts

- Cited by 28 papers in the map (data through 14 Oct 2026).
- Builds on 70 of its 91 references that are in the map.
- First posted on arXiv 18 Sep 2022.

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

- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018)
- [Bootstrap your own latent: A new approach to self-supervised Learning](https://mlsanity.com/paper/2006.07733/) (2020)
- [Continuous control with deep reinforcement learning](https://mlsanity.com/paper/1509.02971/) (2015)
- [Addressing Function Approximation Error in Actor-Critic Methods](https://mlsanity.com/paper/1802.09477/) (2018)
- [High-Dimensional Continuous Control Using Generalized Advantage Estimation](https://mlsanity.com/paper/1506.02438/) (2015)
- [Offline Reinforcement Learning: Tutorial, Review, and Perspectives on Open Problems](https://mlsanity.com/paper/2005.01643/) (2020)
- [Dream to Control: Learning Behaviors by Latent Imagination](https://mlsanity.com/paper/1912.01603/) (2019)
- [Mastering Atari, Go, Chess and Shogi by Planning with a Learned Model](https://mlsanity.com/paper/1911.08265/) (2019)
- [Learning Latent Dynamics for Planning from Pixels](https://mlsanity.com/paper/1811.04551/) (2018)
- [Mastering Atari with Discrete World Models](https://mlsanity.com/paper/2010.02193/) (2020)
- [A Review of Uncertainty Quantification in Deep Learning: Techniques, Applications and Challenges](https://mlsanity.com/paper/2011.06225/) (2020)
- [Recurrent World Models Facilitate Policy Evolution](https://mlsanity.com/paper/1809.01999/) (2018)
- [Deep Reinforcement Learning in a Handful of Trials using Probabilistic Dynamics Models](https://mlsanity.com/paper/1805.12114/) (2018)
- [When to Trust Your Model: Model-Based Policy Optimization](https://mlsanity.com/paper/1906.08253/) (2019)
- [Loss Surfaces, Mode Connectivity, and Fast Ensembling of DNNs](https://mlsanity.com/paper/1802.10026/) (2018)
- [Model-Based Reinforcement Learning for Atari](https://mlsanity.com/paper/1903.00374/) (2019)
- [Learning Invariant Representations for Reinforcement Learning without Reconstruction](https://mlsanity.com/paper/2006.10742/) (2020)
- [Temporal Difference Learning for Model Predictive Control](https://mlsanity.com/paper/2203.04955/) (2022)
- [Mastering Visual Continuous Control: Improved Data-Augmented Reinforcement Learning](https://mlsanity.com/paper/2107.09645/) (2021)
- [Maximum a Posteriori Policy Optimisation](https://mlsanity.com/paper/1806.06920/) (2018)
- [Randomized Ensembled Double Q-Learning: Learning Fast Without a Model](https://mlsanity.com/paper/2101.05982/) (2021)
- [Stochastic Latent Actor-Critic: Deep Reinforcement Learning with a Latent Variable Model](https://mlsanity.com/paper/1907.00953/) (2019)
- [A Survey on Offline Reinforcement Learning: Taxonomy, Review, and Open Problems](https://mlsanity.com/paper/2203.01387/) (2022)
- [Model-Ensemble Trust-Region Policy Optimization](https://mlsanity.com/paper/1802.10592/) (2018)
- [Differentiable MPC for End-to-end Planning and Control](https://mlsanity.com/paper/1810.13400/) (2018)
- [Action-Conditional Video Prediction using Deep Networks in Atari Games](https://mlsanity.com/paper/1507.08750/) (2015)
- [Sample-Efficient Reinforcement Learning with Stochastic Ensemble Value Expansion](https://mlsanity.com/paper/1807.01675/) (2018)
- [Learning Continuous Control Policies by Stochastic Value Gradients](https://mlsanity.com/paper/1510.09142/) (2015)
- [Benchmarking Model-Based Reinforcement Learning](https://mlsanity.com/paper/1907.02057/) (2019)

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

- [Simple Hierarchical Planning with Diffusion](https://mlsanity.com/paper/2401.02644/) (2024)
- [Bridging State and History Representations: Understanding Self-Predictive RL](https://mlsanity.com/paper/2401.08898/) (2024)
- [Seizing Serendipity: Exploiting the Value of Past Success in Off-Policy Actor-Critic](https://mlsanity.com/paper/2306.02865/) (2023)
- [LightZero: A Unified Benchmark for Monte Carlo Tree Search in General Sequential Decision Scenarios](https://mlsanity.com/paper/2310.08348/) (2023)
- [Simplified Temporal Consistency Reinforcement Learning](https://mlsanity.com/paper/2306.09466/) (2023)
- [Inference via Interpolation: Contrastive Representations Provably Enable Planning and Inference](https://mlsanity.com/paper/2403.04082/) (2024)
- [MAD-TD: Model-Augmented Data stabilizes High Update Ratio RL](https://mlsanity.com/paper/2410.08896/) (2024)
- [UniZero: Generalized and Efficient Planning with Scalable Latent World Models](https://mlsanity.com/paper/2406.10667/) (2024)
- [Do Transformer World Models Give Better Policy Gradients?](https://mlsanity.com/paper/2402.05290/) (2024)
- [A Unified View on Solving Objective Mismatch in Model-Based Reinforcement Learning](https://mlsanity.com/paper/2310.06253/) (2023)
- [WHALE: Towards Generalizable and Scalable World Models for Embodied Decision-making](https://mlsanity.com/paper/2411.05619/) (2024)
- [Discrete Codebook World Models for Continuous Control](https://mlsanity.com/paper/2503.00653/) (2025)
- [Robust Visual Imitation Learning with Inverse Dynamics Representations](https://mlsanity.com/paper/2310.14274/) (2023)
- [Can we hop in general? A discussion of benchmark selection and design using the Hopper environment](https://mlsanity.com/paper/2410.08870/) (2024)
- [How to Fine-tune the Model: Unified Model Shift and Model Bias Policy Optimization](https://mlsanity.com/paper/2309.12671/) (2023)
- [DiffTORI: Differentiable Trajectory Optimization for Deep Reinforcement and Imitation Learning](https://mlsanity.com/paper/2402.05421/) (2024)
- [From Observations to States: Latent Time Series Forecasting](https://mlsanity.com/paper/2602.00297/) (2026)
- [First Order Model-Based RL through Decoupled Backpropagation](https://mlsanity.com/paper/2509.00215/) (2025)
- [Stealing That Free Lunch: Exposing the Limits of Dyna-Style Reinforcement Learning](https://mlsanity.com/paper/2412.14312/) (2024)
- [On the Role of Computation in Reinforcement Learning](https://mlsanity.com/paper/2602.05999/) (2026)
- [Navigating Autonomous Vehicle on Unmarked Roads with Diffusion-Based Motion Prediction and Active Inference](https://mlsanity.com/paper/2406.00211/) (2024)
- [Dynamical-VAE-based Hindsight to Learn the Causal Dynamics of Factored-POMDPs](https://mlsanity.com/paper/2411.07832/) (2024)
- [Calibrated Value-Aware Model Learning with Probabilistic Environment Models](https://mlsanity.com/paper/2505.22772/) (2025)
- [λ-models: Effective Decision-Aware Reinforcement Learning with Latent Models](https://mlsanity.com/paper/2306.17366/) (2023)
- [Latent-Predictive Empowerment: Measuring Empowerment without a Simulator](https://mlsanity.com/paper/2410.11155/) (2024)
- [A Forensic Analysis of Synthetic Data in RL: Diagnosing and Solving Algorithmic Failures in Model-Based Policy Optimization](https://mlsanity.com/paper/2510.01457/) (2025)
- [Representation Learning Enables Scalable Multitask Deep Reinforcement Learning](https://mlsanity.com/paper/2606.05555/) (2026)
- [WorldTS: World Modeling for Multimodal Covariate-aware Time Series Forecasting](https://mlsanity.com/paper/2609.31162/) (2026)

## Abstract (arXiv)

While reinforcement learning (RL) methods that learn an internal model of the environment have the potential to be more sample efficient than their model-free counterparts, learning to model raw observations from high dimensional sensors can be challenging. Prior work has addressed this challenge by learning low-dimensional representation of observations through auxiliary objectives, such as reconstruction or value prediction. However, the alignment between these auxiliary objectives and the RL objective is often unclear. In this work, we propose a single objective which jointly optimizes a latent-space model and policy to achieve high returns while remaining self-consistent. This objective is a lower bound on expected returns. Unlike prior bounds for model-based RL on policy exploration or model guarantees, our bound is directly on the overall RL objective. We demonstrate that the resulting algorithm matches or improves the sample-efficiency of the best prior model-based and model-free RL methods. While sample efficient methods typically are computationally demanding, our method attains the performance of SAC in about 50% less wall-clock time.