# Pruning the Way to Reliable Policies:
A Multi-Objective Deep Q-Learning Approach to Critical Care

Ali Shirali Equal contribution University of California, Berkeley  Alexander
Schubert∗ University of California, Berkeley  Ahmed Alaa  University of
California, Berkeley

###### Abstract

Most medical treatment decisions are sequential in nature. Hence, there is
substantial hope that reinforcement learning may make it possible to formulate
precise data-driven treatment plans. However, a key challenge for most
applications in this field is the sparse nature of primarily mortality-based
reward functions, leading to decreased stability of offline estimates. In this
work, we introduce a deep Q-learning approach able to obtain more reliable
critical care policies. This method integrates relevant but noisy intermediate
biomarker signals into the reward specification, without compromising the
optimization of the main outcome of interest (e.g. patient survival). We
achieve this by first pruning the action set based on all available rewards,
and second training a final model based on the sparse main reward but with a
restricted action set. By disentangling accurate and approximated rewards
through action pruning, potential distortions of the main objective are
minimized, all while enabling the extraction of valuable information from
intermediate signals that can guide the learning process. We evaluate our
method in both off-policy and offline settings using simulated environments
and real health records of patients in intensive care units. Our empirical
results indicate that pruning significantly reduces the size of the action
space while staying mostly consistent with the actions taken by physicians,
outperforming the current state-of-the-art offline reinforcement learning
method conservative Q-learning. Our work is a step towards developing reliable
policies by effectively harnessing the wealth of available information in
data-intensive critical care environments.

##  1 Introduction

Specifying a reward function poses a challenge in many (applied) reinforcement
learning (RL) settings. In the presence of multiple desirable factors, it is
often unclear how to obtain a single scalar reward. Additionally, even when a
single outcome is of interest, it may be too rare to learn from directly with
limited data. To address this issue, less accurate but more frequent reward
proxies can be incorporated. However, if these proxies are included in the
reward specification without sufficient consideration, the resulting policy
may deviate from optimizing the true reward.

This challenge frequently arises in the medical context where RL’s promise is
to obtain personalized treatment policies. Many medical decisions require
physicians to periodically select among multiple treatments or to prescribe
the right diagnostic tests at the right time. The sequential nature of these
problems and the increasing availability of health data make it a potential
use case of RL. In many settings, patient survival is a critical outcome of
interest, but it is a rare event that only provides a delayed and sparse
signal. As an alternative, medical risk scores can be used as more frequent
intermediate rewards. However, these scores are only approximations of a
patient’s true individualized severity and may be biased for certain
populations [26]. Hence, determining the most effective way to integrate
information from these intermediate signals into RL policies still remains a
challenge.

In recent years, there has been a surge of interest in exploring the prospects
of RL in critical care settings [12, 31, 32, 33, 4, 17, 20, 29, 7]. Different
studies have taken different approaches towards defining the reward function.
For instance, focusing on the task of managing vasopressor and intravenous
fluids in sepsis patients, [17] solely reward the agent based on $90$-day
survival, while [32] have further included intermediate rewards based on the
SOFA-score [42] and arterial lactate levels.

As mentioned earlier, directly incorporating more frequent but less accurate
intermediate reward signals comes at a cost. Distortions from these rewards
can cause the final policy to deviate from the primary outcome of interest.
Furthermore, the varying complexity in the structures of different rewards can
pose unique learning challenges for an RL agent. As a result, there is a risk
that the agent could primarily optimize toward rewards that are easier to
learn, inadvertently neglecting the genuine, more complex outcomes of
interest.

In order to overcome this challenge, we propose a new two-stage algorithm that
partially disentangles learning from inaccurate intermediate rewards and the
sparse reward of interest. In the first stage, we propose a multi-objective
Q-learning algorithm that enables us to _prune_ the action set based on all
available rewards without committing to an explicit relationship between them.
Pruning is a relaxed version of a policy, that instead of prescribing a single
action to take, removes unpromising actions which are unlikely to make any
improvement to any linearly scalarized reward. This relaxation will minimize
the possible distortion due to less accurate rewards while guiding the second
stage where we run standard Q-learning on the pruned action set.

We evaluate our framework in both off-policy and offline settings using
simulated environments and real data. In the off-policy setting of OpenAI Gym
Lunar Lander environment [2], we show that our algorithm can effectively prune
the action space and that this reduction in complexity significantly improves
learning based on the sparse reward later on.

We further evaluate our framework in an offline setting using real health
records of septic patients in the intensive care unit (ICU). We demonstrate
that simply linearly combining rewards does not effectively improve standard
Q-learning. However, by using our two-stage algorithm, medical scores can
effectively enhance the policy outcome. Additionally, we find that the pruning
step significantly reduces the number of available actions without discarding
the actions chosen by physicians in the majority of cases.

A desirable characteristic of a reliable offline RL policy is to adhere to the
expert policy, in this case, the physician policy, as closely as possible
while incurring minimal performance losses compared to an unrestricted policy.
Conservative Q-learning (CQL) [18] is a prevalent method used to enforce such
similarity. Our results demonstrate that our method achieves a high similarity
to physician actions comparable with CQL while attaining superior performance.
This indicates that our approach is promising in obtaining relevant offline
policies while avoiding out-of-distribution actions.

Reinforcement learning holds the promise of deriving personalized treatment
policies in healthcare that cater to the unique history of each patient.
However, achieving this requires models that can effectively assimilate
information from diverse inputs, while also being able to accurately
distinguish between key outcomes and potentially biased or less accurate
proxies. Our work presents a new approach that takes us closer to achieving
this goal. Our two-stage algorithm allows us to effectively integrate
information from intermediate severity proxies, resulting in better
performance compared to standard conservative Q-learning methods. This is
particularly important in critical care processes, where intermediate
indicators are abundant, and it is crucial that computational methods can
accurately respond to these signals when providing recommendations.

The code to derive the presented results is available at:

https://github.com/alishiraliGit/multi-criteria-deep-q-learning

##  2 Related Work

##### Multi-Objective Reinforcement Learning.

The main body of work on multi-objective (or multi-criteria) reinforcement
learning (RL) can be divided into single-policy methods and multiple-policy
methods. We refer the reader to [21] for a survey of the field. Single-policy
methods [24, 39] reduce multiple objectives into a single scalar reward
assuming a known user-specified or context-driven preference over different
objectives and then seek to maximize the scalarized objective using standard
RL techniques. The major difference between different single-policy approaches
lies in the way in which they attempt to derive preferences over objectives.
These methods, however, cannot be used when the preferences themselves are
unknown.

In contrast, multi-policy methods aim to estimate the Pareto frontier based on
a set of policies. One way to achieve this is by training an ensemble of
policies based on different reward scalarizations [28, 41, 27, 3, 43]. These
methods require exhaustive training and in some cases non-trivial
scalarizations. Other methods extend the standard scalar variables in RL
algorithms to vector variables and use updating rules in the vector space.
Most related to our proposed method of multi-objective Q-learning (Section
4.1) are such value-based methods. Early work in this direction explored the
problem of acquiring all Pareto optimal policies simultaneously [1, 13, 14].
Their methods focused on applications in online settings and on small, finite-
state spaces. [23] extended [1]’s framework to real-valued state features but
with an exponential complexity in time horizon and in the size of the action
space. In contrast, our deep Q-learning method runs for arbitrarily long
trajectories in polynomial time by using an approximate problem formulation.
Furthermore, our familiar formulation enables our method to leverage the
recent advances in offline Q-learning.

Overall, a key focus of this research area lies in either inferring optimal
preference weights or directly uncovering a set of Pareto-optimal policies.
Instead, to the best of our knowledge, this study is the first to leverage
multi-objective RL to prune the action space by identifying and eliminating
inferior state-actions pairs in order to reduce the complexity of the learning
problem for subsequent learning tasks.

##### Non-Deterministic Policies.

While the mathematical framework we employ to devise a pruning function in
phase 1 of our algorithm echoes certain aspects of the research on non-
deterministic policies (or set-valued policies) [6, 36], our methodology
diverges in various ways. Firstly, non-deterministic policies sacrifice
performance to provide a set of potential solutions instead of a single
distinct action. In contrast, our method intends to improve policy
performance. This is achieved as it combines a set-valued policy formulation
to prune the action space in phase 1, allowing for more effective policies in
phase 2. While non-deterministic policies remain prescriptive, this way our
pruning function largely acts as a prohibitive measure that eliminates
suboptimal choices from the action set.

##### Reinforcement Learning in Health.

