Robust agents learn causal world models — part 8/8

Cleaned-text code points [99366, 110730). Bibliography omitted here; source.md includes it.

Published as a conference paper at ICLR 2024
functionU, iii) the decision task is unmediated and iv) domain dependence. Given these assumptions
we can learn an approximation of the underlying CBN given only the policy of the agentπ(σ) under
interventionsσ, with the approximation being exact whenπ(σ) are optimal.
To demonstrate this theoretical result we take the proof for simple Binary decision tasks outlined in
Appendix B and recast it as a causal discovery algorithm (Algorithm 2 below). We test it on CIDs
of the form shown in Figure 6 where we randomly choose the joint distribution overX,Y and their
causal structureG. Note that Algorithm 2 is significantly simpler than the general method outlined
in the proof of Theorem 1, as it exploits the fact thatD,X,Y are binary variables and that|C| = 2.
This causal discovery algorithm requires that we can intervene on the latent variablesX,Y , but only
requires that we can observe the response of a single variable (the decision) to these interventions. To
motivate this setting, we can imagine situations where the latentsX,Y cannot be directly observed
but can be intervened on.
Example. Many diseases cannot be directly observed in patient physiology, but can only be indirectly
observed through the presence of symptoms. Let X,Y ∈{ 0, 1} be two such diseases, for which
there are treatments, i.e. we can intervene to ‘turn off’X andY but cannot observe them.D∈{ 0, 1}
represents a decision to provide a specific pain relief medication, which results in a change in the
symptom severity (utility). The response to pain relief depends on the presence or absence of the
diseases (e.g pain relief is highly effective for patients withX =T , moderately effective forY =T
and less effective forX = F,Y = F ). The doctor’s goal is to minimise symptom severity while
avoiding unnecessary use of pain medication, e.g.U(d,x,y ) =d[s(x,y )−c] wherec is some cost
associated with pain relief and s(x,y ) is the response to pain relief. Following an intervention σ
(e.g. curing a disease σ = do(X = F )), the doctor adapts their treatment policy in the shifted
population. For example, this adaptation could occur by trial and error, with the doctor choosing
random treatment decisionsD and observing the change in symptom severity—a context-free bandit
problem. Although we cannot directly observe the disease statesX,Y , by intervening on the latent
disease state and observing how the doctor’s policy adapts, we can learn both the joint distribution
P (X,Y ) and the causal graph overX,Y .
Figure 9 shows the average error in the learned parametersP (x,y ) andG whenπ(σ) satisfy different
regret bounds. The results are averaged over 1000 randomly generated CBNs where i) the parameters
of the joint distributionP (x,y ) are chosen at random, ii) the DAGG overX,Y is chosen at random
fromX→ Y andX← Y , iii) the utility function U(d,x,y )∈ [0, 1] is chosen at random (see
Appendix A.2 for description of parameters). To simulate the regret-bounded agent we calculate the
optimal policy for each environment and if the sub-optimal decision satisfies the regret bound we
choose randomly from the two decisions when sampling from the policy oracle in Algorithm 1. We
also compare to a random baseline algorithm which estimatesP (x,y ) = 1/4 and randomly selects
fromX→ Y orX← Y with equal probability. In a small number of cases Algorithm 1 fails to
predictP (x,y )∈ [0, 1] due to finite sample errors, and for these cases we replace the output of the
causal discovery algorithm with a random guess.
From Figure 9 it appears that the error rate grows sub-linearly with regret. Note that the relevant
scale for the regret is the difference in expected utility between the two decisions, hence we plot the
normalised regret bound where we divideδ by this expected utility difference. Note that even for
relatively large regret bounds, representing agents that generalise weakly, we can still identify the
causal structure with a high accuracy. For example when the regret bound is 30% of the expected
utility difference, we can still identify the correct causal structure in∼ 90% of the randomly generated
CIDs. This describes an agent that is guaranteed to incur a regret of at most 30% of the expected
utility difference between the decisions before the domain shift. If the domain shift results in the
expected utility difference being less that 30% of the unshifted expected utility difference, the agent
can return a sub-optimal decision.
G A PPENDIX : TRANSPORTABILITY & PEARL’S CAUSAL HIERARCHY
Transportability. The problem of evaluating policies under distributional shifts has been studied
extensively in transportability theory (Pearl & Bareinboim, 2011; Bareinboim & Pearl, 2016; Bellot
& Bareinboim, 2022). For decision tasks as outlined in Section 2.2, transportability aims to provide
necessary and sufficient conditions for identifying the expected utility following a distributional shift,
R = E[u|d, paD;σ], given (partial) knowledge of i) the jointP , causal graphG and interventional
29

