Cost-optimal Sequential Testing via Doubly Robust Q-learning

Doudou Zhou; Yiran Zhang; Dian Jin; Yingye Zheng; Lu Tian; Tianxi Cai

stat.MLApr 13, 2026

Cost-optimal Sequential Testing via Doubly Robust Q-learning

Authors: Doudou Zhou, Yiran Zhang, Dian Jin, Yingye Zheng, Lu Tian, Tianxi Cai

Organizations: Department of Statistics and Data Science, National University of Singapore · Public Health Sciences Division, Fred Hutch · Department of Biomedical Data Science, Stanford University School of Medicine · Department of Biostatistics, Harvard T.H. Chan School of Public Health

Abstract

Clinical decision-making often involves selecting tests that are costly, invasive, or time-consuming, motivating individualized, sequential strategies for what to measure and when to stop ascertaining. We study the problem of learning cost-optimal sequential decision policies from retrospective data, where test availability depends on prior results, inducing informative missingness. Under a sequential missing-at-random mechanism, we develop a doubly robust Q-learning framework for estimating optimal policies. The method introduces path-specific inverse probability weights that account for heterogeneous test trajectories and satisfy a normalization property conditional on the observed history. By combining these weights with auxiliary contrast models, we construct orthogonal pseudo-outcomes that enable unbiased policy learning when either the acquisition model or the contrast model is correctly specified. We establish oracle inequalities for the stage-wise contrast estimators, along with convergence rates, regret bounds, and misclassification rates for the learned policy. Simulations demonstrate improved cost-adjusted performance over weighted and complete-case baselines, and an application to a prostate cancer cohort study illustrates how the method reduces testing cost without compromising predictive accuracy.

Figures & tables

State $s$	Path $(S_{1},S_{2})$	Observed Feature ( $\mathbf{X}_{s}$ )	Cumulative Cost ( $C_{s}$ )
0	$(0,0)$	$\mathbf{X}_{0}=X_{0}$	$0$
1	$(1,0)$	$\mathbf{X}_{1}=(X_{0}^{\top},X_{1}^{\top})^{\top}$	$c_{1}$
2	$(2,0)$	$\mathbf{X}_{2}=(X_{0}^{\top},X_{2}^{\top})^{\top}$	$c_{2}$
12	$(1,2)$ or $(2,1)$	$\mathbf{X}_{12}=(X_{0}^{\top},X_{1}^{\top},X_{2}^{\top})^{\top}$	$c_{1}+c_{2}$

