The Uncertainty Bellman Equation and Exploration

Introduction

We consider the reinforcement learning (RL) problem of an agent interacting with its environment to maximize cumulative rewards over time (Sutton & Barto, 1998). We model the environment as a Markov decision process (MDP), but where the agent is initially uncertain of the true dynamics and mean rewards of the MDP (Bellman, 1957; Bertsekas, 2005). At each time-step, the agent performs an action, receives a reward, and moves to the next state; from these data it can learn which actions lead to higher payoffs. This leads to the exploration versus exploitation trade-off: Should the agent investigate poorly understood states and actions to improve future performance or instead take actions that maximize rewards given its current knowledge?

Separating estimation and control in RL via ‘greedy’ algorithms can lead to premature and suboptimal exploitation. To offset this, the majority of practical implementations introduce some random noise or dithering into their action selection (such as $\epsilon$-greedy). These algorithms will eventually explore every reachable state and action infinitely often, but can take exponentially long to learn the optimal policy (Kakade, 2003). By contrast, for any set of prior beliefs the optimal exploration policy can be computed directly by dynamic programming in the Bayesian belief space. However, this approach can be computationally intractable for even very small problems (Guez et al., 2012) while direct computational approximations can fail spectacularly badly (Munos, 2014).

For this reason, most provably-efficient approaches to reinforcement learning rely upon the optimism in the face of uncertainty (OFU) principle (Lai & Robbins, 1985; Kearns & Singh, 2002; Brafman & Tennenholtz, 2002). These algorithms give a bonus to poorly-understood states and actions and subsequently follow the policy that is optimal for this augmented optimistic MDP. This optimism incentivizes exploration but, as the agent learns more about the environment, the scale of the bonus should decrease and the agent’s performance should approach optimality. At a high level these approaches to OFU-RL build up confidence sets that contain the true MDP with high probability (Strehl & Littman, 2004; Lattimore & Hutter, 2012; Jaksch et al., 2010). These techniques can provide performance guarantees that are ‘near-optimal’ in terms of the problem parameters. However, apart from the simple ‘multi-armed bandit’ setting with only one state, there is still a significant gap between the upper and lower bounds for these algorithms (Lattimore, 2016; Jaksch et al., 2010; Osband & Van Roy, 2016).

One inefficiency in these algorithms is that, although the concentration may be tight at each state and action independently, the combination of simultaneously optimistic estimates may result in an extremely over-optimistic estimate for the MDP as a whole (Osband & Van Roy, 2017). Other works have suggested that a Bayesian posterior sampling approach may not suffer from these inefficiencies and can lead to performance improvements over OFU methods (Strens, 2000; Osband et al., 2013; Grande et al., 2014). In this paper we explore a related approach that harnesses the simple relationship of the uncertainty Bellman equation (UBE), where we define uncertainty to be the variance of the Bayesian posterior of the Q-values of a policy conditioned on the data the agent has collected, in a sense similar to the parametric variance of Mannor et al. (2007). Intuitively speaking, if the agent has high uncertainty (as measured by high posterior variance) in a region of the state-space then it should explore there, in order to get a better estimate of those Q-values. We show that, just as the Bellman equation relates the value of a policy beyond a single time-step, so too does the uncertainty Bellman equation propagate uncertainty values over multiple time-steps, thereby facilitating ‘deep exploration’ (Osband et al., 2017; Moerland et al., 2017).

The benefit of our approach (which learns the solution to the UBE and uses this to guide exploration) is that we can harness the existing machinery for deep reinforcement learning with minimal change to existing network architectures. The resulting algorithm shares a connection to the existing literature of both OFU and intrinsic motivation (Singh et al., 2004; Schmidhuber, 2009; White & White, 2010). Recent work has further connected these approaches through the notion of ‘pseudo-count’ (Bellemare et al., 2016; Ostrovski et al., 2017), a generalization of the number of visits to a state and action. Rather than adding a pseudo-count based bonus to the rewards, our work builds upon the idea that the more fundamental quantity is the uncertainty of the value function and that naively compounding count-based bonuses may lead to inefficient confidence sets (Osband & Van Roy, 2017). The key difference is that the UBE compounds the variances at each step, rather than standard deviation.

