Invariant Causal Prediction for Block MDPs

By Amy Zhang, Clare Lyle, Shagun Sodhani, Angelos Filos, Marta Zofia Kwiatkowska, Joelle Pineau, Yarin Gal, Doina Precup. 2020

159 citations recorded by the source corpus.

Generalization across environments is critical to the successful application of reinforcement learning algorithms to real-world challenges. In this paper, we consider the problem of learning abstractions that generalize in block MDPs, families of environments with a shared latent state space and dynamics structure over that latent space, but varying observations. We leverage tools from causal inference to propose a method of invariant prediction to learn model-irrelevance state abstractions (MISA) that generalize to novel observations in the multi-environment setting. We prove that for certain classes of environments, this approach outputs with high probability a state abstraction corresponding to the causal feature set with respect to the return. We further provide more general bounds on model error and generalization error in the multi-environment setting, in the process showing a connection between causal variable selection and the state abstraction framework for MDPs. We give empirical evidence that our methods work in both linear and nonlinear settings, attaining improved generalization over single- and multi-task baselines.

Subfields: Artificial Intelligence, Management Science and Operations Research

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Management Science and Operations Research

Citations in the corpus

Towards Continual Reinforcement Learning: A Review and Perspectives (2020, background): “In recent work, Zhang et al. (2020); Hansen-Estruch, Zhang, Nair, Yin, and Levine (2022); Ashcraft, Stoler, Ewulum, and Agarwala (2022) connect invariant causal prediction to model-irrelevance state abstractions to learn invariant representations in the Block MDP setting (Du et al., 2019b).”
Learning Invariant Representations for Reinforcement Learning without Reconstruction (2020, methodology): “ can update the bound in Theorem 2 to incorporate L: jV (s) V (˚(s))j 2+2L (1 )(1 c) : MDP dynamics have a strong connection to causal inference and causal graphs, which are directed acyclic graphs [13, 25, 36]. Speciﬁcally, the state and action at time tcausally affect the next state at time t+ 1. ”
AdaRL: What, Where, and How to Adapt in Transfer Reinforcement Learning (2021, methodology): “In such a case, instead of learning the distribution shift over the high-dimensional input, thanks to the parsimonious graphical representation, we introduce a low-dimensional vector θ k to characterize the domain-speciﬁc information in a compact way (Zhang et al., 2020b).”
Multi-Task Reinforcement Learning with Context-based Representations (2021, methodology): “We attach this additional relaxation of the original CMDP definition to define a new setting, a Block Contextual MDP (BC-MDP)2 (Du et al., 2019; Zhang et al., 2020):”
A Relational Intervention Approach for Unsupervised Dynamics Generalization in Model-Based Reinforcement Learning (2022, background): “In addition, (Zhang et al., 2020a;b) improve the generalization ability of state abstraction.”
Automatic Data Augmentation for Generalization in Reinforcement Learning (2021, background): “Several strategies have been proposed to alleviate this problem, such as the use of regularization [19, 71, 11, 28], data augmentation [11, 44, 69, 38, 41], or representation learning [72, 74].”
Contrastive Behavioral Similarity Embeddings for Generalization in Reinforcement Learning (2021, background): “Causality-based methods are a different way of tackling generalization, but current solutions do not scale to high-dimensional observation spaces (e.g., Killian et al., 2017; Perez et al., 2020; Zhang et al., 2020).”
Learning Domain Invariant Representations in Goal-conditioned Block MDPs (2021, methodology): “Moreover, we assume that ME follows the generalized Block structure [6].”