Beyond variance reduction: Understanding the true impact of baselines on policy optimization

Wesley Chung, Valentin Thomas, Marlos C. Machado, Nicolas Le Roux

Introduction

In this paper we focus on techniques derived from stochastic optimization principles, such as EXP3 (Auer et al., 2002; Seldin et al., 2013). Despite the fact that they have higher regret in the non-adversarial setting than techniques explicitly tailored to minimize regret in bandit problems, like UCB (Agrawal, 1995) or Thompson sampling (Russo et al., 2017), they naturally extend to the MDP setting, where they are known as policy gradient methods.

We analyze the problem of learning to maximize the average reward, JJ, by gradient ascent:

with μa\mu_{a} being the average reward of arm aa. In this case, we are mainly interested in outputting an effective policy at the end of the optimization process, without explicitly considering the performance of intermediary policies.

Optimization theory predicts that the convergence speed of stochastic gradient methods will be affected by the variance of the gradient estimates and by the geometry of the function JJ, represented by its curvature. Roughly speaking, the geometry dictates how effective true gradient ascent is at optimizing J(θ)J(\theta) while the variance can be viewed as a penalty, capturing how much slower the optimization process is by using noisy versions of this true gradient. More concretely, doing one gradient step with stepsize α\alpha, using a stochastic estimate gtg_{t} of the gradient, leads to (Bottou et al., 2018):

when JJ is LL-smooth, i.e. its gradients are LL-Lipschitz.

As large variance has been identified as an issue for policy gradient (PG) methods, many works have focused on reducing the noise of the updates. One common technique is the use of control variates (Greensmith et al., 2004; Hofmann et al., 2015), referred to as baselines in the context of RL. These baselines bb are subtracted from the observed returns to obtain shifted returns, r(ai)−br(a_{i})-b, and do not change the expectation of the gradient. In MDPs, they are typically state-dependent. While the value function is a common choice, previous work showed that the minimum-variance baseline for the REINFORCE (Williams, 1992) estimator is different and involves the norm of the gradient (Peters and Schaal, 2008). Reducing variance has been the main motivation for many previous works on baselines (e.g., Gu et al., 2016; Liu et al., 2017; Grathwohl et al., 2017; Wu et al., 2018; Cheng et al., 2020), but the influence of baselines on other aspects of the optimization process has hardly been studied. We take a deeper look at baselines and their effects on optimization.

We show that baselines can impact the optimization process beyond variance reduction and lead to qualitatively different learning curves, even when the variance of the gradients is the same. For instance, given two baselines with the same variance, the more negative baseline promotes committal behaviour where a policy quickly tends towards a deterministic one, while the more positive baseline leads to non-committal behaviour, where the policy retains higher entropy for a longer period.

Furthermore, we show that the choice of baseline can even impact the convergence of natural policy gradient (NPG), something variance cannot explain. In particular, we construct a three-armed bandit where using the baseline minimizing the variance can lead to convergence to a deterministic, sub-optimal policy for any positive stepsize, while another baseline, with larger variance, guarantees convergence to the optimal policy. As such a behaviour is impossible under the standard assumptions in optimization, this result shows how these assumptions may be violated in practice. It also provides a counterexample to the convergence of NPG algorithms in general, a popular variant with much faster convergence rates than vanilla PG when using the true gradient in tabular MDPs (Agarwal et al., 2019).

Further, we identify on-policy sampling as a key factor to these convergence issues as it induces a vicious cycle where making bad updates can lead to worse policies, in turn leading to worse updates. A natural solution is to break the dependency between the sampling distribution and the updates through off-policy sampling. We show that ensuring all actions are sampled with sufficiently large probability at each step is enough to guarantee convergence in probability. Note that this form of convergence is stronger than convergence of the expected iterates, a more common type of result (e.g., Mei et al., 2020; Agarwal et al., 2019).

We also perform an empirical evaluation on multi-step MDPs, showing that baselines have a similar impact in that setting. We observe a significant impact on the empirical performance of agents when using two different sets of baselines yielding the same variance, once again suggesting that learning dynamics in MDPs are governed by more than the curvature of the loss and the variance of the gradients.

Baselines, learning dynamics & exploration

The problem defined in Eq. 1 can be solved by gradient ascent. Given access only to samples, the true gradient cannot generally be computed and the true update is replaced with a stochastic one, resulting in the following update:

where aia_{i} are actions drawn according to the agent’s current policy πθ\pi_{\theta}, α\alpha is the stepsize, and NN, which can be 1, is the number of samples used to compute the update. To reduce the variance of this estimate without introducing bias, we can introduce a baseline bb, resulting in the gradient estimate (r(ai)−b)∇θlog⁡πθ(ai)(r(a_{i})-b)\nabla_{\theta}\log\pi_{\theta}(a_{i}).

While the choice of baseline is known to affect the variance, we show that baselines can also lead to qualitatively different behaviour of the optimization process, even when the variance is the same. This difference cannot be explained by the expectation or variance, quantities which govern the usual bounds for convergence rates (Bottou et al., 2018).

To provide a complete picture of the optimization process, we analyze the evolution of the policy during optimization. We start in a simple setting, a deterministic three-armed bandit, where it is easier to produce informative visualizations.

Fig. 1 presents fifteen learning curves on the probability simplex representing the space of possible policies for the three-arm bandit, when using NPG and a softmax parameterization. We choose ϵ=\nicefrac12\epsilon=\nicefrac{{1}}{{2}} to obtain two baselines with the same variance: bθ+=bθ∗+\nicefrac12b^{+}_{\theta}=b^{*}_{\theta}+\nicefrac{{1}}{{2}} and bθ−=bθ∗−\nicefrac12b^{-}_{\theta}=b^{*}_{\theta}-\nicefrac{{1}}{{2}}.

Inspecting the plots, the learning curves for ϵ=−\nicefrac12\epsilon=-\nicefrac{{1}}{{2}} and ϵ=\nicefrac12\epsilon=\nicefrac{{1}}{{2}} are qualitatively different, even though the gradient estimates have the same variance. For ϵ=−\nicefrac12\epsilon=-\nicefrac{{1}}{{2}}, the policies quickly reach a deterministic policy (i.e., a neighborhood of a corner of the probability simplex), which can be suboptimal, as indicated by the curves ending up at the policy choosing action 2. On the other hand, for ϵ=\nicefrac12\epsilon=\nicefrac{{1}}{{2}}, every learning curve ends up at the optimal policy, although the convergence might be slower. The learning curves also do not deviate much from the curve for the true gradient. Again, these differences cannot be explained by the variance since the baselines result in identical variances.

Additionally, for bθ=bθ∗b_{\theta}=b^{*}_{\theta}, the learning curves spread out further. Compared to ϵ=\nicefrac12\epsilon=\nicefrac{{1}}{{2}}, some get closer to the top corner of the simplex, leading to convergence to a suboptimal solution, suggesting that the minimum-variance baseline may be worse than other, larger baselines. In the next section, we theoretically substantiate this and show that, for NPG, it is possible to converge to a suboptimal policy with the minimum-variance baseline; but there are larger baselines that guarantee convergence to an optimal policy.

We look at the update rules to explain these different behaviours. When using a baseline bb with NPG, sampling aia_{i} results in the update

Thus, supposing we sample action aia_{i}, if r(ai)−br(a_{i})-b is positive, which happens more often when the baseline bb is small (more negative), the update rule will increase the probability πθ(ai)\pi_{\theta}(a_{i}). This leads to an increase in the probability of taking the actions the agent took before, regardless of their quality (see Fig.1(a) for ϵ=−\nicefrac12\epsilon=-\nicefrac{{1}}{{2}}). Because the agent is likely to choose the same actions again, we call this committal behaviour.

While a smaller baseline leads to committal behaviour, a larger (more positive) baseline makes the agent second-guess itself. If r(ai)−br(a_{i})-b is negative, which happens more often when bb is large, the parameter update decreases the probability πθ(ai)\pi_{\theta}(a_{i}) of the sampled action aia_{i}, reducing the probability the agent will re-take the actions it just took, while increasing the probability of other actions. This might slow down convergence but it also makes it harder for the agent to get stuck. This is reflected in the ϵ=\nicefrac12\epsilon=\nicefrac{{1}}{{2}} case (Fig.1(c)), as all the learning curves end up at the optimal policy. We call this non-committal behaviour.

Additional empirical results can be found in Appendix A.1 for natural policy gradient and vanilla policy gradient for the softmax parameterization. Furthermore, we explore the use of projected stochastic gradient ascent and directly optimizing the policy probabilities πθ(a)\pi_{\theta}(a). We find qualitatively similar results in all three cases; baselines can induce committal and non-committal behaviour.

Convergence to suboptimal policies with natural policy gradient (NPG)

We empirically showed that PG algorithms can reach suboptimal policies and that the choice of baseline can affect the likelihood of this occurring. In this section, we provide theoretical results proving that it is indeed possible to converge to a suboptimal policy when using NPG. We discuss how this finding fits with existing convergence results and why standard assumptions are not satisfied in this setting.

Standard convergence results assume access to the true gradient (e.g., Agarwal et al., 2019) or, in the stochastic case, assume that the variance of the updates is uniformly bounded for all parameter values (e.g., Bottou et al., 2018). These assumptions are in fact quite strong and are violated in a simple two-arm bandit problem with fixed rewards. Pulling the optimal arm gives a reward of r1=+1r_{1}=+1, while pulling the suboptimal arm leads to a reward of r0=0r_{0}=0. We use the sigmoid parameterization and call pt=σ(θt)p_{t}=\sigma(\theta_{t}) the probability of sampling the optimal arm at time tt.