The observation that the higher moments of a value function also satisfy a form of Bellman equation is not new and has been observed by some of the early papers on the subject (Sobel, 1982). Unlike most prior work, we focus upon the epistemic uncertainty over the value function, as captured by the Bayesian posterior, i.e., the uncertainty due to estimating a parameter using a finite amount of data, rather than the higher moments of the reward-to-go (Lattimore & Hutter, 2012; Azar et al., 2012; Mannor & Tsitsiklis, 2011; Bellemare et al., 2017). For application to rich environments with complex generalization we will use a deep learning architecture to learn a solution to the UBE, in the style of (Tamar et al., 2016).

Problem formulation

The action-value, or Q-value, at time step $l$ of a particular state under policy $\pi$ is the expected total return from taking that action at that state and following $\pi$ thereafter, i.e., $Q^{l}_{sa}=\mathop{\bf E{}}\left[\sum_{h=l}^{H}r^{h}\mid s^{l}=s,a^{l}=a,\pi\right]$ (we suppress the dependence on $\pi$ in this notation). The value of state $s$ under policy $\pi$ at time-step $h$, $V^{h}(s)=\mathop{\bf E{}}_{a\sim\pi^{h}_{s}}Q^{h}_{sa}$, is the expected total discounted return of policy $\pi$ from state $s$.

The Bellman operator $\mathcal{T}^{h}$ for policy $\pi$ at each time-step $h$ relates the value at each time-step to the value at subsequent time-steps via dynamic programming (Bellman, 1957),

for all $(s,a)$, where $\mu=\mathop{\bf E{}}r$ is the mean reward. The Q-values are the unique fixed point of equation (1), i.e., the solution to $\mathcal{T}^{h}Q^{h+1}=Q^{h}$ for $h=1,\ldots,H$, where $Q^{H+1}$ is defined to be zero. Several reinforcement learning algorithms have been designed around minimizing the residual of equation (1) to propagate knowledge of immediate rewards to long term value (Sutton, 1988; Watkins, 1989). In the next section we examine a similar relationship for propagating the uncertainties of the Q-values, we call this relationship the uncertainty Bellman equation.

The uncertainty Bellman equation

In this section we derive a Bellman-style relationship that propagates the uncertainty (variance) of the Bayesian posterior distribution over Q-values across multiple time-steps. Propagating the potential value of exploration over many time-steps, or deep exploration, is important for statistically efficient RL (Kearns & Singh, 2002; Osband et al., 2017). Our main result, which we state in Theorem 1, is based upon nothing more than the dynamic programming recursion in equation (1) and some crude upper bounds of several intermediate terms. We will show that even in very simple settings this approach can result in well-calibrated uncertainty estimates where common count-based bonuses are inefficient (Osband & Van Roy, 2017).

We consider the Bayesian case, where we have priors over the mean reward $\mu$ and the transition probability matrix $P$ which we denote by $\phi_{\mu}$ and $\phi_{P}$ respectively. We collect some data generated by the policy $\pi$ and use it to derive posterior distributions over $\mu$ and $P$, given the data. We denote by $\mathcal{F}_{t}$ the sigma-algebra generated by all the history up to episode $t$ (e.g., all the rewards, actions, and state transitions for all episodes), and let the posteriors over the mean reward and transition probabilities be denoted by $\phi_{\mu|\mathcal{F}_{t}}$ and $\phi_{P|\mathcal{F}_{t}}$ respectively. If we sample $\hat{\mu}\sim\phi_{\mu|\mathcal{F}_{t}}$ and $\hat{P}\sim\phi_{P|\mathcal{F}_{t}}$, then the resulting Q-values that satisfy

where $\hat{Q}^{H+1}=0$, are a sample from the implicit posterior over Q-values, conditioned on the history $\mathcal{F}_{t}$ (Strens, 2000). In this section we compute a bound on the variance (uncertainty) of the random variable $\hat{Q}$. For the analysis we will require some additional assumptions.

This assumption means that the agent cannot revisit a state within the same episode, and is a common assumption in the literature (Osband et al., 2014). Note that any finite horizon MDP that doesn’t satisfy this assumption can be converted into one that does by ‘unrolling’ the MDP so that each state $s$ is replaced by $H$ copies of the state, one for each step in the episode.

denote the variance of $x$ conditioned on $\mathcal{F}_{t}$. Under the assumptions listed above, the variance of the Q-values under the posterior satisfies the Bellman inequality

