Towards Unknown-aware Deep Q-Learning

1 INTRODUCTION . As reinforcement learning ( RL ) becomes increasingly important for sequential decision-making tasks , it must entail strong safety and reliability guarantees . In particular , RL agents deployed in open-world environments are naturally subject to sudden environmental shifts or data distributional changes . Unfortunately , current RL models commonly make the closed-world assumption , i.e. , the state transitions follow a potentially time-varying Markov Decision Process st+1 ∼ T ( · | st , at ) , where T is the transition model and st , at are the state and action at timestamp t. An RL agent in deployment may encounter and fail to identify unknown out-of-distribution ( OOD ) states — i.e. , state st that is not from the training environment . Such unexpected states can be potentially catastrophic . For example , if a self-driving agent trained for an urban driving environment would be exposed to OOD states such as mountain roads with rocks , it would still output an action but may be a consequential one . This gives rise to the importance of detecting and mitigating unexpected OOD states , which allows the agent to avoid non-safe actions . Currently , the study of handling OOD states in RL remains underexplored . Our paper aims to bridge this critical gap . As a motivating example , we first show that state-of-the-art reinforcement learning networks do not necessarily know what they don ’ t know , and can assign abnormally high Q values for OOD states ( Figure 1 ) . Previous works on OOD detection in RL ( Sedlmeier et al. , 2019 ; Hoel et al. , 2020 ) rely on an ensemble of neural networks to estimate the variances of Q values ( Osband et al. , 2016 ; 2018 ) , which do not fundamentally mitigate the issue of the incorrectly high Q values in the presence of OOD inputs ( see Section 3.1 ) . Moreover , ensemble-based approaches can be computationally expensive as both the training and inference time will increase accordingly with the ensemble size . To this end , we propose a novel unknown-aware deep Q-learning framework—enabling the RL agent to perform OOD uncertainty estimation while guaranteeing the performance on the original in-distribution ( ID ) environment . We formalize ID and OOD states in RL ( Section 2 ) , and tackle the problem of detecting OOD states without extra Bayesian modeling ( Section 3 ) . Specifically , We propose to estimate the OOD uncertainty using the uncertainty score derived directly from the { Q ( st , a ) , a ∈ A } values . This determines the validity of st , which can be either in-distribution or OOD . Our key idea is to regularize the training process so that OOD states will have higher uncertainty scores , while in-distribution states will have relatively lower scores , therefore making them distinguishable . This is in contrast with vanilla Q-learning algorithms that do not take into account unknowns during training . During inference time , our OOD uncertainty estimation score can be flexibly computed on a single deployed model without variance estimation required . We provide theoretical guarantees that our method can improve OOD uncertainty estimation while ensuring the convergence of ID performance ( Section 4 ) . Empirically , we demonstrate state-of-theart performance on six diverse environments from OpenAI gym ( Brockman et al. , 2016 ) , including Atari games ( with discrete state space ) and classic control task ( with continuous state space ) . In all environments , our method consistently achieves near-optimal OOD detection performance , outperforming the baselines by a substantial margin . Our method can be easily plugged into existing Q-learning algorithms including DQN ( Mnih et al. , 2015 ) and QR-DQN ( Dabney et al. , 2018b ) and render stronger OOD detection performance and RL safety ( Section 5 ) . Our main contributions are summarized as follows : • We propose a novel Q-learning framework for improving OOD uncertainty estimation , which aims at mitigating the critical issue of incorrectly high Q values from OOD state inputs in vanilla Q-learning algorithms ; • We provide theoretical guarantees that our regularization improves OOD detection while maintaining the convergence performance of ID in Q-learning ; • Our method establishes state-of-the-art OOD detection performance on a wide range of RL environments , including Atari games and classic control . Extensive ablation studies are conducted to show the efficacy of our algorithm . 2 PRELIMINARIES AND PROBLEM STATEMENT . 2.1 PRELIMINARIES : REINFORCEMENT LEARNING AND Q-LEARNING . Markov decision processes ( MDP ) . In reinforcement learning , the agent explores an environment and maximizes the cumulative reward . Markov decision processes ( MDPs ) are used to model RL . MDPs are defined by the tuples ( S , A , R , T ) , where S indicates the state space , A is the set of actions , T is the transition model between states , andR is the reward model . At each timestep t , the agent observes a state st ∈ S , chooses an action at ∈ A , receives a reward rt ∈ R , and transits to a new state st+1 determined by T ( st+1|st , at ) . The goal is to find a policy π ∈ Π : S → A , which maximizes the discounted return over a potentially infinite horizon : Rt = Σ ∞ k=0γ krt+k , ( 1 ) where γ ∈ [ 0 , 1 ] is the discount factor . For notation brevity , we use s′ = st+1 and omit the timestamp t when the context is clear . Reinforcement learning and Q-learning . To find the optimal policy π∗ , we focus on Qlearning ( Watkins & Dayan , 1992 ) , a commonly used value-based approach . In Q-learning , the action-value function Q : S × A → R describes the expected return when the agent takes action a in state s , and then following policy π for all states afterwards . The optimal function Q∗ is defined as : Q∗ ( s , a ) = max π Eπ [ R|s , a ] . ( 2 ) Given a tuple of one-step transition ( s , a , s′ , r ) , we have the Bellman equation Q∗ ( s , a ) = r + γmaxa′ Q ∗ ( s′ , a′ ) using dynamic programming . The one-step TD error is used to update Q values : TD ( s , a , r , s′ ) = r + γmax a′ Q ( s′ , a′ ) −Q ( s , a ) . ( 3 ) In tabular cases , a table of Q values are maintained for finite state-action pairs , and are updated according to the TD error : Q ( s , a ) ← Q ( s , a ) +α ( r+ γmaxa′ Q ( s′ , a′ ) −Q ( s , a ) ) , where α is the learning rate . This can also be viewed as gradient updates that minimize the L2 norm of TD error . Deep Q-Networks ( DQN ) . In high-dimensional settings or when learning in continuous statespaces , it is common to use parameterized functions such as neural networks to approximate the action-value function Q ( s , a ) . We use the deep Q-networks ( Mnih et al. , 2015 ) parameterized by θ , with corresponding action-value function Qθ . The norm of the TD error is minimized during training to update the parameter θ . 2.2 OUT-OF-DISTRIBUTION DETECTION IN REINFORCEMENT LEARNING . When deploying an RL agent in open-world environments , it is paramount for the agent to detect OOD states and refrain from taking actions according to policy π in the presence of OOD states . Here we provide the definition of OOD states in the setting of RL . In-distribution ( ID ) states . Given an environment and a learning algorithm , we can estimate the distribution of the states encountered by the agent during training . We denote P as in-distribution . Out-of-distribution ( OOD ) states . The general notion of OOD refers to inference-time inputs which deviate substantially from the data encountered during training1 . In other words , OOD states are unexpected and should not have occurred in the ID environment . Importantly , OOD states should not be confused with under-explored states : the latter still belongs to ID with more epistemic uncertainty than ID states , and are crucial for exploration ( Osband et al. , 2016 ) . In contrast , OOD states are not safe to explore , since the agent has no knowledge on such states during training . We expect significantly lower Q values from OOD states , because unlike the underexplored states , the agent should not be encouraged to explore OOD states . OOD state detection . The goal of OOD detection is to define a decision function G such that : G ( s ) = { 1 if s ∼ P , 0 otherwise . In other words , the agent should be able to know what it does not know . The agent should take action a ∼ π ( s ) only if s is in-distribution . 3 METHOD . In this section , we begin with a motivating example that existing Q-learning algorithms can produce incorrectly high Q values for OOD states ( Section 3.1 ) . To mitigate this issue , we introduce our proposed unknown-aware Q-learning method in Section 3.2 . 3.1 A MOTIVATING EXAMPLE . Incorrectly high Q values for OOD states . Here we present a motivating example of possible incorrect estimation of Q values for OOD states using existing Q-learning algorithms , particularly vanilla DQN ( Mnih et al. , 2015 ) and bootstrapped DQN ( Osband et al. , 2016 ; 2018 ) . We show that such incorrectly high extrapolation of Q values can not be mitigated by traditional uncertainty estimation . For simplicity , we use the chain environment to illustrate the problem ; see Figure 1 ( a ) . Action 0 indicates move left and Action 1 indicates move right . The start state is selected uniformly randomly among 5 states { 1 , 2 , 3 , 4 , 5 } . If the agent arrives on state 0 or 6 , a reward of 1 is received and the terminal is reached . Otherwise , the reward is 0 . We train a vanilla DQN and bootstrapped DQN ( with an ensemble size 10 ) , both under the same MLP architecture for 100,000 steps . The MLP has two hidden layers with 256 units each . Figure 1 1By “ inference-time ” , we assume the RL agent has already converged to policy π ; this policy is now fixed and used to behave in the environment , which includes the majority of cases in RL deployment . ( b ) ( c ) show the Q values for action 0 and 1 respectively . We observe that they give the same estimation for in-distribution Q values , which indicates that they both achieve convergence . Vanilla DQN fails to extrapolate reasonably on OOD states ( states < 0 and > 6 ) , where the corresponding Q values are incorrectly high . Bootstrapped DQN ( Osband et al. , 2016 ; 2018 ) ( denoted as the ensemble in the plots ) displays higher variance for OOD state inputs , but the overall estimated distribution of OOD Q values are still incorrectly high . Motivated by this , we now introduce our method , which regularizes the Q values of OOD states and improves OOD detection without compromising the in-distribution Q values ; see the orange lines shown in Figure 1 ( b ) and ( c ) .

Summary: This paper proposed a regularized loss for punishing the Q values for OOD states. They theoretically show this added loss will not influence the performance of vanilla DQN in both the tabular case and the DQN case. The theoretical analyses are intuitive with remarks that help the reader understand the importance of these results. They design experiments for evaluating the OOD detection performance and game performance. The rich ablation study helps the reader understand how the hyper-parameters are tuned.