Data-Automated Policy Learning for Nonlinear Welfare
Published 1 Jun 2026 in econ.EM, stat.ME, and stat.ML | (2606.01659v1)
Abstract: This paper explores policy learning from observational data, focusing on a nonlinear welfare criterion in a binary treatment setting. The nonlinear criterion is inspired by scenarios where policymakers prioritize specific population segments. We model this criterion using a utility function that encompasses potential outcomes and intermediate parameters, with the latter capturing higher moments of the outcome distributions. When formulated in the context of observational data, both the intermediate parameters and the welfare criterion depend on the propensity score, which we estimate using machine-learning techniques. To address bias in machine learning estimates, we introduce a novel reweighting-based debiasing approach that offers a promising alternative to traditional orthogonality-based methods. To tackle the complexities of infinite-dimensional policy spaces, we employ sieve approximations and $K$-fold cross-validation for model selection, thereby fully automating the policy-learning process. Despite these complexities, we demonstrate that both the welfare regret and the average welfare regret of our proposed policy learning method satisfy an oracle inequality, thereby providing theoretical guarantees on the performance of the estimated policy relative to the best possible policy. This finding extends the existing results from linear to nonlinear welfare criteria, from finite-dimensional to infinite-dimensional policy spaces, and from a known propensity score to a machine-learned one.
The paper introduces a flexible framework to optimize policy decisions under nonlinear welfare objectives, moving beyond average treatment effects.
It leverages sieve approximations and cross-validation to navigate infinite-dimensional policy spaces and mitigate overfitting.
The approach uses covariate-balanced reweighting for debiasing, achieving robust oracle guarantees and empirically validated performance.
Data-Automated Policy Learning for Nonlinear Welfare
Introduction and Context
"Data-Automated Policy Learning for Nonlinear Welfare" (2606.01659) addresses the challenge of optimal policy learning from observational data under a genuinely nonlinear welfare criterion in binary treatment settings. Rather than adopting the standard linear welfare objectives that simply aggregate average treatment effects, this paper introduces a general framework that flexibly considers distributional objectives—such as inequality, risk, or tail-based welfare measures—relevant to policymakers targeting specific subpopulations or outcome characteristics. The framework is motivated by numerous applications, including income inequality, risk management, and public health, where nonlinearity arises due to interests in quantiles, coefficients of variation, or tail statistics.
Methodological Contributions
Nonlinear Welfare Objective
The paper formalizes the welfare criterion as a general, policy-dependent functional:
W(π)=E[U(Y∗(π(X)),X,B∗(π))],
where U is a utility function, and B∗(π) represents intermediate parameters such as moments or quantiles of the induced potential outcome distribution. This criterion encompasses nonlinear objects, e.g., Gini coefficients, upper/lower tail ratios, or conditional value-at-risk (CVaR). The nonlinearity critically enters because these intermediate statistics (like quantiles) are themselves functions of potentially complex policy-induced distributions.
Infinite-Dimensional Policy Space
To accommodate rich policy classes (including decision trees and deep nets), the authors leverage sieve approximations, generating a hierarchy of finite-dimensional subsets that densely approximate the infinite-dimensional policy space. Within each subset, empirical policy optimization is performed.
Data-Driven Policy Selection
A K-fold cross-validation (CV) approach is adopted for automatic data-driven model selection. By partitioning the data and penalizing policy class complexity, the procedure identifies the policy class that optimizes out-of-sample welfare, mitigating overfitting risk.
Debiasing via Covariate-Balanced Reweighting
Addressing the well-known challenge of bias from machine-learned nuisance components (propensity scores and conditional models), the paper introduces a reweighting-based debiasing method inspired by covariate balancing literature. Unlike the more common Neyman orthogonalization/double debiasing techniques, this method constructs calibrated entropy-minimizing weights that enforce covariate balance and eliminate first-order estimation biases in both the intermediate statistics and the ultimate welfare criterion. The approach provides both finite-sample and asymptotic bias control under mild regularity conditions.
Theoretical Results
The central theoretical guarantee is an oracle inequality bounding the (average) welfare regret of the learned policy:
where π∗ denotes the optimal policy, π^ the learned policy, and Πl the l-th sieve subclass of policies with VC dimension VC(Πl). This result is robust:
It applies to nonlinear welfare criteria,
Requires only machine-learned (not known) propensity scores,
Handles infinite-dimensional policy spaces, provided these are well-approximated by finite-dimensional sieves.
Strong uniform convergence rates of all nuisance estimators (propensity scores and conditional means/quantiles) are established under high-level regularity on DNNs and moment functions. Theoretical analysis extends existing results for linear policy learning (e.g. [Mbakop & Tabord-Meehan, 2021]; [Athey & Wager, 2021]) to these substantially more general welfare settings.
Empirical Illustration
The procedure is empirically validated using the U.S. National Job Training Partnership Act (JTPA) Study, a canonical policy evaluation dataset. The goal is to allocate training offers to maximize a mean-to-standard-deviation ratio of post-training earnings, a nonlinear "efficiency/equity" welfare objective.
Policies are constrained to be monotone cutoffs in pre-program earnings given education, with a sieve of piecewise-linear cutoffs used for approximation. The learned nonlinear-optimal policy exhibits:
Mean outcome: 16,132.78
Standard deviation: 16,617.18
Relative to a benchmark linear-regret optimal policy (πlin: mean 16,201.57, sd 16,763.98), the nonlinear-optimal policy reduces standard deviation by 0.88% with a very modest 0.42% reduction in mean, directly trading off equity for (average) efficiency as designed.
Practical and Theoretical Implications
The development of policy learning algorithms capable of targeting nonlinear, distribution-sensitive welfare criteria enables a rich set of new applications:
Precision public policy can target reductions in inequality, risk, or tail risk rather than naïve average improvement.
The methods are modular, leveraging advances in flexible function estimation (deep learning) for nuisance components.
Practitioners gain a rigorous cross-validation framework for automating complexity control even in high-dimensional, nonconvex policy spaces.
This work highlights that, in applied optimization of treatment assignment, the welfare objective function's structure (linear vs nonlinear, distributional vs average) fundamentally shapes statistical learning theory: both rates and debiasing strategies must adapt.
Limitations and Directions for Future Research
A key limitation is computational: even with sieve approximation, finding globally optimal policies in high-dimensional, nonconvex settings remains hard, raising the need for more advanced optimization algorithms or convex relaxations. Additionally, while the methods permit policy classes of arbitrary complexity, practical interpretability could become challenging for deep or highly composite rules. Exploring richer forms of constraints (such as fairness or monotonicity) within the nonlinear policy learning framework presents another avenue for research.
Furthermore, generalization to multi-valued or continuous treatments and relaxation of "no unmeasured confounding" may require further methodological innovation.
Conclusion
This paper provides an integrated statistical learning framework for optimal policy learning under general nonlinear welfare objectives, extending prior policy learning theory to accommodate rich, distributional, and equity-driven welfare targets, infinite-dimensional policy spaces, and practical bias correction via reweighting. Theoretical oracle inequalities and empirical illustration underscore the robustness and flexibility of the approach for modern, data-driven policy design that aligns with sophisticated welfare objectives.