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

TL;DR

This work introduces a deep Q-learning approach to obtain more reliable critical care policies by integrating relevant but noisy frequently measured biomarker signals into the reward specification without compromising the optimization of the main outcome.

Abstract

Medical treatments often involve a sequence of decisions, each informed by previous outcomes. This process closely aligns with reinforcement learning (RL), a framework for optimizing sequential decisions to maximize cumulative rewards under unknown dynamics. While RL shows promise for creating data-driven treatment plans, its application in medical contexts is challenging due to the frequent need to use sparse rewards, primarily defined based on mortality outcomes. This sparsity can reduce the stability of offline estimates, posing a significant hurdle in fully utilizing RL for medical decision-making. We introduce a deep Q-learning approach to obtain more reliable critical care policies by integrating relevant but noisy frequently measured biomarker signals into the reward specification without compromising the optimization of the main outcome. Our method prunes the action space based on all available rewards before training a final model on the sparse main reward. This approach minimizes potential distortions of the main objective while extracting valuable information from intermediate signals to guide learning. We evaluate our method in off-policy and offline settings using simulated environments and real health records from intensive care units. Our empirical results demonstrate that our method outperforms common offline RL methods such as conservative Q-learning and batch-constrained deep Q-learning. By disentangling sparse rewards and frequently measured reward proxies through action pruning, our work represents a step towards developing reliable policies that effectively harness the wealth of available information in data-intensive critical care environments.

Figures (8)

• Figure 1: Illustration of our algorithm. The model first leverages all rewards in order to prune the action space for each state. Then another policy is trained based on the sparse main reward but with its action space restricted to the actions available after pruning in the first stage.
• Figure 2: Performance of Pruned QL in the Lunar Lander and Sepsis Simulator environments. (a) and (b) demonstrate that Pruned QL consistently achieves higher returns compared to a baseline Q-learning method that only has access to the sparse reward. Returns are estimated on new rollouts, not in the training data. (c) and (d) highlight that Pruned QL matches oracle Q-learning across varying intermediate reward weights, demonstrating its robustness and ability to leverage reward information effectively.
• Figure 3: Comparison of Pruned CQL and CQL in terms of $\Delta MR$ and WIS-based policy value, for different degrees of overlap with the behavior policy. Dashed lines display linear fits.
• Figure 4: Distribution of physician's actions and pruned actions ($\beta$=40).
• Figure 5: Intermediate and sparse reward signals along the patient trajectory.