Robust agents learn causal world models — part 2/8

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

Published as a conference paper at ICLR 2024
UD
YX
(a) training
UD
YX
(b) testing
UD
YX
(c) known domain shift
Figure 1: CID for a supervised learning task during (a) training and (b) testing following a distribu-
tional (covariate) shift (unsupervised domain adaptation, Wilson & Cook, 2020). The agent chooses
a label prediction D = ˆY given featuresX, with the goal of minimising loss U =−Loss(Y, ˆY ).
Decision variables are depicted as square nodes, chance variables as circular nodes and utilities as
diamond nodes. Information edges (dashed) show the variables the agent conditions their policy on.
In this example the labels cause the featuresY →X (for examples where features cause labels see
Castro et al., 2020; Schölkopf et al., 2012). The black square (‘regime node’ (Correa & Bareinboim,
2020)) in (b) and (c) denotes a distributional shift induced by an intervention on X. Diagram (c)
depicts the idealised case where the agent knows what domain shift has occurred. By theorem 1,
if the agent can return an optimal decision boundary for known covariate and label shifts, then it
must have learned the CBN overC ={X,Y}. Note that even if the agent has sufficient training data
to learnP (X,Y ), the causal structureY →X is in general non-identifiable givenP (X,Y ) and so
domain adaptation requires that the agent solves a non-trivial causal discovery problem.
environment state and the agent’s decision. Examples of unmediated decision tasks include prediction
tasks such as classification and regression, whereas examples of mediated decision tasks that are
not covered by our theorems include Markov decision processes where the agent’s decision (action)
influences the utility via the environment state.
2.3 D ISTRIBUTIONAL SHIFTS
We focus on generalisation that goes beyond the iid assumption, where agents are evaluated in
domains that are distributionally shifted from the training environment. Distributional shifts can
be changes to the environment (domain shifts), as in domain adaptation and domain generalisation
(Farahani et al., 2021; Wilson & Cook, 2020), or changes to the objective (task shifts) as in zero shot
learning (Xian et al., 2018), in-context learning (Brown et al., 2020) and multi-task reinforcement
learning (Reed et al., 2022). Our analysis focuses on domain shifts that involve changes to the causal
data generating process, and hence can be modelled as interventions (Schölkopf et al., 2021). This
does not assume that all shifts an agent will encounter can be modelled as interventions, but requires
that the agent is at least capable of adapting to these shifts.
Examples of interventionally generated shifts include translating objects in images (Engstrom et al.,
2019), noising inputs and adversarial robustness (Hendrycks & Dietterich, 2019), and changes to the
initial conditions or transition function in Markov decision processes (Peng et al., 2018). Examples
of shifts that are not naturally represented as interventions include changing the set of environment
variablesC, and introducing selection biases (Shen et al., 2018). See Appendix A.3 for discussion.
Our main results restrict to local domain shifts, which correspond to local interventions on the chance
variablesC. We do not consider shifts that change the agent’s decisionD, although we include shifts
that drop inputs to the policy PaD→ Pa′
D⊆ PaD (e.g. masking) as local interventions. We do not
consider task shifts i.e. changing the utility function.
As we are interested in determining the capabilities necessary for domain adaptation, we restrict
our attention to decision tasks where domain adaptation is non-trivial, i.e. where the optimal policy
depends on the environment distributionP (C =c).
Assumption 2 (Domain dependence). There existsP (C =c) andP′(C =c) compatible withM
such thatπ∗ = arg maxπ Eπ
P [U] impliesπ∗̸= arg maxπ Eπ
P′[U].
Assumption 2 implies the existence of domain shifts that change the optimal policy.
4