In recent years, RL has been the focus of several studies in healthcare. See
[22] for a review. There has been a particular focus on personalized treatment
plans for sepsis patients in the ICU [12, 9, 32, 17, 34, 30, 37, 33]. However,
these studies either leverage sparse mortality information alone to assign
rewards or combine multiple intermediate sepsis risk proxies based on a
deterministic weight. To the best of our knowledge, our approach is the first
attempt to explicitly address the challenges imposed by sparse rewards in this
context by proposing to prune the action space before estimating a final
policy.

##### Dead-End State Identification.

Given the focus of the pruning stage of our algorithm on eliminating dominated
actions, our method is related to work on dead-end identification in RL [15,
8, 7]. For instance, [7] developed an RL algorithm to identify and avoid
“medical dead-end” states, defined as states from which no action can be made
to achieve a positive terminal outcome (e.g. survival). While the aim of this
method is conceptually close to the goal of our pruning stage, we propose a
distinct technical approach to achieve this goal. [7]’s work focuses on
classifying states that should be avoided; in contrast, our method directly
identifies and excludes dominated state-action pairs. Furthermore, our method
does not stop at the identification of risky actions but instead leverages the
pruned action space for the subsequent training of an optimal policy.

##  3 Notation and Preliminaries

##### Reinforcement Learning.

Following the terminology commonly adopted in the area of RL, we imagine an
agent interacting with an environment. As the agent executes an _action_ ,
denoted as $a\in{\mathcal{A}}$, in a given _state_ , denoted as
$s\in{\mathcal{S}}$, a _reward_ , expressed as $r(s,a)$, materializes, and the
environment’s state subsequently updates to $s^{\prime}$. The agent can then
utilize this reward as feedback to make more optimal future action choices.
The agent’s decision-making process is often elucidated through a _policy_. In
general, a policy $\pi:{\mathcal{S}}\rightarrow\Delta({\mathcal{A}})$ provides
a distribution over potential actions for each state. A deterministic policy
$\pi:{\mathcal{S}}\rightarrow{\mathcal{A}}$ singles out a specific action
given the current state. Our study particularly focuses on _offline_ RL, where
the aim is to devise an effective policy based on previously collected data.
Simultaneously, our algorithms are also applicable in _off-policy_ settings
where the agent interacts with the environment at selected time steps. We
denote the available data as a set of transitions ${\mathcal{D}}$. Each
transition is a tuple $(s,a,s^{\prime},r)$ demonstrating that at state $s$,
the action $a$ is taken, which resulted in a state transition to $s^{\prime}$
and a reward of $r$.

##### Q-Learning.

A standard (scalar) _Q-function_
$Q^{\pi}:{\mathcal{S}}\times{\mathcal{A}}\rightarrow\mathbb{R}$ estimates the
value of policy $\pi$ at state $s$ given that the agent takes action $a$ at
$s$ and follows policy $\pi$ thereafter. This value signifies a discounted sum
of rewards, with any rewards realized at a later timestep being discounted by
a factor of $\gamma$. At times, we omit the superscript $\pi$ as our focus is
on the optimal policy. The deep learning structure that implements the
Q-function is occasionally referred to as a _Q-network_. We use the terms
Q-function and Q-network interchangeably. In the context of Q-Learning, we
sample a batch of transitions ${\mathcal{B}}\in{\mathcal{D}}$ and employ the
Bellman equation to update the Q-network:

$Q(s,a)\leftarrow r+\gamma\max_{a^{\prime}}Q(s^{\prime},a^{\prime}).$ (1)

Let the Q-network be parameterized by $\theta$. Then the above notation is a
shorthand to update $\theta$ from its current value $\theta_{0}$ by minimizing
the loss

${\mathcal{L}}(\theta)=\sum_{(s,a,s^{\prime},r)\in{\mathcal{B}}}\big{(}Q_{\theta}(s,a)-r-\gamma\max_{a^{\prime}}Q_{\theta_{0}}(s^{\prime},a^{\prime})\big{)}^{2}.$
(2)

Commonly, we implement this update via a single or a couple of gradient
descent steps. However, Q-learning based on the above formulation suffers from
a fast-moving target and overestimation. These challenges can be addressed by
employing a target network $Q^{\prime}$ along with an update rule, recognized
as double Q-learning [40]:

$Q(s,a)\leftarrow
r+\gamma\,Q^{\prime}\big{(}s^{\prime},\operatorname*{arg\,max}_{a^{\prime}}Q(s^{\prime},a^{\prime})\big{)}.$
(3)

In the context of double Q-learning, the target network $Q^{\prime}$ is
updated by $Q^{\prime}\leftarrow Q$ after hundreds or thousands of updates to
$Q$. The final deterministic policy is then extracted as
$\pi(s;Q)=\operatorname*{arg\,max}_{a}Q(s,a)$.

##### (Suboptimal) Q-Learning With Softmax Policies.

In the original Bellman equation, the optimal policy is greedy and
deterministic with respect to the Q-values, achieved through the use of the
$\operatorname*{arg\,max}$ function to locate the most effective action.
However, empirical evidence has shown that applying a softmax operator tends
to yield superior policies, particularly in the context of deep Q-networks.
This relaxation has also garnered recent theoretical support [35]. Formally,
given a Q-function $Q$, we can derive a stochastic softmax policy
$\pi^{\beta}$ as

$\pi^{\beta}(a|s;Q)=\frac{\exp\big{(}\beta
Q(s,a)\big{)}}{\sum_{\bar{a}}\exp\big{(}\beta Q(s,\bar{a})\big{)}},$ (4)

where $\beta$ is an inverse temperature parameter. Note that in the limit of
$\beta\rightarrow\infty$, the softmax policy converges to a deterministic
$\operatorname*{arg\,max}$ policy. Let us define

$\operatorname*{softmax}_{a^{\prime}}Q(s^{\prime},a^{\prime})\coloneqq\sum_{a^{\prime}}\pi^{\beta}(a^{\prime}|s^{\prime};Q)\,Q(s^{\prime},a^{\prime}).$
(5)

We dropped the dependence on $\pi^{\beta}$ from the softmax notation for
convenience. Note that
$\operatorname*{softmax}_{a^{\prime}}Q(s^{\prime},a^{\prime})\rightarrow\max_{a^{\prime}}Q(s^{\prime},a^{\prime})$
as $\beta\rightarrow\infty$. A Bellman update using softmax can then be
obtained by substituting the $\max$ with a $\operatorname*{softmax}$ operator
in Equation 1. Similarly, we can implement double Q-learning using a softmax
policy:

$Q(s,a)\leftarrow
r+\gamma\sum_{a^{\prime}}\pi^{\beta}(a^{\prime}|s^{\prime};Q)\,Q^{\prime}\big{(}s^{\prime},a^{\prime}\big{)}.$
(6)

##### Vector-Valued Reward and Q-Function.

In many practical scenarios, the reward from an interaction can have multiple
dimensions. For example, consider the case where there is a sparse main reward
$r_{*}$ as well as a set of noisy but more frequent rewards. In these cases,
we gather all the rewards into a _vector reward_ ${\bm{r}}\in\mathbb{R}^{d}$
and use $r_{i}$ to refer to the $i^{\textrm{th}}$ reward. We introduce a
_vector-valued Q-function_
${\bm{Q}}:{\mathcal{S}}\times{\mathcal{A}}\rightarrow\mathbb{R}^{d}$ which for
any state-action pair outputs a vector of estimated values. $Q_{i}$ denotes
the $i^{\textrm{th}}$ dimension of ${\bm{Q}}$ that encodes the value from
$r_{i}$. We discuss learning vector-valued Q-functions in Section 4.1.

##  4 Methods

The main objective of this study is to formulate an approach that facilitates
the efficient incorporation of frequent yet imprecise reward signals while
preventing such signals from causing the policy to diverge from maximizing the
primary sparse reward of interest. We hypothesize that while some biomarkers
or clinical risk scores might be too noisy to be directly used in deriving an
optimal policy they may still be useful to _identify actions to avoid_. Hence,
we propose a two-stage algorithm: In the first phase, we learn a vector-valued
Q-function relying on (noisy) intermediate rewards. This Q-function will be
the basis to _prune_ the action set at each state. Then in the second phase,
we search for the optimal policy based on the (accurate) sparse reward.
Actions dropped in the first phase won’t be available to the policy of the
second phase, which reduces the complexity of the learning problem and thus
facilitates learning from the sparse reward.

###  4.1 Phase 1: Multi-Objective Deep Q-Learning