Our stochastic estimator of the natural gradient is

where bb is a baseline that does not depend on the action sampled at time tt but may depend on θt\theta_{t}. By computing the variance of the updates, Var[gt]=(1−pt−b)2pt(1−pt)\textrm{Var}[g_{t}]=\tfrac{(1-p_{t}-b)^{2}}{p_{t}(1-p_{t})}, we notice it is unbounded when the policy becomes deterministic, i.e. pt→0p_{t}\to 0 or pt→1p_{t}\to 1, violating the assumption of uniformly bounded variance, unless b=1−ptb=1-p_{t}, which is the optimal baseline. Note that using vanilla (non-natural) PG would, on the contrary, yield a bounded variance. In fact, we prove a convergence result in its favour in Appendix B (Prop. 4).

For NPG, the proposition below establishes potential convergence to a suboptimal arm and we demonstrate this empirically in Fig. 2.

Consider a two-arm bandit with rewards 11 and for the optimal and suboptimal arms, respectively. Suppose we use natural policy gradient starting from θ0\theta_{0}, with a fixed baseline b<0b<0, and fixed stepsize α>0\alpha>0. If the policy samples the optimal action with probability σ(θ)\sigma(\theta), then the probability of picking the suboptimal action forever and having θt\theta_{t} go to −∞-\infty is strictly positive. Additionally, if θ0≤0\theta_{0}\leq 0, we have

All the proofs may be found in the appendix. ∎

The updates provide some intuition as to why there is convergence to suboptimal policies. The issue is the committal nature of the baseline. Choosing an action leads to an increase of that action’s probability, even if it is a poor choice. Choosing the suboptimal arm leads to a decrease in θ\theta by αb1−pt\tfrac{\alpha b}{1-p_{t}}, thus increasing the probability the same arm is drawn again and further decreasing θ\theta. By checking the probability of this occurring forever, P(suboptimal arm forever)=∏t=1∞(1−pt)P(\text{suboptimal arm forever})=\prod_{t=1}^{\infty}(1-p_{t}), we show that 1−pt1-p_{t} converges quickly enough to 1 that the infinite product is nonzero, showing it is possible to get trapped choosing the wrong arm forever (Prop. 1), and θt→−∞\theta_{t}\to-\infty as tt grows.

This issue could be solved by picking a baseline with lower variance. For instance, the minimum-variance baseline b=1−ptb=1-p_{t} leads to variance and both possible updates are equal to +α+\alpha, guaranteeing that θ→+∞\theta\to+\infty, thus convergence. In fact, any baseline b∈(0,1)b\in(0,1) suffices since both updates are positive and greater than αmin⁡(b,1−b)\alpha\min(b,1-b). However, this is not always the case, as we show in the next section.

To decouple the impact of the variance with that of the committal nature of the baseline, Prop. 2 analyzes the learning dynamics in the two-arm bandit case for perturbations of the optimal baseline, i.e. we study baselines of the form b=b∗+ϵb=b^{*}+\epsilon and show how ϵ\epsilon, and particularly its sign, affects learning. Note that, because the variance is a quadratic function with its minimum in b∗b^{*}, both +ϵ+\epsilon and −ϵ-\epsilon have the same variance. Our findings can be summarized as follows:

For the two-armed bandit defined in Prop. 1, when using a perturbed min-variance baseline b=b∗+ϵb=b^{*}+\epsilon, the value of ϵ\epsilon determines the learning dynamics as follows:

For ϵ<−1\epsilon<-1, there is a positive probability of converging to the suboptimal arm.

For ϵ∈(−1,1)\epsilon\in(-1,1), we have convergence in probability to the optimal policy.

For ϵ≥1\epsilon\geq 1, the supremum of the iterates goes to +∞+\infty in probability.

While the proofs can be found in Appendix B.2, we provide here some intuition behind these results.

For ϵ<−1\epsilon<-1, we reuse the same argument as for b<0b<0 in Prop. 1. The probability of drawing the correct arm can decrease quickly enough to lead to convergence to the suboptimal arm.

For ϵ∈(−1,1)\epsilon\in(-1,1), the probability of drawing the correct arm cannot decrease too fast. Hence, although the updates, as well as the variance of the gradient estimate, are potentially unbounded, we still have convergence to the optimal solution in probability.

Finally, for ϵ≥1{\epsilon}\geq 1, we can reuse an intermediate argument from the ϵ∈(0,1)\epsilon\in(0,1) case to argue that for any threshold CC, the parameter will eventually exceed that threshold. For ϵ∈(0,1)\epsilon\in(0,1), once a certain threshold is crossed, the policy is guaranteed to improve at each step. However, with a large positive perturbation, updates are larger and we lose this additional guarantee, leading to the weaker result.

We want to emphasize that not only we get provably different dynamics for ϵ<−1\epsilon<-1 and ϵ≥1{\epsilon}\geq 1, showing the importance of the sign of the perturbation, but that there also is a sharp transition around ∣ϵ∣=1|\epsilon|=1, which cannot be captured solely by the variance.

2 Reducing variance with baselines can be detrimental

As we saw with the two-armed bandit, the direction of the updates is important in assessing convergence. More specifically, problems can arise when the choice of baseline induces committal behaviour. We now show a different bandit setting where committal behaviour happens even when using the minimum-variance baseline, thus leading to convergence to a suboptimal policy. Furthermore, we design a better baseline which ensures all updates move the parameters towards the optimal policy. This cements the idea that the quality of parameter updates must not be analyzed in terms of variance but rather in terms of the probability of going in a bad direction, since a baseline that induces higher variance leads to convergence while the minimum-variance baseline does not. The following theorem summarizes this.

There exists a three-arm bandit where using the stochastic natural gradient on a softmax-parameterized policy with the minimum-variance baseline can lead to convergence to a suboptimal policy with probability ρ>0\rho>0, and there is a different baseline (with larger variance) which results in convergence to the optimal policy with probability 1.

The bandit used in this theorem is the one we used for the experiments depicted in Fig. 1. The key is that the minimum-variance baseline can be lower than the second best reward; so pulling the second arm will increase its probability and induce committal behaviour. This can cause the agent to prematurely commit to the second arm and converge to the wrong policy. On the other hand, using any baseline whose value is between the optimal reward and the second best reward, which we term a gap baseline, will always increase the probability of the optimal action at every step, no matter which arm is drawn. Since the updates are sufficiently large at every step, this is enough to ensure convergence with probability 1, despite the higher variance compared to the minimum variance baseline. The key is that whether a baseline underestimates or overestimates the second best reward can affect the algorithm convergence and this is more critical than the resulting variance of the gradient estimates.

As such, more than lower variance, good baselines are those that can assign positive effective returns to the good trajectories and negative effective returns to the others. These results cast doubt on whether finding baselines which minimize variance is a meaningful goal to pursue. The baseline can affect optimization in subtle ways, beyond variance, and further study is needed to identify the true causes of some improved empirical results observed in previous works. This importance of the sign of the returns, rather than their exact value, echoes with the cross-entropy method (De Boer et al., 2005), which maximizes the probability of the trajectories with the largest returns, regardless of their actual value.

Off-policy sampling

So far, we have seen that committal behaviour can be problematic as it can cause convergence to a suboptimal policy. This can be especially problematic when the agent follows a near-deterministic policy as it is unlikely to receive different samples which would move the policy away from the closest deterministic one, regardless of the quality of that policy.

Up to this point, we assumed that actions were sampled according to the current policy, a setting known as on-policy. This setting couples the updates and the policy and is a root cause of the committal behaviour: the update at the current step changes the policy, which affects the distribution of rewards obtained and hence the next updates. However, we know from the optimization literature that bounding the variance of the updates will lead to convergence (Bottou et al., 2018). As the variance becomes unbounded when the probability of drawing some actions goes to 0, a natural solution to avoid these issues is to sample actions from a behaviour policy that selects every action with sufficiently high probability. Such a policy would make it impossible to choose the same, suboptimal action forever.

Because the behaviour policy changed, we introduce importance sampling (IS) corrections to preserve the unbiased updates (Kahn and Harris, 1951; Precup, 2000). These changes are sufficient to guarantee convergence for any baseline:

Consider a nn-armed bandit with stochastic rewards with bounded support and a unique optimal action. The behaviour policy μt\mu_{t} selects action ii with probability μt(i)\mu_{t}(i) and let ϵt=min⁡iμt(i)\epsilon_{t}=\min_{i}\mu_{t}(i). When using NPG with importance sampling and a bounded baseline bb, if lim⁡t→∞t ϵt2=+∞\lim_{t\to\infty}t\ \epsilon_{t}^{2}=+\infty , then the target policy πt\pi_{t} converges to the optimal policy in probability.

(Sketch) Using Azuma-Hoeffding’s inequality, we can show that for well chosen constants Δi,δ\Delta_{i},\delta and C>0C>0 ,

The condition on μt\mu_{t} imposes a cap on how fast the behaviour policy can become deterministic: no faster than t−1/2t^{-1/2}. Intuitively, this ensures each action is sampled sufficiently often and prevents premature convergence to a suboptimal policy. The condition is satisfied for any sequence of behaviour policies which assign at least ϵt\epsilon_{t} probability to each action at each step, such as ϵ\epsilon-greedy policies. It also holds if ϵt\epsilon_{t} decreases over time at a sufficiently slow rate. By choosing as behaviour policy μ\mu a linear interpolation between π\pi and the uniform policy, μ(a)=(1−γ)π(a)+γK,γ∈(0,1]\mu(a)=(1-\gamma)\pi(a)+\frac{\gamma}{K},\gamma\in(0,1], where KK is the number of arms, we recover the classic EXP3 algorithm (Auer et al., 2002; Seldin et al., 2012).

