# Treatment Effect Estimation for Optimal Decision-Making

arXiv: https://arxiv.org/abs/2505.13092 · HTML: https://mlsanity.com/paper/2505.13092/

## Key facts

- Cited by 8 papers in the map (data through 14 Oct 2026).
- Builds on 32 of its 65 references that are in the map.
- First posted on arXiv 19 May 2025.

## Builds on (most cited first, 30 of 32)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Estimation and Inference of Heterogeneous Treatment Effects using Random Forests](https://mlsanity.com/paper/1510.04342/) (2015)
- [Meta-learners for Estimating Heterogeneous Treatment Effects using Machine Learning](https://mlsanity.com/paper/1706.03461/) (2017)
- [Doubly Robust Off-policy Value Evaluation for Reinforcement Learning](https://mlsanity.com/paper/1511.03722/) (2015)
- [Quasi-Oracle Estimation of Heterogeneous Treatment Effects](https://mlsanity.com/paper/1712.04912/) (2017)
- [Data-Efficient Off-Policy Policy Evaluation for Reinforcement Learning](https://mlsanity.com/paper/1604.00923/) (2016)
- [Estimating Counterfactual Treatment Outcomes over Time Through Adversarially Balanced Representations](https://mlsanity.com/paper/2002.04083/) (2020)
- [Orthogonal Statistical Learning](https://mlsanity.com/paper/1901.09036/) (2019)
- [Fair Inference On Outcomes](https://mlsanity.com/paper/1705.10378/) (2017)
- [Nonparametric Estimation of Heterogeneous Treatment Effects: From Theory to Learning Algorithms](https://mlsanity.com/paper/2101.10943/) (2021)
- [Causal Transformer for Estimating Counterfactual Outcomes](https://mlsanity.com/paper/2204.07258/) (2022)
- [Double Reinforcement Learning for Efficient Off-Policy Evaluation in Markov Decision Processes](https://mlsanity.com/paper/1908.08526/) (2019)
- [Efficiently Breaking the Curse of Horizon in Off-Policy Evaluation with Double Reinforcement Learning](https://mlsanity.com/paper/1909.05850/) (2019)
- [On Inductive Biases for Heterogeneous Treatment Effect Estimation](https://mlsanity.com/paper/2106.03765/) (2021)
- [Causal machine learning for predicting treatment outcomes](https://mlsanity.com/paper/2410.08770/) (2024)
- [Balanced Policy Evaluation and Learning](https://mlsanity.com/paper/1705.07384/) (2017)
- [Doubly Robust Distributionally Robust Off-Policy Evaluation and Learning](https://mlsanity.com/paper/2202.09667/) (2022)
- [Intrinsically Efficient, Stable, and Bounded Off-Policy Evaluation for Reinforcement Learning](https://mlsanity.com/paper/1906.03735/) (2019)
- [Causal Decision Making and Causal Effect Estimation Are Not the Same... and Why It Matters](https://mlsanity.com/paper/2104.04103/) (2021)
- [More Efficient Policy Learning via Optimal Retargeting](https://mlsanity.com/paper/1906.08611/) (2019)
- [Model-agnostic meta-learners for estimating heterogeneous treatment effects over time](https://mlsanity.com/paper/2407.05287/) (2024)
- [More Efficient Off-Policy Evaluation through Regularized Targeted Learning](https://mlsanity.com/paper/1912.06292/) (2019)
- [Reliable Off-Policy Learning for Dosage Combinations](https://mlsanity.com/paper/2305.19742/) (2023)
- [Fair Off-Policy Learning from Observational Data](https://mlsanity.com/paper/2303.08516/) (2023)
- [Efficient Policy Learning from Surrogate-Loss Classification Reductions](https://mlsanity.com/paper/2002.05153/) (2020)
- [Doubly Robust Off-Policy Value and Gradient Estimation for Deterministic Policies](https://mlsanity.com/paper/2006.03900/) (2020)
- [Regret Minimization for Causal Inference on Large Treatment Space](https://mlsanity.com/paper/2006.05616/) (2020)
- [Fair and Robust Estimation of Heterogeneous Treatment Effects for Policy Learning](https://mlsanity.com/paper/2306.03625/) (2023)
- [Causal Machine Learning for Cost-Effective Allocation of Development Aid](https://mlsanity.com/paper/2401.16986/) (2024)
- [Data-Driven Allocation of Preventive Care With Application to Diabetes Mellitus Type II](https://mlsanity.com/paper/2308.06959/) (2023)

## Built on by (most cited first, 8 of 8)

- [Rank-Learner: Orthogonal Ranking of Treatment Effects](https://mlsanity.com/paper/2602.03517/) (2026)
- [Aligning Data-Driven Predictors with Allocation: A Decision-Focused Approach to Survival Analysis](https://mlsanity.com/paper/2606.02671/) (2026)
- [GDR-learners: Orthogonal Learning of Generative Models for Potential Outcomes](https://mlsanity.com/paper/2509.22953/) (2025)
- [Annotation-Assisted Learning of Treatment Policies From Multimodal Electronic Health Records](https://mlsanity.com/paper/2507.20993/) (2025)
- [Reliable Selection of Heterogeneous Treatment Effect Estimators](https://mlsanity.com/paper/2511.18464/) (2025)
- [Uncertainty-aware Causal Decision Making via Effect Bound Decomposition](https://mlsanity.com/paper/2601.22736/) (2026)
- [Should I stay or should I show? Learning to selectively disclose information](https://mlsanity.com/paper/2609.39818/) (2026)
- [Learning Deployable Causal Action Geometry under Temporal Non-Stationarity](https://mlsanity.com/paper/2610.16644/) (2026)

## Abstract (arXiv)

Decision-making across various fields, such as medicine, heavily relies on conditional average treatment effects (CATEs). Practitioners commonly make decisions by checking whether the estimated CATE is positive, even though the decision-making performance of modern CATE estimators is poorly understood from a theoretical perspective. In this paper, we study optimal decision-making based on two-stage CATE estimators (e.g., DR-learner), which are considered state-of-the-art and widely used in practice. We prove that, while such estimators may be optimal for estimating CATE, they can be suboptimal when used for decision-making. Intuitively, this occurs because such estimators prioritize CATE accuracy in regions far away from the decision boundary, which is ultimately irrelevant to decision-making. As a remedy, we propose a novel two-stage learning objective that retargets the CATE to balance CATE estimation error and decision performance. We then propose a neural method that optimizes an adaptively-smoothed approximation of our learning objective. Finally, we confirm the effectiveness of our method both empirically and theoretically. In sum, our work is the first to show how two-stage CATE estimators can be adapted for optimal decision-making.