Model-Based Reinforcement Learning Under Confounding

TL;DR

This paper develops a model-based reinforcement learning approach for confounded environments with unobserved context, using causal inference techniques to enable consistent policy evaluation and planning.

Contribution

It introduces a novel method combining proximal off-policy evaluation and MaxCausalEnt to handle confounding in model-based RL with unobserved contexts.

Findings

Provides a surrogate MDP framework for confounded environments

Enables consistent policy evaluation under unobserved confounding

Integrates causal inference with RL model learning

Abstract

We investigate model-based reinforcement learning in contextual Markov decision processes (C-MDPs) in which the context is unobserved and induces confounding in the offline dataset. In such settings, conventional model-learning methods are fundamentally inconsistent, as the transition and reward mechanisms generated under a behavioral policy do not correspond to the interventional quantities required for evaluating a state-based policy. To address this issue, we adapt a proximal off-policy evaluation approach that identifies the confounded reward expectation using only observable state-action-reward trajectories under mild invertibility conditions on proxy variables. When combined with a behavior-averaged transition model, this construction yields a surrogate MDP whose Bellman operator is well defined and consistent for state-based policies, and which integrates seamlessly with the maximum causal entropy (MaxCausalEnt) model-learning framework. The proposed formulation enables principled model learning and planning in confounded environments where contextual information is unobserved, unavailable, or impractical to collect.