for all $(s,a)$ and $h=1,\ldots,H$, where ${\textstyle{\bf var}_{t}}\hat{Q}^{H+1}=0$ and where we call $\nu^{h}_{sa}$ the local uncertainty at $(s,a)$, and it is given by

We refer to $\nu$ in the above lemma as the local uncertainty since it depends only on locally available quantities, and so can be calculated (in principle) at each state-action independently. Note that even though $\mathop{\bf E{}}(\hat{P}^{h}_{s^{\prime}sa}|\mathcal{F}_{t})$ appears in the denominator above, the local uncertainty is bounded, since for any random variable $X$ on $(0,1]$ we have

and if for any $s^{\prime}$ we have that $\mathop{\bf E{}}(\hat{P}^{h}_{s^{\prime}sa}|\mathcal{F}_{t})=0$ Markov’s inequality implies that $P^{h}_{s^{\prime}sa}=0$ and so we can just remove that term from the sum since it contributes no uncertainty.

With this lemma we are ready to prove our main theorem.

Under assumptions 1 and 2, for any policy $\pi$ there exists a unique $u$ that satisfies the uncertainty Bellman equation

for all $(s,a)$ and $h=1,\ldots,H$, where $u^{H+1}=0$, and furthermore $u\geq{\textstyle{\bf var}_{t}}\hat{Q}$ pointwise.

Let $\mathcal{U}^{h}$ be the Bellman operator that defines the uncertainty Bellman equation, i.e., rewrite equation (2) as

then to prove the result we use two essential properties of the Bellman operator for a fixed policy. Firstly, the solution to the Bellman equation exists and is unique, and secondly the Bellman operator is monotonically non-decreasing in its argument, i.e., if $x\geq y$ pointwise then $\mathcal{U}^{h}x\geq\mathcal{U}^{h}y$ pointwise (Bertsekas, 2005). The proof proceeds by induction; assume that for some $h$ we have ${\textstyle{\bf var}_{t}}\hat{Q}^{h+1}\leq u^{h+1}$, then we have

where we have used the fact that the variance satisfies the Bellman inequality from lemma 1, and the base case holds because ${\textstyle{\bf var}_{t}}\hat{Q}^{H+1}=u^{H+1}=0$. ∎

We conclude with a brief discussion on why the variance of the posterior is useful for exploration. If we had access to the true posterior distribution over the Q-values then we could take actions that lead to states with higher uncertainty by, for example, using Thompson sampling (Thompson, 1933; Strens, 2000), or constructing Q-values that are high probability upper bounds on the true Q-values and using the OFU principle (Kaufmann et al., 2012). However, calculating the true posterior is intractable for all but very small problems. Due to this difficulty prior work has sought to approximate the posterior distribution (Osband et al., 2017), and use that to drive exploration. In that spirit we develop another approximation of the posterior, in this case it is motivated by the Bayesian central limit theorem which states that, under some mild conditions, the posterior distribution converges to a Gaussian as the amount of data increases (Berger, 2013). With that in mind, rather than computing the full posterior we approximate it as $\mathcal{N}(\bar{Q},\mathop{\bf diag}(u))$ where $u$ is the solution to the uncertainty Bellman equation (2), and consequently is a guaranteed upper bound on the true variance of the posterior, and $\bar{Q}$ denotes the mean Q-values under the posterior at episode $t$, i.e., the unique solution to

for $h=1,\ldots,H$, and $\bar{Q}^{H+1}=0$. With this approximate posterior we can perform Thompson sampling as an exploration heuristic. Specifically, at state $s$ and time-step $h$ we select the action using

where $\zeta_{b}$ is sampled from $\mathcal{N}(0,1)$. Our goal is for the agent to explore states and actions where it has higher uncertainty. This is in contrast to the commonly used $\epsilon$-greedy (Mnih et al., 2013) and Boltzmann exploration strategies (Mnih et al., 2016; O’Donoghue et al., 2017; Haarnoja et al., 2017) which simply inject noise into the agents actions. We shall see in the experiments that our strategy can dramatically outperform these naive heuristics.

2 Comparison to traditional exploration bonus

