Year: 2021

Abstract:
Deep reinforcement learning (RL) has achieved remarkable success in known environments where the agents are trained, yet the agents do not necessarily know what they don’t know. In particular, RL agents deployed in the open world are naturally subject to environmental shifts and encounter unknown out-of-distribution (OOD) states—i.e., states from outside the training environment. Currently, the study of handling OOD states in the RL environment remains underexplored. This paper bridges this critical gap by proposing and exploring an unknown-aware RL framework, which improves the safety and reliability of deep Q-learning. Our key idea is to regularize the training of Q-learning so that OOD states will have higher OOD uncertainty, while in-distribution states will have lower OOD uncertainty; therefore making them distinguishable. This is in contrast with vanilla Q-learning which does not take into account unknowns during training. Furthermore, we provide theoretical guarantees that our method can improve OOD uncertainty estimation while ensuring the convergence performance of the in-distribution environment. Empirically, we demonstrate state-of-the-art performance on six diverse environments, achieving near-optimal OOD detection performance.

heading: 1 INTRODUCTION
text: 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.

heading: 2 PRELIMINARIES AND PROBLEM STATEMENT

heading: 2.1 PRELIMINARIES: REINFORCEMENT LEARNING AND Q-LEARNING
text: 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 θ.

heading: 2.2 OUT-OF-DISTRIBUTION DETECTION IN REINFORCEMENT LEARNING
text: 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.

heading: 3 METHOD
text: 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.

heading: 3.1 A MOTIVATING EXAMPLE
text: 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 cannot 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).

heading: 3.2 UNKNOWN-AWARE DEEP Q-LEARNING
text: Energy interpretation of Q-learning. In Q-learning, the agent collects data from the experience and updates Q(s, a) for each state-action pair. We can define the probability2 p(a|s) that a is the chosen action given the state s:
p(a|s) = e Q(s,a)
Σ a∈A
eQ(s,a) . (4)
We can rewrite p(a|s) as:
p(a|s) = e Q(s,a)
Σ a∈A
eQ(s,a) = p(s, a) p(s) . (5)
Directly estimating log p(s) can be computationally intractable as it requires sampling from the entire state space S. On the other hand, the the log of denominator, also known as the negative free energy3 Eneg(s) = log Σ
a∈A eQ(s,a), can be easily computed from the Q values. Note that Eneg is
proportional to log p(s) with some unknown factor that may depend on the state s.
The need for regularization. IfEneg(s) is smaller for OOD states than ID states, we can useEneg(s) for OOD detection. However, this is not guaranteed since the TD error only focuses on the Q values for ID states, Q values for OOD states can be unbounded (see Figure 1 (b)(c)) and thus produce incorrectly high Eneg(s) for OOD states. As a result, Eneg(s) from the vanilla Q-learning cannot be used as an indicator of OOD states, which motivates our regularization in the following.
Energy regularization for Q-learning. We propose regularizing Q values of OOD states, which lowersEneg(s) for OOD states while still optimizing the performance on ID states using the standard TD error. Recall that the original TD error as the loss in Q-learning is given by:
E(s,a,r,s′)∈H[TD(s, a, r, s′)]2 = E(s,a,r,s′)∈H[(r + γmax a′ Q(s′, a′)−Q(s, a))2], (6)
whereH = {(s, a, r, s′)} is an off-policy experience replay set. To regularize the Q values of OOD states, we add a regularization term to the original TD error, which results in the following loss:
L = E(s,a,r,s′)∈H[(r + γmax a′
Q(s′, a′)−Q(s, a))2]
+ λEs̃∈U [(max(log Σ ã∈A
eQ(s̃,ã) − Eout, 0))2], (7)
where U is the set of OOD states for training, and Eout is the energy margin hyperparameter. The loss pulls the energy of OOD states to be lower than Eout. Details of OOD states for training are provided in Section 5, where we show that such states can be conveniently generated using random noises without additional datasets or environments. During inference time, we use Eneg(s) as the uncertainty score for OOD detection, where higher Eneg(s) is classified as ID and vice versa. For the hyperparameter Eout, a natural option is to choose Eout smaller than the lower bound of indistribution Eneg(s), and we also ablate the effect of hyperparameter in Section 5.3 which supports this option.
2Note that one cannot confuse p(a|s) with the policy (which is denoted as π(a|s)). Here we only use p(a|s) to denote probability in the energy-based model to facilitate the interpretation. In Q-learning, π(a|s) is chosen to optimize the Q values (Watkins & Dayan, 1992).
3We refer readers to LeCun et al. (2006) for an extended background on the definition of free energy.

