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

There has been a surge of recent interest in automatically learning policies to target treatment decisions based on rich individual covariates. In addition, practitioners 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 evaluation. 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 estimate of the policy’s improvement passes a significance test during downstream policy evaluation. 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 estimates the Pareto frontier using machine learning models; 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.

💡 Research Summary

The paper tackles the pervasive “winner’s curse” problem that plagues data‑driven policy learning: when a predictive model is trained on historical data and then optimized, the optimizer can exploit systematic prediction errors, yielding a policy whose estimated performance is overly optimistic. Downstream evaluation—typically via inverse propensity weighting (IPW) or doubly robust estimators—then reveals a much lower true performance, undermining practitioner confidence. Existing work focuses on correcting the bias of the evaluation step but does not address the root cause: the optimization itself is blind to how the policy will later be evaluated.

To close this gap, the authors introduce Inference‑Aware Policy Optimization (IAPO), a framework that jointly optimizes two objectives: (1) the expected value of the IPW estimate of the policy’s improvement over the observational (baseline) policy, and (2) the expected z‑score of that IPW estimate, i.e., the probability that the improvement will be statistically significant in a downstream test. By treating these objectives as competing, the authors formulate a Pareto‑optimal trade‑off surface.

The theoretical contribution begins by fixing a target expected improvement λ and minimizing the variance of the IPW estimator under that constraint. This yields a family of convex programs whose Lagrangian can be solved in closed form, providing a parametric representation of a superset of the Pareto frontier. After pruning dominated solutions, the true Pareto frontier is obtained. Key insights emerge: the frontier is bounded by two extreme policies—one that maximizes statistical significance (the highest possible z‑score) and another that maximizes expected outcomes. Between them, any increase in expected outcome necessarily reduces the z‑score, establishing a strict trade‑off.

A striking result is that optimal policies on the frontier are stochastic. Rather than assigning each individual deterministically to a treatment, the optimal solution prescribes treatment probabilities that start at the observational policy’s randomization rates and gradually shift toward 0 or 1 as the data permit. High‑variance individuals are kept close to the baseline randomization, effectively down‑weighting them in the IPW estimator and reducing variance; low‑variance individuals receive more aggressive treatment assignments, boosting expected gain. This stochastic nature is essential for achieving high statistical power.

The authors also derive a simple closed‑form expression for the maximal achievable z‑score, enabling power calculations that tell practitioners whether a given test set is large enough to detect any meaningful improvement. Notably, the z‑score of the outcome‑maximizing policy can behave counter‑intuitively: under certain (though not pathological) conditions it may decline as sample size grows, because the policy pushes treatment probabilities far from the baseline, inflating IPW weights and variance.

Building on the theory, the IAPO algorithm proceeds in three steps: (1) train a predictive model on the training data; (2) compute the Pareto frontier using the derived parametric form and let the decision‑maker select a point reflecting their desired balance of expected gain versus confidence; (3) evaluate the chosen policy on a held‑out test set using standard IPW (or AIPW) methods. Crucially, the optimization stage may use test covariates but never the test outcomes or treatment assignments, preserving the integrity of downstream evaluation.

Empirical simulations compare IAPO against conventional plug‑in methods, naive IPW optimization, and recent bias‑correction techniques (e.g., Gupta et al., 2024; Andrews et al., 2024). Results show that while naive methods often produce policies with large estimated gains that fail to achieve statistical significance, IAPO consistently yields policies that, although sometimes sacrificing a modest amount of expected gain, achieve substantially higher z‑scores and pass significance tests. This demonstrates that incorporating inference considerations during policy learning can pre‑empt the winner’s curse.

In summary, the paper makes three major contributions: (1) formalizing the inference‑aware policy optimization problem; (2) providing a complete analytical characterization of the Pareto frontier between expected improvement and statistical significance; (3) delivering a practical algorithm that enables practitioners to select policies with provable confidence guarantees. The work opens new avenues for policy learning where downstream inference is baked into the optimization, and suggests extensions to multi‑treatment settings, non‑random observational policies, and real‑world applications in healthcare, education, and social services.