Consider a simple decision problem with known deterministic transitions, unknown rewards, and two actions at a root node, as depicted in Figure 1. The first action leads to a single reward $r_{1}$ sampled from $\mathcal{N}(\mu_{1},\sigma^{2})$ at which point the episode terminates, and the second action leads to an chain of length $H$ consisting of states each having random reward $r_{2}$ independently sampled from $\mathcal{N}(\mu_{2}/H,\sigma^{2}/H)$.

Take the case where each action at the root has been taken $n$ times and where the uncertainty over the rewards at each state concentrates like $1/n$ (e.g., when the prior is an improper Gaussian). In this case the true uncertainty about the value of each action is identical and given by $\sigma^{2}/n$. This is also the answer we get from the uncertainty Bellman equation, since for action 1 we obtain $u_{1}=\sigma^{2}/n$ (since $\mathop{\bf var}_{t}P=0$) and for action 2 the uncertainty about the reward at each state along the chain is given by $\sigma^{2}/Hn$ and so we have $u_{2}=\sum_{h=1}^{H}\sigma^{2}/Hn=\sigma^{2}/n$.

Rather than considering the variance of the value as a whole, the majority of existing approaches to OFU provide exploration bonuses at each state and action independently and then combine these estimates via union bound. In this context, even a state of the art algorithm such as UCRL2 (Jaksch et al., 2010) would augment the rewards at each state with a bonus proportional to the standard deviation of the reward estimate at each state (Bellemare et al., 2016). For the first action this would be ${\rm ExpBonus}_{1}=\sigma/\sqrt{n}$, but for the second action this would be accumulated along the chain to be

In other words, the bonus afforded to the second action is a factor of $\sqrt{H}$ larger than the true uncertainty. The agent would have to take the second action a factor of $H$ more times than the first action in order to have the same effective bonus given to each one. If the first action was actually superior in terms of expected reward, it would take the agent far longer to discover that than an agent using the correct uncertainties to select actions. The essential issue is that, unlike the variance, the standard deviations do not obey a Bellman-style relationship.

In Figure 2 we show the results of an experiment showing this phenomenon. Action 1 had expected reward $\mu_{1}=1$, and action 2 had expected reward $\mu_{2}=0$. We set $\sigma=1$ and $H=10$, and the results are averaged over $500$ seeds. We compare two agents, one using the uncertainty Bellman equation to drive exploration and the other agent using a count-based reward bonus. Both agents take actions and use the results to update their beliefs about the value of each action. The agent using the UBE takes the action yielded by Thompson sampling as in equation (3). The exploration-bonus agent takes the action $i$ that maximizes $\hat{Q}_{i}+\beta\log(t){\rm ExpBonus}_{i}$ (the $\log(t)$ term is required to achieve a regret bound (Jaksch et al., 2010), but doesn’t materially affect the previous argument) where $\beta>0$ is a hyper-parameter chosen by a sweep and where $\hat{Q}_{i}$ is the estimate of the value of action $i$. Figure 2 shows the regret of each agent vs number of episodes. Regret measures how sub-optimal the rewards the agent has received so far are, relative to the (unknown) optimal policy, and lower regret is better (Cesa-Bianchi & Lugosi, 2006).

The agent using the uncertainty Bellman equation has well calibrated uncertainty estimates and consequently quickly figures out that the first action is better. By contrast, the exploration bonus agent takes significantly longer to determine that the first action is better due to the fact that the bonus afforded to the second action is too large, and consequently it suffers significantly higher regret.

Estimating the local uncertainty

Section 3 outlined how the uncertainty Bellman equation can be used to propagate local estimates of the variance of $\hat{Q}$ to global estimates for the uncertainty. In this section we present some pragmatic approaches to estimating the local uncertainty $\nu$ that we can then use for practical learning algorithms inspired by Theorem 1. We do not claim that these approaches are the only approaches to estimating the local uncertainty, or even that these simple approximations are in any sense the ‘best’. Investigating these choices is an important area of future research, but outside the scope of this short paper. We present a simple progression from tabular representations, to linear function approximation and then to non-linear neural network architectures.

Consider the case where the posterior over the mean rewards concentrates at least as fast the reciprocal of the visit count, i.e.,