heading: 4 CONVERGENCE GUARANTEE
text: In this section, we provide theoretical understandings of why our proposed regularization can improve OOD uncertainty estimation while ensuring the ID convergence performance of the Qlearning algorithm. For completeness, we discuss two cases separately: tabular Q-learning (Section 4.1) and deep Q-networks (Section 4.2).

heading: 4.1 THE TABULAR CASE
text: We first provide the guarantees under the tabular case, where the Q values for different state-action pairs (s, a) are independently stored in a Q table. We expand the state space and Q table to include the OOD state. We show that we can regularize Q values for OOD states without affecting the convergence of in-distribution Q∗.
Lemma 1 In the tabular case, the global minimizer of the loss function L presented in equation (7) still satisfies Q∗ for (s, a, r, s′) ∈ H, and satisfies Eneg(s̃) ≤ Eout, ∀s̃ ∈ U .
Proof. We can analyze the partial derivative of in-distribution valueQ(s, a), Q(s′, a′) andQ(s̃, ã) value of OOD states respectively:
∂L ∂Q(s, a) = −2P{(s, a, r, s′) ∈ H}(r + γmax a′ Q(s′, a′)−Q(s, a)), (8)
∂L ∂Q(s′, a′) = 2γP{(s, a, r, s′) ∈ H}(r + γmax a′ Q(s′, a′)−Q(s, a)), (9)
and ∂L
∂Q(s̃, ã) = 2λP{s̃ ∈ U}max(log Σ ã∈A eQ(s̃,ã) − Eout, 0)(
eQ(s̃,ã)
Σ ã∈A
eQ(s̃,ã) ). (10)
The loss function is convex for each Q value as independent inputs. To find the global minimum, we only need equation (8), (9) and (10) to be zero. Note that all probabilities above are non-zero. For equation (8) and (9) to be zero, we have that Q(s, a) = r + γmaxa′ Q(s′, a′), which means in-distribution Q values satisfy the Bellman equation and thus reach Q∗. For equation (10) to be zero, we have Eneg(s̃) ≤ Eout,∀s̃ ∈ U . Remark. To see how the Q values of OOD states are regularized, we can look at equation (10). If the energy is larger than Eout, it would be penalized according to the scale of p(a|s). This results in two effects: first, the Q values of OOD states would be pushed down; second, the distribution p(·|s) would be encouraged towards a uniform distribution, which is desirable for the OOD state since the agent should be “confused” about which action to choose under OOD states.
The tabular case is an extreme example of function approximation with perfect localization, i.e., each Q(s, a) can be independently optimized without affecting each other. Next, we proceed to the DQN case, where all Q values are parameterized by a neural network.

heading: 4.2 THE DQN CASE
text: We now consider the DQN case, where loss in equation (7) can be written in the following parameterized form: L(θ) = T (θ) + λEs̃∈U [R(θ)], where
T (θ) = E(s,a,r,s′)∈H[(r + γmax a′
Qθ(s ′, a′)−Qθ(s, a))2] (11)
R(θ) = [max(log Σ ã∈A
eQθ(s̃,ã) − Eout, 0)]2. (12)
The next lemma shows that a local optimum of T (θ) under certain conditions is also a local minimizer for L(θ) (using energy regularization). In other words, we can regularize Q values for OOD states without affecting the convergence of in-distribution Q∗.
Lemma 2 Assume that L(θ) and T (θ) are both twice differentiable functions. Let θ∗ be a local minimizer of T (θ) that satisfies the sufficient conditions:
∇θ∗T (θ∗) = 0, ∇2θ∗T (θ∗) 0.
(13)

