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

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

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Placed in Safe offline RL, one of 13 directions in Offline RL and imitation learning from data.
- Builds on 21 of its 47 references that are in the map.
- First posted on arXiv 15 Nov 2025.

## Placed in

- Offline RL and imitation learning from data / [Safe offline RL](https://mlsanity.com/offlinerl/safe-offline-rl/)

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

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Offline Reinforcement Learning: Tutorial, Review, and Perspectives on Open Problems](https://mlsanity.com/paper/2005.01643/) (2020)
- [Conservative Q-Learning for Offline Reinforcement Learning](https://mlsanity.com/paper/2006.04779/) (2020)
- [Offline Reinforcement Learning with Implicit Q-Learning](https://mlsanity.com/paper/2110.06169/) (2021)
- [Prioritized Experience Replay](https://mlsanity.com/paper/1511.05952/) (2015)
- [A Minimalist Approach to Offline Reinforcement Learning](https://mlsanity.com/paper/2106.06860/) (2021)
- [AWAC: Accelerating Online Reinforcement Learning with Offline Datasets](https://mlsanity.com/paper/2006.09359/) (2020)
- [Diffusion Schr"odinger Bridge with Applications to Score-Based Generative Modeling](https://mlsanity.com/paper/2106.01357/) (2021)
- [Reinforcement Learning in Healthcare: A Survey](https://mlsanity.com/paper/1908.08796/) (2019)
- [Uncertainty-Based Offline Reinforcement Learning with Diversified Q-Ensemble](https://mlsanity.com/paper/2110.01548/) (2021)
- [A survey of the Schr"odinger problem and some of its connections with optimal transport](https://mlsanity.com/paper/1308.0215/) (2013)
- [Likelihood Training of Schr"odinger Bridge using Forward-Backward SDEs Theory](https://mlsanity.com/paper/2110.11291/) (2021)
- [A Note on Shumailov et al. (2024): `AI Models Collapse When Trained on Recursively Generated Data'](https://mlsanity.com/paper/2410.12954/) (2024)
- [Deep Generative Learning via Schr"odinger Bridge](https://mlsanity.com/paper/2106.10410/) (2021)
- [Supervised Reinforcement Learning with Recurrent Neural Network for Dynamic Treatment Recommendation](https://mlsanity.com/paper/1807.01473/) (2018)
- [Comprehensive Exploration of Synthetic Data Generation: A Survey](https://mlsanity.com/paper/2401.02524/) (2024)
- [Wasserstein Proximal Algorithms for the Schr"odinger Bridge Problem: Density Control with Nonlinear Drift](https://mlsanity.com/paper/1912.01244/) (2019)
- [DiffStitch: Boosting Offline Reinforcement Learning with Diffusion-based Trajectory Stitching](https://mlsanity.com/paper/2402.02439/) (2024)
- [Semi-Markov Offline Reinforcement Learning for Healthcare](https://mlsanity.com/paper/2203.09365/) (2022)
- [Towards Real-World Applications of Personalized Anesthesia Using Policy Constraint Q Learning for Propofol Infusion Control](https://mlsanity.com/paper/2303.10180/) (2023)

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

- [Offline Materials Optimization with CliqueFlowmer](https://mlsanity.com/paper/2603.06082/) (2026)

## Abstract (arXiv)

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.