Passive learning of active causal strategies in agents and language models
Authors: Andrew Kyle Lampinen
Stephanie C.Y. Chan
Ishita Dasgupta
Andrew Joohun Nam
Jane X Wang
Venue: NeurIPS 2023 poster (2023)
Keywords: passive; causal; offline; agency; language models
TL;DR: We show formally and empirically (through experiments with BC agents and language models) that it is possible to learn generalize strategies for causal experimentation and goal seeking from purely passive imitation learning.

Abstract: 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 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 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 explanations and reasoning. These results highlight the surprising power of passive learning of active causal strategies, and have implications for understanding the behaviors and capabilities of language models.

Passive learning of active causal strategies in agents and language models
Andrew Lampinen Google DeepMind London, UK lampinenddeepminnd com
Stephanie C. Y. Chan Google DeepMind London, UK
Ishita Dasgupta Google DeepMind London, UK
scychanddeepmind.cmm com
idgådeepmind.com
Andrew J. Nam
Jane X. Wang
Stanford University Stanford, CA a hnam@st taaeenstaaafoo ord. edu
Google DeepMind London, UK
wangj aangjann@deepmin.. com
Abstract
What can be learned about causality and experimentation from passive data? This question is salient given recent successes of mussively-trained language models in interactive domains such as tool use. Passive learning 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 test time. We formally illustrate that, under certain assumptions, learning 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 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 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-confoundee training data. Finally, we show that language models, trained only on passive next-word prediction, can generalize causal inter- vention strategies from few-shot prompt containing examples of experimentation, together with explanations and reasoning. These results highlight the surprising power of passive learning active causal strategies, and may help understand the behaviors and capabilities of language models.
present in the training data; these agents
1 Introduction
Learning from passive observational data only allows learning correlational, not causal, structure. This observation is sometimes cited fundamental limitation of current machine learning research
(60) [1]39]] However, reinforcement learning (RL) agents can intervene on their environment, and
are therefore entirely limited. Indeed, various works have shown that RL agents can learn
to intervene on the environment to discover and exploit its causal structure [50] 14
[6] (3).
However, these prior works leave open the possibility that agent could passively learn generaliz- able strategy for discovering and exploiting causal structure. While it is certainly necessary for an agent to intervene the world 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
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

agent learn to be 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 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 77 35 777, 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. T0, 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. [6). 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 test time.
We demonstrate empirically that an agent trained via imitation (behavioral cloning) on expert data from simple causal intervention task indeed generalize test time to infer and use causal links that are never present in the training data.
We show that passively-traine agents can also generalize strategies for adaptive experimen- tation 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, show that language models, trained only on passive next-word prediction, can generalize causal intervention strategies from a few-shot prompt containing explanations.
2
Related work
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 While in a few cases the causal graph is identifiable from observational data alone [62]. observations are generally not sufficient [c.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 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. model can learn induce causal structures from observational and interventional data. Our work moves beyond theirs in demonstrating that an agent can passively learn both 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.
21 63
3 show that a transformer-basee
Many perspectives on the limitations of LMs correspond causal arguments. For example, the "octopus test" of Bender and Koller 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. Kıeıman et al. 33 study LM knowledge about real- real-world causal structure, and note their potential utility as causal priors. Willig al. [22] 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
2

---

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 RI. has also studied the ability of agents to generalize from passive data 20
38
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 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) 66 68 M. Laskin et al. 4 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.
15 74 but
Instead, work on learning causal strategies has generally relied on online RI. agents. Dasgupta et al. [4] 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 appropriate strategy during training. One way enable this is to teach the agent this strategy in two stages: first, explore a system using an experimentation strategy that is guaranteed 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 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 no in priincicple limitation on its success.
More formally, suppose an agent A receives observations or at time I, and is capable of taking actions a that intervene on system. Assume that the system in question consists of finite set of n random variables (Xo, -.--. X, 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 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 's observations correspond g E G. that is, o, E R x G. The agent's actions correspond to interventions that set the value of variable a; do (X;- 2)) for some E (1,...,,)) n) and x E TR.
current values of the variables and (at some time steps) goal
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:
ag- = - 20) On = do II) 0x+1 = do X; x xx,1= : ...., an+k = do Xx
experimentation: determine causal structure
exploitation: take optimal actions to achieve goal given causal structure
In this example, the first phase the intervention strategy is easy to learn, since it consists of a fixed sequence of actions Thus, as long as the agent successfully learns the training data, it should reproduce this sequence 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-sceking remainder
..g., to accommodate interactions interevene on subsets up to the interaction order. N.B. we discuss this simple strategy because it is clear; it is not the most efficient [17
3

---

Train dola:
100
Optimal
Agant match (%)
Rewards (% of offaptmaa
75
75
Intervention
50
50
Value
Fraluation:
Correlation (lotal)
25
path
Conrelation (Fartiel)
Evall target
0.5
1.0
target
Updates
145
Condition
(a) Train eval. DAGs
(b) Rewards
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, is not allowed to be an ancestor of E (even indirectly), though they may be correlated due confounding variables. In evaluation environments, D the impactful ancestor of E (see text). (b) Rewards obtained by the agent when evaluated in interactive settings, as percentage of the optimal rewards. In both evaluation settings, the agent still achieves close to optimal reward. cc) Analyzing the agent's behavior in more detail, by plotting the proportion of the agent's 's actions that match the optimal behavior, or baselines based on heuristies 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% 95%-Cho -Is,; seeds per condition.)
of the episode, i.c. implementing the exploitation strategy. Thus, the agent is not subject to any in-principle limitations. The two tenets whether a) the agent is able follow the learned exploration strategy in new unobserved causal structures at test time, and b) the agent is able to generalize the exploitation strategy 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.
4 Methods
We used standard agent architecture: 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, used TransformerXL [33 for the agent memory, with sufficiently long memory to cover the expected episode length. For the encoder in the causal DAG environments concatenated the input observations and projected through 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 these environments (41)) For the agent policy, used 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) used single-layer LSTM which received the same input as the policy.
Agent:
Expert policies & datasets: We trained the agent 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 agent in the infinite-date regime (a common assumption language model training 27), that is, the number of episodes in the dataset is larger than the number of episodes the agent trains
the agent does not repeat any episode.
Language model: For experiments with language models, we used the 70 billion parameter Chin- chilla LM [27]: decoder-only transformer that is trained purely language modelling, without any downstream fine-tuning. When evaluating the LM's choice among multiple answers, we calculated the log- liglliieldoood of each as 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 (88 with p 0.8 and T 1.

---

