Censoring-Aware Tree-Based Reinforcement Learning for Estimating Dynamic Treatment Regimes with Censored Outcomes

2025/03/09 by Paul, Animesh Kumar, Greiner, Russell
#Artificial Intelligence (cs.AI) #FOS: Computer and information sciences #Machine Learning (cs.LG) #Methodology (stat.ME)

paper · doi:10.48550/arxiv.2503.06690

Abstract

Dynamic Treatment Regimes (DTRs) provide a systematic approach for making sequential treatment decisions that adapt to individual patient characteristics, particularly in clinical contexts where survival outcomes are of interest. Censoring-Aware Tree-Based Reinforcement Learning (CA-TRL) is a novel framework to address the complexities associated with censored data when estimating optimal DTRs. We explore ways to learn effective DTRs, from observational data. By enhancing traditional tree-based reinforcement learning methods with augmented inverse probability weighting (AIPW) and censoring-aware modifications, CA-TRL delivers robust and interpretable treatment strategies. We demonstrate its effectiveness through extensive simulations and real-world applications using the SANAD epilepsy dataset, where it outperformed the recently proposed ASCL method in key metrics such as restricted mean survival time (RMST) and decision-making accuracy. This work represents a step forward in advancing personalized and data-driven treatment strategies across diverse healthcare settings.

Related

• Dynamic Treatment Regimes using Bayesian Additive Regression Trees for Censored Outcomes
• Multi-stage optimal dynamic treatment regimes for survival outcomes with dependent censoring
• Tree based weighted learning for estimating individualized treatment rules with censored data
• Evaluating Reinforcement Learning Algorithms in Observational Health Settings
• Deep Reinforcement Learning for Dynamic Treatment Regimes on Medical Registry Data