A Distribution Mapping Approach to Counterfactually Fair Reinforcement Learning

Jianhan Zhang, Jitao Wang, John D. Piette, Donglin Zeng (University of Michigan), Chengchun Shi (LSE), and Zhenke Wu (University of Michigan)

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.

The paper is motivated by the PowerED study, which aimed to evaluate an RL-based intervention designed to reduce patients' self-reported opioid analgesic misuse behaviors while conserving counselors' time. In this setting, a fairness-unaware RL agent might systematically decline to assign live counselor calls to patients of certain races, essentially denying those patients' access to high-quality, human-based care. Unfairness can arise even if the sensitive attribute is excluded from the state variables because state and reward variables often encode implicit information about the sensitive attribute.

The paper makes three major contributions:

1. A data-preprocessing algorithm that ensures CF in RL, extending the algorithm introduced in Wang et al. [31] to a strictly larger class of underlying contextual Markov decision processes where the additive noise assumption is allowed to be violated.

2. Theoretical guarantees for the proposed algorithm by deriving bounds for the suboptimality gap and counterfactual unfairness resulting from the policy learned using the proposed algorithm.

3. Empirical testing of the proposed algorithm using numerical experiments and real data analysis.

The paper develops a novel data preprocessing method called counterfactually fair sequential marginal distribution mapping (CFSMDM) which estimates the requisite counterfactuals to construct CF policies. CFSMDM achieves CF in policy learning via a two-stage procedure:

1. Preprocessing stage: Preprocess the training trajectories to obtain augmented training trajectories where the states are estimates S̃̂t of S̃t and the rewards are estimates R̃̂t of R̃t. This step removes information of the sensitive attribute from the training trajectories.

2. Policy learning stage: Perform policy learning on the augmented training trajectories using existing offline RL algorithms, such as fitted Q iteration (FQI).

The 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.

Assumption 1 (Strict monotonicity of the state in the noise variable): For t = 0 and i = 1,..., M, S0(i) = f0 S(Z, U0 S(i)) for some function f0 S that is deterministic and strictly increasing in U0 S(i) on the support of U0 S(i) for all levels of Z. For t ≥ 1 and i = 1,..., M, St(i) = f S(Z, At−1, St−1, Ut S(i)) for some function f S that is deterministic and strictly increasing in Ut S(i) on the support of Ut S(i) for all levels of (Z, At−1, St−1).

Assumption 2 (Strict monotonicity of the reward in the noise variable): For t ≥ 0, Rt = f R(Z, St, At, Ut R) for some function f R that is deterministic and strictly increasing in Ut R on the support of Ut R for all levels of (Z, At, St).

These assumptions are strictly weaker than the additivity assumption made by the alternative method CFSDP proposed by Wang et al. [31]. Similar monotonicity assumptions are common in quantile treatment effect literature.

Theorem 1 (Quantile matching for counterfactual states): Under Assumption 1, for an individual with observed history ht = (z, s̄t, āt−1, r̄t−1), the quantile level of the counterfactual state with respect to the law of the state under the alternative sensitive attribute value does not vary by z′ ∈ Z.

Theorem 2 (Quantile matching for counterfactual rewards): Under Assumption 2, for an individual with observed history ht, the quantile level of the counterfactual reward with respect to the law of the reward under the alternative sensitive attribute value does not vary by z′ ∈ Z.

Theorem 3 (Suboptimality gap bound): Under Assumptions 1-2 and A1-A3, for any ω ∈ N and z ∈ Z:

SG(π̃nᶻ) ≤ 2BL(1-γω)/(1-γ) max t≤ω-1 M t,n + 2BLγω/(1-γ) M S̃ + 2C1dLR maxδ T,n /(λ0(λ0-4dLδ T,n)(1-γ)3) + 2C2dR maxκlog(n)/((1-γ)2λ0√n) + 2γ B R max/((1-γ)3λ0)

with probability at least (1-n(-κ) - d exp(-nλ0/8)) P(δ T,n ≤ λ0/(4dL)).

Theorem 4 (Counterfactual unfairness bound): Under Assumptions 1-2 and A1-A4, for any z′, z″ ∈ Z and t ≥ 0:

E[∥π̃ t,n z′(h̃t) - π̃ t,n z″(h̃t)∥2] ≤ E[O(ξn α) + O((BLδ t,n) α)] + 2[P(δ T,n > λ0/(4dL)) + ηn]

where ξn = C1dLR maxδ T,n /(λ0(λ0-4dLδ T,n)(1-γ)2) + C2dR maxκlog(n)/((1-γ)2λ0√n) + γ B R max/(1-γ) and ηn = n(-κ) + d exp(-nλ0/8).

The paper compares CFSMDM against several baseline methods:

• Full: Standard method using all available covariates including the sensitive attribute

• Unaware: Method using all covariates except the sensitive attribute

• Random: Random policy selecting each action with equal probability

• FLAP M: FLAP algorithm with single-stage marginal distribution mapping adapted to RL

• ECOCF M: Method from Wang et al. [32] adapted to RL

• CFSDP: Method from Wang et al. [31] using neural network for transition kernel learning

Experiment 1 (Nonadditive noise): Under a CMDP with nonadditive noise, CFSMDM outperforms all baselines other than Random in terms of fairness control. Its CF metric decreases as the training sample size N increases, supporting consistency. Full achieves the highest policy value, while CFSMDM demonstrates a tradeoff between fairness and policy value. FLAP M and ECOCF M are unfairer than CFSMDM due to their failure to account for the sequential nature of the environment. CFSDP is relatively unfair and achieves lower policy value due to violation of the additive noise assumption.

