Robust agents learn causal world models — part 1/8

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

Published as a conference paper at ICLR 2024
ROBUST AGENTS LEARN CAUSAL WORLD MODELS
Jonathan Richens∗
Google DeepMind
Tom Everitt
Google DeepMind
ABSTRACT
It has long been hypothesised that causal reasoning plays a fundamental role in
robust and general intelligence. However, it is not known if agents must learn
causal models in order to generalise to new domains, or if other inductive biases are
sufficient. We answer this question, showing that any agent capable of satisfying a
regret bound for a large set of distributional shifts must have learned an approximate
causal model of the data generating process, which converges to the true causal
model for optimal agents. We discuss the implications of this result for several
research areas including transfer learning and causal inference.
1 I NTRODUCTION
What capabilities are necessary for general intelligence (Legg & Hutter, 2007)? One candidate is
causal reasoning, which plays a foundational role in human cognition (Gopnik et al., 2007; Sloman &
Lagnado, 2015). It has even been argued that human-level AI is impossible without causal reasoning
(Pearl, 2018). However, recent years have seen the development of agents that do not explicitly learn
or reason on causal models, but nonetheless are capable of adapting to a wide range of environments
and tasks (Reed et al., 2022; Team et al., 2023; Brown et al., 2020).
This raises the question, do agents have to learn causal models in order to adapt to new domains, or
are other inductive biases sufficient? To answer this question, we have to be careful not to assume that
agents use causal assumptions a priori. For example, transportability theory determines what causal
knowledge is necessary for transfer learning when all assumptions on the data generating process
(inductive biases) can be expressed as constraints on causal structure (Bareinboim & Pearl, 2016).
However, deep learning algorithms can exploit a much larger set of inductive biases (Neyshabur et al.,
2014; Battaglia et al., 2018; Rahaman et al., 2019; Goyal & Bengio, 2022) which in many real-world
tasks may be sufficient to identify low regret policies without requiring causal knowledge.
The main result of this paper is to answer this question by showing that,
Any agent capable of adapting to a sufficiently large set of distributional shifts
must have learned a causal model of the data generating process.
Here, adapting to a distributional shift means learning a policy that satisfies a regret bound following
an intervention on the data generating process—for example, changing the distribution of features or
latent variables. It is known that a causal model of the data generating process can be used to identify
regret-bounded policies following a distributional shift (sufficiency), with more accurate models
allowing lower regret policies to be found. We prove the converse (necessity)—given regret-bounded
policies for a large set of distributional shifts, we can learn an approximate causal model of the data
generating process, with the approximation becoming exact for optimal policies. Hence, learning a
causal model of the data generating process is necessary for robust adaptation.
This has consequences for a number of fields and questions. For one, it implies that causal identi-
fication laws also constrain domain adaptation. For example, we show that adapting to covariate
and label shifts is only possible if the causal relations between features and labels can be identified
from the training data—a non-trivial causal discovery problem. This provides further theoretical
justification for causal representation learning (Schölkopf et al., 2021), showing that learning causal
representations is necessary for achieving strong robustness guarantees. Our result also implies that
we can learn causal models from adaptive agents. We demonstrate this by solving a causal discovery
∗jonrichens@deepmind.com
1

