Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Causal Flow Q-Learning for Robust Offline Reinforcement Learning

Mingxuan Li, Junzhe Zhang, Elias Bareinboim

TL;DR

This paper tackles confounding biases in offline reinforcement learning with pixel-based observations by formulating a Confounded Markov Decision Process (CMDP) and deriving a confounding-robust objective. It then develops Causal Flow Q-Learning (CFQL), which combines a flow-based behavioral policy, a one-step target policy, an ensemble of critics, and a discriminator to estimate a worst-case Q target under confounding, enabling safe policy learning from biased offline data. Through extensive experiments on 25 pixel-based tasks, CFQL substantially improves over confounding-unaware baselines and can match or exceed state-of-the-art Gaussian-policy methods in several scenarios, including offline-to-online fine-tuning. The results demonstrate the practical impact of incorporating causal reasoning and flow-based policy modeling to mitigate unobserved confounding in visual RL, with broad implications for safe, robust decision-making from real-world, imperfect data.

Abstract

Expressive policies based on flow-matching have been successfully applied in reinforcement learning (RL) more recently due to their ability to model complex action distributions from offline data. These algorithms build on standard policy gradients, which assume that there is no unmeasured confounding in the data. However, this condition does not necessarily hold for pixel-based demonstrations when a mismatch exists between the demonstrator's and the learner's sensory capabilities, leading to implicit confounding biases in offline data. We address the challenge by investigating the problem of confounded observations in offline RL from a causal perspective. We develop a novel causal offline RL objective that optimizes policies' worst-case performance that may arise due to confounding biases. Based on this new objective, we introduce a practical implementation that learns expressive flow-matching policies from confounded demonstrations, employing a deep discriminator to assess the discrepancy between the target policy and the nominal behavioral policy. Experiments across 25 pixel-based tasks demonstrate that our proposed confounding-robust augmentation procedure achieves a success rate 120\% that of confounding-unaware, state-of-the-art offline RL methods.

Causal Flow Q-Learning for Robust Offline Reinforcement Learning

TL;DR

This paper tackles confounding biases in offline reinforcement learning with pixel-based observations by formulating a Confounded Markov Decision Process (CMDP) and deriving a confounding-robust objective. It then develops Causal Flow Q-Learning (CFQL), which combines a flow-based behavioral policy, a one-step target policy, an ensemble of critics, and a discriminator to estimate a worst-case Q target under confounding, enabling safe policy learning from biased offline data. Through extensive experiments on 25 pixel-based tasks, CFQL substantially improves over confounding-unaware baselines and can match or exceed state-of-the-art Gaussian-policy methods in several scenarios, including offline-to-online fine-tuning. The results demonstrate the practical impact of incorporating causal reasoning and flow-based policy modeling to mitigate unobserved confounding in visual RL, with broad implications for safe, robust decision-making from real-world, imperfect data.

Abstract

Expressive policies based on flow-matching have been successfully applied in reinforcement learning (RL) more recently due to their ability to model complex action distributions from offline data. These algorithms build on standard policy gradients, which assume that there is no unmeasured confounding in the data. However, this condition does not necessarily hold for pixel-based demonstrations when a mismatch exists between the demonstrator's and the learner's sensory capabilities, leading to implicit confounding biases in offline data. We address the challenge by investigating the problem of confounded observations in offline RL from a causal perspective. We develop a novel causal offline RL objective that optimizes policies' worst-case performance that may arise due to confounding biases. Based on this new objective, we introduce a practical implementation that learns expressive flow-matching policies from confounded demonstrations, employing a deep discriminator to assess the discrepancy between the target policy and the nominal behavioral policy. Experiments across 25 pixel-based tasks demonstrate that our proposed confounding-robust augmentation procedure achieves a success rate 120\% that of confounding-unaware, state-of-the-art offline RL methods.
Paper Structure (28 sections, 2 theorems, 15 equations, 7 figures, 3 tables, 2 algorithms)

This paper contains 28 sections, 2 theorems, 15 equations, 7 figures, 3 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Confounding Robust Policy Learning
  3. Causal Flow Q-Learning
  4. Experiments
  5. Experiment Setup.
  6. Baselines.
  7. Evaluation.
  8. Offline Performance
  9. Offline-to-Online Performance
  10. Hyperparameter Tuning
  11. Discriminator loss coefficient.
  12. Number of the Q-network Ensembles.
  13. Conclusion
  14. Related Work
  15. Off-policy Causal Reinforcement Learning.
  16. ...and 13 more sections

Key Result

Theorem 3.1

For a CMDP environment $\mathcal{M}$ with reward signals $Y_t \in [a, b] \subseteq \mathbb{R}$, fix a policy $\pi$. The state value function $V_{\pi}(s) \geq \underline{V_{\pi}}(s)$ for any state $s \in \mathcal{S}$, where the lower bound $\underline{V_{\pi}}(s)$ is given by as follows,

Figures (7)

  • Figure 1: Confounding biases in offline data causes performance loss. Left-middle: For pixel based tasks in offline RL datasets, expert actions are sampled based on true state vectors but presented with pixel observations during the offline learning stage. Right: SOTA offline RL algorithm performance drops sharply with pixel observations despite using image augmentations and strong neural encoders.
  • Figure 2: Causal diagram representing the data-generating mechanisms in a Confounded Markov Decision Process. Bi-directed arrows represent information used by the expert's policy but unobservable to the learner in the offline dataset.
  • Figure 3: Backup diagram for causal policy gradient.
  • Figure 4: Pixel-based OGBench Tasks.
  • Figure 5: Offline-to-online evaluation. Causal-FQL converges to near optimal success rate in three representative tasks selected from task types that are not solved near-optimal in offline evaluation. All results are averaged over 4 seeds.
  • ...and 2 more figures

Theorems & Definitions (7)

  • Example 1: Confounded Pixel-Based Observations
  • Definition 2.1
  • Example 2: Confounded Pixel Observations (continued)
  • Theorem 3.1
  • Definition 3.1
  • Theorem 3.2: Restatement of \ref{['thm:_3_1']}
  • proof