Title: Successor Uncertainties: exploration and uncertainty in temporal difference learning

Abstract: We consider the problem of balancing exploration and exploitation in sequential decision making problems. This trade-off naturally lends itself to probabilistic modelling. For a probabilistic approach to be effective, considering uncertainty about all immediate and long-term consequences of agent's actions is vital. An estimate of such uncertainty can be leveraged to guide exploration even in situations where the agent needs to perform a potentially long sequence of actions before reaching an under-explored area of the environment. This observation was made by the authors of the Uncertainty Bellman Equation model (O'Donoghue et al., 2018) , which explicitly considers full marginal uncertainty for each decision the agent faces. However, their model still considers a fully factorised posterior over the consequences of each action, meaning that dependencies vital for correlated long-term exploration are ignored. We go a step beyond and develop Successor Uncertainties, a probabilistic model for the state-action value function of a Markov Decision Process with a non-factorised covariance. We demonstrate how this leads to greatly improved performance on classic tabular exploration benchmarks and show strong performance of our method on a subset of Atari baselines. Overall, Successor Uncertainties provides a better probabilistic model for temporal difference learning at a similar computational cost to its predecessors.

1. Introduction:

We consider a sequential decision making problem in which an agent interacts with an unknown environment, modelled as a Markov Decision Process (MDP). The agent's goal is to learn a policy that maximises the expected cumulative reward while keeping the total number of interactions with the environment low. To do so, the agent has to balance exploration of the environment and exploitation of the information it has already acquired.

Because the agent learns sequentially through interaction with a potentially stochastic environment, the consequences of each action are not necessarily known ahead of time and this complicates striking the optimal trade-off between exploration and exploitation. A particular strategy to address this issue is maintaining a probabilistic model that reflects agent's current knowledge, or lack thereof, about the future effects of each action for any given state. Vitally, the model must incorporate uncertainty about not only the immediate but also the downstream effects of agent's actions, and reflect the varying degree of uncertainty associated with different sequences of states encountered and actions taken by the agent.

We will therefore focus on two concepts we believe to be crucial for effective exploration under uncertainty: propagation of uncertainty and incorporation of dependencies. We will say that a probabilistic model of the Q function propagates uncertainty whenever the degree of uncertainty over the Q value for each state-action pair is determined by the cumulative uncertainty of all future state-actions the agent may encounter. Notice that for any joint distribution of the Q function that propagates uncertainty, the model obtained by taking the product of its marginal distributions over Q values also propagates uncertainty. In other words, propagation of uncertainty does not guarantee that any dependencies between individual Q values, like the one implied by the Bellman equation, are captured by the model.

To this end, we will say that a model incorporates dependencies if the dependency structure of individual Q values obeys the Bellman equation. Further details are provided in section 3.

We illustrate importance of propagation of uncertainties and incorporation of dependencies in figure 1 . The agent is faced with a sequence of binary up/down choices and receives a positive reward if and only if it executes a particularly long sequence of uninterrupted up actions. In such a scenario, purely random exploration is extremely ineffective. Similarly, exploration which considers just immediate consequences of actions also performs poorlyshortsightedly reducing one-step uncertainty leads to exploration equivalent in expectation to that of a uniform exploration policy (more on this in section 6). In contrast, a model which propagates uncertainty and incorporates dependencies can be used to guide the agent into under-explored areas irrespective of their distance from the current state, resulting in more effective exploration of the environment. Whilst simple, this MDP captures the key features of classic exploration benchmarks like Montezuma's Revenge (Bellemare et al., 2013) , where exploration is hindered by sparse reward and low connectivity of the state-space. 3. Empirical investigation demonstrating how lack of either propagation of uncertainty or incorporation of dependencies can lead to inefficient exploration on classic exploration benchmarks (Dearden et al., 1998) . Experiments on the Atari benchmark suite (Bellemare et al., 2013) confirming that Successor Uncertainties can be scaled to work in complex domains.

2. Background:

Before proceeding, we briefly review model-free reinforcement learning and Bayesian linear regression. These are key concepts we shall use throughout the remainder of our paper.

Markov Decision Process:

We consider a finite MDP with a horizon H ∈ N, defined on a probability space (Ω, F, P ). An MDP is a tuple (S, A, T ), where S is a finite state space, A a finite action space, and T : S × A → P(S × R) a transition probability kernel with R ⊂ R a bounded reward set, and P(•) the set of probability measures on given space. For each time step t ∈ N, the agent selects an action A t by sampling from a distribution specified by the policy π t : S → P(A) for the current state S t , and receives a new state and reward (S t+1 , R t+1 ) ∼ T (S t , A t ). This gives rise to a random process (S t , A t ) t≥0 , Markov with respect to its natural filtration (F t ) t≥0 , and a sequence of random variables (R t ) t≥1 . The task of solving an MDP is that of finding an optimal policy π = (π t ) 0≤t≤H that maximises the expected return

J π = E[ H τ =0 γ τ R τ +1 ] where γ ∈ [0, 1) is a discount factor.

Model-free reinforcement learning All exploration methods we consider repose on a model-free approach to solving MDPs. These approaches invariably use estimates of state-action values and the corresponding Q function,

Q π t = E[ H τ =t γ τ -t R τ +1 | F t ]

, to guide exploration. Because computation of Q π t is usually intractable, it is estimated for each (s, a) ∈ S × A using a (stationary) parametric model Qπ : S × A → R. This model is trained to obey the Bellman equation, The result of this procedure is an estimate of the optimal policy defined by greedy behaviour with respect to Qπ t , that is π : s → |A| -1 δ A with A = {a : Qπ (s, a) = max a ∈A Qπ (s, a )} and |A| its cardinality.

Q π t = E[R t+1 | F t ] + γE Q π t+1 | F t ,

Bayesian linear regression:

The goal of linear regression is to fit a model of the form r = φ, w using a data set D = {(φ i , r i )} t i=0 . We will treat estimation of w in a fully Bayesian manner, assuming a multivariate normal prior P w = N (0, Σ p ) for w and a Gaussian likelihood P r|φ,w = N ( φ, w , β). By application of the Bayes' rule, the posterior over the weights given D is then

P w|D = N (µ w , Σ w ). Here Σ w = (Σ p + β -1 t i=0 φ i φ T i ) -1 and µ w = Σ w ( t i=0 r i φ i ).

Both µ w and Σ w can be computed in an online fashion.

3. Probabilistic exploration in reinforcement learning:

One of the most popular approaches to probabilistic exploration is posterior sampling (Dearden et al., 1998; Strens, 2000) , also known as Thompson sampling (Thompson, 1933) .

The key feature of this method is that a distribution P Qπ over Qπ is maintained instead of a point estimate. Exploration then proceeds by iterating the following steps: (i) sample Qπ ∼ P Qπ ; (ii) explore the environment using the greedy policy with respect to the sampled Qπ ; (iii) update P Qπ based on the collected data.

Effectiveness of posterior sampling is determined by properties of P Qπ . We will focus on two properties that we believe to be of vital importance.

1. Propagation of uncertainty: the marginal distribution over Qπ (s, a) must reflect the total downstream uncertainty about all Qπ (s , a ) that may be encountered when taking action a in state s and then following policy π.

