Agents Robust to Distribution Shifts Learn Causal World Models Even Under Mediation
Authors: Matteo Ceriscioli
Karthika Mohan
Venue: NeurIPS 2025 poster ()
Keywords: Causality
POMDP
Multi-agent systems
Primary area: probabilistic_methods
TL;DR: We demonstrate how agents robust to domain shifts can infer the causal model of the environment in mediated tasks, multi-agent settings and sequential decision tasks.

Abstract: In this work, we prove that agents capable of adapting to distribution shifts must have learned the causal model of their environment even in the presence of mediation. This term describes situations where an agent's actions affect its environment, a dynamic common to most real-world settings. For example, a robot in an industrial plant might interact with tools, move through space, and transform products to complete its task. We introduce an algorithm for eliciting causal knowledge from robust agents using optimal policy oracles, with the flexibility to incorporate prior causal knowledge. We further demonstrate its effectiveness in mediated single-agent scenarios and multi-agent environments. We identify conditions under which the presence of a single robust agent is sufficient to recover the full causal model and derive optimal policies for other agents in the same environment. Finally, we show how to apply these results to sequential decision-making tasks modeled as Partially Observable Markov Decision Processes (POMDPs).

Agents Robust to Distribution Shifts Learn Causal World Models Even Under Mediation
Matteo Ceriscioli, Karthika Mohan
School of Electrical Engineering and Computer Science (EECS)
Oregon State University Corvallis, OR 97331, USA
Abstract
In this work, we prove that agents capable of adapting to distribution shifts must have learned the causal model of their environment even in the presence media- tion. This term describes situations where agent's actions affect its environment, a dynamic common to most real-world settings. For example, a robot in an indus- trial plant might interact with tools, move through space, and transform products to complete its task. We introduce an algorithm for eliciting causal knowledge from robust agents using optimal policy oracles, with the flexibility to incorporate prior causal knowledge. We further demonstrate its effectiveness in mediated single- agent scenarios and multi-agent environments. We identify conditions under which the presence of a single robust agent is sufficient to recover the full causal model and derive optimal policies for other agents in the same environment. Finally, we show how to apply these results to sequential decision-making tasks modeled as Partially Observable Markov Decision Processes (POMDPs).
Introduction
Consider an Al system designed to control a sprinkler. The goal of this system is to maintain optimal vegetation conditions while minimizing water consumption. Such a system adapts to seasonal changes and shifts in the rainfall distribution in order consistently achieve optimal results across diverse weather patterns. This system can be viewed as an agent since it is an entity that maps percepts (e.g. humidity, month of the year) to actions (activation of the sprinkler), and maximizes expected utility (promoting vegetation growth while conserving water) [Russell and Norvig 1995). Now, suppose we are tasked with developing an Al system control robot responsible for cleaning the windows a building. One might notice that there is an overlap between the environment in which the sprinkler controller operates and that of the window-cleaning robot. For example, along the Mediterranean coasts of Europe, southerly winds such as the sirocco and libeccio transport dust and sand from the Sahara, which, when mixed with rain, settle and leave persistent stains on glass surfaces. Therefore, meteorological factors such as wind direction and rainfall occurrence also play an important role in the cleaning robot decision task. This raises the question of whether the adaptability of the sprinkler controller entails the acquisition of some form of weather-relater knowledge, and whether this knowledge, possibly causal in nature, can be transferred to help us develop robust policies for the window-cleaning robot. In this paper, we show that agents capable of behaving optimally under distribution shifts must indeed possess causal knowledge about their environment, and that this knowledge can be elicited and reused to derive better policies for other agents. Returning to the sprinkler example, its adaptability reflects causal knowledge about the weather, which can be extracted and used to improve the policy of the window-cleaning robot.
Investigating the link between causal understanding and robustness shifts in the environment is essential for building AI systems capable of generalizing across diverse scenarios [Pearl 2018
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

While traditional machine learning excels at pattern recognition within fixed distributions, they often struggle when faced with distribution shifts or interventions that alter the underlying system dynamics. This issue has been extensively studied through diverse methodologies including domain adaptation [Ben-David et al. 2006], transfer learning [Pan and Yang. 2010 Zhuang et al. 2021], federated learning [Konečný et al. 2016), and transportability [Pearl and Bareinboim 2011]. each addressing distinct flavors of the problem. Modern AI systems are expected to meet sev- eral key requirements, including robustness to distribution shifts, reliable generalization, transparent decision-making, and avoidance of unintended consequences [Amodei et al., 2016 Hendrycks et al. 2021 (Shah et al. 2025). Causal models offer a powerful framework that addresses these challenges by providing a formal representation of the mechanisms governing an environment [Pearl 2009]. Causal modeling enables AI systems to generalize more effectively by capturing the underlying causal relationships that persist across scenarios [Schölkopf et al. 2021]. and it enhances system explainability by supporting both interventional and counterfactual reasoning (Bareinboim al.) 2024). Yet, a fundamental question remains: Is causal knowledge truly necessary for generalization, can agents achieve robustness without causal understanding of their environment?
Schölkopi 2022].
Recent work [Richens and Everiti, 2024] has demonstrated that agents capable of adapting to distribution shifts must have learned a causal model of their environment, i.e., if the environment can be represented by a causal model, the model's causal structure can be reconstructed by querying an agent capable of behaving optimally under distribution shifts. However, these results focus single-agent, non- non-sequential tasks, and rely on the strong assumption of no mediation [Pearl 2009]] meaning the agent's 's actions cannot have an effect the utility via environment states. In contrast, many real-world Al applications involve tasks where mediation exists. For example, an autonomous navigating from point A to point B, may affect lane occupancy and, in turn, traffic flow and the behavior of other drivers.
The goal of this paper is to extend this theoretical framework to the more general mediated case and explore the implications for multi-agent systems and sequential decision tasks. See Appendix [A.4] for a comparison with Richens and Everiti (2024). This work makes the following key contributions:
1. We demonstrate that agents robust to distribution shifts must have learned a causal model of their environment, even when the no-mediation assumption in Richens and Everitt (024 is relaxed.
2. We present an algorithm to learn the Causal Influence Diagram (CID) depicting mediated deci- sion tasks, by querying optimal policy oracles. We outline how to incorporate prior knowledge into the causal model (Section 3.1
3. We offer insights into the implications of our findings for multi-agent environments (Section 3.2) and for sequential decision tasks modeled with Partially Observable Markov Decision Pro- cesses (POMDPs) (Section?) 3.3
2 Preliminaries
We are interested in representations of decision tasks that capture both the independence structure and the causal relationships among environmental variables, while explicitly specifying the agent's role, including its observable inputs and the variables subject to its control.
Causal Influence Diagrams (CIDs) [Heckerman] 1995| Everitt et al.) 2021] provide representation that naturally fulfills these requirements. Similar ar to Influence Diagrams (Howard and Matheson 1984), CIDs commonly used to reason about decision-making tasks. CIDs further assume that the graph encodes the causal relationships between the nodes. We denote the set of parents of a node X as Pax, the set of children as Chx,, the set of ancestors Ancx, the set of descendants Descx and instantiations of random variables in lower-case. Each variable X has a finite set of possible values (X).
Definition (Causal influence diagram [Heckerman 1995 Everitt et al. 2021)). A causal influence diagram (CID) is a Causal Bayesian Network M = G= E). P), where P is joint probability distribution compatible with the conditional independences encoded in G. The nodes in V are
partitioned into decision (D), utility (U'), and chance (C) nodes, V = (D, U, C"). Each utility node
U, associated with a real function fa of its parents fa domm(Pa; - RR which is referred as the utility or loss function associated with U,.
2

---

Figure 1: A CID representing mediated decision task with {U}. D (D), and C = {M, R, W), where M denotes the month of the year, R indicates rainfall, and W the lawn wetness. An Al agent controls a sprinkler (D) with the goal of keeping the grass wet ( 1). Directed edges indicate causal relationships except the edge that enters D. The utility function U depends on both wetness and sprinkler use, rewarding W = 1 and discouraging water wastage.
M
R
D
W
The chance nodes correspond to variables describing the environment. A decision task represented with a CID containing a decision node D and a utility node U said to be mediated if DeseD n Ancy + 0. that is, if a decision influences some part of the environment that is relevant to the task. For example, a doctor's treatment choice (decision D) alters a patient's condition (environment variable C), which affects recovery (utility U). Edges pointing into decision nodes are called informational links, they indicate which variables the agent observes when making a decision. All other edges causal links, representing causal relationships. We prove that agent capable of adapting to distribution shifts must have learned the CID of its environment by showing that it is possible to recover the causal structure and the Conditional Probability Tables (CPTs) of the variables describing the environment by observing the agent's 's optimal policies under distribution shifts. Note that this setup is unsuitable for traditional causal discovery algorithms like PC (Spirtes and Glymour [99] and FCI [Spirtes) 2001] because we do not have access to the joint probability distribution of the variables or any sample data. In this paper, the terms adaptable agent and robust agent refer to an agent that maintains optimal performance across all the possible distribution shifts in its environment.
different set of decision nodes and have access distinct subset of observable variables. When making a decision D, the agent only observes the variables that are parents of D. An example of a CID modeling the sprinkler story be found in Figure T
Each agent may correspond
Modeling distribution shifts. A key concept in this setting is domain dependence [Richens and
Everiti 2024), meaning that no single policy is optimal under all possible distribution shifts.
Definition (Domain dependence [Richens and Everitt 2024)). A CID M is said to satisfy domain dependence if there exist P(C c) and P'C c). both compatible with the CID M such that π²" arg max, [U] = π arg max, Ep, [U].
Domain dependence rules out trivial cases where the optimal policy remains the same under all distribution shifts. When it does not hold, any agent that is optimal in the unshifted environment would automatically be robust, making causal reasoning unnecessary for adaptation.
Following the work of Richens and Everitt (024), we represent distribution shifts as mixtures of local interventions. Given random variable X with æ1,... In as possible observable values, local intervention on X is function that maps each observable value I; to a new observable value f(z;) {=1,.... }. In other words, local interventions deterministically reassign random variable's 's outcomes independently of other variables.
Definition 3 (Local intervention [Richens and Everiti 2024)). ). Local intervention a on X involves applying a map the states of X that is not conditional any other endogenous variables, x f(x). The conditional probability distribution of X is modified as follows:
P(xl o)
Σ
P(x' paax
(1)
We use the notation o = do(X = f()) indicating that variable X is assigned the value f(x). Local interventions are subset of soft interventions, with hard interventions special case.
local intervention has limited capacity to model distribution shifts. For instance, it cannot model the shift from coin that always lands on heads to a fair coin because a local intervention must deterministically map the observable value 'head' to another observable value. Therefore, we now report the concept of a mixture of interventions [Richens and Everitt 2024). This mixture is a convex combination o² = Σ, P:o; of interventions o, where each coefficient P. represents the probability that is used to map the observable value of X.
3

---

Definition 4 (Mixture of interventions [Richens and Everiti, 2024 ) A mixture of interventions
o² Eipio, p for E. P: = 1 performs intervention o, with probability pN. Formally:
P(x pazio") Ep. paz:
(2)
We consider mixtures of local interventions the environment variables, that is, the chance nodes of the CID. An example of how a mixture of local interventions can represent distribution shift is provided in Appendix D
We use optimal policy oracles to formalize the agent's understanding of optimal behavior under distribution shifts. A policy for decision D is a set of conditional distributions on dom(D) given each possible instantiation of its parent variables Pa(D). Let D be decision variable taking values in d E dom( D). Given set E of interventions on the environment variables corresponding to chance nodes, an optimal policy oracle is map that for any intervention (distribution shift) a E returns the corresponding optimal policy To (D PAD).
Definition 5 (Policy oracle [Richens and Everitt 2024)). A policy oracle for a set of interventions is map 1ly a + To Pap) Vo E E. A policy oracle 11; optimal if for each a Σ, (1) T. (D) Pap)= = arg max, LEU do(D = π (pan)));
As simple example, consider game where an agent has to guess the outcome of die roll X, the utility function is U = 1 if D = X, and 0 otherwise. Let a be an intervention that corresponds to distribution shift that makes the die always land on 3. Then, querying the policy oracle with would return policy n. (D) so that (D 3)
In most real-world settings, such ideal agent may exist. However, previous results have shown that for unmediated tasks, relaxing the optimality assumption for the policy oracle still requires the agent have learned an approximate causal model [Richens and Everitt 2024]. This suggests that the necessity to learn an exact or approximate causal model is stable property that holds even for suboptimal agents.
In our work, we rely on Algorithm from Richens and Everitt 2024 which takes as input a utility function U, optimal policy oracle, intervention E Σ, and parameter N that controls the number of samples. For any local intervention a E Σ, let d be the deterministic optimal decision under the shift induced by a. By domain dependence, there exists a hard intervention o' such that is no longer optimal, with d2 as the new optimal decision. Considering the mixture o(4) == ga + (1 (1-10) -g)o', there exists value Gerit for q such that d2 and another decision d, are both optimal. The algorithm returns qerit. d1, and d2. The value qerit is used to compute the CPTs of the chance variables which
we use to reconstruct the causal graph. For an example see Appendix E.] and Figure 7
3 Learning Causal Influence Diagrams with Adaptable Agents
we present LearnCII,, an algorithm for recovering the environment's 's CID from an adaptable agent, modeled as an optimal policy oracle. This algorithm establishes that acquiring causal knowledge is necessary for agents to be robust to distribution shifts, even mediated decision tasks. In Section 3.2 we discuss the implications our results for multi-agent environments. Finally, in Section 3.3 we outline the relation between CIDs and POMDPs and explain how our approach extends to sequential decision tasks modeled by POMDPs.
In Section 3.1
3.1 LearnCID algorithm for eliciting causal knowledge from adaptable agents
Under specific assumptions detailed below, the LearnCID algorithm (Algorithm I enables the reconstruction of the underlying causal model by identifying the CID structure and the CPTs for the variables corresponding to chance nodes that are not children of the decision node.
Assumption pertaining causal discovery
A1 The CID satisfies faithfulness (Spirtes et al., 1991 and causal sufficiency Pearl 2009].
Faithfulness means that every conditional independence in the joint probability P of the CID also holds in the graph G. A set of variables in a CID is sufficient it includes all common causes.

---

Assumptions pertaining to the CID structure
A2 Given the CID M (G (V, F.); P), the partition of V into (D,U, C) is known.
A3 The CTD contains exactly one decision node D and one utility node U ( = |0] =))
For CIDs with multiple utility nodes, the optimal policy oracle depends on both the specific CID and decision node, so different utility node selections might correspond to different optimal policy oracles.
run the algorithm once for each utility node. Note that
A4 The set of parents of both D and Chp, and the CPTs of the children of D are known. Motivations for Assumption A4 can be found in Appendix B
A5 D is parent of U.
A6 All chance nodes are ancestors of U. Chance nodes that are not ancestors of U (and consequently, since D) E Pay, not ancestors of D) have no influence on the decision task. LearnCID would ignore these nodes, and their associated causal structure and CPTs would not be recovered.
Assumptions pertaining to utility evaluation and decision policies
A7 The utility function f associated with the utility node is fully specified. The utility function's 's functional form is known, which tells us all the variables involved in calculating the utility. These variables appear in the causal graph as parents of the utility node.
A8 We have access to set E of all possible mixtures of local interventions, along with the optimal policy oracle IT's for decision node D)
A9 There exist no decision d' E dDm(D) that is optimal for all instantiations of Pay \ D. Assumption
A9 states that there does not exist single decision d Dom(D) that is optimal regardless of the value taken by the chance nodes that are parents of U. If Assumption holds. Assumption A9 is testable by substituting different value combinations for the parent nodes of U. While Assumption A9 is equivalent to domain dependence (Definition () in unmediated decision tasks. this equivalence breaks down in the general mediated case. Nevertheless, Assumption A9 sufficient to guarantee domain dependence.
The set of assumptions in Richens and Everitt 2024 identical to that above, except we drop the unmediated task assumption. Under the no-mediation assumption, Assumption A4 correspons to knowing the parents of D. A9 corresponds domain dependence, and Assumption A4 together with Assumption A5 correspond to assuming the knowledge of the parents of D and U. In the unmediated task case, Assumption A5 is implied by domain dependence (Richens and Everitt 2024
In this setting, Assumption A2 implies that the variables comprising the causal model are known. Beyond domain knowledge, there exist techniques to learn these variables; for instance, substantial work in causal representation learning aims, among other goals, to identify latent variables for causal models directly from data (Schölkopf et al. 2021 Variei et al. 2024
The following theorem states that for CID satisfying A9 domain dependence also holds.
Theorem 1. Let M (G, be CID where Descp
Ancy *
A9
Domain dependence.
Proof. See Appendix
While under mediation domain dependence implies A9 Richens and Everiti 2024), this no longer holds in the mediated case. T shows that (A9 equivalent in the unmediated setting. This is further discussed in Appendix C
119iiliies domain dependence, making the two
3.1.1 Algorithm formulation
The LearnCID algorithm reconstructs the structure and CPTs of the CID by iteratively exploring chance nodes connected the utility node U. For each chance node X with known path U, we first identify set of candidate parent chance nodes Pax, also making use of prior knowledge: Vkwn the subset of chance nodes whose parents are fully known, and G' graph containing known edges. Consider the following local intervention:
E, x
if otherwise
fx(X
5

---

Algorithm LearnCID
Nodes = (1D), {0}; CJ. graph G' on V containing known edges, of chance nodes with all known parents Vkwn. set of known CPTs O', and the number of samples N to estimate Jerit- Output: The CID's graph G, and the set of CPTs θ
Input:
1: Mark all chance nodes as unvisited.
2: while there are still unvisited chance nodes with a known path to U, starting from the parents of
U not children of do
3: Let X be unvisited chance node with known path to U.
4: Let Ip be the set of intermediate chance nodes on a directed path p from X to U, or to D
such path to U exists.
5:
Let Cx be the set of chance nodes not in Lp
6:
if X VA then Pa'x +- Pax else Pax t Cx (Dscxx
7:
for each instantiation x of do for each X, Pa' do
8:
9:
for each instantiation of variables e₁ in Cx \ ((Paxx Descx) do Let C2 be valid instantiation of the variables in Cx n Descx Construct c instantiation of Cx ) Pax,, by combining e and C2 Initialize raax. (c) as in Equation for each a E Paxx do
10:
11:
12:
13:
14:
Estimate qerit: d1, d2, pa'y using ALG (U, N. ar). Compute P(X x paz; using and Equations and if Pa'x: - G and Ba, E Eoaax s.t. P(x pax: 0 P(æ: pax; o') then Add Pa + X to G'
15:
16:
17:
18:
for each pax E dom((ax) do
19:
P(x pax) + P(x pax;o) for any hard intervention a compatible with pax. return the updated CID's graph G", and the set of CPTs O'.
20:
where a' is arbitrary observable value for X different from Let Cx represent the set of all chance nodes except X and the intermediate nodes on directed path from X either U or D. For example, in CID X4 X3 X2 - X1 U <- D, then Cx, {Xa) because X2 X2 is on the only directed path between X3 and U. For each candidate parent Pax,; E Pax, we define a family of interventions '
Paax. (c) {do(Pa',, pa'x, Cx
=cx= paxx; E dom (4)
Note that in these interventions, all variables to which we assign a specific value are fixed except one, Pax,; which varied across different interventions. To compute the CPTs, we use AL. (Algorithm 1 in Richens and Everitt [024]). Given intervention d and the corresponding optimal decision d1 obtained from the policy oracle, the algorithm finds another intervention distribution shift) 02 under which dj is no longer optimal. The existence of such an intervention is guaranteed by A9 It then defines a mixture of interventions (g) == qo1 + (1 - g)o2 and identifies Gerits the convex coefficient at which dj performs as well as optimal decision de under a2 In other words, at Gerit the expected utilities of dj and d2 under a (qont) are equal. Rearranging this equality leads to Equations 5aand which use the value of qai compute the interventional distributions.
Let Ct...., Cs be the intermediate chance nodes in directed path from X to U or D. If Ci E Pay let C = (Ci,..., C.... otherwise let C = (C2, Ck)- For both x and x', we compute:
8(x) Σ P(s pac, (U U cd) cEdom(()
(5)
Observe that the right-hand side of Equation Jdependds on X for determining the set containing the path of chance nodes C, and depends on the specific instantiation of X because X is the parent of either chance node C, the decision node D, or the utility node U. Let palz be the instantiation of the variables associated with the chance nodes that are parents of U under σ2. Using Equation

---

can compute P( pax:o) as:
(1- - 1 | (d2, (d2,pd)) U (d1,, pab)]] 3(x')
P(x pax;o)=
Derit
(6)
3(x) (x')
We proceed by testing whether any pair of interventions produces distinct interventional conditional distributions for X. If so, the way the interventions in PP3x, are defined ensures that Pa',; is parent of X. We can then recover the observational conditional distribution of X from the interventional ones by noting that, according to Equation T. for any E OPax, that sets Pax to pax, it holds that P(x pax;o) P(x pax). This process is repeated until all chance nodes that ancestors of U have been processed, returning the learned structure and CPTs.
An example of an application of Algorithm [ (LearnCID) to a single-agent CID is provided in Appendix E.1 The proof of correctness for Algorithm I can be found in Appendix A The algorithmic complexity of Algorithm
D is discussed in Appendix A.3
With Theorem 2 we prove that agents robust to distribution shifts must have learned a causal model of their environment even under mediation and formalize the correctness of Algorithm [|
Theorem 2. Let M be CID satisfying Assumptions A9
A9
Suppose for each C. we know
set Pac, C Pac, and set of nodes Viwn S V where C. Viwn + Poc. = Pac.. Let pap)aex denote the set of optimal policies available where (d(pap)) is an optimal policy
under the intervention a and the set of all mixtures of local interventions. Then, for almost all CIDs, the graph G and the joint distribution P over all the ancestors of the utility node Ancy can be identified.
Proof. See Appendix A
U
In Richens and Everitt 2024). it is claimed that the set of optimal policies is almost al- ways sufficient to identify all causal relationships, where almost always means except on set of parameters of Lebesgue measure zero. While this holds for unmediated tasks, we show that in the mediated case, the children of the decision node D cannot, in general, be identified, which in turn motivates Assumption A4
Theorem 3.
Let M = (G (V,E).)) P) be single decision/single utility CID. assume we know D.U, C, and Pa(). Let E be the set of all mixtures of local interventions, IIş be an optimal policy oracle for D. Then in general Chp cannot be uniquely determined.
Pa(D).
Proof. See Appendix
B
3.2 Implications for multi-agent environments
In this section we show how to apply our results to multi-agent systems represented by multi-decision CIDs. The fact that adaptable agent must necessarily learn a causal model of its environment has significant implications when extending this concept to multi-agen agent systems. Under the assumptions of Algorithm [ the causal model of the nodes in Ancy, the variables influencing the utility function, can be recovered. Now suppose we introduce new agent A' into the same environment, by adding its decision nodes and related edges. The existing causal knowledge over U Ancy can, in principle, inform A' and aid in designing policies that make it adaptable to distribution shifts [Bareinboim al., 2024]- Let U' denote the set of utility functions that A' tries to optimize. Intuitively, the extent to which the original agent's 's causal knowledge can be transferred to A' depends on how much of the environment they share, more precisely, on the intersection Ancy n (Utret Ancy)). We can expect this transferability to be maximal when Ancy = Ureu Ancy. and minimal when Ancy n (Urre Ancyy) = 0. As an example, Consider a factory with two agents: agent A, the production manager, determines how much to produce to maximize profit, while agent B, the maintenance manager, checks the machine conditions and schedules maintenance actions in order to avoid breakdowns and maximize uptime. In this example Ancu, n Ancu,, , is not empty, since breakdown reduces uptime and could prevent agent A from producing the desired quantity of goods. Agent A's causal knowledge about breakdowns can be informative for agent B and vice versa.
Learning multi-decision CIDs. Even LearnCID can be applied only to single-decision CIDs, there exist different approaches to handle multi-decision CIDs. We propose two approaches for when
7

---

M
B
N
DA
M
DA
DB
BB
O
U
DB
0
U
Figure 2: On the left, CID representing a factory pocision-makingg scenario with U {UA, UB}. D = {DA, DB}. and C = {M, B, O, J. Here, M denotes machine condition, B machine breakdown, O the production output, and N the product demand. Agent A (the production manager) controls DA adjusting the production rate in response to N, while Agent B (the maintenance manager) controls Dg, scheduling maintenance actions informed by M. Utilities UA and UB reflect production profit and machine uptime, respectively. On the right, the corresponding transformed CID used to apply LearnCID, with D as the primary decision node and DB secondary decision node, assuming DB is assigned a faithful policy. Edges marked in red denote unknown dependencies.
all agents are able influence the same utility function and where at least agent optimizes it. In both approaches the optimizing agent may correspond to multiple decision nodes. We designate one of these as the primary decision node and require the single-decision assumptions hold for it (knowledge of its children, their CPTs, and its parents), All remaining decision nodes are referred to as secondary decision nodes. In the example described in Figure 2 assuming agent is robust, its decision node DA is designated as the primary decision node and Dg as secondary decision node. CID contains multiple utility nodes, we can prune all but one and then run LearnCID to recover its ancestors and their CPTs. It is then possible to repeat the procedure for other utility nodes in the environment, provided that for each such node there is at least one optimal policy oracle available corresponding to a decision node that optimizes the associated utility function and that the other LearnCID assumptions are satisfied.
If the
1. Let { be the set of all secondary decision nodes. possible approach consists in assuming that the decision associated to each secondary decision node D; is determined by a faithful policy TT. i.e. a policy that actually depends on Pap., Since the policy is fixed, we can convert these decision nodes to chance nodes with the corresponding, potentially unknown, policy as CPT. If we assume the availability of an optimal policy oracle for the primary decision node, we can apply Algorithm to recover the full CID.
As a special case, if we know the parent sets of all secondary decision nodes {D:). It is then possible for to assign any faithful policy to each decision node D;. Then, we can again convert these nodes to chance nodes by using their policies as CPTs. Assuming the existence of an optimal apply Algorithm to learn the full CID.
policy oracle for the primary decision node,
Once the CPTs for all chance nodes are learned, it becomes possible determine optimal policy for every decision node in the original graph under any distribution shift (Shachter 1986 Koller and Milch 200T Hammond et al., 2021 Bareinboim et al. 2022 In Appendix H.2 we provide an example of an application of LeamCID to a system where two agents cooperate.
3.3 Applying LearnCID to Partially Observable Markov Decision Processes
Let M = (S.A.T, R 2, 7,7) be Partially Observable Markov Decision Process (POMDP) Kael- bling et al. 1998 where S is the set of states, the set of actions/decisions, T S x A + II(S) i is
the state-transition function where II(S) is the set of probability distributions over the set of states, the set of observables (values of observable variables), O is the set of conditional observation probabilities, and 7 E [0,1) is the discount factor.
R x A + RR is a reward function,
POMDPs form a class of models for sequential sscision-making that typically involve mediation, as state transitions generally depend on the agent's 's actions. In this section, we demonstrate that under mild assumptions, given a POMDP with unknown state-transition function T and states described by a finite set of discrete variables, POMDP can be represented with a CID with temporal component. Given access to an optimal policy oracle, Algorithm| can then be used to learn both the intra- and inter-temporal causal structure, well the CPTs and the state- caatetrranstion function.

---

Structural constraints.
We make the following assumptions on the POMDP:
Pl homme-eotggentty, i.. the transition probabilities do not change over time.
P2 The set of observable variables is fixed and the agent gets an observation from these variables at every timestep.
Formally, each state in S is described by a finite set of variables V (i.e.. there exists bijective function f dom(V) + S), there exists subset of observable variables Vo S V. and for every state in S and every action a in A it holds that O a, vo) = 1, where % is the instantiation of the observable variables at the current state.
we can model the state space of POMDP using CID with Paα == Vo. If the set of observable variables coincides with the set of all state variables then get Markov Decision Process (MDP) as special case of POMDP.
Note that under Assumptions PI and P2
Causal representation of POMDPs. By identifying the observable variables of the POMDP as the parents of decision nodes, the reward function as the utility function associated with a utility node, and all other unobserved variables chance nodes that are not parents of decision nodes, we can observe that POMDP can be seen as CID unrolled over time. Previous work has explored analogous causal representations of POMDPs [Everitt et al. 2019 Bareinboim et al. 2024].
The mapping from POMDP to a CID with unknown structure described in Algorithm Since we want to apply Algorithm[ to this CID, we further need to assume each variable is discrete, the POMDP reward directly depends on the action (equivalent to E Pay)). In this case, we have slightly weaker condition for the Markov blanket of the decision node: other than the parents of D, we only need to know which chance nodes are children of D in the same timestep, their CPTs, and their parents, while we can ignore the chance nodes that are children of D in the following timesteps (e.g.. considering the POMDP described in Figure 3 we do not need to know that X a child D. because they belong to different timesteps). Due to the Markov property of POMDPs, direct causal relationships can exist only within a single timestep or, when inter-temporal, from past to future between consecutive timesteps. Note that the causal structure, the CPTs, and the state- transitran function associated with the POMDP are independent CID obtained with Algorithm 2 does not contain any information about the original discount factor, as long as the assumptions of LearnCID are satisfied, we can still recover these components.
the discount factor. Therefore, even if the
Algorithm 2 POMDPtoCID
Input: POMDP A, R. S. O, 7) with unknown state- tratts-tannition function T, states described by a finite set of variables V. and the subset of observable variables V. S V.
Output: Corresponding CID's graph G without causal arcs involving chance nodes.
1: Initialize G as empty graph.
2: Add decision node D; to CID's graph G with A as set of decisions. Add chance node Di-1 to CID's graph G.
4: for each variable V E V do
5: Add chance nodes corresponding to random variables V and Vi-1 to CID's graph G.
6: if V E V. then
Add edge V, + D, and Vi-1 + D. CID's graph G. 8: Define LI :oon())))) (V) DD) - as U(v, d)- d) for all d D) and u e dVm
7:
9:
Add utility node Ue with utility function 2/ to G.
10: Add edge D. + U2 with utility function 11 to G.
return CID's graph G.
Learning intra- and inter-temporal causal relationships. We can run LearnCID on the CID produced by Algorithm learn the causal relationships for variables both at the same time step and at different ones. The algorithm maps the variables, decisions, and rewards related to two consecutive timesteps to nodes of a CID. We now refer to these timesteps as / and L. Therefore, for each variable X partially describing the state of the POMDP, this newly defined CID contains two variables X,-1 and X,. Similarly, there will be two decision nodes D,-1 and D, and two utility nodes U,-1 and U.. Observe that prune U,-1, and, similarly to what proposed for multi-agent settings,

---

Figure 3: Example of application of Algo- rithm 2 and Algorithm| [ (LearCIID) to time- homogeneous POMDPs. On the left, causal representation of a POMDP. On the right, the cor- responding CID with unknown edges marked in red. To apply LearnCID, we consider two time steps 1-1 and with decision node Dr and utility node v We select any faithful policy for D:-1, convert it into a chance node, and prune Ut-1. Time-homogeneity ensures that inter-temporal causal relationships are preserved across all time steps, allowing us to recover the full causal graph.
Y2
Y2
Y2
X
X
X
D
De
Di-1
Dz
Ut-1
Ue
U₂
we convert D to a chance node and assign any faithful policy its CPT. This is possible because the chance nodes that are parents of decision nodes correspond to the observable variables which are known. Then, if the assumptions of LearnCID are satisfied, including having an optimal policy oracle available, then we can use it to learn all the missing causal relationships. For time-homogeneous POMDP, like in the example illustrated in Figure 3 this is sufficient to learn all the intra- and inter-temporal causal relationships of variables in Ancy, for any timestep since they do not change when varying t because this would imply a change in the transition probability which contradicts the homeehotggen assumption.
Learning the state-transition function. Once we have learned the CPTs for all the chance nodes and we have fixed policy π for the agent, it is possible to recover the state-transition function T, i.e. for all states s E S and actions a E A we can find probability distribution over the state of the process at the following timestep. Let (,(0) ..... vi(0) be the variables V associated with the chance nodes at timestep we observe that:
T(s((-1), (l(-1)) P(s(0))s((-1) (a(11)) = P(V
Vi(1) pu (1-1)
1 (((-)))
II P(V!!) (pav, (d
II CM
P(V. ppav))
(VECCh dEA
Since all the CPTs are known and the policy is fixed, we can compute the state-transition function for all states and actions.
Non timo-homogeneous case. For non time-homoeencous POMDPs, this method only recovers causal links between two adjacent timesteps, time-varying dynamics can change variable inter- actions. For example, irrigation may affect soil humidity in the dry season but not during the rainy season, so applying the algorithm on two rainy days could falsely suggest no causal link. One fix is to augment the state space to enforce time homogeneity [Puterman| 2005]. Alternatively, the method can be applied separately periodic. Otherwise, it is possible to map each interval to a CID, prune intermediate utility nodes, and use our multi-decision CID approach. Each strategy has different oracle requirements: the augmented model needs time interventions, while the others require one oracle per time interval.
each timestep. which is practical when time dependence is finite or
4 Conclusions and Future Work
In this work, we addressed the challenge of understanding the relationship between robustness to distribution shifts and an agent's causal understanding of its environment. While previous work estab- lished that robust agents encode the causal model in single-agent, unmediated tasks, we demonstrated that this connection also holds in mediated, multi-agent, and sequential settings. We presented an algorithm to show that is possible to elicit the learned causal model from robust agents. In multi- agent systems, we showed how a single robust agent enables the discovery of the complete causal model, and how this could be used to learn optimal policies for other agents in the same environment. We further applied our approach to POMDPs, demonstrating that robust agents necessarily learn intra- and inter-temporal causal relationships and the state-transition function. These findings contribute to a theoretical foundation for world modeling approaches based on the extraction and combination of causal knowledge from robust agents and facilitate the exploration of approximate settings, where the agent's 's optimality assumption is relaxed.
10

---

References
Dario Amodei, Chris Olah, Jacob Steinhardt, Paul Christiano, John Schulman, and Dan Mané. Concrete problems in ai safety. 06 2016. doi: (0448550arX11.16666.66
Elias Bareinboim, Juan D. Correa, Duligur Ibeling, and Thomas Icard. On Pearl's Hierarchy and the Foundations of Causal Inference, page 507-556, Association for Computing Machinery, New York, NY, USA, 1 edition, 2022. ISBN 978 170144939.8..
Elias. Bareinboim, Junzhe Zhang, and Sanghack Lee. Towards causal reinforcement learning. Technical Report R-65, Causal Artificial Intelligence Lab, Columbia University, December 2024.
Shai Ben-David, John Blitzer, Koby Crammer, and Fernando Pereira. Analysis of representations for domain adaptation. In Proceedings of the 19th International Conference on Neural Information Processing Systems, NIPS'06 06, page 137-144, Cambridge, MA, USA, 2006. MIT Press.
Tom Everitt, Rahul Kumar, Victoria Krakovna, and Shane Legg. Modeling agi safety frameworks with causal influence diagrams. In Workshop on Artificial Intelligence Safety, volume 2419 of CEUR Workshop Proceedings, 2019.
Tom Everitt, Ryan Carey, Eric D. Langlois, Pedro A. Ortega, and Shane Legg. Agent incentives: A causal perspective. Proceedings of the AAAI Conference on Artificial Intelligence, 35(13): 11487-11495, May 2021. doi: 10. w300///a....3513....... index Cidee..pp//AA///rriicle//eeeee 17368
URL https://ojs. http///sssaaaaa....aaaaaaaaaa..a org/
Lewis Hammond, James Fox, Tom Everitt, Alessandro Abate, and Michael Wooldridge. Equilibrium refinements for multi-agent influence diagrams: Theory and practice. In Proceedings of the 20th International Conference on Autonomous Agents and MulriAgent Systems, AAMAS '21, page 574-582, Richland, SC, 2021. International Foundation for Autonomous Agents and Multiagent Systems. ISBN 1781330333
David Heckerman. A bayesian approach to learning causal networks. In Proceedings of the Eleventh Conference on Uncertainty in Artificial Intelligence UAr95, page 285-295, San Francisco, CA, USA, 1995. Morgan Kaufmann Publishers Inc. ISBN 1558603859.
Dan Hendrycks, Nicholas Carlini, John Schulman, and Jacob Steinhardt. Unsolved problems in ml safety, 09 2021.
Ronald A. Howard and James E. Matheson. Influence diagrams. Readings on the Principles and Applications of Decision Analysis, Vol. II, 1984.
Leslie Pack Kaelbling, Michael I.. Littman, and Anthony R. Cassandra. Planning and acting in par- tially observable stochastic domains, Artificial Intelligence, 101(1):99-134, 1998. ISSN 0004-3702.
doi: https://doi. htt.aaamccaatttcceeee////////330
URL httpp://wwwt http://www.ccci.... www..scieneediree.oo
Daphne Koller and Brian Milch. Multi-agent influence diagrams for representing and solving games. In Proceedings of the 17th International Join! Conference on Artificial Intelligence Volume IJCAI'01, page 1027-1034, San Francisco, CA, USA, 2001. Morgan Kaufmann Publishers Inc ISBN 1558608125.
Jakub Konečný, H. Brendan McMahan, Felix X. Yu, Peter Richtarik, Ananda Theertha Suresh, and Dave Bacon. Federated learning: Strategies for improving communication efficiency. In NIPS Workshop on Private Multi-Paryy Party Machine Learning. 2016. URL https: http:://axxvvvv. wrg. abs/
1610.05492
Masashi Okamoto. Distinctness of the Figenvalues of a Quadratic form in Multivariate Sample. The Annals of Statistics, 1(4):763 765, 1973. doi: 10.12144ao/1176442477 URL https
//ooi org/ 10. 1214/ aos /ooooo////774477111
Sinno Jialin Pan and Qiang Yang. A survey on transfer learning. IEEE Transactions on Knowledge and Data Engineering. 12910))1345-1159, 2010. doi: 10. 10.1109/KKDE.2099.199

---

Judea Pearl. Causality: Models, Reasoning and Inference. Cambridge University Press, USA, 2nd edition, 2009. ISBN 052189560X.
Judea Pearl. 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, page 3, New York, NY. USA, 2018. Association for Computing Machinery. doi: 10. 145/3199652.316611
ISBN 978 1050355800..
UtLhttps://doi orwgonsttettte
idm. www.clc..ooooo org/ 10. 1141//119/16...... 711.199552...668188 3176182
Judea Pearl and Elias Bareinboim. Transportability of causal and statistical relations: A formal approach. In 2011 IEEE 11th International Conference on Data Mining Workshops, pages 540-547, 2011. doi: 10. 10.110900 10.110//CDM...OO 1.199.
Martin L. Puterman. Markov decision processes : discrete stochastic dynamic programming, 2005.
Jonathan Richens and Tom Everitt. Robust agents learn causal world models. In The Twelfth International Conference on Learning Representations, 2024. URL httpps:://ppeereeiie net/
forum?id=pOok
S. Russell and P. Norvig. Artificial Intelligence: A Modern Approach. Series in Artificial Intelligence. Prentice-Haud Hall, Englewood Cliffs, NJ, 1995.
Bernhard Schölkopf. Causality for Machine Learning, page 765-804. Association for Computing Machinery, New York, NY, USA, edition, 2022. ISBN 978 450395861.
Bernhard Schölkopf, Francesco Locatello, Stefan Bauer, Nan Rosemary Ke, Nal Kalchbrenner, Anirudh Goyal, and Yoshua Bengio. Toward causal representation learning. Proceedings of the IEEE, 109(5):612-33) 2021. doi: 10.1100/PROC.222133359999
Ross D. Shachter. Evaluating influence diagrams, Oper. Res., 34(6):871-882, December 1986. ISSN 0030-364X.
Rohin Shah, Alex Irpan, Alexander Matt Turner, Anna Wang, Arthur Conmy, David Lindner, Jonah Brown-Cohen, Lewis Ho, Neel Nanda, Raluca Ada Popa, Rishub Jain, Rory Greig, Samuel Albanie, Scott Emmons, Sebastian Farquhar, Sébastien Krier, Senthooran Rajamanoharan, Sophie Bridgers, Tobi Ijitoye, Tom Everitt, Victoria Krakovna, Vikrant Varma, Vladimir Mikulik, Zachary Kenton, Dave Orr, Shane Legg, Noah Goodman, Allan Dafoe, Four Flynn, and Anca Dragan. An approach to technical agi safety and security, 2025. URL httpst::: / aa/iv.xiii org/ abs 2504 01849
Peter Spirtes. An anytmee algorithm for causal inference. In Thomas S. Richardson and Tommi S. Jaakkola, editors, Proceedings of the Eighth International Workshop on Artificial Intelligence and Statisties, volume R3 of Proceedings of Machine Learning Research, pages 278 285, PMLR, 04-07 Jan 2001. BtLLhttps:// htttps://proceeddngss mlr pptpss/r//spitteeeeeeeeeeeeeeeee by PMLR on 31 March 2021.
html Reissued
Peter Spirtes and Clark Glymour. An algorithm for fast recovery of sparse causal graphs. Social science computer review, 9()):22-7, 1991.
Peter Spirtes, Clark Glymour, N Scheines, et al. Causation, prediction, and search. 1991.
Burak Varici, Emre Acartürk, Karthikeyan Shanmugam, and Ali Tajer. General identifiability and achievability for causal representation learning. In Sanjoy Dasgupta, Stephan Mandt, and Yingzhen Li, editors, Proceedings of The 27th International Conference on Artificial Intelligence and Statistics, volume 238 of Proceedings of Machine Learning Research, pages 2314-2322. PMLR, 02-04 May 2024. URL /RLLtttp:///proceddiiig
mlr prw.s/ wwt.:/ppesss//eeeeeeeee w238/varici24s html
Fuzhen Zhuang, Zhiyuan Qi, Keyu Duan, Dongbo Xi, Yongchun Zhu, Hengshu Zhu, Hui Xiong, and Qing He. A comprehensive survey on transfer learning. Proceedings of the IEEE, 109(1):43-76, 2021. doi: 10. 10.109//PK0C.2003304555
12

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We carefully presented the paper's claims in the abstract and introduction in way that accurately reflects the paper's contributions.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: The paper includes a transparent list of the assumption made and comments about the precise scope of the results and the corresponding limitations.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model sell-speciiiatiinn asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g. if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
13

---

Answer: [Yes]
Justification:
Each theorem includes precise list of assumptions and a complete proof.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proof's can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Question:
Answer: [NA]
The results of this paper purely theoretical, the paper does not include experiments.
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, descríbing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is new model (e-g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing verifying the results.
5. Open access to data and code
14

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [NA]
Justification:
The paper does not include experimental results, nor the use of data code.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips... cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run reproduce the results. See the NeurlPS code and data submission guidelines (https
/ "pipsscc/puullc//uuu /utee//ggudddeeeedSuuuubmbibbi ssionPolicy) for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [NA]
Justification: This paper does not include experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [NA]
Justification: This paper does not include experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
15

---

The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative enror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [NA]
Justification: This paper does not include experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental rums as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported the paper (e.g., preliminary or failed experiments that didn't 't make it into the paper).
9. Code of ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics
Answer: [Yes]
Justification: We read the NeurIPS Code of Ethics. This paper conforms with the code of ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is a special consid- eration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification:
We included a paragraph in the appendix discussing societal impacts.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
16

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
This paper poses such risks.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Lieenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [NA]
Justification: This paper does not existing assets.
Guidelines:
The answer NA means that the paper does existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
17

---

If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: This paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/mmoel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: This paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question:
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: This paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
16. Declaration of LLM usage
18

---

Question: Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods this research? Note that if the LLM used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA]
Justification: We have read the NeurIPS LLM policy. This research does not involve LLMs as any important component.
Guidelines:
The answer NA
that the method development in this research does involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy for what should or should be described.
19

---

A Correctness Proof for LearnCID
A.1 Setup and Assumptions
Our goal is to show that agents capable of adapting distribution shifts must have learned causal model of their environment. To do this we show it is possible to recover the causal structure and the Conditional Probability Tables (CPTs) of the variables describing the environment in which the agents operate. This environment consists of both observable variables and hidden latent variables. We define a set of interventions modeling distribution shifts, and to learn the causal graph, we query an optimal policy oracle associated with one agent to get the optimal policy for that agent under the specified distribution shift. Observe that this setup is unsuitable for traditional causal discovery algorithms like PC (Spirtes and Glymour, [991] and FCI [Spirtes 2001] because we do not have access to the joint probability distribution of the variables or any sample data.
To model the causal relationships in the environment, we use Causal Influence Diagrams (CIDs) Heck erman 1995 Everitt al., 2021]. Similar Influence Diagrams (Howard and Matheson 1984 CIDs are commonly used to reason about decision-makinn tasks. CIDs further assume that the graph encodes the causal relationships between the nodes. We denote the set of parents of a node X as Pax, the set of children as Chx. the set of ancestors as Ancx, the set of descendants as Descx and instantiations of random variables in lower-case.
Definition (Causal influence diagram [Heckerman 1995 Everitt et al. 2021)). [). A causal influence diagram (CID) is a Causal Bayesian Network M = GG = (V, ET, where P is joint probability distribution compatible with the conditional independences encoded in G. The variables in V are partitioned into decision, utility, and chance variables, V = (D, U, C). Each utility node U; is associated with a real function ft of its parents fi : dom((au:) - R.
Each agent may correspond different set of decision nodes and have access distinct subset of observable variables. A variable observed by one agent may be latent for another. Additionally, considering situation where an agent takes more than one decision, the set of variables that observes when it takes one decision can differ from the one that it observes when taking another decision.
Following the work of [Richens and Everitt 2024]. we represent distribution shifts as mixtures of local interventions. Given random variable X with ø1, 21..... Zn possible observable values, local intervention on X is a function a x2 + fxx) that maps each observable value z; to a new observable value (x;)). In other words, local interventions deterministically reassign random variable's 's outcomes independently of other variables.
Definition 3 (Local intervention [Richens and Everitt 2024 ). Local intervention a on X involves
applying a map the states of X that is not conditional on any other endogenous variables, x + f(x). We use the notation a = do(X = fxx) (variable X is assigned the state f (x)) Formally, this is soft intervention on X that transforms the conditional probability distribution as,
P(x paxio)
Σ
P(a' paax
(8)
In general, local intervention has limited capacity to model distribution shifts. For instance, it cannot model the shift from a coin that always lands on heads to a fair coin because a local intervention must deterministically map the observable value 'head' to another observable value. Therefore, we now report the concept of a mixture of local interventions [Richens and Everiti 2024]. This mixture is a convex combination o² Σ, Pio; of local interventions d;, where each coefficient pa represents the probability that o; is used to map the observable value for X.
Definition 4 (Mixture of interventions [Richens and Everitt 2024 ). A mixture of interventions
0' = P.o; for Σ, P: performs intervention o; with probability pt. Formally, P(x o²) Σ, P(x o.).
We use optimal policy oracles to formalize the agent's understanding of optimal behavior under distribution shifts. Let D be decision variable with observable values d E dom(D), given a set of interventions Σ, an optimal policy oracle is map IT. a + ( pop) for E E, where (d pap) is the optimal policy under the distribution shift induced by the intervention o.
20

---

Definition 5 (Policy oracle [Richens and Everitt 2024)). ) A policy oracle for a set of interventions is map 11y; a + no (D Pap) A E Σ. A policy oracle II; optimal if for each a Σ. (T) =,) Pan)= = arg maxx RU do( = π (pan)): a).
In our work, we rely Algorithm from Richens and Everitt [2024]. which takes as input utility function U, an optimal policy oracle, am intervention 0 E Σ. and a parameter NV that controls the number of samples. For any local intervention under the shift induced by a. By Assumption A9 there exists a hard intervention o' such that is no longer optimal. Let d2 be the deterministic optimal decision under a'. Considering the mixture (q) = qa + (1 - q)a', there exist a value ferit for y such that d2 and another decision dj are both optimal. The algorithm returns qerits d1, and d2.
Σ, let d be the deterministic optimal decision
Now, we list and motivate our assumptions.
Assumption pertaining to causal discovery Assumption 1. The CID is faithful (Spirtes et al. 1991] and sufficient [Peall 2009
Faithfulness implies that every conditional independence in the joint probability of the CID also holds in the graph G. A set of variables in causal model is sufficient when it includes all common causes.
Assumptions pertaining to the CID structure
Assumption 2. Given the CID M (G = =V. with V (D,U,C). the set of nodes and the partition (D, U, C)) is known.
The set of nodes together with the node partition (D, U, C) is known, therefore we know all variables in the system and the type of each node (decision, utility, or chance).
Assumption 3. The CID contains exactly one decision node D and one utility node U.
Despite assumption 3] this algorithm can be applied to multi-decision CIDs. It is possible to find further details the main paper.
Assumption 4. The Markov blanket of decision node D is known. We also know all the edges between these nodes. The CPTs of chance nodes that are children of D are known.
Motivations for Assumption f can be found in Appendix B
Assumption 5. D is parent of U.
We assume D is parent of U. In particular, this plays role in the proof of Lemma I Assumption 6. All chance nodes are ancestors U.
In the presence of chance nodes that are neither ancestors of D or U, these nodes do not have any influence the decision task. LearCIDD would simply not process those nodes and the related causal structure or CPT would not be recovered.
Assumptions pertaining to ulility evaluation and decision policies Assumption 7. The utility function associated with the utility node U fully specified.
The utility function's 's functional form is known, which tells us all the variables involved in calculating the utility. These variables appear in the causal graph as parents of the utility node. Assumption 8. We have access to a set Σ of all possible mixtures of local interventions, along with the optimal policy oracle IT for decision node D.
Assumption 9. There exists no d" E dDm such that d' Earg max. Udd, .)) Vr E dom (D)).
Assumption 9 states that there does not exist a single decision d E Dom(D) that is optimal regardless of the value taken by the chance nodes that are parents of U. When Assumption holds, it is possible to verify if the CID we are considering satisfies Assumption 9 by computing the utility for different instantiations of variables associated with the parents of U.
Observe that under Assumptions 3 and 9 there must be at least one chance node that is a parent of U, otherwise the utility function would only depend on the decision and therefore there would exist at least one optimal decision that would violate Assumption We provide a discussion the Assumption 9 and domain dependence in Appendix
21

---

A.2 Proof
Lemma 1. Under Assumptions 5
and 7 given a CID M (G, P), P).jo for any given local interven- tion a there is a single deterministic optimal policy for almost all P. U.
Proof. When there is no other directed path from to U except D + U, i.e. Desep n Ancy (unmediated case) the statement is proven by Lemma 3 in Richens and Everitt (2024). Assume there exist a path from to U other than D U, i.e. Descp Ance * Let Z Ancy \ Pap, and X Pay {D} Assume 3d1, d2 s.t. both are optimal decisions in the context pup, we divide the proof two cases.
Case 1. Assume Chp &
Since d1 and d2 are both optimal we have:
Elu pap, do(D di); = Elu pap, do(D = d2);
(9)
Since IL Pap in Gp. then according to rule 2 of do-calculus equation as:
Pearl 2009
we can rewrite the
Blu pap, D d1;0] Eluu pap, D d2;
(10)
Equivalently:
Σ U(d1, pap, D D= 01;00) Σ U(22, paD, d2;
(11)
Σ U(d,,))P(2 paD, =d;00) Σ Ud2, pap, D 02: =
(12)
Let us define S₁ = ChD and Pac, = Pac, {D}. We factorize the joint distribution:
P(2 pap, D d1;0)
II
P(Z; paz;;))
II
P(C, pac, D = od;a))
(13)
Z,ES1
CjEChp
Same for d2. Now, we can rewrite Equation 12 as:
Σ EI II P(Z paz;;o)
U(d1,
II P(C: | pac,, d);0)
Z:ESI
C,ECND
(14)
(d22
II P(C,pac, D d2: =
CjEChp
We can define ald, 2) = U(d,z) (c,,ECo P(C,Ipac, D = d;
Observe that:
II Z.ES,
P(Zilpaz; - a(d2, =)] =
(15)
Is polynomial equation with variables Vi. P(Z; (azzi Also notice that if we rewrite each of these variables using Definition (Local intervention)
= >={=((=)==. P(=I[paz,) the equation is still a polynomial equation with the parameters of the CPTs, excluding those related to the children of D, variables. If the polynomial is not trivial (3d, a(d (1,2) 2) a(dg, 2) # 0) then the Lebesgue measure of the solution of this equation zero [Okamoto 1973]. and since the set of parameters that allow for multiple optimal solutions has measure zero along at least one dimension, follows that the whole set has measure zero.
Now, let consider when the polynomial is trivial. Then for all d, we have:
22

---

a(d).2) 2) a(d2, 2) 0 (() II P(Cilpac, D= = d1:0) - U(2,, x) C,ECED
(16)
II C,EC
PlCippac,: D d2; =0 (17)
Again, this is finite number of polynomial equations with some of the network parameters as variables and its coefficients are not trivial because di * d2 and because of Assumption 5 Therefore, it is satisfied only set of Lehesgue measure zero [Okamoto 1973].
Case 2. Now assume Chp nC C. then the factorization of Equation simplifies to:
II
P(C, pac, d1;0)
(18)
CEEck
Therefore we can rewrite 12 as:
xd) II
P(C, D = -1:0))U((da))
-00)-UU(e U(22, C.ECkD
II
P(C D d2; 0
GiCcDD
(19)
That again polynomial equation in some of the network parameters and therefore the set of solutions has Lebesgue measure zero. Again, since the set of parameters that satisfies Equation 12 has measure zero along least one dimension, the whole set has measure zero.
This implies that for almost all P. U and any given local intervention a the optimal decision is unique.
In the following lemma, we consider masking as special case of a local intervention. A local intervention deterministically transforms the states of a random variable X via predefined function. Specifically, we can apply a function that maps all states X E Pap to masked state, rendering X unobservable to the agent under this intervention. We denote by Pa'p C Pan the subset of variables that remain unmasked. Consequently, the set Pap \ Pa', represents the variables that are masked and thus hidden from the agent when making the decision D.
Lemma 2. Given a CID M = (G. P), under Assumptions 7 given an optimal policy oracle IIS where E includes all mixtures of local interventions on C including masking inputs Palp S Pan.
then for any given Pa'p = pa'p such that Pu'p n Pay = 0, we can identify:
P(CC=
=d); d);
P(C cldo(D d' (')
(20)
for some d, d' E dD) where + d' and Z = C C.\Pas. Pa'p
1 for almost all P.U there exist only one optimal decision di arg maxa Dlludd/ = d),pa'pi o) following the shift a. The decision di can be identified using the optimal policy oracle (a).
Proof. By Lemma
By Assumption 9 we know that for every decision d in the context Pap C Pap. there exists at least one instance e = (CI,..., CN) of C where d a arg max U(d',z). Note that we can set the values for X as Pa'p n Pay satisfies dj * arg max, Udd )'), and o' be hard intervention that sets X to x', then there exist d2 arg maxa Udd x'), with dy * di. Note that under a hard intervention like o' we have Eju do = d), papp: Ud,, where x' are the values that the variables Pay [D) take after the intervention. We can pick o' such that it sets Pu'p to the same as in observation.
0. So let x' be the instantiation of Pay {D} that
For E 1] consider the joint distribution over C' under the parametrized family of mixed local interventions ā(q) ga + (1 q)o':
P(C cldo(D d):ā(g)) gP(C cldo ( d):o) + (1 PCC cldo ( = d):
(21)
23

---

Table 1: A partition of the CID's chance nodes C that assigns them to cases in Theorem Cells marked with (1) correspond to nodes that can be identified by case of Theorem in Richens and Everitt (2024). Cells with (2) correspond to case 2. Cells with 0 indicate that either no nodes fall within that intersection or that those nodes can be pruned since we are not able to learn the structure or the parameters for those. The cell marked with (3) corresponds to nodes that exist only in the mediated case and were therefore not considered by the previous paper.
(Aned Descp)
Aned
Descp
Ancu (GD)
(1)
(1)
(3)
Ancy (Gp)
(2)
0
From Assumption 9 it follows that Z C Pap + U. We can write the expected utility as:
BaUapaa do = d):ā(1))) - Σ P(Z = zlpan, do = d);
(22)
P(C cldo( = (q))
E
U(d,z) x)
(23)
P(Pap= = ppallo( = d); (()))
1
Σ qP(C cldo( D = d);o)
P(Pan pap; (q))
(24)
+ (1- = c/do(D D = d): (')((zz)))
Where in Equation 24 P P(Pa pap do d); ()) = P(PaD pap; r())) according to
Rule of do-calculus since D Pap in Gb [Peall] 2009]. Note that dj is the optimal decision for Q = 1, but that is not the case for q = 0. Theretore there exists qerit such that for all q qerit d2 == (ā is a decision in the set {did = arg max, (d, x')), and for 2 qerät the optimal decision is not in this set. Let d3 & {dld arg max, (id, x')). Consider the following equation:
BIU pa do(D d2); (qorit) E pupl dD dz); (garit) =0
(25)
Gerit Σ
= = = ((),,)) P(C = cldo( = dd): o)); x) +
+(1 - 9 ((d2) - U(ds, x'), =0
(26)
Gerit (1 P(C= = - o)))
P(C
U(22, 2')
(27)
Therefore, since the functional relationship between U and its parents is known, if we find geri we can identify:
Σ P(C da): - P(C cldo( D (d);
(28)
Let GD be without the edges leaving D, the mediated case allows for Descp O Ancy (Gp) + Consider the partition of C' proposed in Table the proof of Theorem 1 in Richens and Everitt [0024 can still be used with minor changes in the mediated case for some of the nodes, but new case needs to be introduced for Ancy (Gp) Descp.
24

---

Theorem 2. Let M be a CID satisfying Assumptions 7 1-9 Suppose for each Cs. we know a set Pac, C Pac, and set of nodes Vieun S V where C₁ Vienn = Pac. Pac. Let ( (appatee denote the set of optimal policies available where x (d(dppap) is an optimal policy under the intervention and E is the set of all mixtures of local interventions.
Then, for almost all CIDs, the graph G and the joint distribution P over all the ancestors of the utility node Ancy can be identified.
Proof. Let GD be G without the edges leaving D. Following the CID's chance node partition described in Table 1 we consider three cases:
For the third case, we provide a constructive proof for nodes in Ancu n DesCD. We establish this proof by strong induction. Consider directed path Cs + ..... + Ci where Ci E Pag, i. C, E Desep \ Ch Assume we know Puc. and P(Ck|Pack for all i= -I...,, A 1, we want to learn Pac, and P(Ck|Pac,) ). For each of the nodes Y; in Y = C (C,.......................... Ck) we define the following hard interventions
[Case 3,
oce (Y V. = y, Y, = k)] = do = U......, Y, :..... Yir Un C& = f(x))) where y is an instantiation for Y \ {Yi} and A one for Y, Here f(Ck) is the following local intervention on C₁:
= c/ C C=
(cp ch
S(C))
(29)
otherwise
We also mask all inputs to the policy: Pa'p 0. Assume Cx & ChD. by Lemma 2 we can identify the following query:
Σ cclooo d) (Y Y, = v, x))
(30)
- P(CC = cldo( D = d'): oc. (Y Y; Y, = x))
- P(C₁ = cylpac, dl .c, (Y \ |Y, y, = kd)
Ck
j=1
II P(C ej do) Y, = y. v. x))) x2)
(31)
According to Rule 3 of do-calculus Pearl 2009]. since D 1 Cl,..., O in Gy the expression in Equation 31 is equal to:
& Σ EI. P(C, cylpoc, (Y \ Y y, = - Udd' x)] (32)
CV
j=1
Σ P(Ck = c pacc (Y \ Y, y. Y, = =())((cx))
(33)
=
= =EP(C==
P(C cllc=
(34)
C6
where:
1-1 ((ck) = Σ ... P(C = ccl V. y. = [x)(((dd)))) - U((d,,
CH-1
f=1
(35)
This result is analogous to the one for Case 1. In Equation 34 following the definition of
the intervention oc, (Y) Y = y, Y, = k), we have:
P(C+ = d4 pacyi ; oC, (Y Y = y, Y. = k)) = P(C² = calpace)
(36)
and
P(C₁ cllpacc o0. (Y \ Y, y, Y: k))
1 P(Cx de s/lpacui oc. (Y) \ Y. y, Y, = x))
(37)
= PCC = C2 p))p))
25

---

Therefore there is only one parameter to be identified. If C& & Viwn, we can repeat this procedure with a different leave-one-out intervention for each potential parent Y; of C2 and different Cg. If for some configuration of K1; K2 and C₁ we have
P(C² < pplcocio Cc (Y \Y, = y, Yg = k1))
+
P(Ck d Sploaci OC (Y V y, Y; K2))
then Y; E Pac.- We can exclude from this search all nodes in Pac,, since we already know they are parents of Cz. Cs E Viwn we can skip this step since we already know Pac, Then, for each instantiations of the variables in Pac, and each ck E domn(Cs) we repeat the procedure and recover all the parameters for Cg-
And now we describe the necessary modification to the cases covered in Richens and Everitt 2024 (Case and 2 in Table I]:
1. Ancy (Gp) n [(nnd U Descp) L Anep). The identification problem for nodes in this set is described in Theorem 1 Case I of Richens and Everiti (2024). The proof is based on strong induction on k for directed paths Cs - ... + C1 where C E Pag, where for all .,,..,Kk k we have C # D. The procedure to incorporate prior knowledge is the same as the one specified in the proof for Case 3.
2. Ancy
AncD. This case corresponds to Theorem Case 2 of Richens and Everitt (2024)). The original proof considered strong induction on k for directed paths Cse ->... -.... C1 where Cj E Pap. Again, the procedure to incorporate prior knowledge the same as the one specified in the proof for Case 3.
(Gp) Ancu(GD))
A.3 Algorithmic Complexity of LearnCID
First consider that the complexity of LearnCID (Algorithm depends on the complexity of querying the policy oracle, let call this complexity K. As worst-case scenario there is no prior knowledge about the graph, so Veen is empty. Let n be the number of variables, and b == maxxec |dom X)] be the maximum number of observable values of any chance variable. The most computationally expensive steps correspond to computing forit with ALG (line 10) and the CPTs' entries for all the chance nodes. Each call ALG qorse costs ONN(K + dom((D))) filling one CPT's entry costs O(bb"). Overall the algorithm's 's time complexity is O(n NNKKKKK + (dom(D))) + n²t2n).
A.4 Comparison to Richens and Everitt [2024]
Richens and Everitt introduced an innovative perspective on the connection between robustness to distribution shifts and causal understanding of the environment [Richens and Everiti 2024 While the existence of a link between causal reasoning and generalization across domains was well established (Pearl and Bareinboim] 2011 Schölkopf et al., 2021), it was unclear whether causal understanding was necessary for domain adaptation. They were the first to show that agents must learn a causal model of their environment to achieve robustness to distribution shifts. This work builds on and extends their contributions in several ways. First, while their results are limited to unmediated tasks, this work shows that causal information is learned by robust agents even in mediated settings. Second, this paper includes a general algorithm for eliciting CIDs from optimal policy oracles, whereas Richens and Everitt [2024] presents an example limited
two-variable CID. Third, the main theorem and algorithm of this paper explicitly allow for incorporating prior knowledge about the CID's causal structure, which is dimension not addressed in their work. Fourth, the analysis is extended to multi-agent systems and techniques are discussed for applying the LearnCID algorithm in such settings. Finally, while Richens and Everiti (2024) focuses on single-shot tasks, this work also explores sequential decision-making by representing POMDPs as infinite CIDs, drawing on prior structural ideas from Everitt and Bareinboim [Everitt et al. 2019 Bareinboim et 2024 and shows how LearnCID can be applied to time-homoeee time-homoeeneous POMDPs.
26

---

B Non-identifiability of Chp and their CPTs
Now we prove that, under the same assumptions made for the previous results, Chp cannot be uniquely determined in general. This motivates the LearCIID assumption requiring knowledge of the children of D. their CPTs, and their parents.
Theorem 3. Let M = (G (V,E). P) be single decision/single utility CID. assume we know Pa(D).
D, D,U,C, C, and Pa(D). Let Σ be the set of all mixtures of local interventions, I be an optimal policy oracle for D. Then in general Chp cannot be uniquely determined.
Proof. Consider the following example:
X
X
D
U
D
U
Figure 4: On the left, CID where X & Chi. On the right, CID with the same nodes, but where X Chp.
Let p = P(X 1) = 0.9 and P(X = 0) 0.1. Let (d, a) == if d = a and otherwise. In this case the optimal decision without interventions is d is 1. There four possible local interventions σ1 = do = 1). do = do(X X = 0), did corresponding to an identity map for X, as switching the values and of X. The set of all mixtures of local interventions E includes all the interventions in the parametrized f'amily of mixtures o (po. P1, ps) == Pooo + P1d1 + PidTid PaT, with all the convex coefficient p: E [0, 1] and S p; 1. Note that each of the four local interventions is special case of this mixture, obtained by setting the corresponding coefficient to I and the others to 0.
Let X and D be binary variables and consider the CID illustrated on the left side Figure 4
D
X
P(X
D)
0
0
0.2
0
1
0.8
1
0
1
1
Table 2: Conditional probabilities X given for the CID where E Chp-
Also consider the CID illustrated on the right side of Figure along with its corresponding conditional chance nodes is the same in both CIDs; therefore, the
probability table (CPT) Table 2 The set of all possible mixtures of local interventions, denoted Σ. is also identical. For any given pair of interventions, the critical value Geit is the same both models. For example, the policy oracle II; returns d' = 0 for Jo, d' = 1 for σ1. and for (g), it returns d' = when q < 0.5, d' when q > 0.5, and may return any policy when & = qeri 0.5.
The same behavior holds for other mixtures, such as:
o'(4) == q01 + (1 g) where dem 1
(") qoya + (1 a) where again Jen 0.5.
In general, for all mixtures in Σ, both the critical threshold qen and the optimal decision returned by the oracle are the same in both models. Therefore, not possible to distinguish between the two CIDs based on oracle responses alone.
27

---

Now, let 2 = P(X P(X=11 =1 D 1) and - - 0) in the second CID. Observe that regardless the specific values of 2 and q, the model equivalent model with & Ch(D) and marginal distribution P(X = =
We also show that unlike for other chance nodes, the CPT of the children chance variables of the decision node can not be fully estimated the general case. might possible to estimate only those parts of the CPT that correspond to decisions that are optimal for some distribution shift o, but not for the others.
Corollary 1 (of Lemma 21. Let M = (V, P) be single decision/single utility CID, assume know G. Let 2 be the of all mixtures local interventions, 11; be the optimal policy oracle. Then, using the identification result of Lemma 2 in general the CPTs for C Ch cannot be uniquely determined for a set of parameters with a strictly positive Lebesgue measure.
Proof. Consider the CID described the right side of Figure 4 with the CPT described in Table Assume dom D) = {0. 1, 2} and that we don't know the CPT for X. The set of local interventions E contains 01 == do(X
1). do == do(X 0), dea corresponding to identity map for X, and the parametrized family of mixtures o (po. P1. Pid: ps) = Poou P1 + Pi + PaO, with all the convex coefficient pæ E [0, 1] and 0 p; 1. For all the feasible combinations of po, Pt, Pid, Ps the deterministic policy d = is dominated by = and 1.
D
X
P(X
D)
0
0
0.5
0
1
0.5
1
0
1
2
0
0.4
2
1
0.6
Table 3: Conditional probabilities of X given D, in this example dom(D) = =(0,1,2 1,2).
Therefore, for all E Σ. the deterministic policy corresponding to d 2 is never selected by the policy oracle IT)( Since the identification result of Lemma 2 includes only probabilities P(C do(D = d'); σ) where d' an optimal solution for some a selected by the policy oracle, it follows that, in particular, we cannot identify P(X do(D 3)) P(X D = 3), or, more generally, the portion of the CPT for chance variables that are children corresponding to decisions that are never optimal under any intervention
C Domain dependence in mediated tasks
Previous results [Richens and Everitt 2024 show that for unmediated decision tasks domain de- pendence implies Assumption 9 Here we show that this implication does not hold in the mediated case. Moreover, we prove that Assumption Saciionnn@ii
Pimppiisinniiiiisssseeeei implies domain dependence in the mediated case and therefore that Assumption 9 is equivalent to domain dependence in the unmediated case, which is subcase of the mediated one. We report the definition of Domain dependence.
Definition 2 (Domain dependence [Richens and Everitt 2024). ). A CID M is said to satisfy domain dependence if there exist P(C c)2 PCC é), both compatible with the CID M such that nt" = arg max, [U] n arg max, Ep. (U..
28

---

D
X
D)
D
X
P(XID)
P)
0
0
0
0
0
0.5
0
0
D
X
U
1
0.5
1
0
0
0
0
1
0
Figure 5: An example CID to show that in the mediated case domain dependence does not imply Assumption 9 Starting from the left, specification of the utility function associated with node U, the example's s CID, and the two CPTs for X before the distribution shift (P) and after the distribution shift (P").
Consider the example described in Figure Let xd, == = U(d d = 0 and = 1, and 0 otherwise. For the distribution P' the only optimal policy corresponds to always choosing = 1 because will always correspond to and therefore the expected utility is 2. But this policy is no longer optimal under P' because X will always be observed as 0 and while for example a policy that always chooses D
the expected utility expected utility of 1. Therefore domain dependence holds, but at the same time, Assumption 9 does not hold because d' E arg maxy (dd, x) for all x E (X). Therefore, domain dependence Assumption in the general mediated case.
corresponds to an
*
Now we prove that Assumption g implies domain dependence the mediated case, and consequently is equivalent to domain dependence in the unmediated case.
Theorem 1. Let M (G, be a CID where DescD Ancy + 9 (mediated task).
Assumption
Domain dependence.
Proof. Assume A P'(C = c) compatible with M we have n¹ E arg maxx BIU] arg max; ET, [Uldo = (d)pp papl. Let d E dDm(D) be decision s.t. I" (dd(pap)) > For Assumption 9 there exist non-empty set Xa == fald & arg max U(d'xx)) Let d' d²,x" x=" E arg max U(d', We can write x' as (zf..... .i.... where stantiations of the random variables (X1...., Now, we want to define an alternative distribution p compatible with M by updating the CPTs of the variables corresponding to the parents of U. For each X3 Pa,, \ (D) let pat. paj....,pa, |Pax,1 be instantiations of parents of Xj. Let 2² be an observable value for the variable Xj. We set P(z ..... ..a,, |Pax,1 = paj...... IPax,
in are in- |Pay {D)-
Pay \ (D) and
if x
xj and
P(rj)paa......pp |Pax,] p.l.. = =1 Eije
if at = =z We repeat this procedure for 2
Paxyy
pat ma
Pax,! We call the set of these epsilon parameters E. To
all combinations of zrj and .a...... p preserve faithfulness we require the epsilon parameters to be pair-wise distinct. Observe that if we also require all epsilon parameters for the CID. We repeat the CPT update for all variables (N)).
Σ to be
<
and 0 Eeer é then we obtain valid CPT
Observe that assuming x" arg max [U] for all P' compatible with M implies that V a'' and p compatible with M we have:
Ep [ [P] 20
(38)
Ep((U do T" (d pan) =, papl - [U do x'(d| pan). papl 2 0
(39)
Let r' be deterministic policy where d' always selected and X == Pay \D with x instantiation of X.
- Ep-(U do = n (d pap) =, papl - Ep- (U do(D d'), pap] 2 0
(40)
Σ x" (d' pan) Σ II
P(C: cj pac; ((d'
C.EC ;
(41)
Σ P[C = 0j pac; Ud xx) 2
C.EC ;
29

---

Now we compute the limit for all Σ 3 < + 0:
lim EP - [U]
(42)
MriI0
Σ (d' pa)) - -U(d",
(43)
- po + Σ (d' pap)u(d,x) xx"",x")
(44)
d'ød
<d pad)U U(d')x²))
Σ 7d' ppp) -
(45)
d'etd
where in the last passage we used the fact that d' argmaxx x²).
= Σ (d' pap) - U(d, + U(d,z')+U(d", Uddd",
Σ x' ppap)- - d''d
(46)
d'
= (d pap) (U(d", xx") -d, )) <0
(47)
The last expression is strictly negative because we assumed π (d pap) > 0 and x² E Xa, therefore d & arg max, Ud', xd)) while d' E m (d', x²). Since Ep in the parameters Σ,, we can apply the theorem of permanence of sign and therefore 7 E s.t. the
Ep [U] is a polynomial
inequality 38 is false. Therefore 7p' compatible with M s.t. x² & arg max, ET. f. It follows that domain dependence holds.
From Theorem i directly follows that for unmediated decision tasks, which are subcase of the family of mediated decision tasks, domain dependence is equivalent to Assumption 9 This provides us with very straightforward way to verify domain dependence in these tasks.
Corollary 2. Let M- = P) be an unmediated CID. Assumption g is equivalent to Domain dependence.
Proof. In the unmediated case the implication Domain dependence = Assumption 9 is proven in Richens and Everitt (2024). Since the unmediated case is subcase of the mediated case, from Domain dependence. Therefore the two
Theorem [ it directly follows that Assumption 9
statements are equivalent.
D Example of mixture of local interventions
As an example of how a mixture of local interventions can represent a distribution shift, consider the following (illustrated in Figure we have a random variable X representing the outcome of a biased coin flip that always lands on heads, where X E {H, corresponds to heads and tails, respectively. Let o; == do(X i) with i E (H, then define coefficients Pt, we can map the distribution of the fair coin to any distribution on the observable values set dom(X)
=, P.ot. By changing the
(.. For example, by setting PH = and pr }.2 we can map the original distribution to a new one where heads is observed 3 of the time and tails is observed & of the time. Note that in this example, each local intervention was a hard intervention because, regardless of the value of the coin, each intervention mapped it to a specific value. In general, this is not required, as local intervention can be any deterministic map from the set of observable values to itself.
While hard interventions break the dependence of variable on its parents, local interventions may preserve this dependence. For example, consider an unbiased coin flip Y and a die throw X. If the coin lands on heads, we throw a six-sided die, if it lands on tails, we throw twelve-sided die. In this scenario, Y and X are clearly dependent and in the corresponding causal graph Y E Pax- Now consider the local intervention om == do(X X X mod 12 + 1). This intervention increases the die result by one or sets it to 1 if the result was 12. This local intervention does not make X independent of Y because, for instance:
P(X 4/- T:om) P(X 3/Y T) 12
(48)
30

---

Original Distribution
Shifted Distribution
H
Observed
Observed value
Figure 6: Histograms illustrating the distribution shift induced by mixture of local interventions.
The left histogram shows biased coin distribution P(X), where the coin always lands on heads. The right histogram represents shifted distribution P(X o²) obtained by applying the mixture o² P + pror for a; == do(X = i) with pи = 2 and pr 3
However,
P(X = 3) P(X mm
(49)
12
E
LearnCID examples
In this section, we provide two examples to illustrate the application of LearnCID (Algorithm) D in both single and multi-agent environments.
E.1 Example 1 Single-agent environment
Now will go complete, step-by-step example of a LearnCID application to learn the simple CID illustrated in Figure 7 For this example, we assume Vkwn is empty.
First, let us go over the LearnCID Assumptions
1. By observing the (unknown) CPT of A we can see that for this example faithfulness holds, there are no latent confounders so causal sufficiency also holds.
2. The set of decision nodes is {D), the utility node associated with the utility function U (d, a) is U, and the set of chance nodes is C {A, B3.
3. The CID indeed contains
decision node D and utility node U.
4. The Markov blanket of D is (B), and know that D is parent of B, i.e., the graph contains the edge + B.
5. Since the decision is an input of the utility function Udd, a) we know a parent of U.
6. Both chance nodes and B are ancestors of U.
7. The utility function (d, a) = if and otherwise, is fully specified.
8. We have an optimal policy oracle 11,, where M is the set of all mixtures of local interventions.
9. We can verify that for A 1, arg max, (d, A 1) but arg max, (d, A 0) 0.
For ease of comprehension we summarize the application of Algorithm I to the example CID in Figure 7 in the following steps:
SI Visit A, the only unvisited chance node with known path U.
S2 Define local interventions B as in Equation
S3 Estimate aeiüt using AL.Gqeia (Algorithm 1 in Richens and Everitt 2004
S4 Compute P(A ao/do( B = 6;)).
S5 Repeat steps 2 to for all configurations of a, and b..
S6 Deduce the set of parents of A and its CPT.
31

---

Visualizing Pcrit
1.0
B
A
Expected utility under a(q)
0.8
0.6
D
U
E[U(D=1AA)a(a])
E(U(D=0,A)(({]]
0.4
B
A
P(A
0
0.7
0.2
0
0.3
1
0.1
0.0
1
0.9
0.0
0.2
0.4
0.6
0.8 4crie
1.0
CID. On the bottom-left, the CPT for variable A. The edge marked red is unknown. On the right, following Example 1, let us examine two interventions: oo == do B = 0) with optimal decision d1, and hard intervention o' with optimal decision d2 where d₁ # dz. We define a mixture of local interventions as ( = 4% + (1 q)o'. The plot displays the expected utility of both decisions varies. The value for where both decisions become simultaneously optimal is called ferit-
Figure 7: An example of a single-decision/sinlleeue
Following the aforementioned steps:
Step I: In this example Pap is empty and Pag = {A]. LearnCID visits and finds the set of parents and CPT of all chance nodes that have known path U, since at the beginning the only chance node with a known path to U is A, we start by visiting node A (line 3).
Observe that since A is the only chance node that is not children of D the process will stop after A. On line 4, the variable "Path" is initialized with a set of all chance nodes one path from A, the chance node are currently visiting, and U. Since A Pay, "Path" empty.
On line 5, we initialize the set CA by including all the chance nodes that are not A and not in the path from A to U we selected ("Path"). Since "Path" is empty, the only chance node different from is B and therefore CA (B}.
Step 2: Now, from C = {B} we pick a node that could potentially be a parent A (lines and 8), in this case the only option is B. Since CA {B} is empty. then the only possible instantiation for CA { (B)
Then on line 9, following Equation 4 we initialize a set of interventions o (0) as:
(0) 01)
where de do B 0), and 01 do B =1)
(50)
Also, we pick an observable value for (line 7), for example O, and aim to recover P(A = 0; a) for all dg. Note that since A is binary, the function f in Equation is the identity function and can therefore be ignored.
Step 3: For each E an in ALGgerit we use the oracle to find the optimal decision under (line 11). For example, for on do = 0), IT: (mo) returns the optimal decision do == This evident from the (unknown) CPT because for B 0 the probability that A = 0 is higher than the one for 1, since the utility is maximized when A and D take the same value, the oracle returns the optimal decision D which will be equal to A more often than D 1.
ALGgera then estimate Gerit- It works by finding an intervention o' such that the optimal decision is no longer D = 0 as it was under do- Since in this example D is binary, we already know the new optimal decision must be
1, general, for this we can use the optimal policy oracle. We can find o' by hard intervening on the parents of the utility node U, which is A in this case, such that do 1, therefore
no longer optimal. Specifically, we can define a' as a hard intervention that sets the new optimal decision is dj = 1. Note that this is always possible thanks to Assumption 9 Then, we define the mixture local interventions (g) = goo (1 q)o'. The right side of Figure shows how the expected utility of both decisions varies with We can sample uniformly in the
32

---

interval [0, 1] N times and each time query the optimal policy oracle. Each time the oracle returns an optimal decision for the intervention a' we increment counter 0. Then & is an unbiased estimate N
5 for Gerit. For this example, Gerit 7-
Step 4: Now we can compute P(A 0; oo). Consider the mixture of interventions ((1) from Step 3. We can write:
E do(D = 0))
Σ P(A a, B = 6 do -): ()))
Σ EaP(A=a,B= do 0)) + (1 P(A= a, = do = 0):
a.h
When g = qerat both decisions D and D are optimal:
ElU do(D (D 1): = EUU do(D 0); (
EIU do(D = ( - ElU do(D = 0): (gerit)) =0
=0
- Serit qerit
Σ
P(A a, B b do = 1): - P(A = a, = b bldo(D = (): a)
a,b
+1 (crit) [U 1,a a - 0, 1)]
(51)
Recall 0. = do B 0). Since we know that A is not a child of D and we are hard intervening on all chance nodes non descendants of A, by rule 3 of do-caleulus [Pearl 2009) P(A a do (
d, B b)) P(A = a Bo b)), and we can write:
Σ P(A=a = do(D = 1, B = (1) - a 0, 0, 0, +
Gerit
+ (1 9oritt
[() -(00
- Aerit
E -0 do B
do |do(B=0))((0.))
+
(1 (erit) (U(,,) - UU, 1)]
P((A=al do( ()) u(O.a)] -(0,))
+ (1 aera) (U,,1) UO. =0
Gcrit
Now we can define the following expression, which corresponds to Equation 5 B(a) (1. - a)
(52)
And rewrite Equation 51 as:
P(A- a do[B 0))](a)) + (1- (i-_ U0, 1)] 0
Gerit
Perst [PAA = do(B = ()) + P(A = 1 do(B = ())
(osst) (1,,1) (O, 1)]
+ Gesit (P(A 0 do(B = ()) + P(A-P do ()])))))3)))))
+ (1 Perst (1), U() 0
-
+ P(A=0) 0| do( = 0)) gerit 3(1) + Geria B(1+
-U(1,)) -QorieU (1, + Gorit U (0, 1)
(1 ( U (0, 3(1)
- P(A= 0 do B 0)) =
3(0)-3(1)
33

---

Which corresponds to Equation (O, 1) then:
We can compute g(0) -1, = 1, U(11)) = 1, and
P(A = do(B = 0)) (1 {t) = 0.7
(53)
-2
Which is consistent with the CID and CPT shown in Figure 7
Step 5: We can repeat the same procedure for 01 = do = 1). this case, since is binary we compute P(A on)
- P(A 0; on), in general we can repeat the procedure for all
possible instantiations a of A to obtain P(A = 0; 0) for all a E OB.
Step 6: Since P(A 0ldo B 0)) * P(A Oldo 1)) (line 13), we can conclude that B is a parent of A. In the general case, this approach ensures that B is not just an ancestor but indeed parent of A, because the intervention blocks all other paths from B to A. Thus, if B were not a parent, P(A = A/p do = 0)) would be equal to P(A = alpa Once we have identified can reconstruct the CPT for A using the interventional distributions, e.g. = 0) = 0.3 (lines 14 and 15).
the set of parents,
P(A = dB B = 0)) = P(A
As expected, this process allows us to learn both the correct graph structure and the CPT for A (and, more generally, for all chance nodes that are not children of D).
E.2 Example 2 Multi-agent environment
Examine the multi-decision CID in Figure g It represents a cooperative game between two agents, agents and B, respectively controlling decision variables D and DR. Both agents aim to maximize shared utility function U, and operate in different contexts defined by the parent sets of their decision nodes (PADA = 0 * { = PaD,). ). As usual, we assume knowledge the children for the decision node DA their CPTs, their parents, and the utility function associated with the node U. In this section, we demonstrate how to apply LearnCID to this example by providing a high-level description omitting the details we provided in the single-agent case.
DA
X
DB
DA
X
DB
Y
Z
Y
Z
U
Figure 8: A multi-decision CID that represents environment where agents A and B cooperate to maximize the utility U. The edges marked in red are unknown. Example demonstrates how to adapt this CID to apply Algorithm T and recover the missing edges and CPTs for chance nodes.
Let II's be the optimal policy oracle for DA and x(DB Z) be any given policy that governs Dg-
The nodes for which we need to learn the parents are Y and Z.
Node Y's potential parents are and Z, whereas node Z's only potential parent is Y.
We know X cannot be parent of Z because by Assumption know D B is parent of X, and assumed to know the parents of every decision node. Therefore we know X is a descendant of Z meaning it cannot be its parent. Using Algorithm we determine parental relationships
follows:
1. We check if Z is parent of Y. We consider the instantiation Y = 0. With Algorithm I we can compute P(Y = 0 pay; and P(Y and of == do X 0, Z = 1) respectively. We observe that these two probabilities are equal. We repeat this process with σ = do( ( 1, 0) and 0j == do(X = 1, = 1), and again, P(Y = 0 p1y; P(Y 0 pay: o.). Performing the same procedure
pay; on) using σo := do 0, 0).
34

---

for Y = 1, we find that all pairs of interventions yield the same probabilities. Therefore, Z is not parent of Y.
2. Next, check whether X is parent of Y. Comparing P(Y y pay; with P(Y = y p1y; and P(Y = pay;%) with P(Y= = o for all E (0, 1}. find that at least of these pairs of probabilities differs. This confirms that X is parent of Y.
3. Finally, we consider Z and aim to determine whether is parent of Z. According to the algorithum, would again verify that the conditional probabilities remain unchanged when setting different values for Z. However, we can arrive at the same conclusion also by noticing that since Y found be child of X, Y not be child of Z because this would introduce a cycle in the graph. Z has no other potential parents and therefore we have learned the full CID.
The fact that agents A and B cooperate on the same utility function does not affect the outcome. We would have obtained the same result even agent was competing with agent or if agent B did not try to optimize U but was still able to causally influence it.
F Broader Impacts
This paper presents foundational research on the relationship between an agent's ability to adapt to distribution shifts and its causal understanding of its environment. Our findings contribute to a deeper understanding of what decision making systems, including AI systems, must learn to be robust to distribution shifts. By characterizing the internal representations that such systems develop, our work promotes greater transparency how decisions are made. This understanding lays the groundwork for designing more explainable, interpretable, and robust decision decision-making systems. In the long term, these advancements stand to benefit not only the users these systems but also the broader population affected by their outcomes. If used maliciously, this technology could potentially deecision-makine
systems that are intended to remain private. However, given the theoretical nature of our results and, in particular, the absence of a scalable algorithm to obtain full causal models, such misuse remains very unlikely in practice.
be employed to probe the causal representations
35


Reviews:
[{"review_id": "WfL86VTppG", "reviewer": "Reviewer_pmnc", "summary": "This paper considers the decision-making with Mediation where the actions have an effect on the environment and develops a foundational framework to show that agents robust to distribution shifts must have learned a causal model.  The authors propose the LearnCID algorithm which learns Causal Influence Diagrams (CID). Then, the authors present the application of LearnCID to multi-agent environments and POMDPs.", "questions": "No further questions.", "limitations": "The authors discussed the limitations in the last section.", "rating": 4, "confidence": 1, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 2, "strengths_and_weaknesses": "The paper has the following strengths.\n+ The paper extends the existing results by removing the assumption of no mediation, which builds a foundation for more general use cases.\n+ The paper theoretically show that the causal model is learnable.\n+ The paper considers the concrete applications of POMDP and multi-agent systems.\n\nThe paper has the following limitations.\n- The framework only applies to discrete variables, which may limit its application for general decision-making problems.\n- The framework only focuses on learning the causal model but how to combine it with traditional data-driven RL algorithms is not discussed. \n- Would it be possible to quantify the distribution shift and show the effect of the distribution shift on the performance?", "was_revised": true, "final_justification": "I remain the rating as the author rebuttal has addressed my concerns.", "rebuttal": "Thank you for your review. \n\n**“Would it be possible to quantify the distribution shift and show the effect of the distribution shift on the performance?”**\n\nThe answer is yes to both. As discussed in Section 2 lines 104-125 and Appendix D, it is possible to formalize a distribution shift as an intervention on the causal model. In this case, once you obtain a CID, you can use it to compute an updated probability distribution for any given distribution shift. At that point, it is possible to quantify the shift by, for example, comparing the original probability distributions with those obtained after the shift. This can be done by computing the Wasserstein distance, the KL-divergence, or using any other technique to quantify the difference between probability distributions.\n\nAs for measuring the effect of the distribution shift on performance, given a policy and a utility function, you can compute the expected value of the utility for that policy under a given distribution shift using the CID. Again, if the shift is in the form of an intervention, you can first update the probability distributions, and possibly the graph structure, and then use the updated distribution and the policy to directly compute the expected utility. This can be compared with the expected utility under no distribution shift for the same policy.\n\n\n**“The framework only focuses on learning the causal model but how to combine it with traditional data-driven RL algorithms is not discussed.”**\n\nThank you for raising this point. In this work, we consider an agent that is already robust to distribution shifts, without making any assumptions about how the agent should be trained. This line of work is very novel, as, to the best of our knowledge, there are no other results outside of Richens and Everitt (2024) and this paper. For this reason, we decided to give priority to the development of the core concepts. That said, there are interesting opportunities to explore connections between this framework and RL algorithms. For example, integrating an algorithm like LearnCID into a data-driven RL method could enable a decision system to both learn the task and reveal the causal model simultaneously, or understanding how the use of an RL algorithm might influence the process of learning causal knowledge would also be insightful. These are significant topics and will become important directions for future work once the theoretical foundations of the core concepts have been fully developed."}, {"review_id": "6ZJZP8zKGU", "reviewer": "Reviewer_iDtz", "summary": "Past work shows that agents robust to distributional shifts learn a causal model of their environment. However this result assumes that the agent does not influence its environment with its actions (“no mediation”). This paper extends the previous result to the setting with mediation. The paper presents an algorithm for inferring a causal influence diagram from an optimal policy oracle, and discusses its application in multi-agent and sequential environments.", "questions": "Richens also proves approximate theorems based on regret bounded policy oracles. Do these results carry over to the mediation setting?", "limitations": "No substantial discussion of limitations. The authors should be more upfront about some assumptions (e.g., in the multi agent discussion, the assumption of shared utility)", "rating": 3, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 2, "strengths_and_weaknesses": "I appreciate what the paper is trying to do and think that extending the Richen's results is an important direction. However, the formal work in this paper is not quite tight enough for me, and the contributions are very modest. \n\nHere are my comments:\n\nMajor.\n\n- There is no related work section and several references are missing --- in particular I think it would be good to discuss the literature on CIDs ingrate detail, including their importance to problems in safe AI (e.g., discussing some of the papers from the causal incentives group https://causalincentives.com)\n- There is also this new paper, from Richens, which extends the same Robust Agents results https://arxiv.org/abs/2506.01622 (I haven't read this paper and of course can't blame the authors for not having seen such a new paper, but it's still highly relevant)\n- Aspects of the theory are not properly introduced, for example, policies and CPTs are not defined. \n- The assumptions and main results are not explained clearly enough --- there should be more intuitive discussion of each to help the reader and make the paper more accessible.\n- In a theoretical paper, you should include sketch proofs of theorems in the main paper rather than in the appendix.\n- Some of the work is basically informal (e.g., the multi-agent section) \n\nMinor. \n\nThe first paragraph of the introduction contains too much information. It should try to more concisely state the problem — lots of the detail here should be in related work imo. \n\nThe intro to CIDs isn’t that clear. The two weather examples dont immediately connect to the formal work. You should also start by saying, e.g., “CIDs are a type of probabilistic graphical model…” as of now you start talking about “the graph” without explanation. Further, definition 1 introduces a CID as a CBN but CBNs have not been defined or mentioned.\n\n“showing that it is possible to recover the causal structure and the Conditional\nProbability Tables (CPTs) of the variables describing the environment by observing the agent’s\noptimal policies under distribution shifts.”\n\nCPTs and policies and optimality not defined. \n\nTerminology is often non-standard. Why depart from the literature? “CPTs” —> CPDs; “informational arcs” —> information links \n\n“In this paper, the\nterms adaptable agent and robust agent refer to an agent that maintains optimal performance across\nall the possible distribution shifts in its environment.”  so these mean the same thing? Just stick to robust agent?\n\nIn CID graphs, edges into decisions (i.e., information links) should be dashed. \n\n\n“Each agent may correspond to a different set of decision nodes and have access to a distinct subset of\nobservable variables.” A proper formal analysis of multi-agent settings requires e.g. Causal Games cf Hammond https://arxiv.org/abs/2301.02324\n\nNo need for policy oracles to be optimal by definition. (You already say in many places “an optimal policy oracle”. \n\nMain results section should be split up into different sections. \n\n“multi-agent systems represented with multi-\ndecision CIDs” I don’t think this is a valid model, at least it misses all the game theoretic dynamics. Even if the agents optimise the same utility (a serious limitation) you dont deal with e.g., different observations and memory.", "was_revised": true, "final_justification": "Overall I feel that the paper's formal work is not really tight enough, the novel contribution is modest, more could be done to make the paper intuitive and accessible, and the later parts are weak (e.g., multi agent extension). I think a purely theoretical paper should be more solid on all of these lines to be at the acceptance level for neurips. \n\nI think a future version of the paper could make a solid contribution if it becomes more focused, tightens up the formal work, and perhaps adds experimental validation.", "rebuttal": "Thank you for your review. Starting from the question:\n\n**Q1: “Richens also proves approximate theorems based on regret bounded policy oracles. Do these results carry over to the mediation setting?”**\n\nWe believe these results carry over to the mediation setting, though the formal work to confirm this has not yet been undertaken. The three contributions outlined in the first part of the paper, namely showing that domain dependence is no longer sufficient, establishing unidentifiability results for the children of the decision node, and developing a new procedure to recover the causal model in the optimal case, are prerequisites for extending the approximate results to the mediated case. Exploring the approximate case in the mediation setting is an important direction for future work.\n\n\\\n**Limitations: “The authors should be more upfront about some assumptions (e.g., in the multi agent discussion, the assumption of shared utility)”**\n\nThank you very much for pointing this out. In fact, it is not necessary for different agents to optimize the same utility function. Only one of the agents needs to optimize the utility function, while the others only need to be able to influence it. Specifically, it suffices that all decision nodes under consideration are ancestors of the utility node, and that we have access to a policy oracle for at least one of these decision nodes which is optimal with respect to the utility function associated with the utility node. The proposed methods for applying LearnCID to multi-decision CIDs can be applied to this situation without modification. We will update Section 3.2 to mention this.\n\nRegarding the observations.\n\n**Major 1:** \\\n**“There is no related work section and several references are missing --- in particular I think it would be good to discuss the literature on CIDs ingrate detail, including their importance to problems in safe AI (e.g., discussing some of the papers from the causal incentives group)”**\n\nWhile our current focus has been on citing directly relevant peer-reviewed work, we will expand the related work section to include a more detailed discussion of the broader CID literature, including its relevance to safe AI. We believe we have cited the most relevant existing works, but we will carefully review the literature again (including those from the causal incentives group) and ensure that everything relevant is cited. Does the reviewer have any suggestions for papers we should include that we may have overlooked?\n\n**Major 2:** \\\n**“There is also this new paper, from Richens, which extends the same Robust Agents results (I haven't read this paper and of course can't blame the authors for not having seen such a new paper, but it's still highly relevant)”**\n\nThe paper was not cited because it was first published on June 2nd, and the full paper submission deadline for NeurIPS 2025 was on May 22nd. We have read this paper. It shows that an agent able to adapt to different goals must have learned a predictive model of the environment. The results of the paper are indeed interesting and the approach is similar. We will cite it in our revision. However,  the paper is not about causal models, causal knowledge, or distribution shifts in the environment, so it does not overlap with our work in any way. \n\n**Major 3 and 4:** \\\n**3. “Aspects of the theory are not properly introduced, for example, policies and CPTs are not defined.”** \\\n**4. “The assumptions and main results are not explained clearly enough --- there should be more intuitive discussion of each to help the reader and make the paper more accessible.”**\n\nRegarding policies and CPTs, we understand that accessibility is valuable, and we will take the necessary steps to address this.\nWe also agree that understanding the assumption is very important. We will include a more intuitive description of the assumptions.\n\n**Major 5:** \\\n“In a theoretical paper, you should include sketch proofs of theorems in the main paper rather than in the appendix.”\n\nThank you for the suggestion. To clarify, the appendix contains only complete proofs. Similar to what they do in Richens and Everitt (2024), we chose to keep the main text focused on the high-level ideas, while providing full technical details in the appendix to maintain readability.\n\n**Major 6:** \\\n**“Some of the work is basically informal (e.g., the multi-agent section)”**\n\nCould you please pinpoint which parts you found to be informal? We would greatly appreciate your feedback so we can address them thoroughly.\n\n\\\n**Minor issues**\n\nAbout the minor issues, thank you very much for your corrections and suggestions in the introduction, we will address all of these and also add the CBN definition and turn the edges going into decision nodes from full to dashed. We also agree that it would be more appropriate and equally convenient to separate the notion of policy oracle and that of its optimality, we will update the definition.\n\nAbout CPTs vs CPDs. CPTs are standard terminology in discrete bayesian networks. The two terms are not equivalent. The reason why we used CPTs in this context is that while CPDs can be associated with both discrete and continuous variables, CPTs are only defined for discrete variables, which are those we consider throughout the paper.\n\n\\\n**“ ‘multi-agent systems represented with multi- decision CIDs’ I don’t think this is a valid model, at least it misses all the game theoretic dynamics. Even if the agents optimise the same utility (a serious limitation) you dont deal with e.g., different observations and memory.”**\n\nYou are correct that the mechanised MAID is a richer model when it comes to combining causal and game-theoretical reasoning by adding the decision rules and the parameters of the non-decision variables to the graph. Multi-decision CIDs are a special case of MAID in “Multi-agent influence diagrams for representing and solving games” Koller and Milch (2001), with edges that do not enter decision nodes representing causal relationships, and without explicitly specifying the set of agents and the correspondence between agents and decision nodes. We chose this minimal model because we only need a representation of multi-agent systems that captures the causal dependencies between the environment and decision variables both concisely and flexibly. We believe the results can be specialized for mechanized MAIDs with modest additional work.\n\n\n_Regarding memory and different observations_ \\\nStill, multi-decision CIDs can represent agents with different observations by allowing each decision node to have distinct parent sets and by masking observations through interventions, which is the approach we adopt in this paper.  In a similar way, memory of past observations and decisions can be implemented by including the corresponding nodes as parents of the subsequent decision that relies on this memory.\n\n_On having multiple utility nodes_ \\\nWe agree that it would be interesting to develop a version of LearnCID for multiple utility functions, perhaps by using an oracle that returns a policy optimizing a weighted combination of utilities. However, observe that even in our settings, there can be more than one utility function, and different agents are allowed to optimize different utility functions. When we want to apply LearnCID, we just need to prune all the utility nodes except one and only keep the chance nodes that are its ancestors. We can repeat the same procedure for different utility functions in the same environment if we have at least one corresponding optimal policy oracle available. We will make this clearer by updating Section 3.2."}, {"review_id": "TP5B1bdHDl", "reviewer": "Reviewer_NWB8", "summary": "The paper proposes an algorithm and proves it's consistency for discovering causal world models from robust agents when they affect the environment (mediated setting). The algorithm is further applied in the multi-agent setting and POMDP setting. It expands the work from Richens and Everitt [1].\n\nThe implications of this work are that robust optimal policies need to learn the CID in order to be robust to distribution shifts in the mediated setting.\n\n[1]  Robust agents learn causal world models, ICLR 2024", "questions": "## Questions\n\nl.99 - this sounds more like a theorem or proposition than a definition since it makes a claim about existance of optimal policies in that a re different for CIDs statisfying the conditio?  This cannot be a definition. A follow-up to this question is how does this relate to the notion of identifiability?\n\nl.112 - the way the intervention is written is really easy to confuse with the hard intervention case, but in effect the \"local\" intervention is introducing a transformation of the random variable X? So f is a push-forward map, correct? Shouldn't this just be seen as introducing a \"mechanism\" that introduces another ancestor Y to children of X, i.e. we have a a dependance chain introduced,  X ->  Y -> children(X). Might be simpler to explain it this way.", "limitations": "yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 4, "strengths_and_weaknesses": "## Strengths\n\nSeveral theoretical insights considering causal discovery of the underlying graphical model from agents robust to distribution shifts in the mediated setting, multi-agent setting and POMDP setting.\n\nThe setting of inferring the causal model without the assumption that the agent is not influencing the environment is novel as far as I know. \n\n\n## Weaknesses\n\n\nThe algorithm has similar limitations as classical causal discovery algorithms, it works only in the structured setting where we have nice separation of the variables. The more interesting setting that requires discovery of the latent variables themselves is not tackled.\n\nThere is no experimental validation of the algorithm in any interesting setting for the ML community. \n\nThere is no complexity analysis of the algorithm regarding runtime.", "was_revised": true, "final_justification": "The authors have answered all of my questions and I think that the paper makes a solid technical contribution in the realm of causal world models and decision making under distribution shifts. I am not increasing my score because of the lack of experimental evaluation and significance, I still think it's good work.", "rebuttal": "Thank you for your review and your observations. Regarding your questions:\n\n**Q1: “l.99 - this sounds more like a theorem or proposition than a definition since it makes a claim about existance of optimal policies in that a re different for CIDs statisfying the conditio? This cannot be a definition. A follow-up to this question is how does this relate to the notion of identifiability?”**\n\nThank you for bringing this to our attention. We would like to update it as follows: \\\n\"$\\underline{\\text{A CID }M \\text{ is said to satisfy domain dependence if}}$ there exist $P(C = c)$, $P^\\prime(C = c)$, both compatible with the CID $M$ such that $\\pi^* = argmax_\\pi\\mathbb E_{P}^\\pi[U]\\implies \\pi^*\\not = argmax_\\pi\\mathbb E_{P^\\prime}^\\pi[U]$.” \\\nWe look forward to the reviewer's comments on this.\\\nRegarding identifiability, under the assumptions made in the paper, if domain dependence does not hold, then it is not possible to learn the causal model from the policy oracle. Intuitively, the reason is that if domain dependence does not hold, there exists a policy whose optimality is unaffected by any distribution shift. In this case, the optimality of a decision system is not affected by changes in the environment, and causal knowledge becomes unnecessary.\n\n\\\n**Q2: “l.112 - the way the intervention is written is really easy to confuse with the hard intervention case, but in effect the \"local\" intervention is introducing a transformation of the random variable X? So f is a push-forward map, correct? Shouldn't this just be seen as introducing a \"mechanism\" that introduces another ancestor Y to children of X, i.e. we have a a dependance chain introduced, X -> Y -> children(X). Might be simpler to explain it this way.”**\n\nWe agree that the notation doesn’t make it immediately clear whether the intervention we are applying is hard or not. For these results, we decided to stick with Richens and Everitt’s notation. Yes, f is indeed a push-forward map, and the idea you proposed to add an intermediate node between the intervened variable X and its children is correct and can be an effective way to make the concept more intuitive and easier to understand. Thank you for the suggestion.\n\nAbout the weaknesses that were pointed out:\n\n**W1:** \\\n**“The algorithm has similar limitations as classical causal discovery algorithms, it works only in the structured setting where we have nice separation of the variables. The more interesting setting that requires discovery of the latent variables themselves is not tackled.”**\n\nThank you for this insightful comment. We agree that the setting where latent variables must be discovered directly is interesting. As we understand it, the issue being raised is that the paper assumes knowledge of the latent variables and also assumes causal sufficiency, meaning there is no latent confounding. As you also noted, assuming causal sufficiency is a common approach in the causal discovery literature, especially when proposing new methods to address the problem. Examples include LiM (Zeng et al. 2022), RESIT (Peters et al., 2014), LiNGAM (Shimizu et al., 2006), and PC (Spirtes and Glymour, 1991).\nRegarding the assumption of knowing the set of latent variables, the problem of identifying latent variables and the problem of discovering causal relationships between them are, in principle, separate. There is active research focused on learning latent representations, including latent variables, for causal models. The review suggests that you are probably familiar with techniques like those of causal representation learning that among other things aim at learning the latent variables for causal models from high-dimensional data. All of this is important and interesting work, which we regard as complementary rather than overlapping. A discussion of how these techniques could be integrated with the paper’s framework would be a valuable addition, and we will highlight these possibilities in the revision.\nIt is also important to note that both of these assumptions, along with the others in this paper, were either explicitly stated or implied in Richens and Everitt (2024). In fact, this work relaxes their setting.\n\n**W2:** \\\n**“There is no experimental validation of the algorithm in any interesting setting for the ML community.”**\n\nThank you for emphasizing the importance of empirical validation and we fully appreciate it.\nAs we state in the submission, our work is purely theoretical in nature, aligning with the types of contributions NeurIPS explicitly supports. There is a strong tradition at NeurIPS of high-impact theoretical work without empirical experiments, such as: “A Universal Law of Robustness via Isoperimetry” by Bubeck and Sellke (NeurIPS 2022), “Causal Discovery from Soft Interventions with Unknown Targets: Characterization and Learning” by Jaber et al. (NeurIPS 2020), and “Transportability from Multiple Environments with Limited Experiments: Completeness Results” by Bareinboim and Pearl (NeurIPS 2014).  That said, we agree that empirical validation is a natural next step, but only after the foundational theory is well established. Our work aims to lay that foundation, which future empirical efforts can build upon.\n\n**W3:** \\\n**“There is no complexity analysis of the algorithm regarding runtime.”**\n\nAs mentioned in Section 3.1 line 237, we discuss the runtime complexity of LearnCID in Appendix A.3, starting on line 667."}, {"review_id": "E1CoJ2WOPI", "reviewer": "Reviewer_dEGV", "summary": "The paper provides theoretical results to show that decision-making robust to distribution shifts in the environment implies the learning of a causal model even under mediation, that is, when the decisions taken influence the environment, subject to a number of assumptions. The manuscript shows this by deriving an algorithm that learns the causal influence diagram (CID) and the corresponding conditional probability tables (CPT) from a base CID with some sets of already known information. The paper then extends the algorithm to a sequential decision-making case, by considering an unrolled partially-observable Markov decision process (POMDP) and extending their algorithm with a method to learn a CID from the POMDP.", "questions": "I have given a range of questions with context above, but for ease of access, I copy them below:\n- Do you think your work can really be characterised as being about agents?\n- Does your extension to POMDPs really reflect sequential decision-making? In particular\n  - How does the discount factor come into play in Algorithm 2? \n  - Is there any way this algorithm can model temporal relationships besides encoding everything in the state and relying on the Markov assumption?\n  - Is there a notion of a goal here, e.g. can $U_t$ correspond to a long-term sparse reward signal?\n- Can you really not use causal discovery algorithms, as you are taking some sort of samples in Algorithm 1 from the oracle \n- In Algorithm 1, how are the edges of the CID $V$ updated to give $E'$? \n- How do we or can we know that the base CID assumed in A2 is correct to begin with? \n  - If it is correct, then doesn't that already imply that agents have a causal model of the world to begin with?", "limitations": "I think the paper could benefit a *lot* from a more thorough discussion of its limitations. Throughout reading the paper, I was under the impression that the claims made in the introduction are not fully justified by the results, and in some cases, could potentially be construed as misleading. I do think the results are really valuable, but their interpretation would be much easier if the implied claim that all agents must be learning a CID were put into context by, e.g. providing an accessible (intuitive or non-formal) summary of the assumptions made in the paper. \n\nI also thought that the discussion of related work as it relates to the paper's motivations was not explored to a satisfying extent. The work cites many related papers but doesn't provide a strong case for why CID-style causality would be the best form of causality to explore in this case. To strengthen the paper, it would be nice to see some examples mentioned in the introduction on how causal modelling addressed the challenges of distribution shift, reliable generalisation, and transparent decision-making in practice.", "rating": 3, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 2, "significance": 4, "originality": 3, "strengths_and_weaknesses": "*Note: I want to flag that I am not an expert in \"Pearlian\" causal modelling and did not check the equations for correctness. My comments are, therefore, more high-level, and I welcome clarifications from the authors if I have misunderstood something regarding the technical aspects of the work.*\n\nOverall, I thought the work was really interesting and, if correct, the contributions seem potentially valuable to understanding the nature of agency. However, I felt there was a disconnect between the motivation of the paper, its contributions, and related work, and I am hoping the authors could clarify some points in the discussion phase regarding their work's significance.\n\nOne of the main points I kept revisiting while reading the manuscript is whether the paper really deals with agents. Putting aside the debate on what it means to be an agent and have agency, the main argument of the paper seems to rely on the fact that we are in a one-shot decision-making setting with optimal decisions derived through a relative ordering defined over fixed utilities, as represented by the differences in Eq. 5 and 6. The authors also give an example of a sprinkler as an AI system, which I am not too sure about. In both cases, I am not convinced we are talking about agents, let alone AI. There is no mention of actions, goals, adaptability, or time. In this regard, I do appreciate the authors' extension of their algorithm to POMPDs, which seems to relax some of these assumptions. My concern is whether this extension violates any existing assumptions or introduces new ones (e.g. How does the discount factor come into play in Algorithm 2? Is there any way this algorithm can model temporal relationships besides encoding everything in the state and relying on the Markov assumption?) \n\nOn a technical note, I am wondering whether causal discovery algorithms could really not be used in the context of the manuscript. You are taking some sort of samples in Algorithm 1 from the oracle (though it is not very well explained in the text), so why not just use an existing algorithm? I am also wondering where in Algorithm 1 are the edges of the CID $V$ update to give $E'$?  Algorithm 1 also seems to assume that we have prior knowledge about the causal structure, and Assumption 2 adds quite a bit of prior information on its own. My concern here is: how do we know that the base CID is correct to begin with? And if it is correct, then doesn't that already imply that agents have a causal model of the world to begin with? This seems like assuming the conclusion in some ways.\n\nAnother note is that in current foundational agent discussions (on this see e.g. [Abel et al.](https://arxiv.org/pdf/2502.04403), [Kasirzadeh and Gabriel](https://arxiv.org/pdf/2504.21848), [Bowling and Elelimy](https://arxiv.org/pdf/2504.08161)), generalisation emerges as an important property of continual learning. This work necessarily glosses over this aspect of agency; otherwise, I reckon it would need a few hundred pages of extra content, but I found it curious that there was no mention of it. If the main claim of the paper (\"agents robust to distribution shifts learn causal models...\") is read at face value, then one might be tempted to assume that agents learning continually (or just learning in general) also have a sort of adapting causal graph in their internals, which seems rather strong in light of the fact that the paper makes a range of assumptions and simplifications. For clarity and expectations, it would be useful to include some caveats of the manuscript's results, if not in the title, at least in the introduction. Otherwise, others might cite the work as proof of the titular claim in a real-world setting or in places where the work's claims do not hold.\n\nI also fail to fully understand the implications of the paper's results for future and related work. The motivations seem strongly based on the case that AI safety requires understanding whether and how agents form world models, but the discussion provides no insights into how the work ties in with, e.g. control. Specifically, what if we try to steer agents by trying to learn their causal model and intervene on their decision nodes, only to find out that most assumptions of the manuscript do not hold in practice, and there is actually no one causal model to steer or decision node to steer? I think acknowledging the limitations of the paper and how its results are reflected in the (implied) claim that all agents must learn a causal model would have gone a long way in presenting this work in a more grounded light.", "was_revised": true, "final_justification": "The paper presents a novel result regarding the learnt causal representations of robust agents by relaxing the no-mediation assumption of Richens and Everitt (2024). While the result seems significant, there is a great need for improved clarity in the paper regarding definitions, claims, limitations, and scope. The paper is also trying to address the case of robust agents in POMDPs and multi-agent settings, but these sections are not rigorously treated and feel more like an afterthought rather than the main contribution. \n\nI believe the authors should focus on the main claim of the paper and strengthen this result, e.g. by improving the clarity around their definitions, the connections between those definitions, their relationship to Richens and Everitt, and by showing, with simulated data, that their algorithm works. As such, I think the paper is currently requiring significant revisions, but I encourage the authors to address the reviews and strengthen their paper.", "rebuttal": "Thank you for your review.\n\n**Q1: “Do you think your work can really be characterised as being about agents?”**\n\nWe appreciate the reviewer’s comments on the notion of agency. In this work, we implicitly adopted the definition from “Artificial Intelligence: A Modern Approach” by Russell and Norvig (1995), where an agent is any entity that maps percepts to actions to maximize expected utility. Under this definition, even simple systems such as smart sprinkler controllers qualify as agents, provided they map percepts (month of the year) to actions (activation of the sprinkler) and they maximize expected utility (keeping the lawn humid while minimizing water consumption). This broad definition is a strength as it enables our results to apply across a wide range of systems. We will clarify this point in the revision to avoid confusion about the intended scope of the term agent and to set appropriate expectations about the claims.\n\n**“[...] the main argument of the paper seems to rely on the fact that we are in a one-shot decision-making setting [...]”**\n\nThe LearnCID algorithm is defined for environments that can be described with finite-dimensional CIDs. However, this does not imply that it can only be applied to one-shot decision tasks. It can also be applied to finite-horizon sequential decision tasks using similar techniques to those described in Section 3.2 for multi-decision CIDs and, as we have shown in Section 3.3, it can be utilized to recover the causal structure and the state-transition function of infinite-horizon sequential-decision tasks modeled with POMDPs through Algorithm 2.\n\n\\\n**Q2a: “Does your extension to POMDPs really reflect sequential decision-making? In particular: How does the discount factor come into play in Algorithm 2?”**\n\nThank you for pointing this out, as it allows us to highlight an important detail. Note that in Section 3.3, we use two different causal representations. First, we describe a POMDP using an infinite CID, where each utility node $U_t$ corresponds to the reward at time $t$. The goal is to find sufficient conditions under which LearnCID can be applied to recover the full causal structure and the CPTs, which we can then use to compute the state-transition function. Importantly,  all of these are **independent of the discount factor**.\nWith Algorithm 2, we show that as long as time-homogeneity holds and the number of variables in the POMDP for a single timestep is finite, we can map this infinite CID to a finite CID that includes two consecutive timesteps of the infinite CID.\nThis is the crucial step: to run LearnCID on the CID generated by Algorithm 2 we use an optimal policy oracle that for any given distribution shift returns a policy that maximizes the expected immediate reward. Together with the other assumptions, this is sufficient to recover the full causal structure of the original infinite CID describing the full POMDP and its state-transition function, **regardless of the original discount factor**. We will clarify this in the “Learning intra- and inter-temporal causal relationships” paragraph in Section 3.3, to make this distinction and the role of the discount factor clear.\n\n**Q2b: “Is there any way this algorithm can model temporal relationships besides encoding everything in the state and relying on the Markov assumption?”**\n\nYes, CIDs can model sequential decision tasks for which the Markov assumption does not hold. Suppose we have a discrete time $t$, $V_t$ the set of variables at time $t$ (including the decision). Let $k>1$ be an integer, then if there exists an edge between a node corresponding to a variable in $V_t$ and a chance node corresponding to a variable in $V_{t+k}$, the sequential decision task modelled by this CID does not satisfy the Markov assumption. If the total number of nodes is finite, then, if the assumptions in the paper hold, we can run LearnCID and it would terminate returning all the causal relationships and the CPTs of the ancestors of the utility node. If the size is infinite but there is a structure that is periodically repeated, we can still apply LearnCID in a similar way to how we apply it to POMDPs even if the Markov assumption does not hold.\n\n**Q2c: “Is there a notion of a goal here, e.g. can $U_t$ correspond to a long-term sparse reward signal?”**\n\nIn our framework, goals are represented by utility functions. As long as the CID produced by Algorithm 2 satisfies the LearnCID assumptions, the utility is allowed to correspond to a sparse reward signal, taking the value $0$ for most states and nonzero values only for a limited subset of states. In particular, the utility function must satisfy A9, which states that there does not exist a decision $d$ that is optimal regardless of the instantiation of the other variables that are needed to compute the utility (other parents of $U_t$). This assumption can be satisfied by sparse rewards and can be verified by evaluating the reward function for different combinations of its arguments. \n\nAs an example, consider an environment with one variable “agent’s position” that takes values in the set {$1,2,3,4,5$}, the agent can move right (+1) or left (-1) and it is rewarded if it moves into 3. Therefore the sparse reward $\\mathcal R(d,pos)$ is defined such that $\\mathcal R(right,2)=\\mathcal R(left,4)=1$ and $0$ otherwise. In this case $d=right$ is not optimal for $pos=4$ and  $d=left$ is not optimal for $pos=2$, therefore A9 is satisfied.\n\n\\\n**Q3: “Can you really not use causal discovery algorithms, as you are taking some sort of samples in Algorithm 1 from the oracle”**\n\nThank you for mentioning this, it is important to remark that standard causal discovery methods cannot be used in this setting. The procedure you refer to is part of $ALG_{q_{crit}}$ (Algorithm 1 in Richens and Everitt, 2024). We do not describe it in detail because we adopt it directly from their work. In that algorithm, the optimal policy oracle is repeatedly queried under different interventions (distribution shifts) produced by sampling the convex combination coefficient of a mixture of local interventions from a uniform distribution. Then the optimal policies returned by the oracle are used to compute the CPTs. Note that the oracle only provides policies for the decision node, and we do not obtain samples of any of the variables in the process. Standard causal discovery algorithms such as PC (Spirtes & Glymour, 1991), FCI (Spirtes, 2001), LiNGAM (Shimizu et al., 2006), and RESIT (Peters et al., 2014) require access to samples from the joint distribution, which is not the case in our settings.\n\n\\\n**Q4: “In Algorithm 1, how are the edges of the CID $V$ updated to give $E’$?”**\n\nThe set of edges is updated on line 13 when $Z_i$ is marked as a parent of $X$, meaning that the edge $Z_i \\rightarrow X$ is in $E’$. We agree that not mentioning $E’$ explicitly may be confusing, we will make this clearer in the revision.\n\n\\\n**Q5: “How do we or can we know that the base CID assumed in A2 is correct to begin with?”**\n\nA2 states that the partition of the nodes $V$ into $D$ (decision nodes), $U$ (utility nodes), and $C$ (chance nodes) is known. If you have doubts about assuming the knowledge of $V$, please refer to our response to W1 of reviewer **NWB8**. Once $V$ is specified, the partition follows naturally: each utility function, which we assume to be fully specified (A7), corresponds to a utility node forming the set $U$. The decision nodes $D$ represent agent outputs (e.g. the predicted label in an image classification task). All remaining variables constitute the chance nodes $C$.\n\nIn general, there are several ways to test causal models. Causal models like CIDs encode conditional independence relationships between variables, which can be recovered from the causal graph using the concept of d-separation (Pearl, 2009). These relationships are testable and can be used to reject a model. Other techniques include randomized controlled trials, A/B tests, and comparing the model with domain knowledge.\n\n**Q5a: “If it is correct, then doesn't that already imply that agents have a causal model of the world to begin with?”**\n\nNo. In this context, CIDs are simply a tool for jointly representing causal knowledge and decision-making tasks. The paper does not claim that robust agents have an internal implicit CID that they use to perform causal reasoning. Rather, robust agents encode the information needed to recover the CID, which serves as a tool for us to capture and visualize the causal knowledge we elicit from them. All five assumptions about the CID structure, including A2, make no claims about the agents’ internal representations of the world. Their role is solely to specify the conditions and clarify what must be known in advance about the CID representation of the decision-making task and the environment to apply LearnCID to extract the agent’s causal knowledge.\n\n\\\n**Limitations:** \\\n**“I do think the results are really valuable, but their interpretation would be much easier if the implied claim that all agents must be learning a CID were put into context [...]”**\n\nWhile we believe the paper is already transparent about its assumptions, we agree that readers focusing only on the introduction may misinterpret the results. To address this, we will add more context when presenting our results by also providing an intuitive description of the assumptions, helping ensure the contribution is properly framed and minimizing the risk of overgeneralization or misrepresentation in subsequent citations.\n\nWhy opting for **CID-style causality?**\n\nCIDs build on results from causal inference and decision theory, particularly in complex scenarios like multi-agent environments and sequential decision-making. In such contexts, it is harder to represent the interplay between actions and outcomes with standard causal graphs. We will expand the discussion to better articulate these advantages and clarify why CIDs are a natural choice for our setting."}]

Author rebuttal:

Meta-review: The paper significantly extends the work of Richens and Everitt (2024) to the mediation settings and provide and work out the necessary theory.  Furthermore they provide an algorithm for mediated single-agent scenarios and multi-agent environments.  While not all reviewers are convinced I found many of the critiques of the negative ones somewhat fluffy and not really helpful.  On the contrary, the rebuttal of the authors are quite convincing, so I suggest to accept this paper!
Decision: Accept (poster)