Generalizing Goal-Conditioned Reinforcement Learning with Variational Causal Reasoning
Anonymous Author(s)
AfﬁliationAddress email
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 anddiscovering cause-and-effect relations.
However, how to discover and representcausalities remains a huge gap that hinders the development of causal RL.
In thispaper, we augment Goal-Conditioned RL (GCRL) with Causal Graph (CG) , astructure built upon the relation between objects and events.
We novelly formulatethe GCRL problem into variational likelihood maximization with CG as latent variables.
To optimize the derived objective, we propose a framework with theoreticalperformance guarantees that alternates between two steps: using interventionaldata to estimate the posterior of CG; using CG to learn generalizable models andinterpretable policies.
Due to the lack of public benchmarks that verify generalization capability under reasoning, we design nine tasks and then empirically show theeffectiveness of the proposed method against ﬁve baselines on these tasks.
Furthertheoretical analysis shows that our performance improvement is attributed to thevirtuous cycle of causal discovery, transition modeling, and policy training, whichaligns with the experimental evidence in extensive ablation studies.
Introduction
The generalizability, which enables an algorithm to handle unseen tasks, is fruitful yet challenging inmultifarious decision-making domains.
Recent literature [1, 2, 3] reveals the critical role of reasoningin improving the generalization of reinforcement learning (RL).
However, most off-the-shelf RLalgorithms [4] have not regarded reasoning as an indispensable accessory, thus usually suffer fromdata inefﬁciency and performance degradation due to the mismatch between training and testingsettings.
To attain generalization at the testing stage, some efforts were put into incorporatingdomain knowledge to learn structured information, including sub-task decomposition [5] or programgeneration [6, 7, 8] and execution [9, 10], which guide the model to solve complicated tasks in anexplainable way.
However, such symbolism-dominant methods heavily depend on the re-usability ofsub-tasks and pre-deﬁned 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 recentlyincorporated to compactly represent the aforementioned structured knowledge in RL training [11].
Based on the form of causal knowledge, we divide the related works into two categories, i.e., implicitand explicit causation.
With implicit causal representation, researchers ignore the detailed causalstructure.
For instance, [12] extracts invariant features as one node that inﬂuences the reward function,while the other node consists of task-irrelevant features [13, 14, 15, 16].
This neat structure has goodscalability but requires access to multiple environments that share the same invariant feature [12, 16,17].
One can also turn to the explicit side by estimating detailed causal structures [18, 19, 20, 21], which uses directed graphical models to capture the causality in the environment.
A pre-request forthis estimation is the object-level or event-level abstraction of the observation, which is available inmost tasks and also becoming a frequently studied problem [22, 23, 24].
However, existing explicitcausal reasoning RL models either rely on heuristic design without theoretical guarantees [18] orrequire the true causal graph [25].
In this paper, we propose G ene RA lizing by D iscov ER ing (GRADER) , a causal reasoning method thataugments the RL algorithm with data efﬁciency, interpretability, and generalizability.
We mainlyfocus on Goal-Conditioned RL (GCRL), where different goal distributions during training and testingreﬂect the generalization.
We formulate the GCRL into a probabilistic inference problem [26] with alearnable causal graph as the latent variable.
This novel formulation naturally explains the learningobjective with three components, i.e., transition model learning, planning, and causal graph discovery,leading to an optimization framework that alternates between the causal structure discovery policylearning to gain generalizability.
Under some mild conditions, we prove the unique identiﬁability ofthe causal graph and the theoretical performance guarantee of the proposed framework.
To demonstrate the effectiveness of the proposed method, we conduct comprehensive experimentsin environments that require strong reasoning capability.
Speciﬁcally, we design two types ofgeneralization settings, i.e., spuriousness and composition, and provide an example to illustrate thesesettings in Figure 1.
The evaluation results conﬁrm our method’s advantages in several aspects.
First,the proposed data-efﬁcient discovery method provides an explainable causal graph yet requires muchless data than score-based methods, increasing data efﬁciency and interpretability during task solving.
Second, simultaneously discovering the causal graph during policy learning dramatically increasesthe success rate of solving tasks.
In summary, the contribution of this paper is threefold:
• We use the causal graph as a latent variable to reformulate the GCRL problem and then derive
an iterative training framework from solving this problem.
• We prove that our method uniquely identiﬁes true causal graphs, and the performance of theiterative optimization is guaranteed with a lower bound given converged transition dynamics.• We design nine tasks in three environments that require strong reasoning capability and showthe effectiveness of the proposed method against ﬁve baselines on these tasks.
Problem Formulation and Preliminary
We start by discussing the setting of RL we use in this paper and the assumptions required in causalreasoning.
Then we brieﬂy introduce the necessary concepts related to causality and causal discovery.
2.1 Factorized Goal-conditioned RL
We assume the environment follows the Goal-conditioned Markov Decision Process (MDP) settingwith full observation.
This setting is represented by a tuple M =
( S , A , P , R , G ) , where S is 70 the state space,
A is the action space, P is the probabilistic transition model, G ⇢ S is the goal 71 space which is a set of assignment of values to states, and r ( s, g )
=
1 ( s = g ) 2 R is the sparse 72 deterministic reward function that returns 1 only if the state s match the goal g .
In this paper, wefocus on the goal-conditioned generalization problem, where the goal for training and testing stageswill be sampled from different distributions p train ( g ) and p test ( g ) .
We refer to a goal g 2 G as a taskand use these two terms interchangeably.
To accomplish the causal discovery methods, we make afurther assumption similar to [20, 27] for the state and action space:
Assumption 1 (Space Factorization) .
The state space S =
{S 1 ⇥ · · · ⇥ S M } and action space
A = {A 1 ⇥ · · · ⇥ A N } can be factorized to disjoint components {S i } Mi =1 and {A i } Ni =1 .
The components representing one event or object’s property usually have explicit semantic meaningsfor better interpretability.
This assumption can be satisﬁed by state and action abstraction, which hasbeen widely investigated in [22, 23, 24].
Such factorization also helps deal with the high-dimensionalstates since it might be impossible to treat each dimension as one random variable [18].
2.2 Causal Reasoning with Graphical Models
Reasoning with causality relies on speciﬁc causal structures, which is commonly represented asdirected acyclic graphs (DAGs) [28] over variables.
Consider random variables
X
=
( X 1 , . . .
, X d )with index set V := { 1 , . . .
, d } .
A graph G =
( V , E ) consists of nodes V and edges E ✓ V 2 with 87 ( v, v ) for any v 2 V .
A node i is called a parent of j if e ij 2 E and e ji / 2 E .
The set of parents of jis denoted by PA G j .
We formally discuss the graph representation of causality with the deﬁnitions:
Deﬁnition 1 (Structural Causal Models [28]) .
A structural causal model (SCM) C :
= ( S , U ) consistsof a collection S of d functions,
[Equation]
where PA j ⇢ { X 1 , . . . , X d }\{ X j } are called parents of X j ; and a joint distribution U
=
{ U 1 , . . .
, U d } over the noise variables, which are required to be jointly independent.
Deﬁnition 2 (Causal Graph [28], CG) .
The causal graph G of a SCM is obtained by creating onenode for each X j and drawing directed edges from each parent in PA G j to X j .
We note that CG describes the structure of the causality, and SCM further considers the speciﬁccausation from the parents of X j to X j as well as exogenous noises U j .
To uncover the causalstructure from data distribution, we assume that the CG satisﬁes the Markov Property and Faithfulness([28]), which makes the independences consistent between the joint distribution P ( X 1 , . . .
, X n ) andthe graph G .
We also follow the Causal Sufﬁciency assumption [29] that supposes we have measuredall the common causes of the measured variables.
Existing work [30, 20] believes that two objects have causation only if they are close enough whilethere is no edge between them if the distance is large.
Instead of using such a local view of thecausality, we assume the causal graph is consistent across all time steps.
Still the causal inﬂuencecould be local, and the speciﬁc inﬂuence indicated by edges is estimated by the function f j ( PA j , U j ) .
Generalizing by Discovering (GRADER)
[Figure]
With proper deﬁnitions and assumptions, we
[Figure]
now look at the proposed method.
We ﬁrst derive the framework of GRADER by formulatingthe GCRL problem as a latent variable model,which provides a variational lower bound to optimize.
Then, we divide this objective functioninto three parts and iteratively update them witha speciﬁc performance guarantee.
3.1 GCRL as Latent Variable Models
[Figure]
The general objective of RL is to maximizeexpected reward function w.r.t.
a learnablepolicy model ⇡ .
Particularly, in the goal-conditioned setting, such objective is represented as
[Equation] [Equation]
P
Tt =
r
( s t , g )] , where p ( g ) is the distribution of goals and ⌧ := { s 0 , a 0 , . . .
, s T }the action-state trajectory with maximal time step T .
The trajectory ends only if the goal is achievedg =
s T or
Inspired the question from
"How to ﬁnd actions to achieve the goal?" to "what are the actions if we achieve the goal?" , leadingto a likelihood maximization problem for p ( ⌧ | s ⇤ ) with s ⇤ :
= 1 ( g = s T ) .
Different from previous work [32] that recasts actions as latent variables and inferring actions that result in "observed" highreward, we decompose p ( ⌧ | s ⇤ ) with G as the latent variable to get the evidence lower bound (ELBO)
[Equation]
where the prior p ( G ) and variational posterior q ( G| ⌧ ) represent distributions over a graph structure,i.e, the probability of the existences of edges in the graph.
D KL represents the Kullback–Leibler (KL)divergence between two graphs, which will be explained in Section 3.3.
Recall that the space of goal
G is a subset of the state, here we extend the meaning of g and assume all trajectories achieve thegoal in the ﬁnal state [33], i.e., g = s T .
Such extension makes it possible to further decompose theﬁrst term of (2) as (refer to Appendix A.1.2)
[Equation]
Here, we use the fact that G is effective in both the transition model p ( s t +1 | a t , s t , G ) and thepolicy model log p ( a t | s t , s ⇤ , G ) , g only inﬂuences the policy model, and the initial state s 0 dependsneither on G nor g .
We assume that both initial state log p ( s 0 ) and goal log p ( g ) follow the uniformdistribution; thus, the ﬁrst and last terms of (3) are constants.
The policy term ⇡ , selecting action a taccording to both current state s t and goal g , is implemented with planning method and is furtherdiscussed in Section 3.2.
Finally, we maximize the likelihood p ( ⌧
| s ⇤ ) with the following reformulated
ELBO as the objective
[Equation]
where ✓ is the shared parameter of transition model p ✓ ( s t +1 | a t , s t , G ) and policy ⇡ ✓ ( a t | s t , s ⇤ , G ) ,and  is the parameter of causal graph q  ( G| ⌧ ) .
To efﬁciently solve this optimization problem, weiteratively updates parameter  (causal discovery, Section 3.3) and parameter ✓ (model and policylearning, Section 3.2), as shown in Figure 2.
Intuitively, these processes can be respectively viewedas the discovery of CG and update of SCM, which share tight connections discussed in Section 2.2.
3.2 Model and Policy Learning
Let us start with a simple case where we already obtain a G and use it to guide the learning ofparameter ✓ via max ✓ J ( ✓ ,  ) .
Since the KL divergence of J ( ✓ ,  ) does not involve ✓ , we only needto deal with the ﬁrst expectation term, the likelihood of the transition model and policy.
For thetransition p ✓ ( s t +1 | a t , s t , G ) , we connect it with causal structure by further deﬁning a particular typeof CG and denote it as G in the rest of this paper:
Deﬁnition 3 (Transition Causal Graph) .
We deﬁne a bipartite graph G , whose vertices are dividedinto two disjoint sets U =
{A t , S t } and V = {S t +1 } .
A t represents action nodes at step t , S state 152 nodes at step t , and
S t +1 the state nodes at step t + 1 .
All edges start from set U and end in set V .
Model learning.
This deﬁnition builds the causal graph between two consecutive time steps, whichindicates that the values of states in step t + 1 depend on values in step t .
It also implies that theinterventions [28] on nodes in U are directly obtained since they have no parents.
We denote themarginal distribution of S as p I s ⇡ , which is collected by RL policy ⇡ .
Combined with the Deﬁnition 1of SCM, we ﬁnd that p ✓ ( s t +1 | a t , s t , G ) essentially approximates a collection of functions f j followingthe structure G , which take as input the values of parents of the state node s j and outputs the value s j .
Thus, we propose to model the transition corresponding to G with a collection of neural networksf ✓ ( G ) :
= { f ✓ j } Mj =1 to obtain
[Equation]
where [ PA G j ] t represents the values of all parents of node s t j at time step t and U j follows Gaussiannoise U
j ⇠ N (0 , I ) .
In practice, we use Gated Recurrent Unit [34] as f j because it supports varyingnumbers of input nodes.
We take s tj as the initial hidden embedding and the rest parents [ PA G j \ s j ] tas the input sequence to f j .
The entire model is optimized by stochastic gradient descent with thelog-likelihood log p ✓ ( s t +1 | a t , s t , G ) as objective.
Policy learning with planning.
Then we turn to the policy term ⇡ ✓ ( a t | s t , s ⇤ , G ) in J ( ✓ ,  ) .
Weoptimize it with planning methods that leverage the estimated transition model.
Speciﬁcally, the policyaims to optimize an action-state value function Q ( s t , a t ) =
E hP
Ht
=
t 0 r ( s t 0 + t , a t 0 + t ) | s t , a t i ,which can be obtained by unrolling the transition model with a horizon of H steps and discountfactor  .
In practice, we use model predictive control (MPC) [35] with random shooting [36], whichselects the ﬁrst action in the ﬁxed-horizon trajectory that has the highest action-state value Q ( s t , a t ) ,i.e. ˆ ⇡ ( s t ) =
arg max a t 2 A Q G ✓ ( s t , a t ) .
The formulation we derived so far is highly correlated to themodel-based RL framework [37].
However, the main difference is that we obtain it with variationalinference by regarding the causal graph as a latent variable.
3.3 Data-Efﬁcient Causal Discovery
Input: Trajectory buffer B ⌧ , Causal graph
G , Transition model f ✓ while ✓ not converged do // Policy from planning Sample a goal g ⇠ p train ( g ) while t < T do a t
Planner( f ✓ , s t , g ) s t +1 , r t
Env( a t , g )
B ⌧
B ⌧ [ { a t , s t , s t +1 } //
Estimate causal graph for i  M + N do for j  M do Infer edge e ij q  ( ·|B , ⌘ ) //
Learn transition model
Update f ✓ ( G ) via (5) with B
In this step, we relax the assumption of knowing Gand aim to estimate the posterior distribution q ( G| ⌧ ) 178 to increase the ELBO (4) w.r.t.
parameter  .
In mostscore-based methods [38], likelihood is used to evaluatethe correctness of the causal graph, i.e., a better causalgraph leads to a higher likelihood.
Since the ﬁrst termof (4) represents the likelihood of the transition model,we convert the problem of max  J ( ✓ ,  ) to the causaldiscovery that ﬁnds the true causal graph based oncollected data samples.
As for the second term of (4),the following proposition shows that the KL divergencebetween q  ( G| ⌧ ) and p ( G ) can be approximated by a 188 sparsity regularization (proof in Appendix A.4.2).
Proposition 1 (KL Divergence as Sparsity Regularization) .
With entry-wise independent Bernoulli priorp ( G ) and point mass variational distribution q ( G| ⌧ ) of 192 DAGs,
D KL [ q  k p ] is equivalent to an ` 1 sparsity regularization for the discovered causal graph.
We restrict the posterior q  ( G| ⌧ ) to point mass distribution and use a threshold ⌘ to control the sparsity.
We perform the discovery process from the classiﬁcation perspective by proposing binary classiﬁersq  ( e ij | ⌧ , ⌘ ) to determine the existence of an edge e ij .
This classiﬁer q  ( e ij | ⌧ , ⌘ ) is implemented bystatistic Independent Test [39], and ⌘ is the threshold for the p-value of the hypothesis.
A larger ⌘corresponds to harder sparsity constraints, thus leading to a sparse G since two nodes are more likelyto be considered independent.
According to the deﬁnition 3, we only need to conduct classiﬁcation toedges connecting nodes between U and V .
If two nodes are dependent, we add one edge directed fromthe node in U to the node in V .
This deﬁnition also ensures that we always have q ( G| ⌧ ) 2 Q DAG , 202 where
Q DAG is the class of DAG.
With this procedure, we identify a unique CG G ⇤ under optimality:
Proposition 2 (Identiﬁability) .
Given an oracle independent test, with an optimal interventional datadistribution p
⇤ I s ⇡ , causal discovery obtains  ⇤ that correctly tells the independence between any twonodes, then the causal graph is uniquely identiﬁable, with e ⇤ ij = q  ⇤ ( e ij | ⌧ ) , 8 i 2 [ M + N ] , j 2 [ M ] .
In practice, we use  2 -test for discrete variables and the Fast Conditional Independent Test [39]for continuous variables.
The sample size needed for obtaining the oracle test has been widelyinvestigated [40].
However, testing with ﬁnite data is not a trivial problem, as stated in [41], especiallywhen the data is sampled from a Goal-conditioned MDP.
Usually, the random policy is not enoughto satisfy the oracle assumption because some nodes cannot be fully explored when the task iscomplicated and has a long horizon.
To make this assumption empirically possible, it is necessary tosimultaneously optimize ⇡ ✓ ( a t
| s t , s ⇤ , G ) to access more samples close to ﬁnishing the task, which isfurther analyzed in Section 3.4.
We also empirically support this argument in Section 4.2 and providea detailed theoretical proof in Appendix A.3 and A.4.
3.4 Analysis of Performance Guarantee
The entire pipeline of GRADER is summarized in Algorithm 1.
To analyze the performance of theoptimization of (4), we ﬁrst list important lemmas that connect previous steps and then show that theiteration of these steps in GRADER leads to a virtuous cycle.
By the following lemmas, we show the following performance guarantees step by step.
Lemma 1shows model learning is monotonically better at convergence given a better causal graph from causaldiscovery.
Then the learned transition model helps to improve the lower bound of value functionduring planning according to Lemma 2.
Lemma 3 reveals the connection between policy learningand interventional data distribution, which in turn improves the quality of our causal discovery, as isshown in Lemma 4 and Proposition 2.
Lemma 1 (Monotonicity of Transition Likelihood) .
Assume G
⇤
=
( V, E ⇤ ) be the true CG, for two
CG G 1 = ( V, E 1 ) and G 2 =
( V, E 2 ) , if SHD ( G 1 , G ⇤ ) <
SHD ( G 2 , G ⇤ ) , 9 e, s.t.
E 1 [ { e }
= E 2 , 227 when transition model ✓ converges, the following inequality holds for the transition model in (4):
[Equation]
where SHD is the Structural Hamming Distance deﬁned in Appendix A.2.
Lemma 2 (Bounded Value Function in Policy Learning) .
Given a planning horizon H ! 1 , ifwe already have an approximate transition model D TV (ˆ p ( s 0 | s, a ) , p ( s 0 | s, a ))  ✏ m , the approximatepolicy ˆ ⇡ achieves a near-optimal value function (refer to Appendix A.3 for detailed analysis):
[Equation]
Lemma 3 (Policy Learning Improves Interventional Data Distribution) .
With a step reward r ( s, a ) =
p g 1 ( s = g ) , we show that the value function determines an upper bound for TV divergence betweenthe interventional distribution with its optimal (proof details in Appendix A.3):
[Equation]
where p I s ⇡ is the marginal state distribution in interventional data, and p g is the goal distribution.
Abetter policy with larger V ⇡ ( s ) enforces the distribution of interventional data towards the goal.
Lemma 4 (Interventional Data Beneﬁts Causal Discovery) .
For ✏ g = min p g > 0 p g , D TV ( p I s ⇡ , p g ) <✏ g , the error of our causal discovery is upper bounded with E ˆ G
[ SHD ( ˆ G , G ⇤ )]
 |S|  1 .
After the close-loop analysis of our model, we are now able to analyze the overall performance of theproposed framework.
Under the construction of p ✓ ( s t +1 | a t , s t , G ) with NN-parameterized functions,the following theorem shows that the learning process will guarantee to perform a close estimation oftrue ELBO under the iterative optimization among model learning, planning, and causal discovery.
Theorem 1.
With T-step approximate transition dynamics D TV [Equation] p [Equation]
s
[Equation]
s, [Equation] , p [-( s 
-][Equation]
s
[Equation] s [Equation] ) ⌘  ✏ m , if thegoal distribution satisﬁes ✏ g >  1   ✏ m , and the distribution prior CG is entry-wise independent
Bernoulli ( ✏ G ) , GRADER guarantees to achieve an approximate ELBO ˆ J with the true ELBO J ⇤ :
[Equation]
An intuitive understanding of the performance guarantee is that a better transition model indicatesa better approximation of objective J .
Yet it’s harder for GRADER to approximate a longer trajectory with large state space.
The proof of this theorem and corresponding empirical results are in
Appendix A.5.
Experiments
In this section, we ﬁrst discuss the setting ofour designed environments as well as the baselines used in the experiments.
Then, we providethe numerical results and detailed discussions toanswer the following important research questions: Q1.
Compared to baselines, how does
GRADER gain performance improvement underboth in-distribution and generalization settings?
Q2.
Compared to an ofﬂine random policy, howdoes a well-trained policy improve the results of causal discovery?
Q3.
Compared to score-basedcausal discovery, does the proposed data-efﬁcient causal discovery pipeline guarantee to identifyingtrue causal graph as stated in Section 3.3?
Q4.
Compared to false causal graphs, how does thecorrectness of discovered causal graph inﬂuence the task-solving performance of RL agents?
[Table]
4.1 Environments and Baselines
Since most commonly used RL benchmarks do not explicitly require causal reasoning for generalization, we design three environments shown in Figure 3.
These environments use the true state asobservation to disentangle the reasoning task from visual understanding.
For each environment, wedesign three settings – in-distribution ( I ), spuriousness ( S ), and composition ( C ) – corresponding todifferent goal distributions for generalization.
We use p train ( g ) and p test ( g ) to represent the goal distribution during training and testing, respectively.
I uses the same p train ( g ) and p test ( g ) , S introducesspurious correlations in p train ( g ) but remove them in p test ( g ) , and C contains more similar sub-goalsin p test ( g ) than in p train ( g ) .
The details of these settings in three environments are brieﬂy summarizedin the following (details in in Appendix C.2.1):
• Stack : We design this manipulation task inspired by the CausalWorld [42], where the agent muststack objects to match speciﬁc shapes and colors.
In Stack-S , we let the same shape have the same color in p train ( g ) but randomly sample the color and shape in p test ( g ) .
In Stack-C , the maximumnumber of object is two in p train ( g ) but ﬁve in p test ( g ) .•
Unlock : We design this task for the agent to collect a key to open doors built upon the Minigrid [43].
In Stack-S , the door and the key are always in the same row in p train ( g ) but uniformlysample in p test ( g ) .
In Unlock-C , there are one door in p train ( g ) but two doors in p test ( g ) .•
Crash : 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 [44].
We design such a crashscenario based on highway-env [45], where the goals are to create crashes between a pedestrianand different AVs.
In Stack-S , the initial distance between ego car and pedestrian is a constant inp train ( g ) but irrelevant in p test ( g ) .
In Stack-C , there is one pedestrian in p train ( g ) but two in p test ( g ) .
We use the following methods as our baselines to fairly demonstrate the advantages of GRADER.
SAC: [46] Soft Actor-Critic is a well-known model-free RL method; we augment it with HER [47] inStack because the goal in this task does not have pre-conditions.
ICIN:
[25] It uses DAgger [48] to learn goal-conditioned policy with the causal graph estimated from the expert policy.
We assume it canaccess the true causal graph for supervised learning.
PETS:
[49] We consider the ensemble transition model with random shoot planning as one baseline, which achieves generalization with the uncertaintyaware design.
TICSA:
[18] This is a causal-augmented MBRL method that simultaneously optimizesa soft adjacent matrix (representing the causality) and a transition model.
ICIL:
[16] This methodproposes 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 imitationlearning.
GNN: [50] Since graph neural network (GNN) is 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 andsummarize the results in Table 1 to demonstrate the overall performance.
Generally, our methodoutperforms other baselines in all tasks except Stack-I , which is simple for all methods.
We note thatthe gap between our method and baselines under S and C settings is more signiﬁcant than the I setting,showing that our method still works well in the non-trivial generalization task.
As a model-free method, SAC fails in all three tasks of Unlock and Crash environments since they have 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 explicit or implicit causality of the task,so they are relatively better than SAC and PETS.
However, they are still worse than GRADER in two generalization settings because of the unstable and inefﬁcient causal reasoning mechanism.
We alsoﬁnd that even if the ICIN is given the true causal graph, the policy learning part cannot efﬁcientlyleverage 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 4.
Our method quickly learns to solve the tasks at the beginning of the training, demonstrating dataefﬁciency.
GRADER also outperforms other methods with large gaps in the later training phase.
Theresults of the other two environments can be found in Appendix B.1.
Importance of Policy Learning (Q2)
As we mentioned in Section 3.3, we empirically compare
GRADER and Ofﬂine [51], which uses data from ofﬂine random policy, and plot results in the rightpart of Figure 4.
We use SHD [52] to compute the distance between the estimated causal graph andthe true causal graph, which is obtained by checking the best causal graph with human analysis.
Thetrue causal graph for each environment can be found in Appendix B.2.
When we only use samplescollected ofﬂine by random policy, we cannot obtain variables’ values that require long-horizonreasoning, e.g., the door can be opened if the agent is close to the door and have the key, leadingto unreliable discovery.
As a consequence, the causal graph obtained by Ofﬂine indeed harms theperformance, as shown in Figure 4.
Instead, GRADER gradually explores more regions and quicklyobtains the true causal graph when we iteratively discover the causal graph and update the policy.
Advantage of Data-efﬁcient Causal Discovery (Q3)
To showthe advantage of proposed constraint-based methods, we designa model named Score that optimizes a soft adjacent matrixsimilar to the score-based method in the discovery step.
Scorebased discovery methods, also used in TICSA, are recentlycombined with NN for differentiable training.
According to thediscovery accuracy shown in the right part of Figure 4, we ﬁndscore-based discovery is inefﬁcient.
Based on the performanceof the Score model summarized in Table 1, we also concludethat it is not as good as our constraint-based method and has alarge variance due to the unstable learning of causal graph.
Inﬂuence of Causal Graph (Q4)
To illustrate the importanceof the causal graph, we implement another variant of GRADERnamed Full , which uses a ﬁxed full graph that connects allnodes between the sets U and V .
According to the performance 340 shown in Table 1 and Figure 4, we ﬁnd that the full graph achieves worse results than GRADERbecause of the redundant and spurious correlation.
Intuitively, unrelated information causes additionalnoises to the learning procedure, and the spurious correlation creates a shortcut that makes the modelextract wrong features, leading to worse results in the spuriousness generalization as shown in Table 1.
We then investigate how the correctness of the causal graph inﬂuences the performance.
We useﬁxed graphs interpolating from the best causal graph to the full graph to train a GRADER model on
Unlock-S and summarize the results in Figure
The more correct the graph is, the higher reward theagent obtains, which supports our statements in Section 3.4 that the causal graph is important for thereasoning tasks – a better causal graph helps the model have better task-solving performance.
Related Works
RL Generalization From agent’s view, algorithms focus on actively obtaining the structure of thetask or environment.
Some decompose the given task into sub-tasks [53, 54, 55, 5] and when theyencounter an unseen task, they rearrange those sub-tasks to solve new tasks.
Instead of dividing tasks, Symbolic RL learns a program-based policy consisting of domain language [6, 8] or grammar [7, 56].
The generated program is then executed [9] to interact with the environment, which has the potential tosolve unseen tasks by searching the symbolic space.
From the other view, we generate environments toaugment agents’ experience for better generalization.
One straightforward way is data augmentation ofimage-based observation [57, 58, 59, 60, 61, 62].
When extended to other factors of the environment, Domain Randomization [63] and its variants [64, 65] are proposed.
Goal-Conditioned RL (GCRL)
From the optimization perspective, Universal Value Function [66],reward shaping [67], and latent dynamic model [68] are widely used tools to solve GCRL problem.
Sub-goal generation [69] is another intuitive idea to tackle the long horizon with sparse reward, wherethe core thing is to make sure the generated sub-goals are solvable.
Finally, HER [33], belonging tothe relabelling category, is a ground-breaking yet straightforward method that treats visited statesas "fake" goals when the goal and state share the same space.
Later on, improved versions of
HER [47, 70, 71] were widely studied.
One limitation is that we cannot directly use a visited state asa goal if the goal has pre-conditions.
Similar to our setting, [72] and [73] convert the GCRL problemto variational inference by regarding control as inference [26].
[72] propose a EM framework underthe HER setting and [73] treats the last state as goal and estimate a shaped reward during training.
RL with Causal Reasoning Causality is now frequently discussed in the machine learning ﬁeld tocomplement the interpretability of neural networks [28].
RL algorithms also incorporate causality toimprove the reasoning capability [74].
For instance, [25] and [19] explicitly estimate causal structureswith the interventional data obtained from the environment.
These structures are either combinedwith dynamic model learning in model-based RL [18] or used to adjust the buffer priority [20].
Implicitly, we can abstract the causal structure and formulate it using the Block MDP [12] setting ortraining multiple encoders to extract different kinds of representations [15].
Following the idea ofinvariant risk minimization [75], they assume task-relevant features are invariant and shared acrossall environments, which can be used as the only cause of the reward.
Causal Discovery Causal discovery [76] is a long-stand topic in economics and sociology, where thetraditional methods can be generally categorized into constraint-based and score-based.
Constraintbased methods [29] start from a complete graph and iteratively remove edges with conditionalindependent test [77, 78] as constraints.
Score-based methods [38, 79] use metrics such as Bayesian
Information Criterion [80] as scores and prefer edges that maximize the score given the dataset.
Recently, researchers extend score-based methods with RL [81] or differentiable discovery [82, 83].
The former selects edges with a learned policy, and the latter learns a soft adjacent matrix withobservational or interventional data.
Active intervention methods are also explored [84] to increasethe efﬁciency of data collection and decrease the cost of conducting intervention [85].
Conclusion
This paper proposes a latent variable model that injects a causal graph reasoning process intotransition 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.
Toevaluate the performance of the proposed method, we designed nine tasks in three environments.
Thecomprehensive experiment results show that our method has better data efﬁciency and performancethan 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 wellto the number of nodes.
Developing efﬁcient gradient-based discovery methods could be a promisingdirection. In addition, the factorized state and action space assumption may restrict the usage of thiswork to semantic representations, which need to be processed with abstraction methods.
We furtherdiscuss the potential negative social impact and addition limitation in Appendix [-C.
-]C.3.