Experiment 2 (Varying δ): CFSMDM remains relatively fair as more unfairness is injected into the underlying CMDP by increasing δ, showing robustness against the level of unfairness in training trajectories.

The dataset consists of 207 patients, each participating in the intervention over 12 weeks (one RL decision per week). The analysis considers sex, education, age, and ethnicity (all binarized) as potential sensitive attributes, with weekly pain score and weekly pain interference score as state variables. The reward is based on weekly self-reported opioid medication risk score.

Results show that:

• Full has the worst level of counterfactual fairness for all candidate sensitive attributes

• Random is perfectly counterfactually fair but has low policy values

• CFSMDM achieves the best counterfactual fairness for most sensitive attributes

• CFSMDM noticeably outperforms CFSDP for education and ethnicity

• For sex, CFSMDM is slightly more unfair than CFSDP and Unaware, possibly because excluding the sensitive attribute already removes most unfairness

The paper introduces CFSMDM, a novel and general algorithm of sequential data preprocessing which, when used in tandem with off-the-shelf RL policy learning algorithms, achieves counterfactual fairness. Compared to CFSDP, CFSMDM requires strictly weaker assumptions and is therefore more general and robust. However, CFSDP is easier to implement and less computationally heavy than CFSMDM because it only estimates the means of the conditional counterfactual state and reward distributions, while CFSMDM requires fitting multiple quantile models corresponding to a grid of quantile levels.

Improvement 1: Fair Sequential Decision-Making Under Non-Additive Noise

• What to implement: Integrate CFSMDM's quantile-based distribution mapping into offline RL pipelines (e.g., FQI, DQN) to replace mean-only counterfactual estimation. This removes the additive noise assumption, enabling fair policies in real-world systems where noise interacts nonlinearly with sensitive attributes (e.g., healthcare, lending, hiring).

• Capability gained: The AI can now guarantee counterfactual fairness (equalized outcomes across sensitive groups) even when the underlying data-generating process violates additivity—a common failure mode in practice. It will produce policies that do not systematically disadvantage subgroups (e.g., racial minorities) while maintaining competitive policy value.

Improvement 2: Sequential Counterfactual Data Augmentation for Training

• What to implement: Use the two-stage preprocessing (quantile matching for states and rewards) to generate augmented trajectories where sensitive-attribute information is removed. Feed these into any existing offline RL algorithm without modifying the policy learner.

• Capability gained: The AI system becomes plug-and-play fair—any RL agent (e.g., for chronic disease management, resource allocation) can be trained on counterfactually clean data, eliminating the need for custom fairness constraints during optimization. This reduces engineering overhead and makes fairness accessible to non-specialists.

Improvement 3: Robust Fairness with Theoretical Suboptimality Guarantees

• What to implement: Adopt the theoretical bounds (Theorems 3–4) as a safety certificate. Use the suboptimality gap and unfairness bounds to automatically tune the number of quantile levels (grid size) and training sample size, ensuring the learned policy stays within a user-specified fairness–performance trade-off.

• Capability gained: The AI can now provide provable guarantees on both fairness violation and performance loss before deployment. This is critical for high-stakes domains (e.g., clinical trials, criminal justice) where regulators require auditable fairness and bounded regret.

Improvement 4: Fairness-Aware Exploration in Online RL

• What to implement: Extend CFSMDM's preprocessing to online settings by maintaining a rolling buffer of counterfactually mapped trajectories. Use the quantile matching to dynamically re-weight or filter new experiences, preventing the agent from learning biased policies during exploration.

• Capability gained: The AI system can learn fair policies in real-time from streaming data, adapting to distribution shifts while preserving counterfactual fairness. This enables deployment in dynamic environments (e.g., personalized education, adaptive clinical dosing) where fairness must be maintained as new patient data arrives.

Improvement 5: Multi-Sensitive-Attribute Fairness with Minimal Overhead

• What to implement: Parallelize the quantile mapping across multiple sensitive attributes (e.g., race, sex, age) by fitting separate quantile models per attribute and combining them via a joint counterfactual distribution. Use the theoretical results to bound the cumulative unfairness across attributes.

• Capability gained: The AI can now enforce fairness simultaneously across several protected characteristics without exponential computational cost. This is essential for real-world systems where discrimination is intersectional (e.g., Black women, elderly immigrants), and the system can flag which attribute contributes most to unfairness for targeted intervention.

Improvement 6: Interpretable Fairness Diagnostics for Stakeholders

• What to implement: Expose the intermediate quantile-level estimates (from Theorems 1–2) as a diagnostic tool. For each decision, output the counterfactual state/reward distribution under alternative sensitive attributes, allowing clinicians or policymakers to visually inspect where unfairness arises (e.g., For a Black patient with the same history, the predicted reward is 20% lower at the 75th percentile).

• Capability gained: The AI system becomes transparent and auditable—stakeholders can verify that fairness is achieved not just on average but across the entire distribution of outcomes. This builds trust and facilitates regulatory approval for AI-assisted decisions in medicine and social services.

Improvement 7: Fairness-Preserving Transfer Learning

• What to implement: Use CFSMDM's preprocessed trajectories as a shared representation for transfer learning. Train a base policy on counterfactually clean data from one domain (e.g., clinic A) and fine-tune on a new domain (e.g., clinic B) with minimal additional data, while re-running the quantile mapping only on the new domain's sensitive-attribute distributions.

• Capability gained: The AI can deploy fair policies across heterogeneous populations (e.g., different hospitals, countries) without re-collecting massive datasets. This reduces data requirements by 10–100× while maintaining fairness guarantees, enabling rapid scaling of fair AI in resource-limited settings.

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.