An Illustrated Summary of "Robust Agents Learn Causal World Model"

post by Dalcy (Darcy) · 2024-12-14T15:02:44.828Z · LW · GW · 2 comments

Contents

  Introduction
  Basic Setup
    World
    Agent as a Policy Oracle
    "Robustness" as δ-optimality under interventions
    Assumptions
  Main Theorem
  High-level argument
  Discussion
    Policy oracle is not a good model of an agent.
    The Causal Good Regulator Theorem isn't a structural theorem.
  Conclusion
None
2 comments

This post was written during Alex Altair's agent foundations fellowship program [LW · GW], funded by LTFF. Thank you Alex Altair [LW · GW], Alfred Harwood [LW · GW], Daniel C [LW · GW] for feedback and comments.

Introduction

The selection theorems [LW · GW] agenda aims to prove statements of the following form: "agents selected under criteria $X$ has property $Y$," where $Y$ are things such as [LW · GW] world models, general purpose search, modularity, etc. We're going to focus on world models.

But what is the intuition that makes us expect to be able to prove such things in the first place? Why expect world models?

Because [LW · GW]: assuming the world is a Causal Bayesian Network with the agent's actions corresponding to the $D$ (decision) node, if its actions can robustly control the $U$ (utility) node despite various "perturbations" in the world, then intuitively it must have learned the causal structure of how $U$'s parents influence $U$ in order to take them into account in its actions.

And the same for the causal structure of how $U$'s parents' parents influence $U$'s parents ... and by induction, it must have further learned the causal structure of the entire world upstream of the utility variable.

This is the intuitive argument that the paper Robust Agents Learn Causal World Model by Jonathan Richens and Tom Everitt formalizes.

Informally, its main theorem can be translated as: if an agent responds to various environment interventions by prescribing policies that overall yield low regret, then it's possible to appropriately query the agent to reconstruct an implicit world model that matches up with the ground truth causal structure.

I will refer to this result as the "Causal Good Regulator Theorem". This sequence Thoughts on the Causal Good Regulator Theorem will do the following:

• [1] a precise statement of the theorem alongside a high-level argument of its proof and discussions of the paper's results in the context of the selection theorem agenda - basically a self-complete distillation of the paper. The rest of the posts in this sequence are optional. (this post)
• [2] detailed proof of the paper's theorems, presented in a way that focuses on how one could've discovered the paper's results by themselves, with an emphasis on intuition.
• [3] extension of the paper's theorem by showing how one of the main assumptions can be relaxed without much trouble.

Basic Setup

World

The world is a Causal Bayesian Network $G$ over the set of variables corresponding to the environment $C$, utility node $U$, and decision node $D$. The differences from a normal Causal Bayesian Network is that (1) $U$ is a deterministic function of its parents $U\left({Pa}_U\right)$, and (2) $P(D\mid {Pa}_D)$, the conditional probability distribution for $D$, is undetermined—it's something that our agent will select.

Agent as a Policy Oracle

In this paper's setup, the agent is treated as a mechanism that, informally, takes in an intervention to the world, and returns a policy i.e. it chooses $P(D\mid {Pa}_D)$.

Following this, we formally define an agent via a policy oracle ${\Pi }_{\Sigma }$ which is a function that takes in an intervention $\sigma \in \Sigma$ (where $\Sigma$ represents the set of all allowed interventions over $C$) and returns a policy ${\pi }_{\sigma }(D\mid {Pa}_D)$.

With the $P(D\mid {Pa}_D)$ determined under $\sigma$ as ${\pi }_{\sigma }$, all the conditional probability distributions of $G$ under $\sigma$ are determined, meaning we can e.g., calculate the expected utility, $E^{{\pi }_{\sigma }}\left[U\right]$.

This definition makes intuitive sense - if a human is told about changes in the environment (e.g., The environment has changed to "raining"), they would change their policy accordingly (e.g., greater change of taking the umbrella when going out).

But also, it is unnatural in the sense that it directly receives environmental intervention as input, unlike real agents which have to figure it out with its sensory organs that are also embedded in the world. This will be further discussed later in the post.

"Robustness" as $\delta$-optimality under interventions

By a "robust" agent, intuitively we mean an agent that can consistently maximize its utility despite its environment being subject to various interventions. We formalize this in terms of "regret bounds".

We say a policy oracle ${\Pi }_{\Sigma }$ has a regret $\delta$ if for all interventions allowed $\sigma \in \Sigma$, the ${\pi }_{\sigma }$ the oracle prescribes attains expected utility that is lower than the maximum expected utility attainable in $G_{\sigma }$ by at most $\delta$ (that is, $E^{{\pi }_{\sigma }}\left[U\right]\ge E^{{\pi }^{\ast }}\left[U\right]-\delta$), indicating a bound on suboptimality.