Published as a conference paper at ICLR 2024
task on synthetic data by observing the policy of a regret-bounded agent under distributional shifts.
More speculatively, our results suggest that causal models could play a role in emergent capabilities.
Agents trained to minimise a loss function across many domains are incentivized to learn a causal
world model, which could enable them to solve a much larger set of decision tasks they were not
explicitly trained on.
Outline of paper. In Section 2 we introduce concepts from causality and decision theory used to
derive our results. We present our main theoretical results in Section 3 and discuss their interpretation
in terms of adaptive agents, transfer learning and causal inference. In Section 4 we discuss limitations,
as well as implications for a number of fields and open questions. In section 5 we discuss related work
including transportability (Bareinboim & Pearl, 2016) and the causal hierarchy theorem (Bareinboim
et al., 2022), and recent empirical work on emergent world models. In Appendix B we describe
experiments applying our theoretical results to causal discovery problems.
2 P RELIMINARIES
2.1 C AUSAL MODELS
We use capital letters for random variables V , and lower case for their values v∈ dom(V ). For
simplicity, we assume each variable has a finite number of possible values,|dom(V )|<∞. Bold
face denotes sets of variablesV ={V1,...,V n}, and their valuesv∈ dom(V ) =×idom(Vi). A
probabilistic model specifies the joint distributionP (V ) over a set of variablesV . These models
can support associative queries, for example P (Y =y|X =x) forX,Y ⊆V . Interventions
describe external changes to the data generating process (and hence changing the joint distribution),
for example a hard intervention do(X =x) describes forcing the set of variablesX⊆V to take
valuex. This generates a new distribution P (V | do(X = x)) = P (Vx) whereVx refers to the
variablesV following this intervention. The power of causal models is that they specify not only
P (V ) but also the distribution ofV under all interventions, and hence these models can be used to
evaluate both associative and interventional queries e.g.P (Y =y| do(X =x)).
For the derivation of our results we focus on a specific class of causal models—causal Bayesian
networks (CBNs). There are several alternative models and formalisms that are studied in the
literature, including structural equation models (Pearl, 2009) and the Neyman-Rubin causal models
(Rubin, 2005), and results can be straightforwardly adapted to these.
Definition 1 (Bayesian networks). A Bayesian networkM = (G,P ) over a set of variablesV =
{V1,...,V n} is a joint probability distributionP (V ) that factors according to a directed acyclic
graph (DAG)G, i.e.P (V1,...,V n) =Qn
i=1P (Vi| PaVi), where PaVi are the parents ofVi inG.
A Bayesian network is causal if the graphG captures the causal relationships between the variables
or, formally, if the result of any intervention do(X =x) forX⊆V can be computed from the
truncated factorisation formula:
P (v| do(x)) =
(Q
i:vi̸∈xP (vi| pavi) ifv consistent withx
0 otherwise.
More generally, a soft interventionσvi =P′(Vi| Pa∗
i ) replaces the conditional probability distribu-
tion forVi with a new distributionP′(Vi| Pa∗
i ), possibly resulting in a new parent set Pa∗
i̸= Pai
as long as no cycles are introduced in the graph. We refer to σvi as a domain indicator (Correa &
Bareinboim, 2020) (it has also been called an environment index, Arjovsky et al., 2019). The updated
distribution is denotedP (v;σv′) =Q
i:vi∈v′P′(vi| pa∗
vi)Q
i:vi̸∈v′P (vi| pavi).
In general, soft interventions cannot be defined without knowledge of G. For example, the soft
interventionσY =P′(y|x) is incompatible with the causal structureY →X as it would induce a
causal cycle. As our results are concerned with learning causal models (and hence causal structure),
we focus our theoretical analysis on a subset of the soft interventions, local interventions, that are
compatible with all causal structures and so can be used without tacitly assuming knowledge ofG.
2

Published as a conference paper at ICLR 2024
Definition 2 (Local interventions). Local intervention σ onVi∈ V involves applying a map to
the states of Vi that is not conditional on any other endogenous variables, vi7→ f(vi). We use
the notationσ = do(Vi =f(vi)) (variableVi is assigned the statef(vi)). Formally, this is a soft
intervention onVi that transforms the conditional probability distribution as,
P (vi| pai;σ) =
X
v′
i:f(v′
i)=vi
P (v′
i| pai) (1)
Example: Hard interventions do(Vi =v′
i) are local interventions wheref(vi) is a constant function.
Example: Translations are local interventions as do(Vi = vi + k) = do(Vi = f(vi)) where
f(vi) =vi +k. Examples include changing the position of objects in RL environments (Shah et al.,
2022) and images (Engstrom et al., 2019).
Example: Logical NOT operationX7→¬X for BooleanX is a local intervention.
We also consider stochastic interventions, noting that mixtures of local interventions can also be
defined without knowledge ofG. For example, adding noise to a variableX =X +ϵ,ϵ∼N (0, 1),
is a soft intervention onX described by a mixture over local interventions (translations).
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).
2.2 D ECISION TASKS
Decision tasks involve a decision maker (agent) choosing a policy so as to optimise an objective
function (utility). To give a causal description of decision tasks we use the causal influence diagram
(CID) formalism (Howard & Matheson, 2005; Everitt et al., 2021), which extend a CBN of the
environment (chance) variables by introducing decision and utility nodes (see Figure 1 for examples).
For simplicity we focus on tasks involving a single decision and a single utility function.
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).
Single-decision single-utility CIDs can represent most decision tasks such as classification and
regression as they specify what decision should be made (d∈D), based on what information (paD),
with objective (E[U]). They can also describe some multi-decision tasks such as Markov decision
processes1. The utility is any real-valued function including standard loss and reward functions.
We assume that the environment is described by a set of random variablesC that interact via causal
mechanisms2, and whereC satisfies causal sufficiency (Pearl, 2009) (includes all common causes),
noting that such a choice ofC always exists. We refer to the CBN overC as the ‘true’ or ‘underlying’
CBN. Note we do not assume the agent has any knowledge of the underling CBN, nor do we assume
which variables inC are observed or unobserved by the agent, beyond that the agent can observe
PaD⊆C. We also assume knowledge of the utility functionU(PaU).
The conditional probability distribution for the decision nodeπ(d| paD) (the policy) is not a fixed
parameter of the model but is set by the agent so as to maximise its expected utility, which for a policy
π is Eπ[U] = E[U| do(D =π(paD))]. A policy π∗ is optimal if it maximises Eπ∗
[U]. Typically,
agents do not behave optimally and incur some regretδ, which is the decrease in expected utility
compared to an optimal policyδ := Eπ∗
[U]− Eπ[U].
To simplify our theoretical analysis, we focus on a widely studied class of decision tasks where the
agents decision does not causally influence the environment (e.g. Figure 1).
Assumption 1 (Unmediated decision task). DescD∩ AncU =∅.
In unmediated decision tasks, the agent is provided some (partial) observations of the environment
and chooses a policy, which is then evaluated using the utility function which is a function of the
1Note Markov decision processes can be formulated as a single-decision single-utility CID, by modelling the
choice of policy as a single decision and the cumulative discounted reward as a single utility variable.
2This assumption follows from Reichenbach (1956), and we discuss further in Appendix A.3
3