heading: If θ∗ also satisfies the additional condition below:
text: max(log Σ ã∈A
eQθ∗ (s̃,ã) − Eout, 0) = 0,∀s̃ ∈ U , (14)
then θ∗ is also a local minimizer of L(θ).
Proof. We can derive the first order derivative of the loss: ∇θL(θ) = ∇θT (θ)+λEs̃∈U [∇θR(θ)].
Let y(θ) = ∇θ max(log Σ ã∈A eQθ(s̃,ã) − Eout, 0). We have
∇θR(θ) = 2λmax(log Σ ã∈A eQθ(s̃,ã) − Eout, 0)y(θ). (15)
Plugging in the condition in equation (14), we have∇θ∗R(θ∗) = 0. So∇θ∗L(θ∗) = 0. For the second derivative,∇2θL(θ) = ∇2θT (θ) + λEs̃∈U [∇2θR(θ)]. For ∇2θR(θ) we have
∇2θR(θ) = 2λmax(log Σ ã∈A eQθ(s̃,ã) − Eout, 0)∇θy(θ) + 2λy(θ)y(θ)T . (16)
Notice that y(θ)y(θ)T is positive semidefinite. Plugging in the condition in equation (14), we have ∇2θ∗R(θ∗) < 0. Combined with ∇2θ∗T (θ∗) 0, we have ∇2θ∗L(θ∗) 0. Since L(θ) is twice differentiable and θ∗ satisfies∇θ∗L(θ∗) = 0 and∇2θ∗L(θ∗) 0, θ∗ is a local optimizer of L(θ). Remark. If we only consider the sufficient condition for a local optimizer of T (θ) (equation (13)), leaving the OOD Q values unbounded, there are multiple solutions with the same in-distribution Q values and different OOD Q values. From universal approximation theorems, the more hidden units the neural network has, the more representation power it possesses. It is therefore easy to find θ∗ that satisfies both equation (13) and (14). A visualization of such solutions can be found in Figure 1 (b) and (c), showing that our regularization on DQN can only affect OODQ values without affecting the convergence of in-distribution values.

heading: 5 EXPERIMENTS
text: In this section, we describe our experimental setup (Section 5.1) and demonstrate the effectiveness of our method on several OOD evaluation tasks (Section 5.2). We also conduct extensive ablation studies and comparative analyses that lead to an improved understanding of our method (Section 5.3).