where $\sigma_{r}$ is the variance of the reward process and $n^{h}_{sa}$ is the visit count of the agent to state $s$ and action $a$ at time-step $h$, up to episode $t$. This is the case when, for example, the rewards and the prior over the mean reward are both Gaussian. Furthermore, if we assume that the prior over the transition function is Dirichlet then it is straightforward to show that

where $|S|$ is the number of next states reachable from $s,a$. This holds since the likelihood of the transition function is a categorical distribution, which is conjugate to the Dirichlet distribution and the variance of a Dirichlet concentrates like the reciprocal of the sum of the counts of each category. Under these assumptions we can bound the local uncertainty as

In other words, the local uncertainty can be modeled under these assumptions as a constant divided by the visit count.

Linear value estimate.

for some $\beta$, which in the tabular case (i.e., where $\phi(s)=e_{s}$ and $D=|\mathcal{S}|$) is equal to $\beta^{2}/n^{h}_{sa}$, as expected.

An agent using this notion of uncertainty must maintain and update the matrix $\Sigma_{a}=(\Phi_{a}^{T}\Phi_{a})^{-1}$ as it receives new data. Given new sample $\phi$, the updated matrix $\Sigma_{a}^{+}$ is given by

by the Sherman-Morrison-Woodbury formula (Golub & Van Loan, 2012), the cost of this update is one matrix multiply and one matrix-matrix subtraction per step.

Neural networks value estimate.

If we are approximating our Q-value function using a neural network then the above analysis does not hold. However if the last layer of the network is linear, then the Q-values are approximated as $Q^{h}_{sa}=\phi(s)^{T}w_{a}$, where $w_{a}$ are the weights of the last layer associated with action $a$ and $\phi(s)$ is the output of the network up to the last layer for state $s$. In other words we can think of a neural network as learning a useful set of basis functions such that a linear combination of them approximates the Q-values. Then, if we ignore the uncertainty in the $\phi$ mapping, we can reuse the analysis for the purely linear case to derive an approximate measure of local uncertainty that might be useful in practice.

This scheme has some advantages. As the agent progresses it is learning a state representation that helps it achieve the goal of maximizing the return. The agent will learn to pay attention to small but important details (e.g., the ball in Atari ‘breakout’) and learn to ignore large but irrelevant changes (e.g., if the background suddenly changes). This is a desirable property from the point of view of using these features to drive exploration, because the states that differ only in irrelevant ways will be aliased to (roughly) the same state representation, and states that differ is small but important ways will be mapped to quite different state vectors, permitting a more task-relevant measure of uncertainty.

Deep Reinforcement Learning

Previously we proved that under certain conditions we can bound the variance of the posterior distribution of the Q-values, and we used the resulting uncertainty values to derive an exploration strategy. Here we discuss the application of that strategy to deep-RL. In this case several of the assumptions we have made to derive theorem 1 are violated. This puts us firmly in the territory of heuristic. Specifically, the MDPs we apply this to will not be directed acyclic graphs, the policy that we are estimating the uncertainty over will not be fixed, we cannot exactly compute the local uncertainty, and we won’t be solving the UBE exactly. However, empirically, we demonstrate that this heuristic can perform well in practice, despite the underlying assumptions being violated.

Our strategy involves learning the uncertainty estimates, and then using them to sample Q-values from the approximate posterior, as in equation (3). The technique is described in pseudo-code in Algorithm 1. We refer to the technique as ‘one-step’ since the uncertainty values are updated using a one-step SARSA Bellman backup, but it is easily extendable to the $n$-step case. The algorithm takes as input a neural network which has two output ‘heads’, one which is attempting to learn the optimal Q-values as normal, the other is attempting to learn the uncertainty values of the current policy (which is constantly changing). We do not allow the gradients from the uncertainty output head to flow into the trunk of the network; this ensures the Q-value estimates are not perturbed by the changing uncertainty signal. For the local uncertainty measure we use the linear basis approximation described in section 4. Algorithm 1 incorporates a discount factor $\gamma\in(0,1)$, since deep RL often uses a discount even in the purely episodic case. In this case the Q-learning update uses a $\gamma$ discount and the Uncertainty Bellman equation (2) is augmented with a $\gamma^{2}$ discount factor.

