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: The paper proposes an off-policy RL called NFQI (Nested Fitted Q-Iteration) as an extension of Fitted Q-iteration to estimate group-specific policies and account for group structure in the data having two pre-defined groups of observations (background and foreground). NFQI imposes a structure on the family of function to approximate the Q function with guarantees to converge to Bellman optimal Q-value function. NFQI is trained in two stages: The first stage the shared component is trained with all samples. In the second stage, the foreground component is trained using only the foreground samples. The NFQI method is validated on a nested cart-pole environment where the background is the original environment and foreground includes a constant force that pushes the cart to the left. NFQI is also validated using a real-world clinical data from MIMIC dataset.