heading: 5.1 SETUP
text: Environments. To systematically evaluate our approach, we benchmark on a variety of environments, including Atari games (with discrete state space) and a classic control task (with continuous state space) in OpenAI gym (Brockman et al., 2016). These environments allow assessing OOD states in a controlled and reproducible way. For Atari games, we choose five most commonly used environments: Pong, Breakout, Seaquest, Qbert and RoadRunner. We provide sample pictures of the states in each environment in Appendix E. In addition, we also evaluate a classic control task, Cartpole-v0 from OpenAI gym (Brockman et al., 2016), where the goal is to prevent a pole from falling over a moving cart. Due to space constraints, we present results on Atari games in the main paper (Section 5.2) and we refer readers to the Appendix Section F for results on Cartpole, where our method consistently displays strong performance.
Training. Our method can be directly compatible with existing Q-learning algorithms. In particular, we consider two commonly used methods: vanilla DQN(Mnih et al., 2015) and QR-DQN(Dabney et al., 2018b). We add the regularization term to the loss function in training as shown in equation 7. The regularization weight λ = 0.01 and Eout = 0 for DQN and λ = 0.001 and Eout = 0 for QRDQN in all environments. Hyperparameters for DQN and QR-DQN can be found in Appendix A,
and we maintain the same set of DQN parameters for both vanilla DQN and our method. We train the agent for 100,000 steps for Cartpole, and 10 million steps for Atari games.
Baselines. We compare our approach with various competitive baselines in literature: the first group of baselines are vanilla value-based learning algorithms: A2C (Mnih et al., 2016), PPO (Schulman et al., 2017), and vanilla DQN (Mnih et al., 2015). Another baseline is the ensemble-based Qlearning algorithm Bootstrapped DQN (Osband et al., 2016; 2018), which is used for OOD detection in Sedlmeier et al. (2019) and Hoel et al. (2020). Specifically, we use the variance of Q values for OOD score as in Sedlmeier et al. (2019) and Hoel et al. (2020). Specifically, the detection score is based on −Var[Q(s, amax)], where amax = arg maxa∈AQ(s, a) is the action performed after convergence of training. The variance is estimated on an ensemble of 10 neural networks, same as in Sedlmeier et al. (2019) and Hoel et al. (2020). For value-based learning algorithms, we use Eneg(s) as the score for OOD detection in vanilla DQN, and V (s) in A2C and PPO since they maintain value function instead of Q function during learning4.
Evaluation metrics. We measure the following metrics that are commonly used for OOD detection: (1) the false positive rate (FPR95) of OOD examples when the true positive rate of in-distribution examples is at 95%, and (2) the area under the receiver operating characteristic curve (AUROC). For ID evaluation, we also report the cumulative reward of the agent.

heading: 5.2 EXPERIMENTAL RESULTS
text: OOD states for training and test. For agents trained on each of the 5 Atari environments, we use states generated from the other four environments as test-time OOD states since the states of Atari games are of the same shape and pixel range. For neural network inputs, observations are preprocessed into W ×W × 4 grayscale frames as in Mnih et al. (2015). To create the training-time OOD states s̃ ∈ U , we randomly crop a square of the ID states and replace the pixels with uniform
4The entropy in A2C and PPO can also be used for OOD detection. However, we show in Appendix G that they do not produce desirable OOD detection results.
random pixels in [0, 255]. The cropping width is selected uniformly between [α · W,W ], where α = 1/4. An ablation study on the effect of cropping size and location is provided in Section 5.3.
Detection results. In Table 1, we show comparisons in a series of high-dimensional Atari environments with image inputs. Our method outperforms vanilla DQN, A2C, PPO, and ensemble-based method (Bootstrapped DQN). Moreover, Bootstrapped DQN requires deploying an ensemble of RL agents to estimate the uncertainty, which can be computationally expensive as the inference time increases with ensemble size. In contrast, our method does not incur any additional computation overhead during inference time compared with vanilla DQN. In all environments, our method achieves near-optimal performance (e.g., 0 FPR95). Additional results on QR-DQN are in Appendix C, where our method shows consistent improvement over vanilla QR-DQN.
Convergence performance and learning curves. We further show that our method improves OOD detection while ensuring the convergence performance for the in-distribution environment. Our empirical observation justifies the theoretical insights in Section 4.2, where we show regularizing the Q values of OOD states does not affect the convergence performance of ID. Specifically, Figure 2 shows the learning curves for each Atari environment. We contrast the learning curves of vanilla DQN and our regularized version since they share the same policy and only differ in the regularization term. In most environments like Pong, the convergence performance is comparable before and after regularization. In Seaquest where the vanilla DQN gets stuck in a suboptimal policy, our regularization helps to escape from the suboptimal and converges faster to a better policy, which shows the effect that correcting the wrong extrapolation in OOD spaces can potentially help to learn the correct ID Q values. The final rewards of the ID environment can be found in Table 2, where our method shows comparable or better results than vanilla DQN. Our method has a similar regularization effect on QR-DQN that displays similar or better convergence. Learning curves of QR-DQN and energy-regularized QR-DQN, together with final rewards can be found in Appendix D.