In this section, we first review direct ways to incorporate noisy intermediate
reward signals into a reinforcement learning objective through explicit
scalarization and discuss why this is challenging in our setting. We then
motivate and propose a vector-valued Q-learning algorithm that avoids any
explicit scalarization. We later augment our method for the offline setting by
applying techniques from conservative Q-learning. We conclude the description
of phase 1 by developing a stochastic policy based on the trained vector-
valued Q-function.

##### Challenges of Single-Policy Methods.

A direct method to combine a sparse main reward and intermediate reward
signals is through scalarization of the reward vector ${\bm{r}}$. In the
multi-objective RL literature, this approach is generally referred to as
single-policy methods. In its simplest form, consider a linearly weighted
combination ${\bm{w}}^{\mathbb{T}}{\bm{r}}$ where ${\bm{w}}\in\mathbb{R}^{d}$
determines the relative importance of each reward aspect. We assume $w_{i}\geq
0$ and $\sum_{i=1}^{d}w_{i}=1$. Leveraging the combined reward as the standard
Q-learning objective, we can then obtain the Q-function $Q^{\bm{w}}$ and its
corresponding (possibly stochastic) policy
$\pi^{\beta}(\cdot|\cdot;Q^{\bm{w}})$. However, this approach comes with
multiple challenges: First, we might not have knowledge of the best ${\bm{w}}$
a priori. The selection of the weighting ${\bm{w}}$ entails a trade-off
between the smoothness of learning and the accuracy of approximating the main
reward. The stronger a frequent intermediate reward is represented in the
combined reward, the smoother the learning. But as the intermediate rewards
carry a high level of noise, the final policy might deviate from maximizing
$r_{*}$ in an unwanted way. Second, a good ${\bm{w}}$ to integrate the noisy
intermediate reward signals might not exist at all. In fact, to facilitate
learning, intermediate rewards should be present in the scalarized version
with a minimal power. At this point, the signal from intermediate signals
might be so strong that the Q-network mainly exploits the easy-to-learn
intermediate rewards. Our experiments in the context of treating sepsis in an
ICU setting suggest that such direct incorporation of intermediate reward
signals into the objective leads to a deterioration of policy performance
compared to a policy trained based on the sparse reward alone.

##### Challenges of Multi-Policy Methods With Ensemble of Policies.

One way to get around the challenges of committing to a single weighting
${\bm{w}}$ of rewards is to consider a set of weightings and obtain an
ensemble of policies. Although this method does not obtain a single policy, it
can still prune the action set at each state by removing the actions not
prescribed by any policy in the ensemble. The possibility of having multiple
weightings address our uncertainty about the best weighting but there remain
multiple challenges: First, as mentioned previously, there might not be a good
${\bm{w}}$ at all. In this case, policies in the ensemble are either unstable
if their corresponding ${\bm{w}}$ is heavily concentrated on the true sparse
reward $r_{*}$ or they are largely affected by noisy intermediate rewards.
Second, training a comprehensive ensemble of models for a high-dimensional
reward space is computationally expensive. To address these problems, we
follow [1, 23] and suggest a vector-valued Q-learning approach that
simultaneously learns optimal policies for several weightings of interest.

##### A Vector-Valued Q-Learning Approach.

Given the aforementioned challenges of incorporating intermediate rewards
through explicit scalarization, we present a Q-learning approach to obtain a
vector-valued Q-function ${\bm{Q}}$. We aim to obtain a vector-valued
Q-function such that
$Q^{\bm{w}}(s,a)\approx{\bm{w}}^{\mathbb{T}}{\bm{Q}}(s,a)$ for almost every
${\bm{w}}\in{\mathcal{W}}$. Such a Q-function addresses the previously
mentioned challenges as it does not commit to any explicit ${\bm{w}}$ when
incorporating intermediate rewards into learning and will subsequently be at
the heart of our action-pruning approach in phase 2. In the following, we
first show that a naive extension of the Bellman update to vector-valued
Q-networks does not allow us to derive a universal update rule. Hence, to
obtain a vector-valued Bellman update in the familiar form of Equation 1, we
relax the problem in two ways: First, instead of estimating exact Q-functions,
we conservatively approximate them. We show that our approximation may only
underestimate values from actions unlikely to be taken by any optimal policy,
and thus should have a minimal effect on performance. Second, we limit the
weightings of interest by considering a prior ${\mathcal{P}}$ over
${\mathcal{W}}$. A good prior should be both inclusive and spread-out. It
should assign enough weight to the true reward, ensuring its presence in the
pruning policy with a non-negligible probability, while avoiding concentration
that leads to the same problems as fixing a specific $w$.

First, to show that a naive extension of the Bellman update of $Q^{\bm{w}}$
cannot provide a universal updating rule for ${\bm{Q}}$, consider a fixed
${\bm{w}}$ and rewrite the Bellman update of $Q^{\bm{w}}$ assuming
$Q^{\bm{w}}(s,a)={\bm{w}}^{\mathbb{T}}{\bm{Q}}(s,a)$:

${\bm{w}}^{\mathbb{T}}{\bm{Q}}(s,a)\leftarrow{\bm{w}}^{\mathbb{T}}{\bm{r}}+\gamma\operatorname*{softmax}_{a^{\prime}}{\bm{w}}^{\mathbb{T}}{\bm{Q}}(s^{\prime},a^{\prime}).$
(7)

The above update depends on ${\bm{w}}$ and in general, cannot be true for
every ${\bm{w}}\in{\mathcal{W}}$. To observe this, note that the left-hand
side of Equation 7 is linear in ${\bm{w}}$, but the right-hand side is
not.111In the case of deterministic policies, the right-hand side will be
piecewise linear.

To make the update of Equation 7 universally true for every ${\bm{w}}$, we
propose a linear conservative approximation:

$\operatorname*{softmax}_{a^{\prime}}{\bm{w}}^{\mathbb{T}}{\bm{Q}}(s^{\prime},a^{\prime})\approx\sum_{a^{\prime}}\pi^{\beta}(a^{\prime}|s^{\prime};\hat{{\bm{w}}}^{\mathbb{T}}{\bm{Q}})\;{\bm{w}}^{\mathbb{T}}{\bm{Q}}(s^{\prime},a^{\prime}),$
(8)

where $\hat{{\bm{w}}}\sim{\mathcal{P}}$. To see that this approximation is
conservative for any $\hat{{\bm{w}}}$, note that for large values of $\beta$,
$\pi^{\beta}$ converges to a deterministic policy which chooses
$\hat{a}=\operatorname*{arg\,max}_{a^{\prime}}\hat{{\bm{w}}}^{\mathbb{T}}{\bm{Q}}(s^{\prime},a^{\prime})$
and we have

$\operatorname*{softmax}_{a^{\prime}}{\bm{w}}^{\mathbb{T}}{\bm{Q}}(s^{\prime},a^{\prime})\rightarrow\max_{a^{\prime}}{\bm{w}}^{\mathbb{T}}{\bm{Q}}(s^{\prime},a^{\prime})\geq{\bm{w}}^{\mathbb{T}}{\bm{Q}}(s^{\prime},\hat{a}).$
(9)

Using the approximation in Equation 8, both sides of Equation 7 become linear
in ${\bm{w}}$ and we can thus derive a vector update for ${\bm{Q}}$
independent of ${\bm{w}}$:

${\bm{Q}}(s,a)\leftarrow{\bm{r}}+\gamma\sum_{a^{\prime}}\pi^{\beta}(a^{\prime}|s^{\prime};\hat{{\bm{w}}}^{\mathbb{T}}{\bm{Q}})\;{\bm{Q}}(s^{\prime},a^{\prime}).$
(10)

The approximation of Equation 8 can be loose for a specific ${\bm{w}}$ if the
optimal action for $\hat{{\bm{w}}}$ at state $s^{\prime}$ diverges
significantly from the optimal action for ${\bm{w}}$. We address this
potential issue in the following manner. First, we consider a prior
${\mathcal{P}}$ over ${\mathcal{W}}$. This prior can rule out implausible
weightings. Second, when considering a fixed ${\bm{w}}$, if
$\pi^{\beta}(a|s;{\bm{w}}^{\mathbb{T}}{\bm{Q}})\ll 1$, we do not need to worry
about the approximation being inaccurate as it most likely underestimates the
value of an action unlikely to be taken. Conversely, if
$\pi^{\beta}(a|s;{\bm{w}}^{\mathbb{T}}{\bm{Q}})$ is substantial, we desire a
tight approximation. Hence, we only require $\hat{{\bm{w}}}$ to be a close
approximation for a ${\bm{w}}$ that results in the selection of action $a$ at
state $s$. Formally, given the prior ${\mathcal{P}}$ and the observation that
action $a$ is taken at state $s$, we can procure a posterior over
${\mathcal{W}}$:

${\mathcal{P}}({\bm{w}}|s,a)\sim{\mathcal{P}}({\bm{w}})\;\pi^{\beta}(a|s;{\bm{w}}^{\mathbb{T}}{\bm{Q}}).$
(11)

We then draw $\hat{{\bm{w}}}$ from this posterior. Any posterior sampling
method can be used here. For simplicity, we use a one-iteration of particle
filtering which is essentially sampling particles from ${\mathcal{P}}$,
reweighting them according to $\pi^{\beta}$, normalizing the weights, and then
resampling.

To further improve the stability and performance of the algorithm, we use
double-Q learning [40] by introducing another vector-valued Q-network
${\bm{Q}}^{\prime}$ as the target network. The target network
${\bm{Q}}^{\prime}$ is updated after multiple updates of ${\bm{Q}}$ by copying
the weights from ${\bm{Q}}$. Our update rule for ${\bm{Q}}$ is:

${\bm{Q}}(s,a)\leftarrow{\bm{r}}+\gamma\sum_{a^{\prime}}\pi^{\beta}(a^{\prime}|s^{\prime};\hat{{\bm{w}}}^{\mathbb{T}}{\bm{Q}})\;{\bm{Q}}^{\prime}(s^{\prime},a^{\prime}),$
(12)

where we draw $\hat{{\bm{w}}}$ from posterior ${\mathcal{P}}(\cdot|s,a)$. We
will refer to this method as _multi-objective Q-learning (MQL)_ in the
following sections.

##### Applying Techniques From Conservative Q-Learning.

Offline Q-learning is prone to overestimating Q-values for states and actions
not present in the dataset. Conservative Q-learning (CQL) has been proposed to
mitigate this issue [18]. CQL seeks to prevent the overestimation of Q-values
for out-of-distribution $(s,a)$ pairs by enforcing a conservative loss
function. Note that our proposed approximate update rule of Equation 12 is
already conservative. However, we can further increase its conservativeness
towards actions that are already present in the data by applying a technique
similar to CQL. Specifically, let ${\bm{Q}}$ be parameterized by $\theta$. We
use an augmented loss function

$\displaystyle{\mathcal{L}}_{\alpha}\big{(}\theta;(s,a,s^{\prime},{\bm{r}})\big{)}={\mathcal{L}}\big{(}\theta;(s,a,s^{\prime},{\bm{r}})\big{)}+\frac{\alpha}{d}\sum_{i\in[d]}\Big{[}\log\big{(}\sum_{\bar{a}}\exp{Q_{i,\theta}(s,\bar{a})}\big{)}-Q_{i,\theta}(s,a)\Big{]},$
(13)

where ${\mathcal{L}}\big{(}\theta;(s,a,s^{\prime},{\bm{r}})\big{)}$ is the
loss from the updating rule of Equation 12:

${\mathcal{L}}\big{(}\theta;(s,a,s^{\prime},{\bm{r}})\big{)}=\big{\|}{\bm{Q}}_{\theta}(s,a)-{\bm{r}}-\gamma\sum_{a^{\prime}}\pi^{\beta}(a^{\prime}|s^{\prime};\hat{{\bm{w}}}^{\mathbb{T}}{\bm{Q}})\;{\bm{Q}}^{\prime}(s^{\prime},a^{\prime})\big{\|}_{2}^{2}.$
(14)

For a sampled batch of data ${\mathcal{B}}$, the overall loss thus becomes:

${\mathcal{L}}_{\alpha}=\sum_{(s,a,s^{\prime},{\bm{r}})\in{\mathcal{B}}}{\mathcal{L}}_{\alpha}\big{(}\theta;(s,a,s^{\prime},{\bm{r}})\big{)}.$
(15)

We will refer to this method as _multi-objective conservative Q-learning
(MCQL)_.

##### A Stochastic Policy Out of Vector-Valued Q-Function.

Neither MQL nor MCQL requires obtaining an explicit policy from ${\bm{Q}}$ to
be trained in the offline setting. However, in the off-policy setting, the
agent interacts with the environment and therefore requires to have a policy
at different points during training. We can define a stochastic policy based
on ${\bm{Q}}$ for a prior ${\mathcal{P}}$ over weightings as

$\pi^{\beta}_{\mathcal{P}}(a|s)\coloneqq\mathbb{E}_{{\bm{w}}\sim{\mathcal{P}}}\big{[}\pi^{\beta}(a|s;{\bm{w}}^{\mathbb{T}}{\bm{Q}})\big{]}.$
(16)

In practice, the agent interacts with the environment by sampling from
$\pi^{\beta}_{\mathcal{P}}$ in two steps:

  1. 1.

Draw ${\bm{w}}\sim{\mathcal{P}}$.

  2. 2.

Draw $a\sim\pi^{\beta}(\cdot|s;{\bm{w}}^{\mathbb{T}}{\bm{Q}})$.

The importance of $\pi^{\beta}_{\mathcal{P}}$ extends beyond the off-policy
setting as we will use it to prune actions at each state in the next section.

###  4.2 Phase 2: Q-Learning with Pruning

The new notion of stochastic policy $\pi^{\beta}_{\mathcal{P}}$ in Equation 16
can be interpreted as a relaxation to the notion of policy with explicit
dependence on a specific ${\bm{w}}$. An important property of
$\pi^{\beta}_{\mathcal{P}}$ is its inclusiveness, meaning that as long as the
accurate reward $r_{*}$ is sufficiently strongly represented in some
${\bm{w}}$ with a noteworthy measure, the action maximizing $r_{*}$ will be
present in $\pi^{\beta}_{\mathcal{P}}$ with a non-negligible probability.
Therefore, a reasonable way to prune the action set is to drop actions with
$\pi^{\beta}_{\mathcal{P}}(a|s)$ below a threshold. But calculating
$\pi^{\beta}_{\mathcal{P}}(a|s)$ requires a posterior calculation that can be
computationally hard. Thus, we propose the following pruning function
$\Pi^{\beta}:{\mathcal{S}}\rightarrow 2^{\mathcal{A}}$ by sampling from
$\pi^{\beta}_{\mathcal{P}}$:

  1. 1.

At state $s$, draw $m$ actions $a^{(k)}\sim\pi^{\beta}_{\mathcal{P}}(\cdot|s)$
for $k\in[m]$:

    1. (a)

Draw $m$ samples ${\bm{w}}^{(k)}\sim{\mathcal{P}}$ for $k\in[m]$.

    2. (b)

For each ${\bm{w}}^{(k)}$, draw
$a^{(k)}\sim\pi^{\beta}(\cdot|s;{\bm{w}}^{(k)^{\mathbb{T}}}{\bm{Q}})$.

  2. 2.

Set $\Pi^{\beta}(s)=\\{a^{(k)}\mid k\in[m]\\}$.

The choice of $m$ and $\beta$ determines the expected size of the set
$\Pi^{\beta}(s)$. As a rule of thumb, actions with
$\pi^{\beta}_{\mathcal{P}}(a|s)<1/m$ are unlikely to remain in the action set
after pruning. In practice, we fix $m$ to at least three times
$|{\mathcal{A}}|$ to make sure that all actions have the chance to be in
$\Pi^{\beta}(s)$ but very rare actions do usually not appear. Then tuning the
value of $\beta$ allows us to control pruning strictness as larger values move
the policy closer to a deterministic policy. We treat $\beta$ as a
hyperparameter.

For a choice of $\beta$ we then run standard double-Q learning using only the
sparse main reward and restricting the action set to the actions chosen by
pruning function $\Pi^{\beta}$. In doing so, we revisit Equation 3:

$Q(s,a)\leftarrow
r+\gamma\,Q^{\prime}\big{(}s^{\prime},\operatorname*{arg\,max}_{a^{\prime}\in\Pi^{\beta}(s^{\prime})}Q(s^{\prime},a^{\prime})\big{)}.$
(17)

It is worth noting that a similar modification of the loss function as in our
discussion of conservative Q-learning [18] yields a conservative version of
Q-learning with pruning.

Examining Equation 17, we observe that $\Pi^{\beta}(s^{\prime})$ is the only
place where we incorporate noisy intermediate reward signals. Consequently,
our final policy focuses solely on optimizing the accurate reward, ensuring
that the original objective remains unaltered. We claim that the
disentanglement of noisy and accurate rewards through this two-stage algorithm
enables us to effectively incorporate information from intermediate signals to
simplify the learning problem while minimally manipulating the policy’s true
objective.