Published as a conference paper at ICLR 2024
3 C AUSAL MODELS ARE NECESSARY FOR ROBUST ADAPTATION
We now present our results, showing that learning the underlying CBN is necessary and sufficient for
learning regret bounded policies under a large set of domain shifts. In Section 3.2 we interpret these
results by considering their consequences for transfer learning, causal inference and adaptive agents.
First we focus on the idealised case where we assume optimality. We show for almost all decision
tasks the underlying CBN can be reconstructed given optimal policies for a large set of domain shifts.
Theorem 1. For almost all CIDs M = (G,P ) satisfying Assumptions 1 and 2, we can identify
the directed acyclic graphG and joint distributionP over all ancestors of the utility AncU given
{π∗
σ(d| paD)}σ∈Σ whereπ∗
σ(d| paD) is an optimal policy in the domainσ and Σ is the set of all
mixtures of local interventions. Proof in Appendix C.
The parameters P (vi | pai), U(paU) of the underlying CBN define a parameter space and the
condition for almost all CIDs means that the subset of the parameter space for which the Theorem 1
does not hold is Lebesgue measure zero (see Appendix A.2 for discussion). This condition is
necessary because there exist finely-tuned environments for which the CBN cannot be identified given
the agent’s policy due to variablesX∈ AncU that do not affect the expected utility. For example
considerX→ Y → U,Y =N (0,x ) andU = D +Y , then changing X can only change the
variance ofU while leaving its expected value (and hence the optimal policy) constant. However, this
only occurs for very specific choices of the parametersP andU.
In Appendix B we give a simplified overview of the proof with a worked example. We assume access
to an oracle for optimal policiesπ∗
σ for any given local intervention onC. Note, this assumes the
agent is robust to distributional shifts on a causally sufficient set of variablesC, not that the set of
variables the agent observes is causally sufficient. We devise an algorithm that queries this oracle with
different mixtures of local interventions and identifies the mixtures for which the optimal policies
changes. We then show that these critical mixtures identify the parameters of the CBN, specifying
both the graphG(AncU) and the joint distributionP (AncU).
3.1 R ELAXING THE ASSUMPTION OF OPTIMALITY
We now relax the assumption of optimality, considering the case where the policies πσ satisfy a
regret bound Eπσ[U]≥ Eπ∗
σ[U]−δ. We show that forδ >0 we can recover an approximation of the
environment CBN, with error that grows linearly inδ forδ≪ Eπ∗
[U].
Theorem 2. For almost all CIDs M = (G,P ) satisfying Assumptions 1 and 2, we can iden-
tify an approximate causal model M′ = (P′,G′) given{πσ(d | paD)}σ∈Σ where Eπσ[U] ≥
Eπ∗
σ[U]−δ and Σ is the set of mixtures of local interventions. The parameters of M′ satisfy
|P′(vi| pai)−P (vi| pai)|≤ γ(δ)∀Vi∈ V whereγ(0) = 0 andγ(δ) grows linearly in δ for
small regretδ≪ Eπ∗
[U]. Proof in Appendix D.
The worst-case error boundsγ(δ) for the parameter errors are detailed in Appendix D. Forδ >0 it
may not be possible to identifyG perfectly as some weak causal relations cannot be resolved due
to these error bounds. We describe in Appendix D how we can learn a sub-graphG′⊆G that may
exclude directed edges corresponding to weak causal relations.
Theorem 2 shows that we can learn a (sparse) approximate causal models of the data generating
process from regret bounded policies under domain shifts, where the approximation becoming exact
asδ→ 0. In Appendix F we demonstrate learning the underlying CBN from regret-bounded policies
using simulated data for randomly generated CIDs similar to Figure 1, and explore how the accuracy
of the approximate CBN scales with the regret bound (Figure 3).
Finally, we prove sufficiency, i.e. that having an (approximate) causal model of the data generating
process is sufficient to identify regret-bounded policies. The result is well-known for the non-
approximate case (Bareinboim & Pearl, 2016).
Theorem 3. Given the CBNM = (P,G ) that is causally sufficient we can identify optimal policies
π∗
σ(d| paD) for any givenU where PaU⊆C and for all soft interventionsσ. Given an approximate
causal model M′ = (P′,G′) for which|P′(vi| pai)−P (vi| pai)| ≤ϵ≪ 1, we can identify
regret-bounded policies where the regretδ grows linearly inϵ. Proof in Appendix E.
5

Published as a conference paper at ICLR 2024
(a) Error rate for learned DAG v.s. regret bound
 (b) Mean error forP (x,y ) v.s. regret bound
Figure 3: Comparing the model-average error rates for a) the learned DAGG′ and b) learned joint
distributionP′(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) See Appendix F for details.
Together, Theorems 2 and 3 imply that learning an approximate causal model of the data generating
process is necessary and sufficient for learning regret-bounded policies under local interventions.
3.2 I NTERPRETATION
We interpret Theorems 1 to 3 through three lenses; agents, transfer learning and causal inference.
Agents are adaptive goal-directed systems (Dennett, 1989), meaning they choose actions they expect
to achieve some desired outcome, and would adapt their behaviour if they knew the consequences of
their actions had changed. For example, a firm sets prices to maximise profit, and adapts pricing to
changes to supply and demand (Kenton et al., 2023).
Consider an agent capable of zero-shot generalisation under domain shifts (Kirk et al., 2023). First,
consider the case where the agent has perfect knowledge of the domain shift, i.e. the policy conditions
on the domain indicator3πσ =π(d| paD,σ ). By Theorem 2 ifπσ satisfies a tight regret bound for
all local shiftsσ, we can reconstruct the underlying CBN from the agent’s policy alone (following the
procedure in Appendix C). Therefore, any agent capable of generalising under known local domain
shifts must have learned the underling CBN. Precisely, the agent has learned the policyπ(d| paD,σ ),
which is functionally equivalent to learning the underlying CBN, as any associative or causal query
that can be identified using the CBN can also be identified usingπ(d| paD,σ ). Ifσ is unknown and
has to be inferred by the agent, the underlying CBN can similarly reconstructed from the agent’s
policy as long as it satisfies a tight regret bound4.
Example: Doctors are agents expected to make low regret decisions under a wide range of known
distributional shifts, without re-training in the shifted environment. For example, consider the task of
risk-stratifying patients based on their signs and medical history. The doctor may be transferred to a
new ward where patients have received a treatment (known distributional shift) that has a stochastic
effect on latent variables (mixed intervention) such as curing diseases and causing side effects. The
doctor cannot re-train in this new domain, e.g. taking random decisions and observing outcomes. To
be capable of this, Theorem 2 implies the doctor must have learned the causal relations between the
relevant latent variables—how the treatment affects diseases, how these diseases and their symptoms
are causally related, and so on. Likewise, any medical AI that hopes to replicate this capability must
have learned a similarly accurate causal model, and the better the agent’s performance the more
accurate its causal model must be.
3equivalently an environment index for (in-context) invariant risk minimization (Gupta et al., 2023; Arjovsky
et al., 2019), or the context variable for zero-shot reinforcement learning (Kirk et al., 2023)
4Trivially, any agent capable of returning a regret-bounded policy without inputσ can do so givenσ simply
by discardingσ. If the agent adapts by conditioning their policy on an inferredσ′ (see Kirk et al. (2023) for a
review of implementations), and this policy satisfies a regret bound, it can be used in Theorems 1 and 2.
6