heading: 5.3 ABLATION STUDY ON HYPERPARAMETERS IN REGULARIZATION
text: In this subsection, we ablate the effect of each hyperparameter in our regularization: margin Eout, regularization weight λ, crop width and location, and noise type. For this ablation, we use Pong as the ID, remainder 4 environments as OOD, and another Atari game SpaceInvader for validation. Below we summarize our key findings. The complete results can be found in Appendix B.
(1) Energy margin Eout: We vary the values of Eout in [−10, 10]. Note that all rewards encountered in Atari games are above zero, the lower bound of energy value for ID states is 0. We test values both higher and lower than 0 to show the effect comprehensively. As shown in Table 4 in Appendix B, our method is not sensitive to the hyperparameter as long as Eout ≤ 0. In particular, positive Eout leads to less distinguishable energy distributions between ID and OOD states.
(2) Regularization weight λ: We vary the weights from 0.001 to 0.5, where we show the OOD detection results are not sensitive to the weight parameter (see Table 5 in Appendix).
(3) Cropping size: As we mentioned in Section 5.2, we crop a random square on the original ID input of size W × W and replace it with random noise to construct OOD states for training. Here the crop width and height are randomly chosen from [α ·W,W ]. We ablate on the effect of α ∈ [1/8, 2/8, ..., 7/8]. As seen in Table 6 in Appendix, the results are not very sensitive to α except for when it is too small. It is expected since such states can be similar to ID under a small α.
(4) Cropping location: As shown in Table 7 in Appendix, random cropping location yields the most diverse OOD training states. We contrast this with alternatives using fixed locations, and show that random location leads to better performance due to more data diversity.
(5) Noise type: We evaluate different types of noise: uniform and Gaussian. For Gaussian noise, we addN (0, 1) i.i.d. noise to the normalized image. Table 8 in Appendix show that both ID reward and OOD detection performance are not sensitive to noise types.

heading: 6 RELATED WORKS
text: OOD uncertainty estimation in supervised learning. The phenomenon of neural networks’ overconfidence in out-of-distribution data is revealed by Nguyen et al. (2015). Early works attempt to improve the OOD uncertainty estimation by proposing the ODIN score (Liang et al., 2018), ensemble method (Lakshminarayanan et al., 2017), and Mahalanobis distance-based confidence score (Lee et al., 2018). Recent work by Liu et al. proposed using an energy score for OOD uncertainty estimation, which demonstrated advantages over the softmax confidence score. Different from the classification model, our work focuses on an underexplored setting Q-learning, where the Q values are regressed according to the TD error, rather than cross-entropy loss which does not bound the exact logit values. In addition, we provide new theoretical guarantees which explain rigorously that our training objective does not affect the convergence of the Q-learning (for both tabular and DQN cases). Moreover, while works in the supervised learning setting may require auxiliary outlier data sets for model regularization (Hendrycks et al., 2018; Mohseni et al., 2020), our work shows that a more general and flexible solution of generating random OOD states can be effective in RL.
Uncertainty estimation in deep Q learning. The problem of OOD detection in RL remains largely untapped. Recently, Sedlmeier et al. (2019); Hoel et al. (2020) employed ensemble-based uncertainty estimation for OOD detection, which serves as a strong baseline. Variational inference methods in deep Q-learning as epistemic uncertainty estimation is also explored for OOD detection: Sedlmeier et al. (2019) compared MC-dropout (Kendall & Gal, 2017; Gal et al., 2017) with DQN and show that MC-dropout does not yield viable results. For this reason, we omit it in the baselines. Other works in epistemic estimation for Q-learning, such as Bayesian modeling (Azizzadenesheli et al., 2018), information-directed sampling (Nikolov et al., 2018) and randomized value function (Touati et al., 2020) also do not fundamentally mitigate the issue of incorrectly high Q values in the presence of OOD inputs in Section 3.1. A separate line of works quantifies the aleatoric uncertainty in Q-learning (Tamar et al., 2016; Bellemare et al., 2017; Dabney et al., 2018a). However, aleatoric uncertainty is not suitable for OOD detection, since it focuses on the internal distribution of the ID environment. Our method is also orthogonal and complementary to methods that focus on improving the ID performance of Q-learning (including QR-DQN that we considered).
OOD actions in offline RL. In offline RL, the agent is trained with a constrained set of historical data instead of interacting with the real environment. Current literature focuses on the issue of action distribution shift between the policy used to obtain the historical data and the learned policy, which results in bootstrapping errors and leads to failure in training. Various methods have been proposed to mitigate this issue (Kumar et al., 2019; 2020; Argenson & Dulac-Arnold, 2020; Kostrikov et al., 2021). Our focus is different, where we detect distribution shifts in the state space in deployment, rather than shifts in action space during offline training.