5 Experiments
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 [4]. The environment consists of 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 steps to explore the graph - enough to intervene and a second phase 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
each variable and observe the consequences
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 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 is most strongly influenced via node D. We evaluate the agent's '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 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 permutation-invaariace 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. Indeed, in Fig. Tb 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 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 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. Ic we show that the agent's '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 up experiments, showing that at test time, it is necessary to intervene rather than just observe (Appx. [.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.4p.
5

---

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 do- main knowledge that suggests which variables may be relevant to given phenomenon.
We therefore explored whether the passive agent can learn
to adapt its experimentation procedure based on cues about which variables might matter.
100
Exploration correct (%)
To do so. we increased the variables to = 10, but re- stricted that only are relevant in any given episode. The other 5 are independent from everything, and the goal node maximize is always one of the 5 relevant variables. The agent is given many-hot cue that marks which vari- ables are relevant currently, and the expert policy which
75
50
Train
path
Eval. target
the agent learns to imitate only intervenes variables. 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 above. We test the agent in settings where encounters both held-out set variables, and new relations among them.
the relevant
8.0
0.5
1.0
Updates
1e5
When the agent is passively trained on expert actions that explore (in- tervene on) only a subset of the variables, it can generalize to explore correctly on held-out variable subsets at test time.
Figure 2:
In Fig. 2 we plot over training the proportion of evaluation episodes in which the agent correctly uses an optimal exploration strategy; that is, in which it intervenes on each relevant variable once, without intervening on any irrelevant variable. The agent rapidly learns to generalize the experimentation strategy from the data to correctly experiment when cued with sets of variables never encountered in training. It then goes on to successfully use results of these experiments to achieve the goal, and again aligns with the optimal causal policy more closely than the heuristic baselines (Appx. B.5).
5.2 The odd-one-out environments
Experiment trials
100
100
Success rate (16)
Success rate (96)
trancimn
Exploit trial
25
25
Chance
Train combinations Eval. combinations
Train (0-2 dimensions hard!
Eval. (3 dimensions hard!
Chocu cornet sbjecs
2
0.2
0.4
0.6
0.8
1.0
Updates
165
Updates
(a) Environment & task structure
(b) Feature set generalization
(c) Hard dimension generalization
Figure 3: The odd-one-out meta-learning environments, where the goal is to pick up an object that is unique along the "correct" dimension, which cannot be directly observed. (a) An example of the observations and task structure. At top is an example experimentation trial. The agent is the white square (the view is egocentric, i.e. the agent always in the center); it transforms unique shape, then it will "pick up" that object to check whether shape is the relevant dimension for this episode. After three experimentation trials, the agent proceeds to exploit trial (bottom), where it can no longer transform objects, but where each object is unique along different dimension; it must pick up an object according the latent causal dimension discovered through experimentation. In training, the environment provides explanations of the objects and the results of experiments, which support learning.
object to
(b) Generalization novel feature combinations. (c) Generalization from training data where 0-2 dimensions have a hard initialization (see text) to the case when all 3 dimensions do. (Panel (a) adapted with permission from Lampinen et
3 seeds per condition.)
The above experiments used simple environments, where the agent directly observes the variables and the outcomes of its interventions. Here, extend to more complex set of environments involving

---

causal structure that can only be understood by interpreting relations among objects in high- dimensional pixel observations. Specifically, we experiment on the causal intervention//met-lliin versions of the odd-one-out environments [41]. See Fig. 3a for some example observations.
In these environments, objects vary along multiple dimensions (color, shape, and texture). In each episode there is one dimension that is "correct" - this "correct" dimension is latent causal factor that the agent cannot directly observe. The agent is rewarded for picking up an object that is unique (the odd one out) along the "correct" feature dimension. The agent is first given a series of trials in which to experiment and determine which dimension is currently "correct." Specifically, the agent is given a magic wand that allows it to transform features of objects, thus allowing it to make an object unique along one feature dimension. can then pick up this object, and can infer from whether it receives reward whether the dimension it tried is correct. These experiments allow the agent to succeed on final exploit trial, where it no longer has access to magic wand, but instead faced with multiple objects that are each unique along different dimension. To achieve reward, the agent must use the results of its earlier experiments to choose the object that unique in the correct way.
We evaluated two types of generalization. First, we simply split possible feature combinations into train and test sets. We evaluate whether the agent can learn, from passive data, an experimentation and exploitation strategy that generalizes to novel combinations. In Fig. 3b we show that the agent is able to learn these tasks well, and to generalize equally well to the evaluation conditions. Note that the tasks are designed such that simpler strategies - such as relying on feature values - will not yield above chance performance.
We then evaluated a more challenging generalization condition. The original paper 41 considered harder object initializations that make the experiments more difficult, by forcing the agent to transform one object, and then choose a different object that is made unique by that transformation. We create a more challenging relational generalization split by training the agent in settings where 0-2 of the feature dimensions (randomly sampled in each training episode in the dataset) have this harder initialization. We then evaluate the agent's ability to generalize evaluation in harder setting where all 3 feature dimensions have this difficult initialization. This condition tests the agent's ' generalization to settings more challenging than those observed in the data. In Fig. Be able to learn and generalize these more difficult tasks.
we show that the agent is
Explanations support learning: In these experiments, we followed the method of Lampinen et al. [1] the agent is trained to predict natural language explanations (provided by the environment during training) as auxiliary loss. For example, when the agent performs a successful experiment that identifies the correct dimension as shape, during training this will be accompanied by an explanation like "correct, the latent feature is shape, and this object is uniquely square." In Supp. Fig. [2 we show that these explanations are beneficial even when learning from expert data; without predicting explanations, the agents do not learn or generalize these tasks as readily. Furthermore, explanations can enable other kinds of generalization from passive learning, as we explore in the next section.
5.2.1 Explanations can deconfound for passive learners too
The original odd-one-out work also explored a setting of learning from perfectly confounded experiences, in which a rewarding object is unique along three feature dimensions, so it is unclear which feature(s) are relevant. The authors showed that RL agents can generalize out of distribution to follow any particular dimension in de- confounded test (where different objects are unique along each dimension). if they are trained to predict language explanations that focus on that particular feature dimen- sion. Although these experiments do not directly involve causal intervention, they are relevant to the broad question of what can be learned from confounded data- and are particularly relevant to LMS, because explanations are a common feature of the human language from which these models learn. However, the prior work only experimented with active RL. agents. Can passive learners also generalize out-of-distributin from predicting explanations?
Choice in deconfounded eval. (%))
100
75
500
25
Explain color, choose color
shape
texture
0.0
0.5
1.0
1.5
Updates
1e5
Figure 4: Explanations can deconfound for passive learners trained con- founded data- - the agents generalize OOD in accordance with the explana- tions they were trained to predict.
7

---

We reproduced these deconfounding experiments in our passive learning setup, using the same expert as above. We find that passive learners also generalize consistently OOD according to the explanations they predicted during training (Fig. 1. Passive learners that were trained to predict explanations about shape would tend to choose the agents trained to predict explanations about texture or shape would analogously generalize according to those features. Thus, natural language explanations encountered during passive training can shape the way a system generalizes out-of oul-of-distribution at test time.
object the deconfounded evaluation, and
5.3 Pretrained language models
Finally, we tested whether pretrained language models could generalize causal strategies from a few-shot prompt. To do so, we adapt the odd-one-out intervention experiments from above to language setting, and evaluate whether large language model [70B parameter Chinchilla, from 27] generalize causal strategy in this environment from few-shot prompt. Specifically, we give the LM natural language observations the objects in a scene, and pose choices between actions described in natural language. For the experimentation trials, give choice of which object to intervene on and which feature to transform. In the final trial, we simply give a choice among the objects, each of which is unique along one dimension.
We give the LM four example shots, in all of which the latent "correct" dimension is sam- pled from only two of the three possible dimensions (color, shape, and texture), and eval- uate the ability of the LM to generalize to an episode where the relevant dimension is the third, held-out one. Note that this is a tigated with the agents - generalizing a causal relation to a dimension for which it was never observed in training, from only handful of examples of other dimensions. We also experiment with holding out each of the dimensions, to ensure that the results are robust.
difficult generalization setting than was inves-
LM performance can be increased by incor-
porating factors like reasoning traces/chain- of-thought or explanations (55] 71 69 [77]
In shots
Held-out
(9)
40 35) in the prompt, and these strategies particularly useful in challenging rea- soning cases. We therefore include both reasoning traces and explanations in our prompt examples; at test time we allow the LM generate its own explanations and reasoning traces at the appropriate places in the episode. We also choose the example shots for the prompt by optimizing for per- formance on new samples of tasks involv- ing the two trained dimensions [45] 40) For more details, including an example episode, see Appx. A.2.
50
Evaluation
25
Explonations Reasoning
None
Color
Shape
Texture
Color
Shape
Texture
Hold-out dimension
Hold-out dimension
Figure LMs can generalize the odd-one-out tasks to new dimensions that were not correct in the example shots (right panel), but only if explanations/reasoning are included in the prompt. On new samples of the dimensions included the examples (left), LMs can achieve some performance with or without explanations. (3 prompts per condition; errorbars are 95%
In Fig. we show the results. Given a strong four-shot prompt, the language model can generalize correctly to new stances of the dimensions that were relevant in the prompt; critically, it can also generalize to new instances the held-out dimension in each case (and even to a somewhat harder test with an added object: Appendix B.9) However, explanations are critical for this generalization. In Appendix B.7 we show that either explanations or chain-of-thee of-thought reasoning suffice to provide some benefit, but results are most reliable when both are provided together. Thus, the language model is able to build on its strong prior knowledge to generalize, but only when supported by some form of explanation.
Discussion
What can be learned from passive data about strategies for experimenting to discover and exploit causal structure? Our work illustrates that a non-trivial amount can, at least when the passive data contains the results of experiments. From passive imitation, agents can learn strategies for experimentation that can be generalized to correctly experiment novel causal structures, and learn

---

strategies for exploiting the results of their own experiments to achieve goals under those novel structures, Agents can even generalize in complex environments where they receive high-dimensional observations; likewise, passively-trained language models can generalize from a high-quality few-shot prompt. Furthermore, our results highlight the role that natural language explanations can play in drawing out latent causal structure to enable more generalizable learning. Explanations can even enable passive learners to generalize nearly perfectly from otherwise perfectly-confounde data. Passive learning can be surprisingly powerful under the right circumstances.
These observations provide perspective on the kinds of understanding of causality and experiments that can be acquired by language models. For example, consider the following, from Pearl [60]:
"This hierarchy, and the formal restrictions it entails, explains why statistics-based machine learning systems are prevented from reasoning about actions, experiments and explanations. [...] the hicrarchy denegrades [sic] the impressive achievements
of deep learning to the level of Association [...] Unfortunately. the theoretical barriers that separate the three layers in the hierarchy tell us that the [...] objective function does not matter. As long as our system optimizes some property of the observed data, however noble or sophisticated, while making no reference the world outside the data, we are back to level-1 of the hierarchy with all the limitations that this level entails."
Naively, it might seem that these arguments doom language models to never acquire real knowledge of causality [cf. 72). Our results provide useful formal and empirical nuance. While there are certainly restrictions to learning from passive data (see below for further discussion), it is possible to learn generalizable strategies for causal experimentation and intervention from passive data alone at least if the data include examples of an expert intervening, and perhaps explanations.
Other work has engaged with these issues. In particular, Ortega et al. [56] provide a theoretical account of the effects of confounding on passively trained sequence models, and propose a counterfactual training method to enable an adaptive system to learn effective intervention strategies in the presence of confounders. However, this method relies on having active access to an expert teacher rather than simply offline data; the authors note that "in general, we cannot meta-train an agent only from expert demonstrations to imitate the expert at deployment time" and "[t]o address the problem, is necessary to incorporate assumptions (e.g. a set of hypotheses) about the confounding variable, but currently is unclear how to incorporate this seamlessly into a meta-learning scheme." Our work complements theirs, by illustrating some of the cases in which an agent can successfully generalize causal experimentation and geal-seeking strategies from passive data alone- - even the presence of confounders and by highlighting the role that natural language explanation can play in incorporating expert knowledge to allow active generalization from passive learning.
Explanations & causal structure: Explanations were important for effective learning or general- ization the odd-one-out tasks; explanations can even allow generalizing from otherwise totally- confounded data. These results fit with prior work showing that explanations and reasoning traces help agents 41 533 and LMs 52 24 40 29]. More broadly, explanations are key for human learning because they explicitly highlight generalizable, causal structures [47 46 33]. For example, explanations can constrain the hypothesis space of causal graphs that a system needs to consider [cf. 51. Alternatively, language]] could serve as the "reference to the world outside the data" that Pearl 60 sought in the quote above.
El may obviate the need for explicit causal graphs. That is, explanations
Language models: Our experiments have obvious implications for the causal abilities of language models. Prior works have suggested that LMs can only parrot causal structures observed in training [7]. However, our experiments illustrate that this should not be assumed; the implications of passive learning for causal strategies can be subtle. Passive data is not necessarily strictly observational. In particular, observing other's 'ss interventions, especially with explanations, can enable generalization.
LM training data will contain many examples describing interventions and outcomes together with explanations thereof- - whether scientific articles, or just descriptions of debugging process. Just as the data a scientific paper do not become observational if we read the paper rather than writing it, the data on the internet do not become observational just because an LM consumes it passively. Thus, sequences of interventions, outcomes, and explanations may in some cases be sufficient to allow an
Or other forms of abstraction.
9

---

LM to generalize to proposing experiments and using the results to draw appropriate conclusions, at least if given an appropriate prompt.
Tool use & agency: These findings are of particular interest given the recent focus on language models using tools [e.g- (66. Tool use is clear case of active interaction in the world, which one might assume requires some knowledge of causality. Our results may help explain the successes of these methods with models trained purely from passive supervised learning (though note that many applications of LMs are also actively fine-tuned, which may also contribute, see below). A related question is whether LMs, despite passive learning, can acquire an ability to act as goal-directed agents. Andreas (2] notes that effective language modeling requires modeling the goals of the communicator that produced the language, and thus suggests that through this process LMs learn at least model some aspects of agency. Our results may provide some support for these arguments, as generalizing goal-directed behavior to new causal structures is key aspect of agency.
Experimental limitations: Our experiments are limited in a number of important ways. First, the causal structures we used are small, finite acyclic graphs. While this limitation (i.e., that the causal graph is a DAG over a handful of variables) is assumed by much of the prior causal literature, many real world causal structures are much more complex. For example, they may include cycles, at least over time 4] Changing these assumptions changes the nature of the causal discovery problem 6 19]; thus is an open question whether agents could similarly learn causal strategies in these more complex settings. More generally, even our most complex environments have relatively simple and constrained observation and action spaces compared to the real world, and the space of possible experiments is generally small and heavily constrained. Our agents are also limited to fit with their environments (e.g. having the right number of observation inputs and action outputs for the number of variables in the DAG environment). It is an open question whether agents could in practice learn causal strategies from passive data in environments that relaxed these simplifying assumptions, though our results suggest it would be possible in principle.
What this paper does not imply: We do not intend suggest that passive learning as beneficial as active learning, or that confounding is not challenge for passive learners. Active learning can be substantially more effective or efficient than passive learning in humans [23] 49 animals [25]. and RL agents [88 5 3.57788 [8) It can be difficult to learn from purely positive examples, without some Indeed, while the LM we evaluated was only trained passively, many deployments of LMs incorporate interactive learning through RL from Human Feedback [12 e.g. [59]; one key benefit of this may be the value of training the model on active, on-policy data. We therefore expect that agents or LMs that are fine-tuned in an active setting, or with more sophisticated causal corrections
active correction of pofllppimlllpllicyy behavior [e.g. 65 15
(56), will generally develop more robust and generalizable causal strategies than those trained in purely passive settings.
Furthermore, confounding is clearly a challenge for any system that attempts to discover and use causal structure. The ongoing debates about confounding in the medical literature [e.g. 377 75) illustrate that this is challenge even for causal experts such as human scientists. However, our work contributes to understanding the regimes in which confounding during passive learning can be overcome by experimentation at test time. We consider range of confounds; even our train-test splits effectively correspond to strong confounds which are not directly observable, and for which the agent/L. only experiences training data sampled under one value, and then tested on situations sampled under another. Nevertheless, agents can in some cases generalize despite these confounds, particularly when supported by explanations. These results highlight the role that language and explicit descriptions of the quest for causal structure can play in acquiring causal strategies.
We illustrate the potential for purely passive learning of active causal strategies across various settings, particularly with the support of explanations. These results have implications for understanding the causal limitations of passive learning for a system that can only interact test time, and have particular relevance to understanding the capabilities and behaviour of language models.
Conclusions:
Acknowledgements
We thank Nan Rosemary Ke, Tobias Gerstenberg, and the anonymous reviewers for helpful comments and suggestions.
'Though these can be accommodated in DAG by "unrolling" the graph.
10

---

References
[1] Rishabh Agarwal, Dale Schuurmans, and Mohammad Norouzi. An optimistic perspective on offline reinforcement learning. In International Conference Machine Leaminng pages 104-114. PMLR, 2020.
[2] Jacob Andreas. Language models as agent models. arXiv preprint arXiv:2212.0168/ 2022.
[3] Kai Arulkumaran, Dylan R Ashley, Jürgen Schmidhuber, and Rupesh K Srivastava. All you need is supervised learning: From imitation learning to meta-rl with upside down rl. arXiv preprint 22K2::2202.196000 2022.
[4] Emily M Bender and Alexander Koller. Climbing towards NLU: On meaning, form, and understanding in the age of data. In Proceedings of the 58th annual meeting of the association for computational linguistics, pages 5185-5198, 2020.
[5] Jeannette Bohg, Karol Hausman, Bharath Sankaran, Oliver Brock, Danica Kragic, Stefan Schaal, and Gaurav Sukhatme. Interactive perception: Leveraging action in perception and perception in action. IEEE Transactions on Robotics, 33(6):1233-1291 2017.
[6] Stephan Bongers, Patrick Forré, Jonas Peters, and Joris M Mooij.
Foundations of structural causal models with cycles and latent variables. The Annals of Statistics, 49(5):2885-29) 2915, 2021.
[7] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL
http://github..u cwtp:hhhbb.ooomooooooooeeoo http://ggttuubbccooooog///////// google/jax
[8] Thomas Carta, Clément Romac, Thomas Wolf, Sylvain Lamprier, Olivier Sigaud, and Pierre- Yves Oudeyer. Grounding large language models in interactive environments with online reinforcement learning. arXiv preprint arXiv:2302.02662, 2023.
(9] Stephanie CY Chan, Ishita Dasgupta, Junkyung Kim, Dharshan Kumaran, Andrew Lampinen, and Felix Hill. Transformers generalize differently from information stored in context in weights. MemARI workshop, NeurIPS, 2022.
[10] Stephanic CY Chan, Adam Santoro, Andrew K Lampinen, Jane X Wang, Aaditya Singh, Pierre Richemond, Jay McClelland, and Felix Hill. Data distributional properties drive emergent few-shot learning in transformers. Advances in Neural Information Processing Systems, 2022.
[11] Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems, 14:15084-15097, 2021.
[12] Paul Christiano, Jan Leike, Tom Brown, Miljan Martic, Shane Legg, and Dario Amodei. Deep reinforcement learning from human preferences. Advances in neural information processing systems, 30, 2017.
[13) Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Carbonell, Quoe V Le, and Ruslan Salakhutdinov. Transformer-xl: Attentive language models beyond a fixed-length context. arXiv preprint 19X1::901 2019.
[14] Ishita Dasgupta, Jane Wang, Silvia Chiappa, Jovana Mitrovic, Pedro Ortega, David Raposo, Edward Hughes, Peter Battaglia, Matthew Botvinick, and Zeb Kurth-Nelson. Causal reasoning from meta-reinforrcement learning. arXiv preprint
2019.
[15] Pim De Haan, Dinesh Jayaraman, and Sergey Levine. Causal confusion in imitation learning. Advances in Neural Information Processing Systems, 32, 2019.
[16] Wenhao Ding, Haohong Lin, Li, and Ding Zhao. Generalizing goal-conditioned reinforce- ment learning with variational causal reasoning. arXiv preprint arXiv:2207.0008) 2022.

---

[17] Frederick Eberhardt, Clark Glymour, and Richard Scheines. On the number of experiments sufficient and the worst case necessary identify all causal relations among variables. arXiv preprini arXiv: 1207. 1389, 2012.
[18] Gonçalo Rui Alves Faria, Andre Martins, and Mário AT Figueiredo. Differentiable causal discovery under latent interventions. Conference on Causal Learning and Reasoning, pages 253-274, PMLR, 2022.
[19] Patrick Forré and Joris M Mooij. Constraint-based causal discovery for non- non-linear structural causal models with cycles and latent confounders. arXiv preprint arXiv: 03024, 2018.
[20] Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pages 2052-2062. PMLR, 09-15 Jun 2019. URR/https://proceeddii 9a. html
mlr. paress/99/ffjjmeet
[21] Clark Glymour, Kun Zhang, and Peter Spirtes. Review of causal discovery methods based on graphical models. Frontiers in genetics, 10:524, 2019.
[22] Madelyn Glymour, Judea Pearl, and Nicholas P Jewell. Causal inference in statistics: A primer. John Wiley & Sons, 2016.
[23] Todd M Gureckis and Douglas Markant, Self-directed learning: A cognitive and computa- tional perspective. Perspectives Psychological Science, 7(5):464-481, 2012.
[24] Peter Hase and Mohit Bansal. When can models learn from explanations? formal framework for understanding the roles of explanation data, arXiv preprint arXiv:2 2021.
[25] Richard Held and Alan Hein. Movement- -prouuccd-- stimulation the development of visually guided behavior. Journal of comparative and physiological psychology, 56(5):872, 1963.
[26] Tom Hennigan, Trevor Cai, Tamara Norman, and Igor Babuschkin. Haiku: Sonnet for JAX,
2020. http://: thub. com/ wwtpdeed/mmmmdeeeee dm- haiku
[27] Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al. Training compute-optimal large language models. arXiv preprint arXiv:2203. 2022.
[28] Ari Holtzman, Jan Buys, Li Du, Maxwell Forbes, and Yejin Choi. The curious case of neural text degeneration. In International Conference on Learning Representations, 2020.
[29] Cheng-Yu Yu Hsieh, Chun-Liang Li, Chih-Kuan Yeh, Hootan Nakhost, Yasuhisa Fujii, Alexander Ratner, Ranjay Krishna, Chen- Yu Lee, and Tomas Pfister. Distilling step-by-step! outperforning larger language models with less training data and smaller model sizes. Findings of the Association for Computational Linguistics: ACL 2023, 2023.
[30] Chentian Jiang, Nan Rosemary Ke, and Hado van Hasselt. Learning how to infer partial mdps for in-context adaptation and exploration. arXiv preprint arXiv:2302.04250, 2023.
[31] Nan Rosemary Ke, Silvia Chiappa, Jane Wang, Jorg Bornschein, Theophane Weher, Anirudh Goyal, Matthew Botvinic, Michael Mozer, and Danilo Jimenez Rezende. Learning to induce causal structure. arXiv preprint Xiix:2204.04877 2022.
[32] Emre Kıeıman, Robert Ness, Amit Sharma, and Chenhao Tan. Causal reasoning and large language models: Opening a new frontier for causality. arXiv preprint trXiv:2305.00050 2023.
[33] Diederík P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint 1rXiv:1412.6980, 2014.
[34] Lara Kirfel, Thomas Icard, and Tobias Gerstenberg. Inference from explanation. Journal of Experimental Psychology: General, 151(7):1481, 2022.
[35] Takeshi Kojima, Shixiang Shane Gu, Machel Reid, Yutaka Matsuo, and Yusuke Iwasawa. Large language models are zero-shot reasoners. arXiv preprint arXiv:2205. 11916, 2022.
12

---

[36] Eliza Kosoy, David M Chan, Adrian Liu, Jasmine Collins, Bryanna Kaufmann, Sandy Han Huang, Jessica B Hamrick, John Canny, Nan Rosemary Ke, and Alison Gopnik. Towards understanding how machines can learn causal overhypotheses. arXiv preprint arXiv:2206.08353, 2022.
[37] Chayakrit Krittanawong, Ameesh Isath, Robert S Rosenson, Muzamil Khawaja, Zhen Wang, Sonya E Fogg, Salím S Virani, Lu Qi, Yin Cao, Michelle T Long, et al. Alcohol consumption and cardiovascular health. The American Journal of Medicine, 135(10):1213-12333, 2022.
[38] Aviral Kumar, Justin Fu, Matthew Soh, George Tucker, and Sergey Levine. Stabilizing off- policy q-learning via bootstrapping reduction. Advances in Neural Information Processing Systems, 32, 2019.
[39] Brenden Lake, Tomer D Ullman, Joshua B Tenenbaum, and Samuel Gershman. Building machines that learn and think like people. Behavioral and brain sciences, 40, 2017.
[40] Andrew K Lampinen, Ishita Dasgupta, Stephanie CY Chan, Kory Matthewson, Michael Henry Tessler, Antonia Creswell, James L McClelland, Jane X Wang, and Felix Hill. Can language models learn from explanations in context? Findings of EMNLP, 2022.
[41] Andrew K Lampinen, Nicholas Roy, Ishita Dasgupta, Stephanic CY Chan, Allison Tam, James Mcclelland, Chen Yan, Adam Santoro, Neil C Rabinowitz, Jane Wang, et al. Tell me why! explanations support learning relational and causal structure. In International Conference on Machine Learning, pages 1 11868-11 11688118900. PMLR, 2022.
[42] Michael Laskin, Luyu Wang. Junhyuk Oh, Emilio Parisotto, Stephen Spencer, Richie Steiger- wald, DJ Strouse. Steven Hansen, Angelos Filos, Ethan Brooks, et al. In-context reinforcement learning with algorithm distillation. arXiv preprint arXiv:2210.1422 142111 2022.
[43] Belinda Z Li, Maxwell Nye, and Jacob Andreas. Implicit representations of meaning in neural language models. In Proceedings of the 59th Annual Meeting of the Association for Computational Linguisties and the 1 11th International Joint Conference on Natural Language Processing ( (Volume 1: Long Papers), pages 1813-1827, 2021.
[44] Kenneth Li, Aspen K Hopkins, David Bau, Fernanda Viégas, Hanspeter Pfister, and Martin Wattenberg. Emergent world representations: Exploring sequence model trained synthetic task. arXiv preprint arXiv:2210.13388 2022.
[45] Pengfei Liu, Weizhe Yuan, Jinlan Fu, Zhengbao Jiang, Hiroaki Hayashi, and Graham Neubig. Pre-train, prompt, and predict: systematic survey prompting methods in natural language processing. ACM Comput. Surv., 55(9). jan 2023. ISSN 0360-0300. doi: 10.1 (0.1145/35608.1 15. URL httpt:ttt//oii... org/ 10. 1145/3560815
[46] Tania Lombrozo. The structure and function of explanations. Trends in cognitive sciences, 10 (10):464-470, 2006.
[47] Tania Lombrozo and Susan Carey. Functional explanation and the function of explanation. Cognition, 9972).167-204, 2006.
[48] Andrew L Maas, Awni Hamnun, Andrew Y Ng, et al. Rectifier nonlinearities improve neural network acoustic models. In Proc. ieml, volume 30, page 3. Atlanta, Georgia, USA, 2013.
[49] Douglas B Markant and Todd M Gureckis. Is better select to receive? learning via active and passive hypothesis testing. Journal of Experimental Psychology: General, 143(1):94, 2014.
[50] Suraj Nair, Yuke Zhu, Silvio Savarese, and Li Fei-Fei. Causal induction from visual observations 2019.
for goal directed tasks. arXiv preprint
[51] Andrew Nam, Christopher Hughes, Thomas Icard, and Tobias Gerstenberg. Show and tell: Learning causal structures from observations and explanations.
[52] Sharan Narang, Colin Raffel, Katherine Lee, Adam Roberts, Noah Fiedel, and Karishma Malkan. WT5?! training text-to-text models to explain their predictions. arXiv preprint 2XXv:20044145466 2020.
13

---

[53] Khanh X Nguyen, Dipendra Misra, Robert Schapire, Miroslav Dudik, and Patrick Shafto. Interactive learning from activity description. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Pro- ceedings of Machine Learning Research, pages 8096-8108. PMLR, 18-24 Jul 2021. URL httpp:://rooesedings mlr wwwppr/es/see.. httpss//v13///guen/e/// html
[54] Mizu Nishikawa-Toomas Tristan Deleu, Jithendaraa Subramanian, Yoshua Bengio, and Laurent Charlin. Bayesian learning of causal structure and mechanisms with gflownets and variational bayes. arXiv preprint arXiv:2211.02763, 2022.
[55] Maxwell Nye, Anders Johan Andreassen, Guy Gur-Ari, Henryk Michalewski, Jacob Austin, David Bieber, David Dohan, Aitor Lewkowyez, Maarten Bosma, David Luan, et al. Show your work: Scratchpads for intermediate computation with language models. arXiv preprint arXiv:2112.00114 2021.
[56] Pedro A Ortega, Markus Kunesch, Grégoire Delétang, Tim Genewein, Jordi Grau-Moya, Joel Veness, Jonas Buchli, Jonas Degrave, Bilal Piot, Julien Perolat, et al. Shaking the foundations: delusions in sequence models for interaction and control. arXiv preprint arXiv:2110.10819, 2021.
[57] Georg Ostrovski, Pablo Samuel Castro, and Will Dabney. The difficulty of passive learning in deep reinforcement learning. Advances in Neural Information Processing Systems, 34: 23283-23295, 2021.
[58] Pierre-Y Yves Oudeyer. Frdric Kaplan, and Verena Hafner. Intrinsic motivation systems for autonomous mental development. IEEE transactions on evolutionary computation, 11(2): 265-286, 2007.
[59] Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models to follow instructions with human feedhack. Advances in Neural Information Processing Systems, 25:27730-27744. 2022.
[60] Judea Pearl. Theoretical impediments to machine learning with seven sparks from the causal revolution. arXiv preprint arXiv:/801.044 1801.04066, 2018.
[61] Judea Pearl. The seven tools of causal inference, with reflections machine learning. Commu- nications of the ACM, 62(3):54-60, 2019.
[62] Jonas Peters and Peter Bühlmann. Identifiability of gaussian structural equation models with equal error variances. Biometrika, 101(1):219-22) 228, 2014.
[63] Jonas Peters, Joris Mooij, Dominik Janzing, and Bernhard Schölkopf. Identifiability of causal graphs using functional models. arXiv preprint arXiv: 2012.
[64] Ben Prystawski and Noah D Goodman. Why think step-by-step? reasoning emerges from the locality of experience. arXiv preprint arXiv:2304.088: 03843, 2023.
[65] Stéphane Ross, Geoffrey Gordon, and Drew Bagnell. A reduction of imitation learning and structured prediction to no-regret online learning. In Proceedings of the fourteenth interna- tional conference on artificial intelligence and statistics, pages 627-635. JMLR Workshop and Conference Proceedings, 2011.
[66] Timo Schick, Jane Dwivedi-Yis Yu, Roberto Dessi, Roberta Raileanu, Maria Lomeli, Luke Zettle- moyer, Nicola Cancedda, and Thomas Seialom. Toolformer: Language models can teach themselves to use tools. arXiv preprint arXiv:2302.04761, 2023.
[67] Juergen Schmidhuber. Reinforcement learning upside down: Don't predict rewards-juse map them to actions. arXiv preprint arxiv: 2019.
[68] Rupesh Kumar Srivastava, Pranav Shyam, Filipe Mutz, Wojeiech Jaskowski, and Jürgen Schmidhuber. Training agents using upside-down reinforcement learning. arXiv preprint 1wXiv:1912.02877.2000 2019.
14

---

[69] Mirac Suzgun, Nathan Scales, Nathanael Schärli, Sebastian Gehrmann, Yi Tay, Hyung Won Chung, Aakanksha Chowdhery, Quoe V Le, Ed H Chi. Denny Zhou, et al. Challenging big- bench tasks and whether chain-of-thought can solve them. arXiv preprint arXiv:22 10.0926,, 2022.
[70] Jane Wang, Michael King, Nicolas Pierre Mickael Porcel, Zeb Kurth-Nelson, Tina Zhu, Charlie Deck, Peter Choy, Mary Cassin, Malcolm Reynolds, H Francis Song, et al. Alchemy: benchmark and analysis toolkit for meta-reinforremmet learning agents. Thirty-fifth Conference Neural Information Processing Systems Datasets and Benchmarks Track (Round 2), 2021.
[71] Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Ed Chi, Quoe Le, and Denny Zhou. Chain of thought prompting elicits reasoning in large language models. Advances in Neural Information Processing Systems, 2022.
[72] Moritz Willig, Matej Zečević, Devendra Singh Dhami, and Kristian Kersting, Causal parrots: Large language models may talk causality are not causal. preprint, 2023.
[73] Kevin Xia, Kai- Lee, Yoshua Bengio, and Elias Bareinboim. The causal-neura connection: Expressivenesss, learnability, and inference. Advances in Neural Information Processing Systems, 34:10823-10836 2021.
[74] Junzhe Zhang, Daniel Kumor, and Elias Bareinboim. Causal imitation learning with unobserved confounders. Advances in neural information processing systems, 01112223-12277 2020.
[75] Jinhui Zhao, Tim Stockwell, Tim Naimi, Sam Churchill, James Clay, and Adam Sherk. Asso- ciation between daily alcohol intake and risk of all-cause mortality: A systematic review and meta-analyses. JAMA Network Open, 633)
2023.
[76] Chunting Zhou, Pengfei Liu, Puxin Xu, Srini Iyer, Jiao Sun, Yuning Mao, Xuezhe Ma, Avia Efrat, Ping Yu, Lili Yu, Susan Zhang, Gargi Ghosh, Mike Lewis, Luke Zettlemoyer, and Omer Levy. Lima: Less is more for alignment. arXiv preprint arXiv:2305. /206, 2023.
[77] Hattie Zhou, Azade Nova, Hugo Larochelle, Aaron Courville, Behnam Neyshabur, and Hanie Sedghi. Teaching algorithmic reasoning via in-context learning. arXiv preprint arXiv:2211.09066, 2022.
15


Reviews:
[{"review_id": "ZDjd4Iuysk", "reviewer": "Reviewer_c7Wr", "summary": "The paper deals with the problem of understanding whether causal structures can be learned from passive observational data. The study shows formally that it is possible for an agent to learn a generalizable strategy from passive data, as long as the agent can intervene at test time. Then, the experimental results show empirically that agents trained on expert data can indeed generalize at test time to infer and use causal links which are never present in the training data.", "questions": "None", "limitations": "The paper should describe limitations more explicitly in a dedicated paragraph", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* The paper addresses an interesting and important problem\n* The paper shows that natural language explanations support generalizing causal strategies from passive data\n* The findings of the paper have several applications, including using tools with LLMs\n* The paper is well written and easy to follow", "weaknesses": "* It would be interesting to explore more complex environments", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We greatly appreciate the positive comments. We have added a paragraph to the discussion noting some of the limitations of the environments and tasks used, since these limitations were noted by several reviewers. Although we have not been able to experiment with substantially more complex environments during the brief response period, we have performed a few experiments that slightly increase the challenge of the present tasks (such as generalizing to different types of causal graphs at test time, or reducing coverage of the training distribution), as requested by other reviewers. The results can be found in the global response PDF."}, {"review_id": "vohJnNmP90", "reviewer": "Reviewer_X6cj", "summary": "This paper investigates whether it is possible for an agent to learn how to perform and make use of active causal interventions, despite being trained on purely passive / observational data (e.g. supervised learning, behavioral cloning, etc). The authors first argue that this is possible in principle, if the observed data contains behavioral traces of an expert who performs causal interventions in an environment to learn about its causal structure, and then exploits the learned information to achieve a goal. They then perform experiments in a causal DAG environment, a visual \"odd-one-out\" environment, and linguistic version of the \"odd-one-out\" environment (in order to test pretrained LLMs). In the casual DAG environment, they find that trained agents can learn to generalize causal exploration and exploitation strategies to unseen DAGs, just by passive learning of expert strategies on the training set. They find similar results in the odd-one-out environment, with stronger performance achieved when the training data includes natural language explanations of the interventions performed by the expert. Finally, they find that LLMs prompted with few-shot examples of causal experimentation can generalize to new settings where the latent cause of reward is absent from the few-shot examples, provided that explanation and reasoning traces are also provided in the few-shot examples.", "questions": "See above for various questions and suggestions I have about additional baselines, additional generalization tests, clarifying the role of explanations in the odd-one-out experiments, whether these models can be said to have \"truly learned\" causal reasoning, and whether the term \"perfectly confounded\" is apt.\n\nIn addition, I'm curious whether the authors still endorse the following analysis, given that their post-submission corrections show much better performance of passive learners in the odd-one-out experiment:\n\n> Indeed, although our passive agents performed far above chance in the odd-one-out environments, they did not achieve as robust of performance as the online RL agents in earlier work [30]. This gap may be due to irreversible choices in these settings, which are difficult to learn from purely positive examples, without some active correction of off-optimal-policy behavior [e.g. 48, 11].\n\nPost-Rebuttal Response:\n\nThe authors have addressed my questions and suggestions, and I continue to recommend acceptance.", "limitations": "Some limitations of their findings are discussed (see \"What this paper does not imply:\"), and the authors fairly summarize their findings as showing that passive learning of causal strategies is restricted to the case where expert intervention policies are actually present in the data. I think it'd be valuable to add discussion of exactly how much we can expect the learned causal strategies to generalize, and whether it is possible to establish if they actually constitute sound and complete causal discovery algorithms.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "This was a well-written paper that addresses an important topic of debate that has arisen about the capabilities of models trained purely on observational data, including pretrained (large) language models. I have several questions about the details of the experiments, and about the generality of causal experimentation strategies that are learned from observations of an expert, but overall I think the paper presents clear theoretical arguments and solid empirical demonstrations for its claims, showing that learning causal intervention strategies from offline/data is possible in principle, and is exhibited in practice to at least some degree of generality in actual ML models. The DAG experiments are especially well designed, showing that passively learned behavior is not well explained by approximate heuristics, and better explained by optimal causal reasoning. The odd-one-out and LLM experiments are more open to interpretation due to the absence of heuristic or correlational baselines, and not knowing whether the LLM was trained on similar data, but still provide evidence that passive learners may acquire generalizable causal intervention strategies in non-trivial settings (insofar as it is unlikely for the high performance to be achieved by a learned heuristic).", "weaknesses": "While the experiments have been designed carefully enough to show that observed casual strategies are generalized to some extent, I think they could be extended in the number of ways to probe the limits of this generalization.\n\nIn the DAG experiments for example, it is clear that the agent must have learned some kind of permutation invariant strategy for experimentation and intervention, given that the training data entirely lacks causal links that are present in the test episode. However, I think the experiments still leave open exactly how abstract the learned strategy is -- does it generalize to DAGs with 6 nodes, or exhibit size generalization more generally? (I expect not, due to the network architecture.) Can it generalize to DAGs with non-linear causal relationships? And can it adapt its experimentation strategy when it is *not* cued about which variables are relevant, perhaps by learning to imitate an expert that uses an optimal experimental design that minimizes the number of exploration queries before exploiting? (I found the explicit cues about which variables are relevant to be a fairly strong simplifying assumption -- it would be good to at least vary the number of variables which are cued as relevant to a wider range of numbers.) Similar size generalization tests would be interesting to test in the \"odd-one-out\" and LLM experiments -- e.g. having a test trial with 4 distinct objects, not just 3. For the LLM experiments in particular, it would also be good to say something about whether similar prompts are likely or unlikely to have appeared in Chinchilla 70B training data.\n\nGeneralizability is related to a more general question I have about what exactly \"counts\" as having learned an active causal strategy. In its strongest form, this might require verifying that a learned model actually implements a causal discovery algorithm and causal exploitation algorithm that is optimal, or at least sound and complete. Of course, this is very hard to verify without strong mechanistic interpretability techniques, and in the absence of that, heuristic baselines are helpful to compare against. Because heuristic baselines were included in the DAG experiments, I'm willing to believe that the model has successfully learned a close approximation to the optimal exploitation strategy in the n = 5 case, which presumably requires the agent to simulate the effect of a causal intervention on an internal representation of the DAG structure that was identified from the experimentation phase. But because similar baselines aren't included in the \"odd-one-out\" or LLM experiments, it could be that the learned strategy in those cases is just some \"clever trick\" that happens to do well most of the time without \"true\" causal reasoning. As such, it'd strengthen the paper to show some heuristic or correlational baselines for those experiments, if only to demonstrate that there are no obvious heuristics that can easily achieve the ~90% success rates in the odd-one-out and LLM experiments. For example, one baseline might be to always transform the color of the first/nearest object during experimentation (and to pick the unique colored object in the test trial if that transformation is successful). Another heuristic might be to always pick the object based on the feature that was most often unique in the training data. In addition, I think it'd good to note the limitations of what we can conclude about a model's learned strategies based on observational equivalence alone.\n\nRegarding the need to predict/include linguistic explanations in the \"odd-one-out\" experiments, it was unclear to me whether the results shown in Figure 3 (and in the corrected version, Figure 13) are from a model that was also trained to predict natural language explanations. My guess is that they are, but I think the presentation clarity of this aspect could be improved, e.g. by stating in the caption of Figure 3 that the results shown involve predicting the explanation, by giving an example in the Appendix of what a full training episode (including the natural language explanations) looks like, and by explaining how the natural language explanations are generated (I had to consult the Lampinen et al (2022) paper in order to determine this). \n\nFinally, I think it's somewhat misleading to describe the passive data in section 5.2.1 as \"perfectly confounded\" -- after all, the explanations constitute observational data that makes the setting no longer confounded! After all, it is only in worlds where some feature X causes reward R that choosing a object unique in feature X would result in the explanation E = e_X+ (\"In this game, I was rewarded for unique X\", and choosing an object unique in some other feature would result in the explanation E = e_X- (\"The rewarding dimension must not be Y.\" where Y != X). Hence, the relevant part of the causal DAG is:\n\nR <- X -> E\n\nOnce the agent observes E = e_X, it can reliably infer that there is causal link from X -> R in the true causal graph, since X is a common cause of both R and e_X. So the fact that explanations \"deconfound\" is no surprise at all.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the positive assessment and for the thoughtful and extremely detailed comments and suggestions; we believe they will improve the paper. We have made a number of changes and performed some additional experiments in an attempt to address these comments:\n\n* While agents are architecturally limited to a fixed DAG size (which we have noted in an added limitations paragraph, see below) we performed two DAG generalization experiments, to further explore the robustness of the agents’ learned strategies. Results are plotted in the attached PDF.\n  - Generalizing to nonlinear edges: the nodes of the DAG already had nonlinearities in the main experiments, but at test time (only) we additionally added a nonlinearity on the edges (i.e. applied to each input before summing them in the node). Surprisingly, this had relatively little effect on evaluation performance (in retrospect, we believe this is because the post-node non-linearity masks the effect of the edge ones in many cases).\n  - Generalizing from training on linear DAGs to evaluation on non-linear DAGs. This is a harder generalization task, and agents correspondingly show a larger gap from optimal performance in evaluation. However, the agents still exhibit a fair amount of generalization (~75-80% of optimal performance).\n  - (We also performed a somewhat-relevant experiment exploring how reducing coverage of the training DAG distribution affects generalization to the original evaluation split, as requested by another reviewer.)\n* The design of the odd-one-out tasks, and the setup of our experiments, is such that simple heuristics should yield fairly low performance. We have added some discussion that tries to clarify the experiments, particularly with respect to the following points:\n  - Heuristics such as focusing on individual attributes, or always choosing the object that is unique by color, say, will not yield above-chance performance for the agents.\n  - For the LMs, the fact that the same prompts are yielding high performance along the “training” dimensions included in the prompt examples, and in evaluation on the other dimension, suggest that the model is not using a simple heuristic focusing on a single dimension (which would yield at most 55% performance). \n  - Furthermore, since we performed three separate splits — one where each feature dimension was held out — and observed similar performance across each, the results do not seem to be strongly influenced by the relative presence of the dimensions in the pretraining data (although without explanations, we do note that there are some differences among the dimensions in generalization).\n* In addition to the above, we added the suggested new experiment of evaluating the LLM with four test objects (not plotted in the PDF response due to space constraints). That is, we kept the prompts the same (three objects in the experimentation and test trials) but in addition to evaluating on a held-out dimension, we introduced a fourth object in the final test trial of the evaluation episode. More specifically, we kept three objects in the experimentation trials of the episode, and only surprised the model with a fourth object (which was not unique along any dimension) in the final choice of objects. While this does reduce generalization performance somewhat relative to the original three-object version, the performance is still far above chance (72% accuracy on average, chance is 25%; broken down by held-out dimension: Color:  66% +/- 8.5%, Shape: 71% +/- 12.6% Texture: 78.3% +/- 3.3%). This suggests that the LMs (when prompted with explanations) are relatively robust in generalization, even if the evaluation differs from the prompt in multiple ways. \n* We have made a number of additional edits to improve clarity and accuracy, per the reviewer’s suggestions (and some from the other reviewers), including:\n  - adding more detail about the explanations and their use.\n  - noting the limitations of the observational analyses rather than discovering computational strategy the agent is using.\n  - rephrasing the “perfectly-confounded” statements to clarify that the explanation removes the confounding.\n  - stating the fact that the agent could not generalize to DAGs with more nodes due to architectural limitations as part of a broader paragraph added to the discussion on limitations of the current environments, tasks, and agents.\n  - removing the statement about the gap with RL agents; indeed, our updated results no longer support the claim that it is a large difference. We have kept the more general point about DAGGER etc., however."}, {"review_id": "9M9KmRtuDs", "reviewer": "Reviewer_KP4B", "summary": "The paper studies the possibility of learning to identify new casual structures through purely passive learning, which entails the feasibility of acquiring robust task performance in novel conditions through offline learning. The authors hypothesize that it is possible to learn generalizeable causality exploration strategies from offline data that demonstrate such strategies. An agent can then apply these strategies to explore the casual structures of new environments and effectively achieve goals in them. Experiments on several synthetic environments show that explanations play a crucial role in enabling learning generalizeable exploration strategies. These results provide further insights on the ability of LLMs to quickly acquire interactive skills, even when trained exclusively on offline data.  ", "questions": "See my \"weakness\" section.", "limitations": "The authors compose a very thoughtful and thorough discussion section. I like that the authors anticipate and warn readers to avoid mistakenly interpret the claims made in the paper as passive learning can be as good as active learning.", "rating": 8, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "- The paper studies a fundamental question whose answer can have profound impact on the understanding and development of AI agents. It challenges a long-standing assumption about the limitations of learning from observational data. It suggests that observational data are not all the same; certain forms observations can be more useful than others. The new understanding is powerful. It can help explain the recent successes of LLMs, which are mostly trained through observations of the world. More general, it offer valuable insights to better leverage offline data to reduce the cost of collecting interactive data.\n- The experiments are well-designed and the analyses are insightful. The results can lead to further theoretical and empirical investigation, especially on the role of the explanations in learning. \n- The paper is very clear and well-written. The arguments are nicely presented. ", "weaknesses": "A major weakness of the paper is that it is mostly empirical. \n\nWhile the authors provides intuition on why their hypothesis can be true, the paper is still missing a satisfactory theoretical explanation. The environments studied, while having interesting structures, are artificially constructed. The paper may further benefit from having experiments on real-world tasks. The artificial environment may have odd properties that allow generalization to occur. \n\nThe main hypothesis rests on the assumption that an agent can learn generalizeable exploration strategies. While the empirical results support this assumption, it is unclear why it should hold in theory, because the strategy demonstrations are also observational data. Why is it easier to learn generalizeable exploration strategies than directly learning generalizeable exploitation strategies? I'd strongly suggest the authors add ablation experiments that dissect the importance of the exploration phase (i.e. results without exploration). This would strengthen the claim that the exploration phase is needed. \n\nThe authors emphasize the role of explanations in acquiring robust exploration strategies, but they neither clearly define \"explanations\" nor show examples in the main paper. This term is quite general and can mean many different things. While some examples are shown in the appendix, I think it should be more explicit in the main paper what type of explanations are used, what causal structures/properties they describe. I would also like to see more discussions or experiments on the effect of the explanation medium on the performance. Do explanations have to be in natural language? Or do they just need to have certain properties?\n\nRelated work: \n\n[1] Interactive learning from activity description: provides descriptions of activities as feedback to agents, leading to more efficient and effective learning than RL. \n\n[2] Distilling Step-by-Step! Outperforming Larger Language Models with Less Training Data and Smaller Model Sizes: shows that learning to predict explanations can enhance the pre-training of LLMs. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their thoughtful review; we are glad that they found the paper interesting, and we found their suggestions very useful. We have attempted to address them as follows:\n* We have performed a set of ablation experiments that remove the exploration phase in two ways; performance is substantially degraded in both (results are plotted in the response PDF). These results demonstrate the importance of the exploration phase.\n  - First, we removed the exploration phase entirely (that is, the agent does not have a chance to obtain any information before it needs to exploit and achieve a goal). As expected, the agent performs very poorly without this phase (it learns a very simple strategy of intervening positively on the target node, which achieves consistent, but modest, rewards).\n  - Next, we replaced the exploration with a purely observational phase. Rather than allowing the agent to intervene to explore, it simply observes an equivalent number of samples from the DAG distribution without interventions. While the agent is able to use these observations to achieve slightly better performance, it still performs far worse than an agent that is allowed to intervene. This demonstrates the importance of exploration, and in particular the importance of allowing the agent to intervene at test time.\n* We have added some clarification of the explanations and how they are used in the RL environments (including examples in the text and a note in the figure caption), as well as in the LM experiments. We agree that the benefits of the explanations are not restricted to natural language explanations; rather it is the abstractions that they capture that are important, and we have noted this in the discussion.\n* We have added these references to the relevant portions of the discussion.\n"}, {"review_id": "5V86jBZwwI", "reviewer": "Reviewer_U5ax", "summary": "This paper provides a thorough analysis of whether agents trained to imitate passive data can indeed learn active strategies for intervening during evaluation. In particular, this paper demonstrates that agents can \"learn\" causal structure from passive data in that they correctly intervene on the appropriate variables during testing, even if such interventions were not encountered during training. A formal \"existence proof\" for active causal learning is given. Experiments are demonstrated on a causal DAG environment, a pixel-based shapes environment, and few-shot prompting of a language model. ", "questions": "Were any ablations done to study how the coverage of the training distribution affects the ability of the agent to learn causal strategies? For example, in the adaptive experimentation setting in the DAG environment the expert policy only explores all the relevant nodes. What happens if the expert policy explores more nodes than the causally relevant nodes? The agent still has enough information to identify the causal structure, but it would be interesting to see if it is able to behave optimally in this case. \n\nMore generally, do the authors expect the learning of active causal strategies from passive data to be efficient when presented with more complex environments? ", "limitations": "The limitations are adequately addressed.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "*Originality:* The paper provides an interesting demonstration that refutes some claims from the causality literature. In particular, the paper shows that \"it is possible to learn generalizable strategies for causal experimentation and intervention from passive data alone\" given that the data was generated from an expert's actions. \n\n*Quality:* the experiments appear technically sound and are well-executed.\n\n*Clarity:* Overall the paper is well-written, with understandable diagrams + captions. Moreover, the flow of the paper is impressive, with the authors smoothly anticipating and responding to potential concerns.\n\n*Significance:* The paper qualifies some claims surrounding agents' abilities to learn active causal strategies from passive data. In particular, this extends our understanding of the limits of agent learning. Moreover, this may amplify the concerns for safety of such agents, as they may appear to learn generalizable reasoning strategies even through passive data. ", "weaknesses": "While overall the paper is technically sound, the first claim made in the contribution seems a bit too strong when compared against the setting given by the proof. In particular, the authors claim \"We formally explain why it is possible for an agent to learn a generalizable strategy for exploiting causal structure from passive data, as long as the agent can intervene at test time\". However, their proof relies on an assumption that the system is accurately represented by a finite DAG, which may be unrepresentative of real-world scenarios, which may include cycles due to feedback loops. It would be more accurate to characterize their contribution as an \"intuitive explanation of why agents may be able to learn generalizable strategies for exploiting causal structure under certain settings, namely that the system can be represented as a finite DAG.\"\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their positive assessment of the paper and the results. We agree with the reviewer’s comment about the phrasing of the claims, and we have updated the statements in the abstract and introduction accordingly. We have also added a paragraph to the discussion explicitly focused on the simplicity of the environments and their causal structures (the limitations of finite DAGs, as well as some of the comments from the other reviewers), and noting the possibility that results would be different in more complex settings. We also performed some additional experiments to evaluate generalization to different kinds of causal DAGs (e.g. linear to nonlinear), which show some robustness of the learned strategies, as well as an experiment on generalization of the LLMs to an unexpected extra object in evaluation.\n\nFinally, we performed an experiment where we explored coverage of the training distribution in a slightly different sense, by holding out successively more DAG structures than in the original experiments. In the original experiments, there are 1,277 possible DAG structures that are consistent with the training graph sampling constraints (out of ~29,000 possible DAGs over 5 nodes). We ran two more training runs where we imposed either 3 or 7 more non-dependency constraints on the graph, which reduces the number of possible training DAGs to 428 or 195, respectively. With 428 training DAGs, agents still generalize reasonably well, but with 195 they overfit quite strongly. (Fig. 14 in the global response PDF.) This suggests that some breadth of coverage is required for generalization, but that it is not necessary to thoroughly cover the full distribution."}]

Author rebuttal:
We appreciate the thoughtful and thorough reviews. We are pleased that the reviewers found the paper “interesting” and “important” and felt that it addresses a “fundamental question,” and is “well-written” and “sound.” We have tried to clarify and refine the points noted by the reviewers (see individual responses for more details). In the process, we have run several of the suggested ablations and generalization experiments to clarify the results:
* Ablations that explicitly demonstrate the importance of allowing an interventional exploration phase at evaluation time. If we remove the active (interventional) exploration phase, either replacing it with a purely observational phase (where the agent simply sees samples from the DAG distribution), or removing it entirely, the agent performs much worse. (Fig. 13 in the response PDF.)
* Exploring how causal strategy generalization depends on the variety of causal DAG structures in training, by running two additional experiments that hold out successively more DAG structures. If we reduce the number of possible DAG structures to about 33% of the original training set, the agents still generalize fairly well, but if we reduce it to only 15%, the agents perform much worse. (Fig. 14 in the response PDF.)
* Evaluating the ability of agents to generalize their causal strategies to DAGs with different functional structures at evaluation time, either by testing with additional nonlinearities on the edges, or by training the agents on linear DAGs and testing on nonlinear ones. In either case we see reasonable generalization, though the linear -> nonlinear case is much harder. (Fig. 15 in the response PDF.)
* Evaluating the LLM with a surprise extra test object. That is, we kept the prompts the same (three objects in the experimentation and test trials) but in addition to evaluating on a held-out dimension, we introduced a fourth object in the final test trial of the evaluation episode. More specifically, we kept three objects in the experimentation trials of the episode, and only surprised the model with a fourth object (which was not unique along any dimension) in the final choice of objects. While this does reduce generalization performance somewhat relative to the original three-object version, performance is still far above chance (72% accuracy on average, chance is 25%; broken down by held-out dimension: Color:  66% +/- 8.5%, Shape: 71% +/- 12.6% Texture: 78.3% +/- 3.3%). This suggests that the LMs (when prompted with explanations) are relatively robust in generalization, even if the evaluation differs from the prompt in multiple ways.  (Not plotted in the PDF due to space constraints.)

We believe these experiments, and the edits suggested by the reviewers, have improved the paper, and we thank the reviewers for their valuable suggestions.

Meta-review: This paper proposes an interesting strategy to improve generalization. The ideal is to motivate agents to learn a policy that actively learns the causal structure in the data on deployment, followed by reward optimizing behavior and this policy results in better generalization by virtue of having learned causal dependencies between variables not observed in training.

The overall reception to this paper has been positive. Authors also conducted additional ablations, which I appreciated on the generalization quality vs number of causal DAGs used in training. The interesting aspect is the it does not deteriorate as fast as I expected.

My own criticism of the paper is mainly that I believe it lacks precision in the claims it is making. While empirical performance is great, the paper doesn't answer exactly how the causal structure is benefiting the exploitation phase. I would also have liked more associational baselines.

Considering the overall positive response, I recommend an accept and encourage the authors to include a more nuanced discussion in the camera ready version.
Decision: Accept (poster)