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 considers the problem of detecting out-of-distribution states to avoid taking actions at states that the policy has not been trained on. The general approach is to synthesize "out-of-distribution" states, and penalize the value function at these states to be below some pre-determined threshold, and using this threshold to determine whether or not states at test-time are out-of-distribution. The approach is tested on a synthetic Atari suite.