Planning under Distribution Shifts with Causal POMDPs

Matteo Ceriscioli, Karthika Mohan

Abstract: In the real world, planning is often challenged by distribution shifts. A model of the environment obtained under some circumstances may not remain valid as the distribution of states and the environment dynamics change, causing learned strategies to fail. A distribution shift that alters the environment or its dynamics is known as a domain shift. In this work, we propose a theoretical framework for planning under partial observability, using Partially Observable Markov Decision Processes (POMDPs) equipped with a causal factorization. By representing domain shifts as interventions on the causal model induced by the causal factorization, this framework enables evaluating plans under hypothesized shifts and actively identifying which components of the environment have changed. We show how to maintain and update a belief about the state and the domain and demonstrate that the value function remains piecewise linear and convex in this augmented belief space.