We refer to the two-stage algorithm presented in this paper as _Pruned QL_ ,
which consists of MQL in the first phase and Q-learning with pruning in the
second phase. In the offline setting, we call the conservative version
consisting of MCQL in the first phase and CQL with pruning in the second
phase, _Pruned CQL_.

##  5 Results

We first present the capability of our method within the OpenAI Gym Lunar
Lander environment [2], which provides the advantage of allowing us to observe
the performance of the learned policy when rolled out. Subsequently, we
evaluate our framework in an offline setting using real health records of
septic patients in the intensive care unit (ICU) obtained from the MIMIC
(Medical Information Mart for Intensive Care)-III dataset (v1.4) [16].

###  5.1 Synthetic Experiments on Lunar Lander

We briefly demonstrate the prospects of our action space pruning approach
using the game Lunar Lander. In Lunar Lander, the objective is to successfully
land a small spaceship on the moon. In addition to the ultimate goal, three
intermediate rewards related to the shape and fuel efficiency of the lander
are also available, which can guide the landing process. The availability of a
sparse main objective and informative intermediate signals makes Lunar Lander
an appropriate choice for evaluating our proposed approach.

In particular, we conduct two experiments. First, we examine whether
incorporating intermediate signals through action pruning is effective. To do
so, we train a (double) Q-learning baseline model222Unless otherwise stated,
we use double Q-learning in all experiments. on the sparse reward and compare
it to our Pruned QL method. Please refer to the note at the end of this
section for the model parameters. Figure 1(a) shows the development of the
policy best returns as we interact more with the environment. After $500,000$
iterations the baseline method struggles to learn a good policy solely based
on the sparse reward and yields an average return of $-70$ close to the
minimum of $-100$. In contrast, our Pruned QL approach, which also only
leverages sparse rewards in its final learning phase, achieves a substantially
higher best return of around $50$ within a shorter time. These results hold
for varying pruning strengths $\beta$. It is important to note that Pruned QL
benefits from additional reward signals during its first stage of MQL, making
a direct comparison with the baseline not entirely fair. Nonetheless, these
results are promising in that they underscore the potential of MQL-based
action space pruning.

We next demonstrate how Pruned QL behaves as the intermediate reward signal
becomes less precise. To achieve this we conduct a second experiment where we
add white noise with standard deviations of $0.1$, $1$, and $10$ to the
intermediate rewards. Figure 1(b) depicts the results. We observe that Pruned
QL remains robust until raising the white noise standard deviation to $10$,
where it only does marginally better than the baseline. This result
illustrates two important characteristics of our method. First, it shows that
Pruned QL has generally high robustness to noise. Second, the fact that the
performance drops when the noise standard deviation is increased to $10$
indicates that the intermediate rewards considered in this method must carry
some relevant information about the sparse outcome of interest. At a white
noise standard deviation of $10$, the signal-to-noise ratio (SNR) of the shape
reward is less than $0.5$ and the SNR for the fuel rewards is less than
$0.02$, diluting nearly the entirety of the reward signal. Further evidence
about the noise-robustness of different pruning strengths is provided in
Section A.1.

(a) Varying pruning strength $\beta$, no noise

(b) Varying noise-level, $\beta=40$

Figure 1: Learning curves indicating policy performance of Pruned QL vs. best
baseline Q-Network under (a) different pruning strength, and (b) varying noise
level.

##### Model Parameters.

In both experiments on Lunar Lander we used similar parameters. We used a
Dirichlet distribution with parameter $10$ as the prior over the weighting of
rewards.333We chose a Dirichlet distribution as our prior because it is
flexible, intuitive, and often the first choice for modeling nonnegative
variables that sum to one, like probabilities. Our results are not sensitive
to this parameter choice as long as the chosen prior is not very spread-out or
very concentrated. We did not discount future rewards ($\gamma=1$). The
Q-function is implemented as a 3-layer feed-forward neural network with ReLU
activation. We used the Adam optimizer with a learning rate of $10^{-4}$. The
frequency of updating the target network is once every $3000$ updates of the
Q-network. Please refer to the accompanying code for the remaining Q-learning
parameters.

###  5.2 Real-World Data Experiments on MIMIC-III

####  5.2.1 Cohort and Study Design

We evaluate our framework in a real-life offline learning setting by training
and evaluating policies for the management of vasopressor and intravenous (IV)
fluids in septic patients at intensive care units (ICU). Existing work [17,
30, 7] in this area has mainly focused on reward specifications based on
$90$-day mortality. In contrast, we investigate whether Pruned CQL can
facilitate learning superior policies by incorporating information from
intermediate severity proxies such as the SOFA score [42] and the patient’s
lactate level.

##### Data.

We will use data from a cohort of septic ICU patients in the MIMIC-III
dataset. This dataset consists of deidentified electronic health records of
over $40,000$ patients that were admitted to the critical care units of the
Beth Israel Deaconess Medical Center in Boston between 2001 and 2012. To
construct our sample, we followed the preprocessing steps applied in [17]. We
excluded patients younger than $18$ years old at ICU admission or where
mortality and intravenous fluid intake have not been documented. A trajectory
includes data from $24$ hours before the sepsis onset up until $48$ hours
after the onset, recorded in $4$-hour intervals. For every trajectory, we
extracted a set of $48$ variables, including demographics, Elixhauser
comorbidity index [5], vital signs, laboratory test results, and medication
dosing decisions. All models are trained on $80\%$ of the data, validated on
$5\%$, and tested on $15\%$ for the final evaluation.

##### Action Set and State Space.

Following previous work by [17] and [7], we discretize actions into $25$
treatment choices ($5$ discrete levels for IV fluids and $5$ discrete levels
for vasopressor dosage). To derive the state space, we applied K-means
clustering to the $44$-dimensional features and obtained $752$ clusters. As
the resulting cluster indicators were not numerically meaningful, we
additionally applied the continuous bag of words method [25] to the patients’
state trajectories. This provided us with a $13$-dimensional representation
for each cluster. Thus, our problem formulation is based on a $25$-dimensional
action space and a $13$-dimensional state space. The resulting dataset
consists of $20,912$ ICU stay trajectories, of which $4,917$ resulted in
patient death within $90$ days of the patient’s critical care visit.

##### Reward Specification.

We assign a reward of $100$ to the final state of a trajectory if the patient
survives for at least $90$ days past ICU admission, and a reward of $-100$
otherwise. Achieving a low $90$-day mortality is the ultimate goal of our
learning agent but since this is a sparse and delayed signal, we further
include four medically-motivated rewards that are observed more frequently
throughout the patient’s stay:

  * •

One-period SOFA score change: The negative of the one-period change in SOFA
score.

  * •

Two-period SOFA score change: The negative of the two-period change in SOFA
score.

  * •

One-period lactate level change: The negative of the one-period change in the
lactate level.

  * •

Two-period lactate level change: The negative of the two-period change in the
lactate level.

SOFA [42] is a medical risk score that summarizes the extent of a patient’s
organ failure and in recent years has become a key indicator of the sepsis
syndrome [19]. Arterial lactate levels are an important biomarker for septic
shock because they are closely associated with cellular hypoxia. Sepsis can
cause an imbalance in oxygen supply and demand, resulting in inadequate
delivery of oxygen to the cells and tissues. This can lead to anaerobic
metabolism and the production of lactate as a byproduct. Increased arterial
lactate levels, therefore, indicate that there is a mismatch between oxygen
supply and demand and that the body is experiencing hypoxia [10]. Reward
values are scaled to have a standard deviation of one before feeding to the
algorithms.

##### Model Parameters.

In this offline RL setting, we employ Pruned CQL with a conservativity level
of $\alpha=0.001$ in all experiments. As the trajectories tend to be short in
this setting, we do not discount future rewards ($\gamma=1$). The Q-function
is implemented as a $3$-layer feed-forward neural network with ReLU
activation. We used a Dirichlet distribution with parameter $10$ as the prior
${\mathcal{P}}$ but results are not sensitive to this choice as long as the
prior is not very broad or very concentrated. We treated $\beta$ as a
hyperparameter that takes a value from $20$, $40$, or $160$. We used the Adam
optimizer with a learning rate of $10^{-5}$ in the first phase and $10^{-4}$
in the second phase. The target network is updated after every $1000$ and
$8000$ updates of the Q-network in the first and second phases, respectively.
The baseline model employs the same parameters as the second-phase model.
Please refer to the accompanying code for the remaining Q-learning parameters.

####  5.2.2 Policy Evaluation Approach

