# Shaping Belief States with Generative Environment Models for RL

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

## Key facts

- Cited by 66 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 21 Jun 2019.

## Built on by (most cited first, 30 of 66)

- [Towards Causal Representation Learning](https://mlsanity.com/paper/2102.11107/) (2021)
- [Recurrent Model-Free RL Can Be a Strong Baseline for Many POMDPs](https://mlsanity.com/paper/2110.05038/) (2021)
- [Deep Hierarchical Planning from Pixels](https://mlsanity.com/paper/2206.04114/) (2022)
- [DREAMWALKER: Mental Planning for Continuous Vision-Language Navigation](https://mlsanity.com/paper/2308.07498/) (2023)
- [DreamerPro: Reconstruction-Free Model-Based Reinforcement Learning with Prototypical Representations](https://mlsanity.com/paper/2110.14565/) (2021)
- [BYOL-Explore: Exploration by Bootstrapped Prediction](https://mlsanity.com/paper/2206.08332/) (2022)
- [CACTI: A Framework for Scalable Multi-Task Multi-Scene Visual Imitation Learning](https://mlsanity.com/paper/2212.05711/) (2022)
- [Scaling Instructable Agents Across Many Simulated Worlds](https://mlsanity.com/paper/2404.10179/) (2024)
- [Muesli: Combining Improvements in Policy Optimization](https://mlsanity.com/paper/2104.06159/) (2021)
- [Optimizing Sequential Experimental Design with Deep Reinforcement Learning](https://mlsanity.com/paper/2202.00821/) (2022)
- [Evaluating Long-Term Memory in 3D Mazes](https://mlsanity.com/paper/2210.13383/) (2022)
- [Tell me why! Explanations support learning relational and causal structure](https://mlsanity.com/paper/2112.03753/) (2021)
- [Procedural Generalization by Planning with Self-Supervised World Models](https://mlsanity.com/paper/2111.01587/) (2021)
- [Attention Approximates Sparse Distributed Memory](https://mlsanity.com/paper/2111.05498/) (2021)
- [INFOrmation Prioritization through EmPOWERment in Visual Model-Based RL](https://mlsanity.com/paper/2204.08585/) (2022)
- [Value Function Spaces: Skill-Centric State Abstractions for Long-Horizon Reasoning](https://mlsanity.com/paper/2111.03189/) (2021)
- [Towards Versatile Embodied Navigation](https://mlsanity.com/paper/2210.16822/) (2022)
- [Differentiable Spatial Planning using Transformers](https://mlsanity.com/paper/2112.01010/) (2021)
- [Structured World Belief for Reinforcement Learning in POMDP](https://mlsanity.com/paper/2107.08577/) (2021)
- [Selective Visual Representations Improve Convergence and Generalization for Embodied AI](https://mlsanity.com/paper/2311.04193/) (2023)
- [Neural Recursive Belief States in Multi-Agent Reinforcement Learning](https://mlsanity.com/paper/2102.02274/) (2021)
- [Unsupervised Representation Learning in Deep Reinforcement Learning: A Review](https://mlsanity.com/paper/2208.14226/) (2022)
- [The Benefits of Model-Based Generalization in Reinforcement Learning](https://mlsanity.com/paper/2211.02222/) (2022)
- [Flow-based Recurrent Belief State Learning for POMDPs](https://mlsanity.com/paper/2205.11051/) (2022)
- [Latent State Marginalization as a Low-cost Approach for Improving Exploration](https://mlsanity.com/paper/2210.00999/) (2022)
- [Would I have gotten that reward? Long-term credit assignment by counterfactual contribution analysis](https://mlsanity.com/paper/2306.16803/) (2023)
- [Environment Predictive Coding for Embodied Agents](https://mlsanity.com/paper/2102.02337/) (2021)
- [Temporally Abstract Partial Models](https://mlsanity.com/paper/2108.03213/) (2021)
- [Informed POMDP: Leveraging Additional Information in Model-Based RL](https://mlsanity.com/paper/2306.11488/) (2023)
- [Provable Representation with Efficient Planning for Partial Observable Reinforcement Learning](https://mlsanity.com/paper/2311.12244/) (2023)

## Abstract (arXiv)

When agents interact with a complex environment, they must form and maintain beliefs about the relevant aspects of that environment. We propose a way to efficiently train expressive generative models in complex environments. We show that a predictive algorithm with an expressive generative model can form stable belief-states in visually rich and dynamic 3D environments. More precisely, we show that the learned representation captures the layout of the environment as well as the position and orientation of the agent. Our experiments show that the model substantially improves data-efficiency on a number of reinforcement learning (RL) tasks compared with strong model-free baseline agents. We find that predicting multiple steps into the future (overshooting), in combination with an expressive generative model, is critical for stable representations to emerge. In practice, using expressive generative models in RL is computationally expensive and we propose a scheme to reduce this computational burden, allowing us to build agents that are competitive with model-free baselines.