We denote a $\delta$-optimal policy oracle ${\Pi }_{\Sigma }^{\delta }$, and $0$-optimal one ${\Pi }_{\Sigma }^{\ast }$.

Now we have to choose some class of allowed interventions $\Sigma$. This is put in a collapsible section because it doesn't matter too much for the high-level discussion of the first post.

Choice of $\Sigma$

Now the question is the choice of $\Sigma$ - how broad of a set of perturbation do we want "robustness" to hold in? Ideally we would want it to be small while letting us prove interesting theorems - the broader this class is, the broader the set of environments our oracle is assumed to return a $\delta$-optimal policy, making the agent less realistic.

Recall hard interventions: $\sigma =do(C_i=\alpha )$ replaces $P(C_i\mid {Pa}_i)$ to a delta distribution $P(C_i=c_i\mid {Pa}_i={pa}_i;\sigma ):=1[c_i=\alpha ]$, and so the distribution factorizes to the following:

$P_{do(V=v^{\prime })}\left(c\right)=\{\begin{array}{cc}{\prod }_{i:C_i\notin V}P(C_i=c_i\mid {Pa}_i={pa}_i)& \text{if}c\text{consistent with}{v}^{\prime }\\ 0& \text{otherwise.}\end{array}$

Soft interventions instead change $P(C_i\mid {Pa}_i)$ more generally - not necessarily a delta function. It could even change the parent set ${Pa}_i$ to a different set ${Pa}_i^{\prime }$ as long as it doesn't introduce a cycle, of course.

The set of all soft interventions seem like a good choice of $\Sigma$, until you realize that the condition "as long as it doesn't introduce a cycle" assumes we already know the graph structure of $G$ - but that's what we want to discover!

So the paper considers a restricted form of soft intervention called "local interventions." On top of $P(C_i\mid {Pa}_i)$, they apply a function $f:\mathrm{Val}\left(C_i\right)\to \mathrm{Val}\left(C_i\right)$. Because the function only maps the values of $C_i$ locally, it does not change the fact that $C_i$ depends on ${Pa}_i$. $\sigma =do(C_i=f(c_i\left)\right)$ yields $P(C_i=c_i^{\prime }\mid {Pa}_i={pa}_i;\sigma ):={\sum }_{c_i:f\left(c_i\right)=c_i^{\prime }}P(c_i\mid {pa}_i)$.

Examples:

• $f\left(c_i\right)=\alpha$ makes $P(c_i^{\prime }\mid {\text{pa}}_i;\sigma )={\sum }_{c_i:\alpha =c_i^{\prime }}P(c_i\mid {\text{pa}}_i)=1[c_i=\alpha ]$, corresponding to the hard intervention $C_i=\alpha$.
• $f\left(c_i\right)=c_i+k$ makes $P(c_i^{\prime }\mid {\text{pa}}_i;\sigma )={\sum }_{c_i:c_i+k=c_i^{\prime }}P(c_i\mid {\text{pa}}_i)=P(c_i^{\prime }-k\mid {\text{pa}}_i)$, intuitively corresponding to e.g., shifting $C_i$ by $k$; maybe $C_i$ represents the location of some object in an RL world.

The paper extends this class further, considering "mixtures of local interventions," where we denote a mixed intervention ${\sigma }^{\ast }=\sum p_i{\sigma }_i$ for $\sum p_i=1$, which denotes randomly performing a local intervention ${\sigma }_i$ with probability $p_i$.

Examples:

• Suppose $C_i$ takes $k=\mathrm{Val}\left(C_i\right)$ values $v_1,\dots ,v_k$. Then consider a local intervention $f_j\left(c_i\right)=v_j$, ${\sigma }_j=\mathrm{do}(C_i=f_j(c_i\left)\right)$, and a mixture ${\sigma }^{\ast }=\sum p_j{\sigma }_j$ for some $p_j$ that sums to $1$. This lets us set $C_i$ to literally any arbitrary distribution we want!
  • This now makes it seem like a very broad class of interventions thus bad. But that's what we have for now.

Aside from these "local" interventions that only depend on the value of a node we're intervening on, the paper extends the intervention class by adding a specific form of structural intervention on the decision node $D$.

• Specifically, they consider masking inputs to $D$ such that it only depends on ${P{a}^{\prime }}_D\subseteq {Pa}_D$. This is somewhat reasonable (despite being structural) if we assume prior knowledge of which nodes influence $D$, e.g., perhaps by design, and control over what the decision node reads as the values of these variables.

Note: The paper mentions that this can be implemented by local interventions, but I don't think so since this is a structural intervention that doesn't just depend on values of $D$. A set of hard interventions that set ${Pa}_D\setminus {P{a}^{\prime }}_D$ to a constant wouldn't work, because then we're not just masking inputs to D, but also masking inputs to other descendants of ${Pa}_D\setminus {P{a}^{\prime }}_D$.

Assumptions

1) Unmediated Decision Task states that ${Des}_D\cap {Anc}_U=\varnothing$. This is pretty major.