In the offline RL setting, it is not feasible to roll out the learned policies
and observe their returns. Therefore, we will evaluate the policies using
weighted importance sampling. Additionally, we will use a descriptive measure
to assess whether the learned Q-functions are capable of distinguishing
between high and low mortality trajectories.

There are several approaches to estimating policy values in offline settings,
each with its own limitations. Interested readers may refer to [38] for an
empirical comparison of different methods and [11] for practical
considerations. One effective class of methods for policy evaluation in the
offline setting is the class of importance sampling methods, which estimate
the policy value by re-weighting trajectories based on their relative
likelihood of occurring. However, this requires having a stochastic policy and
knowledge of the physician’s policy. To accommodate our final phase 2
deterministic policy, we thus follow the recommendation of [37] and soften our
policy:
$\tilde{\pi}(a|s)=(1-\epsilon)\mathbbm{1}\\{a=\pi(s)\\}+\frac{\epsilon}{|{\mathcal{A}}|-1}\mathbbm{1}\\{a\neq\pi(s)\\}$.
Here $\pi$ and $\tilde{\pi}$ are the original deterministic policies and the
softened version, respectively, and $\epsilon$ is a hyperparameter set to
$\epsilon=0.01$. To calculate the importance ratio, we approximate the
physician’s policy with a stochastic policy by training a multi-class logistic
regression with cross-entropy loss. We then implement the weighted importance
sampling (WIS) method to estimate policy values as it has a lower variance
than ordinary importance sampling methods. Detailed information on the
formulation of WIS can be found in [11, 38].

One drawback of the WIS estimator is that it is biased toward the behavioral
policy. Hence we further analyze a descriptive measure to provide a
supplemental perspective to the established off-policy evaluation WIS.
Specifically, we aim to evaluate the capability of our Q-function $Q$ to
differentiate between trajectories that resulted in patient survival compared
to those that led to patient death. This is relevant since a high capacity of
a Q-function to distinguish between low and high-risk state-action pairs
increases our confidence that choosing the actions with the highest predicted
Q-value is beneficial. We adopt the following approach to estimate such a
descriptive measure. Consider a trajectory $\tau$ from the test data and let
$r_{*}^{\tau}$ be the mortality-based reward assigned to its final state. As
noted above, $r_{*}^{\tau}=100$ if the patient survives and
$r_{*}^{\tau}=-100$ otherwise. We compute $Q(s,a)$ for each state-action pair
$(s,a)\in\tau$. Note that $Q(s,a)$ is supposed to be the expected return from
the optimal policy at state $s$ followed by action $a$ and not the return of
the physician’s policy $r_{*}^{\tau}$. But if the optimal policy is
sufficiently similar to the physicians, $Q(s,a)$ shall reflect $r_{*}^{\tau}$,
and we would like to see a larger $Q(s,a)$ if $\tau$ results in the patient
survival. To measure this association, we bin the Q-values into quartiles and
calculate the mortality rate in the lower quartile range $MR(Q_{1})$ and the
mortality rate in the upper quartile range $MR(Q_{3})$. We then report the
difference $\Delta MR=MR(Q_{1})-MR(Q_{3})$. A larger $\Delta MR$ indicates
that the policy is more effective in distinguishing survival from death
trajectories.

####  5.2.3 Policy Evaluation Results

We first observe that directly incorporating rewards based on the SOFA score
and arterial lactate level does not improve performance. In fact, testing a
range of relative weights, we found that the policies that performed the best
were those that assigned the lowest weight to the intermediate rewards. This
observation suggests that directly including intermediate signals may lead to
more challenges than benefits in this setting. We further discuss this
phenomenon and potential explanations for it in Section B.1.

We next evaluate whether our two-stage Pruned CQL algorithm can leverage the
additional information provided by the intermediate rewards to achieve a
better $\Delta MR$ and higher policy value estimated by WIS. The results of
our evaluation are summarized in Table 1. We note four key observations from
our results. First, Pruned CQL often achieves better $\Delta MR$ and WIS
values compared to CQL, indicating that our two-stage strategy has been
effective in extracting information from the intermediate rewards. Second,
Pruned CQL with $\beta=160$ is superior to the physician’s policy with the
observed value of $51.9$. Since WIS is conservative in nature, this is an
encouraging performance for the obtained policy. The standard CQL method does
not show such promise. Third, one potential concern regarding action pruning
is that it may eliminate important options from the action space. However,
even when we perform strong pruning with $\beta=160$, the performance of the
final policy is not substantially affected, and in fact, the final policies
became more similar to those of the physicians, suggesting that the pruning
procedure primarily removed less relevant actions. Last but not least, our
method often suffers from less variance, here measured across $10$ random
seeds, which offers higher reliability.

Furthermore, Table 1 demonstrates that more stringent pruning leads to
policies that are more similar to those of physicians. This is expected, as
stricter pruning corresponds to a greater reliance on intermediate severity
indicators, which are important in the physician’s workflow. Similarly, as
expected a higher value of the conservativity $\alpha$ results in a greater
agreement between the CQL policies and the actions taken by physicians.
Notably, however, at similar levels of agreement with the actions taken by
physicians, the Pruned CQL policy outperforms the baseline CQL policy in terms
of $\Delta MR$ and policy value. This effect is also evident in Figure 2,
where we have plotted $\Delta MR$ and the WIS value against the similarity to
physician actions. The different points on the figure correspond to varying
levels of $\alpha$ and $\beta$. Different trends for $\Delta MR$ and WIS-based
value might origin from the fact that WIS is biased towards policies more
similar to the behavior policy. The results indicate that pruning is more
effective than a standard conservative loss in learning effective RL policies
while enforcing closeness to the behavior policy. Overall, this recommends
action pruning as a promising avenue for incorporating human guidelines into
policy optimization.

Table 1: Comparison of different methods in terms of $\Delta MR$, WIS-based policy value, and percentage of actions taken by the policy that match the physician’s actions. Standard errors are calculated based on ten random seeds. The behavior policy has a return of $51.9$. Model |  $\Delta MR$ (%) |  Value (WIS) |  Similarity to physicians (%)
---|---|---|---
CQL($\alpha=0.001$) | $24.6\pm 1.0$ | $14\pm 18$ | $10.4\pm 0.7$
CQL($\alpha=0.005$) | $23.6\pm 1.0$ | $35\pm 21$ | $18.6\pm 0.6$
CQL($\alpha=0.01$) | $22.6\pm 0.9$ | $26\pm 25$ | $26.2\pm 1.0$
Pruned CQL($\alpha=0.001$, $\beta=20$) | $24.5\pm 0.9$ | $32\pm 12$ | $11.3\pm 1.2$
Pruned CQL($\alpha=0.001$, $\beta=40$) | $25.2\pm 0.6$ | $41\pm 15$ | $15.6\pm 0.9$
Pruned CQL($\alpha=0.001$, $\beta=160$) | $24.2\pm 1.1$ | $66\pm 19$ | $22.1\pm 0.9$

Figure 2: Comparison of Pruned CQL and CQL in terms of $\Delta MR$, WIS-based
policy value, for different degrees of overlap with the behavior policy.
Dashed lines correspond to linear fits.

To further analyze the Q-values obtained from pruned CQL we calculated
Q-values for each state-action pair in the test data and divided them into
$25$ bins. For each bin, we calculated the probability of patient survival.
Figure 3 presents the survival rate per bin for pruned CQL with a
conservativity level of $\alpha=0.001$ and pruning strictness of $\beta=40$.
The plot shows a strong positive relationship, indicating a survival rate
below $30\%$ for state-action pairs with the lowest Q-values and of over
$80\%$ for trajectories with the highest Q-values. These results suggest that
the Pruned CQL policy is based on a Q-function that effectively distinguishes
high- and low-risk state-action pairs, which is crucial for identifying
relevant policies. In Section B.2 we demonstrate similar results for pruning
strictness levels of $\beta=20,160$.

Figure 3: Survival rate by Q-value. Rate plotted for $25$ equal-sized bins.

####  5.2.4 Pruning Analysis

In order to evaluate our pruning quality, we investigate two key questions: 1)
Does our pruning procedure in phase $1$ allow us to significantly reduce the
size of the action space? 2) Is the resulting action space consistent with
current medical practice?