We can also confirm that this condition is not satisfied for the simple example we presented when discussing convergence to suboptimal policies. There, ptp_{t} could decrease exponentially fast since the tails of the sigmoid function decay exponentially and the parameters move by at least a constant at every step. In this case, ϵt=Ω(e−t)\epsilon_{t}=\Omega(e^{-t}), resulting in lim⁡t→∞te−2t=0\lim_{t\to\infty}te^{-2t}=0, so Proposition 3 does not apply.

2 Importance sampling, baselines & variance

As we have seen, using a separate behaviour policy that samples all actions sufficiently often may lead to stronger convergence guarantees, even if it increases the variance of the gradient estimates in most of the space, as what matters is what happens in the high variance regions, which are usually close to the boundaries. Fig. 3 shows the ratios of gradient variances between on-policy PG without baseline, on-policy PG with the minimum variance baseline, and off-policy PG using importance sampling (IS) where the sampling distribution is μ(a)=12π(a)+16\mu(a)=\frac{1}{2}\pi(a)+\frac{1}{6}, i.e. a mixture of the current policy π\pi and the uniform distribution. While using the minimum variance baseline decreases the variance on the entire space compared to not using a baseline, IS actually increases the variance when the current policy is close to uniform. However, IS does a much better job at reducing the variance close to the boundaries of the simplex, where it actually matters to guarantee convergence.

12𝜋𝑎16\mu(a)=\frac{1}{2}\pi(a)+\frac{1}{6}. Note that this is not the min. variance sampling distribution and it leads to higher variance than PG without a baseline in some parts of the simplex. This suggests that convergence of PG methods is not so much governed by the variance of the gradient estimates in general, but by the variance in the worst regions, usually near the boundary. While baselines can reduce the variance, they generally cannot prevent the variance in those regions from exploding, leading to the policy getting stuck. Thus, good baselines are not the ones reducing the variance across the space but rather those that can prevent the learning from reaching these regions altogether. Large values of bb, such that r(ai)−br(a_{i})-b is negative for most actions, achieve precisely that. On the other hand, due to the increased flexibility of sampling distributions, IS can limit the nefariousness of these critical regions, offering better convergence guarantees despite not reducing variance everywhere.

Importantly, although IS is usually used in RL to correct for the distribution of past samples (e.g., Munos et al., 2016), we advocate here for expanding the research on designing appropriate sampling distributions as done by Hanna et al. (2017, 2018) and Parmas and Sugiyama (2019). This line of work has a long history in statistics (c.f., Liu, 2008).

3 Other mitigating strategies

We conclude this section by discussing alternative strategies to mitigate the convergence issues. While they might be effective, and some are indeed used in practice, they are not without pitfalls.

First, one could consider reducing the stepsizes, with the hope that the policy would not converge as quickly towards a suboptimal deterministic policy and would eventually leave that bad region. Indeed, if we are to use vanilla PG in the two-arm bandit example, instead of NPG, this effectively reduces the stepsize by a factor of σ(θ)(1−σ(θ))\sigma(\theta)(1-\sigma(\theta)) (the Fisher information). In this case, we are able to show convergence in probability to the optimal policy. See Proposition 4 in Appendix B.

Empirically, we find that, when using vanilla PG, the policy may still remain stuck near a suboptimal policy when using a negative baseline, similar to Fig. 2. While the previous proposition guarantees convergence eventually, the rate may be very slow, which remains problematic in practice. There is theoretical evidence that following even the true vanilla PG may result in slow convergence (Schaul et al., 2019), suggesting that the problem is not necessarily due to noise.

An alternative solution would be to add entropy regularization to the objective. By doing so, the policy would be prevented from getting too close to deterministic policies. While this might prevent convergence to a suboptimal policy, it would also exclude the possibility of fully converging to the optimal policy, though the policy may remain near it.

In bandits, EXP3 has been found not to enjoy high-probability guarantees on its regret so variants have been developed to address this deficiency (c.f. Lattimore and Szepesvári, 2020). For example, by introducing bias in the updates, their variance can be reduced significantly Auer et al. (2002); Neu (2015). Finally, other works have also developed provably convergent policy gradient algorithms using different mechanisms, such as exploration bonuses or ensembles of policies (Cai et al., 2019; Efroni et al., 2020; Agarwal et al., 2020).

Extension to multi-step MDPs

where there is now a discounted distribution over states induced by πθ\pi_{\theta}. Although that distribution depends on πθ\pi_{\theta} in a potentially complex way, the parameter updates are similar to Eq. 2:

where (ai,si)(a_{i},s_{i}) pairs are drawn according to the discounted state-visitation distribution induced by πθ\pi_{\theta} and QQ is the state-action value function induced by πθ\pi_{\theta} (c.f. Sutton and Barto, 2018). To match the bandit setting and common practice, we made the baseline state dependent.

Although our theoretical analyses do not easily extend to multi-step MDPs, we empirically investigated if the similarity between these formulations leads to similar differences in learning dynamics when changing the baseline. We consider a 10x10 gridworld consisting of 4 rooms as depicted on Fig. 4(a). We use a discount factor γ=0.99\gamma=0.99. The agent starts in the upper left room and two adjacent rooms contain a goal state of value 0.60.6 or 0.30.3. The best goal (even discounted), with a value of 11, lies in the furthest room, so that the agent must learn to cross the sub-optimal rooms and reach the furthest one.

Similar to the bandit setting, for a state ss, we can derive the minimum-variance baseline b∗(s)b^{*}(s) assuming access to state-action values Q(s,a)Q(s,a) for πθ\pi_{\theta} and consider perturbations to it. Again, we use baselines b(s)=b∗(s)+ϵb(s)=b^{*}(s)+\epsilon and b(s)=b∗(s)−ϵb(s)=b^{*}(s)-\epsilon, since they result in identical variances. We use a natural policy gradient estimate, which substitutes ∇log⁡π(ai∣si)\nabla\log\pi(a_{i}|s_{i}) by Fsi−1∇log⁡π(ai∣si)F^{-1}_{s_{i}}\nabla\log\pi(a_{i}|s_{i}) in the update rule, where FsiF_{s_{i}} is the Fisher information matrix for state sis_{i} and solve for the exact Q(s,a)Q(s,a) values using dynamic programming for all updates (see Appendix D.6 for details).

In order to identify the committal vs. non-committal behaviour of the agent depending on the baseline, we monitor the entropy of the policy and the entropy of the stationary state distribution over time. Fig.4(b) shows the average returns over time and Fig.4(c) and 4(d) show the entropy of the policy in two ways. The first is the average entropy of the action distribution along the states visited in each trajectory, and the second is the entropy of the distribution of the number of times each state is visited up to that point in training.

The action entropy for smaller baselines tends to decay faster compared to larger ones, indicating convergence to a deterministic policy. This quick convergence is premature in some cases since the returns are not as high for the lower baselines. In fact for ϵ=−1\epsilon=-1, we see that the agent gets stuck on a policy that is unable to reach any goal within the time limit, as indicated by the returns of . On the other hand, the larger baselines tend to achieve larger returns with larger entropy policies, but do not fully converge to the optimal policy as evidenced by the gap in the returns plot.

Since committal and non-committal behaviour can be directly inferred from the PG and the sign of the effective rewards R(τ)−bR(\tau)-b, we posit that these effects extend to all MDPs. In particular, in complex MDPs, the first trajectories explored are likely to be suboptimal and a low baseline will increase their probability of being sampled again, requiring the use of techniques such as entropy regularization to prevent the policy from getting stuck too quickly.

Conclusion

We presented results that dispute common beliefs about baselines, variance, and policy gradient methods in general. As opposed to the common belief that baselines only provide benefits through variance reduction, we showed that they can significantly affect the optimization process in ways that cannot be explained by the variance and that lower variance can even sometimes be detrimental.

Different baselines can give rise to very different learning dynamics, even when they reduce the variance of the gradients equally. They do that by either making a policy quickly tend towards a deterministic one (committal behaviour) or by maintaining high-entropy for a longer period of time (non-committal behaviour). We showed that committal behaviour can be problematic and lead to convergence to a suboptimal policy. Specifically, we showed that stochastic natural policy gradient does not always converge to the optimal solution due to the unusual situation in which the iterates converge to the optimal policy in expectation but not almost surely. Moreover, we showed that baselines that lead to lower-variance can sometimes be detrimental to optimization, highlighting the limitations of using variance to analyze the convergence properties of these methods. We also showed that standard convergence guarantees for PG methods do not apply to some settings because the assumption of bounded variance of the updates is violated.

The aforementioned convergence issues are also caused by the problematic coupling between the algorithm’s updates and its sampling distribution since one directly impacts the other. As a potential solution, we showed that off-policy sampling can sidestep these difficulties by ensuring we use a sampling distribution that is different than the one induced by the agent’s current policy. This supports the hypothesis that on-policy learning can be problematic, as observed in previous work (Schaul et al., 2019; Hennes et al., 2020). Nevertheless, importance sampling in RL is generally seen as problematic (van Hasselt et al., 2018) due to instabilities it introduces to the learning process. Moving from an imposed policy, using past trajectories, to a chosen sampling policy reduces the variance of the gradients for near-deterministic policies and can lead to much better behaviour.