2) Domain dependence states that there exists distributions over the chance variables $P\left(C\right)$ and $P^{\prime }\left(C\right)$ (compatible with $M$) such that $\arg {max}_{\pi }{E}_P^{\pi }\left[U\right]\ne \arg {max}_{\pi }{E}_{P^{\prime }}^{\pi }\left[U\right]$.

• This is very reasonable. If domain dependence does not hold, then the optimal policy is just a constant function.

These together imply:

1. There does not exist a decision $d\in \mathrm{dom}\left(D\right)$ that is optimal, i.e. $\arg {max}_{d}U(d,x)$ across all $x\in X={Pa}_U\setminus \left\{D\right\}$.
2. $D\in {Pa}_U$, i.e. there can't be any intermediate nodes between $D$ and $U$, and all causal effects from $D$ to $U$ must be direct.

Main Theorem

With this basic setup explained, here is the main theorem of the paper (Theorem 1):

For almost all worlds $(G,P)$ satisfying assumption 1 and 2, we can identify the directed acyclic graph $G$ and the joint distribution $P$ over all variables upstream of $U$, given that we have access to a $0$-optimal policy oracle.

and a second theorem, which is the approximate case (Theorem 2):

For almost all worlds $(G,P)$ satisfying assumption 1 and 2 and some other new assumptions (explained below), we can identify the directed acyclic graph $G$ and the joint distribution $P$ over some subset of variables upstream of $U$, and the quality of estimation for each of the conditional distributions scale linearly with $\delta$.

High-level argument

The proof is basically a formalization of the argument given in the introduction, but here we state it in the framework so far introduced.

The world is a Causal Bayesian Network, also containing two nodes each corresponding to an agent's decision node $D$ and utility node $U$.

Suppose we're given an oracle that takes in an intervention $\sigma$ as input and returns a policy (conditional distribution of the decision node given its parents) that attains maximum utility under that intervention in the environment. This oracle is an operationalizes the notion of a "robust agent."

The question of "How robust?" is determined by the class of interventions to be considered. The broader this class is, the broader the set of environments our oracle is assumed to return an optimal policy, hence the agent is less realistic.

So how do we use the oracle?

Suppose you have two interventions $\sigma$ and ${\sigma }^{\prime }$, and you "interpolate"^[1] between them by intervening under $\sigma$ with a probability of $q$, and intervening under ${\sigma }^{\prime }$ with a probability of $1-q$. Denote such an intervention $\tilde{\sigma }\left(q\right)=q\sigma +(1-q){\sigma }^{\prime }$.

If some decision, say, $d_4$ is the optimal decision returned by the oracle under $\sigma$, and $d_1$ is the optimal decision returned under ${\sigma }^{\prime }$, then as you gradually change $q$ from $0$ to $1$, the optimal decision will switch from $d_1$ to other decisions ($d_2,d_3,\dots$) and eventually to $d_4$.

Call the value of $q$ at which the decision eventually switches over to $d_4$, $q_{\text{crit}}$.

The critical insight is that $q_{\text{crit}}$ can be estimated by querying the oracle with various values of q, and it also can expressed as an equation involving a number of terms corresponding to the conditional probability distributions of the Bayes Net (call it $P(C_i\mid {Pa}_i)$ where $C_i$ is some node).

By cleverly setting $\sigma$ and equating the estimate and the expression, all of the $P(C_i\mid {Pa}_i)$ terms for $C_i$ upstream of $U$ can be solved. This is where we involve the earlier intuition of "induction":

• Assume some directed path $C_k\to \dots \to C_1$ where $C_1\in {Pa}_U$
• Step $1$, $\sigma$ is set such that the expression for $q_{\text{crit}}$ only contains $P(C_1\mid {Pa}_1)$, so it can be immediately solved by comparing with the estimate.
• Step k, assume we know $P(C_i\mid {Pa}_i)$ for $i\in 1,\dots ,k-1$, and set $\sigma$ such that the expression contains $P(C_i\mid {Pa}_i)$ for all $i\in 1,\dots ,k-1,k$. The only unknown here is $P(C_k\mid {Pa}_k)$.
• Thus, by induction, we should be able to infer the conditional probability distribution table of $P(C_i\mid {Pa}_i)$ for every variable upstream of $U$.