Propagation of uncertainty ensures that uncertainty at any step is connected to expected uncertainties at subsequent time-steps (O'Donoghue et al., 2018) . It is sufficient for approaches based on marginal statistics, for example upper confidence bound exploration (Auer et al., 2002) , but does not guarantee coherence of individual samples.

2.. Incorporation of dependencies::

the set of random variables { Qπ (s, a)} s∈S,a∈A must (almost surely) obey dependencies specified by the Bellman equation.

Posterior sampling uses greedy policy with respect to a single sample from P Qπ . It is thus crucial that the value at any time-step is connected to the expected value at subsequent time-steps for every such sample, not only in expectation. If a sample which does not satisfy the Bellman equation is used, individual sampled Q values may carry conflicting information at each time-step. Importantly for our discussion later, a model P Qπ propagates uncertainty if it incorporates dependencies but the converse is not necessarily true.

4. Uncertainty propagation through Successor Features:

We present a new method, Successsor Uncertainties, which employs a probabilistic model P Qπ that does not assume a factorisation of the posterior over the state-action space. The model both propagates uncertainties and incorporates dependencies, and thus leads to temporally correlated exploration when combined with posterior sampling.

4.1. The Successor Uncertainty model:

Suppose that an embedding φ :

S × A → R d + exists, d ∈ N, such that ||φ|| 2 = 1 and E [R t+1 | F t ] = φ t , w for some w in R d and φ t = φ(S t , A t ). Then Q π t = E H τ =t γ τ -t R τ +1 | F t = E H τ =t γ τ -t φ τ , w | F t = E π H τ =t γ τ -t φ τ | F t , w ,

where the second equality follows from the tower property of conditional expectation and the third is a simple application of the dominated convergence theorem. We define

ψ π t = E[ H τ =t γ τ -t φ τ | F t ]

which allows us to express the Q function succinctly in terms of this quantity as Q π t = ψ π t , w . The random vectors ψ π t are known as successor features in the literature (Dayan, 1993; Barreto et al., 2017) , and correspond to the (discounted) expected future occurrence of each feature when following a given policy π.

We develop the Successor Uncertainties model for Q π t , denoted by Qπ SF , by modelling the quantities w and ψ π t . We use a stationary estimator of the successor features ψπ which will be obtained by noting that

ψ π t = φ t + γ E[ψ π t+1 | F t ],

meaning any standard temporal difference approach applies. We perform Bayesian linear regression to infer a distribution over w using a Gaussian prior N (0, αI) and likelihood N ( φ, w , β). This makes the computation of the posterior over w given D, P w|D = N (µ w , Σ w ), with D the set of all so far observed states and rewards, analytically tractable. Finally, stacking all successor feature estimates into a matrix Ψπ = [ ψπ (s, a)] s∈S,a∈A , the implied distribution over Qπ SF can be shown to be Law( Qπ SF ) = N ( Ψπ µ w , Ψπ Σ w ( Ψπ ) ) , highlighting the dependence between individual Qπ SF (s, a) introduced by our construction. A limitation of Successor Uncertainties is that only the rewards are explicitly treated in a probabilistic way, whereas a point estimate is used for the transition model in construction of ψπ . This issue may be partially alleviated following an argument analogous to the one made by O' Donoghue et al. (2018) . Specifically, if we are in a tabular setting and φ(s, a) are one-hot encoded representations of individual state-actions, the posterior variance Σ w will be diagonal with β(n sa + β α ) -1 for its non-zero entries, where n sa the number of times the tuple (s, a) has been observed. Combined with the above derived Law( Qπ SF ), the marginal variance over Qπ SF (s, a) will be proportional to βn -1 sa which will decrease at the same rate as if we had incorporated a Dirichlet model over the transitions (O'Donoghue et al., 2018) . Thus, if α = β and β is treated as a hyperparameter (as in our experiments), setting β sufficiently high can adjust our estimate for the uncertainty about the transition model. As in (O'Donoghue et al., 2018) , a similar argument can be made in terms of pseudocounts (Bellemare et al., 2016) for linear Q function models.

Finally, our derivations show that ψ π needs to be relearnt for each policy π which would generally require a significant amount of computational resources. To address this issue, we instead learn ψ π where π = π w dP w|D is the average over policies π w ; here π w is the greedy policy with respect to the particular Q function sample determined by w. This approximation reduces coherence of the dependencies within each Q function sample, but allows the algorithm to scale to large applications.

4.2. Successor Uncertainty with Neural Network embeddings:

One of the main assumptions we made is that the embedding function φ is known a priori. This section considers the scenario where φ has to be learnt jointly with other parameters.

To distinguish from the fixed φ, we use φ : 

S × A → R d + to

where r t+1 is the reward observed after taking action a t in state s t , (•) -means that the quantity (•) is treated as fixed, and ŵ is an auxiliary variable which can be either optimised or replaced with w and integrated out with respect to N (µ w , Σ w ). Please note that the problem is stated in terms of single step t to reduce clutter, but average over mini-batches sampled from a replay buffer will be optimised.

To simplify the optimisation, we will only softly enforce the constraint in equation ( 2) by adding a regularisation term ψπ t -( φt + γ E π(st+1) [ ψπ (s t+1 , a t+1 )]) -2 2 to the objective in equation ( 1). Empirically, the weight of the regularisation term can be set to one for all problems, eliminating what would otherwise be a hyperparameter.

Notice that the first term in equation ( 1) is the standard temporal difference learning loss. Now if a solution achieving zero temporal difference loss exists, then setting each φt = ψπ t -γ ψπ t+1 for a particular policy π will make the second term in equation ( 1) zero. In other words, the additional terms in equations ( 1) and ( 2) do not affect the set of parameters that solve the Q function estimation problem. Because the non-negativity condition is trivially satisfied with ReLU activations, the only change with potentially significant impact on optima is the φt 2 = 1 restriction which impacts the norm of ψt through equation ( 2). In our experience, this constraint simplifies selection of the prior over w as it keeps the scale of the inputs constant across all problems, increasing the robustness of our method.

We illustrate integration of neural network embeddings on our Atari architecture which is a modification of the DDQN model of Van Hasselt et al. (2016) . As you can see in figure 2 , Successor Uncertainties can be combined with existing reinforcement learning frameworks simply by including an additional neural network head for prediction of immediate rewards. The whole model is then optimised by minimising the above described relaxed version of the problem in equations (1-3), with the unit norm constraint on φ enforced by explicit normalisation, and the non-negativity by using the ReLU activation. We would like to emphasise that, unlike in previous work utilising successor embeddings (Kulkarni et al., 2016; Machado et al., 2017; 2018) , our approach does not employ an auxiliary state reconstruction or state-transition prediction model in order to learn meaningful representations. This significantly reduces the implementation effort, and, importantly, makes comparison of Successor Uncertainties to other methods easier as state-prediction is in itself a very strong auxiliary loss even without improved exploration.

5. Related probabilistic exploration methods:

Having introduced Successor Uncertainties, we now discuss several alternative probabilistic models of the Q function, focusing on if and how they propagate uncertainties and incorporate dependencies.

Variational Q-learning Deep Q-learning (Mnih et al., 2015; Watkins & Dayan, 1992 ) is a method which uses a deep neural network (Q-network) to estimate Q function of a particular exploration policy; the network is trained by minimising the temporal difference error on past observations seen by the agent. Some recent methods have attempted to replace the deterministic Q-network with a Bayesian neural network (BNN) and apply posterior sampling using the Bayesian uncertainty estimates (for example Lipton et al., 2016; Gal, 2016) . A common feature of these approaches is that the Q value estimates are treated as if they were real observations Q values in a regression problem. As shown by Riquelme et al. (2018) , if the Q value estimates are inaccurate, which often happens before the environment is sufficiently well explored, variational Q-learning uncertainty estimates can be slow to converge and result in poor exploration.

What is more, uncertainty estimates of these methods rely only on the particular stateaction pair (s, a) under consideration. Whilst the temporal difference updates ensure that the Q value for action a in state s does reflect the rewards that may be obtained by taking action a and a given policy π afterwards, uncertainty about those expected future rewards is not explicitly propagated.

Bayesian linear Q-function prediction Linear Q function models, despite their simplicity (or perhaps because of it), have proven a good modelling choice in reinforcement learning (Levine et al., 2017; Osband et al., 2014; Lagoudakis & Parr, 2003) . One of the possible explanations is that the tractability of the closed form updates often translates into lower variance and better calibrated uncertainty estimates as compared to variational approaches (Riquelme et al., 2018) .

Bayesian linear Q function prediction models, such as Bayesian Deep Q-Networks (BDQN Azizzadenesheli et al., 2018) or Double Uncertain Value Networks (Moerland et al., 2017) , utilise a standard Q-network: state s ∈ S is mapped to a state embedding φ π s , which is then used together with a vector u a for each action a ∈ A to predict Qπ (s, a) = u a , φ π s . An auxiliary Bayesian linear model is then used to independently model uncertainty over each u a . Analogously to Successor Uncertainties, the outputs of the Bayesian linear models define a distribution over Q functions P Qπ , which is used for exploration through posterior sampling. However, as in variational Q-learning, uncertainties are not explicitly propagated.

Uncertainty Bellman Equation:

The Uncertainty Bellman Equation (UBE) model (O'Donoghue et al., 2018) introduces a dynamic programming rule for accumulating Q function uncertainty estimates such that they reflect expected uncertainty over all future consequences. The method performs very well on the Atari benchmark suite (Bellemare et al., 2013) , achieving strong results on the Montezuma's Revenge exploration benchmark.

The UBE model assumes that the MDP is acyclic and thus that R i is independent of R j given π and T for j > i. It further assumes that local uncertainty estimates associated with each state-action tuple are available. These are denoted V t , and are estimated through the use of a Bayesian linear model in the form V : S × A → R + , as within the Bayesian Deep Q-Networks method. UBE then defines the expected future uncertainty

U π t = E H τ =t γ 2(τ -t) V t+1 | F t = E[V t+1 | F t ] + γ 2 E π U π t+1 | F t ,

where the second equality follows by the assumption of independent rewards. The latter expression for U π t allows for efficient learning of a neural network estimator Û : S × A → R + through the use of temporal difference updates.

In effect, UBE takes uncertainty estimates that do not propagate uncertainty, like those of the previously discussed methods, and forms a corrected estimator Û π which does reflect downstream uncertainty. However, as Û π specifies a marginal distribution over each Qπ (s, a), but not a joint distribution over Qπ , it does not incorporate dependencies between different states and actions as is apparent from the discussion before equation ( 3) in (O'Donoghue et al., 2018) where the implicit posterior has diagonal covariance matrix which implies independence in the case of jointly multivariate normal variables.

We present two sets of experiments. First, tabular experiments, where we provide an analysis of a general scenario within which Successor Uncertainties outperforms Bayesian Deep Q-Networks (BDQN) and the Uncertainty Bellman Equation (UBE). Second, we provide a limited set of Atari experiments, where we show that Successor Uncertainties can provide scores competitive with BDQN and UBE.

6.1. Tabular exploration problems:

Our experiments are based on two problems, the binary decision tree introduces first presented in the introduction (figure 1 ) and the Chain problem, introduced in (Dearden et al., 1998) . The chain is illustrated in figure 3 . The version we use follows from (Osband et al., 2016) , consisting of L+2 states and two actions corresponding to a move left and right along the chain respectively. The horizon is L + 9 allowing a maximal reward of 1.0. The simple to discover sub-optimal solution 'left always' makes the problem challenging. These two problems isolate exploration, that is the acquisition of data informative about the MDP, from learning, i.e. the manner in which that data is then used to learn a Q function estimator.

s -1 s 0 s 1 s L 1 1000 . . . 0 0 1 10 a 1 a 2 a 1 a 2 a 1 a 2 a 1 a 2 a 2 a 1 Figure 3: The Chain MDP illustrated

We present the return achieved by the three methods in these tabular problems within figure 4 . The poor performance of BDQN follows from its lack of propagation of uncertainty.

To understand the reason for the similarly poor performance of the Uncertainty Bellman Equation model, note that prior to the reward being reached, µ sa1 = µ s2 = 0 in expectation. Writing Qsai ∼ N (0, σ sai ) for the two actions up and down respectively, we see that P A0|S0 (a 1 | s) = P( Qsa1 -Qsa2 > 0) = 1 /2, irrespective of the variances σ 2 sai . Thus, on this task, posterior sampling with UBE uncertainty estimates results in exploration equivalent to that induced by a uniform random policy. Successor Uncertainties on the other hand distinctly outperforms both approaches. To see why, assume that s i is the furthest state seen thus far during exploration. Sample a weights vector w ∼ N (µ w , Σ w ). Now if the entry w i associated with s i is high, not only does that increase the probability of selecting action up in state s i-1 , but also for all states s 0 , . . . , s i-1 , as these contain a positive entry for s i in their successor representation ψ π . Thus exploration utilising Successor Uncertainties is not equivalent to that of a uniform policy but instead explores in a correlated manner. Under review as a conference paper at ICLR 2019

6.2. Atari experiments:

We've run a series of small experiments (5M steps) on a subset of Atari games. Here, we test the performance Successor Uncertainties using embeddings learnt online through the constrained model-free reinforcement learning approach described in section 4.2.

We present our results in table 1, alongside results from related exploration methods. Test scores reported are obtained using the standard procedure set out in the seminal work on deep reinforcement learning for Atari (Mnih et al., 2015) and utilised subsequently (Van Hasselt et al., 2016; O'Donoghue et al., 2018) . Games were selected by looking at training curves shown in (Azizzadenesheli et al., 2018) and the classification of game exploration difficulty in the appendix of (Ostrovski et al., 2017) . Relatively difficult games where most of the improvement in score can be attained early within training were chosen, allowing for informative experiments with only a modest computational budget. Atari results are more difficult to interpret and analyse than results from tabular experiments. Performance on these games is highly related to how well the method interacts with the convolutional neural network model utilised, and is thus less clear in terms of an evaluation of exploration performance. Our results are comparable to those of other common exploration methods, showing that Successor Uncertainties can be scaled to work in complex domains where state-action embeddings have to be learnt in an online manner.

7. Discussion:

We have introduced Successor Uncertainties, a probabilistic model for the Q function of an MDP that allows for efficient exploration at cost comparable to that of previous Qlearning methods. Our method separates local uncertainties from deterministic features learnt by a temporal difference approach. The benefit is that information is propagated within the state-action space of the MDP and statistical dependencies in the Q-function values are accounted for. Because of this, our method is a truly reinforcement learning exploration solution, rather than an adaptation of a contextual bandit method, like epsilongreedy, to a reinforcement learning problem. Our results show that we compare favourably with UBE and BDQN, two related strong probabilistic modelling techniques, on a subset of Atari problems. We also test on classic exploration benchmarks, which assess the ability of methods to deal with sparse rewards and low connectivity of the state-space, both key aspects of exploration that strong methods must be capable of handling. Here, Successor Uncertanties significantly outperforms both BDQN and UBE, highlighting the importance of incorporating dependencies in probabilistic Q function models.

:

Crazy Climber, Ms Pacman and Q*bert scores were steadily increasing as experiments ended.

Figures and Tables:

Figure 1: Binary tree MDP; in each state, agent moves either up or down. States with odd indices are terminal. Reward is only obtained after a large fixed number of up actions. The main contributions of our work are summarised below: 1. Identification of propagation of uncertainty and incorporation of uncertainties as desirable attributes of models used for probabilistic exploration; 2. Introduction of Successor Uncertainties, a probabilistic model which both propagates uncertainty and incorporates dependencies;

by minimising a monotonically increasing function of the temporal difference error td : (s, a, s , r, π) → | Qπ (s, a) -r -γE a ∼π(s ) [ Qπ (s , a )]| on the data collected through exploration.

Figure 2: Architecture diagram for network used for Atari experiments.

Figure 4: Test reward on Tree (left) and Chain (right), both with L=10. Episodes on x-axis. Average over 250 runs plotted; 95% CI shaded.

refer to the fitted model. Denoting (s t , a t ) the state-action observed at step t, φt = φ(s t , a t ) and ψπ t = ψπ (s t , a t ), we propose to learn φ and ψπ jointly by enforcing the known relationships between φ t , ψ π

Table content: <table/>

Results for selected Atari games for our proposed method Successor Uncertainty, BDQN(Azizzadenesheli et al., 2018), UBE(O'Donoghue et al., 2018) and Double Deep Q-Network (DDQN) with -greedy exploration(Van Hasselt et al., 2016). Numbers in parenthesis indicate the number of interactions with the environment undertaken by each method during training. Entries with dashes indicate no score was given in the source paper for that given game.

Table content: <table><tr><td>Method</td></tr></table>

Peer review (iclr2019):
 Summary/contribution: This paper focuses on the problem of incorporating uncertainty into RL. The primary contribution is exploring the use of successor features for uncertainty prediction over Q-values. The proposed approach builds on O’Donoghue et al. (2018). The authors provide experiments that demonstrate improved performance on a chain MDP environment and a Tree environment. Pros: - I found this paper to be above average in terms of clarity. Cons: - The experiments evaluation is restricted to simplistic environments. The authors make an argument for why using successor features would be more "stable", but I found the experimental evidence to support this claim to be underwhelming. Justification for rating: This paper does a good job of articulating an interesting approach to the exploration problem using successor representations. In the current form however, it is really lacking in experimental evidence to support the main claims/contributions. Currently the domains considered are somewhat toy which I do not find convincing enough to demonstrate the effectiveness of their approach. Other: - I would appreciate a discussion on the relationship to Machado et al. 2018 which explored count based exploration with successor representations.