More broadly, this work suggests that treating bandit and reinforcement learning problems as a black-box optimization of a function J(θ)J(\theta) may be insufficient to perform well. As we have seen, the current parameter value can affect all future parameter values by influencing the data collection process and thus the updates performed. Theoretically, relying on immediately available quantities such as the gradient variance and ignoring the sequential nature of the optimization problem is not enough to discriminate between certain optimization algorithms. In essence, to design highly-effective policy optimization algorithms, it may be necessary to develop a better understanding of how the optimization process evolves over many steps.

Acknowledgements

We would like to thank Kris de Asis, Alan Chan, Ofir Nachum, Doina Precup, Dale Schuurmans, and Ahmed Touati for helpful discussions. We also thank Courtney Paquette, Vincent Liu and Scott Fujimoto for reviewing an earlier version of this paper. Nicolas Le Roux is supported by a Canada CIFAR AI Chair.

References

Organization of the appendix

We organize the appendix into several thematic sections.

The first one, section A contains additional experiments and figures on bandits and MDPs. We have further investigations into committal and non-committal behaviour with baselines. More precisely subsection A.1 contains additional experiments for the 3 arm bandits for vanilla policy gradient, natural policy gradient and policy gradient with direct parameterization and a discussion on the effect the hyperparameters have on the results. In all cases, we find evidence for committal and non-committal behaviours. In the rest of the section, we investigate this in MDPs, starting with a smaller MDP with 2 different goals in subsection A.2 and constant baselines. We also provide additional experiments on the 4 rooms environment in subsection A.3, including the vanilla policy gradient and constant baselines with REINFORCE.

Then, section B contains theory for the two-armed bandit case, namely proofs of convergence to a suboptimal policy (Proposition 1 in Appendix B.1) and an analysis of perturbed minimum-variance baselines (Proposition 2 in Appendix B.2). For the latter, depending on the perturbation, we may have possible convergence to a suboptimal policy, convergence to the optimal policy in probability, or a weaker form of convergence to the optimal policy. Finally, we also show vanilla policy gradient converges to the optimal policy in probability regardless of the baseline in Appendix B.3.

Section C contains the theory for multi-armed bandit, including the proof of theorem 1. This theorem presents a counterexample to the idea that reducing variance always improves optimization. We show that there is baseline leading to reduced variance which may converge to a suboptimal policy with positive probability (see Appendix C.1) while there is another baseline with larger variance that converges to the optimal policy with probability 11 (see Appendix C.2). We identify on-policy sampling as being a potential source of these convergence issues. We provide proofs of proposition 3 in Appendix C.3, which shows convergence to the optimal policy in probability when using off-policy sampling with importance sampling.

Finally, in section D, we provide derivations of miscellaneous, smaller results such as the calculation of the minimum-variance baseline (Appendix D.1), the natural policy gradient update for the softmax parameterization (Appendix D.2) and the connection between the value function and the minimum-variance baseline (Appendix D.3).

Appendix A Other experiments

In this subsection, we provide additional experiments on the three-armed bandit with natural and vanilla policy gradients for the softmax parameterization, varying the initializations. Additionally, we present results for the direct parameterization and utilizing projected stochastic gradient ascent.

The main takeaway is that the effect of the baselines appears more strongly when the initialization is unfavorable (for instance with a high probability of selecting a suboptimal action at first). The effect also are diminished when using small learning rates as in that case the effect of the noise on the optimization process lessens.

While the simplex visualization is very appealing, we mainly show here learning curves as we can showcase more seeds that way and show the effects are noticeable across many runs.

Figure 5 uses the same setting as Figure 1 with 40 trajectories instead of 15. We do once again observe many cases of convergence to the wrong arm for the negative baseline and some cases for the minimum variance baseline, while the positive baseline converges reliably. In this case the value function also converges to the optimal solution but is much slower.

Figure 6 shows a similar setting to Figure 5 but where the initialization parameter is not as extreme. We observe the same type of behavior, but not as pronounced as before; fewer seeds converge to the wrong arm.

In Figure 7 whose initial policy is the uniform, we observe that the minimum variance baseline and the value function as baseline perform very well. On the other hand the committal baseline still has seeds that do not converge to the right arm. Interestingly, while all seeds for the non-committal baseline identify the optimal arm, the variance of the return is higher than for the optimal baseline, suggesting a case similar to the result presented in Proposition 6 where a positive baseline ensured we get close to the optimal arm but may not remain arbitrary close to it.

Vanilla policy gradient

While we have no theory indicating that we may converge to a suboptimal arm with vanilla policy gradient, we can still observe some effect in terms of learning speed in practice (see Figures 8 to 11).

On Figures 8 and 9 we plot the simplex view and the learning curves for vanilla policy gradient initialized at the uniform policy. We do observe that some trajectories did not converge to the optimal arm in the imparted time for the committal baseline, while they converged in all other settings. The mininum variance baseline is slower to converge than the non-committal and the value function in this setting as can be seem both in the simplex plot and learning curves.

On Figures 10 and 11 we plot the simplex view and the learning curves for vanilla policy gradient initialized at a policy yielding a very high probability of sampling the suboptimal actions, 48.7%48.7\% for each. We do observe a similar behavior than for the previous plots with vanilla PG, but in this setting the minimum variance baseline is even slower to converge and a few seeds did not identify the optimal arm. As the gradient flow leads the solutions closer to the simplex edges, the simplex plot is not as helpful in this setting to understand the behavior of each baseline option.

Policy gradient with direct parameterization

Here we present results with the direct parameterization, i.e where θ\theta contains directly the probability of drawing each arm. In that case the gradient update is

where Δ3\Delta_{3} is the three dimensional simplex Δ3={u,v,w≥0,u+v+w=1}\Delta_{3}=\{u,v,w\geq 0,u+v+w=1\}. In this case, however, because the projection step is non trivial and doesn’t have an easy explicit closed form solution (but we can express it as the output of an algorithm), we cannot explicitly write down the optimal baseline. Again, because of the projection step, baselines of this form are not guaranteed to preserve unbiasedness of the gradient estimate. For this reason, we only show experiments with fixed baselines, but keep in mind that these results are not as meaningful as the ones presented above. We present the results in Figures 12 and 13.

Once again in this setting we can see that negative baselines tend to encourage convergence to a suboptimal arm while positive baselines help converge to the optimal arm.

A.2 Simple gridworld

As a simple MDP with more than one state, we experiment using a 5x5 gridworld with two goal states, the closer one giving a reward of 0.8 and the further one a reward of 1. We ran the vanilla policy gradient with a fixed stepsize and discount factor of 0.990.99 multiple times for several baselines. Fig. 14 displays individual learning curves with the index of the episode on the x-axis, and the fraction of episodes where the agent reached the reward of 1 up to that point on the y-axis. To match the experiments for the four rooms domain in the main text, Fig. 15 shows returns and the entropy of the actions and state visitation distributions for multiple settings of the baseline. Once again, we see a difference between the smaller and larger baselines. In fact, the difference is more striking in this example since some learning curves get stuck at suboptimal policies. Overall, we see two main trends in this experiment: a) The larger the baseline, the more likely the agent converges to the optimal policy, and b) Agents with negative baselines converge faster, albeit sometimes to a suboptimal behaviour. We emphasize that a) is not universally true and large enough baselines will lead to an increase in variance and a decrease in performance.

A.3 Additional results on the 4 rooms environment

For the four-rooms gridworld discussed in the main text, we extend the experiments and provide additional details. The environment is a 10x10 gridworld consisting of 4 rooms as depicted on Fig. 4(a) with a discount factor γ=0.99\gamma=0.99. The agent starts in the upper left room and two adjacent rooms contain a goal state of value 0.60.6 (discounted, ≈0.54\approx 0.54) or 0.30.3 (discounted, ≈0.27\approx 0.27). However, the best goal, with a value of 11 (discounted, ≈0.87\approx 0.87), lies in the furthest room, so that the agent must learn to cross the sub-optimal rooms and reach the furthest one.

For the NPG algorithm used in the main text, we required solving for Qπ(s,a)Q_{\pi}(s,a) for the current policy π\pi. This was done using dynamic programming on the true MDP, stopping when the change between successive approximations of the value function didn’t differ more than 0.0010.001. Additionally, a more thorough derivation of the NPG estimate we use can be found in Appendix D.6.

We also experiment with using the vanilla policy gradient with the tabular softmax parameterization in the four-rooms environment. We use a similar estimator of the policy gradient which makes updates of the form:

for all observed si,ais_{i},a_{i} in the sampled trajectory. As with the NPG estimator, we can find the minimum-variance baseline bθ∗b^{*}_{\theta} in closed-form and thus can choose baselines of the form b+=bθ∗+ϵb^{+}=b^{*}_{\theta}+\epsilon and bθ−=bθ∗−ϵb^{-}_{\theta}=b^{*}_{\theta}-\epsilon to ensure equal variance as before. Fig. 17 plots the results. In this case, we find that there is not a large difference between the results for +ϵ+\epsilon and −ϵ-\epsilon, unlike the results for NPG and those for vanilla PG in the bandit setting.

