Provably Efficient Causal Model-Based Reinforcement Learning for Systematic Generalization

TL;DR

This paper proposes a causal, model-based reinforcement learning approach that achieves systematic generalization across diverse environments with provable efficiency, leveraging environment structure and limited reward-free interactions.

Contribution

It introduces a tractable causal framework and a simple algorithm with polynomial sample complexity for systematic generalization in environment families.

Findings

Algorithm guarantees bounded planning error with sub-optimality term.

Achieves polynomial sample complexity under structural assumptions.

Provides a partially positive answer to systematic generalization in RL.

Abstract

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.