Published as a conference paper at ICLR 2024
Transfer learning. In transfer learning (Zhuang et al., 2020), models are trained on a set of source
domains and evaluated on held-out target domains where i) the data distribution differs from the source
domains, and ii) the data available for training is restricted compared to the source domains (Wang
et al., 2022). For example, in unsupervised domain adaptation the learner is restricted to samples
of the input features from the target domains paD∼P (PaD;σ), whereas in domain generalisation
typically no data from the target domain is available during training (Farahani et al., 2021).
LetDS denote the training data from the source domains andDσ denote the training data available
from a given target domainσ. Let there exist a transfer learning algorithm that returns a policyπσ
satisfying a regret bound for a given target domainσ, provided this training data. Asπσ is a function
of the training data, then by Theorems 1 and 2 the existence of this algorithm implies that we can
identify the underlying CBN fromDS∪{Dσ}σ∈Σ. To see that this imparts non-trivial constraints on
the existence of the transfer learning algorithm, we can consider the following simple example.
Example: Consider the CID for the supervised learning task depicted in Figure 1. Let DS =
{(xi,yi)∼ P (X,Y )}n
i=1, so for sufficiently large n the agent can learn the P (X,Y ) fromDS.
However,Y →X must also be identifiable from the training dataDS∪{Dσ}σ∈Σ. In other words,
the transfer learning problem contains a hidden causal discovery problem. IfDσ =∅ thenY →X
must be identifiable fromP (x,y ) alone, which is impossible unless the causal data generating process
obeys additional assumptions (see for example Hoyer et al., 2008). If unlabelled features from the
target domain are included in the training dataDσ ={xi∼P (X;σ)}nσ
i=1,Y →X can in principle
be identified asP (X;σY )̸=P (X).
Causal inference. Theorem 1 can also be interpreted purely in terms of causal inference. We can
compare to the causal hierarchy theorem (CHT) (Bareinboim et al., 2022), which states that an oracle
for L1 queries (observational) is almost always insufficient to evaluate all L2 queries (interventional).
Our Theorem 1 can be stated in an analogous way; an oracle for optimal policies under mixtures of
local interventions Π∗
Σ :σ7→π∗(σ), can evaluate all L2 queries, which follows from the fact that the
oracle identifies the underlying CBN which in turn identifies all L2 queries. Note Π∗
Σ is a strict subset
of L2, and we describe a subset of L2 as being L2-complete if evaluating these queries is sufficient to
evaluate all L2 queries. Hence Theorem 1 can be summarised as Π∗
Σ is L2-complete. It would be
interesting in future work to determine what other strict subsets of L2 are L2-complete, as identifying
these queries is sufficient to identify all interventional queries.
Why is this surprising? Firstly, we may expect the optimal policies to encode a relatively small
number of causal relations, as they can be computed from E[u| d, paD;σ], which describes the
response of a single variable U to intervention σ. However, Theorem 1 shows that the optimal
policies encode all causal and associative relations in AncU , including causal relations between latent
variables, for exampleP (Yx) for anyX,Y ⊆ AncU . Secondly, Theorems 2 and 3 combined imply
that learning to generalise under domain shifts is equivalent to learning a causal model of the data
generating process—problems that on the surface are conceptually distinct.
Figure 4: The set L2 (Bareinboim et al., 2022) contains all casual queries including the set of optimal
policies under domain π∗
C and task shifts π∗
U . We show that π∗
C also contains L2, i.e. learning
optimal policies under all shifts for a single utilityU is sufficient to identify L2
7