heading: 7 CONCLUSION AND OUTLOOK
text: In this work, we propose an unknown-aware deep Q-learning framework for detecting out-ofdistribution states in open-world RL deployment—an important yet underexplored problem. We show that energy regularization during Q-learning can effectively mitigate the issue of incorrectly high extrapolation of Q values for OOD states, and as a result, lead to strong separability between ID and OOD states in deployment. We provide theoretical guarantees, provably showing that our method can ensure the convergence performance of the ID environment while improving OOD detection. Empirically, our method establishes state-of-the-art OOD detection performance on diverse environments including Atari games and classic control. We hope our work will increase the research awareness and inspire future research toward a broader view of OOD uncertainty estimation for deep reinforcement learning.
REPRODUCIBILITY STATEMENT
The authors of the paper are committed to reproducible research. Below we summarize our efforts that facilitate reproducible results:
1. Environments. We use publicly available RL environments from OpenAI gym, which are described in detail in Section 5.1.
2. Assumption and proof. The complete proof of our theoretical contribution is provided in Section 4, which supports our theoretical claims made in Lemma 1 and Lemma 2.
3. RL training. We follow the original implementation in Raffin (2020) for DQN and QRDQN. Details of training configurations are fully summarized in Appendix A (Table 3).
4. Hyperparameters of our method. The hyperparameters used in energy regularization are described in Section 5.1. We also ablate the effect of each hyperparameter in our method, and provide results in Section 5.3 and Appendix B.
5. Baselines. The description of baseline methods and hyperparameters are specified in Section 5.1 and Appendix A.
6. Open source code. The codebase will be publicly released for reproducible research. As a reference, we will also attach the code in OpenReview (visible to ACs and reviewers on the discussion forum).

heading: A HYPERPARAMETERS OF RL TRAINING
text: We follow the original implementation in Raffin (2020) for DQN and QR-DQN. We maintain the same hyperparameters from their implementation for both original vs. our regularized networks. Details of hyperparameters used in DQN and QR-DQN are as follows:
For the ensemble-based method, we maintain the common hyperparameter from DQN, and follow https://github.com/johannah/bootstrap_dqn for hyperparameters of Bootstrapped DQN. For A2C and PPO, we use the same hyperparameters as Raffin (2020).

