Medical Decision Transformer (MeDT)

• Medical Decision Transformer (MeDT) is an offline reinforcement learning framework that integrates long-horizon return-to-go with per-step acuity targets for tailored sepsis treatment.
• It utilizes a causal decoder-only transformer to encode clinical states, actions, and objectives, enabling counterfactual ‘what-if’ analysis and interactive dosing recommendations.
• MeDT demonstrates competitive performance with lower terminal SAPS-II scores on MIMIC-III data, offering interpretable, clinician-guided decision support in complex critical care scenarios.

The Medical Decision Transformer (MeDT) is a specialized, goal-conditioned offline reinforcement learning (RL) framework that leverages a Decision Transformer backbone to address the challenges of clinical decision support in sepsis treatment. MeDT incorporates both long-horizon and clinician-specified short-term targets, facilitating interpretable, customizable, and interactive policy recommendations for intravenous (IV) fluid and vasopressor dosing. It is designed to handle sparse terminal rewards characteristic of survival in sepsis while enabling “what-if” counterfactual exploration for clinical stakeholders (Rahman et al., 2024).

1. Problem Formulation and Markov Decision Process Structure

MeDT models sepsis management as a discrete-time Markov Decision Process (MDP), where the environment is defined as follows:

• State Space $\mathcal{S}$: At each 4-hour interval $t$, the state vector $s_t \in \mathbb{R}^{43}$ comprises 5 static demographic features and 38 time-varying clinical measurements (vital signs, laboratory results, ventilation status, fluid balance).
• Action Space $\mathcal{A}$: Actions are discretized as 25 joint bins derived from 5 IV fluid dosages and 5 vasopressor dosages, encoding $a_t \in \{1, ..., 25\}$ as the policy's recommendation per step.
• Reward Function $\mathcal{R}$: Terminal reward $R_T\in\{+1,-1\}$ encodes patient survival ($+1$) or mortality ($-1$) at the end of the trajectory, with $r_t=0$ for all $t<T$.

The return-to-go at time $t$ is $R_t = \sum_{t' = t}^T r_{t'}$, with return conditioning motivated by sparse outcomes in critical-care settings. To address reward sparsity and enable nuanced, short-term control, MeDT introduces per-step “Acuity-to-Go” (ATG) targets, $k_t \in \mathbb{R}^7$, comprising the seven SAPS-II organ-system subscores (cardiovascular, respiratory, neurological, renal, hepatic, haematologic, and other).

2. Model Architecture and Input Encoding

MeDT adopts a causal decoder-only transformer architecture with six layers. The key modeling components are:

• Input Tokenization: For each time step $t$, the model processes the following input tokens:
  1. Return-to-go $R_t$
  2. Acuity-to-go $k_t \in \mathbb{R}^7$
  3. State $s_t \in \mathbb{R}^{43}$
  4. Action $a_t$ (as a one-hot vector during training)

• Token Embedding: Each token is projected via a linear layer and layer normalization to a $d$-dimensional embedding. Learned positional encodings $p_t \in \mathbb{R}^d$ are added to capture temporal information.

The input sequence at each step is ordered as $$(R_t, k_t, s_t, a_{t-1}, \ldots, R_1, k_1, s_1, a_0)$$, and is supplied to the transformer without specialized cross-attention or gating mechanisms. Both long-term ($R_T$) and granular ($k_t$) goals are interleaved as tokens for direct conditioning.

3. Training Objectives and Optimization

MeDT employs a behavior cloning objective to predict the clinician’s action $a_t$ at each step given historical context. The primary loss is cross entropy over the action prediction:

$$L_\text{action} = -\frac{1}{T} \sum_{t=1}^T \log P_\theta(a_t^{\text{GT}} \mid R_{1:t}, k_{1:t}, s_{1:t}, a_{0:t-1})$$

Additionally, a parallel state estimator transformer models patient transitions by predicting $s_{t+1}$ from previous states and actions, optimized via mean squared error:

$$L_\text{state} = \frac{1}{T} \sum_{t=1}^T \| s_{t+1} - \hat{s}_{t+1} \|^2$$

Incorporation of ATG tokens ($k_t$) at each step addresses the reward sparsity problem, effectively functioning as reward shaping and providing denser supervision than the purely terminal $R_T$.

4. Interactive Clinical Conditioning and Counterfactual Roll-Outs

MeDT uniquely allows clinician or heuristic input of desired short-term acuity goals $k_{1:T}$ at inference. Users can specify targets per organ system (e.g., “reduce renal score by 1,” “maintain hepatic bilirubin below 2”), alongside a terminal survival goal ($R_T=+1$). This conditioning mechanism enables interactive, patient-specific scenario planning.

A counterfactual roll-out mechanism operates as follows:

1. For each step $t$, MeDT ingests the current context $(R_{1:t}, k_{1:t}, s_{1:t}, a_{0:t-1})$ and outputs a recommended action $\hat{a}_t$.
2. The auxiliary state estimator predicts the next state $\hat{s}_{t+1}$.
3. SAPS-II formulas are applied to $\hat{s}_{t+1}$ to update predicted acuity $\hat{k}_{t+1}$.
4. This new information is appended, and the sequence is repeated.

This process generates hypothetical dosing trajectories (“what-if” scenarios) and supports modification of future acuity goals in real time.

5. Data Processing and Evaluation Protocols

Training and validation are conducted on the MIMIC-III database, focusing on ~$19,633$ septic patients (9% ICU mortality) identified according to Sepsis-3 criteria. Data are windowed into non-overlapping 4-hour intervals; within-windows, measurements are averaged, and missing values imputed via $k$-nearest neighbors. The resulting state representation is $5$ static and $38$ dynamic variables; actions are the 25-binned IV fluid/vasopressor joint options.

Offline policy evaluation (OPE) uses three primary estimators:

Metric	Definition	Purpose
Weighted Importance Sampling (WIS)	$$WIS=\frac{1}{N} \sum_{n=1}^N \frac{\rho_{1:T}^{(n)}}{w_T} \sum_{t=1}^{T^{(n)}} \gamma^{t-1} r_t^{(n)}$$	Re-weights observed returns by importance ratios to estimate counterfactual policy value
Fitted Q-Evaluation (FQE)	Bellman backups estimate $Q^\pi$, then $FQE$-value averages initial-state action-values	Batch policy value estimation leveraging approximate Q-functions
Weighted Doubly Robust (WDR)	Combines IS weights and fitted Q/V values to lower variance	Reduces estimator bias and variance relative to IS and FQE alone

Comparative performance is summarized below (terminal SAPS-II, lower is better):

Method	Predicted SAPS-II $\pm$ SE
BCQ	42.10 ± 0.03
NFQI	44.21 ± 0.05
DDQN	43.47 ± 0.04
CQL	40.42 ± 0.03
BC (transformer)	40.50 ± 0.03
Decision Transformer	40.38 ± 0.03
MeDT	40.31 ± 0.03

MeDT outperforms or is competitive with all baselines under FQE and WDR, with lower variance observed compared to pure RL methods (Rahman et al., 2024).

6. Interpretability, Personalization, and Ablation Insights

Transformer relevance visualization, following Chefer et al. (2021), is used to compute gradient-weighted attention over attention maps:

$$\overline{A} = \mathbb{E}_{\text{heads}}[(\partial y/\partial A \odot A)^+], \quad R \leftarrow R + \overline{A} R$$

This enables heatmap generation on input token importance for individual dosing decisions. Analysts can confirm, for instance, that a dominant role was played by a particular organ-system subscore (e.g., hepatic ‘$k_h$’) in a dosing recommendation at a given time.

Ablation experiments indicate that MeDT gains most on positive-return (survival) trajectories, demonstrating ATG conditioning’s ability to leverage rich positive samples. In contrast, standard Decision Transformers see more improvement on negative trajectories, implying that pure return-to-go conditioning is advantageous in “rescue” scenarios.

7. Limitations and Implications

Operational limitations of MeDT include lower sample efficiency relative to conventional Decision Transformers—more data are needed to support ATG conditioning. Biases present in MIMIC-III and the approximate state estimator may propagate through hypothetical rollouts, particularly in long-term counterfactuals where model error can accumulate. Nonetheless, MeDT advances offline RL in medicine by providing direct “what-if” interactivity via per-step acuity conditioning and exposing token-level explanation heatmaps for key recommendations.

A plausible implication is that such frameworks may accelerate the integration of interactive, clinician-in-the-loop RL systems for other safety-critical applications, provided robust external validation and bias assessment strategies are pursued (Rahman et al., 2024).