# Improving Model-Based Reinforcement Learning with Internal State Representations through Self-Supervision

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

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 12 of its 21 references that are in the map.
- First posted on arXiv 10 Feb 2021.

## Builds on (most cited first, 12 of 12)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Asynchronous Methods for Deep Reinforcement Learning](https://mlsanity.com/paper/1602.01783/) (2016)
- [OpenAI Gym](https://mlsanity.com/paper/1606.01540/) (2016)
- [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 Chess and Shogi by Self-Play with a General Reinforcement Learning Algorithm](https://mlsanity.com/paper/1712.01815/) (2017)
- [Visual Foresight: Model-Based Deep Reinforcement Learning for Vision-Based Robotic Control](https://mlsanity.com/paper/1812.00568/) (2018)
- [Deep Successor Reinforcement Learning](https://mlsanity.com/paper/1606.02396/) (2016)
- [Contextual Imagined Goals for Self-Supervised Robotic Learning](https://mlsanity.com/paper/1910.11670/) (2019)
- [Universal Planning Networks](https://mlsanity.com/paper/1804.00645/) (2018)
- [An Online Learning Approach to Model Predictive Control](https://mlsanity.com/paper/1902.08967/) (2019)
- [Curious Meta-Controller: Adaptive Alternation between Model-Based and Model-Free Control in Deep Reinforcement Learning](https://mlsanity.com/paper/1905.01718/) (2019)

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

- [Demystifying MuZero Planning: Interpreting the Learned Model](https://mlsanity.com/paper/2411.04580/) (2024)
- [Model Predictive Control with Self-supervised Representation Learning](https://mlsanity.com/paper/2304.07219/) (2023)

## Abstract (arXiv)

Using a model of the environment, reinforcement learning agents can plan their future moves and achieve superhuman performance in board games like Chess, Shogi, and Go, while remaining relatively sample-efficient. As demonstrated by the MuZero Algorithm, the environment model can even be learned dynamically, generalizing the agent to many more tasks while at the same time achieving state-of-the-art performance. Notably, MuZero uses internal state representations derived from real environment states for its predictions. In this paper, we bind the model's predicted internal state representation to the environment state via two additional terms: a reconstruction model loss and a simpler consistency loss, both of which work independently and unsupervised, acting as constraints to stabilize the learning process. Our experiments show that this new integration of reconstruction model loss and simpler consistency loss provide a significant performance increase in OpenAI Gym environments. Our modifications also enable self-supervised pretraining for MuZero, so the algorithm can learn about environment dynamics before a goal is made available.