Provably Efficient Causal Model-Based Reinforcement Learning for Systematic Generalization

In the sequential decision making setting, an agent aims to achieve systematic generalization over a large, possibly infinite, set of environments. Such environments are modeled as discrete Markov decision processes with both states and actions represented through a feature vector. The underlying structure of the environments allows the transition dynamics to be factored into two components: one that is environment-specific and another one that is shared. Consider a set of environments that share the laws of motion as an illustrative example. In this setting, the agent can take a finite amount of reward-free interactions from a subset of these environments. The agent then must be able to approximately solve any planning task defined over any environment in the original set, relying on the above interactions only. Can we design a provably efficient algorithm that achieves this ambitious goal of systematic generalization? In this paper, we give a partially positive answer to this question. First, we provide the first tractable formulation of systematic generalization by employing a causal viewpoint. Then, under specific structural assumptions, we provide a simple learning algorithm that allows us to guarantee any desired planning error up to an unavoidable sub-optimality term, while showcasing a polynomial sample complexity.

In the sequential decision making setting, an agent aims to achieve systematic generalization over a large, possibly infinite, set of environments. Such environments are modeled as discrete Markov decision processes with both states and actions represented through a feature vector. The underlying structure of the environments allows the transition dynamics to be factored into two components: one that is environment-specific and another one that is shared. Consider a set of environments that share the laws of motion as an illustrative example. In this setting, the agent can take a finite amount of reward-free interactions from a subset of these environments. The agent then must be able to approximately solve any planning task defined over any environment in the original set, relying on the above interactions only. Can we design a provably efficient algorithm that achieves this ambitious goal of systematic generalization? In this paper, we give a partially positive answer to this question. First, we provide the first tractable formulation of systematic generalization by employing a causal viewpoint. Then, under specific structural assumptions, we provide a simple learning algorithm that allows us to guarantee any desired planning error up to an unavoidable sub-optimality term, while showcasing a polynomial sample complexity.

