Treatment Stitching with Schrödinger Bridge for Enhancing Offline Reinforcement Learning in Adaptive Treatment Strategies

AI summary

Overview

• Research area: Offline reinforcement learning for adaptive treatment strategies (ATS) in clinical decision-making, combined with data augmentation via trajectory stitching and the Schrödinger bridge (optimal transport).
• Technical level: Intermediate — the paper assumes familiarity with Markov decision processes, offline RL (specifically Conservative Q-Learning), generative modeling, and optimal transport, but its central idea (stitching real patient trajectories) is intuitive.
• Scope: The paper proposes Treatment Stitching (TreatStitch), a data augmentation framework that builds new, clinically valid treatment trajectories by cutting and joining segments of real trajectories, using a Schrödinger bridge to connect states that are too dissimilar to stitch directly.

What This Paper Is About

Offline RL can learn treatment policies from historical patient data without risky online trial-and-error, but clinical datasets are typically too small and too homogeneous for RL to train well, and existing augmentation methods that synthesize trajectories from scratch tend to produce low-quality data. TreatStitch instead recombines existing trajectories: it finds pairs of intermediate patient states that are similar across different trajectories and swaps the segments that follow them, producing new trajectories that never existed in the original data. When no sufficiently similar state pairs exist (as in sparse or heterogeneous clinical data), it inserts a short, smooth bridging trajectory generated by a Schrödinger bridge so that dissimilar states can still be connected.

Key Contributions

1. A new data augmentation framework for ATS. The authors state that TreatStitch is the first data augmentation framework for adaptive treatment strategies to use trajectory stitching — generating clinically valid synthetic trajectories out of existing offline treatment data rather than synthesizing data from scratch.
2. A Schrödinger bridge mechanism for dissimilar states. When direct stitching is impossible because no intermediate states are similar enough, the framework uses the Schrödinger bridge method to construct smooth, energy-efficient bridging trajectories that enable stitching between dissimilar states, extending the method to sparse and heterogeneous datasets.
3. Empirical validation across multiple treatment datasets. TreatStitch is evaluated on the EpiCare benchmark (8 environment settings) and on the real-world MIMIC-III sepsis dataset, using CQL as the backbone offline RL algorithm and comparing against GAN, DDPM, SynthER, GTA, ATraDiff, and RTDiff.
4. Theoretical justification of clinical validity. A theorem shows that stitched trajectories remain within an O(L√(2(1−δ)))-neighborhood of transitions already present in the offline dataset, i.e., they stay close to the original data support and avoid out-of-distribution (OOD) transitions.

Main Findings

• Stitching beats from-scratch synthesis in the full-data setting. On EpiCare with 131,072 episodes, TreatStitch reached a mean average cumulative reward of 57.93 versus 54.37 for the CQL backbone, 27.56 for GAN, 38.53 for DDPM, 38.45 for SynthER, 53.08 for GTA, 53.46 for ATraDiff, and 55.03 for RTDiff. TreatStitch with the Schrödinger bridge reached 58.06.
• The Schrödinger bridge helps most when data is scarce. With only 1,024 episodes (128 times fewer than the full setting), TreatStitch reached a mean of 35.63 and TreatStitch w/ SB reached 40.07, compared with 15.79 for the CQL backbone, 9.56 for GAN, 12.11 for DDPM, 12.70 for SynthER, 17.91 for GTA, 17.66 for ATraDiff, and 20.09 for RTDiff. The authors attribute the SB gain to the scarcity of valid stitching points in sparse datasets.
• Generative models trained from scratch can hurt performance. GAN and DDPM fell well below the CQL backbone in both data settings (for example, GAN mean 9.56 in the restricted setting versus 15.79 for the backbone), which the authors connect to reported model collapse from erroneous, low-quality synthetic data.
• Real-world MIMIC-III results improve on clinically meaningful metrics. TreatStitch w/ SB reported 578.90 and 0.28 on the two reported severity/risk metric columns, versus 611.30 and 0.39 for the CQL backbone, and also reported the lowest values among listed methods across the remaining reported columns (10.86, −0.30, 526.71, 0.24, −3.33, −0.56, versus 13.27, −0.37, 566.64, 0.33, −4.70, −0.69 for the backbone). Competing methods GTA (634.81, 0.42), ATraDiff (607.75, 0.39), and RTDiff (601.77, 0.32) did not match it on the leading metric column.
• Schrödinger bridge produces shorter bridges than diffusion. A kernel density estimate of bridging trajectory lengths showed the SB method with a pronounced peak at shorter lengths, which the authors attribute to its optimal-transport basis seeking smooth, energy-efficient paths; a diffusion-based variant produced comparable but weaker performance overall.
• Only one competing method improved in the full-data setting. Among the compared augmentation methods, RTDiff was the only one to show a performance improvement over the backbone in the full data setting, and the authors describe that gain as limited because RTDiff targets general RL tasks rather than clinical validity.

