# Model-Invariant State Abstractions for Model-Based Reinforcement Learning

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

## Key facts

- Cited by 18 papers in the map (data through 14 Oct 2026).
- Builds on 27 of its 48 references that are in the map.
- First posted on arXiv 19 Feb 2021.

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

- [A Simple Framework for Contrastive Learning of Visual Representations](https://mlsanity.com/paper/2002.05709/) (2020)
- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018)
- [Invariant Risk Minimization](https://mlsanity.com/paper/1907.02893/) (2019)
- [OpenAI Gym](https://mlsanity.com/paper/1606.01540/) (2016)
- [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)
- [MOPO: Model-based Offline Policy Optimization](https://mlsanity.com/paper/2005.13239/) (2020)
- [MOReL : Model-Based Offline Reinforcement Learning](https://mlsanity.com/paper/2005.05951/) (2020)
- [Causal inference using invariant prediction: identification and confidence intervals](https://mlsanity.com/paper/1501.01332/) (2015)
- [Learning Invariant Representations for Reinforcement Learning without Reconstruction](https://mlsanity.com/paper/2006.10742/) (2020)
- [Neural Network Dynamics for Model-Based Deep Reinforcement Learning with Model-Free Fine-Tuning](https://mlsanity.com/paper/1708.02596/) (2017)
- [Information-Theoretic Considerations in Batch Reinforcement Learning](https://mlsanity.com/paper/1905.00360/) (2019)
- [Inductive Biases for Deep Learning of Higher-Level Cognition](https://mlsanity.com/paper/2011.15091/) (2020)
- [The Risks of Invariant Risk Minimization](https://mlsanity.com/paper/2010.05761/) (2020)
- [Stochastic Latent Actor-Critic: Deep Reinforcement Learning with a Latent Variable Model](https://mlsanity.com/paper/1907.00953/) (2019)
- [Model-Ensemble Trust-Region Policy Optimization](https://mlsanity.com/paper/1802.10592/) (2018)
- [Representation Learning via Invariant Causal Mechanisms](https://mlsanity.com/paper/2010.07922/) (2020)
- [Provably efficient RL with Rich Observations via Latent State Decoding](https://mlsanity.com/paper/1901.09018/) (2019)
- [Model-Based Value Estimation for Efficient Model-Free Reinforcement Learning](https://mlsanity.com/paper/1803.00101/) (2018)
- [Learning explanations that are hard to vary](https://mlsanity.com/paper/2009.00329/) (2020)
- [Kinematic State Abstraction and Provably Efficient Rich-Observation Reinforcement Learning](https://mlsanity.com/paper/1911.05815/) (2019)
- [Invariant Causal Prediction for Block MDPs](https://mlsanity.com/paper/2003.06016/) (2020)
- [Objective Mismatch in Model-based Reinforcement Learning](https://mlsanity.com/paper/2002.04523/) (2020)
- [Exploring Model-based Planning with Policy Networks](https://mlsanity.com/paper/1906.08649/) (2019)
- [On the model-based stochastic value gradient for continuous reinforcement learning](https://mlsanity.com/paper/2008.12775/) (2020)
- [Near-optimal Reinforcement Learning in Factored MDPs](https://mlsanity.com/paper/1403.3741/) (2014)
- [Off-policy evaluation for MDPs with unknown structure](https://mlsanity.com/paper/1502.03255/) (2015)

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

- [Causal Dynamics Learning for Task-Independent State Abstraction](https://mlsanity.com/paper/2206.13452/) (2022)
- [Dynamic Sparse Training for Deep Reinforcement Learning](https://mlsanity.com/paper/2106.04217/) (2021)
- [AdaRL: What, Where, and How to Adapt in Transfer Reinforcement Learning](https://mlsanity.com/paper/2107.02729/) (2021)
- [A Survey on Causal Reinforcement Learning](https://mlsanity.com/paper/2302.05209/) (2023)
- [Generalizing Goal-Conditioned Reinforcement Learning with Variational Causal Reasoning](https://mlsanity.com/paper/2207.09081/) (2022)
- [The Essential Role of Causality in Foundation World Models for Embodied AI](https://mlsanity.com/paper/2402.06665/) (2024)
- [Learning Representations for Pixel-based Control: What Matters and Why?](https://mlsanity.com/paper/2111.07775/) (2021)
- [A Step Toward World Models: A Survey on Robotic Manipulation](https://mlsanity.com/paper/2511.02097/) (2025)
- [Causal Semantic Communication for Digital Twins: A Generalizable Imitation Learning Approach](https://mlsanity.com/paper/2304.12502/) (2023)
- [Provably Efficient Causal Model-Based Reinforcement Learning for Systematic Generalization](https://mlsanity.com/paper/2202.06545/) (2022)
- [Fine-Grained Causal Dynamics Learning with Quantization for Improving Robustness in Reinforcement Learning](https://mlsanity.com/paper/2406.03234/) (2024)
- [Masked Imitation Learning: Discovering Environment-Invariant Modalities in Multimodal Demonstrations](https://mlsanity.com/paper/2209.07682/) (2022)
- [Exploiting Causal Graph Priors with Posterior Sampling for Reinforcement Learning](https://mlsanity.com/paper/2310.07518/) (2023)
- [Curious Causality-Seeking Agents Learn Meta Causal World](https://mlsanity.com/paper/2506.23068/) (2025)
- [Beyond Sliding Windows: Learning to Manage Memory in Non-Markovian Environments](https://mlsanity.com/paper/2512.19154/) (2025)
- [Model-Advantage and Value-Aware Models for Model-Based Reinforcement Learning: Bridging the Gap in Theory and Practice](https://mlsanity.com/paper/2106.14080/) (2021)
- [Efficient and Generalizable Environmental Understanding for Visual Navigation](https://mlsanity.com/paper/2506.15377/) (2025)
- [Disentangled Representations for Causal Cognition](https://mlsanity.com/paper/2407.00744/) (2024)

## Abstract (arXiv)

Accuracy and generalization of dynamics models is key to the success of model-based reinforcement learning (MBRL). As the complexity of tasks increases, so does the sample inefficiency of learning accurate dynamics models. However, many complex tasks also exhibit sparsity in the dynamics, i.e., actions have only a local effect on the system dynamics. In this paper, we exploit this property with a causal invariance perspective in the single-task setting, introducing a new type of state abstraction called textitmodel-invariance. Unlike previous forms of state abstractions, a model-invariance state abstraction leverages causal sparsity over state variables. This allows for compositional generalization to unseen states, something that non-factored forms of state abstractions cannot do. We prove that an optimal policy can be learned over this model-invariance state abstraction and show improved generalization in a simple toy domain. Next, we propose a practical method to approximately learn a model-invariant representation for complex domains and validate our approach by showing improved modelling performance over standard maximum likelihood approaches on challenging tasks, such as the MuJoCo-based Humanoid. Finally, within the MBRL setting we show strong performance gains with respect to sample efficiency across a host of other continuous control tasks.