PaperlayerView exact revision on arXiv

Source-linked AI summary

Counterfactual Off-Policy Evaluation with Gumbel-Max Structural Causal Models

Michael Oberst, David Sontag

arXiv:1905.05824v3cs.LGstat.ML

TL;DR

The paper addresses the challenge of evaluating RL policies from observational data in high-risk settings, where aggregate off-policy estimates may be misleading and difficult to inspect. It introduces Gumbel-Max SCMs and a counterfactual off-policy procedure for generating and pairing alternative trajectories, demonstrating how this can flag dangerous policy behavior in synthetic sepsis management. The approach depends on choosing an SCM because the MDP alone does not identify counterfactual trajectories.

  • Problem

    Off-policy evaluation from observational data is difficult in high-risk RL applications, and aggregate estimates provide limited insight into which individual episodes would differ under a proposed policy.

  • Method

    The paper introduces Gumbel-Max SCMs and a post-hoc Monte Carlo procedure that generates counterfactual POMDP trajectories and pairs them with observed trajectories.

  • Results

    The synthetic sepsis example shows that counterfactual trajectory inspection can reveal dangerous implicit assumptions even when model-based and counterfactual off-policy evaluation appear optimistic.

  • Takeaways & Limitations

    Divergent counterfactual and observed outcomes can be used to flag individual trajectories for domain-expert review and identify potentially sub-optimal policy behavior.

  • Takeaways & Limitations

    The approach requires specifying an SCM because knowing the MDP is insufficient, and the sensitivity of counterfactual trajectories to the SCM choice remains open.

Abstract

from arXiv · show

We introduce an off-policy evaluation procedure for highlighting episodes where applying a reinforcement learned (RL) policy is likely to have produced a substantially different outcome than the observed policy. In particular, we introduce a class of structural causal models (SCMs) for generating counterfactual trajectories in finite partially observable Markov Decision Processes (POMDPs). We see this as a useful procedure for off-policy "debugging" in high-risk settings (e.g., healthcare); by decomposing the expected difference in reward between the RL and observed policy into specific episodes, we can identify episodes where the counterfactual difference in reward is most dramatic. This in turn can be used to facilitate review of specific episodes by domain experts. We demonstrate the utility of this procedure with a synthetic environment of sepsis management.

1. Introduction

The paper addresses counterfactual evaluation of RL policies in high-risk settings, where observational off-policy evaluation is difficult and potentially misleading. It proposes SCM-based counterfactual trajectories to identify episodes with substantial disagreement between observed and proposed policies for expert review.

  • Motivation: Observational off-policy evaluation of RL policies is challenging because of confounding, small effective sample sizes, and limited introspection.Safety concerns also make training RL agents directly on real patients impractical.
  • Motivation: The paper reframes policy evaluation as asking which individual observed episodes would have had different outcomes under the learned policy.This targets cases of disagreement rather than relying only on aggregate policy comparisons.
  • Approach: It develops a post-hoc method for model-based RL with discrete states that pairs counterfactual and observed trajectories for domain-expert sanity checks.The method uses a learned dynamics model and can incorporate additional information such as patient chart review.
  • Theoretical challenge: The paper identifies non-identifiability as a fundamental obstacle: even a fully specified finite MDP can correspond to SCMs with different counterfactual outcomes.Binary monotonicity identifies some counterfactuals, but the paper notes no analogous condition was known for categorical variables.
  • Contribution: The Gumbel-Max SCM satisfies the proposed counterfactual stability condition and supports Monte Carlo counterfactual trajectories for discrete probabilistic dynamics models.This allows counterfactual trajectories to be drawn post hoc from any discrete conditional probability distribution.

2. Preliminaries

The preliminaries define SCMs, interventions, and counterfactual distributions, then connect these concepts to POMDPs. They emphasize that counterfactual conclusions require assumptions beyond observed or interventional distributions, especially with stochastic categorical transitions.

  • Structural causal models: An SCM combines independent latent variables, structural functions, and observed variables whose parent relationships form a causal DAG.The latent-variable distribution and structural functions determine the induced data distribution.
  • Interventions: Interventions replace structural mechanisms, producing interventional distributions that describe expected outcomes under imposed actions.The do-operator denotes such interventions, including atomic assignments such as do(X_i = x).
  • Counterfactual distributions: Counterfactual distributions use observed assignments to infer latent-variable posteriors before propagating them through modified structural mechanisms.This is the retrospective counterpart to forward-looking interventional estimation.
  • Identification: Counterfactual quantities are generally not identifiable because multiple SCMs can agree with the same interventional distribution while producing different counterfactual estimates.Binary monotonicity restricts SCMs enough to identify quantities such as the Probability of Necessity.
  • POMDPs: An episodic POMDP can be represented as an SCM containing states, actions, observations, rewards, histories, and stochastic policy mechanisms.Extending counterfactual reasoning from deterministic transitions to stochastic categorical transitions requires assumptions about the functional relationships.