Methodology in Plain English

1. Start from historical trajectories. The offline dataset consists of trajectories of states (patient symptoms), actions (treatments), and rewards (responses), drawn from a clinical database.
2. Split trajectories by outcome. Trajectories whose discounted cumulative reward is at or above the q-th percentile of the dataset go into a high-reward group; the rest go into a low-reward group.
3. Prioritize informative pairings. A Boltzmann distribution samples trajectories with the highest rewards from the high group and the lowest rewards from the low group, on the intuition that pairing adverse-outcome data with beneficial-outcome data yields more informative synthetic trajectories. The temperature parameter is gradually increased during training, shifting from strongly prioritized to more uniform sampling for broader coverage.
4. Test a candidate stitching point. Cosine similarity is computed between every pair of intermediate states across the two sampled trajectories. If the similarity meets the threshold δ (set to 0.95 in the experiments), that pair becomes the stitching point; otherwise the process resamples.
5. Cut and join. The early segment of one trajectory up to the stitching point is concatenated with the later segment of the other, forming a new stitched trajectory that is added to the dataset.
6. Bridge when nothing is similar enough. If no state pair reaches the threshold, the framework picks the most similar candidate pair anyway and trains a neural network to solve the Schrödinger system, generating a short sequence of bridging states that smoothly connects them. The number of bridging states is deliberately minimized to limit error accumulation from synthetic data.
7. Fill in actions and rewards. An inverse dynamics model predicts the action between each pair of successive states, and a reward model predicts the reward for each state-action pair; both are trained jointly with a mean-squared-error objective on real trajectories. Together with the bridge states, they complete a bridging trajectory.
8. Concatenate the full bridging trajectory. The early segment of one real trajectory, the bridging trajectory, and the later segment of the other real trajectory are joined into a single stitched trajectory.
9. Train offline RL on the augmented dataset. The augmented dataset, with its greater diversity and coverage, is used to train the offline RL algorithm — in these experiments, CQL, which penalizes overly optimistic Q-values on out-of-distribution state-action pairs.

Why This Matters

• Research impact: The paper reframes clinical data scarcity as a data-side problem rather than only an algorithmic one. It also supplies a formal bound showing that stitched trajectories stay near the support of the original data, which is a concrete argument for why stitching is safer than synthesizing treatment data from scratch.
• Real-world applications:
  • Longitudinal management of chronic diseases, where patients return repeatedly for assessment and treatment adjustment — the scenario the paper opens with.
  • Sepsis treatment in the intensive care unit, the setting of the MIMIC-III experiments, where the reported metrics include SOFA (organ dysfunction severity) and NEWS2 (mortality risk).
  • Clinical decision support systems that recommend treatment sequences by learning from historical electronic health records.
  • Sparse or heterogeneous hospital datasets, where direct stitching points are rare and the Schrödinger bridge provides the additional stitching opportunities.
• Industry relevance: The work targets healthcare providers, hospital systems, and clinical AI developers who hold large electronic health record archives but cannot run prospective trials for every policy. Because the method reuses existing data and avoids online experimentation on patients, it fits regulatory and ethical constraints that block conventional online RL in clinical settings.

Future Directions

• Scaling the Schrödinger bridge mechanism to other offline RL backbones. The paper uses CQL as the main backbone and reports results with other offline RL models in Appendix C; how broadly the gains transfer is left to that appendix and beyond.
• Reducing reliance on hand-set thresholds. The similarity threshold δ = 0.95 and the temperature schedule for priority sampling are chosen a priori (the paper provides further analysis of priority sampling in Appendix D); automating or adapting these settings is an open question.
• Extending beyond the evaluated datasets. Evaluation covers the EpiCare benchmark's 8 environment settings and MIMIC-III sepsis treatment; the paper does not report results on other disease areas or other real-world record systems.
• Bounding error accumulation more tightly. The framework deliberately minimizes the number of bridging states to limit error accumulation, but the paper does not report a quantitative analysis of cumulative error as synthetic data volume grows. The paper's stated future work is not reported in the available content.

Target Audience

Researchers and practitioners at the intersection of offline reinforcement learning, healthcare AI, and optimal transport — particularly those working on clinical decision support, treatment policy learning from electronic health records, and data augmentation for sequential decision-making. It is also relevant to clinicians and health-system data scientists evaluating whether offline RL is viable under real data scarcity constraints. Readers without background in RL or optimal transport will need to consult the appendices (A through G) referenced in the paper for the extended related work, backbone comparisons, cross-model results, priority sampling analysis, Schrödinger bridge analysis, training details, and dataset details.

Authors’ 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ödinger 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.