Here we present results of Algorithm (1) on the Atari suite of games (Bellemare et al., 2012), where the network is attempting to learn the Q-values as in DQN (Mnih et al., 2013, 2015) and the uncertainties simultaneously. The only change to vanilla DQN we made was to replace the $\epsilon$-greedy policy with Thompson sampling over the learned uncertainty values, where the $\beta$ constant in (3) was chosen to be $0.01$ for all games, by a parameter sweep. We used the exact same network architecture, learning rate, optimizer, pre-processing and replay scheme as described in Mnih et al. (2015). For the uncertainty sub-network we used a single fully connected hidden layer with 512 hidden units followed by the output layer. We trained the uncertainty head using a separate RMSProp optimizer (Tieleman & Hinton, 2012) with learning rate $10^{-3}$. The addition of the uncertainty head and the computation associated with it, only reduced the frame-rate compared to vanilla DQN by about 10% on a GPU, so the additional computational cost of the approach is negligible.

We compare two versions of our approach: a $1$-step method and an $n$-step method where we set $n$ to $150$. The $n$-step method accumulates the uncertainty signal over $n$ time-steps before performing an update which should lead to the uncertainty signal propagating to earlier encountered states faster, at the expense of increased variance of the signal. Note that in all cases the Q-learning update is always $1$-step; our $n$-step implementation only affects the uncertainty update.

We compare our approaches to vanilla DQN, and also to an exploration bonus intrinsic motivation approach, where the agent receives an augmented reward consisting of the extrinsic reward and the square root of the linear uncertainty in equation (4), which was scaled by a hyper-parameter chosen to be $0.1$ by a sweep. In this case a stochastic policy was still required for good performance and so we used $\epsilon$-greedy with the DQN annealing schedule.

We trained all strategies for 200M frames (about 8 days on a GPU). Each game and strategy was tested three times per method with the same hyper-parameters but with different random seeds, and all plots and scores correspond to an average over the seeds. All scores were normalized by subtracting the average score achieved by an agent that takes actions uniformly at random. Every 1M frames the agents were saved and evaluated (without learning) on 0.5M frames, where each episode is started from the random start condition described in (Mnih et al., 2015). The final scores presented correspond to first averaging the evalution score in each period across seeds, then taking the max average episodic score observed during any evalution period. Of the tested strategies the $n$-step UBE approach was the highest performer in 32 out of 57 games, the $1$-step UBE approach in 14 games, DQN in 1 game, the exploration bonus strategy in 7 games, and there were 3 ties. In Table 1 we give the mean and median normalized scores as percentage of an expert human normalized score across all games, and the number of games where the agent is ‘super-human’, for each tested algorithm. Note that the mean scores are significantly affected by a single outlier with very high score (‘Atlantis’), and therefore the median score is a better indicator of agent performance. In Figure 3 we plot the number of games at super-human performance against frames for each method, and in Figure 4 we plot the median performance across all games versus frames, where a score of $1.0$ denotes human performance. The results across all 57 games, as well as the learning curves for all 57 games, are given in the appendix.

Of particular interest is the game ‘Montezuma’s Revenge’, a notoriously difficult exploration game where no one-step algorithm has managed to learn anything useful. Our $1$-step strategy learns in 200M frames a policy that is able to consistently get about 500 points, which is the score the agent gets for picking up the first key and moving into the second room. In Figure 5 we show the learning progress of the agents for 500M frames where we set the Thompson sampling parameter slightly higher; $0.016$ instead of $0.01$ (since this game is a challenging exploration task it stands to reason that a higher exploration parameter is required). By the end of 500M frames the $n$-step agent is consistently getting around 3000 points, which is several rooms of progress. These scores are close to state-of-the-art, and are state-of-the-art for one-step methods (like DQN) to the best of our knowledge.

In the recent work by Bellemare et al. (2016), and the follow-up work by Ostrovski et al. (2017), the authors add an intrinsic motivation signal to a DQN-style agent that has been modified to use the full Monte Carlo return of the episode when learning the Q-values. Using Monte Carlo returns dramatically improves the performance of DQN in a way unrelated to exploration, and due to that change we cannot compare the numerical results directly. In order to have a point of comparison we implemented our own intrinisic motivation exploration signal, as discussed above. Similarly, we cannot compare directly to the numerical results obtained by Bootstrap DQN (Osband et al., 2016) since that agent is using Double-DQN, a variant of DQN that achieves a higher performance in a way unrelated to exploration. However, we note that our approach achieves a higher evaluation score in 27 out of the 48 games tested in the Bootstrap DQN paper despite using an inferior base DQN implementation, and it runs at a significantly lower computational and memory cost.