Published as a conference paper at ICLR 2024
(a) Misclassification rate for G scaling with regret
bound
(b) Mean parameter error for P(x, y) scaling with
regret bound
(c) Worst-case error for P(x, y) scaling with regret
bound
Figure 9: Comparing the model-average error rates for a) the learned DAG and b) the mean error
for parameters P (x,y ) and c) the worst-case error for parameters P (x,y ), v.s. the (normalised)
regret boundδ/|E[u|D = 1]− E[u|D = 0]|. Average error taken over 1000 randomly generated
environments with binary decisionD and two binary latent variablesX,Y . Comparison to error rate
for random guess (green). Results appear to show sub-linear growth in error rate with regret bound.
Note that even weakly generalising agents can be used to identify causal structure significantly better
than the random baseline.
30

Published as a conference paper at ICLR 2024
Algorithm 2 Graph Learner for simple CID
1: function GRAPH LEARNER (Πδ
Σ,U,δ,N)
2: d1,d 2,x′,y′,q crit← Algorithm 1(U, Πδ
Σ,N,σ 1 = do(Y = 0)) ▷ Identifyqcrit for
do(Y = 0)
3: Exp. U difference = (U(d2,x′,y′)−U(d1,x′,y′))∗ (1− 1/qcrit)
4: ∆0 =U(0, 0,d 2)−U(0, 0,d 1)
5: ∆1 =U(1, 0,d 2)−U(1, 0,d 1)
6: P (XY =0 = 0) = (Exp. U difference− ∆1)/(∆0− ∆1)
7:
8: d1,d 2,x′,y′,q crit← Algorithm 1(U, Πδ
Σ,N,σ 1 = do(Y = 1)) ▷ Identifyqcrit for
do(Y = 1)
9: Exp. U difference = (U(d2,x′,y′)−U(d1,x′,y′))∗ (1− 1/qcrit)
10: ∆0 =U(0, 1,d 2)−U(0, 1,d 1)
11: ∆1 =U(1, 1,d 2)−U(1, 1,d 1)
12: P (XY =1 = 0) = (Exp. U difference− ∆1)/(∆0− ∆1)
13:
14: d1,d 2,x′,y′,q crit← Algorithm 1(U, Πδ
Σ,N,σ 1 = do(X = 0)) ▷ Identifyqcrit for
do(X = 0)
15: Exp. U difference = (U(d2,x′,y′)−U(d1,x′,y′))∗ (1− 1/qcrit)
16: ∆0 =U(0, 0,d 2)−U(0, 0,d 1)
17: ∆1 =U(0, 1,d 2)−U(0, 1,d 1)
18: P (YX=0 = 0) = (Exp. U difference− ∆1)/(∆0− ∆1)
19:
20: d1,d 2,x′,y′,q crit← Algorithm 1(U, Πδ
Σ,N,σ 1 = do(X = 1)) ▷ Identifyqcrit for
do(X = 1)
21: Exp. U difference = (U(d2,x′,y′)−U(d1,x′,y′))∗ (1− 1/qcrit)
22: ∆0 =U(1, 0,d 2)−U(1, 0,d 1)
23: ∆1 =U(1, 1,d 2)−U(1, 1,d 1)
24: P (YX=1 = 0) = (Exp. U difference− ∆1)/(∆0− ∆1)
25:
26: ifP (YX=0 = 0) =P (YX=1 = 0) then ▷ IdentifyG andP from interventionals
27: ifP (XY =0 = 0) =P (XY =1 = 0) then
28: G← ()
29: P (x,y ) =P (XY =0 =x)P (YX=0 =y)
30: else
31: G← (Y →X)
32: P (x,y ) =P (YX=0 =y)P (XY =y =x)
33: end if
34: else
35: G← (X→Y )
36: P (x,y ) =P (XY =0 =x)P (YX=x =y)
37: end if
38: returnG,P (x,y )
39: end function
31

