PaperlayerView exact revision on arXiv

Source-linked AI summary

Provably Efficient Causal Model-Based Reinforcement Learning for Systematic Generalization

Mirco Mutti, Riccardo De Santi, Emanuele Rossi, Juan Felipe Calderon, Michael Bronstein, Marcello Restelli

arXiv:2202.06545v3cs.LG

TL;DR

The paper asks how an agent can generalize across a potentially infinite set of environments using reward-free interactions from only some of them. It formulates the problem through shared causal structure and gives a model-based algorithm with polynomial sample complexity, while retaining an unavoidable sub-optimality term.

  • Problem

    The problem is to use finite reward-free experience from a universe of environments to approximately solve any task in that universe without further interactions.

    Source: §1 · p. 1, §3 · p. 3

  • Method

    The approach learns dependencies common across environments, estimates their shared Bayesian-network transition model, and uses it for planning in an unknown environment.

    Source: §1 · p. 2, §3.1 · p. 3, §4.1 · p. 5

  • Results

    The algorithm has polynomial sample complexity, but systematic generalization retains an unavoidable sub-optimality term ελ because the shared model cannot fully explain every environment’s dynamics.

    Source: §1 · p. 2, §4.1 · p. 5

  • Takeaways & Limitations

    Within the stated formulation, learned universal causal relations support approximately solving tasks in other environments without further interactions.

    Source: §1 · p. 1

    The training environments are assumed diverse by design; selecting such a class through active interactions is left for future work.

    Source: §3.2 · p. 4

Abstract

from arXiv · show

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 that is shared. Consider a set of environments that share the laws of motion as an 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 a tractable formulation of systematic generalization by employing a causal viewpoint. Then, under specific structural assumptions, we provide a simple learning algorithm that guarantees any desired planning error up to an unavoidable sub-optimality term, while showcasing a polynomial sample complexity.

1 Introduction

The introduction frames systematic generalization as learning causal relations shared across environments so an agent can plan for new tasks without further interactions. It presents a tractable causal formulation and a model-based algorithm with polynomial sample complexity, subject to an unavoidable error term.

1 Introduction

  • Task-specific behavioral correlations may not transfer, whereas shared causal relations such as laws of motion can apply across environments with different dynamics.The introduction notes that such knowledge can help solve new tasks with few or even zero interactions.Source: §1 · p. 1
  • Systematic generalization formalizes learning universal causal relations from a few environments to approximately solve tasks in other environments without further interactions.The formulation extends reward-free planning to a potentially infinite universe of discrete environments with shared causal structure.Source: §1 · p. 1
  • A causal model-based algorithm learns the universe’s dependency graph and shared transition model, then plans for an unknown environment with polynomial sample complexity and an unavoidable error term.The error cannot be overcome without additional samples from the test environment.Source: §1 · p. 2
  • The contributions include a tractable formulation, an efficient algorithm, and sample-complexity analyses for causal-structure and Bayesian-network estimation.The analyses adapt causal-discovery and Bayesian-network estimation results to the MDP setting as building blocks for the generalization rate.Source: §1 · p. 2
  • The paper also aims to connect model-based and reward-free RL, causal RL, factored MDPs, independence testing, and experimental design in a common framework.The stated aim is for progress in these areas to be enhanced beyond the sum of their parts.Source: §1 · p. 2

2 Preliminaries

The preliminaries establish notation for feature spaces, graphs, Bayesian networks, and episodic MDPs. They also define the feature-based representation that includes tabular MDPs as a special case.

2 Preliminaries

  • The section introduces set and probability-simplex notation, scopes for feature subsets, and L1 distance and KL divergence for discrete distributions.A scope selects coordinates of a product space.Source: §2 · p. 2
  • It defines directed graphs, DAGs, graph skeletons and edit distance, then represents causal structure with a DAG and a Bayesian network’s conditional distributions.In a Bayesian network, the joint distribution factors into conditional distributions given each variable’s parents.Source: §2 · p. 2
  • An episodic MDP specifies states, actions, transitions, horizon, and rewards; policies induce value functions, and the optimal policy provides the reference for an ε-close policy.The discrete-MDP formalism represents states and actions as vectors of discrete features and can encode tabular MDPs using one-hot vectors.Source: §2 · pp. 2–3

