A Distribution Mapping Approach to Counterfactually Fair Reinforcement Learning

Jianhan Zhang; Jitao Wang; John D. Piette; Donglin Zeng; Chengchun Shi; Zhenke Wu

A Distribution Mapping Approach to Counterfactually Fair Reinforcement Learning

Machine Learning 2026-08-09 v1 Computers and Society Machine Learning Methodology

Authors: Jianhan Zhang , Jitao Wang , John D. Piette , Donglin Zeng , Chengchun Shi , Zhenke Wu

Abstract

Reinforcement learning (RL) seeks to optimize sequential decisions to maximize population-level benefits over time. However, when deployed in high-stakes settings such as healthcare, RL decisions might systematically restrict some subpopulation's access to valuable services in a manner contrary to the values and goals of stakeholders. Counterfactual fairness (CF) offers a promising framework to address this problem based on causal reasoning. This paper develops a data preprocessing algorithm that, when used in tandem with policy learning, enables CF in RL. Our algorithm relies on a novel quantile distribution mapping method for sequentially estimating the counterfactual states and rewards in the data preprocessing step, subsuming common additivity assumptions used for counterfactual prediction as a special case. We theoretically prove that the per-step level of counterfactual unfairness and infinite-horizon suboptimality gap can be bounded under mild regularity conditions. We also empirically test our algorithm in numerical experiments as well as in application to a real-world interventional digital health dataset.

Keywords

counterfactual fairness reinforcement learning distribution mapping data preprocessing causal reasoning healthcare

Cite

@article{arxiv.2608.08743,
  title  = {A Distribution Mapping Approach to Counterfactually Fair Reinforcement Learning},
  author = {Jianhan Zhang and Jitao Wang and John D. Piette and Donglin Zeng and Chengchun Shi and Zhenke Wu},
  journal= {arXiv preprint arXiv:2608.08743},
  year   = {2026}
}