Nested Policy Reinforcement Learning for Clinical Decision Support

Off-policy reinforcement learning ( RL ) has proven to be a powerful framework for guiding agents ’ actions in environments with stochastic rewards and unknown or noisy state dynamics . In many real-world settings , these agents must operate in multiple environments , each with slightly different dynamics . For example , we may be interested in developing policies to guide medical treatment for patients with and without a given disease , or policies to navigate curriculum design for students with and without a learning disability . Here , we introduce nested policy fitted Q-iteration ( NFQI ) , an RL framework that finds optimal policies in environments that exhibit such a structure . Our approach develops a nested Q-value function that takes advantage of the shared structure between two groups of observations from two separate environments while allowing their policies to be distinct from one another . We find that NFQI yields policies that rely on relevant features and perform at least as well as a policy that does not consider group structure . We demonstrate NFQI ’ s performance using an OpenAI Gym environment and a clinical decision making RL task . Our results suggest that NFQI can develop policies that are better suited to many real-world clinical environments . 1 INTRODUCTION . Off-policy reinforcement learning ( RL ) has proven to be a powerful framework for guiding agents ’ actions in environments with stochastic rewards and unknown or noisy state dynamics ( 13 ; 2 ) . Previous off-policy RL algorithms ( 9 ) learn a single optimal policy designed to maximize agents ’ longterm expected rewards . However , such a learning scheme is not always ideal . In many real-world settings , agents ’ environments can be divided into multiple groups for the same task . These groups may share rewards but have different state dynamics , which often entails distinct optimal policies . For example , say that we are designing a clinical decision support algorithm to treat ischemic stroke in two groups of patients : one group has a large number of patients with healthy kidneys and the other has a smaller number of patients with renal ( kidney ) disease . Note that there are similarities in the treatment plans for these two groups ; in particular , both sets of patients exhibit the effects of disturbed blood flow to the brain due to stroke , and may require similar procedures . However , patients with renal disease must follow a treatment plan that considers the inability of their kidneys to properly filter out toxins . Traditional off-policy RL algorithms would learn a single policy that does not consider the underlying presence of renal disease . This policy is likely to overfit to the patients with functioning kidneys and overlook any known differences in optimal treatment plans . In such scenarios , it would be beneficial to learn distinct policies for pre-defined groups of patients , while leveraging any shared environmental structure and reward functions between the groups to ensure a large and effective sample size . To address this problem , we introduce nested policy fitted Q-iteration ( NFQI ) , an off-policy RL algorithm that extends fitted Q-iteration ( FQI ) ( 9 ) to account for nested environments . Here , nested refers to an environment structured with two predefined groups that share some characteristics ; we will refer to the two groups as background and foreground . The background group often contains far more samples . In the clinical example above , the background group includes patients with healthy kidneys and the foreground group includes renal disease patients . Both groups share a reward function , in this case successfully treating a stroke , but have different state dynamics . The contributions of this paper to off-policy RL in the presence of nested environments are threefold : 1 . We introduce NFQI , a supervised learning method that uses a nested Q-value function to learn policies that are suited to nested environments . NFQI concurrently models both the background and foreground datasets and is agnostic to modeling choices . 2 . We develop a training procedure inspired by transfer learning ( 31 ) to fit this nestedQ-value function and allow for group imbalance in the data . 3 . We show that NFQI generates policies that outperform baseline methods on a set of quantitative and qualitative metrics and rely on relevant state features to do so . This paper proceeds as follows . Section 2 discusses related work . In Section 3 , we present a framework for NFQI . In Section 4 , we show results for NFQI in simulations and hospital data . Section 5 concludes and discusses potential future work . 2 RELATED WORK . Multi-task RL . The goal of multi-task RL is to train an agent to perform optimally at multiple tasks ( 8 ) . Recent work has developed a method to efficiently learn a task-conditioned policy π ( a|s , z ) , where a is an action , s is a state , and z is a task embedding ( 30 ) . This approach has been observed to perform better than single-task learning because it uses the relationships between tasks to improve performance for individual tasks ( 27 ) . However , the goals of multi-task RL differ from those of NFQI : Multi-task RL optimizes performance for multiple tasks , each with a different Markov decision process ( MDP ) , whereas NFQI is most appropriate to compare two environments with identical tasks and reward functions but different transition dynamics . Additionally , the neural networks typically used for multi-task RL ( 30 ) require large training sets that are not always available in the medical and learning applications NFQI was built to address . As such , these networks are likely to suffer from overparameterization in these applications . Meta RL . Meta RL algorithms train agents to solve new tasks efficiently based on experience with other similar tasks . These approaches infer the task from a distribution p ( T ) using a small set of observations , rather than through an explicit task embedding . Much of the effort in meta RL is focused on task design as opposed to model tuning , to make this inference process easier ( 12 ) . Other work has shown that meta RL can be appropriate for offline RL ( 21 ) . Like multi-task RL , meta RL is different from NFQI because it trains agents to perform multiple tasks with different state and action spaces . Unlike meta RL and multi-task RL , NFQI expects all samples to arise from MDPs that have the same state and action spaces and follow the same reward function . Since the dynamics of the two environments vary , the MDPs will have different state transition functions , entailing different optimal policies for the groups . Transfer learning . Transfer learning methods fine-tune a pre-trained model for a target dataset . Often , the target dataset is relatively small ; transfer learning using a larger and similar dataset may result in better performance than training a model on just the target dataset ( 6 ) . A common way to use these methods is to first train a neural network using the large dataset ; then , some of the network layers are frozen , and the rest of the network is fine-tuned by training on the target dataset . In this way , the network is adapted to the target problem by learning a nonlinear map from the solution estimated for the larger dataset to the outcome for the target dataset . This strategy can be applied to reinforcement learning ( 31 ) and may improve resource usage when the target dataset is small ( 4 ) . The two datasets that NFQI uses are likely to have substantial sample size imbalance ; for this reason , we train the parameters for NFQI using a procedure inspired by transfer learning . 3 METHODS . In this section we present nested policy fitted Q-iteration ( NFQI ) , an off-policy RL algorithm for environments that have a nested group structure . We first review Q-learning and fitted Q-iteration ( FQI ) and then describe how to extend them to NFQI . 3.1 NOTATION AND PRELIMINARIES . A Markov decision process ( MDP ) is a discrete-time stochastic control process defined by a 4- tuple ( S , A , P , R ) , where S is the state space ; A is the set of actions ; P ( st+1|st , at ) is a function governing state transition probabilities from time t to t + 1 ; and R : S → R is a reward function . We assume a finite number of actions . A policy π : S → A is a mapping from states to actions that describes an agent ’ s strategy within an MDP . In general , we assume st ∈ Rp and at ∈ Rq . 3.2 Q-LEARNING AND FITTED Q-ITERATION ( FQI ) Fitted Q-iteration ( FQI ) is an off-policy RL algorithm based on Q-learning ( 9 ; 29 ) that is designed to find an optimal policy for a given MDP . Q-learning approaches define a function Q : S ×A → R that maps state-action pairs to an estimate of the cumulative expected reward obtained by taking a particular action in a particular state . For a given policy π at time step t , Q is defined as Qπ ( st , at ) = rt + γEs′ [ V π ( s′ ) ] , where the expectation is taken over the possible next states given the transition probabilities P , V π is the value function corresponding to π , and γ ∈ [ 0 , 1 ] is a discount factor . The principle of Bellman optimality ( 5 ) states that a policy π ? is optimal for a given MDP if and only if , for every state s ∈ S , it holds that V π∗ ( s ) = maxaR ( s , a ) + γEs′ [ V π∗ ( s′ ) ] . Broadly , Q-learning algorithms attempt to approximate the Q function using agents ’ observed interactions with an environment , as the state transition function is typically not directly observable . This approximation is used to recommend optimal actions for a given state . The data used to estimate the Q function typically consist of samples representing steps taken within an MDP . In particular , a sample at time t is a 4-tuple ( st , at , st+1 , rt ) representing the agent ’ s state , chosen action , observed next state , and observed reward . FQI approximates Q using an offline batch optimization approach ( 9 ) . Specifically , it first posits a family of functions fθ ( s , a ) , parameterized by θ , that are meant to approximate the expected future reward after taking action a from state s. Common choices for f are regression trees ( 9 ) and neural networks ( 25 ) . FQI then employs an iterative algorithm that alternates between updating the current Q-values and fitting a regression problem that optimizes the parameters θ . Suppose we have a set of T samples from an MDP , represented as a set of one-step transition tuples { ( st , at , st+1 , rt ) } Tt=1 . On step k of the algorithm , FQI first updates its estimate of the Q-function according to the Bellman equation : Q̂k ( st , at ) = rt+1 + γmaxa′∈A fθk−1 ( st+1 , a ′ ) for t = 1 , . . . , T . The second step minimizes a loss measuring the discrepancy between the current Q-function estimate and the approximating function f . The optimization problem is minθk ∑N n=1 L ( Q̂k ( s n t , a n t ) , fθk ( s n t , a n t ) ) , where L ( y , ŷ ) is the loss function comparing the “ true ” value y to predicted value ŷ . After K iterations , FQI ’ s approximation to the optimal policy is given by π ( s ) = argmaxa′∈A fθK ( s , a ′ ) .

Summary: This paper considers decision making in nested MDPs where there are two distinct groups having partially shared (but unknown) state dynamics. A new algorithm called NFQI (nested fitted-Q iteration) is presented, and it’s based on a simple modification to FQI, a commonly used offline RL algorithm. Baseline comparisons and sensitivity analyses were done in a set of simulated experiments and using OpenAI gym CartPole. The proposed approach is then applied to a real-world RL task formulated from MIMIC-IV EHR data.