3 Problem Formulation

Systematic generalization asks an agent to use finite experience in a potentially infinite universe of reward-free MDPs to approximately solve any environment-and-reward task without further interaction. The formulation relies on causal structure shared across environments and a finite training class with specific structural properties.

3 Problem Formulation

  • An agent uses finite experience from a potentially infinite universe of reward-free MDPs to approximately solve any environment-and-reward task without further interaction.The target is a policy within a desired sub-optimality level.Source: §3 · p. 3

3.1 Causal Structure of the Transition Dynamics

  • The universe shares a causal structure linking current state-action features to next-state features, even though individual environments have distinct transition dynamics.Source: §3 · p. 3, §3.1 · p. 3
  • The shared structure G is the intersection of environment-specific dependency graphs and is assumed time-consistent and sparse.Sparseness bounds the number of state-action features on which each next-state feature depends.Source: §3.1 · p. 3
  • Each transition model factors into the shared causal transition model PG and environment-specific factors, with λ-sufficiency bounding their maximum L1 distance.λ = 0 means PG suffices to explain every environment's dynamics, while λ = 1 implies no shared structure.Source: §3.1 · pp. 3–4

3.2 A Class of Training Environments

  • The agent trains on a finite class M assumed diverse enough to reveal G and balanced enough for non-universal factors to even out.These conditions make the finite training class informative about the full universe.Source: §3.2 · p. 4
  • The training class is assumed diverse and even by design; selecting such a class through active interactions is left for future work.Source: §3.2 · p. 4

3.3 Learning Systematic Generalization

  • The causal structure G can be identified from purely observational data, without specific interventions.The bipartite graph directs edges from past state-action features to future state features.Source: §3.3 · p. 4
  • The analysis considers online estimation of G and PG using a generative model, and analogous offline results require convenient dataset-coverage assumptions.A generative model lets the agent set an MDP state before sampling a transition.Source: §3.3 · p. 4

4 Sample Complexity Analysis

The analysis establishes sample-complexity guarantees for systematic generalization and for estimating the causal structure and Bayesian network of discrete MDPs. Its main planning guarantee uses a learned shared causal transition model, with an unavoidable error term when that model cannot fully explain target-environment dynamics.

4 Sample Complexity Analysis

  • The sample-complexity analysis is a core contribution alongside the problem formulation, with causal-structure and Bayesian-network estimation also presented as results of independent interest.The section treats systematic-generalization complexity first, followed by the two estimation problems.Source: §4 · p. 4

4.1 Sample Complexity of Systematic Generalization with a Generative Model

  • A planning oracle uses the estimated causal transition model to approximately solve tasks across unknown MDPs and reward functions, under the target-transition factorization assumption.The model is estimated from interactions with a finite class of MDPs; the guarantee is subject to the unavoidable ϵλ term.Source: §4.1 · p. 5
  • For discrete MDPs, systematic generalization requires eO(MH6d3 SZ2n5Z+3) samples, reducing to eO(MH6S4A2Z2) for tabular MDPs, with an unavoidable error term ϵλ.The residual term reflects that the shared model PG may not fully explain each environment's dynamics, even when estimated exactly.Source: §4.1 · p. 5

4.2 Sample Complexity of Learning the Causal Structure of a Discrete MDP

  • Causal-structure estimation targets the ϵ-dependency subgraph, retaining dependencies whose distance from independence is at least ϵ, and provides a sample-complexity bound for discrete MDPs.The analysis adapts causal-discovery results to MDP structure; its tabular case is also stated as a corollary.Source: §4.2 · pp. 5–6

4.3 Sample Complexity of Learning the Bayesian Network of a Discrete MDP

  • For Bayesian-network estimation, a Z-sparse MDP graph permits polynomial sample complexity, while a dense graph can make parameter estimation mostly intractable.The analysis refines known estimation results for the specific MDP structure; a tabular corollary is also provided.Source: §4.3 · p. 6

5 Numerical Validation