Furthermore, all of this can be relaxed to approximately-optimal policy oracles.

The claim for the approximate case is that given an oracle that is "imperfect by $\delta$ amount" it is possible to identify a subgraph of the graph corresponding to variables upstream of $U$, and that the conditional distribution it estimates $\hat{P}\left(C_i|{Pa}_i\right)$ will differ from the true conditional distribution $P\left(C_i|{Pa}_i\right)$ by an amount that scales linearly with $\delta$ for small values of $\delta$.

Discussion

When viewed in terms of the selection theorem agenda, I think the results of this paper signify real advancement - (1) in its attempt at formalizing exactly the intuitive argument often used to argue for the existence of world models, (2) incorporation of causal information, (3) providing an explicit algorithm for world model discovery, and (4) relaxing its proof to the approximately optimal oracle case - many of which are important advancements on their own.

However, this work still leaves many room for improvement.

Policy oracle is not a good model of an agent.

Real agents respond to changes in the world through its sensory organs that are embedded in the world. However, the use of policy oracles imply that the agent can directly perceive the environmental intervention, where the policy oracle is sort of a sensory organ that lives outside of the causal graph of the world.

That's why even in cases like the above, where D has no input, the policy oracle can produce policies sensitive to environmental interventions. Loosely speaking, the agent's sensory nodes are only partially embedded into the environment (D's connection to its parents), and the rest (its ability to sense interventions) are outside of the causal graph of the environment.

The Causal Good Regulator Theorem isn't a structural theorem.

I think the term "world model" in the paper's title is misleading.

The existence of a "world model" is an inherently structural claim. It is a question about whether a mind's cognition explicitly uses/queries a modular subsystem that "abstracts away" the real world in some sense.

But the theorems of this paper are purely about behavioral properties - namely that some algorithm, given access to the policy oracle and the set of variables in $G$, can query the policy oracle in a particular way to reconstruct $G$ accurately. This says nothing about whether this reconstructed $G$ is actually used internally within the agent's cognition!

• If we take the output of this algorithm and view it literally as a "world model" of the agent, then that implies the agent's "world model" is perfectly isomorphic to a subset of the environment's causal model. That would be very wrong. It shouldn't be that the "world model" recovered by the algorithm should depend on choice of variables used in $G$, the world model should have its own ontology!

This is akin to how e.g., the VNM representation theorem is not structural: While it is possible to reconstruct a utility function from the betting behavior of agents satisfying certain axioms, that does not imply e.g., that the agent internally represents this utility function and argmaxes it.

I think this paper's result is better understood as a way to derive an implicit behavioral causal world model of an agent over a given set of variables, in the sense of answering the question: "given that I represent the world in terms of these variables $C$, what causal relationship among these variables do the agent believe in?"

This itself is a very cool result!! The algorithm is inferring the agent's implied belief about causality with respect to any choice of variable ontology (i.e. choice of $C$ used to represent the world).

For example, I can literally imagine ...

• ... having $C_1$ to consist of variables like "sodium intake, blood pressure, arterial stiffness, ... vascular inflammation", and $C_2$ to consist of variables like "dietary habits, cardiovascular risk, life expectancy"
• ... then running the proof's algorithm on a human/LLM by treating them as a policy oracle (where $\sigma$ can be formulated as a text input, like providing them the text "your dietary habits are set to X ... ")
• ... and being able to infer the human/LLM's implied causal beliefs on these two different ontologies about healthcare issues!

But again, it would be misleading to call these $C$ alongside their inferred causal relationship the "world model" of the human/LLM.

Conclusion

The paper Robust Agents Learn Causal World Model signifies a real advancement in the selection theorem agenda, proving that it is possible to derive an implicit behavioral causal world model from policy oracles (agents) with low regret by appropriately querying them. But it has a lot of room for improvement, especially in making its claims more "structural."

I believe this post was a self-contained explanation of the paper along with my take on it.

However, my primary motivation for reading this paper was to more rigorously understand the proof of the theorem, with the goal of identifying general proof strategies that might be applied to proving future selection theorems. This is especially relevant because this paper appears to be the first that proves a selection theorem of substantial content while being rooted in the language of causality, and I believe causality will play a critical role in future selection theorems.

For this, wait for the next post in the sequence soon to be published.

2 comments

Comments sorted by top scores.