Treatment Stitching with Schr\"odinger Bridge for Enhancing Offline Reinforcement Learning in Adaptive Treatment Strategies

TL;DR

This paper introduces TreatStitch, a data augmentation framework using Schrödinger bridges to improve offline reinforcement learning for adaptive treatment strategies by generating clinically valid synthetic treatment trajectories.

Contribution

TreatStitch is a novel method that stitches treatment data segments and uses Schrödinger bridges to create realistic, diverse trajectories for better offline RL performance in clinical settings.

Findings

TreatStitch significantly improves offline RL policy quality.

Synthetic trajectories enhance data diversity and model robustness.

Method maintains clinical validity by avoiding out-of-distribution states.

Abstract

Adaptive treatment strategies (ATS) are sequential decision-making processes that enable personalized care by dynamically adjusting treatment decisions in response to evolving patient symptoms. While reinforcement learning (RL) offers a promising approach for optimizing ATS, its conventional online trial-and-error learning mechanism is not permissible in clinical settings due to risks of harm to patients. Offline RL tackles this limitation by learning policies exclusively from historical treatment data, but its performance is often constrained by data scarcity-a pervasive challenge in clinical domains. To overcome this, we propose Treatment Stitching (TreatStitch), a novel data augmentation framework that generates clinically valid treatment trajectories by intelligently stitching segments from existing treatment data. Specifically, TreatStitch identifies similar intermediate patient states across different trajectories and stitches their respective segments. Even when intermediate states are too dissimilar to stitch directly, TreatStitch leverages the Schr\"odinger bridge method to generate smooth and energy-efficient bridging trajectories that connect dissimilar states. By augmenting these synthetic trajectories into the original dataset, offline RL can learn from a more diverse dataset, thereby improving its ability to optimize ATS. Extensive experiments across multiple treatment datasets demonstrate the effectiveness of TreatStitch in enhancing offline RL performance. Furthermore, we provide a theoretical justification showing that TreatStitch maintains clinical validity by avoiding out-of-distribution transitions.