3. Gumbel-Max SCMs

The section develops counterfactual reward decompositions for POMDP policies and addresses their non-identifiability with categorical transitions. It introduces counterfactual stability and the Gumbel-Max SCM as a principled basis for generating counterfactual trajectories.

  • The proposed policy’s expected reward difference from the observed policy can be decomposed across observed episodes using counterfactual trajectories generated by an SCM.This enables episode-level inspection of where policy outcomes differ.
  • 3.1. Non-Identifiability of Categorical Counterfactuals: Non-deterministic categorical transitions admit multiple SCM parameterizations with identical transition and reward distributions but different counterfactual outcomes.Thus, the episode-level decomposition is not identified without additional structural assumptions.
  • 3.2. Counterfactual Stability: Counterfactual stability rules out categorical counterfactual outcomes whose relative likelihood fails to increase under the alternative intervention.It also implies that an outcome remains unchanged when only its own probability increases.
  • 3.2. Counterfactual Stability: For binary variables, counterfactual stability implies Pearl’s monotonicity condition.The paper states this relationship as Theorem 1.
  • 3.3. Gumbel-Max SCM: The Gumbel-Max SCM satisfies counterfactual stability and supports a Monte Carlo procedure for drawing counterfactual trajectories under a proposed policy.The construction is based on independent Gumbel variables and the Gumbel-Max trick for categorical sampling.

4. Related Work

Related work applies causal ideas to reinforcement learning, potential-outcome estimation, structure learning, and categorical sampling. This paper instead treats the counterfactual distribution itself as the primary object and gives the Gumbel-Max mechanism a causal interpretation.

  • Prior work uses causal inference to augment RL training in bandits and model-based reinforcement learning.The paper identifies Buesing et al. (2019) as the closest related work.
  • Unlike closely related work that approximates interventional draws, this paper treats the counterfactual distribution as the primary object of interest.
  • Potential outcomes are often used to predict quantities related to the interventional distribution, such as CATE, rather than individual counterfactual trajectories.
  • Structure-learning assumptions distinguish causal directions through differences in observational distributions, whereas this paper’s assumptions distinguish counterfactual distributions compatible with observed and interventional behavior.
  • Monotonicity has been formalized and used in epidemiology to estimate counterfactual effects, including water sanitation in Kenya.
  • Although Gumbel variables are established for categorical sampling, this paper proposes their causal and counterfactual interpretation.

5. Experiments

The sepsis experiment evaluates a target RL policy using learned-model and counterfactual methods, then uses paired individual trajectories to identify and inspect disagreements between observed and counterfactual outcomes.

  • Experimental setup: The experiment learns a finite MDP from 1000 observed trajectories with hidden glucose and diabetes state, then learns a target policy using Policy Iteration.Unobserved state variables and limited samples are expected to make the target policy perform poorly relative to the behavior policy.
  • Off-policy evaluation: Weighted importance sampling and model-based off-policy evaluation make the target policy appear superior, although newly simulated target-policy trajectories show it is actually inferior.The learned MDP is used as a generative model for model-based evaluation, while new simulator samples provide the target-policy outcome under deployment-like conditions.
  • Counterfactual analysis: Counterfactual off-policy evaluation also gives optimistic results, but pairing counterfactual and observed outcomes supports investigation of why the learned MDP favors the target policy.Counterfactual trajectories are generated with the same MDP and the Gumbel-Max SCM.
  • Identification of informative trajectories: Patients are assigned to nine categories using their observed outcome and the most frequent counterfactual outcome across five draws, highlighting trajectories for further analysis.The outcome categories are death, no change, and discharge under each observed-versus-counterfactual pairing.
  • Identification of informative trajectories: 10% of observed trajectories are cases where patients who died would most likely have lived under the target policy, providing a focused subset for expert review.The paper presents these cases as plausible targets for domain-expert investigation and illustrates one with observed and counterfactual trajectories.
  • Insights from individual trajectories: The illustrated failure occurs because the target policy halts treatment while glucose is dangerously low, which the learned MDP misses because glucose and diabetes are unobserved and rare.The learned MDP therefore assigns high discharge probability when the other three vitals are normal and all treatments are stopped.