Table 2 presents the average number of available actions after pruning and the
corresponding recall of these action sets. As there is no ground truth best
policy available, we define recall based on whether the pruned action sets
contain the physician actions taken in each respective state in the test set.
Since physicians are highly trained professionals and may have access to
additional indicators not captured in our records, we aim to retain their
decisions among the available options as much as possible. We observe that our
pruning procedure significantly reduces the size of the original
$25$-dimensional action space. Moderate pruning with $\beta=40$ reduces the
mean action set size to less than half, while stricter pruning with
$\beta=160$ reduces the average number of available actions to almost one-
sixth. However, even with such aggressive pruning, our method manages to
maintain a high recall, ranging from $50\%$ to above $90\%$, well beyond the
chance level.

Table 2: Mean action set size and recall for different pruning levels. The initial action set size was $25$. $\beta$ |  Num. of available actions after pruning | Recall (%)
---|---|---
$20$ | $19.7\pm 0.3$ | $94.7\pm 0.3$
$40$ | $11.6\pm 0.7$ | $83.3\pm 1.4$
$160$ | $4.1\pm 0.3$ | $49.4\pm 1.9$

To investigate the pruning behavior further, Figure 4 displays the
distribution of removed actions compared to the actions taken by physicians in
the test set for $\beta=40$. A similar observation holds for less or more
strict pruning levels, as described in Section B.3. The figure suggests that
our procedure primarily prunes extreme dosing regimes or incoherent decisions,
such as a low intravenous fluid dose, while assigning a high vasopressor dose.
These insights further support the idea that the pruning procedure reduces the
complexity of the RL problem while retaining relevant actions for the second
learning stage.

Figure 4: Distribution of observed physician actions and pruned actions
($\beta$=40).

##  6 Discussion

In this study, we have introduced a novel reinforcement learning method that
utilizes action space pruning to facilitate learning when the main reward
signal is sparse, intermediate signals are approximate, and there is ambiguity
about the weights that should be given to optimal reward components. Our work
introduces new algorithms in order to effectively integrate more frequent but
imprecise reward proxies into learning.

We demonstrate in the off-policy setting of the Lunar Lander environment that
action-space pruning enables Q-learning from the sparse reward alone when
standard Q-learning approaches fail. Furthermore, in an offline setting aiming
to design vasopressor and intravenous fluid dosing policies for septic
patients in the ICU, our results indicate the effectiveness of the pruning
approach to reduce the size of the action space while preserving relevant
actions. Our trained policy behaves safely by staying mostly consistent with
the physician’s policy while achieving better performance than a conservative
Q-learning policy for the same level of agreement with the physician’s
actions. These results indicate that our learning framework facilitates the
efficient incorporation of recurrent yet imprecise reward signals while
preventing such signals from causing the policy to diverge from maximizing the
primary sparse reward of interest.

Though motivated by the healthcare setting our approach is applicable to
further domains in which the true reward signal of interest is sparse and
available proxies are frequent but only provide imperfect indicators of the
ultimate outcome.

##### Limitations.

While our study provides valuable insights, it has some limitations. First,
our results are based on data up to the year 2012, and sepsis treatment
guidelines have been evolving since then. Hence, the estimated physician
policies in this work may reflect a slightly outdated standard of care.
Second, our method relies on the availability of intermediate reward signals
that can serve as proxies for the final outcome of interest. In cases where
such meaningful signals are not available, our method may not be applicable.
Finally, in this study, we have utilized weighted importance sampling (WIS) as
an off-policy evaluation method. Although it is a widely accepted technique
for evaluating RL policies in offline settings, it is also known to have high
variance and is dependent on accurate estimation of the behavior policy [38].
We partially addressed this issue by evaluating our policy using a related
descriptive measure that highlights consistent results with the provided WIS
estimates.