The numerical validation uses a synthetic setting in which environments represent people and actions affect weight and academic performance. Across the reported quantities, estimation improves with more samples, though transition-model error remains above zero under the stated assumption mismatch.

Experimental Setup

  • The synthetic experiment treats each environment as a person and models how training, sleep, study, diet, and caffeine relate to weight and academic performance.It uses three environments to estimate the causal transition model and reports averages and standard deviations over 10 repetitions.Source: §5 · p. 6

Causal Structure Estimation

  • As sample counts increase, causal-structure edit distance converges to zero and the exact graph can be recovered.The structure is estimated from independence tests in each training environment, followed by intersecting their edges.Source: §5 · p. 6, Figure 4 · p. 7

Causal Transition Model Estimation

  • The causal transition-model L1-distance shrinks toward 0.05 as samples grow, attributed to environments not fully respecting the evenness assumption.Figure 4b plots this error against the number of samples.Source: §5 · pp. 6–7

Value Function Estimation

  • With a reasonable number of samples, the algorithm approximates the optimal value function for an unseen environment up to a small error.The value-function comparison is shown in Figure 4c.Source: §5 · p. 7

6 Related Work

The related-work discussion positions the paper at the intersection of causal model-based RL, factored MDPs, and systematic generalization. It distinguishes its focus on learning shared causal structure across infinitely many MDPs from prior work addressing different structures or settings.

Hidden Structures in RL

  • Unlike prior hidden-structure RL work focused on latent state representations or other MDP structural properties, this paper targets causal structure shared across infinitely many MDPs while assuming state features are available.The authors describe their structural assumption as more general than the linear structures in the cited work.Source: §6 · p. 7

Model-Based RL

  • Earlier model-based RL theory generally analyzes value functions from learned models, while this paper targets estimation of shared transition dynamics across an infinite set of MDPs.Tarbouriech et al. (2020) also provide model-estimation guarantees, but for a single MDP.Source: §6 · p. 7

Factored MDPs

  • Prior factored-MDP work studies known or learned transition factorization, but the authors know of no work combining it with reward-free learning and systematic generalization.They identify factorization discovery and transition-factor estimation as challenges of that combination.Source: §6 · p. 7

Causal RL

  • Among causal model-based RL studies, the paper claims the first polynomial sample complexity for systematic generalization, while noting Zhang et al. (2020) uses different structural assumptions and yields non-overlapping results.Feng et al. (2022) study non-stationary environmental changes, whereas this work targets systematic generalization.Source: §6 · p. 7

A Proofs

The causal graph can be identified from observational data alone because its edge orientations are determined by time.

Proofs of Section 3.3

  • Observational data identify the causal graph entirely here because time determines edge orientations, beyond identifying its skeleton and v-structures.In the general case described, observational data alone identify only the Markov equivalence class; interventions are needed to determine orientations.Source: Proofs of Section 3.3 · p. 10

Proofs of Section 4.1: Causal Transition Model Estimation

These proofs establish sample-complexity guarantees for estimating causal transition models, first for a single discrete MDP and then for classes of MDPs. They control approximation error by separating causal-structure estimation, Bayesian-network estimation, and approximation from removing weak dependencies.

Approximation bound

  • Removing weakly dependent edges yields an ϵ-dependency subgraph whose Bayesian network is close in L1 distance to the network over the original Z-sparse graph.The proof uses the weak conditional-dependence bound for omitted edges and Z-sparseness to bound the aggregate difference.Source: Proofs of Section 4.1: Causal… · p. 10

Singleton class

  • For a singleton class of discrete MDPs, Algorithm 1 returns a Bayesian-network estimate whose L1 error is at least ϵ with probability at most δ, under the derived sample complexity.Source: Proofs of Section 4.1: Causal… · pp. 10–11

Class of discrete MDPs

  • For a class of M discrete MDPs, Algorithm 1 returns a causal transition model satisfying Pr(||P_hat_G − P_G||_1 ≥ ε) ≤ δ, with the sample-complexity guarantee derived under ε^(4/3) ≪ ε^2.K′ ensures Pr(||P_hat_G − P_Gε′||_1 ≥ ε/3) ≤ δ/2; Gε′ is the intersection of the causal structures of the single class components.Source: Proofs of Section 4.1: Causal… · pp. 11–12