Published as a conference paper at ICLR 2024
distributionsI in the source domain, and ii) (partial) knowledge of the joint P∗ and causal graph
G∗ in the target domain (Pearl & Bareinboim, 2011; Bareinboim & Pearl, 2012b). Hence, these
results differ from Theorems 1 and 2 in that they restrict to the case where all assumptions on the data
generating process (i.e. inductive biases) can be expressed as (partial) knowledge of the underlying
CBN. For example, Bareinboim & Pearl (2016) claim the problem is essentially solved in the case
where ‘assumptions are expressible in DAG form’. This does not constrain possible approaches to
domain generalisation that make use of non-causal assumptions and heuristics 5, and indeed deep
learning algorithms exploit a much wider set of inductive biases than causal assumptions alone
(Neyshabur et al., 2014; Battaglia et al., 2018; Rahaman et al., 2019; Goyal & Bengio, 2022; Cohen &
Welling, 2016). In many real-world tasks these may be sufficient to identify ‘good enough’ (i.e. regret-
bounded) policies without requiring knowledge of the causal structure of the data generating process.
Our aim has been to establish if learning causal models is necessary for domain generalisation in
general. Hence assuming that agents are restricted to using inductive biases that amount to (partial)
knowledge of the underlying CBN would be begging the question.
Causal hierarchy’s theorem (CHT). The celebrated causal hierarchy theorem (Bareinboim et al.,
2022; Ibeling & Icard, 2021) shows that for almost all environments there are causal relations
between environment variables that cannot be identified from observational data without additional
assumptions. Does this imply that a causal model is necessary for identifying optimal policies?
First, note that the CHT is an insufficiency result, and only implies trivial necessity results. For
example, is a causal model necessary for identifying all causal and associative relations between
environment variables? Yes, but only because this set of observational and interventional distributions
is a causal model. Formally, we can identify the underlying causal model (up to latent confounders) by
assuming causal faithfulness, which holds for almost all causal models (Meek, 2013). The difference
here is that the CHT is concerned with the identifiability of all causal and associative relations
between environment variables. This sets a much higher bar than domain generalisation, which
focuses on identifying a strict subset of these (regret-bounded policies) (Figure 4).
Secondly, the CHT is concerned with the collapse (or lack thereof) of the causal hierarchy. For
example, that observational data is insufficient for identifying all causal queries. We do not restrict
agents to having observational training data—in fact, typically we assume that agents have access to
both observational and interventional data in the online learning setting that we consider (e.g. agents
can intervene to fix the decision nodeD by assumption).
Finally, we can imagine a refinement of the CHT which states that observational data is insufficient
for identifying regret-bounded policies without additional assumptions, bringing it in line with
Theorems 1 and 2. If this was implied by the CHT, it would not imply our results unless we restrict
to the case where all assumptions as constraints on the causal structure (similar to transportability).
Likewise, it is simple to show that Theorem 1 does not imply the CHT. In deriving Theorem 1 we do
not restrict to observational distributions (or make any restrictions on the data available to the agent
when generating its policy).
5Indeed, notable examples of causal assumptions that go beyond those expressible in DAG form include
restricting the classes of structural equations Mooij et al. (2016) and assuming cause-effect asymmetry (Mitrovic
et al., 2018)
32