## References

  * [1] Leon Barrett and Srini Narayanan  “Learning all optimal policies with multiple criteria”  In _Proceedings of the 25th international conference on Machine learning - ICML ’08_ ACM Press, 2008  DOI: 10.1145/1390156.1390162
  * [2] Greg Brockman et al.  “OpenAI Gym”, 2016  eprint: arXiv:1606.01540
  * [3] Li-Fang Cheng, Niranjani Prasad and Barbara E Engelhardt  “An optimal policy for patient laboratory tests in intensive care units”  In _BIOCOMPUTING 2019: Proceedings of the Pacific Symposium_ , 2018, pp. 320–331  World Scientific
  * [4] Li-Fang Cheng, Niranjani Prasad and Barbara E. Engelhardt  “An Optimal Policy for Patient Laboratory Tests in Intensive Care Units”  In _Biocomputing 2019_ WORLD SCIENTIFIC, 2018  DOI: 10.1142/9789813279827“–˙˝0029
  * [5] Anne Elixhauser, Claudia Steiner, D Robert Harris and Rosanna M Coffey  “Comorbidity measures for use with administrative data”  In _Medical care_ JSTOR, 1998, pp. 8–27
  * [6] M Milani Fard and Joelle Pineau  “Non-deterministic policies in Markovian decision processes”  In _Journal of Artificial Intelligence Research_ 40, 2011, pp. 1–24
  * [7] Mehdi Fatemi, Taylor W Killian, Jayakumar Subramanian and Marzyeh Ghassemi  “Medical Dead-ends and Learning to Identify High-risk States and Treatments”  In _Advances in Neural Information Processing Systems_ 34, 2021, pp. 4856–4870
  * [8] Mehdi Fatemi, Shikhar Sharma, Harm Van Seijen and Samira Ebrahimi Kahou  “Dead-ends and Secure Exploration in Reinforcement Learning”  In _International Conference on Machine Learning_ 97, 2019, pp. 1873–1881  URL: http://proceedings.mlr.press/v97/fatemi19a/fatemi19a.pdf
  * [9] Joseph Futoma et al.  “An Improved Multi-Output Gaussian Process RNN with Real-Time Validation for Early Sepsis Detection”  In _arXiv: Machine Learning_ , 2017  URL: https://arxiv.org/pdf/1708.05894.pdf
  * [10] G Gernardin et al.  “Blood pressure and arterial lactate level are early indicators of short-term survival in human septic shock”  In _Intensive care medicine_ 22 Springer, 1996, pp. 17–25
  * [11] Omer Gottesman et al.  “Evaluating reinforcement learning algorithms in observational health settings”  In _arXiv preprint arXiv:1805.12298_ , 2018
  * [12] Katharine E Henry, David N Hager, Peter J Pronovost and Suchi Saria  “A targeted real-time early warning score (TREWScore) for septic shock”  In _Science translational medicine_ 7.299 American Association for the Advancement of Science, 2015, pp. 299ra122–299ra122
  * [13] Kazuyuki Hiraoka, Manabu Yoshida and Taketoshi Mishima  “Parallel reinforcement learning for weighted multi-criteria model with adaptive margin”  In _Cognitive Neurodynamics_ 3.1 Springer ScienceBusiness Media LLC, 2008, pp. 17–24  DOI: 10.1007/s11571-008-9066-9
  * [14] Hitoshi Iima and Yasuaki Kuroe  “Multi-objective reinforcement learning for acquiring all Pareto optimal policies simultaneously - Method of determining scalarization weights”  In _2014 IEEE International Conference on Systems, Man, and Cybernetics (SMC)_ IEEE, 2014  DOI: 10.1109/smc.2014.6974022
  * [15] Alex Irpan et al.  “Off-Policy Evaluation via Off-Policy Classification”  In _Neural Information Processing Systems_ 32, 2018, pp. 5437–5448  URL: https://arxiv.org/pdf/1906.01624.pdf
  * [16] Alistair EW Johnson et al.  “MIMIC-III, a freely accessible critical care database”  In _Scientific data_ 3.1 Nature Publishing Group, 2016, pp. 1–9
  * [17] Matthieu Komorowski et al.  “The artificial intelligence clinician learns optimal treatment strategies for sepsis in intensive care”  In _Nature medicine_ 24.11 Nature Publishing Group, 2018, pp. 1716–1720
  * [18] Aviral Kumar, Aurick Zhou, George Tucker and Sergey Levine  “Conservative q-learning for offline reinforcement learning”  In _Advances in Neural Information Processing Systems_ 33, 2020, pp. 1179–1191
  * [19] Simon Lambden, Pierre Francois Laterre, Mitchell M Levy and Bruno Francois  “The SOFA score—development, utility and challenges of accurate assessment in clinical trials”  In _Critical Care_ 23.1 BioMed Central, 2019, pp. 1–9
  * [20] Rongmei Lin, Matthew D. Stanley, Mohammad M. Ghassemi and Shamim Nemati  “A Deep Deterministic Policy Gradient Approach to Medication Dosing and Surveillance in the ICU”  In _2018 40th Annual International Conference of the IEEE Engineering in Medicine and Biology Society (EMBC)_ IEEE, 2018  DOI: 10.1109/embc.2018.8513203
  * [21] Chunming Liu, Xin Xu and Dewen Hu  “Multiobjective Reinforcement Learning: A Comprehensive Overview”  In _IEEE Transactions on Systems, Man, and Cybernetics: Systems_ 45.3 Institute of ElectricalElectronics Engineers (IEEE), 2015, pp. 385–398  DOI: 10.1109/tsmc.2014.2358639
  * [22] Siqi Liu et al.  “Reinforcement Learning for Clinical Decision Support in Critical Care: Comprehensive Review”  In _Journal of Medical Internet Research_ 22.7 JMIR Publications Inc., 2020, pp. e18477  DOI: 10.2196/18477
  * [23] Daniel J Lizotte and Eric B Laber  “Multi-objective Markov decision processes for data-driven decision support”  In _The Journal of Machine Learning Research_ 17.1 JMLR. org, 2016, pp. 7378–7405
  * [24] Shie Mannor and Nahum Shimkin  “The Steering Approach for Multi-Criteria Reinforcement Learning”  In _Neural Information Processing Systems_ 14, 2001, pp. 1563–1570
  * [25] Tomas Mikolov, Kai Chen, Greg Corrado and Jeffrey Dean  “Efficient estimation of word representations in vector space”  In _arXiv preprint arXiv:1301.3781_ , 2013
  * [26] William Dwight Miller et al.  “Accuracy of the sequential organ failure assessment score for in-hospital mortality by race and relevance to crisis standards of care”  In _JAMA network open_ 4.6 American Medical Association, 2021, pp. e2113891–e2113891
  * [27] Hossam Mossalam, Yannis M. Assael, Diederik M. Roijers and Shimon Whiteson  “Multi-Objective Deep Reinforcement Learning”  In _arXiv: Artificial Intelligence_ , 2016  URL: https://arxiv.org/pdf/1610.02707.pdf
  * [28] Sriraam Natarajan and Prasad Tadepalli  “Dynamic preferences in multi-criteria reinforcement learning”  In _Proceedings of the 22nd international conference on Machine learning - ICML ’05_ Bonn, Germany: ACM Press, 2005, pp. 601–608  DOI: 10.1145/1102351.1102427
  * [29] Arne Peine et al.  “Development and validation of a reinforcement learning algorithm to dynamically optimize mechanical ventilation in critical care”  In _npj Digital Medicine_ 4.1 Springer ScienceBusiness Media LLC, 2021  DOI: 10.1038/s41746-021-00388-6
  * [30] Xuefeng Peng et al.  “Improving Sepsis Treatment Strategies by Combining Deep and Kernel-Based Reinforcement Learning.”  In _American Medical Informatics Association Annual Symposium_ 2018, 2018, pp. 887–896
  * [31] Niranjani Prasad et al.  “A Reinforcement Learning Approach to Weaning of Mechanical Ventilation in Intensive Care Units”  In _arXiv: Artificial Intelligence_ , 2017  URL: https://arxiv.org/pdf/1704.06300.pdf
  * [32] Aniruddh Raghu et al.  “Continuous state-space models for optimal sepsis treatment: a deep reinforcement learning approach”  In _Machine Learning for Healthcare Conference_ , 2017, pp. 147–163  PMLR
  * [33] Aniruddh Raghu, Matthieu Komorowski and Sumeetpal Singh  “Model-based reinforcement learning for sepsis treatment”  In _arXiv preprint arXiv:1811.09602_ , 2018
  * [34] Suchi Saria  “Individualized sepsis treatment using reinforcement learning”  In _Nature Medicine_ 24.11 Springer ScienceBusiness Media LLC, 2018, pp. 1641–1642  DOI: 10.1038/s41591-018-0253-x
  * [35] Zhao Song, Ron Parr and Lawrence Carin  “Revisiting the softmax bellman operator: New benefits and new perspective”  In _International conference on machine learning_ , 2019, pp. 5916–5925  PMLR
  * [36] Shengpu Tang, Aditya Modi, Michael Sjoding and Jenna Wiens  “Clinician-in-the-loop decision making: Reinforcement learning with near-optimal set-valued policies”  In _International Conference on Machine Learning_ , 2020, pp. 9387–9396  PMLR
  * [37] Shengpu Tang, Aditya Modi, Michael W. Sjoding and Jenna Wiens  “Clinician-in-the-Loop Decision Making: Reinforcement Learning with Near-Optimal Set-Valued Policies”  In _International Conference on Machine Learning_ 1, 2020, pp. 9387–9396  URL: http://proceedings.mlr.press/v119/tang20c/tang20c.pdf
  * [38] Shengpu Tang and Jenna Wiens  “Model selection for offline reinforcement learning: Practical considerations for healthcare settings”  In _Machine Learning for Healthcare Conference_ , 2021, pp. 2–35  PMLR
  * [39] Gerald Tesauro et al.  “Managing Power Consumption and Performance of Computing Systems Using Reinforcement Learning”  In _Neural Information Processing Systems_ 20, 2007, pp. 1497–1504  URL: http://lass.cs.umass.edu/~shenoy/courses/fall09/691gc/papers/Kephart-2007-1-NIPS2007_0906.pdf
  * [40] Hado Van Hasselt, Arthur Guez and David Silver  “Deep reinforcement learning with double q-learning”  In _Proceedings of the AAAI conference on artificial intelligence_ 30, 2016
  * [41] Kristof Van Moffaert, Madalina M. Drugan and Ann Nowe  “Scalarized multi-objective reinforcement learning: Novel design techniques”  In _2013 IEEE Symposium on Adaptive Dynamic Programming and Reinforcement Learning (ADPRL)_ IEEE, 2013  DOI: 10.1109/adprl.2013.6615007
  * [42] J.-L. Vincent et al.  “The SOFA (Sepsis-related Organ Failure Assessment) score to describe organ dysfunction/failure”  In _Intensive Care Medicine_ 22.7, 1996, pp. 707–710  DOI: 10.1007/BF01709751
  * [43] Jie Xu et al.  “Prediction-Guided Multi-Objective Reinforcement Learning for Continuous Robot Control”  In _International Conference on Machine Learning_ 1, 2020, pp. 10607–10616  URL: http://proceedings.mlr.press/v119/xu20h/xu20h.pdf

##  Appendix A Further Insights From LunarLander Experiments

###  A.1 Robustness to Noise

To augment the findings illustrated in Figure 1, we depict the effect of noise
on Pruned QL performance in for $\beta=20,160$ in Figure 5. Overall, the main
results discussed continue to hold, Pruned QL behaves robustly to the noise
until raising the noise standard deviation to $10$.

(a) $\beta=20$

(b) $\beta=160$

Figure 5: Learning curves indicating policy performance of Pruned QL vs. best
baseline Q-Network under different noise levels.

##  Appendix B Further Insights From MIMIC-III Experiments

###  B.1 Learning From Individual Rewards

We run standard Q-learning on each individual reward. Figure 6 shows learning
during training with different settings. The learning rate of $10^{-4}$ and
target network update frequency (TUF) of $8000$ are the optimal choice of
parameters to learn from the sparse reward. However, as the figure suggests,
learning from intermediate rewards like SOFA-1 and LAC-1 requires a slower
learning rate and decreased TUF. Further, the optimal choice of the number of
iterations before overfitting also differs for two types of reward. Another
way to see this is in Figure 6 where we have plotted the best $\Delta MR$ for
each setting. One can see that changing the learning rate and TUF has opposite
effects on Q-learning from the sparse reward and Q-learning from intermediate
rewards. This observation can explain why linearly mixing sparse and
intermediate rewards might not be a good idea.

Figure 6: Running Q-learning based on different rewards and settings. The
learning curves show intermediate rewards require a slower learning rate and
lower target update frequency. The bar plot shows the best $\Delta MR$ in each
setting. The darker bars correspond to the learning rate of $10^{-4}$ and TUF
of $8000$.

###  B.2 Survival Rate Versus Q-Value for Different Pruning Strengths

To augment the findings illustrated in Figure 3, we depict the survival curve
for pruning strength $\beta=20,160$ in Figure 7.

(a) $\beta=20$

(b) $\beta=160$

Figure 7: Relationship between Q-values and patient survival for different
$\beta$.

###  B.3 Pruning for Different Pruning Strengths

To augment the findings illustrated in Figure 4, we depict the effect of
pruning strength on the distribution of removed actions in Figure 8 for
$\beta=20,160$.

(a) $\beta=20$

(b) $\beta=160$

Figure 8: Distribution of pruned actions for different pruning strengths.