Robust agents learn causal world models — part 5/8

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

Published as a conference paper at ICLR 2024
Example: Logical NOT operationX→¬X for BooleanX
We also consider mixtures of interventions, which can also be described without knowledge ofG.
Definition 3 (Mixtures of interventions). A mixed interventionσ∗ =P
ipiσi forPpi = 1 performs
interventionσi with probabilitypi. Formally,P (v|σ∗) =P
ipiP (v|σi).
Example: Adding Gaussian noise is a mixture over local operations (translations) σϵ = do(X =
X +ϵ) whereϵ∼N (0, 1).
In common to most decision making tasks such as prediction, classification, and reinforcement
learning, is that a decision should be outputted based on some information to optimise some objective.
The exact terms vary: decisions are sometimes called outputs, actions, predictions, or classifications;
information is sometimes called features, context, or state; and objectives are sometimes called
utility functions or loss functions. However, all of these setups can be described within the causal
influence diagram (CID) framework (Howard & Matheson, 2005; Everitt et al., 2021). CIDs are
causal Bayesian networks where the variables are divided into decision D, utility U, and chance
variablesV , and no conditional probability distribution is specified for the decision variables. The
task of the agent is to select the distributionπ =P (D =d| PaD = paD), also known as the policy
or decision rule. An optimal policyπ∗ is defined as a policyπ∗ that maximizes the expected value of
the utility Eπ∗[U].
Definition 4 (Causal influence diagram). A (single-decision, single-utility) causal influence diagram
(CID) is a CBNM = (G,P ) where the variablesV are partitioned into decision, utility, and chance
variables,V = ({D},{U},C). The utility variable is a real-valued function of its parents,U(paU).
By convention, decision nodes are drawn square, utility nodes diamond, and chance nodes round. The
parents ofD, PaD, can be interpreted as the information the decision is allowed to depend on, and
are depicted as dashed lines. See Figure 5 for an example. In the following we restrict our attention
to a class of CIDs we refer to as ‘unmediated decision tasks’, where where the agent’s decision does
not causally influence any chance variables that go on to influence the utility. This simplifies our
theoretical analysis, although it is likely that our results extend to the general case.
Assumption 1 (Unmediated decision task). DescD∩ AncU =∅.
Examples of unmediated decision tasks include all standard classification and regression tasks, and
generative AI tasks where the output is not included in the training set. For example, in classification
typically the choice of label does not influence the data generating process. Problems that are
mediated rather than unmediated decision tasks includes most control and reinforcement learning
tasks, where the agent’s decision is an action that influences the state of the environment. Furthermore,
we will focus on non-trivial unmediated decision tasks i.e. whereU∈ ChD, as the case ChD =∅
describes trivial decision tasks (the agent’s action does not influence the utility). Figure 5 is an
example of a non-trivial unmediated decision task.
In transfer learning we are typically interested in problems where generalising from the source
to target domain(s) is non-trivial, and in the trivial case we cannot expect agents to have to learn
anything about their environments in order to generalise. If this is not the case, then generalising
under distributional shifts is trivial. Therefore we restrict our attention to decision tasks where the
distribution of the environment is relevant to the agent when determining its policy. Specifically, we
say that a decision task is domain independent if there exists a single policy that is optimal for all
choices of 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].
Lemma 1. Domain dependence implies that;
i) There exists no d∈ dom(D) such thatd∈ arg maxdU(d,c )∀c∈ dom(C).
ii) PaD ⊊ AncU
iii) D∈ PaU
Proof. i) For any P′ we have EP′[U | do(D = d), paD] = P
cP′(Cd = c | paD)U(d,c)
= P
cP′(C = c| paD)U(d,c ) where we have used DescD∩ AncU =∅. Therefore if ∃ d∗
15

