SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Google DeepMind London, UK pinen@deepmind.com
Stephanie C. Y. Chan
Google DeepMind London, UK scychan@deepmind.com
Ishita Dasgupta
Google DeepMind London, UK idg@deepmind.com
Andrew J. Nam
Stanford University
Stanford, CA
ajhnam@stanford.edu
Jane X. Wang
Google DeepMind London, UK wangjane@deepmind.com
What can be learned about causality and experimentation from passive data? This question is salient given recent successes of passively-trained language models in interactive domains such as tool use. Passive learning is inherently limited. However, we show that purely passive learning can in fact allow an agent to learn generalizable strategies for determining and using causal structures, as long as the agent can intervene at test time. We formally illustrate that, under certain assumptions, learning a strategy of first experimenting, then seeking goals, can allow generalization from passive learning in principle. We then show empirically that agents trained via imitation on expert data can indeed generalize at test time to infer and use causal links which are never present in the training data; these agents can also generalize experimentation strategies to novel variable sets never observed in training. We then show that strategies for causal intervention and exploitation can be generalized from passive data even in a more complex environment with high-dimensional observations, with the support of natural language explanations. Explanations can even allow passive learners to generalize out-of-distribution from otherwise perfectly-confounded training data. Finally, we show that language models, trained only on passive next-word prediction, can generalize causal intervention strategies from a few-shot prompt containing examples of experimentation, together with explanations and reasoning. These results highlight the surprising power of passive learning of active causal strategies, and may help to understand the behaviors and capabilities of language models.
Learning from passive observational data only allows learning correlational, not causal, structure. This observation is sometimes cited as a fundamental limitation of current machine learning research [60, 61, 39]. However, reinforcement learning (RL) agents can intervene on their environment, and are therefore not entirely limited. Indeed, various works have shown that RL agents can (meta-)learn to intervene on the environment to discover and exploit its causal structure [50, 14, 41, 16, 30].
However, these prior works leave open the possibility that an agent could passively learn a generalizable strategy for discovering and exploiting causal structure. While it is certainly necessary for an agent to intervene on the world at test time to discover causal structure, it may be possible for the agent to learn such a strategy from purely passive, offline data. Metaphorically, we ask "could an
agent learn to be a scientist by only reading a textbook detailing previous experiments and discoveries, and then successfully perform new experiments to discover and exploit new causal structures?"
The answer to this question is interesting for several reasons. From the perspective of causal learning, it is useful to investigate what limitations passive learning actually induces with respect to the knowledge an agentic system can later gain of the world's causal structure. Moreover, the answer is particularly salient at present because language models (LMs) — even when trained on purely passive imitation — show emergent capabilities such as adapting or reasoning in-context [71, 55, 77], or using tools to interact with the world [66]. These abilities are presumably driven by distributional properties of the passive data on which the LMs are trained [cf. 10, 64]. While LMs are often fine-tuned in active settings, e.g. using RL from Human Feedback [12], it is unclear whether this step is necessary to unlock causal abilities; the vast majority of learning occurs during the passive pretraining phase [e.g. 76]. Thus, understanding whether passive learning can lead to generalizable causal strategies can contribute to understanding the capabilities and limitations of language models.
In this work, we make the following contributions:
- We formally show that it is possible for an agent to learn a generalizable strategy for exploiting certain causal structures from passive data, as long it can intervene at test time.
- We demonstrate empirically that an agent trained via imitation (behavioral cloning) on expert data from a simple causal intervention task indeed can generalize at test time to infer and use causal links that are never present in the training data.
- We show that passively-trained agents can also generalize strategies for adaptive experimentation to experiment optimally in a novel situation.
- We show that strategies for causal intervention and exploitation can be learned from passive data even in a more complex environment with pixel observations and relational structure.
- We further show that natural language explanations support learning and generalizing causal strategies from passive data, and can enable generalization out-of-distribution, even from otherwise totally confounded training data.
- Finally, we show that language models, trained only on passive next-word prediction, can generalize causal intervention strategies from a few-shot prompt containing explanations.
Here, we attempt to briefly situate our contributions within the extensive prior literature on causal structure induction, language models, and online and offline RL agents.
A variety of works have considered the problem of causal graph structure induction from data [21, 54, 63]. While in a few cases the causal graph is identifiable from observational data alone [62], observations are generally not sufficient [e.g. 73], so most works also incorporate at least some interventional data. However, there is some flexibility in the interventional information required; for example, Faria et al. [18] show that a deep learning system can infer causal relations from data in which interventions are present, but latent (not directly observed). Perhaps most relevantly, Ke et al. [31] show that a transformer-based model can learn to induce causal structures from observational and interventional data. Our work moves beyond theirs in demonstrating that an agent can passively learn both to actively collect its own interventional data, and to use implicit inferences about causal structure to achieve goals — and by connecting these observations to recent progress in LMs.
Many perspectives on the limitations of LMs correspond to causal arguments. For example, the "octopus test" of Bender and Koller [4] is an argument about causal structure in the world that is latent in language. Other works have expressed skepticism about the possibility of LMs (or other models) learning causal strategies from passive data [e.g. 56, 60]. Even works that take more optimistic perspectives tend to focus on memorization of observed causal structure — e.g. Kiciman et al. [32] study LM knowledge about real-world causal structure, and note their potential utility as causal priors. Willig et al. [72] make some overlapping observations, but argue that these are simply examples of discovering correlations between causal structures in the real world and their descriptions in language, and that therefore LLMs "simply recite the causal knowledge embedded in the data" and thus "we cannot expect any sort of generalization in terms of causal prowess." However, some works have suggested that LMs can learn implicit world models from the training data [43, 44], suggesting that there is potential for more than purely superficial memorization. Correspondingly, we will show
that it is possible to learn generalizable causal strategies from passive data, especially when given explanations that link observations to causal structure (which is particularly relevant for LMs).
Work in offline RL has also studied the ability of agents to generalize from passive data [20, 38, 1], but has generally used more complex offline algorithms rather than pure imitation. Some work on pure imitation has considered the causal challenges for learning an optimal policy [65, 15, 74], but generally under a fixed causal structure. Other work has explored the generalization capabilities of agents trained from pure imitation when conditioned on information about the trajectories (such as returns) [67, 68, 3, 11]. Laskin et al. [42] show that passive offline imitation can learn a meta-RL algorithm that can improve an RL policy at test time. However, none of these works focused on tasks of causal experimentation and exploitation.
Instead, work on learning causal strategies has generally relied on online RL agents. Dasgupta et al. [14] showed that causal reasoning can emerge implicitly from meta-RL; other approaches learn to infer causal structures explicitly as part of an meta-RL process [50, 30]. Lampinen et al. [41] demonstrate that explanations can support meta-learning of causal strategies. Other works have proposed benchmarks for meta-learning causal experimentation [70]. However, all of these works relied on active RL. Kosoy et al. [36] recently considered learning of causal overhypotheses, and in addition to online RL agents evaluated behavioral cloning, but found it performed poorly; a decision transformer performed somewhat better, but only if forced to explore. Thus, past works leave open the question of whether agents can generalize causal strategies from passive data.
3 Passive learning can suffice
The key idea is that an agent that can intervene at test time could correctly infer and use the causal structure, as long as it has acquired an appropriate strategy during training. One way to enable this is to teach the agent this strategy in two stages: first, explore a system using an experimentation strategy that is guaranteed to identify the causal structure; second, learn an exploitation strategy to use this identified causal structure to achieve a goal. If the agent is successfully able to execute the experimentation strategy in the first part, it will necessarily have acquired the information required. Generalizing the exploitation strategy is directly amenable to passive learning – given that all the relevant information has been uncovered, there is no in-principle limitation on its success.
More formally, suppose an agent  $A$  receives observations  $o_{t}$  at time  $t$ , and is capable of taking actions  $a_{t}$  that intervene on a system. Assume that the system in question consists of a finite set of  $n$  random variables  $\{X_0,\dots,X_n\}$  that are connected to form a causal Directed Acyclic Graph (DAG; see e.g. Glymour et al. [22], Ch. 2). We will assume for simplicity that the variables are real-valued and that edges in the graph correspond to linear effects. We assume that independent variables (i.e., those without ancestors) are sampled from a normal distribution with mean 0 and a random variance, followed by a nonlinearity; nodes with ancestors are sampled from a normal distribution with mean corresponding to the linear effects of its inputs and again a random variance, and final nonlinearity. It is easy to adapt our argument in order to change many of these assumptions. The agent's observations correspond to current values of the variables and (at some time steps) a goal  $g\in \mathcal{G}$ , that is,  $o_{t}\in \mathbb{R}^{n}\times \mathcal{G}$ . The agent's actions correspond to interventions that set the value of a variable  $a_{t} = do(X_{i} = x)$  for some  $i\in \{1,\ldots ,n\}$  and  $x\in \mathbb{R}$ .
Now suppose we have a dataset of expert demonstrations on varied causal structures, in which the expert first intervenes on every variable to determine the causal structure; then uses that structure to achieve some goal. For example, the trajectories of the expert might have the following form:
$$
\underbrace {a _ {0} = d o \left(X _ {0} = x _ {0}\right) , \dots , a _ {n} = d o \left(X _ {n} = x _ {n}\right)} _ {\text {e x p e r i m e n t a t i o n : d e t e r m i n e c a u s a l s t r u c t u r e}} \underbrace {a _ {n + 1} = d o \left(X _ {i} = x _ {n + 1} ^ {g}\right) , \dots , a _ {n + k} = d o \left(X _ {j} = x _ {n + k} ^ {g}\right)} _ {\text {e x p l o i t a t i o n : t a k e o p t i m a l a c t i o n s t o a c h i e v e g a l o g} g \text {g i v e n c a u s a l s t r u c t u r e}}
$$
In this example, the first phase — the intervention strategy — is easy to learn, since it consists of a fixed sequence of actions.2 Thus, as long as the agent successfully learns the training data, it should reproduce this sequence on new causal structures. Once the agent has completed these interventions, it will have acquired all the information necessary to correctly infer the causal structure. With this information, the agent can in principle behave optimally in the goal-seeking remainder
$^{1}$ E.g., to accommodate interactions interevene on subsets up to the interaction order.
$^{2}$ N.B. we discuss this simple strategy because it is clear; it is not the most efficient [17].
(a) Train vs. eval. DAGs
(b) Rewards
(c) Action analysis
Figure 1: The causal DAG environment and experimental results. (a) The constraints imposed on the causal DAG structures when creating the training dataset, and when evaluating the agent on test structures. During training, D is not allowed to be an ancestor of E (even indirectly), though they may be correlated due to confounding variables. In evaluation environments, D is the most impactful ancestor of E (see text). (b) Rewards obtained by the agent when evaluated in interactive settings, as a percentage of the optimal rewards. In both evaluation settings, the agent still achieves close to optimal reward. (c) Analyzing the agent's behavior in more detail, by plotting the proportion of the agent's actions that match the optimal behavior, or baselines based on heuristics from the interventions or correlational statistics. The agent matches the optimal strategy significantly more closely than it matches the heuristic baselines. (Error bars are  $95\%$  -CIs; 5 seeds per condition.)
of the episode, i.e. implementing the exploitation strategy. Thus, the agent is not subject to any in-principle limitations. The two core tenets are whether a) the agent is able to follow the learned exploration strategy in new unobserved causal structures at test time, and b) the agent is able to generalize the exploitation strategy to account for the novel outcomes of those interventions. Below, we show empirically that an agent can indeed exhibit this kind of generalization from passive data.
Agent: We used a standard agent architecture: an input encoder that processes the observations, a memory that tracks and integrates information over time, and a policy that predicts actions. In analogy to the language model setting, we used a TransformerXL [13] for the agent memory, with a sufficiently long memory to cover the expected episode length. For the encoder in the causal DAG environments we concatenated the input observations and projected through a single linear layer. For the odd-one-out environments we used a CNN for visual observation encoding and an LSTM for language (following prior work on these environments [41]). For the agent policy, we used a single linear layer which receives as input the state output from the Transformer memory, and produces logits over the set of possible actions. For explanation prediction output (where applicable) we used a single-layer LSTM which received the same input as the policy.
Expert policies & datasets: We trained the agent on a dataset of trajectories generated by an expert policy, which behaves (near-)optimally. The specific implementation of the expert policy depends on the task, see below and Appendix A. We place our agent in the infinite-data regime (a common assumption in language model training [e.g., 27]), that is, the number of episodes in the dataset is larger than the number of episodes the agent trains on, so the agent does not repeat any episode.
Language model: For experiments with language models, we used the 70 billion parameter Chinchilla LM [27]; a decoder-only transformer that is trained purely on language modelling, without any downstream fine-tuning. When evaluating the LM's choice among multiple answers, we calculated the log-likelihood of each as a continuation of the current prompt, and selected the one with the highest likelihood. When generating from the model, to produce explanations and reasoning traces at test time, we used nucleus sampling [28] with  $p = 0.8$  and  $T = 1$ .
5.1 A simple causal DAG environment
Our first experiments take place in a simple causal DAG setting, inspired by the argument above and prior work on RL agents learning causal strategies [14]. The environment consists of a causal DAG over  $n = 5$  variables (A, B, C, D, E), with the causal graph, edge weights, and noise variance for each node resampled on each episode. Each episode consists of several segments: first, the agent is given  $n$  steps to explore the graph — enough to intervene on each variable and observe the consequences — and a second phase in which the agent is cued to try to maximize a particular variable. During this second phase, the agent is rewarded according to the value of the cued variable after its actions. In both phases, the agent observes its previous intervention and the value of all variables before and after. In the second phase, it also observes the goal cue. For details see Appx. A.1.
The training dataset is generated by an optimal expert policy that first explores (by choosing a random node to start with, then iterating through subsequent nodes in order, looping back to the first node after the last, and intervening on each to determine its effect). Then, the expert plans using the inferred causal structure to identify the intervention that will optimally achieve the given goal, and takes that action. The learning agent is trained to imitate this expert dataset via behavioral cloning — that is, by passively learning to maximize the probability of the expert actions.
The training data are generated from a subset of the possible causal DAGs, holding out some causal dependency links entirely for testing. We train the agent on expert data from DAGs in which node E is never descended (even indirectly) of node D. We then test the agent (interactively) on DAGs in which node E is descended from node D, with sufficiently large edge weights along paths from D to E (and sufficiently small from other ancestors) that node E is most strongly influenced via node D. We evaluate the agent's ability to determine this dependency through its experiments, and then, when cued to maximize node E, to optimally intervene. Specifically, we consider two test conditions: "eval. target," where node D is the optimal target to intervene on, and "eval. path," where node D is somewhere on the path from the optimal target to node E (but an ancestor of node D might be a better intervention). The "eval. target" condition is perhaps the most difficult, as in the training data the agent never sees any structure in which intervening on node D affects node E at all. Note that this generalization split is challenging because neither the environment observations nor the agent architecture encode permutation-invariance of the variables; their values are presented in a fixed order in distinct observation entries, and the agent uses distinct actions for intervening on each variable.
Thus, the question is whether the agent can learn to infer this novel causal dependency path from its experiments, and to use it to achieve the goal — i.e. can an agent generalize its exploitation strategy to novel causal structures? (Generalization of exploration strategy is investigated below in Sec 5.1.1). Indeed, in Fig. 1b we show that the agent rapidly learns to achieve high reward under interactive evaluation in all conditions (both training and evaluation structures).
However, simpler heuristic strategies could achieve some rewards. We therefore compare the agent's actions during the exploitation phase to the optimal strategy, and four baselines. The change and value baseline strategies use simpler heuristics, rather than relying on the full causal structure to achieve the goal of maximizing a given node. The value baseline simply repeats the action that during the exploration phase resulted in the highest value for the goal node. The change baseline is slightly more sophisticated; it chooses the intervention that changed the value of the goal node the most, and then flips its sign if the change was negative. We also consider two strategies that observe all the node values that the agents do, but consider only correlational statistics — either the total correlation between two nodes, or the correlation when controlling for all other nodes. Note that these baselines and the optimal expert policy are not mutually exclusive; these heuristics are effective because they often are match the expert. However, in Fig. 1c we show that the agent's actions match the expert policy significantly better than any baseline, across both training and evaluation. These observations suggest that the agent is learning to use more complex causal strategies, rather than simpler heuristics or correlational statistics.
In the supplement we perform a number of follow-up experiments, showing that at test time, it is necessary to intervene rather than just observe (Appx. B.2); how causal strategy generalization depends on the variety of causal DAG structures in training (Appx. B.3); and that agents can generalize causal strategies to DAGs with different functional structures, such as linear DAGs in training and nonlinear in test (Appx. B.4).
5.1.1 Adaptive experimentation strategies
It is limiting to assume that the optimal policy requires intervening on every causal factor; that does not scale to realistic scenarios. Scientists are guided by domain knowledge that suggests which variables may be relevant to a given phenomenon.
We therefore explored whether the passive agent can learn to adapt its experimentation procedure based on cues about which variables might matter.
To do so, we increased the variables to  $n = 10$ , but restricted so that only 5 are relevant in any given episode. The other 5 are independent from everything, and the goal node to maximize is always one of the 5 relevant variables. The agent is given a many-hot cue that marks which variables are relevant currently, and the expert policy which the agent learns to imitate only intervenes on the relevant variables. In the training data, we hold out some subsets of variables such that the agent never sees experiments on these subsets, as well as holding out certain ancestor structures as above. We test the agent in settings where it encounters both a held-out set of variables, and new relations among them.
Figure 2: When the agent is passively trained on expert actions that explore (intervene on) only a subset of the variables, it can generalize to explore correctly on held-out variable subsets at test time.
In Fig. 2, we plot over training the proportion of evaluation
5.2 The odd-one-out environments