Table 1: Information State (

s

), Acquisition Path (

(S_{1},S_{2})

), Observed Feature Vector (

\mathbf{X}_{s}

), and Cumulative Costs (

C_{s}

Figure 1: Scenario 1. Violin plots of average total loss across 50 repetitions at representative sample sizes

n=400,3200,12800

. Rows correspond to settings A/B/C.

Figure 2: Scenario 2. Violin plots of average total loss across 50 repetitions at representative sample sizes

n=400,3200,12800

. Rows correspond to settings A/B/C.

Method	Total Loss	Prediction Loss	Average Cost
COST-Q	0.4212	0.3925	0.0287
One-Time	0.4637	0.4340	0.0297
Only-Complete	0.4593	0.4323	0.0270
BOWL	0.4489	0.4259	0.0230
Always-stop	0.4704	0.4704	0.0000
Always-test-all	0.4284	0.3924	0.0360

Table 2: Loss decomposition on real data (mean over 10 seeds). Bold and underlined indicate the best and second-best values, respectively.

Method	Total Loss	Prediction Loss	Average Cost	AUC
COST-Q	0.4041	0.3771	0.0269	0.8181
One-Time	0.4224	0.3951	0.0273	0.7936
Only-Complete	0.4200	0.3929	0.0271	0.8087
BOWL	0.4286	0.4037	0.0249	0.7895

Table 3: Comparison under an approximately matched testing budget of

0.027

. Bold and underlined indicate the best and second-best values, respectively.

Method	AUC	Spec@90%	G-Mean@90%	Spec@95%	G-Mean@95%
COST-Q	0.8506	0.5949	0.7170	0.5186	0.6937
One-Time	0.8328	0.5650	0.6997	0.4915	0.6773
Only-Complete	0.8074	0.4972	0.6564	0.4011	0.6119
BOWL	0.7933	0.5198	0.6537	0.4181	0.6096
Always-stop	0.6330	0.2429	0.4409	0.2090	0.4311
Always-test-all	0.8075	0.6158	0.6921	0.4746	0.6413

Table 4: AUC and operating-point metrics on real data (mean over 10 seeds). Thresholds for recall targets are selected on training data and evaluated on test data. Bold and underlined indicate the best and second-best values, respectively.

Method	Stop at $X_{0}$	Take $X_{1}$ only	Take $X_{2}$ only	Take both	Avg.# tests
COST-Q	0.000	0.329	0.257	0.414	1.414
One-Time	0.167	0.153	0.518	0.162	0.995
Only-Complete	0.126	0.090	0.532	0.252	1.126
BOWL	0.009	0.000	0.928	0.063	1.054
Always-stop	1.000	0.000	0.000	0.000	0.000
Always-test-all	0.000	0.000	0.000	1.000	2.000

Table 5: Empirical policy behavior on the real-data test set. Each entry is the proportion of subjects assigned to the corresponding terminal testing path; Avg.# tests is the mean number of acquired tests per subject.

Figure 3: Distribution of COST-Q terminal testing paths across quintiles of the baseline PCPT risk score in the test set. Bars show the proportions of subjects who stop at baseline, receive the blood test only, receive the urine test only, or receive both tests.

Q-Learning

Q-learning estimates long-term action values from experience so agents can improve decisions without a learned transition model.

See all 98 papers

Near-Equivalent Q-learning Policies for Dynamic Treatment Regimes
Mar 19, 2026Sophia Yazzourh, Erica E. M. MoodieQ-Learning Policy Learning
Department of Epidemiology and Biostatistics McGill University Montréal, Québec, Canada
Beyond Differences: Doubly Robust Meta-Learners for Ratio-Based Treatment Effects
May 25, 2026Michael Fuchs, Dominik KreissCausal Effect Estimation Conditional Average Treatment Effect Estimation
Actuarial Department Allianz Versicherungs-AG Munich, Germany
Learning Who to Treat When Treatment is Missing
Jul 15, 2026Johnna Sundberg, Rayid Ghani, Eli Ben-Michael +1Causal Effect Estimation Conditional Average Treatment Effect Estimation
Carnegie Mellon University, Pittsburgh, Pennsylvania

Mar 19, 2026J/K move · Enter open · S save

Near-Equivalent Q-learning Policies for Dynamic Treatment Regimes

Sophia Yazzourh, Erica E. M. Moodie

Department of Epidemiology and Biostatistics McGill University Montréal, Québec, Canada

Q-Learning Policy Learning Sequential Decision Making

Open paper AskPDF arXiv

Precision medicine aims to tailor therapeutic decisions to individual patient characteristics. This objective is commonly formalized through dynamic treatment regimes, which use statistical and machine learning methods to derive sequential decision rules adapted to evolving clinical information. In most existing formulations, these approaches produce a single optimal treatment at each stage, leading to a unique decision sequence. However, in many clinical settings, several treatment options may yield similar expected outcomes, and focusing on a single optimal policy may conceal meaningful alternatives. We extend the Q-learning framework for retrospective data by introducing a worst-value tolerance criterion controlled by a hyperparameter

\varepsilon

, which specifies the maximum acceptable deviation from the optimal expected value. Rather than identifying a single optimal policy, the proposed approach constructs sets of

\varepsilon

-optimal policies whose performance remains within a controlled neighborhood of the optimum. This formulation shifts Q-learning from a vector-valued representation to a matrix-valued one, allowing multiple admissible value functions to coexist during backward recursion. The approach yields families of near-equivalent treatment strategies and explicitly identifies regions of treatment indifference where several decisions achieve comparable outcomes. We illustrate the framework in two settings: a single-stage problem highlighting indifference regions around the decision boundary, and a multi-stage decision process based on a simulated oncology model describing tumor size and treatment toxicity dynamics.

Abstract

Q-Learning

Explore similar work

Near-Equivalent Q-learning Policies for Dynamic Treatment Regimes

Beyond Differences: Doubly Robust Meta-Learners for Ratio-Based Treatment Effects

Learning Who to Treat When Treatment is Missing