Introduction
Whereas recent breakthroughs have established Reinforcement Learning (RL) (Sutton & Barto, 2018) as a powerful tool to address a wide range of sequential decision making problems, the curse of generalization (Kirk et al., 2021 still a main limitation of commonly used techniques. RL algorithms deployed on a given task are usually effective in discovering the correlation between an agent's behavior and the resulting performance from large amounts of labeled samples (Jaksch et al., 2010;Lange et al., 2012). However, those algorithms are usually unable to discover basic cause-effect relations between the agent's behavior and the environment dynamics. Crucially, the aforementioned correlations are oftentimes specific to the task at hand, and they are unlikely to be of any use for addressing different tasks or environments. Instead, some universal causal relations generalize over the environments, and once learned can be exploited for solving any task. Let us consider as an illustrative example an agent interacting with a large set of physical environments. While each of these environments can have its specific dynamics, we expect the basic laws of motion to hold across the environments, as they encode general causal relations. Once they are learned, there is no need to discover them again from scratch when facing a new task, or an unseen environment. Even if the dynamics over these relations can change, such as moving underwater is different than moving in the air, or the gravity can change from planet to planet, the underlying causal structure still holds. This knowledge alone often allows the agent to solve new tasks in unseen environments by taking a few, or even zero, interactions.
We argue that we should pursue this kind of generalization in RL, which we call systematic generalization, where learning universal causal relations from interactions with a few environments allows us to approximately solve any task in any other environment without further interactions. Although this problem setting might seem overly ambitious or even far-fetched, in this paper we provide the first tractable formulation of systematic generalization, thanks to a set of structural assumptions that are motivated by a causal viewpoint. The problem formulation is partially inspired by reward-free RL (Jin et al., 2020), in which the agent can take unlabelled interactions with an environment to learn a model that allows approximate planning for any reward function. Here, we extend this formulation to a large, potentially infinite, set of reward-free environments, or a universe, the agent can freely interact with. We consider discrete environments, such that both their states and actions can be described through vectors of discrete features. Crucially, these environments share a common causal structure that explains a significant portion, but not all, of their transition dynamics. Can we design a provably efficient algorithm that guarantees an arbitrarily small planning error for any possible task that can be defined over the set of environments, by taking reward-free interactions with a generative model?
In this paper, we provide a partially positive answer to this question by presenting a simple but principled causal modelbased approach (see Fig. 1). This algorithm interacts with a finite subset of the universe to learn the causal structure underlying the set of environments in the form of a causal dependency graph G (Wadhwa & Dong, 2021). Then, the causal transition model, which encodes the dynamics that is common across the environment, is obtained by estimating the Bayesian network P G over G (Dasgupta, 1997) from a mixture of the environments. Finally, the causal transition model is employed by a planning oracle to provide an approximately optimal policy for a latent environment and a given reward function. We can show that this simple recipe, with a sample complexity that is polynomial in all the relevant quantities of the problem, allows achieving any desired planning error up to an unavoidable error term. The latter is inherent to the setting, which demands generalization over an infinite set of environments, and cannot be overcome without additional samples from the test environment.
The sample complexity rates that we derived are detailed in Table 1, while the main contributions of this paper include: • The first tractable formulation of the systematic generalization problem in RL, thanks to structural assumptions motivated by causal considerations (Section 3); • A provably efficient algorithm to learn systematic generalization over an infinite set of environments with polynomial sample complexity (Section 4.1); • The sample complexity of estimating the causal structure underlying a discrete MDP (Section 4.2); • The sample complexity of estimating the Bayesian network underlying a discrete MDP (Section 4.3).

Preliminaries
We start with some basic notions about graphs, causality, and Markov decision processes that will be useful thereafter. We will denote a set of integers We will denote by U A the uniform distribution over A.

Graphs
We define a graph G as a pair G := (V, E), where V is a set of nodes and E ⊆ N × N is a set of edges between them. We call G a directed graph if all of its edges E are directed (i.e., ordered pairs of nodes). We also define the in-degree of a node to be its number of incoming edges: degree in (A) = |{(B, A) : (B, A) ∈ E, ∀B}|. G is said to be a Directed Acyclic Graph (DAG) if it is a directed graph without cycles. We call G a bipartite graph if there exists a partition X ∪ Y = V such that none of the nodes in X and Y are connected by an edge, i.e., E ∩ ( The skeleton of a graph G is the undirected graph that is obtained from G by replacing all the directed edges in E with undirected ones. Finally, the graph edit distance between two graphs is the minimum number of graph edits (addition or deletion of either a node or an edge) necessary to transform one graph into the other.

Causal Graphs and Bayesian Networks
For a set X of random variables, we represent the causal structure over X with a DAG G X = (X , E), 1 which we call the causal graph of X . For each pair of variables A, B ∈ X , a directed edge (A, B) ∈ G X denotes that B is conditionally dependent on A. For every variable A ∈ X , we denote as Pa(A) the causal parents of A, i.e., the set of all the variables B ∈ X on which A is conditionally dependent, (B, A) ∈ G X .
A Bayesian network (Dean & Kanazawa, 1989) over the set X is defined as N := (G X , P ), where G X specifies the structure of the network, i.e., the dependencies between the variables in X , and the distribution P : X → ∆ X specifies the conditional probabilities of the variables in X , such that

Markov Decision Processes
A tabular episodic Markov Decision Process (MDP) (Puterman, 2014) is defined as M := (S, A, P, H, r), where S is a set of |S| = S states, A is a set of |A| = A actions, P is a transition model such that P (s |s, a) gives the conditional probability distribution of the next state s having taken action a in state s, H is the horizon of an episode, and r : S × A → [0, 1] is a deterministic reward function.
The strategy of an agent interacting with M is represented by a non-stationary, stochastic policy, a collection of functions (π h : S → ∆ A ) h∈ [H] where π h (a|s) denotes the conditional probability of taking action a in state s at step h. The value function V π h : S → R associated to π is defined as the expected sum of the rewards that will be collected, under the policy π, starting from s at step h, i.e., For later convenience, we further define P V π h+1 (s, a) := E s ∼P (·|s,a) [V π h+1 (s )] and V π 1 := E s∼P [V π 1 (s)]. We will write V π M,r to denote V π 1 in the MDP M with reward function r (if not obvious from the context). For an MDP M with finite states, actions, and horizon, there always exists an optimal policy π * that gives the value V * h (s) = sup π V π h (s) for every s, a, h. The goal of the agent is to find a policy π that is -close to the optimal one, i.e., V * 1 − V π 1 ≤ . Finally, we define a discrete Markov decision process as M := ((S, d S , n), (A, d A , n), P, H, r), where S, A, P, H, r are specified as before, and where the states and actions spaces admit additional structure, such that every s ∈ S can be represented through a d S -dimensional vector of discrete features taking value in [n], and every a ∈ A can be represented through a d A -dimensional vector of discrete features taking value in [n]. Note that any tabular MDP can be formulated under this alternative formalism by taking n = 2, d S = S, and d A = A, so that the states and actions are specified through one-hot encoding.

Problem Formulation
In our setting, a learning agent aims to master a large, potentially infinite, set U of environments modeled as discrete MDPs without rewards, , which we call a universe. The agent can draw a finite amount of experience by interacting with the MDPs in U. From these interactions alone, the agent aims to acquire sufficient knowledge to approximately solve any task that can be specified over the universe U. Specifically, a task is defined as any pairing of an MDP M ∈ U and a reward function r, whereas solving it refers to providing a slightly sub-optimal policy via planning, i.e., without taking additional interactions. We call this problem systematic generalization, which we can formalize as follows. Definition 1 (Systematic Generalization). For any latent MDP M ∈ U and any given reward function r : S × A → [0, 1], the systematic generalization problem requires the agent to provide a policy π, such that V * M,r − V π M,r ≤ up to any desired sub-optimality > 0.
Since the set U is infinite, we clearly require additional structure to make the problem feasible. On the one hand, the state space (S, d S , n), action space (A, d A , n), and initial state distribution µ are shared across M ∈ U. The transition dynamics P i is instead specific to each MDP M i ∈ U. However, we assume the presence of a common causal structure that underlies the transition dynamics of the universe, and relates the single transition models P i .

The Causal Structure of the Transition Dynamics
The transition dynamics of a discrete MDP gives the conditional probability of next state features s given the current [n] d S + d A Figure 2. Illustration of the causal structure G of U.
state-action features (s, a). To ease the notation, from now on we will denote the current state-action features with a random vector X = ( , and the next state features with a random For each environment M i ∈ U, the conditional dependencies between the next state features Y and the current state-action features X are represented through a bipartite dependency graph . Clearly, each environment can display its own dependencies, but we assume there is a set of dependencies that represent general causal relationships between the features, and that appear in any M i ∈ U. In particular, we call the intersection G := ∩ ∞ i=0 G i the causal structure of U, which is the set of conditional dependencies that are common across the universe. In Figure 2, we show an illustration of such a causal structure. Since it represents universal causal relationships, the causal structure G is time-consistent, i.e., G (h) = G (1) for any step h ∈ [H], and we further assume that G is sparse, which means that the number of features X[z] on which a feature Y [j] is dependent on is bounded from above.
Given a causal structure G, and without losing generality, 2 we can express each transition model P i as in which P G is the Bayesian network over the causal structure G, whereas F i includes environment-specific factors affecting the conditional probabilities. 3 Since it represents the conditional probabilities due to universal causal relations in U, we call P G the causal transition model of U. Thanks to the structure G, this model can be further factored as where the scopes Z j are the set of indices z such that . For the purpose of this work, we make the crucial assumption that the causal transition model is non-vacuous and that it explains a significant part of the transition dynamics of M i ∈ U. Assumption 2 (λ-sufficiency). Let λ ∈ [0, 1] be a constant. The causal transition model P G is causally λ-sufficient if Notably, the parameter λ controls the amount of the transition dynamics that is due to the universal causal relations G (λ = 0 means that P G is sufficient to explain the transition dynamics of any M i ∈ U, whereas λ = 1 implies no shared structure between the transition dynamics of the M i ∈ U).
In this paper, we argue that learning the causal transition model P G is a good target for systematic generalization and we provide theoretical support for this claim in Section 4.

A Class of Training Environments
Even if the universe U admits the structure that we presented in the last section, it is still an infinite set. Instead, the agent can only interact with a finite subset of discrete MDPs which we call a class of size M . Crucially, the causal structure G is a property of the full set U, and if we aim to infer it from interactions with a finite class M, we have to assume that M is informative enough on the universal causal relations of U.
Analogously, if we aim to infer the causal transition model P G from interactions with the transition models P i of the single MDPs M i ∈ M, we have to assume that M is balanced in terms of the conditional probabilities displayed by its components, so that the factors that do not represent universal causal relations even out while learning.
Assumption 4 (Evenness). Let M ⊂ U be class of size M . We say that M is causally even if Whereas in this paper we assume that M is diverse and even by design, we leave as future work the interesting problem of selecting such a class from active interactions with U, which would add to our problem formulation flavors of active learning and experimental design (Hauser & Bühlmann, 2014;Kocaoglu et al., 2017;Ghassami et al., 2018).

Learning Systematic Generalization
Before addressing the sample complexity of systematic generalization, it is worth considering the kind of interactions that we need in order to learn the causal transition model P G and its underlying causal structure G. Especially, thanks to the peculiar configuration of the causal structure G, i.e., a bipartite graph in which the edges are necessarily directed from the state-action features X to the next state features Y , as a causation can only happen from the past to the future, learning the skeleton of G is equivalent to learning its full structure. Crucially, learning the skeleton of a causal graph does not need specific interventions, as it can be done from observational data alone (Hauser & Bühlmann, 2014).
Proposition 1. The causal structure G of U can be identified from purely observational data.
In this paper, we will consider the online learning setting with a generative model for estimating G and P G from sampled interactions with a class M of size M . A generative model allows the agent to set the state of an MDP before sampling a transition, instead of drawing sequential interactions from the process. Finally, analogous results to what we obtain here can apply to the offline setting as well, in addition to convenient coverage assumptions on the dataset.

Sample Complexity Analysis
We provide a sample complexity analysis of the problem, which stands as a core contribution of this paper along with the problem formulation itself that we provided in Section 3.
First, we consider the sample complexity of systematic generalization (Section 4.1). Then, we provide ancillary sample complexity results on the estimation of the causal structure (Section 4.2) and the Bayesian network (Section 4.3) of an MDP, which can be of independent interest.

Sample Complexity of Systematic Generalization with a Generative Model
We have access to a class M of discrete MDPs within a universe U, from which we can draw interactions with a generative model P (X). We would like to solve the systematic generalization problem as described in Definition 1. This problem requires to provide, for any combination of a (latent) MDP M ∈ U, and a given reward function r, a planning policy π such that V * M,r − V π M,r ≤ . Especially, can we design an algorithm that guarantees this requirement with high probability by taking a number of samples K that is polynomial in and the relevant parameters of M?
Here we give a partially positive answer to this question, by providing a simple but provably efficient algorithm that guarantees systematic generalization over U up to an unavoidable sub-optimality term λ that we will later specify.
The algorithm implements a model-based approach into two separated components. The first, for which we provide the pseudocode in Algorithm 1, is the procedure that actually interacts with the class M to obtain a principled estimation P G of the causal transition model P G of U. The second, is a planning oracle that takes as input a reward function r and the estimated causal transition model, and returns an optimal policy π operating on P G as an approximation of the transition model P i of the true MDP M i . 5 First, we provide the sample complexity of the causal transition model estimation (Algorithm 1), which in turn is based on repeated causal structure estimations (Algorithm 2) to obtain G, and an estimation procedure of the Bayesian network over G (Algorithm 3) to obtain P G .
be a class of M discrete MDPs, let δ ∈ (0, 1), and let > 0. The Algorithm 1 returns a causal transition model P G such that P r( P G − P G 1 ≥ ) ≤ δ with a sample complexity An analogous result can be derived for the tabular MDP setting, as stated by the following lemma.
MDPs. The sample complexity of Lemma 4.1 reduces to Having established the sample complexity of the causal transition model estimation, we can now show how the learned model P G allows us to approximately solve, via a planning oracle, any task defined by a combination of a latent MDP M i ∈ U and a given reward function r. To provide this result in the discrete MDP setting, we have to further assume that the transition dynamics P i of the target MDP M i admits factorization analogous to (1), such that we can write where the scopes Z j are given by the environment-specific causal structure G i , which we assume to be 2Z-sparse (Assumption 1).
Theorem 4.3. Let δ ∈ (0, 1) and > 0. For a latent discrete MDP M ∈ U, and a given reward function r, a planning oracle operating on the causal transition model P G as an approximation of M returns a policy π such that Without the additional factorization of the environmentspecific transition model, the result of Theorem 4.3 reduces to the analogous for the tabular MDP setting. Theorem 4.3 and Corollary 4.4 establish the sample complexity of systematic generalization through Lemma 4.1 and Lemma 4.2 respectively. For the discrete MDP setting, we have that O(M H 6 d 3 S Z 2 n 5Z+3 ), which reduces to O(M H 6 S 4 A 2 Z 2 ) in the tabular setting. Unfortunately, we are only able to obtain systematic generalization up to an unavoidable sub-optimality term λ . This error term is related to the λ-sufficiency of the causal transition model (Assumption 2), and it accounts for the fact that P G cannot fully explain the transition dynamics of each M ∈ U, even when it is estimated exactly. This is inherent to the ambitious problem setting, and can be only overcome with additional interactions with the test MDP M.

Sample Complexity of Learning the Causal Structure of a Discrete MDP
As a byproduct of the main result in Theorem 4.3, we can provide a specific sample complexity result for the problem of learning the causal structure G underlying a discrete MDP M with a generative model. We believe that this problem can be of independent interest, mainly in consideration of

Algorithm 1 Causal Transition Model Estimation
Input: class of MDPs M, error , confidence δ Output: causal transition model P G previous work on causal discovery of general stochastic processes (e.g., Wadhwa & Dong, 2021), for which we refine known results to account for the structure of an MDP, which allows for a tighter analysis of the sample complexity.
Instead of targeting the exact dependency graph G, which can include dependencies that are so weak to be nearly impossible to detect with a finite number of samples, we only address the dependencies above a given threshold.
Before presenting the sample complexity result, we state the existence of a principled independence testing procedure.
We can now provide an upper bound to the number of samples required by a simple estimation procedure to return an ( , δ)-estimate G of the causal dependency graph G.
The Algorithm 2 returns a dependency graph G such that P r( G = G ) ≤ δ with a sample complexity Finally, we can state an analogous result for a tabular MDP setting, by taking n = 2, d S = S, d A = A. Corollary 4.7. Let M be a tabular MDP. The result of Theorem 4.6 reduces to K = O log(S 2 A/δ)/ 2 .

Sample Complexity of Learning the Bayesian Network of a Discrete MDP
We present as a standalone result an upper bound to the sample complexity of learning the parameters of a Bayesian network P G with a fixed structure G. Especially, we refine known results (e.g., Friedman & Yakhini, 1996;Dasgupta, 1997;Cheng et al., 2002;Abbeel et al., 2006;Canonne et al., 2017) by considering the specific structure G of an MDP.
If the structure G is dense, the number of parameters of P G grows exponentially, making the estimation problem mostly intractable. Thus, we consider a Z-sparse G (Assumption 1), as in previous works (Dasgupta, 1997). Then, we can provide a polynomial sample complexity for the problem of learning the Bayesian network P G of a discrete MDP M.
The Algorithm 3 returns a Bayesian network P G such that P r( P G − P G 1 ≥ ) ≤ δ with a sample complexity Analogously, we can state the sample complexity of learning the fixed structure Bayesian network of a tabular MDP.
Corollary 4.9. Let M be a tabular MDP. The result of Theorem 4.8 reduces to

Numerical Validation
We empirically validate the theoretical findings of this work by experimenting on a synthetic example where each environment is a person, and the MDP represents how a series of actions the person can take influences their weight (W ) and academic performance (A). As actions we consider hours of physical training (P ), hours of sleep (S), hours of study (St), amount of vegetables in the diet (D), and the amount of caffeine intake (C). The obvious use-case for such a model would be a tracking device that monitors how the actions of a person influence their weight and academic performance and provides personalized recommendations to

Algorithm 3 Bayesian Network Estimation for an MDP
Output: Bayesian network P G reach the person's goals. While the physiological responses of different individuals can vary, there are some underlying mechanisms shared by all humans, and therefore deemed causal in our terminology. Examples of such causal links are the dependency of weight on the type of diet, and the dependency of academic performance on the number of hours of study. Other links, such as the dependency of weight on the amount of caffeine can be present in some individuals, but are generally not shared and therefore not causal. We report in the Appendix ( Figure 5) an illustration of the described domain. For simplicity, all variables are treated as discrete with possible values 0 (below average), 1 (average) or 2 (above average). See Appendix B for details on how transition models of different environments are generated. A class M of 3 environments is used to estimate the causal transition model. All experiments are repeated 10 times and report the mean and standard deviation of the desired quantities.
Causal Structure Estimation We first empirically investigate the graph edit distance between estimated and groundtruth causal structure, GED(G, G), as a function of number of samples (K in Algorithm 1). The causal structure is estimated by obtaining the causal graph for each training environment (using a series of independence tests), and taking the intersection of their edges. As expected, as we increase the number of samples, the distance converges to zero, and we can perfectly recover the causal graph (Figure 4a). Figure 4b presents the L 1 -distance between the estimate and ground truth causal transition model, as a function of the total number of samples (K + K in Algorithm 1). As the number of samples grows, the L 1 -distance shrinks, converging at around 0.05, an estimation error due to the environments  not respecting the evenness assumption.

Causal Transition Model Estimation
Value Function Estimation Finally, we investigate how well we can approximate the optimal value function for a new unseen environment (Figure 4c). We observe that our algorithm is able to effectively approximate the optimal value function up to a small error with a reasonable number of samples.

Related Work
We revise the relevant literature and we discuss how it relates with our problem formulation and the reported results. whenever U is small, but leveraging the inner structure of U becomes more and more important as U grows to infinity, while M remains constant. However, our approach pays this further generality with the additional error term λ , which is unavoidable. Finally, it is an interesting direction for future work to establish whether the additional factors in S, A, H w.r.t. reward-free RL are also unavoidable.

Reward-Free Reinforcement Learning
Multi-Task Reinforcement Learning In the multi-task RL setting (Lazaric & Ghavamzadeh, 2010;Lazaric & Restelli, 2011;Brunskill & Li, 2013;Liu et al., 2016), the agent interacts with a set of MDPs (tasks) pursuing generalization from one task to the other, which is in a way similar to the aim of systematic generalization. Although sample complexity results have been derived for the multitask setting (e.g., Brunskill & Li, 2013), we are not aware of previous works in this stream that considered the problem of interacting with a set of reward-free MDPs, nor the presence of a common causal structure underlying the set of MDPs.

Model-Based Reinforcement Learning
The modelbased RL (Sutton & Barto, 2018) methodology prescribes learning an approximate model of transition dynamics in order to learn an optimal policy. Theoretical works (e.g., Jaksch et al., 2010;Ayoub et al., 2020) generally concern with the estimation of the approximate value functions obtained through the learned model, rather than the estimation of the model itself. A notable exception is the work by Tarbouriech et al. (2020), which targets point-wise high probability guarantees on the model estimation as we do in Lemma 4.1, 4.2. However, they address the model estimation of a single tabular MDP M, instead of the shared transition dynamics of an infinite set of MDPs U that we target in this paper. Other works (e.g., Zhang et al., 2020;Tomar et al., 2021) have also addressed model-based RL from a causal perspective. To the best of our knowledge, we are the first to prove a polynomial sample complexity for causal model-based RL in systematic generalization.

Factored Markov Decision Processes
The factored MDP formalism (Kearns & Koller, 1999) allows encoding transition dynamics that are the product of multiple independent factors. This formalism is closely related to how we define the causal transition model as a product of independent factors in (1), which can be actually seen as a factored MDP.

Conclusion
In this work, we formulate the problem of systematic generalization in reinforcement learning from a causal viewpoint. We then propose a provably efficient algorithm which partially solves the problem under certain structural assumptions motivated by causal considerations.

A. Proofs
Proofs of Section 3.3 Proposition 1. The causal structure G of U can be identified from purely observational data.
Proof. First, recall that with observational data alone, a causal graph can be identified up to its Markov equivalence class (Hauser & Bühlmann, 2014). This means that its skeleton and v-structure are properly identified, meanwhile determining the edge orientations requires interventional data in the general case. Since in the considered causal graph G the edges orientations are determined a priori (as they follow the direction of time), the causal graph can be entirely determined by using only observational data.

Proofs of Section 4.1: Causal Transition Model Estimation
Before reporting the proof of the main result in Theorem 4.3, it is worth considering a set of lemmas that will be instrumental to the main proof.
First, we provide an upper bound to the L1-norm between the Bayesian network P G over a given structure G and the Bayesian network P G over the structure G , which is the -dependency subgraph of G as defined in Definition 2.
Lemma A.1. Let G a Z-sparse dependency graph, and let G its corresponding -dependence subgraph for a threshold > 0. The L1-norm between the Bayesian network P G over G and the Bayesian network P G over G can be upper bounded as Proof. The proof is based on the fact that every edge (X i , Y j ) such that (X i , Y j ) ∈ G and (X i , Y j ) / ∈ G corresponds to a weak conditional dependence (see Definition 2), which means that P Yj |Xi − P Yj 1 ≤ .
We denote with Z j the scopes of the parents of the node Y [j] in G, i.e., Pa G (Y [j]) = X[Z j ], and with Z j, the scopes of the parents of the node Y [j] in G , i.e., Pa G (Y [j]) = X[Z j, ]. As a direct consequence of Definition 2, we have Z j, ⊆ Z j for any j ∈ d S , and we can write Then, we let Z j \ Z j, = [I] overwriting the actual indices for the sake of clarity, and we derive in which we employed the property µ − ν 1 ≤ i µ i − i ν i 1 ≤ i µ i − ν i 1 for the L1-norm between product distributions µ = i µ i , ν = i ν i to write (2), we repeatedly applied the triangle inequality µ−ν 1 ≤ µ−ρ 1 + ρ−ν 1 to get (3) from (2), we upper bounded each term of the sum in (3) with thanks to Definition 2, and we finally employed the Z-sparseness Assumption 1 to upper bound I with Z in (4).
Next, we provide a crucial sample complexity result for a provably efficient estimation of a Bayesian network P G over an estimated -dependency subgraph G, which relies on both the causal structure estimation result of Theorem 4.6 and the Bayesian network estimation result of Theorem 4.8.
Lemma A.2. Let M be a discrete MDP, let M = {M} be a singleton class, let δ ∈ (0, 1), and let > 0. The Algorithm 1 returns a Bayesian network P G such that P r( P G − P G 1 ≥ ) ≤ δ with a sample complexity Proof. We aim to obtain the number of samples K = K + K for which Algorithm 1 is guaranteed to return a Bayesian network estimate P G over a causal structure estimate G such that P r( P G − P G 1 ≥ ) ≤ δ in a setting with a singleton class of discrete MDPs. First, we derive the following decomposition of the error in which we employed the triangle inequality µ − ν 1 ≤ µ − ρ 1 + ρ − ν 1 . Then, we can write Bayesian network estimation ( ) Bayesian network subgraph ( ) through the decomposition (5) and a union bound to isolate the three independent sources of error ( ), (•), ( ). To upper bound the latter term ( ) with 0, we invoke Lemma A.1 to have d s Z ≤ 3 , which gives ≤ 3d S Z . Then, we consider the middle term (•), for which we can write We can now upper bound (•) ≤ δ/2 through (6) by invoking Theorem 4.6 with threshold = 3d S Z and confidence δ = δ 2 , which gives Next, we can upper bound ( ) ≤ δ/2 by invoking Theorem 4.8 with threshold = 3 and confidence δ = δ 2 , which gives Finally, through the combination of (7) and (8), we can derive the sample complexity that guarantees P r( P G − P G 1 ≥ ) ≤ δ under the assumption 4/3 2 , i.e., which concludes the proof.
Whereas Lemma A.2 is concerned with the sample complexity of learning the Bayesian network of a singleton class, we can now extend the result to account for a class M composed of M discrete MDPs.
be a class of M discrete MDPs, let δ ∈ (0, 1), and let > 0. The Algorithm 1 returns a causal transition model P G such that P r( P G − P G 1 ≥ ) ≤ δ with a sample complexity Proof. We aim to obtain the number of samples K = M K + K for which Algorithm 1 is guaranteed to return a Bayesian network estimate P G over a causal structure estimate G such that P r( P G − P G 1 ≥ ) ≤ δ in a setting with a class of M discrete MDPs. First, we can derive an analogous decomposition as in (5), such that we have P r P G − P G 1 ≥ ≤ P r P G − P G 1 ≥ 3 Bayesian network estimation ( ) + P r P G − P G 1 ≥ 3 causal structure estimation (•) + P r P G − P G 1 ≥ 3 Bayesian network subgraph ( ) through a union bound. Crucially, the terms ( ), ( ) are unaffected by the class size, which leads to K = (8) by upper bounding ( ), and ≤ 3d S Z by upper bounding ( ), exactly as in the proof of Lemma A.2. Instead, the number of samples K has to guarantee that (•) = P r( P G − P G 1 ≥ /3) ≤ δ/2, where the causal structure G is now the intersection of the causal structures of the single class components M i , i.e., G = ∩ M i=1 G ,i . Especially, we can write through a union bound on the estimation of the single causal structures G i . Then, we can upper bound (•) ≤ δ/2 through (9) by invoking Theorem 4.6 with threshold = 3d S Z and confidence δ = δ 2M , which gives Finally, through the combination of (10) and (8), we can derive the sample complexity that guarantees P r( P G − P G 1 ≥ ) ≤ δ under the assumption 4/3 2 , i.e., which concludes the proof.
It is now straightforward to extend Lemma 4.1 for a class M composed of M tabular MDPs.
Proof. To obtain K = M K + K , we follows similar steps as in the proof of Lemma 4.1, to have the usual decomposition of the event P r( P G − P G 1 ≥ ) in the ( ), (•), ( ) terms. We can deal with ( ) as in Lemma 4.1 to get ≤ 3SZ . Then, we upper bound (•) ≤ δ/2 by invoking Corollary 4.7 (instead of Theorem 4.6) with threshold = 3SZ and confidence δ = δ 2M , which gives Similarly, we upper bound ( ) ≤ δ/2 by invoking Corollary 4.9 (instead of Theorem 4.5) with threshold = 3 and confidence δ = δ 2 , which gives Finally, we combine 11 with 12 to obtain

Proofs of Section 4.1: Planning
For the sake of notational clarity within the following proofs we express Theorem 4.3. Let δ ∈ (0, 1) and > 0. For a latent discrete MDP M ∈ U, and a given reward function r, a planning oracle operating on the causal transition model P G as an approximation of M returns a policy π such that where λ = 2λH 3 d S n 2Z+1 , and P G is obtained from Algorithm 1 with δ = δ and = /2H 3 n Z+1 .
Proof. Consider the MDPs with transition model P and P G . We refer to the respective optimal policies as π * and π * . Moreover, since the reward r is fixed, we remove it from the expressions for the sake of clarity, and refer with V to the value function of the MDP with transition model P G . As done in (Jin et al., 2020, Theorem 3.5), we can write the following decomposition, where V * := V π * .
where in the last step we have set to 0 the approximation due to the planning oracle assumption, and we have bounded the evaluation errors according to Lemma A.3. In order to get 2n Z+1 H 3 = we have to set = 2n Z+1 H 3 . Considering the sample complexity result in Lemma 4.1 the final sample complexity will be: Under the preconditions of Theorem 4.3, with probability 1 − δ, for any reward function r and policy π, we can bound the value function estimation error as follows.
where V is the value function of the MDP with transition model P G , is the approximation error between P G and P G studied in Lemma 4.1, and λ stands for the λ-sufficiency parameter of P G .
Proof. The proof will be along the lines of that of Lemma 3.6 in (Jin et al., 2020). We first recall (Dann et al., 2017, Lemma E.15), which we restate in Lemma A.5. In this prof we consider an environment specific true MDP M with transition model P , and an mdp M that has as transition model the estimated causal transition model P G . In the following, the expectations will be w.r.t. P . Moreover, since the reward r is fixed, we remove it from the expressions for the sake of clarity. We can start deriving We now bound a single term within the sum above as follows: We can now bound each term. Let us start considering the first term: We can now bound the first term of (16): Due to the uniform sampling and Z-sparseness assumptions, we have P G (X[Z j ]) = 1 n Z , hence: = n Z Therefore: Replacing this in (17) and marginalizing over Y \Y [j] we obtain: Where d S is the approximation term of each component. By plugging this bound into (16) we get: where in the last step we have recursively bounded the right terms as in (24). By considering 2Z-sparseness, λ-sufficiency, and that the transition model P factorizes, we can apply the same procedure to bound the second term of equation (15) as: Therefore the initial expression in (14) becomes: Corollary 4.4. For a tabular MDP M ∈ M, the result of Theorem 4.3 holds with λ = 2λSAH 3 , = /2SAH 3 .
Proof. Consider the MDPs with transition model P and P G . We refer to the respective optimal policies as π * and π * . Moreover, since the reward r is fixed, we remove it from the expressions for the sake of clarity, and refer with V to the value function of the MDP with transition model P G . As done in (Jin et al., 2020, Theorem 3.5), we can write the following decomposition, where V * := V π * .
where in the last step we have set to 0 the approximation due to the planning oracle assumption, and we have bounded the evaluation errors according to Lemma A.4. In order to get 2SAH 3 = we have to set = 2SAH 3 . Considering the sample complexity result in Lemma 4.2 the final sample complexity will be: Lemma A.4. Under the preconditions of Corollary 4.4, with probability 1 − δ, for any reward function r and policy π, we can bound the value function estimation error as follows.
where V is the value function of the MDP with transition model P G , is the approximation error between P G and P G studied in Lemma 4.1, and λ stands for the λ-sufficiency parameter of P G .
Proof. The proof will be along the lines of that of Lemma 3.6 in (Jin et al., 2020). We first recall (Dann et al., 2017, Lemma E.15), which we restate in Lemma A.5. In this prof we consider an environment specific true MDP M with transition model P , and an mdp M that has as transition model the estimated causal transition model P G . In the following, the expectations will be w.r.t. P . Moreover, since the reward r is fixed, we remove it from the expressions for the sake of clarity. We can start deriving We now bound a single term within the sum above as follows: where in the last step we have used the fact that there must exist an optimal deterministic policy. Due to the uniform sampling assumption, we have P (s, a) = 1 SA , hence: max π † P π † (s, a) P (s, a) ≤ 1 P (s, a) = SA Therefore: P π † (s, a) ≤ SA · P (s, a) Moreover, notice that, since π is deterministic we have P π (s) = P π (s) = P π (s, a) ≤ SA · P (s, a). Replacing it in the expression above we get where is the approximation error between P G and P G studied in Lemma 4.1, and in the penultimate step we have used the following derivation: Hence, due to this recursive unrolling, we have: Notice that the same argument holds also for the second term of (22), replacing with λ.
By plugging the result in equation (23) into the initial expression we get: In the following we restate (Dann et al., 2017, Lemma E.15) for the case of stationary transition model.
Lemma A.5. For any two MDPs M and M with rewards r and r and transition models P and P , the difference in value functions V , V w.r.t. the same policy π can be written as: Proofs of Section 4.2 We provide the proof of the sample complexity result for learning the causal structure of a discrete MDP with a generative model.
Theorem 4.6. Let M be a discrete MDP with an underlying causal structure G, let δ ∈ (0, 1), and let > 0. The Algorithm 2 returns a dependency graph G such that P r( G = G ) ≤ δ with a sample complexity Proof. We aim to obtain the number of samples K for which Algorithm 2 is guaranteed to return a causal structure estimate G such that P r( G = G ) ≤ δ in a discrete MDP setting. First, we can upper bound the probability of the bad event P r( G = G ) in terms of the probability of a failure in the independence testing procedure I(X z , Y j ) for a single pair of nodes X z ∈ G , Y z ∈ G , i.e., where we applied an union bound to obtain the last inequality. Now we can look at the probability of a single independence test failure. Especially, for a provably efficient independence test (the existence of such a test is stated by Lemma 4.5, whereas the Algorithm 2 in Diakonikolas et al. (2021) reports an actual testing procedure), we have P r(test I(X z , Y j ) fails) ≤ δ , for any choice of δ ∈ (0, 1), > 0, with a number of samples  (27), we obtain P r( G = G ) with a sample complexity under the assumption 2 4/3 , which concludes the proof.
The proof of the analogous sample complexity result for a tabular MDP setting (Corollary 4.7) is a direct consequence of Theorem 4.6 by letting n = 2, d S = S, d A = A.

Proofs of Section 4.3
We first report a useful concentration inequality for the L1-norm between the empirical distribution computed over K samples and the true distribution (Weissman et al., 2003, Theorem 2.1). Lemma A.6 (Weissman et al. (2003)). Let X 1 , . . . , X K be i.i.d. random variables over [n] having probabilities P r(X k = i) = P i , and let P K (i) = 1 K K k=1 1(X k = i). Then, for every threshold > 0, it holds P r P K − P 1 ≥ ≤ 2 exp(−K 2 /2n).
We can now provide the proof of the sample complexity result for learning the Bayesian network of a discrete MDP with a given causal structure. Theorem 4.8. Let M be a discrete MDP, let G be its underlying causal structure, let δ ∈ (0, 1), and let > 0. The Algorithm 3 returns a Bayesian network P G such that P r( P G − P G 1 ≥ ) ≤ δ with a sample complexity Proof. We aim to obtain the number of samples K for which Algorithm 3 is guaranteed to return a Bayesian network estimate P G such that P r( P G − P G 1 ≥ ) ≤ δ in a discrete MDP setting. First, we note that in which we employed the property µ − ν 1 ≤ i µ i − i ν i 1 ≤ i µ i − ν i 1 for the L1-norm between product distributions µ = i µ i , ν = i ν i to write (28), and we applied a union bound to derive (31) from (30). Similarly, we can write P r P j − P j 1 ≥ d S ≤ P r x∈[n] |Z j | P j (·|x) − P j (·|x) 1 ≥ d S n |Zj | (32) ≤ x∈[n] |Z j | P r P j (·|x) − P j (·|x) 1 ≥ d S n |Zj | (33) by applying a union bound to derive (33) from (32), and by employing Assumption 1 to bound |Z j | with Z in (34). We can now invoke Lemma A.6 to obtain the sample complexity K that guarantees P r( P j (·|x) − P j (·|x) 1 ≥ ) ≤ δ , i.e., K = 2n log(2/δ ) ( ) 2 = 2 d 2 S n 2Z+1 log(2d S n Z /δ) 2 , where we let = d S n Z , δ = δ d S n Z . Finally, by summing K for any x ∈ [nm] |Zj | and any j ∈ [d S ], we obtain which proves the theorem.
To prove the analogous sample complexity result for a tabular MDP we can exploit a slightly tighter concentration on the KL divergence between the empirical distribution and the true distribution in the case of binary variables (Dembo & Zeitouni, 2009, Theorem 2.2.3) 6 , which we report for convenience in the following lemma.
We can now provide the proof of Corollary 4.9.
Proof. We aim to obtain the number of samples K for which Algorithm 3 is guaranteed to return a Bayesian network estimate P G such that P r( P G − P G 1 ≥ ) ≤ δ in a tabular MDP setting. We start by considering the KL divergence d KL ( P G ||P G ). Especially, we note Then, for any > 0 we can write P r d KL P G ||P G ≥ ≤ P r in which we applied a first union bound to get (36) from (35), a second union bound to get (38) from (37), and Assumption 1 to bound |Z j | with Z in (39). We can now invoke Lemma A.7 to obtain the sample complexity K that guarantees P r(d KL ( P j (·|x)||P j (·|x)) ≥ ) ≤ δ , i.e., where we let = S2 Z , and δ = δ S2 Z for any choice of δ ∈ (0, 1). By summing K for any x ∈ [2] |Zj | and and j ∈ [S], we obtain the sample complexity K that guarantees P r(d KL ( P G ||P G ) ≥ ) ≤ δ , i.e., Finally, we employ the Pinsker's inequality P G − P G 1 ≤ 2d KL ( P G ||P G ) (Csiszár, 1967) to write which gives the sample complexity K that guarantees P r( P G − P G 1 ≥ ) ≤ δ by letting = 2 2 and δ = δ in (40), i.e.,

B. Numerical Validation
Let A, W , P , S, D, C and St represent academic performance, weight, physical activity, sleep, diet and study respectively. We start by defining the causal transition model as follows: The transition model for a specific environment is then generated by adding independent white noise to each coefficient in the above equations for µ A and µ W , including the ones not shown because set to 0: µ A = (1 + 1 )A + 2 W + 3 P + (0.2 + 4 )D + (0.5 + 5 )S + 6 C + (0.8 + 7 )St − 0.8 µ W = 8 A + (1 + 9 )W + (−0.5 + 10 )P + (−0.5 + 11 )D + 12 S + 13 C + 14 St + 0. causal edge e ∈ G environment specific edge e / ∈ G Figure 5. Bayesian network for a specific environment (person) of our synthetic example. Grey edges are causal and therefore shared by other environments in the same universe U, while blue dashed edges are environment-specific dependencies.