heading: B FULL ABLATION RESULTS ON HYPERPARAMETERS IN REGULARIZATION
text: Here we present the full ablation results and final rewards of ID performance. The training environment is Pong. While ablating each hyperparameter, we keep others unchanged. Note that for Pong, the highest cumulative reward for ID evaluation is 21.
(1) Ablation on energy margin Eout: We vary the values of Eout in [−10, 10]. Note that all rewards encountered in Atari games are above zero, the lower bound of energy value for ID states is 0. We vary Eout to be both higher and lower than 0 and show the effect comprehensively. As shown in Table 4, our method is not sensitive to the hyperparameter as long asEout ≤ 0. In particular, positive Eout leads to less distinguishable energy distributions between ID and OOD states.
(2) Ablation on regularization weight λ: We vary the weights from 0.001 to 0.5, where we show the OOD detection results are not sensitive to the weight parameter (see Table 5).
5MLP with 2 hidden layers; 256 units for each hidden layer. 6MLP with 2 hidden layers; 256 units for each hidden layer; number of quantiles = 25. 7The same CNN structure as in Mnih et al. (2015). 8The same CNN structure as in Mnih et al. (2015); number of quantiles = 200.
(3) Ablation on crop width: As we mentioned in Section 5, we crop a random square on the original ID input of size W × W and replace it with random noise to construct OOD states for training. Here the crop width and height are randomly chosen from [α ·W,W ]. We ablate on the effect of α ∈ [1/8, 2/8, ..., 7/8]. As seen in Table 6, the results are not very sensitive to α except for when it is too small (e.g., 1/8). It is expected since the resulting input can be similar to ID under a small α.
(4) Ablation on crop location: As shown in Table 7, random cropping location yields the most diverse OOD training states. We contrast this with alternatives using fixed locations, and show that random location leads to better performance overall.
(5) Ablation on noise types: We evaluate different types of noise: uniform and Gaussian. For Gaussian noise, we add N (0, 1) i.i.d. noise to the normalized image. Table 8 shows that both ID reward and OOD detection performance are not sensitive to noise types.

heading: C OOD DETECTION RESULTS ON QR-DQN IN ATARI GAMES
text: Here we present the OOD detection results in Table 9 when we apply the regularization term to another common Q-learning algorithm QR-DQN (Dabney et al., 2018b), where our method also shows significant improvement with near-optimal results.

heading: D EVALUATION ON FINAL REWARD AND FULL LEARNING CURVES IN ATARI GAMES
text: Here we present the final rewards of QRDQN and our regularized DQN in Table 10, where our method also maintains a similar ID performance compared to vanilla QR-DQN.
We present the learning curves of QR-DQN and our regularized version in five Atari games in Figure 3. Combined with final rewards in Table 10, we show that our regularization also maintains the similar ID reward while improving the OOD detection performance, or even help the agent to achieve better ID performance (e.g., Seaquest).

heading: E SAMPLE PICTURES OF STATES FROM ATARI GAMES
text: (a) Pong (b) Breakout (c) Seaquest (d) Qbert (e) RoadRunner

heading: F CLASSIC CONTROL TASK: CARTPOLE
text: OOD states for training and test. The state space of Cartpole is a four-dimension vector, each specifying the physical property with the following meaning and range: Cart position [−4.8, 4.8],
Cart Velocity [-inf , inf], Pole Angle [−41.8, 41.8], Tip Velocity [-inf , inf]. For training, we generate independent zero-mean random values from Gaussian distributions, then clamp them within the ranges above9. This breaks the physical constraint of each dimension, which makes these states OOD. For the test, we consider a different type of OOD: we first generate a random value within [−5, 5], and apply the same value in each dimension to build the OOD states. This generates states that are outside the ID environments.
Detection and Convergence Results. Our training is based on the same architecture as the vanilla DQN, both of which are trained for the same number of steps. Table 11 shows the OOD detection and convergence result. Overall, our regularization can outperform the uncertainty-based algorithm
9The standard deviations used for each dimension are 2.0, 5.0, 0.2, 5.0.
in OOD detection while achieving the same final convergence performance (200 is the highest possible cumulative reward). We also show that our method is effective on an alternative Q-learning algorithm QR-DQN.

heading: G ENTROPY IN A2C AND PPO
text: The entropy in the policy used in A2C and PPO can also be used for OOD detection. However, the entropy extracted from policy generally does not yield desirable results in OOD detection. We take Atari games as an example and present the OOD test results below.

Review summaries:
This paper proposes a new method based on Q-learning that is able to detect out-of-distribution states during deployment time. In order to accomplish this, an energy-based regulizeer has been introduced in which it pushes down Q-values for out-of-distribution states during training time. In order to create out-of-distribution states during training, this paper utilizes some heuristic such as cropping applied to random locations to input images. To evaluate the proposed method, this paper utilizes 4 Atari games and Cartpole environment. 

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. 

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.