# Provably Efficient Causal Model-Based Reinforcement Learning for Systematic Generalization

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

## Key facts

- Cited by 12 papers in the map (data through 14 Oct 2026).
- Builds on 28 of its 60 references that are in the map.
- First posted on arXiv 14 Feb 2022.

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

- [Provably Efficient Reinforcement Learning with Linear Function Approximation](https://mlsanity.com/paper/1907.05388/) (2019)
- [Model-Based Reinforcement Learning with Value-Targeted Regression](https://mlsanity.com/paper/2006.01107/) (2020)
- [FLAMBE: Structural Complexity and Representation Learning of Low Rank MDPs](https://mlsanity.com/paper/2006.10814/) (2020)
- [Provably Efficient Exploration in Policy Optimization](https://mlsanity.com/paper/1912.05830/) (2019)
- [Reward-Free Exploration for Reinforcement Learning](https://mlsanity.com/paper/2002.02794/) (2020)
- [Unifying PAC and Regret: Uniform PAC Bounds for Episodic Reinforcement Learning](https://mlsanity.com/paper/1703.07710/) (2017)
- [Provably efficient RL with Rich Observations via Latent State Decoding](https://mlsanity.com/paper/1901.09018/) (2019)
- [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)
- [Fast active learning for pure exploration in reinforcement learning](https://mlsanity.com/paper/2007.13442/) (2020)
- [On Reward-Free Reinforcement Learning with Linear Function Approximation](https://mlsanity.com/paper/2006.11274/) (2020)
- [Adaptive Reward-Free Exploration](https://mlsanity.com/paper/2006.06294/) (2020)
- [Sample Complexity of Multi-task Reinforcement Learning](https://mlsanity.com/paper/1309.6821/) (2013)
- [Provably Efficient Reward-Agnostic Navigation with Linear Value Iteration](https://mlsanity.com/paper/2008.07737/) (2020)
- [Near-optimal Reinforcement Learning in Factored MDPs](https://mlsanity.com/paper/1403.3741/) (2014)
- [Causal Reinforcement Learning using Observational and Interventional Data](https://mlsanity.com/paper/2106.14421/) (2021)
- [Budgeted Experiment Design for Causal Structure Learning](https://mlsanity.com/paper/1709.03625/) (2017)
- [Efficient Local Planning with Linear Function Approximation](https://mlsanity.com/paper/2108.05533/) (2021)
- [Testing Conditional Independence of Discrete Distributions](https://mlsanity.com/paper/1711.11560/) (2017)
- [Active Model Estimation in Markov Decision Processes](https://mlsanity.com/paper/2003.03297/) (2020)
- [On Reward-Free RL with Kernel and Neural Function Approximations: Single-Agent MDP and Markov Game](https://mlsanity.com/paper/2110.09771/) (2021)
- [Model-Invariant State Abstractions for Model-Based Reinforcement Learning](https://mlsanity.com/paper/2102.09850/) (2021)
- [Optimal Testing of Discrete Distributions with High Probability](https://mlsanity.com/paper/2009.06540/) (2020)
- [Towards Minimax Optimal Reinforcement Learning in Factored Markov Decision Processes](https://mlsanity.com/paper/2006.13405/) (2020)
- [On the Sample Complexity of Learning Bayesian Networks](https://mlsanity.com/paper/1302.3579/) (2013)
- [On the Sample Complexity of Causal Discovery and the Value of Domain Expertise](https://mlsanity.com/paper/2102.03274/) (2021)
- [Oracle-Efficient Regret Minimization in Factored MDPs with Unknown Structure](https://mlsanity.com/paper/2009.05986/) (2020)
- [Improved Exploration in Factored Average-Reward MDPs](https://mlsanity.com/paper/2009.04575/) (2020)

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

- [Factored Adaptation for Non-Stationary Reinforcement Learning](https://mlsanity.com/paper/2203.16582/) (2022)
- [Fine-Grained Causal Dynamics Learning with Quantization for Improving Robustness in Reinforcement Learning](https://mlsanity.com/paper/2406.03234/) (2024)
- [Exploiting Causal Graph Priors with Posterior Sampling for Reinforcement Learning](https://mlsanity.com/paper/2310.07518/) (2023)
- [Towards Empowerment Gain through Causal Structure Learning in Model-Based RL](https://mlsanity.com/paper/2502.10077/) (2025)
- [BECAUSE: Bilinear Causal Representation for Generalizable Offline Model-based Reinforcement Learning](https://mlsanity.com/paper/2407.10967/) (2024)
- [Causal Knowledge Transfer for Multi-Agent Reinforcement Learning in Dynamic Environments](https://mlsanity.com/paper/2507.13846/) (2025)
- [Boosting Efficiency in Task-Agnostic Exploration through Causal Knowledge](https://mlsanity.com/paper/2407.20506/) (2024)
- [Adaptive Data Exploitation in Deep Reinforcement Learning](https://mlsanity.com/paper/2501.12620/) (2025)
- [A Roadmap Towards Improving Multi-Agent Reinforcement Learning With Causal Discovery And Inference](https://mlsanity.com/paper/2503.17803/) (2025)
- [Disentangled Representations for Causal Cognition](https://mlsanity.com/paper/2407.00744/) (2024)
- [Transferable Delay-Aware Reinforcement Learning via Implicit Causal Graph Modeling](https://mlsanity.com/paper/2605.12312/) (2026)
- [When Do Causal World Models Help Modular LLM Agents](https://mlsanity.com/paper/2610.00012/) (2026)

## Abstract (arXiv)

In the sequential decision making setting, an agent aims to achieve systematic generalization over a large, possibly infinite, set of environments. Such environments are modeled as discrete Markov decision processes with both states and actions represented through a feature vector. The underlying structure of the environments allows the transition dynamics to be factored into two components: one that is environment-specific and another that is shared. Consider a set of environments that share the laws of motion as an example. In this setting, the agent can take a finite amount of reward-free interactions from a subset of these environments. The agent then must be able to approximately solve any planning task defined over any environment in the original set, relying on the above interactions only. Can we design a provably efficient algorithm that achieves this ambitious goal of systematic generalization? In this paper, we give a partially positive answer to this question. First, we provide a tractable formulation of systematic generalization by employing a causal viewpoint. Then, under specific structural assumptions, we provide a simple learning algorithm that guarantees any desired planning error up to an unavoidable sub-optimality term, while showcasing a polynomial sample complexity.