Conclusion

In this paper we derived a Bellman equation for the uncertainty over the Q-values of a policy. This allows an agent to propagate uncertainty across many time-steps in the same way that value propagates through time in the standard dynamic programming recursion. This uncertainty can be used by the agent to make decisions about which states and actions to explore, in order to gather more data about the environment and learn a better policy. Since the uncertainty satisfies a Bellman recursion, the agent can learn it using the same reinforcement learning machinery that has been developed for value functions. We showed that a heuristic algorithm based on this learned uncertainty can boost the performance of standard deep-RL techniques. Our technique was able to significantly improve the performance of DQN across the Atari suite of games, when compared against naive strategies like $\epsilon$-greedy.

Acknowledgments

We thank Marc Bellemare, David Silver, Koray Kavukcuoglu, Daniel Kasenberg, and Mohammad Gheshlaghi Azar for useful discussion and suggestions on the paper.

References

Appendix

For ease of exposition in this derivation we shall use the notation ${\textstyle{\mathop{\bf E{}}_{t}}}(\cdot)$ to denote the expectation of a random variable conditioned on the history $\mathcal{F}_{t}$, rather than the usual $\mathop{\bf E{}}(\cdot|\mathcal{F}_{t})$.

Recall that if we sample $\hat{\mu}$ from the posterior over the mean rewards $\phi_{\mu|\mathcal{F}_{t}}$, and $\hat{P}$ from the posterior over the transition probability matrix $\phi_{P|\mathcal{F}_{t}}$ then the Q-values that are the unique solution to

are a sample from the (implicit) posterior over Q-values, conditioned on $\mathcal{F}_{t}$ (Strens, 2000).

Using the definition of the conditional variance

where we have used the fact that $\hat{\mu}^{h}_{sa}$ is conditionally independent (conditioned on $\mathcal{F}_{t}$) of $\hat{P}^{h}_{s^{\prime}sa}\hat{Q}^{h+1}_{s^{\prime}a^{\prime}}$ because assumption 1 implies that $\hat{Q}^{h+1}_{s^{\prime}a^{\prime}}$ depends only on downstream quantities.

Now we make the assumption that ${\textstyle{\mathop{\bf E{}}_{t}}}\hat{P}^{h}_{s^{\prime}sa}>0$ for all $h,s^{\prime},s,a$. Note that this is not a restriction because if ${\textstyle{\mathop{\bf E{}}_{t}}}\hat{P}^{h}_{s^{\prime}sa}=0$, then the fact that $P$ is nonnegative combined with Markov’s inequality implies that $P^{h}_{s^{\prime}sa}=0$, in which case we can just remove that term from the sum, since it contributes no variance (or equivalently $s^{\prime}$ is not reachable after taking $s,a$). Furthermore, note that any sample must satisfy $\sum_{s^{\prime}}\hat{P}^{h}_{s^{\prime}sa}=1$ which implies that $\sum_{s^{\prime}}{\textstyle{\mathop{\bf E{}}_{t}}}\hat{P}^{h}_{s^{\prime}sa}=1$ and therefore $\sum_{s^{\prime}a^{\prime}}\pi_{s^{\prime}a^{\prime}}{\textstyle{\mathop{\bf E{}}_{t}}}\hat{P}^{h}_{s^{\prime}sa}=1$. In other words $\pi^{h}_{s^{\prime}a^{\prime}}{\textstyle{\mathop{\bf E{}}_{t}}}\hat{P}^{h}_{s^{\prime}sa}$ defines a probability distribution over $s^{\prime},a^{\prime}$. With this we can bound the second term as

by applying Jensen’s inequality to the quadratic. Combining these we have

Assumption 1 also implies that that $\hat{P}^{h}_{s^{\prime}sa}$ and $\hat{Q}^{h+1}_{s^{\prime}a^{\prime}}$ are conditionally independent, and so we can write the middle expectation in the second term as

Now using the conditional independence property again and the fact that the Q-values are bounded, as implied by assumption 2, we have

Putting this together with equation (6) we obtain

where $\nu^{h}_{sa}$ is the local uncertainty, and is given by

Atari suite scores