Successor Uncertainties: exploration and uncertainty in temporal difference learning
Year: 2018

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.

heading: 1 Introduction
text: 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 poorly — shortsightedly 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.
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 propa-
gates uncertainty and incorporates dependencies; 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.

heading: 2 Background
text: 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 At by sampling from a distribution specified by the policy πt : S → P(A) for the current state St, and receives a new state and reward (St+1, Rt+1) ∼ T (St, At). This gives rise to a random process (St, At)t≥0, Markov with respect to its natural filtration (Ft)t≥0, and a sequence of random variables (Rt)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 γ
τ−tRτ+1 | Ft], 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, Qπt = E[Rt+1 | Ft] + γE [ Qπt+1 | Ft ] , by minimising a monotonically increasing function of the temporal difference error `td : (s, a, s′, r, π) 7→ |Q̂π(s, a)− r − γEa′∼π(s′)[Q̂π(s′, a′)]| on the data collected through exploration. The result of this procedure is an estimate of the optimal policy defined by greedy behaviour with respect to Q̂πt , that is π? : s 7→ |A|−1δA with A = {a : Q̂π(s, a) = maxa′∈A Q̂π(s, a′)} and |A| its cardinality.
Bayesian linear regression The goal of linear regression is to fit a model of the form r = 〈φ,w〉 using a data set D = {(φi, ri)}ti=0. We will treat estimation of w in a fully Bayesian manner, assuming a multivariate normal prior Pw = N (0,Σp) for w and a Gaussian likelihood Pr|φ,w = N (〈φ,w〉, β). By application of the Bayes’ rule, the posterior over the weights given D is then Pw|D = N (µw,Σw). Here Σw = (Σp + β−1 ∑t i=0 φiφ T i )−1 and
µw = Σw( ∑t i=0 riφi). Both µw and Σw can be computed in an online fashion.

heading: 3 Probabilistic exploration in reinforcement learning
text: 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 PQ̂π over Q̂π is maintained instead of a point estimate. Exploration then proceeds by iterating the following steps: (i) sample Q̂π ∼ PQ̂π ; (ii) explore the environment using the greedy policy with respect to the sampled Q̂π; (iii) update PQ̂π based on the collected data. Effectiveness of posterior sampling is determined by properties of PQ̂π . 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 PQ̂π . 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 PQ̂π propagates uncertainty if it incorporates dependencies but the converse is not necessarily true.

heading: 4 Uncertainty propagation through Successor Features
text: We present a new method, Successsor Uncertainties, which employs a probabilistic model PQ̂π 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.

heading: 4.1 The Successor Uncertainty model
text: Suppose that an embedding φ : S × A → Rd+ exists, d ∈ N, such that ||φ||2 = 1 and E [Rt+1 | Ft] = 〈φt, w〉 for some w in Rd and φt = φ(St, At). Then
Qπt = E [ H∑ τ=t γτ−tRτ+1 | Ft ] = E [ H∑ τ=t γτ−t〈φτ , w〉 | Ft ] = 〈 Eπ [ H∑ τ=t γτ−tφτ | Ft ] , 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φτ | Ft] 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 | Ft], 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, Pw|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 β(nsa + βα )
−1 for its non-zero entries, where nsa 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−1sa 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 π̄? = ∫ π?wdPw|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.

heading: 4.2 Successor Uncertainty with Neural Network embeddings
text: 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 → Rd+ to refer to the fitted model. Denoting (st, at) the state-action observed at step t, φ̂t = φ̂(st, at) and ψ̂πt = ψ̂π(st, at), we propose to learn φ̂ and ψ̂π jointly by enforcing the known relationships between φt, ψπt and E[Rt+1|Ft]: minimise ‖〈ŵ, ψ̂πt 〉 − γ(〈ŵ, ψ̂πt+1〉)− − rt+1‖22 + ‖〈ŵ, φ̂t〉 − rt+1‖22 w.r.t. φ̂t, ψ̂πt , ŵ (1) subject to ψ̂πt = φ̂t + γ Eπ(st+1)[ψ̂
π(st+1, a′t+1)] , (2) φ̂t, ψ̂ π t ≥ 0 elementwise and ‖φ̂t‖2 = 1 , (3)
where rt+1 is the reward observed after taking action at in state st, (·)− means that the quantity (·) is treated as fixed, and ŵ is an auxiliary variable which can be either optimised or replaced with w and integrated out with respect toN (µ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)[ψ̂π(st+1, a′t+1)])−‖22 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
check if the abbrev already introduced not previously introduced
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.

heading: 5 Related probabilistic exploration methods
text: 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
5
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 following 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 ua for each action a ∈ A to predict Q̂π(s, a) = 〈ua, φπs 〉. An auxiliary Bayesian linear model is then used to independently model uncertainty over each ua. Analogously to Successor Uncertainties, the outputs of the Bayesian linear models define a distribution over Q functions PQ̂π , 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 Ri is independent of Rj given π and T for j > i. It further assumes that local uncertainty estimates associated with each state-action tuple are available. These are denoted Vt, 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)Vt+1 | Ft ] = E[Vt+1 | Ft] + γ2Eπ [ Uπt+1 | Ft ] ,
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 Û : 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 Ûπ which does reflect downstream uncertainty. However, as Ûπ 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.

heading: 6 Experiments
text: 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 QNetworks (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.

heading: 6.1 Tabular exploration problems
text: 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.
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 Q̂sai ∼ N (0, σsai) for the two actions up and down respectively, we see that PA0|S0(a1 | s) = P(Q̂sa1 − Q̂sa2 > 0) = 1/2, irrespective of the variances σ2sai . 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 si is the furthest state seen thus far during exploration. Sample a weights vector w ∼ N (µw,Σw). Now if the entry wi associated with si is high, not only does that increase the probability of selecting action up in state si−1, but also for all states s0, . . . , si−1, as these contain a positive entry for si 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.

heading: 6.2 Atari experiments
text: 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.

heading: 7 Discussion
text: 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.
1Crazy Climber, Ms Pacman and Q*bert scores were steadily increasing as experiments ended.

Review summaries:
The paper proposes an exploration approach (either based on posterior sampling or optimism) based on successor features representation. A high probability ellipsoid confidence set (defined by the Gram matrix \Sigma_t) is estimated based on linear regression (using some features \phi) to the immediate reward function. Now for any policy \pi, a confidence interval for the Q-value of any policy \pi can be derived by application of the \psi transformation, where \psi = expected sum of discounted future \phi under \pi. The algorithm selects action by posterior sampling (or UCB) in the \psi-space.