Technical Summary: Set-Valued Policy Learning

Problem Statement

Conventional treatment policy learning in precision medicine typically maps patient covariates to a single recommended intervention to maximize expected clinical outcomes. However, this point-valued approach faces significant limitations:

1. Sensitivity to Uncertainty: Recommendations are highly sensitive to estimation uncertainty, model specification, and finite-sample variability.
2. Lack of Confidence Guidance: Standard policies provide little indication of how confident one should be in the recommended action.
3. Ambiguity Ignorance: When multiple treatments yield comparable outcomes, point-valued policies fail to capture this ambiguity, potentially conflicting with practitioner preferences regarding unquantified costs or adverse events.
4. Ground Truth Unavailability: In causal inference, only the outcome of the administered treatment is observed; counterfactual outcomes (and thus the true optimal treatment) are unobserved, complicating the construction of policies with formal guarantees.

The authors propose a set-valued policy learning paradigm for multi-treatment settings ($K > 1$). Instead of a single recommendation, the policy outputs a subset of plausible treatments. The primary objective is to construct policies $C$ such that the set contains at least one optimal treatment with high probability:
$$P(\Pi^\star(X) \cap C(X) \neq \emptyset) \geq 1 - \alpha$$
where $\Pi^\star(X)$ is the set of optimal treatments for covariates $X$, and $\alpha$ is a user-specified error level. This shifts the model's role from prescribing a decision to supporting it by providing a filtered set of options, naturally quantifying uncertainty via set cardinality.

Methodology

The paper introduces two primary frameworks for learning set-valued policies with formal coverage guarantees:

1. Greatest Lower Bound (GLB) Method

This approach extends the "learning-to-defer" framework to multiple treatments using conditional mean outcome estimators.

• Mechanism: It relies on trusted lower ($\ell_n$) and upper ($u_n$) confidence bounds for the conditional mean outcomes $\mu(x, a)$.
• Construction:
  1. Identify the treatment $a_{maxmin}$ that maximizes the lower bound: $a_{maxmin}(x) \in \arg\max_a \ell_n(x, a)$.
  2. Construct the set $C_\alpha(x)$ by including all treatments whose upper bound exceeds the greatest lower bound benchmark:
     $$C_\alpha(x) = \{a \in A : u_n(x, a) \geq \ell_n(x, a_{maxmin}(x))\}$$
• Guarantee: Under the assumption that valid confidence bounds exist, this method achieves conditional coverage, implying marginal coverage.
• Limitation: It requires the non-trivial construction of valid bounds and is sensitive to estimation inaccuracies.

2. Conformal Set-Valued Policy Learning (Noisy Label Regime)

Recognizing that optimal treatments are unobserved, the authors frame policy learning as a classification task under a noisy label regime, where estimated optimal treatments serve as proxies for ground truth.

• Meta-Algorithm:

  1. Black-box Label Generation: Train a policy learner (e.g., Q-learning, policy trees) on a subset of data to generate noisy labels $\hat{A}^*$ (estimates of optimal treatments).
  2. Nonconformity Score: Learn a score function $s_n(x, a)$ (e.g., margin-based) using a training subset to quantify the alignment of a candidate treatment relative to the calibration data.
  3. Noisy Calibration: Compute scores on a calibration set using the noisy labels $\hat{A}^*$ and derive a quantile threshold $\hat{q}_{1-\alpha}$.
  4. Prediction Set: Define $C_\alpha(x) = \{a : s_n(x, a) < \hat{q}_{1-\alpha}\}$.

• Addressing Coverage Gaps: Standard conformal prediction on noisy labels does not guarantee coverage for the true optimal treatments without assumptions. The authors propose two strategies to ensure marginal coverage:

  1. Rank-Preserving Assumption: Assumes the estimator preserves the ordering of treatment means. If satisfied, coverage is guaranteed.
  2. Randomness Injection: A novel approach to enforce stochastic dominance without strict structural assumptions.
     • Perturb the noisy labels by mixing them with purely random labels with probability $r$: $\hat{A}^*_{r,i} = R_i \cdot A^{rd}_i + (1-R_i) \cdot \hat{A}^*_i$.
     • This controlled randomness increases nonconformity scores stochastically, ensuring the prediction sets are conservative enough to cover the true optimal treatments. Theorem 1 provides a condition on $r$ to guarantee coverage.

Evaluation Metrics

The paper introduces Set-Policy Value to evaluate these policies, as standard policy value metrics are insufficient for sets.

• Choice Function ($\delta$): A mapping selecting a single treatment from the set.
• Uniform Set-Policy Value ($\bar{V}_P(C)$): The expected outcome if a treatment is selected uniformly at random from the set. This serves as a lower bound for the value achieved by a decision-maker who performs better than random selection.

Key Results

Synthetic Experiments

• Coverage: The experiments demonstrate that without randomness injection ($r=0$), coverage guarantees are often violated due to estimation errors. Introducing a moderate level of randomness (e.g., $r=0.2$) successfully restores marginal coverage.
• Trade-offs: Increasing $r$ increases the mean cardinality of the predicted sets (making them larger and more conservative) to ensure coverage. The GLB method also achieves coverage but tends to over-cover due to sensitivity to estimation errors.
• Value: Higher confidence levels (lower $\alpha$) or higher randomness levels are required to maintain high set-policy values, balancing the trade-off between set size (informativeness) and coverage reliability.

Real-World Application: In-Vitro Fertilization (IVF)

• Dataset: 18,538 ovarian stimulation cycles with ordinal gonadotropin dosages and outcomes (follicular yield and estradiol levels).
• Findings:
  • Set-valued policies successfully filter extreme doses, preserving the ordinal structure of treatments.
  • Decision Support: The framework allows for the integration of secondary objectives. For instance, a "minimal dose" strategy ($\delta_{lower}$) applied to the set-valued output effectively controlled estradiol levels (reducing OHSS risk) with only a marginal reduction in follicular yield compared to uniform sampling.
  • Cardinality: As in synthetic settings, higher randomness levels led to larger sets, offering safer but potentially more burdensome decision spaces.

Significance and Claims

The paper claims to be the first to learn set-valued policies with formal coverage guarantees within the causal policy learning framework. Its significance lies in:

1. Paradigm Shift: Moving from "prescribing" a single decision to "supporting" decision-makers with a set of plausible options, thereby generalizing the "learning-to-defer" framework to multi-treatment settings.
2. Intrinsic Uncertainty Quantification: The size of the predicted set naturally reflects the degree of decision ambiguity, providing clinicians with actionable confidence levels.
3. Model Agnosticism: The conformal approach is model-agnostic and can be combined with arbitrary black-box policy learning methods without requiring specialized robustness modifications.
4. Handling Unobserved Ground Truth: By leveraging insights from the noisy-label literature and introducing randomness injection, the method bridges the gap between unobserved optimal treatments and estimated rules, providing rigorous coverage guarantees without relying on the robustness properties of the underlying estimator.
5. Clinical Relevance: The IVF application demonstrates that these methods produce robust, actionable policies that can incorporate clinical constraints (e.g., dose de-escalation) and balance performance with safety.

The authors conclude that reliable individualized decision-making requires methods that remain trustworthy in the presence of statistical uncertainty and modeling error. Their framework supports rather than replaces human decision-makers, facilitating context-aware choices that incorporate secondary objectives and clinical judgment.