Generalizing Goal-Conditioned Reinforcement Learning with Variational Causal Reasoning
Authors: Wenhao Ding
Haohong Lin
Bo Li
Ding Zhao
Venue: NeurIPS 2022 Accept (2022)
Keywords: Reinforcement Learning
Generalization
Causal Reasoning

Abstract: As a pivotal component to attaining generalizable solutions in human intelligence, reasoning provides great potential for reinforcement learning (RL) agents' generalization towards varied goals by summarizing part-to-whole arguments and discovering cause-and-effect relations. However, how to discover and represent causalities remains a huge gap that hinders the development of causal RL. In this paper, we augment Goal-Conditioned RL (GCRL) with Causal Graph (CG), a structure built upon the relation between objects and events. We novelly formulate the GCRL problem into variational likelihood maximization with CG as latent variables. To optimize the derived objective, we propose a framework with theoretical performance guarantees that alternates between two steps: using interventional data to estimate the posterior of CG; using CG to learn generalizable models and interpretable policies. Due to the lack of public benchmarks that verify generalization capability under reasoning, we design nine tasks and then empirically show the effectiveness of the proposed method against five baselines on these tasks. Further theoretical analysis shows that our performance improvement is attributed to the virtuous cycle of causal discovery, transition modeling, and policy training, which aligns with the experimental evidence in extensive ablation studies.