Published as a conference paper at ICLR 2024
C1
C2
C3C4
C5 C6
D U
Figure 5: The CID for an unmediated decision task, where D has no causal influence on the
environment stateC. Our main theorem implies that an agent that is robust to distributional shifts
onC must learn the CBN over AncU ={C1,C 2,C 3,C 4,C 6}, noting that C5̸∈ AncU . C4 is an
example of a variable that is only an ancestor ofU viaD and so has no direct causal effect on the
utility, but is still relevant to the decision task as it is a proxy for C1 which is a cause of U. C6 is
a cause of U but not of D, and naively one might assume that distributional shifts on C6 cannot
influence the agent’s decision. However, the optimal policy can change under distributional shifts on
C6 as these effect the utility, and hence the agent will have to learn a CBN includingC6 if it is to be
robust to shifts onC6.
s.t. d∗ = arg maxdU(d,c)∀ C = c then EP′[U | do(D = d∗), paD] ≥ P
cP′(C = c |
paD)U(d′,c)EP′[U| do(D =d), paD]∀d̸=d∗, and soD =d∗ is optimal for allP′(C =c) and
we violate domain dependence.
ii) AsD∈ PaU (iii), then PaU⊆ AncU . If AncU = PaD then EP [u|d, paD] =U(d, paD) which
is independent of P (C =c), and hence there is a single optimal policy for all P and we violate
domain dependence.
iii) IfD̸∈ AncU then the CID is trivial, in the sense that E[U| do(D =d)] = E[U], and hence all
decisions are optimal for all distributionsP (C), which violates domain dependence (Assumption 2).
ThereforeD∈ AncU which with DescD∩ AncU =∅ impliesD∈ PaU .
A.2 P ARAMETERISATION OF CID S
The joint distributionP is defined for all environment variablesC, and the CID is defined by the
parameters forP (C) andU(PaU). We restrict our attention toC that are categorical, and without loss
of generality we label statesci = 0, 1,..., dimi−1 where dimi is the dimension of variableCi. Firstly,
the jointP (C) is parameterised by the conditional probability distributions (CPDs) in the Markov
factorization with respect to G,θP ={P (ci| pai)∀ci∈{ 0,..., dimi− 2}, pai∈ Pai, Ci∈C.
Note that the CPDsp(Ci = dimi− 1| pai) are not included inθP as they are fully constrained by
normalizationP (Ci = dimi− 1| pai) = 1−Pdimi−1
j=0 P (ci| pai). Secondly, the utility function is
simply parameterised by its value given the state of its parentsθU ={U(paU)∀ PaU = paU}. For
simplicity we work with the normalized utility function,
U(paU)→
U(paU)− minpa′
U
U(PaU = pa′
U)
maxpa′
U
U(PaU = pa′
U)− minpa′
U
U(PaU = pa′
U) (2)
with values between 0 and 1. Noting that as this is a positive affine transformation of the utility
function the set of optimal policies invariant, and we can re-scale regret bounds accordingly. LetθM
denote the set of all parameters for the CID,θM =θP∪θU , and note that the elements ofθM in the
[0, 1] interval and are logically independent, i.e. we can independently choose any [0, 1] value for
each parameter and this defines a valid parameterization of the CID for the baseline environment. In
the following when we refer to ‘the parametersP,U ’ we are referring toθM .
We follow the method outlined in (Meek, 2013) to prove that certain constraints onP,U hold ‘for
almost allP,U ’ and hence for almost all decision tasks. This involves converting a given constraint
into polynomial equations overθM and applying the following Lemma,
Lemma 2 (Okamoto, 1973). The solutions to a (nontrivial) polynomial are Lebesgue measure zero
over the space of the parameters of the polynomial.
16

Published as a conference paper at ICLR 2024
A polynomial inn variables is non-trivial (not an identity) if not all instantiations of then variables
are solutions of the polynomial. For example, the equation poly(θM) = 0 is trivial if and only if all
coefficients of the polynomial expression poly(θM) are zero. Therefore, any constraint onP,U that
can be converted into a polynomial equation overθM must either hold for allθM or for a Lebesgue
measure zero subset of instantiations ofθM .
Operationally, this means that if we have any smooth distribution over the parameter space (for
example, describing the distribution of environments we expect to encounter), the probability of
drawing an environment from this distribution for which the condition does not hold is 0.
A.3 D ISTRIBUTIONAL SHIFTS & POLICY ORACLES
In the derivation of our results we restrict out attention to distributional shifts that can be modelled
as (soft) interventions on the data generating process. We note that by Reichenbach’s principle
(Reichenbach, 1956), which states that all statistical associations are due to underlying causal
structures, we can assume the existence of a causal data generating process that can be described
in terms of a CBN M = (P,G ). Therefore there is a causal factorization of the joint P (C =
c) =Q
iP (ci| Pai). By allowing for mixtures of interventions, we can reach any distribution over
C, which can be seen trivially by noting that we can perform a soft intervention to achieve any
deterministic distributionP (C =c) =δ(C =c′), and then take a mixture over these deterministic
distributions to achieve an arbitrary distribution over C. The set of distributions that cannot be
generated by interventions include those that change the set of variablesV including the decision and
utility variables, and introducing selection biases (which are causally represented with the introduction
of additional nodes that are conditioned on Bareinboim & Pearl, 2012a). For further discussions on
the relation between distributional shifts and interventions see Schölkopf et al. (2021); Meinshausen
(2018).
In the following proofs we use policy oracles to formalise knowledge of regret-bounded behaviour
under distributional shifts.
Definition 5 (Policy oracle). A policy oracle for a set of interventions Σ is a map Πδ
Σ :σ7→πσ(d|
paD)∀σ∈ Σ where Σ is a set of domains. It isδ-optimal ifπσ(d| paD) achieves an expected utility
Eπσ[U]≥ Eπ∗
[U]−δ in the CIDM(σ) whereδ≥ 0.
Hereδ is the regret upper bound, which is satisfied under all distributional shiftsσ∈ Σ. We refer
toδ-optimal policy oracles for δ = 0 as optimal policy oracles. For the proof of our main result
we restrict our attention to policy oracles with Σ that includes mixtures over all local interventions
(def. 2).
Note that the policy oracle specifies only what policy the agent returns in a distributionally shifted
environmentM(σ). It does not specify how this policy is generated, which will depend on the
specific setup. For example, in domain generalisation that agent typically receives no additional data
from the target domains, and is expected to produce a policy (decision boundary) that achieves a low
regret across all target domains. On the other hand in domain adaptation and few shot learning, the
agent is provided with some new data from each target domain with which to adjust its policy. As we
hope to accommodate all of these perspectives we specify only the agent’s policy, not the data used
to generate it. This is discussed further in Section 3.2.
What distributional shifts do we consider? In our proofs, we assume the agent is robust to any
domain shifts that can be described as a mixture of local interventions on the environment variables
C. We do not consider interventions that change the utilityU or the agent’s decisionD, though we
do include dropping inputs to the policy (masking) PaD→ Pa′
D⊆ PaD as local interventions.
B A PPENDIX : S IMPLIFIED PROOF
In this section we outline the proof of Theorem 1 for a simple binary decision task with binary
latent variables. As mentioned in Section 4, the method used to identify the CBN in Theorem 1 can
be viewed as an algorithm for learning the CBN over latent variables by observing the policy of a
regret-bounded agent under various distributional shifts. To demonstrate this, in Appendix F we use
an implementation of the algorithm on randomly generated CIDs, showing empirically that we can
17

Published as a conference paper at ICLR 2024
learn the underlying CBN in this way, and explore how the agent’s regret bound affects the accuracy
of the learned CBN.
Consider the CID in Figure 6, describing a binary decision taskD∈{ 0, 1} with two binary latent
variablesX,Y ∈ PaU .
U
X Y
D
Figure 6: Example CID describing a context-free mutli-armed bandit with binary latent variables
X,Y .
Consider an agent that selects a policyπD such that it maximises the expected utility. That is, the
CID describes a context-free bandit problem, whereX,Y are latent variables that influence the arm
values E[u|d] =P
x,yP (x,y )U(x,y,d ).
Our aim is to learn this CID given only knowledge of the agent’s policy under distributional shifts,
and knowledge that it satisfies a regret bound. We assume knowledge of i) the set of chance variables
C ={X,Y}, ii) the utility functionU(d,x,y ), and iii) the policyπD(σ) under distributional shifts
σ (other variables (U,X,Y ) are unobserved). To learn the CID the aim is therefore to learn the
parameters of the joint distribution over latentsP (x,y ) and the unknown causal structure. As we know
the utility function we knowD,X,Y ∈ PaU , and by assuming the CID is unmediated (Assumption 1)
we knowX,Y ̸∈ DescD. Likewise the decision task is context free hence D̸∈ DescX∪ DescY .
Hence the only unknown causal structure is the DAG over the latent variablesC ={X,Y}.
The expected utility difference betweenD = 0 andD = 1 following a hard intervention onX is
given by
E[u|D = 0; do(X = 0)]− E[u|D = 1; do(X = 0)] =
X
y
P (YX=0 =y)[U(0, 0,y )−U(1, 0,y )]
(3)
=P (YX=0 = 0)[U(0, 0,Y = 0)−U(1, 0, 0)] + (1−P (YX=0 = 0))[U(0, 0, 1)−U(1, 0, 1)]
(4)
As we know U(d,x,y ) we can therefore identify P (YX=0 = 0) if we can identify this expected
utility difference. We do this using the agent’s policy under distributional shifts, and in this simple
case we can restrict our attention to hard interventions. Following the steps outlined in Lemma 4,
domain dependence insures that we can identify a hard interventionσ2 = do(X =x′,Y =y′) that
results in a different optimal policy to the optimal policy under σ1 = do(X = 0). For a mixture
of these two interventions σ3 = qσ1 + (1−q)σ2 the expected utility is E[u| d,σ 3] = qE[u|
d,σ 1] + (1−q)E[u|d,σ 2]. This is a linear function with respect to q, and forq = 1 the optimal
decision (d1) is different than forq = 0 (d2̸=d1). Therefore, there is a single indifference pointqcrit
for which both decisions are optimal. It is simple to show that this indifference point is given by,
qcrit =

1− E[u|D =d1; do(X = 0)]− E[u|D =d2; do(X = 0)]
U(d1,x′,y′)−U(d2,x′,y′)
−1
(5)
D =d1 is optimal forq≤qcrit andD =d2 is optimal forq≥qcrit. We can estimateqcrit by randomly
sampling values ofq uniformly over [0, 1] and observing the optimal decision under the resulting
mixed intervention (Algorithm 1). That is, qcrit is the probability that D = d1 is returned by the
policy oracle for a randomly sampledq. In this way we learnqcrit and as we knowU(d,x,y ) we can
identify the expected utility difference under do(X = 0) in the numerator of Equation (5) and so
identifyP (YX=0 = 0).
Similarly we identify P (YX=1 = 0),P (XY =0 = 0) and P (XY =1 = 0) , which encode both
the causal relation between X and Y (e.g. there is a directed path from X to Y if and only if
P (YX=0)̸=P (YX=1) for almost all CBNs), and determine the parameters of the CBN asP (Ci =
ci| do(C\Ci)) =P (Ci =ci| Pai = pai).
18