Tabular class

Proofs of Section 4.1: Planning

The planning proofs bound the policy error incurred when planning with an estimated causal transition model. They relate model-estimation error and λ-sufficiency to the resulting value and planning guarantees, with a corresponding tabular specialization.

Theorem 4.3

  • Theorem 4.3 bounds the planning error for a policy returned by an oracle using the estimated causal transition model, in terms of model error and λ-sufficiency.The model approximation error is the quantity studied in Lemma 4.1, while λ is the λ-sufficiency parameter of the true transition model.Source: Proofs of Section 4.1: Planning · pp. 12–13

Tabular planning

  • Under the preconditions of Corollary 4.4, the tabular specialization uses ϵ_λ = 2λSAH^3 and ϵ′ = ϵ/(2SAH^3).The corollary states that Theorem 4.3 holds with these tabular parameters.Source: Proofs of Section 4.1: Planning · p. 15

Value estimation

  • The proof bounds value-function estimation error by decomposing policy loss into evaluation and optimization errors and recursively bounding the evaluation terms.The stated value-error lemmas apply with probability at least 1−δ for any reward function and policy under their preconditions.Source: Proofs of Section 4.1: Planning · pp. 13, 15

Proofs of Section 4.2

These proofs establish sample-complexity guarantees for learning a discrete MDP’s causal dependency graph. The tabular result follows by substituting the tabular dimensions into the general bound.

General discrete MDP

  • Algorithm 2 estimates the ϵ-dependency graph Gϵ with failure probability at most δ under a sample complexity stated as K = O(…).The proof controls individual independence-test failures and combines them with a union bound.Source: Proofs of Section 4.2 · p. 18

Proof strategy

  • The independence-testing argument allocates failure probability across node pairs and uses a provably efficient test for each pair.The proof chooses thresholds ϵ′ = ϵ and δ′ = δ/d_S^2 before combining the test bound with the union bound.Source: Proofs of Section 4.2 · p. 18

Tabular MDP

Proofs of Section 4.3

These proofs establish sample-complexity guarantees for estimating Bayesian-network transition models from a known causal structure. They derive the general discrete-MDP result using concentration and union bounds, then obtain a tabular result using a tighter binary-variable concentration bound.

Theorem 4.8

  • For a discrete MDP with causal structure G, Algorithm 3 returns a Bayesian network bPG with Pr(∥bPG − PG∥1 ≥ ϵ) ≤ δ, for ϵ > 0 and δ ∈ (0, 1).The proof sums K′ over any x and any j.Source: Proofs of Section 4.3 · pp. 18–19

Proof strategy

  • The general proof controls the aggregate error by applying union bounds across conditional distributions and using Z-sparseness to bound parent-set size.Source: Proofs of Section 4.3 · p. 19

Tabular result

  • For a tabular MDP, Corollary 4.9 specializes the Bayesian-network estimation guarantee using a tighter concentration bound for binary variables.The proof uses KL-divergence concentration and converts the resulting bound to an L1 guarantee.Source: Proofs of Section 4.3 · pp. 19–20

B Numerical Validation

The synthetic example defines a shared causal transition model and generates environment-specific dynamics by adding independent noise to its coefficients. A Bayesian network distinguishes shared causal edges from environment-specific dependencies.

B Numerical Validation

  • The shared equations model A and W as normal variables with means determined by the displayed relationships among the variables.The equations give µA = A + 0.2D + 0.5S + 0.8St − 0.8 and µW = W − 0.5D − 0.5P + 1.Source: B Numerical Validation · p. 20
  • Independent white noise is added to each coefficient in the shared equations to generate a specific environment’s transition model.This includes coefficients omitted from the displayed equations because they are set to zero.Source: B Numerical Validation · pp. 20–21
  • In the Bayesian network, grey edges are causal and shared across environments, while blue dashed edges are environment-specific dependencies.Source: Figure 4 · p. 21
Loading 2202.06545v3…