The reason for this discrepancy may be due to the magnitudes of the perturbations ϵ\epsilon relative to the size of the unperturbed update Qπ(si,ai)−bθ∗Q_{\pi}(s_{i},a_{i})-b^{*}_{\theta}. The magnitude of Qπ(si,ai)−b∗Q_{\pi}(s_{i},a_{i})-b^{*} varies largely from the order of 0.0010.001 to 0.10.1, even within an episode. To investigate this further, we try another experiment using perturbations ϵ=c(max⁡aQπ(si,a)−bθ∗\epsilon=c(\max_{a}Q_{\pi}(s_{i},a)-b^{*}_{\theta} for various choices of c>0c>0. This would ensure that the magnitude of the perturbation is similar to the magnitude of Qπ(si,ai)−b∗Q_{\pi}(s_{i},a_{i})-b^{*}, while still controlling for the variance of the gradient estimates. In Fig. 16, we see that there is a difference between the +ϵ+\epsilon and −ϵ-\epsilon settings. As expected, the +ϵ+\epsilon baseline leads to larger action and state entropy although, in this case, this results in a reduction of performance. Overall, the differences between vanilla PG and natural PG are not fully understood and there may be many factors playing a role, possibly including the size of the updates, step sizes and the properties of the MDP.

We consider an alternative visualization for the experiment of vanilla policy gradient with constant baselines: Figures 19(a), 19(b) and 19(c). Each point in the simplex is a policy, and the position is an estimate, computed with 1,0001,000 Monte-Carlo samples, of the probability of the agent reaching each of the 3 goals. We observe that the starting point of the curve is equidistant to the 2 sub-optimal goals but further from the best goal, which is coherent with the geometry of the MDP. Because we have a discount factor of γ=0.99\gamma=0.99, the agent first learns to reach the best goal in an adjacent room to the starting one, and only then it learns to reach the globally optimal goal fast enough for its reward to be the best one.

In these plots, we can see differences between b=−1b=-1 and b=1b=1. For the lower baseline, we see that trajectories are much more noisy, with some curves going closer to the bottom-right corner, corresponding to the worst goal. This may suggest that the policies exhibit committal behaviour by moving further towards bad policies. On the other hand, for b=1b=1, every trajectory seems to reliably move towards the top corner before converging to the bottom-left, an optimal policy.

Appendix B Two-armed bandit theory

In this section, we expand on the results for the two-armed bandit. First, we show that there is some probability of converging to the wrong policy when using natural policy gradient with a constant baseline. Next, we consider all cases of the perturbed minimum-variance baseline (b=b∗+ϵ)b=b^{*}+\epsilon) and show that some cases lead to convergence to the optimal policy with probability 1 while others do not. In particular there is a difference between ϵ<−1\epsilon<-1 and ϵ>1\epsilon>1, even though these settings can result in the same variance of the gradient estimates. Finally, we prove that the vanilla policy gradient results in convergence in probability to the optimal policy regardless of the baseline, in contrast to the natural policy gradient.

pt=σ(θt)p_{t}=\sigma(\theta_{t}) is the probability of sampling the optimal arm (arm 1).

gtg_{t} is a stochastic unbiased estimate of ∇θJ(θt)\nabla_{\theta}J(\theta_{t}). It will take different forms depending on whether we use vanilla or natural policy gradient and whether we use importance sampling or not.

For {αt}t\{\alpha_{t}\}_{t} the sequence of stepsizes, the current parameter θt\theta_{t} is a random variable equal to θt=∑i=1tαigi+θ0\theta_{t}=\sum_{i=1}^{t}\alpha_{i}g_{i}+\theta_{0} where θ0\theta_{0} is the initial parameter value.

We will also be making use of Azuma-Hoeffding’s inequality to show that the iterates stay within a certain region with high-probability, leading to convergence to the optimal policy.

For {Xt}\{X_{t}\} a martingale, if ∣Xt−Xt−1∣≤ct|X_{t}-X_{t-1}|\leq c_{t} almost surely, then we have ∀t,ϵ≥0\forall t,\epsilon\geq 0

B.1 Convergence to a suboptimal policy with a constant baseline

For the proofs in this subsection, we assume that the step size is constant i.e. αt=α\alpha_{t}=\alpha for all tt and that the rewards are deterministic.

First, we deal with the case where θ0<0\theta_{0}<0.

Next, we use the bound 1−x≥exp⁡(−x1−x)1-x\geq\exp(\frac{-x}{1-x}). This bound can be derived as follows:

Continuing with x=exp⁡(θ0−αbt)x=\exp(\theta_{0}-\alpha bt), the bound holds when x∈[0,1)x\in[0,1), which is satisfied assuming θ0≤0\theta_{0}\leq 0.

For now we ignore t=0t=0 and we will just multiply it back in at the end.

The last line follows by considering the integrand as the right endpoints of rectangles approximating the area above the curve.

Solving this integral by substituting y=−θ0+αbty=-\theta_{0}+\alpha bt, multiplying the numerator and denominator by eye^{y} and substituting u=eyu=e^{y}, we get:

If θ0>0\theta_{0}>0, then there is a positive probability of reaching θ<0\theta<0 in a finite number of steps since choosing action 2 makes a step of size αb\alpha b in the left direction and we will reach θt<0\theta_{t}<0 after m=θ0−0αbm=\frac{\theta_{0}-0}{\alpha b} steps leftwards. The probability of making mm left steps in a row is positive. So, we can simply lower bound the probability of picking left forever by the product of that probability and the derived bound for θ0≤0\theta_{0}\leq 0. ∎

The regret for the previously described two-armed bandit is linear.

Letting RtR_{t} be the reward collected at time tt,

The second line follows since choosing the left action at each step incurs a regret of 11 and this is one term in the entire expectation. The third line follows since choosing left TT times is a subset of the event of choosing left forever. The last line implies linear regret since we know Pr(left forever)>0Pr(\text{left forever})>0 by the previous theorem. ∎

B.2 Analysis of perturbed minimum-variance baseline

In this section, we look at perturbations of the minimum-variance baseline in the two-armed bandit, i.e. baselines of the form b=1−pt+ϵb=1-p_{t}+\epsilon. In summary:

For ϵ<−1\epsilon<-1, convergence to a suboptimal policy is possible with positive probability.

For ϵ∈(−1,1)\epsilon\in(-1,1), we have convergence almost surely to the optimal policy.

For ϵ≥1\epsilon\geq 1, the supremum of the iterates goes to ∞\infty (but we do not have convergence to an optimal policy)

It is interesting to note that there is a subtle difference between the case of ϵ∈(−1,0)\epsilon\in(-1,0) and ϵ∈(0,1)\epsilon\in(0,1), even though both lead to convergence. The main difference is that when θt\theta_{t} is large, positive ϵ\epsilon leads to both updates being positive and hence improvement is guaranteed at every step. But, when ϵ\epsilon is negative, then only one of the actions leads to improvement, the other gives a large negative update. So, in some sense, for ϵ∈(−1,0)\epsilon\in(-1,0), convergence is less stable because a single bad update could be catastrophic.

Also, the case of ϵ=−1\epsilon=-1 proved to be difficult. Empirically, we found that the agent would incur linear regret and it seemed like some learning curves also got stuck near p=0p=0, but we were unable to theoretically show convergence to a suboptimal policy.

For the two-armed bandit with sigmoid parameterization, natural policy gradient and a perturbed minimum-variance baseline b=1−pt+ϵb=1-p_{t}+\epsilon, with ϵ<−1\epsilon<-1, there is a positive probability of choosing the suboptimal arm forever and diverging.

We can reuse the result for the two-armed bandit with constant baseline b<0b<0. Recall that for the proof to work, we only need θ\theta to move by at least a constant step δ>0\delta>0 in the negative direction at every iteration.

In detail, the update after picking the worst arm is θt+1=θt+α(1+ϵ1−pt)\theta_{t+1}=\theta_{t}+\alpha(1+\frac{\epsilon}{1-p_{t}}). So, if we choose ϵ<−1−δ\epsilon<-1-\delta for some δ>0\delta>0, we get the update step magnitude is δ+p1−p>δ\frac{\delta+p}{1-p}>\delta and hence the previous result applies (replace αb\alpha b by δ\delta). ∎

For the two-armed bandit with sigmoid parameterization, natural policy gradient and a perturbed minimum-variance baseline b=1−pt+ϵb=1-p_{t}+\epsilon, with ϵ∈(−1,0)\epsilon\in(-1,0), the policy converges to the optimal policy in probability.

Recall that the possible updates when the parameter is θt\theta_{t} are:

θt+1=θt+α(1−ϵσ(θt))\theta_{t+1}=\theta_{t}+\alpha(1-\frac{\epsilon}{\sigma(\theta_{t})}) if we choose action 1, with probability σ(θt)\sigma(\theta_{t})

θt+1=θt+α(1+ϵ1−σ(θt))\theta_{t+1}=\theta_{t}+\alpha(1+\frac{\epsilon}{1-\sigma(\theta_{t})}) if we choose action 2, with probability 1−σ(θt)1-\sigma(\theta_{t}).

First, we will partition the real line into three regions (AA, BB, and CC with a<b<ca<b<c for a∈A,b∈B,c∈Ca\in A,b\in B,c\in C), depending on the values of the updates. Then, each region will be analyzed separately.

We give an overview of the argument first. For region AA (θ\theta very negative), both updates are positive so θt\theta_{t} is guaranteed to increase until it reaches region BB.

For region CC (θ\theta very positive), sampling action 2 leads to the update α(1+ϵ1−σ(θt))\alpha(1+\frac{\epsilon}{1-\sigma(\theta_{t})}), which has large magnitude and results in θt+1\theta_{t+1} being back in region AA. So, once θt\theta_{t} is in CC, the agent needs to sample action 1 forever to stay there and converge to the optimal policy. This will have positive probability (using the same argument as the divergence proof for the two-armed bandit with constant baseline).

For region BB, the middle region, updates to θt\theta_{t} can make it either increase or decrease and stay in BB. For this region, we will show that θt\theta_{t} will eventually leave BB with probability 1 in a finite number of steps, with some lower-bounded probability of reaching AA.

Once we’ve established the behaviours in the three regions, we can argue that for any initial θ0\theta_{0} there is a positive probability that θt\theta_{t} will eventually reach region CC and take action 1 forever to converge. In the event that does not occur, then θt\theta_{t} will be sent back to AA and the agent gets another try at converging. Since we are looking at the behaviour when t→∞t\xrightarrow{}\infty, the agent effectively gets infinite tries at converging. Since each attempt has some positive probability of succeeding, convergence will eventually happen.

We now give additional details for each region.

To define region AA, we check when both updates will be positive. The update from action 1 is always positive so we are only concerned with the second update.

Hence, we set A=(−∞,σ−1(1+ϵ))A=(-\infty,\sigma^{-1}(1+\epsilon)). Since every update in this region increases θt\theta_{t} by at least a constant at every iteration, θt\theta_{t} will leave AA in a finite number of steps.

For region CC, we want to define it so that an update in the negative direction from any θ∈C\theta\in C will land back in AA. So C=[c,∞)C=[c,\infty) for some c≥σ−1(1+ϵ)c\geq\sigma^{-1}(1+\epsilon). By looking at the update from action 2, α(1+ϵ1−σ(θ))=α(1+ϵ(1+eθ))\alpha(1+\frac{\epsilon}{1-\sigma(\theta)})=\alpha(1+\epsilon(1+e^{\theta})), we see that it is equal to 0 at θ=σ−1(1+ϵ)\theta=\sigma^{-1}(1+\epsilon) but it is a decreasing function of θ\theta and it decreases at an exponential rate. So, eventually for θt\theta_{t} sufficiently large, adding this update will make θt+1∈A\theta_{t+1}\in A.

So let c=inf⁡{θ:θ+α(1−ϵ1−σ(θ)),θ≥σ−1(1+ϵ)}c=\inf\{\theta:\theta+\alpha\left(1-\frac{\epsilon}{1-\sigma(\theta)}\right),\theta\geq\sigma^{-1}(1+\epsilon)\}. Note that it is possible that c=σ−1(1+ϵ)c=\sigma^{-1}(1+\epsilon). If this is the case, then region BB does not exist.

When θt∈C\theta_{t}\in C, we know that there is a positive probability of choosing action 1 forever and thus converging (using the same proof as the two-armed bandit with constant baseline).

Finally, for the middle region B=[a,c)B=[a,c) (a=σ−1(1+ϵ)a=\sigma^{-1}(1+\epsilon)), we know that the updates for any θ∈B\theta\in B are uniformly bounded in magnitude by a constant uu.

We define a stopping time τ=inf⁡{t;θt≤a or θt≥c}\tau=\inf\{t;\theta_{t}\leq a\text{ or }\theta_{t}\geq c\}. This gives the first time θt\theta_{t} exits the region BB. Let “∧\land” denote the min operator.

The second line follows from substituting λ=−αt+c\lambda=-\alpha t+c. Note that the RHS goes to 0 as tt goes to ∞\infty.

Next, we continue from the LHS. Let θt∗=sup⁡0≤n≤tθn\theta^{*}_{t}=\sup_{0\leq n\leq t}\theta_{n}

Hence the probability the stopping time exceeds tt goes to and it is guaranteed to be finite almost surely.

Now, if θt\theta_{t} exits BB, there is some positive probability that it reached CC. We see this by considering that taking action 1 increases θ\theta by at least a constant, so the sequence of only taking action 11 until θt\theta_{t} reaches CC has positive probability. This is a lower bound on the probability of eventually reaching CC given that θt\theta_{t} is in BB.

Finally, we combine the results for all three regions to show that convergence happens with probability 1. Without loss of generality, suppose θ0∈A\theta_{0}\in A. If that is not the case, then keep running the process until either θt\theta_{t} is in AA or convergence occurs.

Let EiE_{i} be the event that θt\theta_{t} returns to AA after leaving it for the ii-th time. Then Ei∁E_{i}^{\complement} is the event that θt→∞\theta_{t}\xrightarrow{}\infty (convergence occurs). This is the case because, when θt∈C\theta_{t}\in C, those are the only two options and, when θt∈B\theta_{t}\in B we had shown that the process must exit BB with probability 1, either landing in AA or CC.

Next, we note that P(Ei∁)>0P(E_{i}^{\complement})>0 since, when θt\theta_{t} is in BB, the process has positive probability of reaching CC. Finally, when θt∈C\theta_{t}\in C, the process has positive probability of converging. Hence, P(Ei∁)>0P(E_{i}^{\complement})>0.

To complete the argument, whenever EiE_{i} occurs, then θt\theta_{t} is back in AA and will eventually leave it almost surely. Since the process is Markov and memoryless, Ei+1E_{i+1} is independent of EiE_{i}. Thus, by considering a geometric distribution with a success being EiCE^{C}_{i} occurring, EiCE_{i}^{C} will eventually occur with probability 1. In other words, θt\theta_{t} goes to +∞+\infty.

For the two-armed bandit with sigmoid parameterization, natural policy gradient and a perturbed minimum-variance baseline b=1−pt+ϵb=1-p_{t}+\epsilon, with ϵ=0\epsilon=0, the policy converges to the optimal policy with probability 1.

By directly writing the updates, we find that both updates are always equal to the expected natural policy gradient, so that θt+1=θt+α\theta_{t+1}=\theta_{t}+\alpha for any θt\theta_{t}. Hence θt→∞\theta_{t}\xrightarrow{}\infty as t→∞t\xrightarrow{}\infty with probability 1. ∎

For the two-armed bandit with sigmoid parameterization, natural policy gradient and a perturbed minimum-variance baseline b=1−pt+ϵb=1-p_{t}+\epsilon, with ϵ∈(0,1)\epsilon\in(0,1), the policy converges to the optimal policy in probability.

The overall idea is to ensure that the updates are always positive for some region A={θ:θ>θA}A=\{\theta:\theta>\theta_{A}\} then show that we reach this region with probability 1.

Recall that the possible updates when the parameter is θt\theta_{t} are:

θt+1=θt+α(1−ϵσ(θt))\theta_{t+1}=\theta_{t}+\alpha(1-\frac{\epsilon}{\sigma(\theta_{t})}) if we choose action 1, with probability σ(θt)\sigma(\theta_{t})

θt+1=θt+α(1+ϵ1−σ(θt))\theta_{t+1}=\theta_{t}+\alpha(1+\frac{\epsilon}{1-\sigma(\theta_{t})}) if we choose action 2, with probability 1−σ(θt)1-\sigma(\theta_{t}).

First, we observe that the update for action 2 is always positive. As for action 1, it is positive whenever p≥ϵp\geq\epsilon, equivalently θ≥θA\theta\geq\theta_{A}, where θA=σ−1(ϵ)\theta_{A}=\sigma^{-1}(\epsilon). Call this region A={θ:θ>θA(=σ−1(ϵ))}A=\{\theta:\theta>\theta_{A}(=\sigma^{-1}(\epsilon))\}. If θt∈A\theta_{t}\in A, then we can find a δ>0\delta>0 such that the update is always greater than δ\delta in the positive direction, no matter which action is sampled. So, using the same argument as for the ϵ=0\epsilon=0 case with steps of +δ+\delta, we get convergence to the optimal policy (with only constant regret).

In the next part, we show that the iterates will enter the good region AA with probability 1 to complete the proof. We may assume that θ0<θA\theta_{0}<\theta_{A} since if that is not the case, we are already done. The overall idea is to create a transformed process which stops once it reaches AA and then show that the stopping time is finite with probability 1. This is done using the fact that the expected step is positive (+α+\alpha) along with Markov’s inequality to bound the probability of going too far in the negative direction.

Since we only stop the process {θt∧τ}\{\theta_{t\land\tau}\} after reaching AA, then we need to compute the largest value θt∧τ\theta_{t\land\tau} can take after making an update which brings us inside the good region. In other words, we need to compute sup⁡θ{θ+α(1+ϵ1−σ(θ)):θ∈A∁}\sup_{\theta}\{\theta+\alpha(1+\frac{\epsilon}{1-\sigma(\theta)}):\theta\in A^{\complement}\}. Fortunately, since the function to maximize is an increasing function of θ\theta, the supremum is easily obtained by choosing the largest possible θ\theta, that is θ=σ−1(ϵ)\theta=\sigma^{-1}(\epsilon). This gives us that C=θA+UAC=\theta_{A}+U_{A}, where UA=α(1+ϵ1−ϵ)U_{A}=\alpha(1+\frac{\epsilon}{1-\epsilon}).

Applying Markov’s inequality, for λ>0\lambda>0 we have:

Note that the RHS goes to 0 as t→∞t\xrightarrow{}\infty. We then manipulate the LHS to eventually get an upper bound on P(t≤τ)P(t\leq\tau).

Since the first line goes to , the last line goes to and hence we have that θt\theta_{t} will enter the good region with probability 1.

We follow the same argument as in the ϵ∈(0,1)\epsilon\in(0,1) case with a stopping time defined as τ=inf⁡{t:θt>c}\tau=\inf\{t:\theta_{t}>c\} and using θA=c\theta_{A}=c, to show that

B.3 Convergence with vanilla policy gradient

We now proceed to prove the necessary requirements.

Assuming bounded rewards and a bounded baseline, the martingale {Xt}\{X_{t}\} associated with vanilla policy gradient has bounded increments

Then, the stochastic gradient estimate is

Assuming bounded rewards and a bounded baseline, the martingale {Xt}\{X_{t}\} associated with policy gradient with importance sampling distribution qq such that min⁡{q,1−q}≥ϵ>0\min\{q,1-q\}\geq\epsilon>0 has bounded increments

Let us also call ϵ>0\epsilon>0 the lowest probability of sampling an arm under qq.

Then, the stochastic gradient estimate is

As the rewards are bounded, ∃Ri>0\exists R_{i}>0 such that ∣ri∣≤Ri|r_{i}|\leq R_{i} for all ii

We call non-singular importance sampling any importance sampling distribution so that the probability of each action is bounded below by a strictly positive constant.

For vanilla policy gradient and policy gradient with nonsingular importance sampling, the expected parameter θt\theta_{t} has infinite limit. i.e. if μ1≠μ0\mu_{1}\neq\mu_{0},

In other words, the expected parameter value converges to the optimal arm.

We reason by contradiction. The contradiction stems from the fact that on one hand we know θt\theta_{t} will become arbitrarily large with tt with high probability as this setting satisfies the convergence conditions of stochastic optimization. On the other hand, because of Azuma’s inequality, if the average θt\theta_{t} were finite, we can show that θt\theta_{t} cannot deviate arbitrarily far from its mean with probability 1. The contradiction will stem from the fact that the expected θt\theta_{t} cannot have a finite limit.

We have θt−θ0=∑i=0tαigi\theta_{t}-\theta_{0}=\sum_{i=0}^{t}\alpha_{i}g_{i}. Thus

where Δ=μ1−μ0>0\Delta=\mu_{1}-\mu_{0}>0 the optimality gap between the value of the arms. As it is a sum of positive terms, its limit is either positive and finite or +∞+\infty.

As ∑i=0∞αi2=γ\sum_{i=0}^{\infty}\alpha_{i}^{2}=\gamma, using Azuma-Hoeffing’s inequality

where ci=αiCc_{i}=\alpha_{i}C like in the proposition above. And for M>∣θ0∣+β+2Cγlog⁡2M>|\theta_{0}|+\beta+2C\sqrt{\gamma\log 2} we have

As ∑i=0∞ci=γC2\sum_{i=0}^{\infty}c_{i}=\gamma C^{2} , we have

i.e for any MM large enough, the probability that {θt}\{\theta_{t}\} is bounded by MM is bigger than a strictly positive constant.

Because policy gradient with diminishing stepsizes satisfies the convergence conditions defined by Bottou et al. , we have that

Here we show that θt\theta_{t} cannot be bounded by any constant with non-zero probability at t→∞t\to\infty. This contradicts the previous conclusion.

Policy gradient with stepsizes satisfying the Robbins-Monro conditions (∑tαt=∞,∑tαt2<∞\sum_{t}\alpha_{t}=\infty,\sum_{t}\alpha_{t}^{2}<\infty) converges to the optimal arm.

Note that this convergence result addresses the stochastic version of policy gradient, which is not covered by standard results for stochastic gradient algorithms due to the nonconvexity of the objective.

Appendix C Multi-armed bandit theory

The example of convergence to a suboptimal policy for the minimum-variance baseline and convergence to the optimal policy for a gap baseline are outlined in the next two subsections. ∎

Consider a three-armed bandit with rewards of 1, 0.7 and 0. Let the policy be parameterized by a softmax (πi∝eθi\pi_{i}\propto e^{\theta_{i}}) and optimized using natural policy gradient paired with the mininum-variance baseline. If the policy is initialized to be uniform random, there is a nonzero probability of choosing a suboptimal action forever and converging to a suboptimal policy.

The policy probabilities are given by πi=eiθ∑jejθ\pi_{i}=\frac{e^{\theta}_{i}}{\sum_{j}e^{\theta}_{j}} for i=1,2,3i=1,2,3. Note that this parameterization is invariant to shifting all θi\theta_{i} by a constant.

Next, we compute the minimum-variance baseline. Here, we have two main options. We can find the baseline that minimizes the variance of the sampled gradients gig_{i}, the “standard” choice, or we can instead minimize the variance of the sampled natural gradients, F−1giF^{-1}g_{i}. We analyze both cases separately.

where wi=((1−πi)2+πj2+πk2)πiw_{i}=((1-\pi_{i})^{2}+\pi_{j}^{2}+\pi_{k}^{2})\pi_{i}.

The proof idea is similar to that of the two-armed bandit. Recall that the rewards for the three actions are 1, 0.7 and 0. We will show that this it is possible to choose action 2 (which is suboptimal) forever.

To do so, it is enough to show that we make updates that increase θ2\theta_{2} by at least δ\delta at every step (and leave θ1\theta_{1} and θ3\theta_{3} the same). In this way, the probability of choosing action 2 increases sufficiently fast, that we can use the proof for the two-armed bandit to show that the probability of choosing action 2 forever is nonzero.

In more detail, suppose that we have established that, at each step, θ2\theta_{2} increases by at least δ\delta. The policy starts as the uniform distribution so we can choose any initial θ\theta as long as three components are the same (θ1=θ2=θ3\theta_{1}=\theta_{2}=\theta_{3}). Choosing the initialization θi=−log⁡(\nicefrac12)\theta_{i}=-\log(\nicefrac{{1}}{{2}}) for all ii, we see that π2=eθ2∑i=13θi=eθ21+eθ2=σ(θ2)\pi_{2}=\frac{e^{\theta_{2}}}{\sum_{i=1}^{3}\theta_{i}}=\frac{e^{\theta_{2}}}{1+e^{\theta_{2}}}=\sigma(\theta_{2}) where σ(.)\sigma(.) is the sigmoid function. Since at the nn-th step, θ2>θ0+nδ\theta_{2}>\theta_{0}+n\delta, we can reuse the proof for the two-armed bandit to show Pr(action 2 forever)>0Pr(\text{action 2 forever})>0.

To complete the proof, we need to show that the updates are indeed lower bounded by a constant. Every time we sample action 2, the update is θ←θ+α(r2−b∗)(λe+1π2e2)\theta\xleftarrow{}\theta+\alpha(r_{2}-b^{*})(\lambda e+\frac{1}{\pi_{2}}e_{2}). We can choose any value of λ\lambda since they produce the same policy after an update due to the policy’s invariance to a constant shift of all the parameters. We thus choose λ=0\lambda=0 for simplicity. In summary, an update does θ2←θ2+α(r2−b∗)1π2\theta_{2}\xleftarrow{}\theta_{2}+\alpha(r_{2}-b^{*})\frac{1}{\pi_{2}} and leaves the other parameters unchanged.

In the next part, we use induction to show the updates are lower bounded at every step. For the base case, we need r2−b∗>δr_{2}-b^{*}>\delta for some δ>0\delta>0. Since we initialize the policy to be uniform, we can directly compute the value of b∗≈0.57b^{*}\approx 0.57, so the condition is satisfied for, say, δ=0.1\delta=0.1.

For the inductive case, we assume that r2−b∗>δr_{2}-b^{*}>\delta for δ>0\delta>0 and we will show that r2−b+∗>δr_{2}-b^{*}_{+}>\delta also, where b+∗b^{*}_{+} is the baseline after an update. It suffices to show that b+∗≤b∗b^{*}_{+}\leq b^{*}.

To do so, we examine the ratio w2w1\frac{w_{2}}{w_{1}} in b∗b^{*} and show that this decreases. Let (w2w1)+\left(\frac{w_{2}}{w_{1}}\right)_{+} be the ratio after an update and let c=r2−b∗c=r_{2}-b^{*}.

The last line follows by considering the function f(z)=ex−z+ey−zf(z)=e^{x-z}+e^{y-z} for a fixed x≤yx\leq y. f′(z)=−ex−z+ey+z>0f^{\prime}(z)=-e^{x-z}+e^{y+z}>0 for all zz, so f(z)f(z) is an increasing function. By taking x=2θ2x=2\theta_{2} and y=2θ3y=2\theta_{3} (θ2≥θ3\theta_{2}\geq\theta_{3}), along with the fact that cπ2>δ\frac{c}{\pi_{2}}>\delta (considering these as zz values), then we we see that the denominator has increased in the last line and the inequality holds.

By the same argument, recalling that δ>0\delta>0, we have that the last ratio is less than 11. Hence, (w2w1)+<(w2w1)\left(\frac{w_{2}}{w_{1}}\right)_{+}<\left(\frac{w_{2}}{w_{1}}\right).

Returning to the baseline, b∗=w1r1+w2r2+w3r3w1+w2+w3b^{*}=\frac{w_{1}r_{1}+w_{2}r_{2}+w3r_{3}}{w_{1}+w_{2}+w_{3}}. We see that this is a convex combination of the rewards. Focusing on the (normalized) weight of r2r_{2}:

The first line follows since w1=w3w_{1}=w_{3} and the second by dividing the numerator and denominator by w1w_{1}. This is an increasing function of \nicefracw2w1\nicefrac{{w_{2}}}{{w_{1}}} so decreasing the ratio will decrease the normalized weight given to r2r_{2}. This, in turn, increases the weight on the other two rewards equally. As such, since the value of the baseline is under r2=0.7r_{2}=0.7 (recall it started at b∗≈0.57b^{*}\approx 0.57) and the average of r1r_{1} and r3r_{3} is 0.50.5, the baseline must decrease towards 0.50.5.

Thus, we have shown that the gap between r2r_{2} and b∗b^{*} remains at least δ\delta and this completes the proof for the minimum-variance baseline of the gradients.

Next, we tackle the minimum-variance baseline for the updates. Recall that the natural gradient updates are of the form xi=λe+1πieix_{i}=\lambda e+\frac{1}{\pi_{i}}e_{i} for action ii where ee is a vector of ones and eie_{i} is the ii-th standard basis vector.

The minimum-variance baseline for updates is given by

We have that ∣∣xi∣∣2=2λ2=(λ+1πi)2||x_{i}||^{2}=2\lambda^{2}=(\lambda+\frac{1}{\pi_{i}})^{2}. At this point, we have to choose which value of λ\lambda to use since it will affect the baseline. The minimum-norm solution is a common choice (corresponding to use of the Moore-Penrose pseudoinverse of the Fisher information instead of the inverse). We also take a look at fixed values of λ\lambda, but we find that this requires an additional assumption 3λ2<\nicefrac1π123\lambda^{2}<\nicefrac{{1}}{{\pi_{1}^{2}}}.

First, we consider the minimum-norm solution. We find that the minimum-norm solution gives 23πi2\frac{2}{3\pi_{i}^{2}} for λ=−13πi2\lambda=\frac{-1}{3\pi_{i}^{2}}.

We will reuse exactly the same argument as for the minimum-variance baseline for the gradients. The only difference is the formula for the baseline, so all we need to check is the that the ratio of the weights of the rewards decreases after one update, which implies that the baseline decreases after an update.

So we have the weights wi=1πiw_{i}=\frac{1}{\pi_{i}} and the ratio is

for c=α(r2−b∗)c=\alpha(r_{2}-b^{*}), which is less than the initial ratio. This completes the case where we use the minimum-norm update.

The weights are given by wi=(2λ2+(λ+1πi)2)πiw_{i}=(2\lambda^{2}+(\lambda+\frac{1}{\pi_{i}})^{2})\pi_{i}

We know that after an update π2\pi_{2} will increase and π1\pi_{1} will decrease. So, we check the partial derivative of the ratio to assess its behaviour after an update.

We need this to be an increasing function in π1\pi_{1} so that a decrease in π1\pi_{1} implies a decrease in the ratio. This is true when 3λ2<\nicefrac1π123\lambda^{2}<\nicefrac{{1}}{{\pi_{1}^{2}}}. So, to ensure the ratio decreases after a step, we need an additional assumption on λ\lambda and π1\pi_{1}, which is that 3λ2<\nicefrac1π123\lambda^{2}<\nicefrac{{1}}{{\pi_{1}^{2}}}. This is notably always satisfied for λ=0\lambda=0.

C.2 Convergence with gap baselines

For a three-arm bandit with deterministic rewards, choosing the baseline bb so that r1>b>r2r_{1}>b>r_{2} where r1r_{1} (resp. r2r_{2}) is the value of the optimal (resp. second best) arm, natural policy gradient converges to the best arm almost surely.

Let us define Δi=ri−b\Delta_{i}=r_{i}-b which is striclty positive for i=1i=1, stricly negative otherwise. Then the gradient on the parameter θi\theta^{i} of arm ii

For the martingale Xt=αΔ1t+θ01−θt1X_{t}=\alpha\Delta_{1}t+\theta^{1}_{0}-\theta^{1}_{t}, we have

thus satisfying the bounded increments assumption of Azuma’s inequality. We can therefore show

This shows that θt1\theta^{1}_{t} converges to +∞+\infty almost surely while the θti,i>1\theta^{i}_{t},i>1 remain bounded by θ0i\theta^{i}_{0}, hence we converge to the optimal policy almost surely.

C.3 Convergence with off-policy sampling

We show that using importance sampling with a separate behaviour policy can guarantee convergence to the optimal policy for a three-armed bandit.

Suppose we have an nn-armed bandit where the rewards for choosing action ii are distributed according to PiP_{i}, which has finite support and expectation rir_{i}. Assume at the tt-th round the behaviour policy selects each action ii with probability μt(i)\mu_{t}(i). Then, if we draw action ii, the stochastic estimator for the natural policy gradient with importance sampling is equal to

with probability μt(i)\mu_{t}(i) and RiR_{i} drawn from PiP_{i}.

By subtracting the expected updates, we define the multivariate martingale Xt=θt−θ0−αΔtX_{t}=\theta_{t}-\theta_{0}-\alpha\Delta t. Note that the ii-th dimension XtiX^{i}_{t} is a martingale for all ii.

Suppose we have bounded rewards and a bounded baseline and a behaviour policy selecting all actions with probability at least ϵt\epsilon_{t} at round tt. Then, the martingale {Xt}\{X_{t}\} associated with natural policy gradient with importance sampling has bounded increments

for all dimensions ii and some fixed constant CC.

The updates and XtX_{t} are defined as above.

Thus ∣Xti−Xt−1i∣≤Cαϵt|X_{t}^{i}-X_{t-1}^{i}|\leq\frac{C\alpha}{\epsilon_{t}} for all ii. ∎

Next, we choose δ∈(0,1)\delta\in(0,1) such that (1−δ)Δ1>(1+δ)Δj(1-\delta)\Delta_{1}>(1+\delta)\Delta_{j}. We apply Azuma’s inequality to Xt1X^{1}_{t}, the martingale associated to the optimal action, with ϵ=αδΔit\epsilon=\alpha\delta\Delta_{i}t.

Similarly, we can apply Azuma’s inequality to actions i≠1i\neq 1 and obtain

Letting AA be the event θt1≤θ01+α(1−δ)Δ1t\theta_{t}^{1}\leq\theta_{0}^{1}+\alpha(1-\delta)\Delta_{1}t and BiB_{i} be the event that θti−θ0i≥α(1+δ)Δit\theta_{t}^{i}-\theta_{0}^{i}\geq\alpha(1+\delta)\Delta_{i}t for i≠1i\neq 1, we can apply the union bound to get

The RHS goes to when ∑t≥0tϵt2=∞\sum_{t\geq 0}t\epsilon_{t}^{2}=\infty.

Notice that A∁A^{\complement} is the event θt1>θ01+α(1−δ)Δ1t\theta_{t}^{1}>\theta_{0}^{1}+\alpha(1-\delta)\Delta_{1}t and B∁B^{\complement} is the event θti<θ0i+α(1+δ)Δit\theta_{t}^{i}<\theta_{0}^{i}+\alpha(1+\delta)\Delta_{i}t. Then, inspecting the difference between θt1\theta^{1}_{t} and θti\theta^{i}_{t}, we have

By our assumption on δ\delta, the term within the parenthesis is positive and hence the difference grows to infinity as t→∞t\xrightarrow{}\infty. Taken together with the above probability bound, we have convergence to the optimal policy in probability.

Appendix D Other results

For completeness, we include a derivation of the minimum-variance baseline for the trajectory policy gradient estimate (REINFORCE) and the state-action policy gradient estimator (with the true state-action values).

The second equality follows since the baseline doesn’t affect the bias of the estimator. Thus, since the second term does not contain bb, we only need to optimize the first term.

Taking the derivative with respect to bb, we have:

The minimum of the variance can then be obtained by finding the baseline b∗b^{\ast} for which the gradient is , i.e

So that we only need to take into account the first term.

Therefore the baseline that minimizes the variance for each state is

Note that for the natural policy gradient, the exact same derivation holds and we obtain that

D.2 Natural policy gradient for softmax policy in bandits

We derive the natural policy gradient estimator for the multi-armed bandit with softmax parameterization.

D.3 Link between minimum variance baseline and value function

We show here a simple link between the minimum variance baseline and the value function. While we prove this for the REINFORCE estimator, a similar relation holds for the state-action value estimator.

D.4 Variance of perturbed minimum-variance baselines

Here, we show that the variance of the policy gradient estimator is equal for baselines b+=b∗+ϵb_{+}=b^{*}+\epsilon and b−=b∗−ϵb_{-}=b^{*}-\epsilon, where ϵ>0\epsilon>0 and b∗b^{*} is the minimum-variance baseline. We will use the trajectory estimator here but the same argument applies for the state-action estimator.

We have g=R(τ)−b)∇log⁡π(τ)g=R(\tau)-b)\nabla\log\pi(\tau) and the variance is given by

