# Stage-Aware Learning for Dynamic Treatments

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

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 14 of its 66 references that are in the map.
- First posted on arXiv 30 Oct 2023.

## Builds on (most cited first, 14 of 14)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [SGDR: Stochastic Gradient Descent with Warm Restarts](https://mlsanity.com/paper/1608.03983/) (2016)
- [Neural Machine Translation by Jointly Learning to Align and Translate](https://mlsanity.com/paper/1409.0473/) (2014)
- [Playing Atari with Deep Reinforcement Learning](https://mlsanity.com/paper/1312.5602/) (2013)
- [Residual Weighted Learning for Estimating Individualized Treatment Rules](https://mlsanity.com/paper/1508.03179/) (2015)
- [Learning Optimal Distributionally Robust Individualized Treatment Rules](https://mlsanity.com/paper/2006.15121/) (2020)
- [Estimating Optimal Infinite Horizon Dynamic Treatment Regimes via pT-Learning](https://mlsanity.com/paper/2110.10719/) (2021)
- [Random Forest Missing Data Algorithms](https://mlsanity.com/paper/1701.05305/) (2017)
- [Using Decision Lists to Construct Interpretable and Parsimonious Treatment Regimes](https://mlsanity.com/paper/1504.07715/) (2015)
- [Efficient augmentation and relaxation learning for individualized treatment rules using observational data](https://mlsanity.com/paper/1901.00663/) (2019)
- [A General System for Heuristic Solution of Convex Problems over Nonconvex Sets](https://mlsanity.com/paper/1601.07277/) (2016)
- [Multicategory Angle-based Learning for Estimating Optimal Dynamic Treatment Regimes with Censored Data](https://mlsanity.com/paper/2001.04629/) (2020)
- [Robust Learning for Optimal Treatment Decision with NP-Dimensionality](https://mlsanity.com/paper/1510.04378/) (2015)
- [Finding the Optimal Dynamic Treatment Regime Using Smooth Fisher Consistent Surrogate Loss](https://mlsanity.com/paper/2111.02826/) (2021)

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

- [Learning Robust Treatment Rules for Censored Data](https://mlsanity.com/paper/2408.09155/) (2024)
- [Cost-optimal Sequential Testing via Doubly Robust Q-learning](https://mlsanity.com/paper/2604.11165/) (2026)

## Abstract (arXiv)

Recent advances in dynamic treatment regimes (DTRs) facilitate the search for optimal treatments, which are tailored to individuals' specific needs and able to maximize their expected clinical benefits. However, existing algorithms relying on consistent trajectories, such as inverse probability weighting estimators (IPWEs), could suffer from insufficient sample size under optimal treatments and a growing number of decision-making stages, particularly in the context of chronic diseases. To address these challenges, we propose a novel individualized learning method which estimates the DTR with a focus on prioritizing alignment between the observed treatment trajectory and the one obtained by the optimal regime across decision stages. By relaxing the restriction that the observed trajectory must be fully aligned with the optimal treatments, our approach substantially improves the sample efficiency and stability of IPWE-based methods. In particular, the proposed learning scheme builds a more general framework which includes the popular outcome weighted learning framework as a special case of ours. Moreover, we introduce the notion of stage importance scores along with an attention mechanism to explicitly account for heterogeneity among decision stages. We establish the theoretical properties of the proposed approach, including the Fisher consistency and finite-sample performance bound. Empirically, we evaluate the proposed method in extensive simulated environments and a real case study for the COVID-19 pandemic.