Paper deep dive
Planning under Distribution Shifts with Causal POMDPs
Matteo Ceriscioli, Karthika Mohan
Intelligence
Status: succeeded | Model: Gemma-4-26B-A4B | Prompt: intel-v1 | Confidence: 93%
Last extracted: 7/20/2026, 8:28:23 AM
Summary
This paper proposes a theoretical framework for planning under partial observability using Partially Observable Markov Decision Processes (POMDPs) formulated with causal knowledge. It addresses the challenge of distribution shifts by representing environmental changes as interventions on a causal POMDP. The authors demonstrate that the value function remains piecewise linear and convex (PWLC) in an augmented belief space covering both latent states and domain shifts, thereby preserving the tractability of standard alpha-vector-based planning algorithms.
Entities (8)
Relation Signals (7)
Distribution Shift → representedas → Intervention
confidence 96% · By representing shifts in the environment as interventions on this causal POMDP
Causal POMDP → handles → Distribution Shift
confidence 95% · By representing shifts in the environment as interventions on this causal POMDP, the framework enables evaluating plans under hypothesized changes
Causal POMDP → hasproperty → Piecewise Linear and Convex (PWLC)
confidence 94% · we prove that the value function remains piecewise linear and convex (PWLC) in this augmented belief space.
Stochastic Shift Intervention → istypeof → Intervention
confidence 93% · We introduce stochastic shifts, a family of soft interventions.
Piecewise Linear and Convex (PWLC) → enables → Alpha-vector-based POMDP
confidence 92% · Preservation of PWLC under distribution shifts has the advantage of maintaining the tractability of planning via alpha-vector-based POMDP methods.
Planning under Unknown Distribution Shift → requires → Joint Belief
confidence 91% · it is appropriate to keep track of them using a joint belief b(s, sigma).
Causal Influence Diagram → usedin → Causal POMDP
confidence 90% · In this paper, we study factored POMDPs whose intra- and inter-temporal causal structure is expressed using Causal Influence Diagrams (CIDs)
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract:In the real world, planning is often challenged by distribution shifts. As such, a model of the environment obtained under one set of conditions may no longer remain valid as the distribution of states or the environment dynamics change, which in turn causes previously learned strategies to fail. In this work, we propose a theoretical framework for planning under partial observability using Partially Observable Markov Decision Processes (POMDPs) formulated using causal knowledge. By representing shifts in the environment as interventions on this causal POMDP, the framework enables evaluating plans under hypothesized changes and actively identifying which components of the environment have been altered. We show how to maintain and update a belief over both the latent state and the underlying domain, and we prove that the value function remains piecewise linear and convex (PWLC) in this augmented belief space. Preservation of PWLC under distribution shifts has the advantage of maintaining the tractability of planning via $\alpha$-vector-based POMDP methods.
Tags
Links
- Source: https://arxiv.org/abs/2602.23545v2
- Canonical: https://arxiv.org/abs/2602.23545v2
Trouble viewing inline? Open PDF directly →
Full Text
38,555 characters extracted from source content.
Expand or collapse full text
Planning under Distribution Shifts with Causal POMDPs Matteo Ceriscioli, Karthika Mohan School of Electrical Engineering and Computer Science (EECS) Oregon State University Corvallis, OR 97331, USA ceriscim,karthika.mohan@oregonstate.edu Abstract In the real world, planning is often challenged by distribution shifts. As such, a model of the environment obtained under one set of conditions may no longer remain valid as the distribution of states or the environment dynamics change, which in turn causes previously learned strategies to fail. In this work, we propose a theoretical framework for planning under partial observability using Partially Observable Markov Decision Processes (POMDPs) formulated using causal knowledge. By representing shifts in the environment as interventions on this causal POMDP, the framework enables evaluating plans under hypothesized changes and actively identifying which components of the environment have been altered. We show how to maintain and update a belief over both the latent state and the underlying domain, and we prove that the value function remains piecewise linear and convex (PWLC) in this augmented belief space. Preservation of PWLC under distribution shifts has the advantage of maintaining the tractability of planning via α-vector-based POMDP methods. Introduction Planning aims to compute an optimal way for an agent to act, assuming access to a complete and correct model of the environment. In stochastic settings, this often means identifying a policy that maps information states to actions and maximizes long-term reward [Kaelbling et al., 1998]. Difficulties arise once the agent is deployed in a different environment and the underlying state distribution or transition dynamics change. Distribution shifts undermine the assumption that the transition and observation processes remain fixed and accurately describe how states evolve. To illustrate this issue, consider a delivery rover whose policy and world model were obtained in an environment with specific conditions, such as urban areas with temperate weather. When the same rover is deployed in a different setting, relevant factors, such as friction, visibility, or mechanical stress, may change. A conventional POMDP planner cannot identify which mechanisms have shifted or why performance has degraded. A causally informed model can represent such changes as interventions on specific components, reason about their effects on the environment and the sensors, and support effective planning under the resulting distribution shift. This paper examines whether planning remains tractable in the presence of distribution shifts. The tractability of finite-horizon POMDPs relies on the fact that their value functions are PWLC in the belief state [Smallwood and Sondik, 1973]. When the environment is affected by an unknown shift, the resulting ambiguity in the transition dynamics introduces an additional source of uncertainty into the model. Although certain structured forms of uncertainty are known to preserve the PWLC structure [Osogami, 2015], general forms of ambiguity can cause finite-horizon POMDPs to lose their PWLC properties [Saghafian, 2018]. We consider a causal formulation of POMDPs in which distribution shifts are represented as interventions. Under this formulation, we show that the value function retains its PWLC structure despite the presence of unknown shifts. Thus this ensures that standard α-vector–based POMDP planning algorithms remain applicable. Causal modeling is widely used to support robust decision making under perturbations of the data-generating process [Pearl and Bareinboim, 2011, Pearl, 2018, Schölkopf, 2022, Ceriscioli and Mohan, 2025], as it offers a principled framework for analyzing how distribution shifts affect outcomes. In particular, such shifts can be naturally represented as interventions in a causal model [Richens and Everitt, 2024, Ceriscioli and Mohan, 2025], which allows systematic reasoning about their consequences. The primary contribution of this paper is a causal framework for planning under partial observability in the presence of distribution shifts, which preserves the PWLC structure of the value function. The framework enables the evaluation of policies under specified shifts and supports the detection and identification of changes in the environment. All proofs are included in Appendix A. 1 Preliminaries Yt−1Y_t-1Xt−1X_t-1Dt−1D_t-1Ut−1U_t-1YtY_tXtX_tDtD_tUtU_tYt+1Y_t+1Xt+1X_t+1Dt+1D_t+1Ut+1U_t+1⋯·s⋯·s (a) Non-time-homogeneous Causal POMDP Yt−1Y_t-1Xt−1X_t-1Dt−1D_t-1Ut−1U_t-1YtY_tXtX_tDtD_tUtU_tYt+1Y_t+1Xt+1X_t+1Dt+1D_t+1Ut+1U_t+1⋯·s⋯·s (b) Time-homogeneous Causal POMDP YtY_tXtX_tDtD_tUtU_tYt+1Y_t+1Xt+1X_t+1 (c) Compact representation . Figure 1: On the left is an example of a non–time-homogeneous causal POMDP. If the causal model is faithful, that is, every edge in the causal graph corresponds to an actual direct dependency between two variables, then any change in the graph structure across timesteps implies a change in the transition function. In the center is a CID that is compatible with a time-homogeneous causal POMDP. On the right is a compact representation of a time-homogeneous causal POMDP compatible with the CID shown in the center. Let X be a node in a graph G. In this paper we denote the set of parent nodes of X as Pa(X)Pa(X). Random variables are represented with upper-case letters, e.g. X, and the corresponding instantiations with lower-case letters, e.g. x. The set of observable values of a random variable X is called dom(X)dom(X). Given a set X, Π() (X) is the set of all probability distributions over X. Factored POMDPs. Partially Observable Markov Decision Processes (POMDPs) [Åström, 1965, Kaelbling et al., 1998] model sequential decision-making in which the agent lacks full observability of the environment. When each state in a POMDP is represented as a vector s=(v1,…,vn)s=(v_1,…,v_n) of instantiations of random variables V1,…,VnV_1,…,V_n, the model is called a factored POMDP [Boutilier and Poole, 1996]. This structure allows the transition function to be decomposed into sub-functions involving only subsets of the state variables. Causal Influence Diagrams. In this paper, we study factored POMDPs whose intra- and inter-temporal causal structure is expressed using Causal Influence Diagrams (CIDs) [Heckerman, 1995, Everitt et al., 2021]. CIDs extend Causal Bayesian Networks (CBNs) [Pearl, 2009] by incorporating decision-theoretic elements, offering a structured way to represent causal relationships while specifying what the agent observes and what it influences through its actions. Definition 1 (Causal Influence Diagram [Everitt et al., 2021]). A Causal Influence Diagram (CID) is a Causal Bayesian Network M=(G=(,),P)M=(G=(V,E),P), where P is a joint probability distribution compatible with the conditional independences encoded in the DAG G. The nodes in V are partitioned into decision (D), utility (U), and chance (C) nodes, =(,,)V=(D,U,C). Each utility node UiU_i is assigned a utility function fi:dom(Pa(Ui))→ℝf_i:dom(Pa(U_i)) . In a CID, the environment is described by the set of chance nodes C. Each chance node C∈C corresponds to a random variable, The set of decision nodes D contains all the variables which value is set by an agent, for each D∈D it is possible to assign a policy π:dom(Pa(D))→π:dom(Pa(D))→ A. 2 Causal POMDPs There exist various causal representations of POMDPs, including models that integrate causal structure into online planning or representation learning. For example, CAR-DESPOT incorporates causal information into the AR-DESPOT planner to reason about confounding in robotic environments [Cannizzaro and Kunze, 2023]. Earlier work has proposed a causal POMDP model for causal representation learning and has shown its usefulness for zero-shot learning in complex tasks [Sontakke et al., 2021]. Other approaches use causal modeling to recover hidden causal dynamics within partially observable environments [Liang and Boularias, 2021], or employ causal abstractions to improve long-horizon reasoning [Gao et al., 2025]. In this work, we adapt the formulation of Ceriscioli and Mohan [2025], which models a POMDP as an infinite CID unrolled over time containing a decision node and a utility node at each timestep t, corresponding respectively to the action and reward at time t. Definition 2 (Causal POMDP). A causal POMDP is a tuple (,,,R,o,O,γ)(V, A,T,R,V_o,O,γ) where: 1. =(V1,…,Vn)V=(V_1,…,V_n) is the ordered set of state variables. 2. A is the finite set of actions. 3. T:dom()×→Π(dom())T:dom(V)× A→ (dom(V)) is the state-transition function. Given (t)V^(t) the set of state variables before the transition, and (t+1)V^(t+1) the set of state variables after the transition, it is possible to decompose T as follows: T(v(t),a(t),v(t+1))=∏i=1nTi(vi(t+1)∣a(t),pa(vi(t+1)))T(v^(t),a^(t),v^(t+1))= _i=1^nT_i(v_i^(t+1) a^(t),pa(v_i^(t+1))) (1) with Pa(Vi(t+1))⊆(t)∪(⋃k<iVk(t+1))Pa(V_i^(t+1)) ^(t)∪( _k<iV^(t+1)_k). Each function TiT_i governs the transition of the state variable ViV_i. 4. R:R×→ℝR:V_R× A is a reward function, where R⊆V_R . 5. o=Vo,1,…,Vo,pV_o=\V_o,1,…,V_o,p\ is a set of observable variables, forming the agent’s observation. 6. =O1,…,OpO=\O_1,…,O_p\ O:dom()×dom(o)→[0,1]O:dom(V)× A× dom(V_o)→[0,1] is the set of conditional observation probabilities. 7. γ∈[0,1)γ∈[0,1) is the discount factor. The set of state variables induces the set of states ≔dom()=dom(V1)×⋯×dom(Vn)S dom(V)=dom(V_1)×…× dom(V_n). Similarly, the set of observations is ≔dom(o) dom(V_o). The agent receives a reward R(s,a)R(s,a) after taking an action a∈a∈ A in state s∈s . Observable variables can be included in the state (o⊆V_o ) without loss of expressiveness: for any causal POMDP where o∩=∅V_o = , there exists an equivalent causal POMDP with o′⊆V_o where transitions and rewards do not depend on oV_o. Assumption 1. The observable variables are also state variables, i.e. o⊆V_o . When Assumption 1 holds, the state fully determines the observation. Therefore, the conditional observation probability function simplifies to O(v′,a,vo′)=1O(v ,a,v_o )=1 if vo′v_o is compatible with v′v and 0 otherwise, are the new state and observation after action a. Since v′v determines vo′v _o, we can write O(v′,vo′)O(v ,v _o) instead of O(v′,a,vo′)O(v ,a,v _o). Under Assumption 1, a causal POMDP (,,,R,o,O,γ)(V, A,T,R,V_o,O,γ) induces a CID with an infinite DAG G=(′,E)G=(V ,E), where =⋃t=0∞(t)V= _t=0^∞V^(t) and (t)≔V^(t) (see Figure 1(b)). Time-homogeneity of the transition function allows a compact CID representation including only two consecutive timesteps t and t+1t+1, with decisions and utilities at t, and edges involving D and U, between timesteps, and within t+1t+1 (see Figure 1(c)). If all state variables are observable, this reduces to a causal MDP. Distribution shifts generally alter the state transition function, and if the distribution shift is unknown, then estimating it requires planning using a non-time-homogeneous POMDP, non-time-homogeneous causal POMDPs can also be represented with infinite CIDs, as illustrated in Figure 1(a). 3 Planning under Distribution Shifts Being able to discern the causal relationships governing the environment allows us to infer the state dynamics under changing conditions. In a causal model, interventions are deliberate alterations of components or mechanisms of the model. A distribution shift is any change in the probability distribution of a random variable. In this paper, we focus on shifts that affect the variables describing the environment in which we plan, specifically, the variables used to factorize the POMDP state in the causal POMDP model. A domain denotes the probabilistic configuration of the environment under a particular shift, including the original one. In this section, we develop the use of causal POMDPs for planning under potential distribution shifts. We first formalize distribution shifts as interventions within the underlying causal model. A key consequence of modeling distribution shifts as interventions is that they can be embedded directly into the agent’s belief and planning process. We proceed by describing how to evaluate a policy under a specified distribution shift. Finally, we address the problem of planning while concurrently identifying the distribution shift, showing how to update the agent’s belief over states and domains and that the resulting belief value function remains piecewise linear and convex. 3.1 Modeling Distribution Shifts as Interventions As shown in previous work [Richens and Everitt, 2024, Ceriscioli and Mohan, 2025], interventions are an effective way to represent distribution shifts in a causal model. Applying an intervention σ to a set of variables may change the joint distribution P()P(V). When the model and intervention are known, it is possible to compute the updated distribution P(;σ)P(V;σ) which is called the interventional distribution, opposed to the observational distribution P()P(V), which is the one we observe when no intervention is applied. We introduce stochastic shifts, a family of soft interventions. Definition 3 (Stochastic Shift Intervention). Let M=(G,Θ)M=(G, ) be a CBN, and X be a discrete random variable with dom(X)=x1,…,xmdom(X)=\x_1,…,x_m\. A Stochastic Shift Intervention σ is an intervention associated with a matrix: Aσ=(p11…p1m⋮⋱⋮pm1…pmm)with ∑jpij=1 for every i.A_σ= ( matrixp_11&…&p_1m\\ & & \\ p_m1&…&p_m matrix ) $ _jp_ij=1$ for every $i$. (2) such that when applied to a random variable X, its conditional distribution is updated as follows: P(X=xj∣pa(X);σ)=∑i=1mpijP(X=xi∣pa(X))P(X=x_j pa(X);σ)= _i=1^mp_ijP(X=x_i pa(X)) (3) The condition ∑jpij=1 _jp_ij=1 is equivalent to requiring that each row in AσA_σ sums up to one. Observe that if the probability distribution of X is represented with a probability vector, e.g. P(X|pa(X))=(P(x1∣pa(X)),…,P(xm∣pa(X)))TP(X|pa(X))=(P(x_1 pa(X)),…,P(x_m pa(X)))^T, then we can apply the stochastic shift intervention σ by matrix multiplication, i.e. P(X;σ)=AσTP(X)P(X;σ)=A_σ^TP(X) (4) AσA_σ is the identity matrix iff σ is the identity intervention σid _id, which leaves the domain and distribution unchanged. Example. Let X∼Unif(1,2,3)X Unif(\1,2,3\) and σ be a stochastic shift s.t. p11=p22=1p_11=p_22=1, and p31=p32=12p_31=p_32= 12. Each time X takes the value 33, σ remaps it to 11 or 22 with equal probability. So P(X=x;σ)=12P(X\!=\!x;σ)\!=\! 12 if x∈1,2x\!∈\!\1,2\ and 0 if x=3x=3. A stochastic shift intervention σ maps probability distributions over a finite set to other distributions. The following result shows that, for any starting distribution, there exists a σ that maps it to any target distribution. Proposition 1. Given a conditional distribution P(X∣pa(X))P(X pa(X)) and an arbitrary target conditional distribution P′(X∣pa(X))P (X pa(X)), it is possible to define a stochastic shift intervention σ s.t. P′(X∣pa(X))=P(X∣pa(X);σ)P (X pa(X))=P(X pa(X);σ). Note that even if the original distribution P(X∣pa(X))P(X pa(X)) and the shifted distribution P(X∣pa(X);σ)P(X pa(X);σ) are both observed, in the general case it is not always possible to uniquely determine the distribution shift σ, as it is possible that two different shifts applied to the same distribution generate the same distribution. Suppose X is the outcome of a fair coin flip, i.e., X∼Bern(12)X Bern( 12), also suppose that after altering the coin we observe Bern(34)Bern( 34) as the new distribution, then both Aσ=(100.50.5)Aσ′=(0.50.510)A_σ= ( matrix1&0\\ 0.5&0.5 matrix ) A_σ = ( matrix0.5&0.5\\ 1&0 matrix ) (5) correspond to shifts that map Bern(12)Bern( 12) to Bern(34)Bern( 34). 3.2 Evaluating Policies under Distribution Shifts Let M=(,,,R,o,O,γ)M=(V, A,T,R,V_o,O,γ) be a causal POMDP, π:dom(o)→π:dom(V_o)→ A a policy, and σ a stochastic shift intervention representing a distribution shift. When planning under partial observability, the state is generally not completely available to the agent and therefore it maintains a belief bSb_S about the state that updates whenever it receives an observation [Kaelbling et al., 1998]. Evaluating a policy consists in computing its expected return, represented by the belief value function Vπ(b)V^π(b) using the Bellman equation for causal POMDPs: Vπ(bS)=∑s∈SR(s,π(bS))bS(s)+γ∑s′,o′O(s′,o′)∑sbS(s)∏Vi′∈(t+1)Ti(vi′∣π(bS),pa(Vi′))Vπ(bSπ(bS),o′)V^π(b_S)= _s∈ SR(s,\!π(b_S))b_S(s)+γ\! _s ,o \!O(s \!,o )\! _sb_S(s)\!\!\!\!\!\!\!\!\! _V_i ^(t+1)\!\!\!\!\!\!\!\!T_i(v _i\! \!π(b_S),pa(V _i))V^π(b_S^π(b_S),o \!) (6) Known shift σ. We write Vπ(bS;σ)V^π(b_S;σ) to denote the value function when the environment is affected by the shift σ. Its expression is identical to that in Equation 6, except that each transition term Ti(vi∣π(bS),pa(vi))T_i(v_i π(b_S),pa(v_i)) is replaced by Ti(vi∣π(bS),pa(vi);σ)T_i(v_i π(b_S),pa(v_i);σ). Note that if σ=σidσ= _id then Vπ(bS;σ)=Vπ(bS)V^π(b_S;σ)=V^π(b_S). Once the shift σ is fixed and the transition functions are updated via Equation 3, the problem of computing the value function of a causal POMDP reduces to that of a standard POMDP. 3.3 Planning under Unknown Distribution Shifts Now we consider the task of planning in an environment affected by an unknown distribution shift. 3.3.1 State and Domain Estimation. As part of planning with causal POMDP under an unknown distribution shift, it is no longer sufficient for the agent to keep track of its belief about the state, as it also needs to keep a belief about the unknown domain. It is possible to define a prior on the state and the domain separately as bS(s)b_S(s) and bΣ(σ)b_ (σ), however, since in the general case at each timestep the observation depends both on the previous state and the domain, the state belief and the domain belief become coupled and therefore it is appropriate to keep track of them using a joint belief b(s,σ)b(s,σ). If we have no prior knowledge we can express that as uniform priors on both the states and the domain b(s,σ)=b(σ|s)b(s)=bΣ(σ)bS(s)b(s,σ)=b(σ|s)b(s)=b_ (σ)b_S(s) where bS(s)∼Unif(S)b_S(s) (S), and bΣ(σ)∼Unif(Σ)b_ (σ) ( ). The following proposition shows that the joint belief over states and domains admits a Bayesian update analogous to the standard POMDP filter, but extended to account for domain-dependent transition dynamics. Proposition 2. [State-Domain Joint Belief Update] Let b be the current joint belief over states and domains, a be the action taken at time t, and o′o the observation received after performing a. Then the updated state-domain joint belief b′b is: bo′,a′(s′,σ)=O(s′,o′)P(o′∣a,b)∑sb(s,σ)∏Vi∈(+)T(vi∣pa(Vi);σ) splitb _o ,a(s ,σ)= O(s ,o )P(o a,b) _sb(s,σ)\!\!\!\!\!\! _V_i ^(t+1)\!\!\!\!\!T(v_i\! \!pa(V_i);σ) split (7) such that s′=(v1,…,vn)s =(v_1,…,v_n) for V1,…,Vn=(t+1)\V_1,…,V_n\=V^(t+1). 3.4 Preservation of PWLC under Distribution Shift We establish that, despite the presence of an unknown shift, the structural properties that support tractable POMDP planning remain intact. Value Function for Casual POMDPs under an Unknown Shift. To demonstrate that the value function of a causal POMDP subject to a distribution shift and planning horizon n is piecewise linear and convex, and that it admits a finite vector representation analogous to that of standard POMDPs, we propose an approach similar to Porta et al. [2005] and present a constructive proof based on the state–action value function. The following lemma shows that the value function under an unknown distribution shift can still be expressed as a maximum over linear functionals of the joint belief. Lemma 1. Let Σ be a set of stochastic shift interventions, the value function of a causal POMDP with set of states S subject to an unknown distribution shift in the set Σ with planning horizon n can be expressed as: Vn(b)=maxαnii∑s∈∫Σαni(s,σ)b(s,σ)σV_n(b)= _\ _n^i\_i _s _ α^i_n(s,σ)b(s,σ)\,dσ (8) for some functions αni:S×Σ→ℝ _n^i:S× . Using the characterization of Vn(b)V_n(b) provided by Lemma 1, the following theorem completes our analysis by proving that the value function under distribution shifts is still piecewise linear and convex, thereby preserving the structural form that underlies α-vector–based planning. Theorem 1. Let Σ be a set of stochastic shift interventions, the value function of a causal POMDP subject to an unknown distribution shift in the set Σ with planning horizon n is piecewise linear and convex in the joint state-domain belief b. This confirms that distribution shifts do not compromise the representational tractability of the planning problem. Conclusions Recognizing distribution shifts as a critical source of uncertainty in planning, this work introduces a framework for decision-making under such shifts using causal POMDPs. By representing the state in a factorized form and distribution shifts as interventions on the underlying causal model, causal POMDPs enable both the evaluation of policies under specified shifts and planning in an unknown domain by maintaining a joint belief over states and domains. We show that, even under the uncertainty in the transition dynamics that arises from an unknown distribution shift, the finite-horizon value function remains piecewise linear and convex with respect to the joint state–domain belief, thereby preserving the structural properties that support α-vector–based planning in standard POMDPs. References K. J. Åström (1965) Optimal control of markov processes with incomplete state information. Journal of Mathematical Analysis and Applications 10, p. 174–205. External Links: Link Cited by: §1. C. Boutilier and D. Poole (1996) Computing optimal policies for partially observable decision processes using compact representations. In Proceedings of the Thirteenth National Conference on Artificial Intelligence - Volume 2, AAAI’96, p. 1168–1175. External Links: ISBN 026251091X Cited by: §1. R. Cannizzaro and L. Kunze (2023) CAR-despot: causally-informed online pomdp planning for robots in confounded environments. In 2023 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS), Vol. , p. 2018–2025. External Links: Document Cited by: §2. M. Ceriscioli and K. Mohan (2025) Agents robust to distribution shifts learn causal world models even under mediation. In The Thirty-ninth Annual Conference on Neural Information Processing Systems, External Links: Link Cited by: §2, §3.1, Introduction. T. Everitt, R. Carey, E. D. Langlois, P. A. Ortega, and S. Legg (2021) Agent incentives: a causal perspective. Proceedings of the AAAI Conference on Artificial Intelligence 35 (13), p. 11487–11495. External Links: Link, Document Cited by: §1, Definition 1. H. Gao, T. Xu, T. Zhang, Y. Guo, C. Zhao, J. Ren, Y. Jiang, S. Guo, and F. Chen (2025) Causal dreamer for partially observable model-based reinforcement learning. Neurocomputing 652, p. 131012. External Links: ISSN 0925-2312, Document, Link Cited by: §2. D. Heckerman (1995) A bayesian approach to learning causal networks. In Proceedings of the Eleventh Conference on Uncertainty in Artificial Intelligence, UAI’95, San Francisco, CA, USA, p. 285–295. External Links: ISBN 1558603859 Cited by: §1. L. P. Kaelbling, M. L. Littman, and A. R. Cassandra (1998) Planning and acting in partially observable stochastic domains. Artificial Intelligence 101 (1), p. 99–134. External Links: ISSN 0004-3702, Document, Link Cited by: §1, §3.2, Introduction. J. Liang and A. Boularias (2021) Inferring time-delayed causal relations in pomdps from the principle of independence of cause and mechanism. In Proceedings of the Thirtieth International Joint Conference on Artificial Intelligence, IJCAI-21, Z. Zhou (Ed.), p. 1944–1950. Note: Main Track External Links: Document, Link Cited by: §2. T. Osogami (2015) Robust partially observable markov decision process. In Proceedings of the 32nd International Conference on Machine Learning, F. Bach and D. Blei (Eds.), Proceedings of Machine Learning Research, Vol. 37, Lille, France, p. 106–115. External Links: Link Cited by: Introduction. J. Pearl and E. Bareinboim (2011) Transportability of causal and statistical relations: a formal approach. Proceedings of the AAAI Conference on Artificial Intelligence 25 (1), p. 247–254. External Links: Link, Document Cited by: Introduction. J. Pearl (2009) Causality: models, reasoning and inference. 2nd edition, Cambridge University Press, USA. External Links: ISBN 052189560X Cited by: §1. J. Pearl (2018) Theoretical impediments to machine learning with seven sparks from the causal revolution. In Proceedings of the Eleventh ACM International Conference on Web Search and Data Mining, WSDM ’18, New York, NY, USA, p. 3. External Links: ISBN 9781450355810, Link, Document Cited by: Introduction. J. M. Porta, M. T. J. Spaan, and N. Vlassis (2005) Robot planning in partially observable continuous domains. In Proceedings of Robotics: Science and Systems, Cambridge, USA. External Links: Document Cited by: §3.4. J. Richens and T. Everitt (2024) Robust agents learn causal world models. In The Twelfth International Conference on Learning Representations, External Links: Link Cited by: §3.1, Introduction. S. Saghafian (2018) Ambiguous partially observable markov decision processes: structural results and applications. Journal of Economic Theory 178, p. 1–35. External Links: ISSN 0022-0531, Document, Link Cited by: Introduction. B. Schölkopf (2022) Causality for machine learning. In Probabilistic and Causal Inference: The Works of Judea Pearl, p. 765–804. External Links: ISBN 9781450395861 Cited by: Introduction. R. D. Smallwood and E. J. Sondik (1973) The optimal control of partially observable markov processes over a finite horizon. Operations Research 21 (5), p. 1071–1088. External Links: Document, Link Cited by: Introduction. S. A. Sontakke, A. Mehrjou, L. Itti, and B. Schölkopf (2021) Causal curiosity: rl agents discovering self-supervised experiments for causal representation learning. In Proceedings of the 38th International Conference on Machine Learning, M. Meila and T. Zhang (Eds.), Proceedings of Machine Learning Research, Vol. 139, p. 9848–9858. External Links: Link Cited by: §2. Appendix A Proof Proposition 1. Given a conditional distribution P(X∣pa(X))P(X pa(X)) and an arbitrary target conditional distribution P′(X∣pa(X))P (X pa(X)), it is possible to define a stochastic shift intervention σ s.t. P′(X∣pa(X))=P(X∣pa(X);σ)P (X pa(X))=P(X pa(X);σ). Proof. Let dom(X)=x1,…,xmdom(X)=\x_1,…,x_m\. We define a stochastic shift interventions σ with parameters pji=P′(X=xi∣pa(X))p_ji=P (X=x_i pa(X)). Observe that for all i∈1…,mi∈\1…,m\: P(X=xi∣pa(X);σ)=∑jpjiP(X=xj∣pa(X))=P′(X=xi∣pa(X))∑jP(X=xj∣pa(X))=P′(X=xi∣pa(X)) splitP(X&=x_i pa(X);σ)= _jp_jiP(X=x_j pa(X))\\ &=P (X=x_i pa(X)) _jP(X=x_j pa(X))\\ &=P (X=x_i pa(X)) split (9) Hence, the intervention σ transforms the original conditional distribution P(X∣pa(X))P(X pa(X)) exactly into the target distribution P′(X∣pa(X))P (X pa(X)), which concludes the proof. ∎ Proposition 2 (State-Domain Joint Belief Update). Let b be the current joint belief over states and domains, a be the action taken at time t, and o′o the observation received after performing a. Then the updated state state-domain joint belief b′b is: bo′,a′(s′,σ)=O(s′,o′)P(o′∣a,b)∑sb(s,σ)∏Vi∈(+)T(vi∣pa(Vi);σ) splitb _o ,a(s ,σ)= O(s ,o )P(o a,b) _sb(s,σ)\!\!\!\!\!\! _V_i ^(t+1)\!\!\!\!\!T(v_i\! \!pa(V_i);σ) split (10) such that s′=(v1,…,vn)s =(v_1,…,v_n) for V1,…,Vn=(t+1)\V_1,…,V_n\=V^(t+1). Proof. By definition: bo′,a′(s′,σ)=P(s′,σ∣o′,a,b)b _o ,a(s ,σ)=P(s ,σ o ,a,b) (11) where b is the current belief, a is the action taken at the current timestep, and o′o is the observation received after executing a. By Bayes rule: =P(o′∣s′,a,σ,b)P(o′∣a,b)P(s′,σ∣a,b) = P(o s ,a,σ,b)P(o a,b)P(s ,σ a,b) (12) Note that o′⟂a,σ,b∣s′o \!\!\!\!\!\! \a,σ,b\ s , so P(o′∣s′,a,σ,b)=P(o′∣s′)=O(s′,o′)P(o s ,a,σ,b)=P(o s )=O(s ,o ). =O(s′,o′)P(o′∣a,b)∑s∈P(s′∣s,a,σ,b)P(s,σ∣a,b) = O(s ,o )P(o a,b) _s P(s s,a,σ,b)P(s,σ a,b) (13) As s′⟂b∣s,a,σs \!\!\!\!\!\! b \s,a,σ\ and s,σ⟂a∣b\s,σ\ \!\!\!\!\!\! a b = = O(s′,o′)P(o′∣a,b)∑s∈P(s′∣s,a,σ)P(s,σ∣b) O(s ,o )P(o a,b) _s P(s s,a,σ)P(s,σ b) (14) = = O(s′,o′)P(o′∣a,b)∑s∈b(s,σ)∏Vi∈(+)T(vi∣pa(Vi);σ) O(s ,o )P(o a,b) _s b(s,σ)\!\!\!\! _V_i ^(t+1)\!\!\!\!T(v_i pa(V_i);σ) (15) which corresponds to the desired expression. ∎ Lemma 1. Let Σ be a set of stochastic shift interventions, the value function of a causal POMDP with set of states S subject to an unknown distribution shift in the set Σ with planning horizon n can be expressed as: Vn(b)=maxαnii∑s∈∫Σαni(s,σ)b(s,σ)σV_n(b)= _\ _n^i\_i _s _ α^i_n(s,σ)b(s,σ)\,dσ (16) for some functions αni:S×Σ→ℝ _n^i:S× . Proof. We prove the statement by induction on the planning horizon n. Base case. Let n=0n=0, since the plan ends after the execution of a single action, then the value function corresponds to: V0(b)=maxa∈Q0(b,a)=maxa∈∑s∈∫ΣR(s,a)b(s,σ)σV_0(b)= _a∈ AQ_0(b,a)= _a∈ A _s _ R(s,a)b(s,σ)\,dσ (17) We define α0i(s,σ)=R(s,a)a∈\ _0^i(s,σ)\=\R(s,a)\_a∈ A and hence the value function is in the desired form. Induction step. Assume the statement holds for planning horizon n−1n-1. Let bo′,a′b_o ,a be the belief obtained by updating the belief b after observing o′o and executing action a according Equation 7. Observe that the value function for a planning horizon n can be defined recursively as: Vn(b)=maxa∈∑s∈∫Σb(s,σ)R(s,a)σ+γ∑o′∈P(o′∣a,b)Vn−1(bo′,a′) splitV_n(b)= _a∈ A\ _s _ b(s,σ)R(s,a)\,dσ+γ _o ∈ P(o a,b)V_n-1(b_o ,a )\ split (18) by inductive hypothesis: Vn−1(bo′,a′)=maxαn−1j∑s′∈∫Σαn−1j(s′,σ)bo′,a′(s′,σ)σV_n-1(b_o ,a )= _\ _n-1^j\ _s _ _n-1^j(s ,σ)b_o ,a (s ,σ)\,dσ (19) by substituting the expression from Equation 19 into Equation 18, we obtain: Vn(b)=maxa∑s∈∫ΣR(s,a)b(s,σ)σ++γ∑o′∈P(o′∣a,b)maxαn−1j∑s′∈∫Σαn−1j(s′,σ)bo′,a′(s′,σ)dσ splitV_n(b)= _a\& _s _ R(s,a)b(s,σ)\,dσ+\\ +γ& _o ∈ P(o a,b) _\ _n-1^j\ _s _ _n-1^j(s ,σ)b_o ,a (s ,σ)\,dσ\ split (20) by updating the belief according to Proposition 2 we get: =maxa∈∑s∈∫ΣR(s,a)b(s,σ)σ++γ∑o′∈maxαn−1j∑s∈∫Σ[∑s′∈O(s′,o′)P(s′∣s,a,σ)αn−1j(s′,σ)]b(s,σ)dσ split= _a∈ A\& _s _ R(s,a)b(s,σ)\,dσ+\\ +γ& _o ∈ _\ _n-1^j\ _s _ [ _s O(s ,o )P(s s,a,σ) _n-1^j(s ,σ) ]b(s,σ)\,dσ\ split (21) Let αa,o′j(s,σ)≔∑s′∈O(s′,o′)P(s′∣s,a,σ)αn−1j(s′,σ) _a,o ^j(s,σ)\! \! _s \!O(s ,o )P(s \! \!s,a,σ) _n-1^j(s ,σ), then Equation 21 can be rewritten as: =maxa∑s∈∫ΣR(s,a)b(s,σ)σ+γ∑o′∈maxαn−1j∑s∈∫Σαa,o′j(s,σ)b(s,σ)σ split= _a\ _s _ R(s,a)b(s,σ)\,dσ+γ _o ∈ _\ _n-1^j\ _s _ _a,o ^j(s,σ)b(s,σ)\,dσ\ split (22) Let αa,o′,b≔argmaxαa,o′j∑s∈∫Σαa,o′j(s,σ)b(s,σ)σ _a,o ,b *arg\,max_\ _a,o ^j\ _s _ _a,o ^j(s,σ)b(s,σ)\,dσ, then: Vn(b)=maxa∈∑s∈∫Σ[R(s,a)+γ∑o′∈αa,o′,b]b(s,σ)σ splitV_n(b)= _a∈ A\! \ _s _ \! [R(s,a)\!+\!γ\! _o ∈ \! _a,o ,b ]b(s,σ)\,dσ\! \ split (23) Finally, we define: αnii≔R(s,a)+γ∑o′∈αa,o′,ba∈\ _n^i\_i \R(s,a)+γ _o ∈ _a,o ,b\_a∈ A (24) Then: Vn(b)=maxαnii∑s∈∫Σαni(s,σ)b(s,σ)σV_n(b)= _\ _n^i\_i _s _ _n^i(s,σ)b(s,σ)\,dσ (25) proving the induction step. ∎ Theorem 1. Let Σ be a set of stochastic shift interventions, the value function of a causal POMDP subject to an unknown distribution shift in the set Σ with planning horizon n is piecewise linear and convex in the joint state-domain belief b: Proof. By Lemma 1 for every planning horizon n we know that: Vn(b)=maxαnii∑s∈∫Σαin(s,σ)b(s,σ)σV_n(b)= _\ _n^i\_i _s _ _i^n(s,σ)b(s,σ)\,dσ (26) Let Vni(b)≔∑s∫Σαin(s,σ)b(s,σ)σV_n^i(b) _s _ _i^n(s,σ)b(s,σ)\,dσ and b1,b2b_1,b_2 be two joint state-domain beliefs. Let λ1,λ2∈ℝ _1, _2 , then: Vni(λ1b1+λ2b2)=∑s∈∫Σαin(s,σ)(λ1b1(s,σ)+λ2b2(s,σ))σ=λ1∑s∈∫Σαin(s,σ)b1(s,σ)σ+λ2∑s∈∫Σαin(s,σ)b2(s,σ)σ=λ1Vni(b1)+λ2Vni(b2) splitV_n^i( _1b_1+ _2b_2)&= _s _ _i^n(s,σ)( _1b_1(s,σ)+ _2b_2(s,σ))\,dσ\\ &= _1 _s _ _i^n(s,σ)b_1(s,σ)\,dσ+ _2 _s _ _i^n(s,σ)b_2(s,σ)\,dσ\\ &= _1V_n^i(b_1)+ _2V_n^i(b_2) split (27) Therefore Vni(b)V_n^i(b) is linear in b. Since, for any belief b, the value function Vn(b)V_n(b) is given by the Vni(b)V_n^i(b) with the largest value, and because there is a finite number of linear functions Vni(b)V_n^i(b), it follows that Vn(b)V_n(b) is piecewise linear. Since linear functions are convex and Vn(b)V_n(b) is obtained as the pointwise maximum of linear functions, Vn(b)V_n(b) is convex as well. ∎