6. Conclusion

The proposed technique is intended to help experts inspect RL policies in high-risk settings by flagging trajectories with divergent observed and counterfactual outcomes. Its use depends on specifying a structural causal model, since the MDP alone does not identify counterfactuals.

  • Conclusion: Five counterfactual draws per observed trajectory are used illustratively, while concentration arguments can quantify the draws needed for desired precision.The example specifically mentions quantities such as the probability of death.
  • Conclusion: Figure 5 illustrates the review process with an observed trajectory ending in death and five counterfactual trajectories ending in discharge.All counterfactual trajectories match the observed trajectory until actions diverge at t = 2; glucose has no counterfactual trajectory because it was excluded from the model.
  • Conclusion: The technique flags individual trajectories for domain-expert review when their observed and counterfactual outcomes diverge.The sepsis example uses this process to identify sub-optimal and potentially dangerous behavior in a proposed policy.
  • Conclusion: Applying the method requires a specified structural causal model because knowing only the MDP is insufficient for counterfactual analysis.The paper identifies sensitivity to the choice of SCM as an open follow-up question.

A. Omitted Proofs

The omitted proofs establish counterfactual reward decomposition, connect counterfactual stability to binary monotonicity, and prove that the Gumbel-Max SCM satisfies stability.

  • Counterfactual decomposition: The counterfactual reward decomposition rewrites target-policy expected reward as an aggregation over observed trajectories and their posterior counterfactual outcomes.The derivation changes variables from exogenous noise to trajectories while preserving the shared-noise posterior.
  • Counterfactual decomposition: Counterfactual trajectories under a target policy are deterministic functions of the policy and shared exogenous variables U.The observed trajectory constrains the posterior over U, enabling counterfactual evaluation under another policy.
  • Theoretical guarantees: Counterfactual stability implies monotonicity for binary outcomes and treatments.The proof shows that the discordant counterfactual event has probability zero under the stability condition.
  • Theoretical guarantees: The Gumbel-Max SCM satisfies the counterfactual stability condition.The proof uses fixed Gumbel variables across interventions and monotonicity of the logarithm to establish the required ordering property.

B. Non-Identifiability Example

Figure 6 visually depicts the unidentifiability example presented in Section 3.1.

  • Figure 6 provides a visual depiction of the unidentifiability example from Section 3.1.

C.1. Sepsis Simulator

The appendix points readers to the released code for reproducibility and further details about the simulator setup.

  • C.1. Sepsis Simulator: The experiments’ code, including the appendix figures, is available online for reproducing the results.The repository is also referenced for more in-depth information about the simulator setup.

C.2. Impact of hidden state

The hidden-state analysis examines how withholding glucose and diabetes variables affects off-policy evaluation and policy quality. It compares repeated training and held-out evaluations with partially observed versus full-state dynamics models, while noting several interacting sources of uncertainty.

  • C.2. Impact of hidden state: The main experiments hide glucose and diabetes state from the dynamics model used for the RL policy.The appendix evaluates the consequences of this modeling choice for off-policy estimates and policy quality.
  • C.2. Impact of hidden state: The synthetic results may reflect confounding, MDP-learning sample complexity, and small training and held-out sample sizes.The authors state that the relative strength and interaction of these factors cannot be determined conclusively; observing all state information increases the state space from 144 to 1440 states.
  • C.2. Impact of hidden state: Figure 7 reports medians and 95% intervals across 100 independently resimulated evaluations with separate training and held-out WIS datasets.Each evaluation uses 1000 training episodes and 1000 held-out episodes; WIS is computed on both sets.
  • C.2. Impact of hidden state: Figure 8 repeats the Figure 7 setup while allowing the dynamics model used by the MDP to observe the full state.The corresponding results are also reported in Table 1.
  • C.2. Impact of hidden state: A more comprehensive simulated environment could disentangle these factors and support broader evaluation of off-policy methods.The authors identify this as a direction for future work.
Loading 1905.05824v3…