where the third line follows since the baseline does not affect the bias of the policy gradient.

D.5 Baseline for natural policy gradient and softmax policies

For a softmax policy, this is: g=(Ri−b)(1πθ(i)ei+λe)g=(R_{i}-b)(\frac{1}{\pi_{\theta}(i)}e_{i}+\lambda e), where eie_{i} is a vector containing a 1 at position ii and 0 otherwise, ee is a vector of all one and λ\lambda is an arbitrary constant. Checking the expectation, we see that

So the baseline only causes a constant shift in all the parameters. But for the softmax parameterization, adding a constant to all the parameters does not affect the policy, so the updates remained unbiased. In other words, we can always add a constant vector to the update to ensure the expected update to θ\theta does not change, without changing the policy obtained after an update.

D.6 Natural policy gradient estimator for MDPs

In this section, we provide a detailed derivation of the natural policy gradient with QQ-values estimate used in the MDP experiments.

Together, our estimate of the policy gradient is

Since we have a tabular representation, F(si)F(s_{i}) is a block diagonal matrix where each block corresponds to one state and F(si)F(s_{i}) contains nonzero entries only for the block corresponding to state sis_{i}. Hence, the sum is a block diagonal matrix with nonzero entries corresponding to the blocks of states s0,...,sT−1s_{0},...,s_{T-1} and we can invert the sum by inverting the blocks. It follows that the inverse of the sum is the sum of the inverses.

Finally, we notice that ∇log⁡π(ai∣si)\nabla\log\pi(a_{i}|s_{i}) is a vector of zeros except for the entries corresponding to state sis_{i}. So, F(sj)−1∇log⁡π(ai∣si)F(s_{j})^{-1}\nabla\log\pi(a_{i}|s_{i}) is nonzero only if i=ji=j giving us our final estimator

Note that this is the same as applying the natural gradient update for bandits at each sampled state ss, where the rewards for each action is given by Qπ(s,a)Q_{\pi}(s,a).