Generalizing Goal-Conditioned Reinforcement Learning with Variational Causal Reasoning
Wenhao Ding', Haohong Lin', Bo Li?, Ding Zhao'
Carnegie Mellon University
University of Illinois Urbana-Champaiin
{wenhaod, hachongl Maabongg1Gaa Candrew. edu, 1bo@illinois. dingzhaodemu.ed edu
Abstract
As a pivotal component attaining generalizable solutions in human intelligence, reasoning provides great potential for reinforcement learning (RL) agents' gen- eralization towards varied goals by summarizing part-to- whole arguments and discovering cause-and-effece relations. However, how to discover and represent causalities remains a huge gap that hinders the development of causal RI.. In this paper, we augment foal-Conditioneed RI. (GCRL) with Causal Graph (CG), a structure built upon the relation between objects and events. We novelly formulate the GCRL problem into variational likelihood maximization with CG latent vari- ables. To optimize the derived objective, we propose a framework with theoretical performance guarantees that alternates between two steps: using interventional data to estimate the posterior of CG; using CG to learn generalizable models and interpretable policies. Due to the lack of public benchmarks that verify gener- alization capability under reasoning, we design nine tasks and then empirically show the effectiveness of the proposed method against five baselines these tasks. Further theoretical analysis shows that our performance improvement is attributed the virtuous eycle of causal discovery, transition modeling, and policy training, which aligns with the experimental evidence in extensive ablation studies. Code is available on https:
pettps///githee http://gttttbbbtee wtwcchu...oomccc www.cimmmmiiiii hatmm/Gllgamessh/G.e.e.
Introduction
The generalizability, which enables algorithm to handle unseen tasks, is fruitful yet challenging in multifarious decision-making domains. Recent literature 2 reveals the critical role of reasoning in improving the generalization of reinforcement learning (RL). However, most off-the-shelf RL algorithms [4] have not regarded reasoning as an indispensable accessory, thus usually suffering from data inefficiency and performance degradation due to the mismatch between training and testing settings. To attain generalization at the testing stage, some efforts were put into incorporating domain knowledge to learn structured information, including sub-task decomposition [ and program generation 6 7 & [0]. which guide the model to solve complicated tasks in an explainable way. However, such symbolism-dominant methods heavily depend on the re-usability of sub-tasks and pre-defined grammars, which may not always be accessible in decision-making tasks.
Inspired by the close link between reasoning and the cause-and-effect relationship. causality is recently incorporated to compactly represent the aforementioned structured knowledge RL training ([T]. Based the form of causal knowledge, we divide the related works into two categories, i.c., implicit and explicit causation. With implicit causal representation, researchers ignore the detailed node that influences the
causal structure. For instance, [12] extracts invariant features
reward function, while the other node consists of task-irrelevant features [ 15
16] This
neat structure has good scalability but requires access to multiple environments that share the
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

Causal Graph
Training
Testing
stop
Goal
Goal #2
Goal
color
shape
select
select
color
shape
distribtton
Spurccusness
Composition
Generalizzaion
The causal graph of this task. Right: Training setting: the hexagon is always pink and the box is always blue. Three testing settings: (1) In distribution: the same as the training setting. (2) Spuriousness: swap the color and shape to break the spurious correlation. (3) Composition: increase the height of the goal.
Figure I: The robot picks and places objects to achieve given goals. Left:
In contrast, one can turn the explicit side by estimating 21, which uses directed graphical models to capture the causality the environment. A pre-request for this estimation the object-level or event-level abstraction of the observation, which is available in most tasks and also becoming frequently studied problem [22 24.. However, existing explicir causal reasoning RL models either require the true
16 detailed causal structures 18
same invariant feature 12
19
20
causal graph 25 or rely heuristic design without thcoretical guarantees [[8].
In this paper, we propose GeneRAlizing by DiscovERing (GRADER), a causal reasoning method that augments the RL algorithm with data efficiency, interpretability, and generalizability. We mainly focus on Goal-Conditiones RL (GCRL) (26]. where different goal distributions during training and testing reflect the generalization. We formulate the GCRL into a probabilistic inference problem [27 with learnable causal graph as the latent variable. This novel formulation naturally explains the learning objective with three components transition model learning, planning, and causal graph discovery leading to an optimization framework that alternates between causal discovery and policy learning to gain generalizability. Under some mild conditions, we prove the unique identifiability of the causal graph and the theoretical performance guarantee the proposed framework.
To demonstrate the effectiveness of the proposed method, we conduct comprehensive experiments in environments that require strong reasoning capability. Specifically, generalization settings, i.e., spuriousness and composition, and provide an example to illustrate these settings Figure The evaluation results confirm the advantages of our method in two aspects. First, the proposed data-efficient discovery method provides an explainable causal graph yet requires much fewer data than previous methods, increasing data efficiency and interpretability during task solving. Second, simultaneously discovering the causal graph during policy learning dramatically increases the success rate of solving tasks. In summary, the contribution of this paper is threefold:
design two types of
We use the causal graph as latent variable to reformulate the GCRL problem and then derive an iterative training framework from solving this problem.
We prove that our method uniquely identifies true causal graphs, and the performance of the iterative optimization is guaranteed with a lower bound given converged transition dynamics.
We design nine tasks in three environments that require strong reasoning capability and show the effectiveness of the proposed method against strong baselines on these tasks.
2 Problem Formulation and Preliminary
We start by discussing the setting consider in this paper and the assumptions required in causal reasoning. Then we briefly introduce the necessary concepts related to causality and causal discovery.
2.1 Factorized Goal-conditioned RL
We assume the environment follows the conlicondi Markov Decision Process (MDP) setting with full observation. This setting is represented by tuple M (S, A. P. R, G). where S is the state space, .A is the action space, P is the probabilistic transition model, G C S is the goal space which is set of assignment values to states, and r(s.g) 11 (s = g) E R is the sparse deterministic reward function that returns I only if the state match the goal g. In this paper, we focus on the goal-conditioned generalization problem, where the goal for training and testing stages will be sampled from different distributions Purain (g) and (g) We refer to goal E G as task

---

and use these two terms interchangeably. To accomplish the causal discovery methods, we make further assumption similar [20 28 for the state and action space:
Assumption 1 (Space Factorization). The state space A = {A1 x ... x } can be factorized to disjoint components {S and {A,]
(Si x... x SM and action space
The components representing one event or object's property usually have explicit semantic meanings for better interpretability. This assumption can be satisfied by state and action abstraction, which has been widely investigated in [22 23 24 Such factorization also helps deal with the high-dimensional states since it could be intractable to treat each dimension one random variable [8]
2.2 Causal Reasoning with Graphical Models
Reasoning with causality relies on specific causal structures, which are commonly represented as directed acyelic graphs (DAGs) [29] over variables. Consider random variables X i..,... Xa) with index set V == {I.....d) d}. A graph 9 = (V,E) consists of nodes V and edges & S V2 with (i,j) for any i.j E V. A node is called a parent of j if eij and Cjt É E. The set of parents of j is denoted by PAY. We formally discuss the graph representation of causality with two definitions:
Definition 1 (Structural Causal Models (29). A structural causal model (SCM) C := (S,U) consists of collection S of d functions X3 (A Uý), E [d], where PA (X1,..., Mxx) called parents of X3; and joint distribution U -U1.....UU Ua) over the noise variables, which are required to be jointly independent.
Definition 2 (Causal Graph 229, CG). The causal graph 9 of SCM is obtained by creating one node for each X, and drawing directed edges from each parent in PA to X,-
We note that CG describes the structure of the causality, and SCM further considers the specific causation from the parents of X, to X, via fj as well as exogenous noises Uj. To uncover the causal assume that the CG satisfies the Markov Property and Faithful- ness 129), which make the independences consistent between the joint distribution P(XI,.... X,) and the graph G. We also follow the Causal Sufficiency assumption 30 measured all the common causes of the measured variables.
structure from data distribution,
that supposes have
believes that two objects have causality only if they close enough while there is no edge between them if the distance is large. Instead of using such a local view of the causality, we assume the causal graph is consistent across all time steps, which also handles the local causal influence. The specific influence indicated by edges is estimated by the function (PA Uz).
Existing work 31 20
3 Generalizing by Discovering (GRADER)
With proper definitions and assumptions, we now look at the proposed method. We first de- rive the framework GRADER by formulating the GCRL problem latent variable model, which provides variational lower bound to op- timize. Then, divide this objective function into three parts and iteratively update them with performance guarantee.
action execution
Environment
interventional data
A1
S1
planner
<-...
model
S2
Model and Policy Learning update
Causal Discovery
(Section 3.2)
update ø (Section 3.3)
3.1 GCRL as Latent Variable Models
Action
State
Next State
Goal
The general objective of RI. is to maximize
Figure 2: The paradigm of GRADER.
the expected reward function w.r.t. learnable
policy model n. Particularly, in the goal-conditioned setting, such objective is represented as max, Ey ())() Er r(s", g)], where p(g) is the distribution of goal and (s a sT) is the action- state trajectory with maximal time step T. The trajectory ends only if the goal achieved g S or the maximal step is reached.
Inspired by viewing the RL problem as probabilistic inference "How to find actions to achieve the goal?" to "what are the actions we achieve the goal?". ", lcading
32
27
we replace the objective from
3

---

to likelihood maximization problem for p(r]s") with s= : =()= Different from previous work [333 that recasts actions as latent variables and infers actions that result in "observed" high reward, we decompose p(r(s") with 9 as the latent variable to get the evidence lower bound (ELBO)
logp(r)s") log p(T (p) yd 2 B.to s*) D -xLlg(()))ll(())) (1)
where the prior p()) and variational posterior (( represent distributions over graph structures, ie, the probability of the existences of edges in graphs. 11x denotes the Kullback-Leibler (KI.) divergence between two graphs, which will be explained in Section 3.3 Recall that the space of goal G is subset of the state, we extend the meaning of g and assume all trajectories achieve the goal in the final state [d i.e., 9 5² Such extension makes it possible to further decompose the first term of I as (refer to Appendix A.1.2
T-1
logp(T|9, 8*) = logg(s")
Σ
log 1 1 a., 9) + log (a (a'ls', ,8,88 + logp(g). (2)
Σ
Here, we use the fact that 9 is effective in both the transition model p(s"+1aa', s 9) and the policy model log p(a" |s", ,,s",,, g), g only influences the policy model, and the initial state so depends neither on G nor g. We also assume that both initial state logp(s)) and goal log p(g) follow the uniform distribution. Thus, the first and last terms of 2 are constants. The policy term ग,, selecting action a' according to both current state st and goal g. implemented with the planning method and is further discussed in Section 3.2 Finally, we maximize the likelihood P(T(s") with the following reformulated ELBO as the objective
J(0,0) = (91)) [log po "'','',9) + (a'lss s" st,9)] (xxllo(((r)]Ib(G) (3)
where 0 is the shared parameter of transition model ps (a s², g) and policy Ne (a' g). and Φ is the parameter of causal graph qo (9]T). To efficiently solve this optimization problem, we iteratively updates parameter & (causal discovery, Section and parameter 9 (model and policy learming, Section 3.2). as shown in Figure 2 Intuitively, these processes can be respectively viewed as the discovery of graph and the update of ft. which share tight connections as discussed in Section 2.2
3.2 Model and Policy Learning
Let us start with a simple case where we already obtain 9 and use it to guide the learning of parameter 0 via max [(,,)) Φ). Since the KL divergence of [[, Φ) does not involve we only need to deal with the first expectation term, i.e., the likelihood of transition model and policy. For the transition po (s'+1laa aa s², 9), we connect it with causal structure by further defining particular type of CG and denote it as G in the rest of this paper:
Definition 3 (Transition Causal Graph). We define a bipartite graph G, whose vertices are divided into two disjoint sets U = (A',S') and V = {8(+1). A' represents action nodes at step I, St state nodes at step t, and +1+ the state nodes at step t + 1. All edges start from set l and end in set V.
Model learning. This definition builds the causal graph between two consecutive time steps, which indicates that the values of states in step t ++ depend on values in step t. It also implics that the interventions [2 on nodes in U are directly obtained since they have no parent nodes. We denote the marginal distribution of S as PI:, which is collected by RL policy æ. Combined with the Definition of SCM, we find that p=(s*+1 la,, s',9) essentially approximates a collection of functions f, following the structure G, which take input the values of parents of the state node 8y and outputs the value 8y. Thus, we propose to model the transition corresponding to G with a collection of neural networks fo(g) = (fo, 2=11 M,M. to obtain
5 for (PAS11 Ug),
where (P represents the values of all parents of node si at time step t and U, follows Gaussian noise U, ~ (0, 1). In practice, we use Gated Recurrent Unit 355 as f3 because it supports varying numbers of input nodes. We take sf as the initial hidden embedding and the rest parents [PA/\s.)]

---

as the imput sequence to fz. The entire model is optimized by stochastic gradient descent with the log-likelihood Pog (st+1|a', la', s², 9) as objective.
Policy learning with planning. Then we turn to the policy term ( s* 9) in [(0, 6). We optimize it with planning methods that leverage the estimated transition model. Specifically, the policy aims to optimize an action-state value function Q(s", a²) = EE En'= 7" r(st'rt +t, a +) a²) which can be obtained by unrolling the transition model with horizon of И steps and discount factor 7. In practice, we use model predictive control (MPC) (36] with random shooting 37 which selects the first action in the fixed-horizon trajectory that has the highest action-state value Q(s", a²), Le. *s*) = arg max, AA (8'' aa)) The formulation we derived so far is highly correlated to the model-based RL framework [38]. However, the main difference is that we obtain with variational inference by regarding the causal graph as latent variable.
3.3 Data-Efficient Causal Discovery
In this step, relax the assumption of knowing G and aim to estimate the posterior distribution (G]TT) to
Algorithm 1: GRADER Training
Input: Trajectory buffer B,, Causal graph G. Transition model fo. causal discovery threshold n
optimize ELBO 33 w.r.t. parameter Φ. In most score- based methods [3) likelihood is used to evaluate the correctness of the causal graph, i.e., better causal graph leads to a higher likelihood. Since the first term of 3 represents the likelihood of the transition model, we convert the problem of max [0, Φ) to the causal discovery that finds the true causal graph based on collected data samples. As for the second term of ( the following proposition shows that the KL divergence between ( and p(g) can be approximated by sparsity regularization (proof in Appendix A.4.2).
while 0 not converged do // Policy from planning Sample goal g ~ Purain while t < T do
a' + -Planner(Se, St, s*,g) 8 8 x- Env(a', g)
B, L.
s
1/ Estimate causal graph
for i i M + N do
Proposition 1 (KL Divergence as Sparsity Regular- ization). With entry-wise independent Bernoulli prior p(G) and point mass variational distribution q(G]T) of DAGs, ID DxL [yo [ip] is equivalent to an l sparsity regu- larization for the discovered causal graph.
for j < M do
Infer edge Gij + go [B, n)
11 Learn transition model Update fo(G) via with B
We restrict the posterior qu (G)T) to point mass distribution and use a threshold 1) to control the sparsity, We perform the discovery process from the classification perspective by proposing binary classifiers qs (e .7) to determine the existence of an edge Ezj. This classifier 4s (e [T. n) is implemented by statistic Independent Test (0) and 1) is the threshold for the p-value of the hypothesis. A larger n corresponds to harder sparsity constraints, leading to sparse G since two nodes are more likely to be considered independent. According to the definition g] we only need to conduct classification to edges connecting nodes between 21 and V. If two nodes are dependent, we add one edge directed from the node in li to the node in V. This definition also ensures that we always have ((GTT) E QD where QD is the class of DAG. With this procedure, we identify a unique CG G under optimality:
Proposition (Identifiability). Given an oracle independent test, with an optimal interventional data distribution pr.. causal discovery obtains 6² that correctly tells the independence between any two nodes, then the causal graph is uniquely identifiable, with et, (0-((e))7)) Vi E [M + N1.j E [M]]
In practice, we use -tett for discrete variables and the Fast Conditional Independent Test 40 for continuous variables. The sample size needed for obtaining the oracle test has been widely investigated [I].. However, testing with finite data is not a trivial problem, as stated in [4) especially when the data is sampled from cool-conniiioned MDP. Usually, the random policy is not enough to satisfy the oracle assumption because some nodes cannot be fully explored when the task is complicated and has a long horizon. To make this assumption empirically possible, it is necessary to simultaneously optimize Ta ma(a")s Is", *",) to access more samples close to finishing the task, which is
further analyzed in Section 3.4 We also empirically support this argument Section 4.2 and provide
detailed theoretical proof in Appendix
A.3 and A.4
5

---

3.4 Analysis of Performance Guarantee
The entire pipeline of GRADER summarized in Algorithm To analyze the performance of the optimization of Bb, we first list important lemmas that connect previous steps and then show that the iteration of these steps in GRADER leads virtuous cycle.
By the following lemmas, we show the following performance guarantees step by step. Lemma I shows model learning is monotonically better convergence given better causal graph from causal discovery. Then the learned transition model helps to improve the lower bound of the value function during planning according to Lemma 2 Lemma Deinaarreveee reveals the connection between policy learning
and interventional data distribution, which in turn improves the quality of our causal discovery, as is
shown in Lemma E and Proposition 2
Lemma (Monotonicity of Transition Likelihood), Assume G* = (V. E²) be the true CG. for two CG 91 = (V, E and G2 = (V, E2). if (HD(GI,9*) < SHD(G2, 9*), 3 e s.t. Ei U [e} E2. when transition model 0 converges, the following inequality holds for the transition model in
log pe (s"+ s 9*) 2 log (s s 91) 2 (s , s4,92)
(5)
where SHD is the Structural Hamming Distance defined in Appendix
Lemma 2 (Bounded Value Function in Policy Learning). Given planning horizon oc, if we already have approximate transition model DTv((p ((''( a..pap'psp (..)..a))))))) a)) S the approximate
policy it achieves near-optimal value function (refer to Appendix A.3 for detailed analysis):
I - V S (1-7)² Cm
(6)
Lemma (Policy Learning Improves Interventional Data Distribution). With a step reward r(s, a) 18((s g). we show that the value function determines upper bound for TV divergence between the interventional distribution with its optimal (proof details in Appendix A.3
Dir(PI: .Py) S - (1T))
where pış the marginal state distribution in interventional data, and Py is the goal distribution. better policy with larger V (s) enforces the distribution of interventional data toward the goal.
Lemma 4 (Interventional Data Benefits Causal Discovery). For <g = min >0 Pg. DDv((PI; Pg) <
Eg. the error of our causal discovery is upper bounded with Eg[SHD(G. g*)] |
After the close-loop analysis of our model, we are now able to analyze the overall performance of the proposed framework. Under the construction of pe(s+tlaa Iaa,, ,s (.) with AN-parameterized functions. the following theorem shows that the learning process will guarantee to perform a close estimation of true ELBO under the iterative optimization among model learning, planning, and causal discovery.
(Tv(pps (((s' |s, a), p(s'(s,a)) S Em, if the Theorem 1. With T-step approximate transition dynamics WTv (iss
goal distribution satisfies Eg -2mEm and the distribution prior CG is entry-wise independent Bernoulli(Eg), GRADER guarantees achieve approximate ELBO I with the true ELBO J
(1-68) (15) - 1),
- - (1- C T + log
(8)
Eg
An intuitive understanding of the performance guarantee is that better transition model indicates better approximation of objective J. The proof of this theorem and corresponding empirical results
are in Appendix A.5
4 Experiments
In this section, we first discuss the setting of our designed environ- ments well as the baselines used in the experiments. Then, we provide the numerical results and detailed discussions to answer the following important research questions: Q1. Compared to existing strong baselines, how does GRADER gain performance improve- ment under both in-distribution and generalization settings? Q2.
(1) Slack
[2] Unlock
cousell graph
(3) Cresh
(4) Chemistry

---

Compared to an offline random policy, how does a well-traines policy improve the results of causal discovery? Q3. Compared score-based causal discovery, does the proposed data-efficient causal discovery pipeline guarantee identifying the true causal graph as stated in Section 3.3 Q4. Considering the correctness of causal graphs, how does the imperfect causal graph influence the task- solving performance of GCRL agents?
4.1 Environments and Baselines
Since most commonly used RL benchmarks do not explicitly require causal reasoning for generaliza- tion, we design three new environments, which are shown Figure ( (excluding Chemistry [K These environments use the true state as observation to disentangle the reasoning task from visual understanding. For each environment, we design three settings in-distribution (/), spuriousness (S), and composition (C) corresponding to different goal distributions for generalization. We use Prein (9) and (g) represent the goal distribution during training and testing, respectively. uses the same Pirain (g) and Prea (g), S introduces spurious correlations in Prain (g) but remove them in hes)(g); and C contains more similar sub-goals in Prea(g) than in Ptrain (g) The details of these
settings in are briefly summarized in the following (details in Appendix C.2.1
where the agent must stack objects match specific shapes and colors. In Stack-S, we let the same shape have the same color in Ptain (g) but randomly sample the color and shape in (g))(g). In Stack-C, the maximum number of object is two in Prain (g) but five in (g))
Stack:
We design this manipulation task inspired by the Corla
Unlock: We design this indoor house-holding task for the agent to collect a key open doors. This environment is built upon the Minigrid A5]. In Stack-S, the door and the key are always in the same row in Pusin (g) but uniformly sample in (g)) In Unlock-C, there are one door in Pusin (g) but two doors in Prest (g).
The occurrence of accidents usually relies on causality, e.g., an autonomous vehicle (AV) collides with a jaywalker because its view is blocked by another car [46] We design such a crash scenario based on highway-env [47]. where the goals are to create crashes between pedestrian and different AVs. In Stack-S, the initial distance between AV and pedestrian is a constant in Puain (g) but irrelevant in (g) In Stack-C, there is one pedestrian in (gain(g) but two in (r))(g).
Crash:
Chemistry 4 There are 10 nodes with different colors. An underlying causal graph controls the color-changing mechanism of all nodes. In one step, the agent changes the color of one node. The goal is to match the given colors of all nodes. In the spuriousness setting, we let all nodes have the same target color. There is no composition setting in this environment.
We use the following methods as our baselines to fairly demonstrate the advantages of GRADER. SAC: 48] Soft Actor-Critte is well-known model-free RL method that uses entropy to increase the diversity of action. ICIN:
25 It uses DAgger 49 to learn goal-conditioned policy with the causal graph estimated from the expert policy. We assume it can access the true causal graph for supervised learning. PETS: (50) We consider the ensemble transition model with random shoot planning as one baseline, which achieves generalization with the uncertainty-awant design. TICSA: [8] This is causal- MBRL method that simultancously optimizes a soft adjacent matrix (representing the causality) and transition model. ICIL: [6]] This method proposes an invariant feature learning structure that captures the implicit causality of multiple tasks. We only use it for transition model learning since the original method is designed for imitation learning. GNN: [51]] Since graph neural networks are good at learning structural information, we implement a GNN-based baseline using Relational Graph Convolutional Network.
4.2 Results Discussion
Overall Performance (Q1) We compare the testing reward of all methods under nine tasks and summarize the results Table to demonstrate the overall performance. Generally, our method outperforms baselines in all tasks except Stack-/ because this task is too simple for all methods. We note that the gap between our method and baselines in S and settings more significant than in the I setting, showing that our method still works well in the non-trivial generalization task. As
7

---

Table 1: Success rate (%) for nine settings in three environments. Bold font means the best.
Method
Stack-1
Stack-S
Stack-C
Unlock-1
Unlock-S
Unlock-C
Crasb-1
Crash-S
Crasb-C
SAC
34.7=16.1
22.1=:14.0
31.7+5.1
0.1
0.0+0.2
0,4
22.5+1766
18,6+8,7
6.74338
ICIN
71.8
71.0=7.
58.64883
31.7.99.
31,548.5 32.7
27.946.1
15.8
7.8.888
PETS
97.2:169
77.7=135
73.7
59.5
20.6:55.
28.31100
52.3
44.62125
37.1
TICSA
85.9
10.
76.2=8.3
18.5
33.
29.8+8.3
68.9459
56.8486
15.0+82
ICIL
93.7+5.9
81.2:1144
62.8=13.0
67.1+11.6
15.9:1477
53.6+11..
55.3+2099
21.7+11777
14.347.
GNN
45.7:.9.
39.0:110.4
41.7=x66
3.4123
3,4124
4.543.0
4.2+4.0
5.115.1
3.8428
Score
92.7=7..
90.5+7.
73.9
44.9428.1
23.1
36.2430.1
42.3+175
13.4
8,446.1
Full
92.9#63
86.0+9.
15.
63,8+9.2
18.3+74
53.7+14.3
69.8=14.
52.6+128
42,0+17,2
Offline
96.81558
95.4116.1
81.41778
13.818.11
13.917.5
11.716,9
13.1
30.2.116.5
14.9
GRADER
95.6115.
97.6.16.0
93.7.18.4
64.219.1
61.4
82.1
91.5
84.3
84.7
Indisttibuion
3purious
Compositon
Causal Discovery
GMBEE
KCL
Fel
2.0
PGTS
0.8
0.8-
T14
Offire
1.6
0.6
0.6
1.2
M2 2110
0.4
0.0
0.0
1.8
2.4
0.0
2.4
Time steps Ix101)
Time steps (=110))
Time steps (x10%)
Number aftimm
Figure 4: Left: Test reward of the Crash environment calculated with 30 trails. Right: The accuracy of causal graph discovery with samples from GRADER, Score, and Offline.
model-free method, SAC fails in all three tasks of Unlock and Crash environments since they have very sparse rewards. Without learning the causal structure of the environment, PETS even cannot fully solve Unlock-S, Unlock-C, and all Crash tasks. Both TICSA and ICIL learn the causality underlying the task so that they are relatively better than SAC and PETS. However, they are still worse than GRADER in two generalization settings because of the unstable and inefficient causal reasoning mechanism. We also find that even ICIN is given the true causal graph, the policy learning part cannot efficiently leverage the causality, leading to worse performance in generalization settings.
To further analyze the tendency of learning, we plot the curves of all methods under Crash in Figure Our method quickly learns solve tasks at the beginning of the training, demonstrating high data efficiency. GRADER also outperforms other methods with large gaps the later training phase. The training figures of the other two environments
be found in Appendix B.
Importance of Policy Learning (Q2) As we mentioned in Sec- tion 3.3 we empirically compare GRADER and Offline (52). which uses data collected from offline random policy, and plot results in the right part of Figure We use SHD [53] to com- pute the distance between the estimated causal graph and the true causal graph. The true causal graph for each environment can found in Appendix B.2 When we only use samples collected offline by random policy, we cannot obtain variables' values that require long-horizon reasoning, c.g., the door can be opened only if the agent is close to the door and has the key. As a consequence, the causal graph obtained by Offline harms the performance, gradually explores more regions and quickly obtains the true causal graph when iteratively discover the causal graph and update the policy.
Best Causal Graph
Full Causal Graph
0.6
Test reward
0.4
0.2
0.0
shown in Figure 4 Instead, GRADER
0.0
0.4
0.8
1.2
1.6
Time steps (x10*)
Figure 5:
Influence of different causal graphs in Unlock-S.
Advantage of Data-efficient Causal Discovery (Q3) To show the advantage of proposed constraint- based methods, we design a model named Score that optimizes a soft adjacent matrix using score- based method [54], which is recently combined with NN for differentiable training, for example,
8

---

Table 2: Discovery results on Chemistry environment (GRADER / Score). Bold font means the best.
Metric
Collider
Chain
Jungle
Full
SHD (L)
3,70=1179/15,44:
7.00::2.19/13. 0.40
2.40==12001 1000:11.20111001 10+:3331
Accuracy (1)
0.99: 20.06
0.99:...000.888-0000
0.98 -9.
0.99: 000000
Precision (T)
0.90::0.05/0.775. 10.
0.94110.4400. 20000
0.86:.0.000.88::.001
1.00 000:0:.00
Recall (t)
0.999...02/0.83..00
0.96::0.020.733-..00
0.96:
13
F-score (t) 0.94: 0.95=t0.03/0.7::000
0.91:
0.98:
in TICSA. According to the discovery accuracy shown in the right part of Figure 4 find that score-base based discovery is inefficient. Based on the performance of the Score model summarized in Table II we also conclude that it is not as good as constraint-basee method and has large variance due to the unstable learning of the causal graph.
Influence of Causal Graph (Q4) To illustrate the importance of the causal graph, we implement another variant of GRADER named Full, which uses a fixed full graph that connects all nodes between the sets U and V. According to the performance shown in Table and Figure E we find that the full graph achieves worse results than GRADER because of the redundant and spurious correlation. Intuitively, unrelated information causes additional noises the learning procedure, and the spurious correlation creates a shortcut that makes the model extract wrong features, leading to worse results in the spuriousness generalization shown in Table !
We then investigate how the correctness of the causal graph influences the performance. We use fixed graphs interpolating from the best causal graph to the full graph to train GRADER model in Unlock-5 and summarize the results in Figure 5 The more correct the graph is, the higher reward the agent obtains, which supports our statements in Section 5.4 that the causal graph is important for the reasoning tasks a better causal graph helps the model have better task-solving performance.
4.3 Further Analysis of Causal Discovery
Collider
Jungle
Finally, we conduct further analysis of the dis- covery performance on the Chemistry environ-
Discovores GRADER
ment 43, which is standard benchmark for evaluating causal discovery methods. In this en-
03 32 03
vironment, the colors of nodes are controlled by the causal graph, therefore, finding the true causal graph makes much easier to achieve the goal that requires matching all target colors. The agent can discover the graph by doing in- terventions via interacting with the environment. We consider four types of causal graphs (Col- lider, Chain, Jungle, Full) with 10 nodes in this experiment.
1.0
6.0
The discovery performance is shown in Table 2 with five metrics indicating the classification eиror. We can see that GRADER outperforms the Score method in all types of graphs. In Figure 6 we show the discovered graphs from GRADER (averaged over 10 seeds) and the true causal graphs of Collider and Jungle settings. We also show that GRADER achieves advantages over other baselines in solving the color-matching downstream task, which can be found with detailed experiment results in Appendix B.4
Tran
Figure Top:
Discovered causal graph from GRADER. Bottom: true causal graph.
5 Related Works
RL Generalization From the agent's view, algorithms focus actively obtaining the structure of the task or environment. Some decompose the given task into sub-tasks [5 when they encounter an unseen task, they rearrange those sub-tasks to solve new tasks. Instead of
56 57 3
and

---

dividing tasks, Symbolic RL learns program- program-based based policy consisting of domain language G g or grammar 158 The generated program is then executed (9] to interact with the environment, which has the potential to solve unseen tasks by searching the symbolic space. From the other view, we can generate environments to augment agents' experience for better generalization. One straightforward
way is data augmentation of image-based observation
(59 60
61
12
63
64
When extended to
other factors of the environment, Domain Randomization 665 and its variants 66
67
are proposed.
Considering the interaction between agent and environment, Carriculum Learning 68 also gradually generates difficult tasks to train generalizable agents.
Coad-Conditioned RL (GCRL) The generalization problem is naturally related to GCRL 12 which aims to train an agent for multiple tasks. From the optimization perspective, Universal Value Function [69], reward shaping [70], and latent dynamic model [Z]] are widely used tools to solve GCRL problem. Sub-goal generation (r
another intuitive idea to tackle the long horizon with sparse reward, where the core thing is to make sure the generated sub-goals are solvable. Finally, Hindsight Experience Replay (HER) 60 belonging to the relabelling category, is a ground-breaking
yet straightforward method that treats visited states "fake" goals when the goal and state share the same space, Later on, improved versions of HER 7
14 75 were widely studied. One limitation is that we cannot directly use a visited state goal if the goal has pre-conditions. Similar to our setting, [76 and 77 convert the GCRL problem to variational inference by regarding control as inference 127. 76
propose EM framework under the HER setting and 7ir treats the last state as the goal and estimates a shaped reward during training.
RL with Causal Reasoning Causality is now frequently discussed in the machine learning field to RL algorithms also incorporate causality to 25 and [19] explicitly estimate causal structures
complement the interpretability of neural networks (2 improve the reasoning capability (78.. For instance,
with the interventional data obtained from the environment. These structures can be used to constraint
Building dynamic models in model-based RL 18 80 2 based on causal graphs is also widely studied recently. Implicitly, we can abstract the causal structure and formulate it using the Block MDP [22] setting or training multiple encoders to extract different kinds of representations [13]. Following the idea of invariant risk minimization [8], they assume task-relevant features are invariant and shared across all environments, which can be used as the only cause of the reward.
output space [79] or adjust the buffer priority 200
Causal Discovery Causal discovery 82 is long-stand topic in economics and sociology, where the traditional methods can be generally categorized into constraint-basee and score-based. Constraint- based methods 30 start from complete graph and iteratively remove edges with conditional independent test 34 as constraints. Score-based methods 39 35
use metrics such as Bayesian Information Criterion 86 as scores and prefer edges that maximize the score given the dataset. Recently, researchers extend score-based methods with RL (87 or differentiable discovery [54] 88
89]. The former selects edges with learned policy, and the latter learns soft adjacent matrix with observational or interventional data, Active intervention methods are also explored [90] to increase the efficiency of data collection and decrease the cost of conducting intervention 91]
6
Conclusion
This paper proposes a latent variable model that injects a causal graph reasoning process into transition model learning and planning to solve GCRL problems under the generalization setting. We theoretically prove that our iterative optimization process can obtain the true causal graph. To evaluate the performance of the proposed method, we designed nine tasks in three environments. The comprehensive experiment results show that our method has better data efficiency and performance than baselines. Our method also provides interpretability by the explicitly discovered causal graph.
The main limitation of this work is that the explicit estimation of causal structure does not scale well to the number of nodes. Developing efficient gradient-based discovery methods could be a promising direction. In addition, the factorized state and action space assumption may restrict the usage of this work to semantic representations, which need to be processed with abstraction methods. We further discuss the potential negative social impact and additional limitations in Appendix C
Acknowledgements. We gratefully acknowledge support from the National Science Foundation under grant CAREER CNS-2047454.
10

---

References
[1] Kai Xu, Akash Srivastava, Dan Gutfreund. Felix Sosa, Tomer Ullman, Josh Tenenbaum, and Charles Sutton. A bayesian-symbolie approach to reasoning and learning in intuitive physics. Advances in Neural Information Processing Systems, 34, 2021.
[2] Miles Cranmer, Alvaro Janchez-Gonzalle Peter Battaglia, Rui Xu, Kyle Cranmer, David Spergel, and Shirley Ho. Discovering symbolic models from deep learning with inductive biases (2020). arXiv preprint arXiv:2006. 1287, 2006.
[3] Qing Li, Siyuan Huang, Yining Hong, Yixin Chen, Ying Nian Wu, and Song- Mong-Chun. Zhu. Closed loop neural-symbolic learning via integrating neural perception, grammar parsing, and symbolic reasoning. In International Conference on Machine Learning (ICML), 2020.
[4] Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
[5] Yuchen Lu, Yikang Shen, Siyuan Zhou, Aaron Courville, Joshua B Tenenbaum, and Chuang Gan. Learning task decomposition with ordered memory policy network. arXiv preprint arXiv:2103.10972 2021.
[6] Yichen Yang. Jeevana Priya Inala, Osbert Bastani, Yewen Pu, Armando Solar-Lezama, and Martin Rinard. Program synthesis guided reinforcement learning for partially observed environ- ments. Advances in Neural Information Processing Systems, 34, 2021.
[7] Mikel Landajuela, Brenden K Petersen, Sookyung Kim, Claudio P Santiago, Ruben Glatt, Nathan Mundhenk, Jacob Pettit, and Daniel Faissol. Discovering symbolic policies with deep reinforcement learning. In International Conference Machine Learning, pages 5979-5989. PMLR, 2021.
[8] Jiankai Sun1 Hao Sun1 Tian Han and Bolei Zhou. Neuro-symbolic program search for au- tonomous driving decision module design. In Conference on Robot Learning (CoRL), 2020.
[9] Zelin Zhao, Karan Samel, Binghong Chen, et al. Proto: Program-guided transformer for program program-guided tasks. Advances in Neural Information Processing Systems, 34, 2021.
[10] Shao-Hua Sun, Te-Lin Wu, and Joseph J Lim. Program guided agent. In International Conference on Learning Representations, 2019.
[11] Samuel J Gershman. Reinforcement learning and causal models. The Oxford handbook of causal reasoning, 1:295, 2017.
[12] Amy Zhang, Clare Lyle, Shagun Sodhani, Angelos Filos, Marta Kwiatkowska, Joelle Pineau, Yarin Gal, and Doina Precup. Invariant causal prediction for block mdps. In International Conference on Machine Learning, pages 11214-1 1224. PMLR, 2020.
[13] Manan Tomar, Amy Zhang, Roberto Calandra, Matthew E Taylor, and Joelle Pineau. Model-invariant state abstractions for model-based reinforcement learning. arXiv preprint 1rXiv:2102.09850, 2021.
[14] Sumedh A Sontakke, Arash Mehrjou, Laurent Itti, and Bernhard Schölkopf. Causal curios- ity: RI agents discovering self self-supe experiments for causal representation learning. In International Conference Machine Learning, pages 9848-9858. PMLR, 2021.
[15] Shagun Sodhani, Sergey Levine, and Amy Zhang, Improving generalization with approximate factored value functions. In ICLR2022 Workshop on the Elements of Reasoning: Objects, Structure and Causality, 2022.
[16] Ioana Bica, Daniel Jarrett, and Mihaela van der Schaar. Invariant causal imitation learning for generalizable policies. Advances in Neural Information Processing Systems, 34, 2021.
[17] Beining Han, Chongyi Zheng, Harris Chan, Keiran Paster, Michael Zhang, and Jimmy Ba. Learning domain invariant representations in goal-conditioned block mdps. Advances Neural Information Processing Systems, 34, 2021.

---

[18] Zizhao Wang, Xuesu Xiao, Yuke Zhu, and Peter Stone. Task indeeenndenn causal state abstraction. Workshop on Robot Learning: Seff-Superviee and Lifelong Learning, NeurIPS, 2021.
[19] Sergei Volodin, Nevan Wichers, and Jeremy Nixon. Resolving spurious correlations in causal models of environments via interventions. arXiv preprini arXiv:2002.05217, 2020.
[20] Maximilian Seitzer, Bernhard Schölkopf, and Georg Martius. Causal influence detection for improving efficiency in reinforcement learning. Advances in Neural Information Processing Systems, 34, 2021.
[21] Maxime Gasse, Damien Grasset, Guillaume Gaudron, and Pierre- Oudeyer. Causal reinforcement learning using observational and interventional data. arXiv preprint arXiv:2106.14421, 2021.
[22] David Abel. A theory of abstraction in reinforcement learning, arXiv preprint arXiv:2203.00399, 2022.
[23] Murray Shanahan and Melanie Mitchell. Abstraction for deep reinforcement learning. arXiv preprint urXiv:2202.058339 2022.
[24] David Abel, Dilip Arumugam, Lucas Lehnert, and Michael Littman. State abstractions for lifelong reinforcement learning. In International Conference on Machine Learning, pages 10-19, PMLR, 2018.
[25] Suraj Nair, Yuke Zhu, Silvio Savarese, and Li Fei-Fei. Causal induction from visual observations for goal directed tasks. arXiv preprint arXiv:191000175,, 2019.
[26] Minghuan Liu, Menghui Zhu, and Weinan Zhang. Goal-conditioned reinforcement learning: Problems and solutions. arXiv preprint arXiv:2201.08229, 2022.
[27] Sergey Levine. Reinforcement learning and control as probabilistic inference: Tutorial and review. arXiv preprint
2018.
[28] Craig Boutilier, Richard Dearden, and Moisés Goldszmidt. Stochastic dynamic programming with factored representations. Artificial intelligence, 121(1-2)499-007, 2000.
[29] Jonas Peters, Dominik Janzing, and Bernhard Schölkopf. Elements of causal inference: founda- tions and learning algorithms. The MIT Press, 2017.
[30] Peter Spirtes, Clark N Glymour, Richard Scheines, and David Heckerman. Causation, prediction. and search. MIT press, 2000.
[31] Silviu Pitis, Elliot Creager, and Animesh Garg. Counterfactual data augmentation using locally factored dynamies. Advances in Neural Information Processing Systems, 33:3976-3990, 2020.
[32] Abbas Abdolmaleki, Jost Tobias Springenberg, Yuval Tassa, Remi Munos, Nicolas Heess, and Martin Riedmiller. Maximum posteriori policy optimisation, arXiv preprint arXiv:1806.069)) 2018.
[33] Joseph Marino and Yisong Yue. An inference perspective on model-based reinforcement learning. In ICML Workshop on Generative Modeling and Model-Basee Reasoning for Robotics and AI, 2019.
[34] Marein Andrychowicz, Filip Wolski, Alex Ray, Jonas Schneider, Rachel Fong, Peter Welinder, Bob McGrew, Josh Tobin, OpenAl Pieter Abbeel, and Wojeiech Zaremba. Hindsight experience replay. Advances in neural information processing systems, 30, 2017.
[35] Junyoung Chung, Caglar Gulcchre, KyungHyun Cho, and Yoshua Bengio. Empirical evaluation of gated recurrent neural networks on sequence modeling. arXiv preprint arXiv:1412.3555, 2014.
[36] Eduardo F Camacho and Carlos Bordons Alba. Model predictive control. Springer science business media, 2013.
12

---

[37] Arthur George Richards. Robust constrained model predictive control. PhD thesis, Mas- sachusetts Institute of Technology, 2005.
[38] Tingwu Wang, Xuchan Bao, Igmasi Clavera, Jerrick Hoang, Yeming Wen, Eric Langlois, Shunshi Zhang, Guodong Zhang, Pieter Abbeel, and Jimmy Ba. Benchmarking model-based reinforcement learning. arXiv preprint arXiv:1907.020577 2019.
[39] David Maxwell Chickering. Optimal structure identification with greedy search. Journal of machine learning research, 30N0v): 2002.
[40] Krzysztof Chalupka, Pietro Perona, and Frederick Eberhardt. Fast conditional independence test for vector variables with large sample sizes. arXiv preprint arXiv: 02747, 2018.
[41] Clément I. Canonne, Ilias Diakonikolas, Daniel M Kane, and Alistair Stewart. Testing condi- tional independence of discrete distributions. 2018 Information Theory and Applications Workshop (ITA), pages 1-57. IEEE, 2018.
[42] Rajen D Shah and Jonas Peters. The hardness of conditional independence testing and the generalised covariance measure. The Annals of Statistics, 48(3):1514-1538 2020.
[43] Nan Rosemary Ke, Aniket Didolkar, Sarthak Mittal, Anirudh Goyal, Guillaume Lajoie, Stefan Bauer, Danilo Rezende, Yoshua Bengio, Michael Mozer, and Christopher Pal. Systematic evaluation of causal discovery in visual model based reinforcement learning. arXiv preprint arXiv:2107.00848, 2021.
[44] Ossama Ahmed, Frederik Träuble, Anirudh Goyal, Alexander Neitz, Yoshua Bengio, Bern- hard Schölkopf, Manuel Wüthrich, and Stefan Bauer. Causalworld: A robotic manipulation benchmark for causal structure and transfer learning. urXiv preprint urXiv:2010.04426, 2020.
[45] Maxime Chevalier- Chevalier-Boisvert, Lucas Willems, and Suman Pal. Minimalistic gridworld environ- ment for openai gym. https://giis httpss///gihhbb...onxxxiiimbbbb www.w/wwwww..c/..o//e.oe http://gggmmmiiiii...m 2018.
[46] Zenna Tavares, James Koppel, Xin Zhang, Ria Das, and Armando Solar-Lezama. A language for counterfactual generative models. In International Conference on Machine Learning, pages 10173-10182. PMLR, 2021.
[47] Edouard Leurent. An environment for autonomous driving decision-making https://githu: com/ httm//elerrenttighaaa-eeeee 2018.
[48] Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off- policy maximum entropy deep reinforcement learning with a stochastic actor. In International conference machine learning, pages 1861-1870. PMLR, 2018
[49] Stéphane Ross, Geoffrey J Gordon, and J Andrew Bagnell. No-regret reductions for imitation learning and structured prediction. In In AISTATS. Citeseer, 2011.
[50] Kurtland Chua, Roberto Calandra, Rowan McAllister, and Sergey Levine. Deep reinforcement learning in a handful of trials using probabilistic dynamics models. Advances in neural information processing systems, 31, 2018.
[51] Michael Schlichtkrull, Thomas N Kipf, Peter Bloem, Rianne van den Berg, Ivan Titov, and Max Welling. Modeling relational data with graph convolutional networks. In European semantic web conference, pages 593-607. Springer, 2018.
[52] Zheng-Mao Zhu, Xiong-Hui Chen, Hong-Long Tian, Kun Zhang and Yang Yu. Offline reinforcement learning with causal structured world models. arXiv preprint 2rX1v:2006.01474, 2022.
[53] Ioannis Tsamardinos, Laura Brown, and Constantin F Aliferis. The max-min hill-ciiming bayesian network structure learning algorithm. Machine learning, 65(1):3177, 2006.
[54] Philippe Brouillard, Sébastien Lachapelle, Alexandre Lacoste, Simon Lacoste-Julien, and Alexandre Drouin. Differentiable causal discovery from interventional data, arXiv preprint arXiv:2007.01754. 2020.
13

---

[55] Thomas Kipf, Yujia Li, Hanjun Dai, Vinicius Zambaldi, Alvaro Sanchez-Gonzalee Edward Grefenstette, Pushmeet Kohli, and Peter Battaglia. Compile: Compositional imitation learning and execution. In International Conference on Machine Learning, pages 3418-3422 PMLR, 2019.
[56] Danfei Xu, Roberto Martín-Martín, De-An Huang, Yuke Zhu, Silvio Savarese, and Li F Fei-Fei. Regression planning networks. Advances Neural Information Processing Systems, 32, 2019.
[57] De-An Huang, Suraj Nair, Danfei Xu, Yuke Zhu, Animesh Garg, Li Fei-Fei, Silvio Savarese, and Juan Carlos Niebles. Neural task graphs: Generalizing to unseen tasks from single video demonstration. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 8565 8574, 2019.
[58] Marta Garnelo, Kai Arulkumaran, and Murray Shanahan. Towards deep symbolic reinforcement learning. arXiv preprint arXiv: 2016.
[59] llya Kostrikov, Denis Yarats, and Rob Fergus. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. arXiv preprint arXiv:2004. 13649, 2020.
[60] Nicklas Hansen, Hao Su, and Xiaolong Wang. Stabilizing deep q-learning with convnets and vision transformers under data augmentation. Advances in Neural Information Processing Systems, 34, 2021.
[61] Aravind Srinivas, Michael Laskin, and Pieter Abbeel. Curl: Contrastive unsupervised represen- tations for reinforcement learning, arXiv preprint #rXiv:2004.04136 2020.
[62] Nicklas Hansen and Xiaolong Wang- Generalization reinforcement learning by soft data augmentation. In 202/ IEEE International Conference on Robotics and Automation (ICRA), pages 13611-13617. IEEE, 2021.
[63] Kimin Lee, Kibok Lee, Jinwoo Shin, and Honglak Lee. Network randomization: A simple technique for generalization in deep reinforcement learning. arXiv preprint arXiv: 190.0539, 2019.
[64] Roberta Raileamu, Maxwell Goldstein, Denis Yarats, Ilya Kostrikov, and Rob Fergus. Automatic data augmentation for generalization in reinforcement learning. Advances in Neural Information Processing Systems, 34, 2021.
[65] Josh Tobin, Rachel Fong, Alex Ray, Jonas Schneider, Wojeiech Zaremba, and Pieter Abbeel. Domain randomization for transferring deep neural networks from simulation the real world. In 2017 IEEE/RSJ international conference intelligent robots and systems (IROS), pages 23-30. IEEE, 2017.
[66] Bhairav Mehta, Manfred Diaz, Florian Golemo, Christopher J Pal, and Liam Paull. Active domain randomization. In Conference on Robot Learning, pages 1162-1176. PMLR, 2020.
[67] Aayush Prakash, Shaad Boochoon, Mark Brophy, David Acuna, Eric Cameracci, Gavriel State, Omer Shapira, and Stan Birchfield. Structured domain randomization: Bridging the reality gap by context-aware synthetic data. In 2019 International Conference on Robotics and Automation (ICRA), pages 7249-7255. IEEE, 2019.
[68] Sanmit Narvekar, Bei Peng, Matteo Leonetti, Jivko Sinapov, Matthew E Taylor, and Peter Stone. Curriculum learning for reinforcement learning domains: framework and survey. arXiv preprint arXiv:2003.0496/ 2020.
[69] Tom Schaul, Daniel Horgan, Karol Gregor, and David Silver. Universal value function ap- proximators. In International conference machine learning, pages 1312-1320. PMLR, 2015.
[70] Alexander Trott, Stephan Zheng, Caiming Xiong, and Richard Socher. Keeping your dis- tance: Solving sparse reward tasks using self-balancing shaped rewards. Advances Neural Information Processing Systems, 32, 2019.
14

---

[71] Suraj Nair. Silvio Savarese, and Chelsea Finn. Goal-aware prediction: Learning to model what matters. In International Conference on Machine Learning. pages 7207-7219. PMLR, 2020.
[72] Carlos Florensa, David Held, Xinyang Geng, and Pieter Abbeel. Automatic goal generation for reinforcement learning agents. In International conference on machine learning, pages 1515-1528. PMLR, 2018.
[73] Zhizhou Ren, Kefan Dong, Yuan Zhou, Qiang Liu, and Jian Peng. Exploration via hindsight goal generation. Advances in Neural Information Processing Systems, 32, 2019.
[74] Silviu Pitis, Harris Chan, Stephen Zhao, Bradly Stadie, and Jimmy Ba. Maximum entropy gain exploration for long horizon multi-goal reinforcement learning. In International Conference on Machine Learning, pages 7750-7761. PMLR, 2020.
[75] Meng Fang, Tianyi Zhou, Yali Du, Lei Han, and Zhengyou Zhang. Curriculum-guided hindsight experience replay. Advances in neural information processing systems, 32, 2019.
[76] Yunhao Tang and Alp Kucukelbir. Hindsight expectation maximization for gonl-cond reinforcement learning. In International Conference on Artificial Intelligence and Statistics, pages 2863-2871. PMLR, 2021.
[77] Tim GJ Rudner, Vitchyr Pong, Rowan McAllister, Yarin Gal, and Sergey Levine, Outcome- driven reinforcement learning via variational inference. Advances in Neural Information Processing Systems, 34, 2021.
[78] Prashan Madumal, Tim Miller, Liz Sonenberg, and Frank Vetere. Explainable reinforcement learning through a causal lens. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pages 2493-2500, 2020.
[79] Wenhao Ding, Haohong Lin, Bo Li, and Ding Zhao. Causalaf: Causal autoregressive flow for goal-directed safety-critical scenes generation. arXiv preprint arXiv:2110.13932 2021.
[80] Zizhao Wang, Xuesu Xiao, Zifan Xu, Yuke Zhu, and Peter Stone. Causal dynamics learning for task-independent state abstraction. arXiv preprint arXiv:2206. 13452, 2022.
[81] Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz. Invariant risk mini- mization, arXiv preprint arXiv: 1xXXiv:907.0000 0.289,0 2019.
[82] Clark Glymour. Kun Zhang, and Peter Spirtes. Review of causal discovery methods based on graphical models. Frontiers in genetics, 10:524, 2019.
[83] Karl Pearson. X. on the criterion that a given system of deviations from the probable in the case of correlated system of variables is such that it can be reasonably supposed have arisen from random sampling. The London, Edinburgh, and Dublin Philosophical Magazine and Journal of Science, 50(302):157-175, 1900.
[84] Kun Zhang, Jonas Peters, Dominik Janzing, and Bernhard Schölkopf. Kernel-based conditional independence test and application in causal discovery. arXiv preprint arXiv:1202.3375. 2012.
[85] Alain Hauser and Peter Bühlmann. Characterization and greedy learning of interventional markov equivalence classes of directed acyelic graphs. The Journal of Machine Learning Research, (311)22409 2464, 2012.
[86] Andrew A Neath and Joseph E Cavanaugh. The bayesian information criterion: background, derivation, and applications. Wiley Interdisciplinary Reviews: Computational Statisties, 192):199-203, 2012.
[87] Shengyu Zhu, Ignavier Ng, and Zhitang Chen. Causal discovery with reinforcement learning. arXiv preprint arXiv: /9067044777 2019.
[88] Nan Rosemary Ke, Olexa Bilaniuk, Anirudh Goyal, Stefan Bauer, Hugo Larochelle, Bernhard Schölkopf, Michael C Mozer, Chris Pal, and Yoshua Bengio. Learning neural causal models from unknown interventions. arXiv preprint 17000:1110 2019.
15

---

[89] Yunzhu Li, Antonio Torralba, Anima Anandkumar, Dieter Fox, and Animesh Garg. Causal discovery physical systems from videos. Advances in Neural Information Processing Systems, 33:9180-9192, 2020.
[90] Nino Scherrer, Olexa Bilaniuk, Yashas Annadani, Anirudh Goyal, Patrick Schwab, Bernhard Schölkopf, Michael Mozer, Yoshua Bengio, Stefan Bauer, and Nan Rosemary Ke. Learning neural causal models with active interventions. arXiv preprint arXiv:2109.02429, 2021.
[91] Erik Lindgren, Murat Kocaoglu, Alexandros G Dimakis, and Sriram Vishwanath. Experimental design for cost-aware learning of causal graphs. Advances in Neural Information Processing Systems, 31, 2018.
[92] Dar Gilboa, B. Chang, Minmin Chen, Greg Yang, Samuel S. Schoenholz, Ed H. Chi, and Jeffrey Pennington. Dynamical isometry and a mean field theory of Istms and grus. ArXiv, (bss/1901.0998/ 2019.
[93] Karren D. Yang, Abigail Katoff, and Caroline Uhler. Characterizing and learning equivalence classes of causal dags under interventions. In ICML, 2018.
[94] Raghavendra Addanki, Shiva Kasiviswanathan, Andrew Megregor, and Cameron Musco. Effi- cient intervention design for causal discovery with latents. Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 63-73. PMLR, 13-18 Jul 2020.
[95] Wenhao Ding, Chejian Xu, Haohong Lin, Bo Li, and Ding Zhao. A survey on safety-critical scenario generation from methodological perspective. arXiv preprint arXiv:2202.0222 02011000..222... 2022.
16

---

Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's
contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes] In the conclusion section.
(c) Did you discuss any potential negative societal impacts of your work? [Yes] In the Appendix.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [Yes]
(b) Did you include complete proofs of all theoretical results? [Yes] In the Appendix.
3. If you ran experiments..
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material or URL)? [Yes] In Supplemen- tary.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes] In the Appendix.
(c) Did you report error bars (e-g., with respect the random seed after running experi- ments multiple times)? [Yes] In the Appendix.
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)? [Yes] In the Appendix.
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) your work uses existing assets, did you cite the creators? [Yes] the Supplementary.
(b) Did you mention the license of the assets? [Yes] In the Supplementary.
(c) Did you include any new assets either the supplemental material or as URL? [Yes] In the Supplementary.
(d) Did you discuss whether and how consent obtained from people whose data you're using/curating? [Yes] In the Supplementary
(c) Did you discuss whether the data you using/curating contains personally identifiable information or offensive content? [Yes] the Supplementary.
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given participants and screenshots, if applicable? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent participant compensation? [N/A]
17


Reviews:
[{"review_id": "m89YdHwbca5", "reviewer": "Reviewer_qSKr", "summary": "The paper proposes to augment Goal-Conditioned RL (GCRL) with Causal Graph (CG), and treats the causal graph as a latent variable. \nThe paper then proposes to use  interventional data to estimate the posterior over the underlying causal graph, and then proposes to use the causal graph to learn factorized action conditioned model and agent's policy. \n\n\n=====\n\nAfter rebuttal: I've read the rebuttal, and I remain by my rating. ", "questions": "**Model learning**\n\n- The authors proposes to represent the transition model, with an ensemble of recurrent models, such that they can take into account variable number of \"causal parents\" for each node of the causal graph. Such a design choice is commonly used in causality papers (already cited by authors [1, 2], where the authors learn structural and functional parameters. Structural parameters corresponds to a soft-adjacency matrix, and functional parameters correspond to a set of MLPs which take in input a set of causal parents. This design choice has also been used as an inductive bias for learning architectures which consists of ensemble of recurrent models, such that each model can decide which other models to pay attention to via the use of attention [3, 4]. \n\n[1] Differentiable causal discovery from interventional data, https://arxiv.org/abs/2007.01754  \\\n[2] Learning neural causal models from unknown interventions, https://arxiv.org/abs/1910.01075 \\\n[3] Neural Production Systems, https://arxiv.org/abs/2103.01937  \\\n[4] Recurrent Independent Mechanisms, https://arxiv.org/abs/1909.10893 \n\n**Chosen Method for  learning causal graph**\n\nIt may be useful to employ gradient based methods of learning causal graph as compared to methods based on independence testing. Recently, few methods have been proposed which allows inferring posterior over the learned causal graph [1, 2, 3].\n\n[1] DiBS: Differentiable Bayesian Structure Learning, https://arxiv.org/abs/2105.11839 \\\n[2] BCDNets, https://proceedings.neurips.cc/paper/2021/hash/39799c18791e8d7eb29704fc5bc04ac8-Abstract.html \\\n[3] Bayesian Structure Learning with Generative Flow Networks, https://arxiv.org/abs/2202.13903 \n\n\n**RL benchmarks for causal reasoning**\n\nIt may be useful to evaluate the proposed method on set of problems which are specifically tailored towards measuring generalization in RL (in causal sense) [1]. It may also be useful to evaluate an \"attention\" (or GNN) based baseline which learns interactions between different causal variables to predict the future state for planning. \n\n[1] Systematic Evaluation of Causal Discovery in Visual Model Based Reinforcement Learning, https://arxiv.org/abs/2107.00848 \n\n", "limitations": "**Limitations**\n\nThe paper discusses two limitations: (a) the method may not scale well to large number of nodes, (b) assuming factorized state and action space restricts the scope to scenarios where only observables are low level information i.e., pixels in the case of image data. It may also be useful to discuss that **independent testing** i.e., the method used to learn causal graph may not scale well as compared to gradient based approaches of learning causal graph.\n\n\n", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Clarify**\n\n- The paper is easy to read and well written. \n\n**Originality**\n\n- The individual elements are known in causality literature, and various other works (most of them are cited by the authors). The contribution lies in using these ideas (for ex. planning with a factorial action conditioned model) in the context of goal conditioned policies. \n\n**Significance**\n\n- The reviewer like the experiments where the authors focussed on using the transition models for doing planning. The reviewer believes it's a step in the right direction. \n\n**Quality**\n\n- The reviewer likes the experimental setup as it's tailored to test generalization (robustness to spurious correlations and compositional generalization).\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "e1X3bs0VeGX", "reviewer": "Reviewer_7kq2", "summary": "Paper investigates the problem of model-based RL or planning. It proposes to decompose the transition model into a causal graph and some parameterized functions. The causal discovery here is formulated as learning a latent variable model using VI, where the causal graph is the latent variable. Then the transition loss is used to find the best set of parameters. Experiments on several synthetic domains including one designated for causality show some promises over state-of-the-art model-based RL avenues.  ", "questions": "See [Weaknesses]\n", "limitations": "The authors mentioned that their limitation could be the inability of scaling to causal graphs with more nodes. However, I also doubt whether this approach could be applied to real-world domains with non-symbolic input, ex. pixel-based or other raw sensory input. Extracting symbolic input from these sources could impose new challenges on observation uncertainty, which cold be pivotal to causal discovery and model learning.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "[Stengths]\n\n+Paper is overall clear and well-written. The authors also provide the code and indeed it helps with understanding the idea here.\n\n+I found the formulation intuitive and reasonable. The math on identifiability and performance guarantee could be a bit dense but I didn't find any major issue.   \n\n+The advantages over MBRL baselines are impressive, especially in the challenging car crash domain. With the code already there, I'm sure the results are reproducible if the authors can release it at some point.\n\n[Weaknesses]\n\nI don't have major concerns at this point. But I will appreciate it if the following questions can be addressed in a rebuttal.\n\n-The authors claim that their approach is for \"goal-based RL\". But it seems that goal is not necessary for all their formulations and implementations. Why not just call it model-based RL or planning in general? Some clarifications here are needed.\n\n-The causal discovery part might be clear in the formulation but the implementation details are yet to be unveiled. I also check the code, but as far as I could find the causal graphs are **pre-defined**. The authors are encouraged to elaborate more on this and some pointers to the code will be helpful.\n\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "HulbKN77VGr", "reviewer": "Reviewer_zrjV", "summary": "Summary: the paper proposes an algorithm for learning causl strutures from the environment and using it to act and plan. The papers presents experiments on OOD generalization on either spurious correlations or compositional generalization.", "questions": "- It's unclear if the underlying causal structure of the environment for the stacking and door opening experiments are \"interesting\", in the sense that changing the color or backgrounds of the object has no causal influence on anything else in the env. \n\n- Can't the spurious correlation experiment be solved by only focusing on learning what is a causal variables (compared to learning the causal mechanisms).\n\n", "limitations": "One of the major limitations of this paper is that the theoretical results the paper presents is based on somewhat unrealistic assumptions. However, this is not discussed properly in the paper.\n\n One of the main claims of the paper is the identifiability proof in that the algorithm can always identify the true graph. However, this proof is based on assumptions that generally may not hold true, espiecally in high-dimensional data. Another major reason why causal identifiability is so challenging in low-level data is that the causal variables are often latent, making the interventions unknown in the latent space. this is not discussed in the paper.\n\nHowever, the proof that the authors presents in the paper assumes 2 things, one being that the mapping from the low-level observational space to latent causal variables are trivial and can be easily obtained, and second, that the independence test has perfect information, in that it can tell perfectly which variable cases which other variable upon intervention.\n\nI think both of these assumptions generally do not hold true. Discovering latent states is non-trivial, and Im not sure if faithfulness generally holds either. \n\nAnother point that the authors have not discussed is that the interventions are unknown. The intervention is only observed on low-level images, but it's unknown to the abstract (latent) causal variables. I do not see a discussion on this anywhere in the paper.", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Pros:\n\n- Causal models hold promise to better generalization, especially in RL. Causal models focuses on the right causal variables as well as causal mechanisms between variables.  Hence, this paper addresses an very important problem, espeically in RL.\n\n- The experimental results are interesting, especially the OOD generalization experiments. \n\nCons:\n\nThere are some limitations of the work that would be nice to be addressed.\n\n- One of the main claims of the paper is the identifiability proof in that the algorithm can always identify the true graph. However, this proof is based on assumptions that generally may not hold true, espiecally in high-dimensional data. Another major reason why causal identifiability is so challenging in low-level data is that the causal variables are often latent, making the interventions unknown in the latent space. this is not discussed in the paper.\n\nHowever, the proof that the authors presents in the paper assumes 2 things, one being that the mapping from the low-level observational space to latent causal variables are trivial and can be easily obtained, and second, that the independence test has perfect information, in that it can tell perfectly which variable cases which other variable upon intervention.\n\nI think both of these assumptions generally do not hold true. Discovering latent states is non-trivial, and Im not sure if faithfulness generally holds either. \n\nAnother point that the authors have not discussed is that the interventions are unknown. The intervention is only observed on low-level images, but it's unknown to the abstract (latent) causal variables. I do not see a discussion on this anywhere in the paper.\n\n- The experiments are interesting. However, the spurious correlation experiments can be solved without learning the underlying causal mechanisms. The model would only need to discover the causal variables (but not necessarily mechanisms). The compositional generalization tasks that the authors proposed is more interesting, but again, the agent doesnt  necessarily have to know the undelrying causal graph in order to generalize. For example, stacking 2 objects and stacking 5 objects as proposed by the authors. A model that has not learned the underlying causal structure of the environment might also be able to generalize to this.\n\n- It's also unclear if the underlying causal structure of the environment for the stacking and door opening experiments are \"interesting\", in the sense that changing the color or backgrounds of the object has no causal influence on anything else in the env. \n\n- In order to test to see if the model has indeed captured the underlying causal structure, I would propose that the authors use either \n\n1. \"Alchemy: A structured task distribution for meta-reinforcement learning\"\n2. \"Systematic Evaluation of Causal Discovery in Visual Model Based Reinforcement Learning\"\n\nBoth of these environments have more complicated and interesting causal graphs. There is also a set of interesting OOD experiments that can be performed on these experiments.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept