Offline Inverse Constrained Reinforcement Learning for Safe-Critical Decision Making in Healthcare

## Abstract

Abstract Reinforcement Learning (RL) applied in healthcare can lead to unsafe medical decisions and treatment, such as excessive dosages or abrupt changes, often due to agents overlooking common-sense constraints. Consequently, Constrained Reinforcement Learning (CRL) is a natural choice for safe decisions. However, specifying the exact cost function is inherently difficult in healthcare. Recent Inverse Constrained Reinforcement Learning (ICRL) is a promising approach that infers constraints from expert demonstrations. ICRL algorithms model Markovian decisions in an interactive environment. These settings do not align with the practical requirement of a decision-making system in healthcare, where decisions rely on historical treatment recorded in an offline dataset. To tackle these issues, we propose the Constraint Transformer (CT). Specifically, 1) we utilize a causal attention mechanism to incorporate historical decisions and observations into the constraint modeling, while employing a Non-Markovian layer for weighted constraints to capture critical states. 2) A generative world model is used to perform exploratory data augmentation, enabling offline RL methods to simulate unsafe decision sequences. In multiple medical scenarios, empirical results demonstrate that CT can capture unsafe states and achieve strategies that approximate lower mortality rates, reducing the occurrence probability of unsafe behaviors.

## 1 Introduction

In recent years, the doctor-to-patient ratio imbalance has drawn attention, with the U.S. having only 223.1 physicians per 100,000 people (Petterson et al., ). AI-assisted therapy emerges as a promising solution, offering timely diagnosis, personalized care, and reducing dependence on experienced physicians. Therefore, the development of an effective AI healthcare assistant is crucial.

**Table 1: Proportion of unsafe vaso doses recommended by physician and DDPG policy. Typical vaso dosages range from $0.1$ to $0.2\mu g/(kg\cdot min)$, with doses above $0.5$ considered high (Bassi et al., ). A critical threshold of $0.75$ is associated with increased mortality (Auchet et al., ).**
| Actions $(\mu g/(kg\cdot min))$ | Physician policy | DDPG policy |
| --- | --- | --- |
| vaso $\textgreater 0.75$ | $2.27\%$ | $7.44\%$ $\uparrow$ |
| vaso $\textgreater 0.9$ | $1.71\%$ | $7.40\%$ $\uparrow$ |
| $\Delta$ vaso $\textgreater 0.75$ | $2.45\%$ | $21.00\%$ $\uparrow$ |
| $\Delta$ vaso $\textgreater 0.9$ | $1.88\%$ | $20.62\%$ $\uparrow$ |

Reinforcement learning (RL) offers a promising approach to develop AI assistants by addressing sequential decision-making tasks. However, this method can still lead to unsafe behaviors, such as administering excessive drug dosages, inappropriate adjustments of medical parameters, or abrupt changes in medication dosages. These actions, including “too high” or “sudden change”, may significantly endanger patients, potentially resulting in acute hypotension, arrhythmias, and organ damage, with fatal consequences (Jia et al., ; Shi et al., ). For example, in sepsis treatment, vasopressor (vaso) doses above $1\mu g/(kg\cdot min)$ are linked to a 90% mortality rate (Martin et al., ), and sudden changes in vaso can cause dangerous blood pressure fluctuations (Fadale et al., ). Our experiments show that Huang et al. () use of the DDPG algorithm in sepsis, which exhibits “too high” (^1^11“too high” refers to a lethal drug dose for a particular patient; however, this is not a single exact value, as it can vary depending on the patient’s individual condition.) and “sudden change” (^2^22“sudden change” indicates that the change in dosage between two-time points exceeds the threshold.) in vaso recommendations, as seen in Table 1. Moreover, if the dosage is clipped using simple thresholding, it will not account for the individualized tolerance of each patient.

This paper aims to achieve safe healthcare policy learning to mitigate unsafe behaviors. The most common method for learning safe policies is Constrained Reinforcement Learning (CRL) (Liu et al., ; ), with the key to its success lying in the constraints representation. However, in healthcare, we can only design the cost function based on prior knowledge, which limits its application due to a lack of personalization, universality, and reliance on prior knowledge. For more details about issues, please refer to Appendix B. Therefore, Inverse Constrained Reinforcement Learning (ICRL) (Malik et al., ) emerges as a promising approach, as it can infer the constraints adhered to by experts from their demonstrations. However, existing ICRL methods face the following challenges:

Figure: Figure 1: The distribution of vaso for patients with the same state. Physicians make different decisions by referencing historical information, whereas agents, based on the Markov assumption, can only make the same decision.
Refer to caption: x1.png

1) The Markov decision is not compatible with medical decisions.
ICRL algorithms model Markov decisions, where the next state depends only on the current state and not on the history (Kijima, ; Zhang et al., ). However, in healthcare, the historical states of patients are crucial for medical decision-making (Plaisant et al., ), as demonstrated in the experiments shown in Figure 1. Therefore, ICRL algorithms based on Markov assumption can not capture patient history, and ignore individual patient differences, thereby limiting effectiveness.
2) Interactive environment is not available for healthcare or medical decisions.
ICRL algorithms (Malik et al., ; Gaurav et al., ) follow an online learning paradigm, allowing agents to explore and learn from interactive environments. However, exploration in healthcare often entails unsafe behaviors that could breach constraints and result in substantial losses. Therefore, it is necessary to infer constraints using only offline datasets.

In this paper, we introduce offline Constraint Transformer (CT), a novel ICRL framework that incorporates patients’ historical information into constraint modeling and learns from offline data to infer constraints in healthcare. Specifically,

1) Inspired by the recent success of sequence modeling (Zheng et al., ; Chen et al., ; Kim et al., ), we incorporate historical decisions and observations into constraint modeling using a causal attention mechanism. To capture key events in trajectories, we introduce a Non-Markovian transformer to generate constraints and cost weights, and then define constraints using weighted sums. CT takes trajectories as input, allowing for the observation of patients’ historical information and evaluation of key states.

2) To learn from an offline dataset, we introduce a model-based offline RL method that simultaneously learns a policy model and a generative world model via auto-regressive imitation of the actions and observations in medical decisions. The policy model employs a stochastic policy with entropy regularization to prevent it from overfitting and improve its robustness. Utilizing expert datasets, the generative world model uses an auto-regressive exploration generation paradigm to effectively discover a set of violating trajectories. Then, CT can infer constraints in healthcare through these unsafe trajectories and expert trajectories.

In the medical scenarios of sepsis and mechanical ventilation, we conduct experimental evaluations of offline CT. Experimental evaluations demonstrate that offline CT can capture patients’ unsafe states and assign higher penalties, thereby providing more interpretable constraints compared to previous works (Huang et al., ; Raghu et al., ; Peng et al., ). Compared to unconstrained, custom constraints and LLMs constraints (designed by Large Language Models (LLMs)), CT achieves strategies that closely approximate lower mortality rates with a higher probability (improving by $8.85\%$ compared to DDPG). To investigate the avoidance of unsafe behaviors with offline CT, we evaluate the probabilities of “too high” and “sudden changes” occurring in the sepsis. The experimental results show that CRL with CT can reduce the probability of unsafe behaviors to zero.

## 2 Related Works

Reinforcement Learning in Healthcare. RL has made great progress in the realm of healthcare, such as sepsis treatment (Huang et al., ; Raghu et al., ; Peng et al., ; Do et al., ), mechanical ventilation (Kondrup et al., ; Gong et al., ; Yu et al., ), sedation (Eghbali et al., ) and anesthesia (Calvi et al., ; Schamberg et al., ). However, the works mentioned above have not addressed potential safety issues such as sudden changes or too high medications. Therefore, the development of policies that are both safe and applicable across various healthcare domains is crucial.

Inverse Constrained Reinforcement Learning.
Previous works inferred constraint functions by determining the feasibility of actions under current states. In discrete state-action space, Chou et al. () and Park et al. () learned constraint sets to differentiate constrained state-action pairs. Scobee & Sastry () proposed inferring constraint sets based on the principle of maximum entropy, while some studies (McPherson et al., ; Baert et al., ) extended this approach to stochastic environments using maximum causal entropy (Ziebart et al., ). However, discrete approaches often face limitations when scaling to high-dimensional problems. As the state-action space increases, the computational cost rises significantly. This makes inference in large, discrete spaces challenging, requiring additional optimization techniques or assumptions. In continuous domains, Malik et al. (), Gaurav et al. (), and Qiao et al. () used neural networks to approximate constraints. Some works (Liu et al., ; Chou et al., ) applied Bayesian Monte Carlo and variational inference to infer the posterior distribution of constraints in high-dimensional state space. Xu & Liu () modeled uncertainty perception constraints for arbitrary and epistemic uncertainties. However, these methods can only be applied online and lack historical dependency.

Transformers for Reinforcement Learning.
Transformer has produced exciting progress on RL sequential decision problems (Zheng et al., ; Chen et al., ; Janner et al., ; Liu et al., ). These works no longer explicitly learn Q-functions or policy gradients, but focus on action sequence prediction models driven by target rewards. Chen et al. () and Janner et al. () perform auto-regressive trajectories modeling to achieve offline policy learning. Furthermore, Zheng et al. () unify offline pretraining and online fine-tuning within the Transformer framework. Liu et al. () and Kim et al. () integrate the transformer architecture into constraint learning and preference learning.
With its sequence modeling capability and independence from the Markov assumption, the transformer architecture can capture temporal dependencies in medical decision-making. Thus, it is well-suited for trajectory learning and personalized learning in medical settings.

## 3 Problem Formulation

Constrained Reinforcement Learning (CRL).
We model the medical environment with a Constrained Non-Markov Decision Process (Constrained Non-MDP) $\mathcal{N}^{c}$, which can be defined by a tuple $\left(\mathcal{S},\mathcal{A},h_{t},\mathcal{P},\mathcal{R},\mathcal{C},\gamma
,\kappa,\rho_{0},T\right)$ where: 1) $s\in\mathcal{S}$ denotes the state indicators of the patient at each time step. 2) $a\in\mathcal{A}$ denotes the administered drug doses or instrument parameters of interest. 3) $h_{t}=\{s_{0},a_{0},s_{1},a_{1},...,s_{t}\}$ is the treatment history, where $t$ represents the current time step. 4) $\mathcal{P}\left(s_{t+1}\mid h_{t},a_{t}\right)$ defines the transition probabilities.
5) The reward function $\mathcal{R}(h_{t},a_{t})$ is used to describe the quality of the patient’s condition and provided by experts based on prior work (Huang et al., ; Kondrup et al., ). 6) The constraint function $\mathcal{C}(h_{t},a_{t})$ describes the risk or cost associated with taking a particular action given the current historical information. 7) $\gamma$ denotes the discount factor. 8) $\kappa\in\mathbb{R}_{+}$ denotes the bound of cumulative costs. 9) $\rho_{0}$ defines the initial state distribution. 10) $T$ is the length of the trajectory $\tau$. At each time step $t$, an agent acts $a_{t}$ at a patient’s state $s_{t}$. This process generates the reward $r_{t}\sim\mathcal{R}(h_{t},a_{t})$, the cost $c_{t}\sim\mathcal{C}(h_{t},a_{t})$ and the next state $s_{t+1}\sim\mathcal{P}\left(s_{t+1}\mid h_{t},a_{t}\right)$. The goal of the CRL policy $\pi$ is to maximize expected discounted rewards while limiting the cost in a threshold $\kappa$:

$$ $\displaystyle\arg\max_{\pi}\mathbb{E}_{\pi,\rho_{0}}[{\textstyle\sum_{t=1}^{T} }\gamma^{t}r_{t}],\quad\text{ s.t.}\quad\mathbb{E}_{\pi,\rho_{0}}[{\textstyle \sum_{t=1}^{T}}\gamma^{t}c_{t}]\leq\kappa.$ (1) $$

CRL commonly assumes that constraint signals are directly observable. However, in healthcare, such signals are often difficult to obtain due to the variability in individual patient characteristics and the need for multi-objective evaluation.
Therefore, our objective is to infer reasonable constraints for CRL to achieve safe policy learning in healthcare.

Inverse Constrained Reinforcement Learning (ICRL).
ICRL (Malik et al., ) based on Markov Decision Process (MDP) $\mathcal{M}$ aims to recover the cost function $\mathcal{C}^{*}$ by leveraging a set of trajectories $\mathcal{D}=\{\tau^{(i)}\}_{i=1}^{N}$ sampled from an expert policy $\pi_{e}$, where $N$ denotes the number of the trajectories. ICRL is commonly based on the Maximum Entropy framework (Scobee & Sastry, ), and the likelihood function is articulated as:

$$ $p(\mathcal{D}\mid\mathcal{C})=\frac{1}{Z^{N}}\prod_{i=1}^{N}\exp\left[\mathcal {R}(\tau^{(i)})\right]\mathbb{I}(\tau^{(i)})$ (2) $$

Here, 1) $\tau=\left\{s_{0},a_{0},r_{0},s_{1},...\right\}$ is the trajectory from $\mathcal{D}$. 2) $Z=\int\exp(\beta r(\tau))\mathbb{I}(\tau)d\tau$ is the normalizing term, where $\beta\in[0,\infty)$ is a parameter that measures the proximity of the agent to the optimal distribution. As $\beta\to\infty$, the agent approaches optimal behavior, whereas as $\beta\to 0$, the agent’s actions become increasingly random. 3) The indicator $\mathbb{I}(\tau^{(i)})$ signifies the extent to which the trajectory $\tau^{(i)}$ satisfies the constraints. It can be approximated using a neural network $\zeta_{\theta}(\tau^{(i)})$ parameterized with $\theta$, defined as $\zeta_{\theta}(\tau^{(i)})=\prod_{t=0}^{T}\zeta_{\theta}(s_{t}^{(i)},a_{t}^{(i
)})$. Consequently, the cost function can be formulated as $C_{\theta}=1-\zeta_{\theta}$. Substituting the neural network for the indicator, we can update $\theta$ through the gradient of the log-likelihood function:

$$ $\displaystyle\nabla_{\theta}\mathcal{L}\left(\theta\right)=\mathbb{E}_{\tau \sim\pi_{e}}\left[\nabla_{\theta}\log[\zeta_{\theta}(\tau)]\right]-\mathbb{E}_ {\hat{\tau}}\sim\pi_{\mathcal{M}^{\hat{\zeta}_{\theta}}}\left[\nabla_{\theta} \log[\zeta_{\theta}(\hat{\tau})]\right]$ (3) $$

where $\tau$ is sampled from the expert policy $\pi_{e}$, while $\hat{\tau}$ is sampled from the executing policy $\pi_{\mathcal{M}^{\hat{\zeta}_{\theta}}}$. $\mathcal{M}^{\hat{\zeta}_{\theta}}$ denotes the MDP derived by augmenting the original MDP with the network $\hat{\zeta}_{\theta}$. The policy $\pi_{\mathcal{M}^{\hat{\zeta}_{\theta}}}$ is used to execute this augmented MDP.

Safe-Critical Decision Making with Constraint Inference in Healthcare.
Our general goal for our policy is to approximate the optimal policy, which refers to the strategy under which the patient’s mortality rate is minimized (achieving a zero mortality rate is often difficult since there are patients who can not recover, regardless of all potential future treatment sequences (Fatemi et al., )). Decision-making with constraints can formulate safer strategies by discovering and avoiding unsafe states, thereby approaching the optimal policy.

However, most offline RL algorithms rely on online evaluation, where the agent is evaluated in an interactive environment, whereas in medical scenarios, only offline policy evaluation (Luo et al., ) can be utilized. Besides, some works (Jia et al., ; Huang et al., ; Raghu et al., ; Komorowski et al., ) analyze by comparing the differences (DIFF) between the drug dosage recommended by the estimated policy $\pi$ and the dosage administered by clinical physicians $\hat{\pi}$, and the relationship of DIFF with mortality rates, through graphical analysis. In the graph depicting the relationship between the DIFF and mortality rate, at the point when DIFF is zero, the lower the mortality rate of patients, the better the performance of the policy (Raghu et al., ).

To provide a more accurate quantitative evaluation, we introduce the concept of the probability of approaching the optimal policy, defined as $\omega$:

$$ $\omega=\frac{\text{ Number of survivors among the top }N\text{ patients }}{N}$ (4) $$

1) We randomly select $2N$ patients from the dataset, where $N$ patients survived under the doctor’s treatment, and $N$ patients died.
2) For each patient in the real dataset, we have access to their state and the drug dosage administered by the doctor ($a$). Using an estimated policy, we compute an alternative drug dosage ($b$) for the same patient state.
3) For each patient, we calculate the difference between the dosages, defined as $\text{DIFF}=b-a$. This gives us $2N$ DIFF values across all patients.
4) We then sort the $2N$ DIFF values in ascending order. Next, we examine the survival status of the top $N$ patients based on the sorted DIFF values (using the survival data from the original doctor’s policy).
5) The top $N$ patients represent those for whom the difference between our policy and the doctor’s policy is smallest. If the survival rate of these top $N$ patients (denoted as $\omega$) is higher, it suggests that our policy is closer to the optimal policy. If our policy is the ideal optimal policy, the survival rate of the top $N$ patients will be the highest.
6) Additionally, we need to account for the magnitude of the DIFF values. For patients who survived, a smaller DIFF value is more desirable, as it indicates a closer alignment between our policy and the doctor’s policy.

## 4 Method

To infer constraints and achieve safe decision-making in healthcare, we introduce the Offline Constraint Transformer (shown in Figure 2), a novel ICRL framework.

In practice, ICRL can be conceptualized as a bi-level optimization task (Liu et al., ). We can 1) update this policy based on Equation 1, and 2) employ Equation 3 for constraint learning. Intuitively, the objective of Equation 3 is to distinguish between trajectories generated by expert policies and imitation policies that may violate the constraints. Specifically, task 1 involves updating the policy using advanced CRL methods. Significant progress has been made in some works such as BCQ-Lagrangian (BCQ-Lag) (Fujimoto et al., ), COpiDICE (Lee et al., ), VOCE (Guan et al., ), and CDT (Liu et al., ). Task 2 focuses on learning the constraint function, as shown in Figure 2. Our research primarily improves the latter process, addressing two key challenges that ICRL faces in healthcare: challenge 1 pertains to the limitations of the Markov property, and challenge 2 involves the issue of inferring constraints only from offline datasets. To address these challenges, we propose the offline CT as our solution.

Figure: Figure 2: The overview of the safe healthcare policy learning with offline CT.
Refer to caption: x2.png

Offline Constraint Transformer.
To address the first challenge, we delve into the inherent issues of applying the Markov property to healthcare and draw inspiration from sequence modeling tasks, redefining the representation of the constraints. To realize the offline training, we consider the essence of ICRL updates, proposing a model-based RL to generate unsafe behaviors used to train CT. We outline three parts: establishing the constraint representation model (Section 4.1), creating an offline RL for violating data (Section 4.2), and learning safe policies (Section 4.3).

### 4.1 Constraint Transformer

Figure: Figure 3: The structure of the CT.
Refer to caption: x3.png

ICRL methods relying on the Markov property overlook patients’ historical information, focusing only on the current state. However, both current and historical states, along with vital sign changes are crucial for a human doctor’s decision-making process (Plaisant et al., ). To emulate the observational approach of humans, we draw inspiration from the existing historical sequence model (such as Long Short-Term Memory (LSTM) (Graves & Graves, ) and Transformer (Vaswani, )) to incorporate historical information into constraints for a more comprehensive observation and judgment. Compared to LSTM, the Transformer effectively captures long-range dependencies and complex time series patterns due to its self-attention mechanism (Chen et al., ), without the need for sequential processing, which improves computational efficiency. Additionally, its dynamic attention weights provide better model interpretability, helping to understand the basis of the model’s decisions. Therefore, we propose a constraint modeling approach based on a causal attention mechanism, as shown in Figure 3. The structure comprises a causal Transformer for sequential modeling and a Non-Markovian layer for weighted constraints learning.

Sequential Modeling for Constraints Inference. For a trajectory segment of length $T$, $2T$ input embeddings are generated, with each position containing state $s$ and action $a$ embeddings, which are learned by a linear layer and a normalization layer. And the state and action at the same timestep share the same positional embedding, which is also learned. Then the input embeddings are fed into the causal Transformer, which produces output embeddings $\{d_{t}\}_{t=0}^{T}$ determined by preceding input embeddings from $\{s_{0},a_{0},...,s_{T},a_{T}\}$. Here, $d_{t}$ depends only on the previous $t$ states and actions.

Modeling Non-Markovian for Weighted Constraints Learning. Although $d_{t}$ represents the cost function $c_{t}$ derived from observations over long trajectories, it does not pinpoint which previous key actions or states led to its increase. In healthcare, identifying key actions or states is vital for analyzing risky behaviors and status, and enhancing model interpretability. To address this, we draw inspiration from the design of the preference attention layer in (Kim et al., ) and introduce an additional attention layer. This layer is employed to define the cost weight for Non-Markovian. It takes the output embeddings $\{d_{t}\}^{T}_{t=0}$ from the casual transformer as input and generates the corresponding cost and the cost weights. The output of the attention layer (i.e., the cost function) is computed by weighting the values through the normalized dot product between the query and other keys:

$$ $C(\tau)=\frac{1}{T}\sum_{i=0}^{T}\sum_{t=0}^{T}\operatorname{softmax}\left( \left\{\left\langle q_{i},k_{t^{\prime}}\right\rangle\right\}_{t^{\prime}=0}^{ T}\right)_{t}\cdot c_{t}=\sum_{t=0}^{T}w_{t}\cdot c_{t}$ (5) $$

Here, 1) the key $k_{t}\in\mathbb{R}^{m}$, query $q_{t}\in\mathbb{R}^{m}$, and value $c_{t}\in\mathbb{R}$ are derived from the $t$-th input $d_{t}$ through linear transformations, where $\mathbb{R}$ represents the set of real numbers and $m$ denotes the embedding dimension. Since $d_{t}$ depends only on the previous state-action pairs $\{(s_{i},a_{i})\}^{t}_{i=0}$ and serves as the input embedding for the Non-Markovian Layer, $c_{t}$ is also associated solely with the preceding $t$ time steps. 2) $w_{t}=\frac{1}{T}\sum_{i=0}^{T}\operatorname{softmax}\left(\left\{\left\langle
q
_{i},k_{t^{\prime}}\right\rangle\right\}_{t^{\prime}=0}^{T}\right)_{t}$ depends on the full sequence $\{(s_{t},a_{t})\}^{T}_{t=0}$ to model the cost importance weight.
Introducing the newly defined cost function, we redefine Equation 3 for CT as:

$$ $\nabla_{\phi}\mathcal{L}\left(\phi\right)=\mathbb{E}_{\tau_{v}\sim\mathcal{D}_ {v}}\left[\nabla_{\phi}\log[C_{\phi}(\tau_{v})]\right]-\mathbb{E}_{\tau_{e} \sim\mathcal{D}_{e}}\left[\nabla_{\phi}\log[\ C_{\phi}(\tau_{e})]\right]$ (6) $$

where $\phi$ is the parameter of CT. $\mathcal{D}_{e}$ and $\mathcal{D}_{v}$ represent the expert data and the violating data, respectively, while $\tau_{e}$ and $\tau_{v}$ are the trajectories from these datasets. This formulation implies that the constraint should be minimized on the expert policy and maximized on the violating policy. We construct an expert dataset and a violating dataset to evaluate Equation 6 offline. The expert data can be acquired from existing medical datasets or hospitals. Regarding the violating dataset, we introduce a generative model to establish it, as detailed in Section 4.2.

Figure: Figure 4: The structure of the model-based offline RL.
Refer to caption: x4.png

### 4.2 Model-based Offline RL

To train CT offline, we introduce a model-based offline RL method (shown in Figure 4) to generate violating data that refers to unsafe behavioral data and can be represented as $\tau_{v}=\{s_{0},a_{0},r_{0},s_{1},...\}\in\mathcal{D}_{v}$. The model simultaneously learns a policy model and a generative world model via auto-regressive imitation of the actions and observations in healthcare.
The model processes a trajectory, $\tau_{e}\in\mathcal{D}_{e}$, as a sequence of tokens encompassing the return-to-go, states, and actions, defined as $\{\hat{R}_{0},s_{0},a_{0},...,\hat{R}_{T},s_{T},a_{T}\}$. Notably, the return-to-go $\hat{R}_{t}$ at timestep $t$ is the sum of future rewards, calculated as $\hat{R}_{t}=\sum_{t^{\prime}=t}^{T}r_{t^{\prime}}$. At each timestep $t$, it employs the tokens from the preceding $K$ timesteps as its input, where $K$ represents the context length. Thus, the input tokens for it at timestep $t$ are denoted as $o_{t}=\{\hat{R}_{-K:t},s_{-K:t},a_{-K:t-1}\}$, where
$\hat{R}_{-K:t}=\{\hat{R}_{K},...,\hat{R}_{t}\}$, $s_{-K:t}=\{s_{K},...,s_{t}\}$ and $a_{-K:t-1}=\{a_{K},...,a_{t-1}\}$.

Policy Model. The input tokens are encoded through a linear layer for each modality. Subsequently, the encoded tokens pass through a casual transformer to predict future action tokens. To explore a diverse set of actions and improve performance, we employ a stochastic Gaussian policy (Liu et al., ). Furthermore, we incorporate a Shannon entropy regularizer $\mathcal{H}\left[\pi_{\vartheta}(\cdot\mid o)\right]$ to prevent policy overfitting and enhance robustness. The optimization objective is to minimize the negative log-likelihood loss while maximizing the entropy with weight $\lambda$:

$$ $\min_{\vartheta}\max_{\lambda}\quad\mathbb{E}_{o_{t}\sim\mathcal{D}_{e}}[-\log \pi_{\vartheta}(a_{t}\mid o_{t})-\lambda\mathcal{H}\left[\pi_{\vartheta}(\cdot \mid o_{t})\right]]$ (7) $$

where the policy $\pi_{\vartheta}\left(\cdot\mid o_{t}\right)=\tilde{a}_{t}=\mathcal{N}\left(\mu
_{\vartheta}\left(o_{t}\right),\Sigma_{\vartheta}\left(o_{t}\right)\right)$ adopts the stochastic Gaussian policy representation and $\vartheta$ is the policy parameter.

Generative World Model. To predict states and rewards, we use $x_{t}=\{o_{t}\cup a_{t}\}=\{\hat{R}_{-K:t},s_{-K:t},a_{-K:t}\}$ as the input, which is encoded by linear layers and passes through the casual transformer with two additional decoded layers to predict the current reward $\tilde{r}_{t-1}$ and the next state $\tilde{s}_{t}$. The optimization objective for the generative world model is to minimize the mean squared error for the current reward and next state, defined as:

$$ $\min_{\varphi,\mu}\mathbb{E}[{(r_{t}-L^{\tilde{r}}_{\varphi}(x_{t}))^{2}}+{(s_ {t+1}-L^{\tilde{s}}_{\mu}(x_{t}))^{2}}]$ (8) $$

where $L^{\tilde{r}}_{\varphi}$ and $L^{\tilde{s}}_{\mu}$ are the reward and state generation network for the generative world model, with parameters $\varphi$ and $\mu$.

In the model-based offline RL framework, the policy model and the generative world model have the objectives of generating actions, rewards, and the next state, respectively. The causal transformer structure is used to extract historical information for both the policy and the generative world models. During training, the causal transformer is trained alongside the above models, with the goal of simultaneously minimizing Equations 7 and 8 until the convergence of the models.

Generating Violating Data. In RL, excessively high rewards, surpassing those provided by domain experts, may incentivize agents to violate the constraints to maximize the total reward (Liu et al., ; ). Therefore, we set a high initial target reward $\hat{R}_{1}$ to obtain violation data. We feed $\hat{R}_{1}$ and initial state $s^{(i)}_{1}$ into the model-based offline RL to generate $\tau^{(i)}_{v}$ in an auto-regressive manner, as depicted in model-based offline RL of Figure 2, where $\widetilde{a}$, $\widetilde{r}$ and $\widetilde{s}$ are predicted by the model. The target reward $\hat{R}$ decreases incrementally and can be represented as $\hat{R}_{t+1}=\hat{R}_{t}-\widetilde{r}_{t}$. Considering the average error in trajectory prediction, we generate trajectories with the length $K=10$. Repeating $N$ initial states, we can get violating data $\mathcal{D}_{v}=\{\tau_{v}^{(i)}\}_{i=1}^{N}$. The detailed parameter settings and sensitivity analysis can be found in Appendix C.

Note that certain other generative models, such as Variational Auto-Encoder (VAE) (Kim et al., ), Generative Adversarial Networks (GAN) (Hsu et al., ; Iyer et al., ), and Denoising Diffusion Probabilistic Models (DDPM) (Croitoru et al., ), may be better at generating data. We introduce the model-based offline RL primarily because it has been shown to generate violating data with exploration (Liu et al., ) and possess the ability to process time-series features efficiently.

### 4.3 Safe-Critical Decision Making with Constraints

To train offline CT, we gather the medical expert dataset $\mathcal{D}_{e}$ from the environment. Then, we employ gradient descent to train the model-based offline RL, guided by Equation 7 and Equation 8, continuing until the model converges. Using this RL model, we automatically generate violating data denoted as $\mathcal{D}_{v}$. Subsequently, CT is optimized based on Equation 6 to get the cost function $C$, leveraging samples from both $\mathcal{D}_{e}$ and $\mathcal{D}_{v}$. To learn a safe policy, we train the policy $\pi$ using $C$ until it converges based on Equation 1. The detailed training procedure is presented in Algorithm 1.

Figure: Algorithm 1 Safe Policy Learning with Offline CT

## 5 Experiment

In this section, we first provide a brief overview of the task, as well as data extraction and preprocessing. Subsequently, in Section 5.1, we demonstrate that CT can describe constraints in healthcare and capture critical patient states. We emphasize its applicability to various CRL methods and its ability to approach the optimal policy for reducing mortality rates in Section 5.2. Section 5.3 discusses the realization of the objective of safe medical policies. Finally, we use offline policy evaluation (OPE) methods to estimate our policy in the field of dynamic treatment regimes in Section D.2.2.

Tasks. We primarily use the sepsis task that is commonly used in previous works (Huang et al., ; Raghu et al., ; Komorowski et al., ; Do et al., ), and supplement some experiments on the mechanical ventilator task (Kondrup et al., ; Peine et al., ).
The detailed definition of the two tasks mentioned above can be found in Appendix A.1 and A.2. For detailed experiments on the mechanical ventilator task, please refer to Appendix D.2.

Data Extraction and Pre-processing.
Our medical dataset is derived from the Medical Information Mart for Intensive Care III (MIMIC-III) database (Johnson et al., ). For each patient, we gather relevant physiological parameters, including demographics, lab values, vital signs, and intake/output events. Data is grouped into 4-hour windows, with each window representing a time step. In cases of multiple data points within a step, we record either the average or the sum. We eliminate variables with significant missing values and use the $k$-nearest neighbors method to fill in the rest.

Model-based Offline RL Evaluation.
To ensure the rigor of the experiments, we evaluate the validity of the model-based offline RL, as detailed in Appendix C.

### 5.1 Can Offline CT Learn Effective Constraints?

Figure: Figure 5: The relationship between cost and mortality.
Refer to caption: x5.png

In this section, we primarily assess the efficacy of the cost function learned by offline CT in sepsis, focusing on its capability to evaluate patient mortality rates and capture critical events. First, we employ the cost function to compute cost values for the validation dataset. Subsequently, we statistically analyze the relationship between these cost values and mortality rates. As shown in Figure 5, there is an increase in patient mortality rates with rising cost values. It is noteworthy that such increases in mortality rates are often attributed to suboptimal medical decisions. Therefore, these experimental findings affirm that the cost values effectively reflect the quality of medical decision-making. To observe the impact of the attention layer (Non-Markovian layer), we conduct experiments by removing the attention layer from CT. The results reveal that the penalty values do not correlate proportionally with mortality rates (shown in Figure 5). This indicates that the attention layer plays a crucial role in assessing constraints.

Figure: Figure 6: The relationship between physiological indicators and cost values. As SOFA and lactate levels become increasingly unsafe, the cost increases. Mean BP and HR at lower values within the safe range incur a lower cost, but as they move into unsafe ranges, the cost increases, penalizing previous state-action pairs. The cost can differentiate between relatively safe and unsafe regions.
Refer to caption: x6.png

To assess the capability of the cost function to capture key events, we analyze the relationship between physiological indicators and cost values. We focus on four key indicators in sepsis treatment: Sequential Organ Failure Assessment (SOFA) score (Li et al., ), lactate levels (Ryoo et al., ), Mean Arterial Pressure (MeanBP) (Pandey et al., ), and Heart Rate (HR) (Carrara et al., ). The SOFA score and lactate levels are critical indicators for assessing sepsis severity, with higher values indicating greater patient risk. MeanBP and HR are essential physiological metrics, typically ranging from $70$ to $100$ mmHg and $60$ to $100$ beats, respectively. Deviations from these ranges can signify patient risk. As depicted in Figure 6, the cost values effectively distinguish between high-risk and safe conditions, reflecting changes in patient status. Moreover, we demonstrate that the cost function can capture the dangerous states of other feature variables, including hidden variables. For more detailed information, refer to Appendix D.2.

### 5.2 Can Offline CT Improve the Performance of CRL?

Baselines.
We adopt the DDPG method as the baseline in sepsis research (Huang et al., ). Since no other offline inverse reinforcement learning works are available for reference, we have included three additional settings: no cost, custom cost, and LLMs cost. In the case of no cost, the cost is set to zero, while the design of custom constraints and LLMs cost are outlined in Appendix B. These settings help evaluate whether CT can infer effective constraints.

Metrics. To assess effectiveness, we use $\omega$ to indicate the probability that the policy is optimal and analyze the relationship between DIFF and mortality rate through a graph.

**Table 2: Performance of sepsis strategies under various offline CRL models and different constraints.**
| Model | Cost | $\omega_{\text{IV DIFF}}$% $\uparrow$ | $\omega_{\text{VASO DIFF}}$% $\uparrow$ | $\omega_{\text{ACTION DIFF}}$% $\uparrow$ |
| --- | --- | --- | --- | --- |
| DDPG | - | 50.95$\pm$1.34 | 51.45$\pm$0.75 | 51.15$\pm$1.15 |
|  | No cost | 47.45$\pm$0.52 | 46.35$\pm$1.82 | 51.00$\pm$0.86 |
|  | Custom cost | 46.45$\pm$0.46 | 52.00$\pm$0.98 | 49.40$\pm$1.04 |
|  | LLMs cost | 48.15$\pm$1.23 | 48.90$\pm$0.77 | 50.70$\pm$1.68 |
| VOCE | CT | 53.33$\pm$0.94 | 59.04$\pm$1.13 | 56.15$\pm$1.08 |
|  | No cost | 48.30$\pm$0.91 | 60.10$\pm$0.60 | 51.25$\pm$0.70 |
|  | Custom cost | 53.05$\pm$1.35 | 55.20$\pm$0.24 | 53.90$\pm$1.04 |
|  | LLMs cost | 51.05$\pm$1.50 | 58.95$\pm$0.38 | 54.35$\pm$0.89 |
| CopiDICE | CT | 51.95$\pm$0.41 | 60.85$\pm$1.08 | 54.60$\pm$0.60 |
|  | No cost | 47.50$\pm$1.32 | 51.05$\pm$0.61 | 49.35$\pm$1.08 |
|  | Custom cost | 51.54$\pm$0.16 | 56.23$\pm$1.43 | 53.69$\pm$1.62 |
|  | LLMs cost | 56.44$\pm$0.75 | 53.59$\pm$1.15 | 57.88$\pm$0.72 |
| BCQ-Lag | CT | 52.45$\pm$1.01 | 55.34$\pm$1.20 | 54.49$\pm$0.86 |
|  | No cost | 56.50$\pm$0.81 | 62.45$\pm$1.20 | 58.90$\pm$1.34 |
|  | Custom cost | 54.70$\pm$1.12 | 59.85$\pm$1.51 | 57.80$\pm$1.00 |
|  | LLMs cost | 52.45$\pm$0.80 | 60.15$\pm$1.17 | 56.35$\pm$1.59 |
| CDT | CT | 57.15$\pm$1.67 | 65.20$\pm$1.22 | 60.00$\pm$1.49 |
| CDT | Without CT | 56.50$\pm$0.81 | 62.45$\pm$1.20 | 58.90$\pm$1.34 |
| CDT | No attention layer | 56.70$\pm$0.64 | 62.50$\pm$1.57 | 59.10$\pm$0.44 |
| Generative<br>Model | - | 55.49$\pm$2.55 | 56.60$\pm$1.33 | 57.00$\pm$2.06 |

Results. We combine our method CT with CRL algorithms (e.g., VOCE, COpiDICE, BCQ-Lag, and CDT), and compare them with both no-cost and custom cost settings. Each CRL model is trained using no cost, custom cost, and CT separately, with other parameters set the same during training. For evaluation metrics, we use IV difference (IV DIFF), vaso difference (VASO DIFF), and combined [IV, VASO] difference (ACTION DIFF) as the metrics to be ranked. We measure the mean and variance of $\omega\%$ in $10$ sets of random seeds, and the results are shown in Table 2. From the results, we can conclude: 1) CT makes the strategies in the VOCE, CopiDICE, and CDT methods closer to the lower mortality strategies. 2) CDT+CT achieves better results on all three metrics. CDT is also a transformer-based method, which indicates that transformer-based architecture indeed exhibits more outstanding performance in healthcare.

Figure: Figure 7: The relationship between DIFF and the mortality rate. The x-axis represents the DIFF. The y-axis indicates the mortality rate of patients at a given DIFF. The solid line represents the mean, while the shaded area indicates the Standard Error of the Mean (SEM).
Refer to caption: x7.png

Figure 7 shows the relationship between IV and VASO DIFF with mortality rates under the DDPG and CDT+CT methods in sepsis. In VASO DIFF, when the gap is zero, the mortality rate under CDT+CT is lower than that under DDPG, indicating that following the former strategy could lead to a lower mortality rate. Similarly, in IV DIFF, the same trend is observed. Notably, for the IV strategy, the lowest mortality rate for DDPG does not occur at the point where the difference is zero, indicating a significant estimation bias.

### 5.3 Can CRL with Offline CT Learn Safe Policies?

We have confirmed the existence of two unsafe strategy issues, namely “too high” and “sudden change” in the treatment of sepsis, particularly in vaso in Section 1. To validate whether the CRL+CT approach could address these concerns, we employ the same statistical methods to evaluate our methodology, shown in Table 3. To elucidate the efficacy of CT, we compare it with CDT+No-cost and CDT+Custom-cost approaches. We find that only the custom cost and CT methods successfully mitigated these unsafe behaviors. However, the custom cost mitigates risks by avoiding medication ($max=0$). This is an overly conservative strategy. The CDT+CT approach can give a more appropriate drug dosage and is not an overly conservative strategy.
In addition, there may be other safety issues that we have not yet verified. Future testing can be conducted on simulation systems such as DTR-Bench (Luo et al., ), detailed in Appendix E.

### 5.4 Off-policy evaluation

**Table 3: Proportion of unsafe actions recommended by policies.**
| Drug dosage<br>$(\mu g/(kg\cdot min))$ | Physician | DDPG | CDT |  |  |
| --- | --- | --- | --- | --- | --- |
| No cost | Custom cost | CT |  |  |  |
| vaso $\textgreater 0.75$ | $2.27\%$ | $7.44\%$ | $0.13\%$ |  |  |
| vaso $\textgreater 0.9$ | $1.71\%$ | $7.40\%$ | $0.09\%$ | 0% $\downarrow$<br>(max $=0.00$) | 0% $\downarrow$<br>(max $=0.11$) |
| $\Delta$ vaso $\textgreater 0.75$ | $2.45\%$ | $21.00\%$ | $0.64\%$ |  |  |
| $\Delta$ vaso $\textgreater 0.9$ | $1.88\%$ | $20.62\%$ | $0.48\%$ | 0% $\downarrow$<br>(max $\Delta=0.00$) | 0% $\downarrow$<br>(max $\Delta=0.10$) |

Baselines. 1) Naive baselines. A naive baseline can provide worst-case scenario benchmarks for algorithm evaluation (Luo et al., ), including random policy $\pi_{r}$, zero-drug policy $\pi_{\text{min}}$, max-drug policy $\pi_{\text{max}}$, alternating policy $\pi_{\text{alt}}$, and weight policy $\pi_{\text{weight}}$.
2) RL methods baselines. We select common RL methods such as Deep Q-Network (DQN), Conservative Q-Learning (CQL), Implicit Q-Learning (IQL), and Batch Constrained Q-Learning (BCQ) as baseline models.

Metrics. A recent series of studies have applied offline policy evaluation techniques to dynamic treatment regimes, including Weighted Importance Sampling (WIS) (Kidambi et al., ; Nambiar et al., ) and Doubly Robust (DR) estimators (Raghu et al., ; Wu et al., ; Wang et al., ). To more accurately evaluate the policy, we use metrics such as RMSE and F1 score to describe the deviation from the clinician’s policy.

We used the same reward function to compare the policy results under different evaluation metrics, as shown in Table 4. Our findings present that the CDT+CT method outperforms other methods in terms of $\text{RMSE}_{\text{IV}}$, WIS, $\text{WIS}_{b}$, and $\text{WIS}_{\text{bt}}$ evaluation metrics. Since CDT+CT produces more safe and conservative policies, there is a certain distance from the clinician’s policy, so it does not perform as well in terms of $\text{RMSE}_{\text{vaso}}$ and F1 score.

**Table 4: Comparison across policies on the sepsis test set. The best algorithms are highlighted in red. $\text{RMSE}_{\text{IV}}$ and $\text{RMSE}_{\text{VASO}}$ mean the RMSE loss for the IV fluid treatment and vasopressor treatment. P.F1 and S.F1 denote the patient-wise F1 and sample-wise F1.**
| Metric | alt | max | min | random | weight | DQN | CQL | IQL | BCQ | CDT+CT |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| $\text{RMSE}_{\text{IV}}\downarrow$ | $763.89$ | $861.51$ | $645.83$ | $671.39$ | $645.83$ | $638.51\pm 8.63$ | $541.67\pm 5.74$ | $578.96\pm 10.06$ | $626.2\pm 9.56$ | $433.55\pm 7.20$ |
| $\text{RMSE}_{\text{VASO}}\downarrow$ | $0.67$ | $0.89$ | $0.32$ | $0.5$ | $0.59$ | $0.44\pm 0.07$ | $0.30\pm 0.01$ | $0.31\pm 0.01$ | $0.31$ | $1.13\pm 0.01$ |
| WIS $\uparrow$ | $-4.58$ | $-4.62$ | $-4.58$ | $-3.84$ | $-3.78$ | $-3.79\pm 0.01$ | $-4.10\pm 1.43$ | $-5.83$ | $-4.58$ | $-3.51\pm 0.11$ |
| $\text{WIS}_{b}$ $\uparrow$ | $-5.43$ | $-4.81$ | $-5.76$ | $-4.4$ | $-4.73$ | $-3.88\pm 0.73$ | $-4.48\pm 0.77$ | $-5.31\textpm 0.06$ | $-5.41\pm 0.17$ | $-3.52\pm 0.17$ |
| $\text{WIS}_{t}$ $\uparrow$ | $-4.58$ | $-4.62$ | $-4.58$ | $-3.97$ | $-3.78$ | $-3.84\pm 0.11$ | $-4.10\pm 1.43$ | $-5.83$ | $-4.58$ | $-3.51\pm 0.11$ |
| $\text{WIS}_{bt}$ $\uparrow$ | $-5.64$ | $-4.69$ | $-5.61$ | $-4.5$ | $-4.5$ | $-3.87\pm 0.67$ | $-4.38\pm 0.98$ | $-5.27\pm 0.05$ | $-5.55\pm 0.19$ | $-3.52\pm 0.17$ |
| DR $\uparrow$ | $-0.54$ | $-0.19$ | $-1.55$ | $-0.35$ | $-0.3$ | $-0.14\pm 0.04$ | $-0.71\pm 0.05$ | $-0.51\pm 0.04$ | $-1.54\pm 0.01$ | $-3.08$ |
| P.F1 $\uparrow$ | $0.2$ | $0.02$ | $0.2$ | $0.2$ | $0.0$ | $0.06\pm 0.02$ | $0.33\pm 0.01$ | $0.34\pm 0.01$ | $0.23\pm 0.01$ | $0.17\pm 0.02$ |
| S.F1 $\uparrow$ | $0.19$ | $0.02$ | $0.19$ | $0.19$ | $0.0$ | $0.06\pm 0.02$ | $0.32\pm 0.01$ | $0.33\pm 0.01$ | $0.22\pm 0.01$ | $0.16\pm 0.02$ |

Ablation Study.
To investigate the impact of each component on the model’s performance, we conducted experiments by sequentially removing each component from the CDT+CT model. The results are presented in the lower half of Table 2. Both CT and its Non-Markovian layer (attention layer) are essential components; removing either one results in a decrease in performance. Additionally, we observed that even a pure generative model outperforms DDPG in terms of performance. This is primarily because it inherently operates as a sequence-based reinforcement learning model, possessing exploration and consideration for long-term history. Therefore, this further underscores the effectiveness of sequence-based approaches in healthcare applications. To further analyze the performance of different sequence models, we conduct offline policy evaluation on models based on LSTM and transformer architectures. We found that the latter performs better, see Appendix D.3.

## 6 Conclusion

In this paper, we propose offline CT, a novel ICRL algorithm designed to address safety issues in healthcare. This method utilizes a causal attention mechanism to observe patients’ historical information, similar to the approach taken by actual doctors, and employs Non-Markovian importance weights to effectively capture critical states. To achieve offline learning, we introduce a model-based offline RL for exploratory data augmentation to discover unsafe decisions and train CT. Experiments in sepsis and mechanical ventilation demonstrate that our method avoids risky behaviors while achieving strategies that closely approximate the lowest mortality rates.

Limitations.
There are also several limitations of offline CT:
1) Lack of rigorous theoretical analysis. We did not precisely define the types of constraint sets, thereby conducting rigorous theoretical analysis on constraint sets remains challenging. 2) Need for more computational resources. Due to the Transformer architecture, more computational resources are required. 3) Unrealistic assumptions of expert demonstrations. We assume that expert demonstrations are optimal in both constraint satisfaction and reward maximization. However, in reality, this assumption may not always hold. Therefore, researching a more effective approach to address the aforementioned issues holds promise for the field of secure medical reinforcement learning.

## Appendix A Problem Define

### A.1 Sepsis Problem Define

Our definition is similar to (Raghu et al., ). We extract data from adult patients meeting the criteria for sepsis-3 criteria (Singer et al., ) and collect their data within the first 72 hours of admission.

State Space.
We use a 4-hour window and select 48 patient indicators as the state for a one-time unit of the patient. The state indicators include Demographics/Static, Lab Values, Vital Signs, and Intake and Output Events, detailed as follows (Raghu et al., ):

- •
Demographics/Static: Shock Index, Elixhauser, SIRS, Gender, Re-admission, GCS - Glasgow Coma Scale, SOFA - Sequential Organ Failure Assessment, Age
- •
Lab Values Albumin: Arterial pH, Calcium, Glucose, Hemoglobin, Magnesium, PTT - Partial Thromboplastin Time, Potassium, SGPT - Serum Glutamic-Pyruvic Transaminase, Arterial Blood Gas, BUN Blood Urea Nitrogen, Chloride, Bicarbonate, INR - International Normalized Ratio, Sodium, Arterial Lactate, CO2, Creatinine, Ionised Calcium, PT - Prothrombin Time, Platelets Count, SGOT Serum Glutamic-Oxaloacetic Transaminase, Total bilirubin, White Blood Cell Count
- •
Vital Signs: Diastolic Blood Pressure, Systolic Blood Pressure, Mean Blood Pressure, PaCO2, PaO2, FiO2, PaO/FiO2 ratio, Respiratory Rate, Temperature (Celsius), Weight (kg), Heart Rate, SpO2
- •
Intake and Output Events: Fluid Output - $4$ hourly period, Total Fluid Output, Mechanical Ventilation

Action Space.
Regarding the treatment of sepsis, there are two main types of medications: intravenous fluids and vasopressors. We select the total amount of intravenous fluids for each time unit and the maximum dose of vasopressors as the two dimensions of the action space, defined as ($\text{sum}\text{(IV)},\max\text{(Vaso)}$). Each dimension is a continuous value greater than 0 0.

Reward Function.
We refer to the reward function used in (Huang et al., ), as shown in the following equation:

$$ $\left.r\left(s_{t},s_{t+1}\right)=\lambda_{1}\tanh\left(s_{t}^{\mathrm{SOFA}}- 6\right)+\lambda_{2}\left(s_{t+1}^{\mathrm{SOFA}}-s_{t}^{\mathrm{SOFA}}\right)\right)$ (9) $$

Where $\lambda_{0}$ and $\lambda_{1}$ are hyperparameters set to $-0.25$ and $-0.2$, respectively. This reward function is designed based on the SOFA score, as it is a key indicator of the health status of sepsis patients and is widely used in clinical settings. The formula describes a penalty when the SOFA score increases and a reward when the SOFA score decreases. We set $6$ as the cutoff value because the mortality rate sharply increases when the SOFA score exceeds $6$ (Ferreira et al., ).

### A.2 Mechanical Ventilation Treatment Problem Define

The RL problem definition for Mechanical Ventilation Treatment is referenced from (Kondrup et al., ).

State Space. We also use a 4-hour window and select 48 patient indicators as the state for a one-time unit of the patient. The state indicators are as follows:

- •
Demographics/Static: Elixhauser, SIRS, Gender, Re-admission, GCS, SOFA, Age
- •
Lab Values Albumin: Arterial pH, Glucose, Hemoglobin, Magnesium, PTT, BUN Blood Urea Nitrogen, Chloride, Bicarbonate, INR, Sodium, Arterial Lactate, CO2, Creatinine, Ionised Calcium, PT, Platelets Count, White Blood Cell Count, Hb
- •
Vital Signs: Diastolic Blood Pressure, Systolic Blood Pressure, Mean Blood Pressure, Temperature, Weight (kg), Heart Rate, SpO2
- •
Intake and Output Events: Urine output, vasopressors, intravenous fluids, cumulative fluid balance

Action Space.
The action space mainly consists of Positive End Expiratory Pressure (PEEP) and Fraction of Inspired Oxygen (FiO2), which are crucial parameters in ventilator settings. Here, we consider a discrete space configuration, with each parameter divided into $7$ intervals. Therefore, our action space is $7\times 7$, depicted as 5.

**Table 5: The action space of the mechanical ventilator.**
| Action | 0 | 1 | 2 | 3 | 4 | 5 | 6 |
| --- | --- | --- | --- | --- | --- | --- | --- |
| PEEP($cmH_{2}0$) | 0-5 | 5-7 | 7-9 | 9-11 | 11-13 | 13-15 | $>$15 |
| FiO2(Percentage(%)) | 25-30 | 30-35 | 35-40 | 40-45 | 45-50 | 50-55 | $>$55 |

Reward Function.
The primary objective of setting respiratory parameters is to ensure the patient’s survival. We adopt the same reward function design as the work (Kondrup et al., ), defined as Equation 10. This reward function first considers the terminal reward: if the patient dies, the reward $r$ is set to $-1$; otherwise, it is $+1$ in the terminal state. Additionally, to provide more frequent rewards, intermediate rewards are considered. Intermediate rewards mainly focus on the Apache II score, which evaluates various parameters to describe the patient’s health status. This reward function utilizes the increase or decrease in this score to reward the agent.

$$ $r\left(s_{t},a_{t},s_{t+1}\right)=\left\{\begin{array}[]{ll}+1&\text{ if }t=T \text{ and }m_{t}=1\\ -1&\text{ if }t=T\text{ and }m_{t}=0\\ \frac{\left(A_{t+1}-A_{t}\right)}{\max_{A}-\min_{A}}&\text{ otherwise }\end{ array}\right.$ (10) $$

In Equation 10, $T$ represents the length of the patient’s trajectory, $m$ indicates whether the patient ultimately dies, $A$ denotes the Apache II score, and $\max_{A}$ and $\min_{A}$ respectively denote the maximum and minimum values.

## Appendix B Design and Analysis of the Custom and LLMs cost Function

### B.1 Custom cost function

We base our design on prior knowledge that intravenous (IV) intake exceeding $2000mL/4h$ or vasopressor (Vaso) dosage surpassing $1\textmu g/(kg\cdot min)$ is generally considered unsafe in sepsis treatment (Shi et al., ). To design a reasonable constraint function, we refer to the constraint function designed by Liu *et al.* in the Bullet safety gym environments (Liu et al., ). We define the cost function as shown in Equation 11. Thus, during the treatment of sepsis, if the agent exceeds the maximum dosage thresholds of the two medications, it incurs a cost due to constraint violation.

$$ $c\left(s,a\right)\ =\bm{1}\left(a_{\text{IV}}>a_{\text{IV}_{\text{max}}}\right )+\bm{1}\left(a_{\text{Vaso}}>a_{\text{Vaso}_{\text{max}}}\right)$ (11) $$

where, $s$ and $a$ represent the patient’s state and action, respectively. $a_{IV\ \max}=2000$ indicates that the maximum fluid intake through IV is $2000mL$, and $a_{Vaso\ \max}=1$ signifies that the maximum Vaso dosage is $1\mu g/(kg\cdot min)$.

We applied our custom constraint function in the CDT (Liu et al., ) method, and the results are shown in Figure 8. Compared to the Vaso dosage recommended by doctors, our strategy exhibits excessive suppression of the Vaso. The maximum dosage of Vaso is $0.0011\mu g/(kg\cdot min)$, which is minimal and insufficient to provide the patient with effective therapeutic effects.

Therefore, Equation 11 is not suitable. The primary issues may include uniform constraint strength for excessive drug dosages, for instance, the cost for IV exceeding $2000$ mL and IV exceeding $3000$ mL is the same at $1$; lack of generalization, where the constraint cost does not vary with the patient’s tolerance. If a patient has an intolerance to VASO, the maximum value for VASO maybe 0 0, which cannot be captured by the self-imposed constraint function. Moreover, it lacks generalization, requiring redesign of the constraint function when addressing other unsafe medical issues; and it’s essential to ensure the correctness of the underlying medical knowledge premises.

Figure: Figure 8: Drug dosage distribution under custom constraint functions in sepsis.
Refer to caption: x8.png

### B.2 LLMs cost function

We provide prior knowledge to GPT-4.0, and the cost function it designs is as shown in Equation 12. Based on the self-designed constraint function, LLMs added a penalty for Vaso doses mutations, giving the agent a certain penalty when the change in Vaso doses exceeds the threshold.

$$ $c\left(s,a\right)\ =\bm{1}\left(a_{\text{IV}}>a_{\text{IV}_{\text{max}}}\right )+\bm{1}\left(a_{\text{Vaso}}>a_{\text{Vaso}_{\text{max}}}\right)+\bm{1}((a_{ \text{vaso}}-a_{\text{Vaso}_{\text{prev}}})>a_{\text{vaso}_{\text{change_ threshold}}})$ (12) $$

## Appendix C The Evaluation of Model-based Offline RL

### C.1 The length of a trajectory.

Regarding the selection of trajectory length, we consider the relationship between the average prediction error, the error of the last point in the trajectory, and the trajectory length. We use the model-based offline RL to generate trajectories and compare them with expert data using the Euclidean distance to measure their differences. We evaluate the average error and the error of the last point in the trajectory, as shown in Figure 10. We observe that with an increase in trajectory length, the average prediction error at each time step decreases, while the state error stabilizes. Taking into account the observation length and prediction accuracy, we ultimately choose to generate trajectories with lengths ranging from $10$ to $15$.

Figure: Figure 9: The relationship between average prediction error and trajectory length.
Refer to caption: x9.png

### C.2 Generating data within a reasonable range.

To validate model-based offline RL, we first check whether the values it produces fall within the legal range. The results are depicted in Figure 10. After analyzing the generated data, we find that the majority of state values have a probability of over $99\%$ of being within the legal range. A few values related to gender and re-admission range between $60\%$ and $70\%$. This could be due to these two indicators having limited correlation with other metrics, making them more challenging for the model to assess.

### C.3 Generating violating data.

In addition, we evaluate the violating actions generated by the model, as shown in Figure 11. When compared with expert strategies and penalty distributions, we find that the actions generated by the model mostly fall within the legal range. However, it occasionally produces behaviors that are inappropriate for the current state, constituting violating data. This indicates that our generative model can produce legally violating data.

Figure: Figure 11: The distribution and penalty values of violating data and expert data.
Refer to caption: extracted/5924295/figure/generation_action_cost3.png

### C.4 The sensitivity of CT models to generative models and reward setting.

We designed the following experiment to explore the sensitivity of the estimated policy to the generative world model. Since the quality of data generated by the generative world model depends on the target reward, higher target rewards lead the world model to generate more aggressive data to obtain more rewards. We set the target rewards to 1, 5, 10, 40, and 50, and observed the impact of the generated data on the policy, as shown in Table 6. The policy performance improves as the target reward increases, but it reaches an upper bound and does not increase indefinitely.

**Table 6: The impact of generative world models with different target rewards on policy estimation.**
| Target Reward | IV DIFF | VASO DIFF | ACTION DIFF |
| --- | --- | --- | --- |
| 1 | $51.60\pm 1.78$ | $58.8\pm 2.74$ | $54.25\pm 1.79$ |
| 5 | $52.50\pm 1.46$ | $58.84\pm 3.24$ | $54.45\pm 1.65$ |
| 10 | $52.25\pm 1.33$ | $56.85\pm 4.20$ | $55.00\pm 1.80$ |
| 40 | $52.05\pm 1.30$ | $56.75\pm 3.13$ | $55.80\pm 1.76$ |
| 50 | $52.00\pm 1.31$ | $57.35\pm 2.09$ | $55.05\pm 1.93$ |

## Appendix D THE EVALUATION OF COST FUNCTION

### D.1 The Evaluation of Cost function in Sepsis

#### D.1.1 Capture unsafe variables

To validate that the CT method captures key states, we conduct statistical analysis on the relationship between state values and penalty values. We collect penalty values under different state values for all patients, and the complete information is shown in Figure 13. We find that the CT method successfully captures unsafe states and imposes higher penalties accordingly. The safe range of state values is shown in Table 7.

**Table 7: State indicators and their normal ranges.**
| Indicator | Safe Range | Indicator | Safe Range | Indicator | Safe Range |
| --- | --- | --- | --- | --- | --- |
| Albumin | 3.5$\sim$5.1 | HCO3 | 25$\sim$40 | SGOT | 0$\sim$40 |
| Arterial_BE | -3$\sim$+3 | Glucose | 70$\sim$140 | SGPT | 0$\sim$40 |
| Arterial_lactate | 0.5$\sim$1.7 | HR | 60$\sim$100 | SIRS | $\downarrow$ |
| Arterial_PH | 7.35$\sim$7.45 | Hb | 12$\sim$16 | SOFA | $\downarrow$ |
| BUN | 7$\sim$22 | INR | 0.8$\sim$1.5 | Shock_Index | $\downarrow$ |
| CO2_mEqL | 20$\sim$34 | MeanBP | 70$\sim$100 | Sodium | 135$\sim$145 |
| Calcium | 8.6$\sim$10.6 | PT | 11$\sim$13 | SpO2 | 95$\sim$99 |
| Chloride | 96$\sim$106 | PTT | 23$\sim$37 | SysBP | 90$\sim$139 |
| Creatinine | 0.5$\sim$1.5 | PaO2_FiO2 | 400$\sim$500 | Temp_C | 36.0$\sim$37.0 |
| DiaBP | 60$\sim$89 | Platelets_count | 125$\sim$350 | WBC_count | 4$\sim$10 |
| FiO2 | 0.5$\sim$0.6 | Potassium | 4.1$\sim$5.6 | PaCO2 | 35$\sim$45 |
| GCS | $\uparrow$ | RR | 12$\sim$20 | PaO2 | 80$\sim$100 |

#### D.1.2 Capture unsafe hidden variables

In a medical context, mortality rates may be influenced by various factors. The dataset often contains numerous unaccounted features (hidden variables), such as epinephrine, dopamine, medical history, and phenotypes. As noted in (Jeter et al., ), clinicians typically set a mean arterial pressure (MAP) target (e.g., 65) and administer vasopressors until the patient reaches a safe pressure level. Additionally, Luo et al. () suggest using the NEWS2 score as evidence for clinical rewards. To validate whether our penalty function captures changes in hidden variables (NEWS and MAP), we conducted supplementary experiments, as shown in Figure 12. When the NEWS score is excessively high, the penalty value increases accordingly; similarly, when MAP falls outside the normal range, the penalty also rises. This indicates that the penalty function successfully captures changes in hidden variables and compensates for the reward function’s omission of certain parameter variables. Therefore, we can rely on a simple reward function and use the penalty function to achieve safe policy learning.

Figure: Figure 12: The relationship between the NEWS2 and MAP indicators with cost values.
Refer to caption: x11.png

Figure: Figure 13: The relationship between all states and cost values
Refer to caption: x12.png

#### D.1.3 Ablation Study: The Role of the Attention Layer.

To validate the role of the attention layer in capturing states in CT, we conducted tests, and the experimental results are presented in Figure 14 and 13. We found that the attention layer plays a crucial role in state capture. For instance, in the case of an increase in the SOFA score, without the attention layer, this increase cannot be captured, while with the attention layer, it clearly captures the change. Thus, this indicates that SOFA, as a key diagnostic indicator of sepsis, with the help of the attention layer, CT can accurately capture its changes.

Figure: Figure 14: The performance contrast between CT with and without an attention layer. The blue line represents the absence of an attention layer, while the green line indicates the presence of an attention layer.
Refer to caption: x13.png

### D.2 The Evaluation of Cost function in mechanical ventilator

#### D.2.1 Can Offline CT Improve the Performance of CRL?

Baselines. We adopt the Double Deep Q-Learning (DDQN) and Conservative Q-Learning (CQL) methods as baselines in ventilator research (Kondrup et al., ).

Corresponding experiments are conducted on the mechanical ventilator, as shown in Figure 15. Compared to previous methods DDQN and CQL, under the CDT+CT approach, a noticeable trend is observed where the proportion of mortality rates increases with increasing differences. When there is a significant difference in DIFF, the results may be unreliable, possibly due to the limited data distribution in the tail.

Figure: Figure 15: The relationship between the DIFF of actions and mortality in mechanical ventilator. The actions mainly consist of Positive End Expiratory Pressure (PEEP) and Fraction of Inspired Oxygen (FiO2), which are crucial parameters in ventilator settings.
Refer to caption: x14.png

#### D.2.2 Off-policy evaluation in mechanical ventilator

Baselines. 1) Naive baselines. A naive baseline can provide worst-case scenario benchmarks for algorithm evaluation (Luo et al., ), including random policy $\pi_{r}$, zero-drug policy $\pi_{\text{min}}$
, max-drug policy $\pi_{\text{max}}$, alternating policy $\pi_{\text{alt}}$ and weight policy $\pi_{\text{weight}}$.
2) RL methods baselines. We select common RL methods such as Deep Q-Network (DQN), Conservative Q-Learning (CQL), Implicit Q-Learning (IQL), and Batch Constrained Q-Learning (BCQ) as baseline models.

Metrics. A recent series of studies have applied offline policy evaluation techniques to dynamic treatment regimes, including Weighted Importance Sampling (WIS) (Kidambi et al., ; Nambiar et al., ) and Doubly Robust (DR) estimators (Raghu et al., ; Wu et al., ; Wang et al., ). To more accurately evaluate the policy, we use metrics such as RMSE and F1 score to describe the deviation from the clinician’s policy.

We used the same reward function to compare the policy results under different evaluation metrics in mechanical ventilators, as shown in Table 8. Our findings present that the CDT+CT method outperforms other methods in terms of $\text{RMSE}_{\text{PEEP}}$, WIS, $\text{WIS}_{b}$, and $\text{WIS}_{\text{bt}}$ evaluation metrics.

**Table 8: Comparison across policies on the mechanical ventilator test set. The best algorithms are highlighted in red. $\text{RMSE}_{\text{PEEP}}$ and $\text{RMSE}_{\text{FiO2}}$ mean the RMSE loss for the PEEP and FiO2. P.F1 and S.F1 denote the patient-wise F1 and sample-wise F1.**
| Metric | alt | max | min | random | weight | DQN | CQL | IQL | BCQ | CDT+CT |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| $\text{RMSE}_{\text{PEEP}}\downarrow$ | $8.15$ | $8.96$ | $7.30$ | $6.01$ | $7.16$ | $6.15\pm 0.48$ | $5.51\pm 0.06$ | $5.51\pm 0.03$ | $6.15\pm 0.04$ | $2.88\pm 0.04$ |
| $\text{RMSE}_{\text{FiO2}}\downarrow$ | $21.80$ | $14.09$ | $27.28$ | $18.56$ | $26.34$ | $15.81\pm 1.36$ | $13.08\pm 0.17$ | $13.72\pm 0.18$ | $16.69\pm 0.16$ | $13.13\pm 0.15$ |
| WIS $\uparrow$ | $0.66$ | $1.01$ | $0.66$ | $0.66$ | $0.84$ | $-1.13$ | $0.84\pm 0.07$ | $0.66\pm 0.29$ | $0.81\pm 0.09$ | $0.86\pm 0.07$ |
| $\text{WIS}_{b}$ $\uparrow$ | $0.16$ | $0.7$ | $0.14$ | $0.7$ | $0.83$ | $-0.81\pm 0.13$ | $0.78\pm 0.07$ | $0.54\pm 0.23$ | $0.78\pm 0.06$ | $0.83\pm 0.10$ |
| $\text{WIS}_{t}$ $\uparrow$ | $0.66$ | $1.01$ | $0.66$ | $0.66$ | $0.84$ | $-1.13$ | $0.84\pm 0.07$ | $0.66\pm 0.29$ | $0.81\pm 0.09$ | $0.86\pm 0.07$ |
| $\text{WIS}_{bt}$ $\uparrow$ | $0.11$ | $0.73$ | $-0.02$ | $0.72$ | $0.83$ | $-0.81\pm 0.15$ | $0.58\pm 0.14$ | $0.51\pm 0.24$ | $0.77\pm 0.04$ | $0.84\pm 0.08$ |
| DR $\uparrow$ | $-0.14$ | $-0.1$ | $-0.1$ | $-0.47$ | $-0.02$ | $-0.06\pm 0.02$ | $-0.80\pm 0.04$ | $-1.30\pm 0.03$ | $-0.69\pm 0.05$ | $-0.15\pm 0.02$ |
| P.F1 $\uparrow$ | $0.01$ | $0.01$ | $0.01$ | $0.01$ | $0.0$ | $0.01$ | $0.24$ | $0.28$ | $0.18$ | $0.03$ |
| S.F1 $\uparrow$ | $0.01$ | $0.01$ | $0.01$ | $0.01$ | $0.0$ | $0.02$ | $0.25$ | $0.25$ | $0.21$ | $0.04$ |

### D.3 The evaluation of Different Sequence models

To further analyze the performance of different sequence models, we conduct offline policy evaluation on models based on LSTM and transformer architectures. In sepsis and mechanical ventilator environments, the transformer-based models outperform LSTM-based models in a greater number of evaluation metrics, as shown in Table 9.

**Table 9: LSTM vs Attention. The best algorithms are highlighted in red.**
| Metric | Sepsis | Mechanical ventilator |  |  |
| --- | --- | --- | --- | --- |
| CT(LSTM) | CT(Attention) | CT(LSTM) | CT(Attention) |  |
| $\text{RMSE}_{\text{action1}}\downarrow$ | $505.06\pm 12.27$ | $433.55\pm 7.20$ | $8.74$ | $2.88\pm 0.04$ |
| $\text{RMSE}_{\text{action2}}\downarrow$ | $1.57$ | $1.13\pm 0.01$ | $19.05$ | $13.13\pm 0.15$ |
| WIS $\uparrow$ | $-3.03\pm 0.31$ | $-3.51\pm 0.11$ | $-1.05$ | $0.86\pm 0.07$ |
| $\text{WIS}_{b}\uparrow$ | $-3.55\pm 0.38$ | $-3.52\pm 0.17$ | $-0.38\pm 0.04$ | $0.83\pm 0.10$ |
| $\text{WIS}_{t}\uparrow$ | $-3.03\pm 0.31$ | $-3.51\pm 0.11$ | $-1.05$ | $0.86\pm 0.07$ |
| $\text{WIS}_{bt}\uparrow$ | $-3.64\pm 0.48$ | $-3.52\pm 0.17$ | $-0.30\pm 0.08$ | $0.84\pm 0.08$ |
| DR $\uparrow$ | $-3.05\pm 0.38$ | $-3.08$ | $-0.13$ | $-0.15\pm 0.02$ |
| P.F1 $\uparrow$ | $0.10\pm 0.11$ | $0.17\pm 0.02$ | $0.05$ | $0.03$ |
| S.F1 $\uparrow$ | $0.09\pm 0.10$ | $0.16\pm 0.02$ | $0.04$ | $0.04$ |

## Appendix E Online testing methods

Currently, some studies (Yoon et al., ; Luo et al., ; Brophy et al., ) have proposed simulation modeling approaches to address the challenges of directly testing RL-based dynamic treatment regimes in clinical environments. However, since the existing online testing systems (such as DTR-Bench (Luo et al., )) do not provide expert data in their simulation environment, and the offline method proposed in this paper requires expert datasets to train a safe policy, we are unable to use online testing systems for evaluation. In the future, we can establish an offline testing system to enable the testing of offline reinforcement learning strategies.

## Appendix F Experimental Settings

To train the CRL+CT model, we use a total of 3 NVIDIA GeForce RTX 3090 GPUs, each with 24GB of memory. Training a CRL+CT model typically takes 5-6 hours. We employ 5 random seeds for validation. We use the Adam optimization algorithm to optimize all our networks, updating the learning rate using a decay factor parameterization at each iteration. The main hyperparameters are summarized in Table 10 and 11.

**Table 10: List of the utilized hyperparameters in CT.**
| Offline CT Parameters | values |
| --- | --- |
| Genetivate Model |  |
| Embedding_dim | 128 |
| Layer | 3 |
| Head | 8 |
| Learning rate | 1e-4 |
| Pre-train steps | 5000 |
| Batch size | 256 |
| CT |  |
| Embedding_dim | 64 |
| Layer | 3 |
| Head | 1 |
| Learning rate | 1e-6 |
| Update steps | 30000 |
| Batch size | 512 |
| CDT |  |
| Learning rate | 1e-4 |
| Embedding_dim | 128 |
| Layers | 3 |
| Heads | 8 |
| Update steps | 60000 |

**Table 11: List of the utilized hyperparameters in CRL.**
| Parameters | Sepsis | Parameters | Mechanical Ventilation |
| --- | --- | --- | --- |
| General |  | General |  |
| Expert data patient number | 14313 | Expert data patient number | 13846 |
| Validation data patient number | 6275 | Validation data patient number | 5954 |
| Max Length | 10 | Max Length | 10 |
| Action_dim | 2 | Action_dim | 2 |
| State_dim | 48 | State_dim | 36 |
| Gamma | 0.99 | Gamma | 0.99 |
| DDPG |  | DDQN |  |
| Learning rate | 1e-3 | Learning rate | 1e-4 |
| Policy Network | 256,256 | Policy Network | 64,64 |
| Replay memory size | 20000 | Update steps | 500000 |
| Update steps | 20000 |  |  |
| VOCE |  | CQL |  |
| Learning rate | 1e-3 | Learning rate | 1e-4 |
| Policy Network | 256,256 | Policy Network | 64,64 |
| Alpha scale | 10 | Update steps | 500000 |
| KL constraint | 0.01 | Alphas | 0.05,0.1,0.5,1,2 |
| Dual constraint | 0.1 |  |  |
| Update steps | 4000 |  |  |
| CopiDICE |  |  |  |
| Learning rate | 1e-4 |  |  |
| Policy Network | 256,256 |  |  |
| Alpha | 0.5 |  |  |
| Cost limit | 10 |  |  |
| Update steps | 100000 |  |  |
| BCQ-Lag |  |  |  |
| Learning rate | 1e-3 |  |  |
| Policy Network | 256,256 |  |  |
| Cost limit | 10 |  |  |
| Lambda | 0.75 |  |  |
| Beta | 0.5 |  |  |
| Update steps | 100000 |  |  |

## Peer Reviews

confidence: 2
contribution: 2
initial_rating: 6
presentation: 3
questions: Please see weaknesses.
soundness: 3
strengths: - The work is well motivated and investigates an important problem, where online test or interaction is unavailable and risky for ICRL.
- The paper is well written
- The proposed method achieved good performance on MIMIC-III derived dataset which is more realistic
- Limitation, especially that expert demonstrations are assumed as optimal, is well discussed.
summary: The paper proposed a constraint transformer method to tackle the problem that only historical data is provided in offline inverse reinforcement learning. The proposed method exhibits effectiveness in a sepsis task.
weaknesses: - I’m curious any justification for only investigating one dataset for a specific task, sepsis treatment, though I agree sepsis is an important task and MIMIC-III has been broadly investigated. But given the proposed method is proposed targeting a general problem in healthcare, either more strictly scope the claims and investigated settings, or providing broader investigation could be helpful for further demonstrating the effectiveness of the proposed method.

---

confidence: 3
contribution: 3
initial_rating: 3
presentation: 3
questions: - Table 1: the definition of "too high" is somewhat problematic. Claiming "too high" is associated with increased mortality making it an "unsafe behavior" without conditioning on the patient state seems inappropriate. Even though the footnote says this should be patient-dependent, the table still shows a single threshold. 
- Sec 3 problem formulation defines a non-Markov decision process (non-MDP) but the remaining text all references an MDP (L173-L175).  
- L173: What does "the MDP derived by augmenting the original MDP with the network" mean exactly? Why do we augment the MDP with a network? Or do you mean augmented with the constraint? 
- L176-L206 evaluation metric defined using "DIFF": this graphical analysis known as the "U-curve" is deprecated and widely criticized because it cannot distinguish a good policy from a bad policy (even a random one). See Figure 4 in Gottesman et al. "Evaluating Reinforcement Learning Algorithms in Observational Health Settings" (https://arxiv.org/abs/1805.12298) 
  - also "We randomly select 2N patients from the dataset where N patients survived under the doctor’s treatment, and N patients died" do not seem justified. Depending on N this may not be representative of the dataset. Why not use some form of weighted average in the ranking? 
- L244: what does "creating an offline RL for violating data" mean? 
- L190: "probability of approaching the optimal policy" L438 "use ω to indicate the probability that the policy is optimal" - in offline setting, we can never obtain the optimal policy, so I don't believe this terminology is appropriate.
soundness: 3
strengths: This paper studies an interesting and potentially important problem.
summary: This paper studies inverse RL with learnable constraints in the offline setting, focusing on practical applications in healthcare tasks. The main approach appears to be combining the decision transformer architecture in the offline RL literature with inverse constrained RL with max entropy framework. Experiments were conducted on two healthcare tasks: sepsis and mechanical ventilation.
weaknesses: There are some inconsistencies throughout the text in terms of mathematical descriptions, and the experimental evaluation also suffers from some issues.

---

confidence: 5
contribution: 2
initial_rating: 3
presentation: 2
questions: 1) What are your assumptions in this work and why are they plausible?
2) How does the distribution of clinician actions impact/bias the work? If the clinician decisions are all centred around administering very few drugs does that not overconstrain the system or end up with you not learning any meaningful constraints at all? How do you deal with uneven clinician action distributions?
3) What is the purpose of each component in the story and what role does it play? Why dont you compare the representations you learn with other representation learning approaches.
4) What makes the learnt representation causal? Is it robust against various biases? Do you have clinical validation to support what you have learnt?
5) What role does the Safe RL literature play here. How do you compare to these? 
6) Do you have another domain in which you can validate and support your findings?
soundness: 2
strengths: Overall I think the paper tries to address a very relevant and significant problem that is not only limited to healthcare. The paper is written well in general and I like the way the experiments are laid out in the sense that the authors clearly divide their analysis into questions that tackle various aspects of the evaluation. The authors also do a good job in motivating the problem overall.
summary: The paper presents an offline approach for inverse RL in constrained settings for safety critical decision making. The overall premise of the paper is the traditional RL methods might overlook common-sense constraints which might lead to potentially harmful decisions in a high-risk setting. Moreover, such healthcare decisions must usually be inferred from offline batch data, where there is no online/limited interaction with the environment and the exact constraints might be unknown as is typically the case in healthcare. The authors introduce a two step approach to overcome these issues: i) first they employ an attention-based architecture to identify long-range dependencies and critical aspects of a patient's history; ii) based on the learnt representation from i), the authors use a generative model of the environment for data augmentation that enables unrolling/simulating potentially unsafe decision sequences. The authors try demonstrate on MIMIC III that their approach recovers the relevant constraints accurately, is performant in offline settings and enables safe policy learning.
weaknesses: 1) The authors dont make explicit what their assumptions are and I had to do a lot of sifting through details multiple times, to get a clearer picture of the exact problem setting: it is assumed that there is a reward and a set of constraints (costs) within the MDP; the reward is observable but the constraints are unknown which is why an IRL approach is warranted. 

2) The authors also assume that the expert is optimal which may not be the case in practice so the authors should make this explicit in an assumptions section. The authors should in particular make very explicit what the distribution of clinician actions is for their applications and allow us to visualise this. For instance, if the most common clinician action choice is to provide very small doses of vasopressors and fluids, how does this impact the quality of the constraints learnt? I would think that if the distribution of clinician actions is strongly biased towards certain decisions, this might have significant impact on the constraints learnt/might end up overconstraining the problem or end up not learning anything meaningful.

3) My main criticism for this paper is I feel the method contains a LOT of components and it is not always clear what the purpose of each component is. While the purpose of the transformer is to retain relevant aspects of the patient history and capture critical states, I would have liked to have seen some analysis into what is learnt/whether these points are in fact critical, why they are important etc. There is a huge body of work on learning relevant representations e.g. through decision/task-focused learning (Sharma et al 2024 https://arxiv.org/pdf/2110.13221) or otherwise through information-based methods (Liu et al 2024, https://arxiv.org/pdf/2405.09308) or similar, but the authors dont compare the representations they learn to any other method which makes it hard to understand whether what is captured is indeed relevant. There is also a large amount of work related to capturing decision-regions or important points in a patient trajectory that might impact doctor decision making in terms of how much they agree on what to do (Zhang et al 2022 https://www.nature.com/articles/s41746-022-00708-4) which should be discussed here. If nothing else, the authors should compare what is learnt to a standard LSTM representation. Subsequently, a generative model is used to unroll or simulate downstream decision sequences and perform safe decision-making under constraints. There is little discussion on how the simulator behaves as a causal model and serves as an adequate means of simulating possible trajectories. I will elaborate on this next.

4) The authors also mention that the generative model used to simulate trajectories is causal in the abstract and elsewhere. How can this be verified? What makes the representation causal? Is it robust against confounding biases? Are there experiments that validate this where you demonstrate explicitly that the effects of such biases are minimised and do not have downstream consequences on the constraints learnt and the simulations produced? Did you validate the simulations with domain expertise to check whether these were indeed valid potential trajectories/simulations of how a patient might evolve?

5) Throughout the paper, there is extensive mention of the terms Safe RL yet the entire literature on safe RL is ignored/not discussed meaningfully in terms of related work or in terms of establishing baselines for the approach to compare to (see Garcia et al 2015 https://www.jmlr.org/papers/volume16/garcia15a/garcia15a.pdf as a starting point). I understand the difference is now you are learning C but one could think of C as being part of a multiobjective reward R in which case all the works on multi-objective RL need to be discussed in a lot more detail in the related work. Given that there are various notions of what it means to be safe e.g. risk sensitive, risk-directed/aware etc. Where does this approach sit? 

6) The authors dont provide any other experiments apart from MIMIC III. It would be good to have either some toy demonstrations to validate the utility of each component of the method and why it serves a purpose or to have another clinical data set with another clinical task where you demonstrate the performance of the approach (hopefully where the action distribution differs significantly)

---

confidence: 5
contribution: 3
initial_rating: 5
presentation: 3
questions: NA
soundness: 3
strengths: This paper shows its strengths in the following aspects:

The paper addresses the novel angle of ensuring safety in offline reinforcement learning (RL) for healthcare, a critical and previously underexplored issue.

It incorporates Inverse Constrained Reinforcement Learning (ICRL) into offline reinforcement learning (RL) for healthcare, introducing a novel approach to inferring constraints from expert demonstrations in a non-interactive environment.

The implementation of a causal transformer to learn the constraint function is interesting, allowing the integration of historical patient data and capturing critical states more effectively.

Extensive results on 2 well-known dynamic treatment regime datasets are presented. The proposed Constraint Transformer (CT) framework is shown to reduce unsafe behaviours and approximates lower mortality rates.
summary: This paper introduces the Constraint Transformer (CT) framework to enhance safe decision-making in healthcare. The proposed CT model uses transformers to incorporate historical patient data into constraint modelling and employs a generative world model to create exploratory data for offline RL training. The authors supported their points by presenting experimental results in scenarios like sepsis treatment, showing that CT effectively reduces unsafe behaviours and approximates lower mortality rates, outperforming existing methods in both safety and interoperability.
weaknesses: I think the paper is generally solid and I vote for its acceptance. However, there are quite a lot of presentation issues and technical details to be addressed.

1. result

The RMSE_VASO of the proposed method is greatly larger than that of other policies. Also, it seems that CDT+CT outperforms other baselines on the importance sampling-based method but not on DR. Could you explain why? The authors said, "Since CDT+CT produces safer and more conservative policies, there is a certain distance from the clinician’s policy, so it does not perform as well in terms of RMSEvaso and F1 score." From my understanding, more safe and conservative policies should result in a smaller distance from the clinician's policy. Is this statement inconsistent with the result?

2. presentation issue

This paper has multiple presentation issues. Due to the space limit, I will only point out a few of them. *a)* Figure 1 is quite confusing. If an agent learns stochastic policy, its action choice can still be distributional. The authors misunderstood some fundamental concepts of MDPs in RL. *b)* "The Markov decision is not compatible with medical decisions." is technically wrong. The Markov decision **IS** compatible with medical decisions if you define the MDP correctly. To do RL, we must use the MDP assumption in a certain way. If I understand the authors correctly, you meant to say 'strictly first-order Markov' is not a fit for medical decisions. But why bother mention it in the introduction? Wouldn't a RNN solve this problem? *c)* Line 63-77: these 2 challenges are not in ICRL, but in applying ICRL in healthcare.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
questions: See above.
- How does the evaluation procedure described in L190 relate to off-policy evaluation methods / propensity-weighting schemes?
- Fig 1: what is a state?
soundness: 3
strengths: - The paper addresses the interesting problem of learning safe decision-making policies from observational data, which is an important challenge for applying reinforcement learning to healthcare applications.
- Considering the inherent challenge of evaluating such methods, authors consider multiple metrics to assess the performance of policies trained with their method / different baselines.
summary: This paper proposes an Offline Inverse Constrained RL method (ICRL) for medical applications, which aims to (1) learn decision-making constraints from observational data using a history-dependent model architecture (‘Constraint Transformer’, CT) as well as simulated unsafe decision-making trajectories, and (2) using these within a constrained offline RL optimization procedure.
The goal is to avoid unsafe behavior in the output policy.
weaknesses: 1. The method proposed (and in particular its complexity, as it involves a range of different modeling steps) needs to be better motivated. Below are some example modeling choices that would require a more in-depth explanation + detailed experimental analysis:
- Is the generative model in L302 not biased by partial observability? As the policy changes from expert to violating policy, does this not induce confounding in the generative model? 
- The fact that medical decision-making problems are history-dependent or POMDPs is not new (Yu et al., 2020, Reinforcement Learning in Healthcare: A Survey), yet it is presented as a major contribution. 
- Authors find that the cost function correlates with the mortality rate. Isn’t this expected from the fact that doctors’ policy aims to avoid mortality? Is the cost function truly learning something distinct from the reward function, or is it just providing additional reward shaping based on imitation learning? Table 4 would benefit from additional results comparing to TD3+BC (RL and imitation learning regularization) and to CDT without cost or with custom cost.

2. Experiment results do not demonstrate a clear added value for the proposed method, considering its added complexity.
- Evaluation: constraints are inferred and optimized, but we have no ground-truth knowledge of what they should look like. I believe this could be evaluated in a toy experiment.
- Overlapping results: (1) Table 2: CDT with and without CT cost overlap. What is the added value of CT, considering the significant additional modeling complexity required? (2) Fig 7: Authors claim a difference at DIFF=0 between their methods and the baseline, but both curves overlap in confidence interval.
- How does CDT without cost work -- is it simply a decision transformer? If so, why does it perform so much better than DDPG?
- L504 Why is CDT+CT producing a policy that is more “distant from the clinician’s policy” than other offline RL baselines? Are we not specifically learning constraints based on what actions are *different* from the behavioral policy -- ie doing a form of regularization to imitation learning?


3. Clarity is poor overall, and this observation relates to my multiple questions above. Some additional comments:
- Is $\zeta_{\theta}$ dependent on $\tau$ or (s,a)?
- Vertical white space is reduced excessively throughout the paper, which makes for an unpleasant reading experience.
- L321 where does a_t come from, is it sampled from the dataset D_e? So is the objective simply imitation learning? Or is there a missing reward signal to make it an RL objective?
- what is the strategy for bolding results in Table 2, when they clearly overlap with others? Same comment for red numbers in Table 4.
- Table 2: How does the ‘generative model’ output a policy?
- Table 3: how can proportions for CDT + CT be 0% if at least one value exceeds the thresholds considered? Is 0% an approximation?
- Fig 7: How is mortality estimated here?

4. Additional details
- L262: Transformers are not ‘interpretable’ per se beyond the first layer.

Meta-review: This paper proposes a transformer based model to overcome some challenges with safety in healthcare RL by using the decision-transformer framework combined with a reward design (IRL) setup where markovian assumptions may be violated. Empirical evaluation compares results in  sepsis (mechanical vent) management in MIMIC. 

Reviewers have raised several concerns regarding, lack of clarity in claims and overall writing, insufficient empirical evaluation with experiments only demonstrating MIMIC-III, lack of comparison with other safe RL frameworks, including the claim of benefits of transformer based representation learning but lacking rigorous evaluation of other representation learning settings in healthcare RL, as well as choice and comprehensiveness of the evaluation measures. 

Authors have clarified several concerns while also updating the draft. However, reviewers have lingering concerns that are not sufficiently addressed in the updated draft or the author response, therefore most reviewers have not updated their score. I therefore recommend a reject in broad consensus with reviewer concerns that there are still some claims that need to be well justified through writing, and rigorous empirical evaluation.

Decision: Reject