Beating the Winner's Curse via Inference-Aware Policy Optimization

TL;DR

This paper addresses the winner's curse in policy learning by introducing inference-aware policy optimization, which explicitly trades off expected improvement against the statistical significance of improvement when evaluated with IPW. It derives a complete Pareto frontier between these objectives by solving a family of convex relaxations parameterized by a dual variable $\xi$, and shows that the Pareto-optimal policies can be traced as $\pi^*(\xi)$ from the observational policy toward the mean-maximizing policy. A practical pipeline is proposed: estimate counterfactual means and variances from ML models on training data, construct an estimated frontier $\hat{\Pi}_{\Lambda}$, and select a policy by choosing $\xi$ to meet a target improvement $\lambda$ and a minimum $z$-score $Z_{min}$ (e.g., $\xi=\frac{2\lambda}{Z_{min}^2}$). Simulations demonstrate that IAPO can consistently achieve statistically significant improvements with modest training data and can outperform several baseline policy optimization methods, especially in complex causal settings. The work provides both theoretical insights into the frontier structure and a scalable plug-in algorithm for inference-aware policy design with downstream evaluation in mind.

Abstract

There has been a surge of recent interest in automatically learning policies to target treatment decisions based on rich individual covariates. A common approach is to train a machine learning model to predict counterfactual outcomes, and then select the policy that optimizes the predicted objective value. In addition, practitioners also want confidence that the learned policy has better performance than the incumbent policy according to downstream policy evaluation. However, due to the winner's curse-an issue where the policy optimization procedure exploits prediction errors rather than finding actual improvements-predicted performance improvements are often not substantiated by downstream policy optimization. To address this challenge, we propose a novel strategy called inference-aware policy optimization, which modifies policy optimization to account for how the policy will be evaluated downstream. Specifically, it optimizes not only for the estimated objective value, but also for the chances that the policy will be statistically significantly better than the observational policy used to collect data. We mathematically characterize the Pareto frontier of policies according to the tradeoff of these two goals. Based on our characterization, we design a policy optimization algorithm that uses machine learning to predict counterfactual outcomes, and then plugs in these predictions to estimate the Pareto frontier; then, the decision-maker can select the policy that optimizes their desired tradeoff, after which policy evaluation can be performed on the test set as usual. Finally, we perform simulations to illustrate the effectiveness of our methodology.

Figures (2)

• Figure 1: The performance of the estimated Pareto frontier improves with the accuracy of the predictive model. We see this via increased performance of policies on the frontier as the size of the training dataset increases and decreased performance as the model complexity increases.
• Figure 2: Given sufficient training samples and a simple enough causal model, most policy optimization methods can extract significant policy improvements. However IAPO can do so consistently with fewer trainings samples and complex causal models.

Theorems & Definitions (20)

• Proposition 1
• Proposition 2
• Proposition 3
• Proposition 4
• Theorem 1
• Proposition 5
• Theorem 2
• Proposition 6
• Proposition 7
• Theorem 3