Set-Valued Policy Learning

TL;DR

This paper introduces set-valued policy learning for multiple treatments, providing uncertainty quantification and robust recommendations, demonstrated through synthetic and real IVF data.

Contribution

It develops a novel set-valued policy framework with conformal methods and a greatest lower bound approach for multiple treatments, enhancing robustness and interpretability.

Findings

Set-valued policies improve decision robustness in treatment selection.

The methods achieve reliable coverage without strong assumptions.

Application to IVF data shows practical utility and improved decision confidence.

Abstract

Conventional treatment policies map patient covariates to a single recommended intervention in order to maximize expected clinical outcomes. Although a rich body of causal inference methods has been developed to estimate such policies, point-valued recommendations can be highly sensitive to estimation uncertainty, model specification, and finite-sample variability, while typically providing little guidance about how confident one should be in the recommended action. In this work, we propose a set-valued policy learning paradigm for the multiple-treatment setting, in which policies output a set of plausible treatments rather than a single recommendation. This formulation enables intrinsic uncertainty quantification, with the size of the predicted set reflecting the degree of decision ambiguity. We extend the learning-to-defer framework to multiple treatments via a novel \textit{greatest Lower Bound} method, and introduce \textit{conformal policy learning}, which bridges the gap between unobserved ground-truth optimal treatments and estimated optimal treatment rules. Drawing on insights from the noisy-label literature, we develop a randomness-injection approach that guarantees marginal coverage without requiring assumptions on underlying black-box optimal treatment rules. Through experiments on synthetic data and a real-world application to In-Vitro Fertilization (IVF), we demonstrate that our methods produce robust and actionable policies that naturally incorporate clinical considerations while effectively balancing performance and reliability.