Improved Regret Bound and Experience Replay in Regularized Policy Iteration

Model-free reinforcement learning (RL) algorithms combined with powerful function approximation have achieved impressive performance in a variety of application domains over the last decade. Unfortunately, the theoretical understanding of such methods is still quite limited. In this work, we study single-trajectory learning in infinite-horizon undiscounted Markov decision processes (MDPs), also known as average-reward MDPs, which capture tasks such as routing and the control of physical systems.

One line of works with performance guarantees for the average-reward setting follows the “online MDP” approach proposed by Even-Dar et al. (2009), where the agent selects policies by running an online learning algorithm in each state, typically mirror descent. The resulting algorithm is a version of approximate policy iteration (API), which alternates between (1) estimating the state-action value function (or Q-function) of the current policy and (2) setting the next policy to be optimal w.r.t. the sum of all previous Q-functions plus a regularizer. Note that, by contrast, standard API sets the next policy only based on the most recent Q-function. The policy update can also be written as maximizing the most recent Q-function minus KL-divergence to the previous policy, which is somewhat similar to recently popular versions of API (Schulman et al., 2015, 2017; Achiam et al., 2017; Abdolmaleki et al., 2018; Song et al., 2019a).

The original work of Even-Dar et al. (2009) studied this scheme with known dynamics, tabular representation, and adversarial reward functions. More recent works (Abbasi-Yadkori et al., 2019; Hao et al., 2020; Wei et al., 2020a) have adapted the approach to the case of unknown dynamics, stochastic rewards, and value function approximation. With linear value functions, the Politex algorithm of Abbasi-Yadkori et al. (2019) achieves $O(T^{3/4})$ high-probability regret in ergodic MDPs, and the results only scale in the number of features rather than states. Wei et al. (2020a) later show an $O(\sqrt{T})$ bound on expected regret for a similar algorithm named MDP-EXP2. In this work, we revisit the analysis of Politex and show that it can be sharpened to $O(\sqrt{T})$ under nearly identical assumptions, resulting in the first $O(\sqrt{T})$ high-probability regret bound for a computationally efficient algorithm in this setting.

In addition to improved analysis, our work also addresses practical implementation of Politex with neural networks. The policies produced by Politex in each iteration require access to the sum of all previous Q-function estimates. With neural network function approximation, exact implementation requires us to keep all past networks in memory and evaluate them at each step, which is inefficient in terms of memory and computation. Some practical implementation choices include subsampling Q-functions and/or optimizing a KL-divergence regularized objective w.r.t. a parametric policy at each iteration. We propose an alternative approach, where we approximate the average of all past Q-functions by training a single network on a replay buffer with past data. We demonstrate that this choice often outperforms other approximate implementations, especially in terms of run-time. When available memory is constrained, we propose to subsample transitions using the notion of coresets (Bachem et al., 2017).

Our work also provides a novel perspective on the benefits of experience replay. Experience replay is a standard tool for stabilizing learning in modern deep RL, and typically used in off-policy methods like Deep Q-Networks (Mnih et al., 2013), as well as “value gradient” methods such as DDPG (Lillicrap et al., 2016) and SVG (Heess et al., 2015). A different line of on-policy methods typically does not rely on experience replay; instead, learning is stabilized by constraining consecutive policies to be close in terms of KL divergence (Schulman et al., 2015; Song et al., 2019a; Degrave et al., 2019). We observe that both experience replay and KL-divergence regularization can be viewed as approximate implementations of Politex. Thus, we provide a theoretical justification for using experience replay in API, as an approximate implementation of online learning in each state. Note that this online-learning view differs from the commonly used justifications for experience replay, namely that it “breaks temporal correlations” (Schaul et al., 2016; Mnih et al., 2013). Our analysis also suggests a new objective for subsampling or priority-sampling transitions in the replay buffer, which differs priority-sampling objectives of previous work (Schaul et al., 2016).

In summary, our main contributions are (1) an improved analysis of Politex, showing an $O(\sqrt{T})$ regret bound under the same assumptions as the original work, and (2) an efficient implementation that also offers a new perspective on the benefits of experience replay.

We consider learning in infinite-horizon undiscounted ergodic MDPs $(\mathcal{X},\mathcal{A},r,P)$, where $\mathcal{X}$ is the state space, $\mathcal{A}$ is a finite action space, $r:\mathcal{X}\times\mathcal{A}\to[0,1]$ is an unknown reward function, and $P:\mathcal{X}\times\mathcal{A}\to\Delta_{\mathcal{X}}$ is the unknown probability transition function. A policy $\pi:\mathcal{X}\to\Delta_{\mathcal{A}}$ is a mapping from a state to a distribution over actions. Let $\{(x_{t}^{\pi},a_{t}^{\pi})\}_{t=1}^{\infty}$ denote the state-action sequence obtained by following policy $\pi$. The expected average reward of policy $\pi$ is defined as

Let $\mu_{\pi}$ denote the stationary state distribution of a policy $\pi$, satisfying $\mu_{\pi}=\mathbb{E}_{x\sim\mu_{\pi},a\sim\pi}[P(\cdot|x,a)]$. We will sometimes write $\mu_{\pi}$ as a vector, and use $\nu_{\pi}=\mu_{\pi}\otimes\pi$ to denote the stationary state-action distribution. In ergodic MDPs, $J_{\pi}$ and $\mu_{\pi}$ are well-defined and independent of the initial state. The optimal policy $\pi_{*}$ is a policy that maximizes the expected average reward. We will denote by $J_{*}$ and $\mu_{*}$ the expected average reward and stationary state distribution of $\pi_{*}$, respectively.

The value function of a policy $\pi$ is defined as:

The state-action value function $Q_{\pi}(x,a)$ is defined as

Notice that

Equations (2.3) and (2.4) are known as the Bellman equation. If we do not use the definition of $V_{\pi}(x)$ in Eq. (2.2) and instead solve for $V_{\pi}(x)$ and $Q_{\pi}(x,a)$ using the Bellman equation, we can see that the solutions are unique up to an additive constant. Therefore, in the following, we may use $V_{\pi}(x)$ and $Q_{\pi}(x,a)$ to denote the same function up to a constant. We will use the shorthand notation $Q_{\pi}(x,\pi^{\prime})=\mathbb{E}_{a\sim\pi^{\prime}(\cdot|x)}[Q_{\pi}(x,a)]$; note that $Q_{\pi}(x,\pi)=V_{\pi}(x)$.

The agent interacts with the environment as follows: at each round $t$, the agent observes a state $x_{t}\in\mathcal{X}$, chooses an action $a_{t}\sim\pi_{t}(\cdot|x_{t})$, and receives a reward $r_{t}:=r(x_{t},a_{t})$. The environment then transitions to the next state $x_{t+1}$ with probability $P(x_{t+1}|x_{t},a_{t})$. Recall that $\pi_{*}$ is the optimal policy and $J_{*}$ is its expected average reward. The regret of an algorithm with respect to this fixed policy is defined as

The learning goal is to find an algorithm that minimizes the long-term regret $R_{T}$.

Our analysis will require the following assumption on the mixing rate of policies.

Assumption 2.1 implies that for any pair of distributions $\nu,\nu^{\prime}$, $\|(\nu-\nu^{\prime})^{\top}H_{\pi}\|_{1}\leq\gamma\|\nu-\nu^{\prime}\|_{1}$; see Lemma A.1 in Appendix A for a proof.

Regret bounds for average-reward MDPs. Most no-regret algorithms for infinite-horizon undiscounted MDPs are only applicable to tabular representations and model-based (Bartlett, 2009; Jaksch et al., 2010; Ouyang et al., 2017; Fruit et al., 2018; Jian et al., 2019; Talebi and Maillard, 2018). For weakly-communicating MDPs with diameter $D$, these algorithms nearly achieve the minimax lower bound $\Omega(\sqrt{D|\mathcal{X}||\mathcal{A}|T})$ (Jaksch et al., 2010) with high probability. Wei et al. (2020b) provide model-free algorithms with regret bounds in the tabular setting. In the model-free setting with function approximation, the Politex algorithm (Abbasi-Yadkori et al., 2019) achieve $O(d^{1/2}T^{3/4})$ regret in uniformly ergodic MDPs, where $d$ is the size of the compressed state-action space. Hao et al. (2020) improve these results to $O(T^{2/3})$. More recently, Wei et al. (2020a) present three algorithms for average-reward MDPs with linear function approximation. Among these, FOPO achieves $O(\sqrt{T})$ regret but is computationally inefficient, and OLSVI.FH is efficient but obtains $O(T^{3/4})$ regret. The MDP-EXP2 algorithm is computationally efficient, and under similar assumptions as in Abbasi-Yadkori et al. (2019) it obtains a $O(\sqrt{T})$ bound on expected regret (a weaker guarantee than the high-probability bounds in other works). Our analysis shows a high-probability $O(\sqrt{T})$ regret bound under the same assumptions.

KL-regularized approximate policy iteration. Our work is also related to approximate policy iteration algorithms which constrain each policy to be close to the previous policy in the sense of KL divergence. This approach was popularized by TRPO (Schulman et al., 2015), where it was motivated as an approximate implementation of conservative policy iteration (Kakade and Langford, 2002). Some of the related subsequent works include PPO (Schulman et al., 2017), MPO (Abdolmaleki et al., 2018), V-MPO (Song et al., 2019a), and CPO (Achiam et al., 2017). While these algorithms place a constraint on consecutive policies and are mostly heuristic, another line of research shows that using KL divergence as a regularizer has a theoretical justification in terms of either regret guarantees (Abbasi-Yadkori et al., 2019; Hao et al., 2020; Wei et al., 2020b, a) or error propagation (Vieillard et al., 2020a, b).

Experience replay. Experience replay (Lin, 1992) is one of the central tools for achieving good performance in deep reinforcement learning. While it is mostly used in off-policy methods such as deep Q-learning (Mnih et al., 2013, 2015), it has also shown benefits in value-gradient methods (Heess et al., 2015; Lillicrap et al., 2016), and has been used in some variants of KL-regularized policy iteration (Abdolmaleki et al., 2018; Tomar et al., 2020). Its success has been attributed to removing some temporal correlations from data fed to standard gradient-based optimization algorithms. Schaul et al. (2016) have shown that non-uniform replay sampling based on the Bellman error can improve performance. Unlike these works, we motivate experience replay from the perspective of online learning in MDPs (Even-Dar et al., 2009) with the goal of approximating the average of past value functions well.

Continual learning. Continual learning (CL) is the paradigm of learning a classifier or regressor that performs well on a set of tasks, where each task corresponds to a different data distribution. The tasks are observed sequentially, and the learning goal is to avoid forgetting past tasks without storing all the data in memory. This is quite similar to our goal of approximating the average of sequentially-observed Q-functions, where the data for approximating each Q-function has different distribution. In general, approaches to CL can be categorized as regularization-based (Kirkpatrick et al., 2017; Zenke et al., 2017; Farajtabar et al., 2020; Yin et al., 2020), expansion-based (Rusu et al., 2016), and replay-based (Lopez-Paz and Ranzato, 2017; Chaudhry et al., 2018; Borsos et al., 2020), with the approaches based on experience replay typically having superior performance over other methods.

Our algorithm is similar to the Politex schema (Abbasi-Yadkori et al., 2019). In each phase $k$, Politex obtains an estimate $\widehat{Q}_{\pi_{k}}$ of the action-value function $Q_{\pi_{k}}$ of the current policy $\pi_{k}$, and then sets the next policy using the mirror descent update rule:

where $\mathcal{H}(\cdot)$ is the entropy function. When the functions $\{\widehat{Q}_{\pi_{i}}\}_{i=1}^{k}$ are tabular or linear, the above update can be implemented efficiently by simply summing all the table entries or weight vectors. However, with neural network function approximation, we need to keep all networks in memory and evaluate them at each step, which quickly becomes inefficient in terms of storage and computation. Some of the efficient implementations proposed in literature include keeping a subset of the action-value functions (Abbasi-Yadkori et al., 2019), and using a parametric policy and optimizing the KL-regularized objective w.r.t. the parameters over the available data (Tomar et al., 2020).

Our proposed method, presented in Algorithm 1, attempts to directly approximate the average of all previous action-value functions

To do so, we only use a single network, and continually train it to approximate $Q_{k}$. At each iteration $k$, we obtain a dataset of tuples $\mathcal{D}_{k}=\{(x_{t},a_{t},R_{t})\}$, where $R_{t}$ is the empirical return from the state-action pair $(x_{t},a_{t})$. We initialize $\widehat{Q}_{k}$ to $\widehat{Q}_{k-1}$ and update it by minimizing the squared error over the union of $\mathcal{D}_{k}$ and the replay buffer $\mathcal{R}$. We then update the replay buffer with all or a subset of data in $\mathcal{D}_{k}$.

In the sequel, in Section 5, we first show that by focusing on estimating the average Q-function, we can improve the regret bound of Politex under nearly identical assumption; in Section 6, we instantiate this bound for linear value functions, where we can estimate the average Q-function simply by weight averaging; in Section 7, we focus on practical implementations, in particular, we discuss the limitations of weight averaging, and provide details on how to leverage replay data when using non-linear function approximation; in Section 8 we present our experimental results; and in Section 9 we make final remarks.

In this section, we revisit the regret analysis of Politex (Abbasi-Yadkori et al., 2019) in ergodic average-reward MDPs, and show that it can be improved from $O(T^{3/4})$ to $O(\sqrt{T})$ under similar assumptions. Our analysis relies in part on a simple modification of the regret decomposition. Namely, instead of including the estimation error of each value-function $Q_{\pi_{k}}$ in the regret, we consider the error in the running average $Q_{k}$. When this error scales as $O(1/\sqrt{k})$ in a particular weighted norm, the regret of Politex is $O(\sqrt{T})$. As we show in the following section, this bound can be instantiated for linear value functions under the same assumptions as Abbasi-Yadkori et al. (2019).

For the purpose of our algorithm, functions $\widehat{Q}_{\pi_{k}}$ are unique up to a constant. Thus we can equivalently assume that $\|\widehat{Q}_{\pi_{k}}(x,\cdot)\|_{\infty}\leq Q_{\max}$ for all $x$.

Define $\widehat{V}_{\pi_{k}}(x):=\widehat{Q}_{\pi_{k}}(x,\pi_{k})$. Let $V_{k}:=\frac{1}{k}\sum_{i=1}^{k}V_{\pi_{i}}(x)$ and $\widehat{V}_{k}(x):=\sum_{i=1}^{k}\widehat{V}_{\pi_{i}}(x)$ be the average of the state-value functions and its estimate. We will require the estimation error of the running average to scale as in the following assumption.

Define $S_{\delta}(|\mathcal{A}|,\mu_{*})$ as in Abbasi-Yadkori et al. (2019):

We bound the regret of Politex in the following theorem. Here, recall that $\tau$ is the length of each phase, and $\gamma$ and $\eta$ are defined in Assumption 2.1 and Eq. (4.1), respectively.

In this section, we show that the estimation error condition in Assumption 5.2 (and thus $O(\sqrt{T})$ regret) can be achieved under similar assumptions as in Abbasi-Yadkori et al. (2019) and Wei et al. (2020a), which we state next.

We now describe a simple procedure for estimating the average action-value functions $Q_{k}(x,a)=\phi(x,a)^{\top}w_{k}$, where $w_{k}=\frac{1}{k}\sum_{i=1}^{k}w_{\pi_{i}}$, such that the conditions of Assumption 5.2 are satisfied. Essentially, we estimate each $Q_{\pi_{i}}$ using least-squares Monte Carlo and then average the weights. We will use the shorthand notation $\phi_{t}=\phi(x_{t},a_{t})$. Let $\mathcal{H}_{i}$ and ${\mathcal{T}}_{i}$ be subsets of time indices in phase $i$ (defined later). We estimate $Q_{\pi_{i}}$ as follows:

where $R_{t}$ are the empirical $b$-step returns ($b$ is specified later), computed as

We then estimate $w_{k}$ as $\widehat{w}_{k}=\frac{1}{k}\sum_{i=1}^{k}\widehat{w}_{\pi_{i}}$. Note that for this special case of linear value functions, we do not need to use the replay buffer in Algorithm 1. For analysis purposes, we divide each phase of length $\tau$ into $2m$ blocks of size $b$ and let $\mathcal{H}_{i}$ (${\mathcal{T}}_{i}$) be the starting indices of odd (even) blocks in phase $i$. Due to the gaps between indices and fast mixing, this makes the data almost independent (we make this precise in Appendix C) and the error easier to analyze. In practice, one may simply want to use all data.

For a distribution $\mu$, let $\|x\|_{\mu}$ denote the distribution-weighted norm such that $\|x\|_{\mu}^{2}=\sum_{i}\mu_{i}x_{i}^{2}$. Using Jensen’s inequality, we have that $(\mathbb{E}_{x\sim\mu}[q(x)])^{2}\leq\mathbb{E}_{x\sim\mu}[q(x)^{2}]$. Thus, it suffices to bound the $Q$-function error in the distribution-weighted norm, $\|Q_{K}-\widehat{Q}_{K}\|_{\mu_{*}\otimes\pi_{*}}$. Furthermore, given bounded features,

so it suffices to bound the error in the weights. We bound this error in the following Lemma, proven in Appendix C.

Furthermore, in Appendix D, we show that the error in the average state-value function satisfies the following:

We have demonstrated that Assumption 5.2 can be satisfied with linear value functions. For Assumption 5.1, it suffices for the weight estimates $\{\widehat{w}_{\pi_{i}}\}$ to be bounded. This will be true, since we assume that the true weights are bounded and we can bound the error in the weight space. Thus Politex has an $O(\sqrt{T})$ regret in this setting (though note that we incur an extra $|\mathcal{A}|$ factor coming from the $\widehat{V}_{K}$ error bound).

As mentioned earlier, the key idea in our policy update is to obtain an estimate $\widehat{Q}_{k}$ of the average of all the Q-functions in previous phases. We have seen that when we use linear functions to approximate the Q-functions, we can simply average the weights in order to get an estimate of the average Q-function. However, in practice, we often need to use non-linear functions, especially neural networks, to approximate complex Q-functions. In this section, we discuss how to efficiently implement our algorithm with non-linear function approximation.

In this section, we evaluate our implementations empirically. We make comparisons with several baselines in two control environments.

Environments. We use two control environments with the simulators described in Tassa et al. (2018). Both environments are episodic with episode length $1000$. The environments we evaluate are:

• •

  Cart-pole (Barto et al., 1983): The goal of this environment is to balance an unactuated pole by applying forces to a cart at its base. We discretize the continuous force to $5$ values: $\{-2,-1,0,1,2\}$. The reward at each time step is a real number in $[0,1]$. We end episodes early when the pole falls (we use rewards less than $0.5$ an indicator for falling), and assume zero reward for the remaining steps when reporting results.

• •

  Ball-in-cup: Here, an actuated planar receptacle can translate in the vertical plane in order to swing and catch a ball attached to its bottom. This task has a sparse reward: $1$ when the ball is in the cup, and $0$ otherwise. We discretize the two-dimensional continuous action space to a $3\times 3$ grid, i.e., $9$ actions in total.

Algorithms. We compare the following algorithms: our proposed implementation of Politex using experience replay, Politex using weight averaging, the original Politex algorithm, the variant of Politex that averages $10$ randomly selected Q-functions, the mirror descent policy optimization (MDPO) algorithm (Tomar et al., 2020), the constrained policy optimization (CPO) (Achiam et al., 2017), and the V-MPO algorithm (Song et al., 2019b).

For all the algorithms, we extract Fourier basis features (Konidaris et al., 2011) from the raw observations for both environments: for Cart-pole, we use $4$ bases and for Ball-in-cup we use $2$ bases. For variants of Politex we construct the state-action features $\phi(x,a)$ by block one-hot encoding, i.e., we partition the $\phi(x,a)$ vector into $|\mathcal{A}|$ blocks, and set the $a$-th block to be the Fourier features of $x$ and other blocks to be zero. We approximate $Q$-functions using neural networks with one hidden layer and ReLU activation: the width of the hidden layer is $50$ for Cart-pole and $250$ for Ball-in-cup. For MDPO, CPO and V-MPO, the algorithms use a policy network whose input and output are the state feature and action probability, respectively. We use one hidden layer networks with width $50$ for Cart-pole and $250$ for Ball-in-cup. These three algorithms also need a value network, which takes the state feature as input and outputs its value for the current policy. For both environments, we use one hidden layer network with width $50$.

For Cart-pole, we choose phase length $\tau=10^{4}$ and for Ball-in-cup, we choose $\tau=2\times 10^{4}$. Since the environments are episodic, we have multiple episodes in each phase. For Cart-pole, since the training usually does not need a large number of phases, we do not set limit on the size of the replay buffer, whereas for Ball-in-cup, we set the limit as $2\times 10^{6}$. For our algorithm and the variants of Politex, we choose parameter $\eta$ in $\{5,10,20,40,80,160\}$ and report the result of each algorithm using the value of $\eta$ that produces the best average reward. For MDPO, CPO and V-MPO, we note that according to Eq. (4.1), the KL regularization coefficient between policies is the reciprocal of the $\eta$ parameter in our algorithm, we also run these baseline algorithms in the same range of $\eta$ and report the best result. In addition, in CPO, we set the limit in KL divergence between adjacent policies $0.001\eta$, which, in our experiments, leads to good performance. We treat the length of Monte Carlo estimate $b$ as a hyper parameter, so we choose it in $\{100,300\}$ and report the best performance.

Results. We run each algorithm in each environment $20$ times and report the average results as well as the standard deviation of the $20$ runs (shaded area). We report the average reward in each phase before training on the data collected during the phase. The results are presented in Figure 1. Every run is conducted on a single P100 GPU.

From Fig. 1(a), we can see that the iteration complexity (average reward vs number of phases) and final performance of the experience replay implementation using all the data from past phases (blue curve) are similar to many state-of-the-art algorithms, such as Politex and V-MPO. Meanwhile, according to Fig. 1(b), the experience replay implementation achieves strong performance in training speed, i.e., best performance at $30$ minute in Cart-pole and second best performance at $200$ minute in Ball-in-cup. We note here that the training time includes both the time for data collection, i.e., interacting with the environment simulator and the training of a new policy. The observation that the experience replay based implementation is faster than the original Politex and the implementation that uses $10$ Q-functions can be explained by the fact that the replay-based algorithm only uses a single Q-function and thus is faster at data collection, as discussed in Section 7.2. In addition, that the replay-based algorithm runs faster than MDPO, CPO, and V-MPO can be explained by its simplicity: in variants of Politex, we only need to approximate the Q-function, whereas in MDPO, CPO, and V-MPO, we need to train both the policy and value networks.

The Politex weight averaging scheme achieves the best performance on Ball-in-cup in both iteration complexity and training speed, but does not converge to a solution that is sufficiently close to the optimum on Cart-pole. Notice that for weight averaging, we used the initialization technique mentioned in Section 7.1. Considering the consistent strong performance of experience replay in both environments, we still recommend applying experience replay when using non-linear function approximation.

In Fig. 1(c), we can see that when we only sample a small subset ($1\%$ or $4\%$) of the data to store in the replay buffer, the coreset technique described in Section 7.2 achieves better performance due to the variance reduction effect. This improvement is quite significant in the Ball-in-cup environment, which has sparse rewards. Notice that sampling coreset only adds negligible computational overhead during training, and thus we recommend the coreset technique when there exists a strict memory constraint.

The main contributions of our work are an improved analysis and practical implementation of Politex. On the theoretical side, we show that Politex obtains an $O(\sqrt{T})$ high-probability regret bound in uniformly mixing average-reward MDPs with linear function approximation, which is the first such bound for a computationally efficient algorithm. The main limitation of these result is that, similarly to previous works, they hold under somewhat strong assumptions which circumvent the need for exploration. An interesting future work direction would be to relax these assumptions and incorporate explicit exploration instead.

On the practical side, we propose an efficient implementation with neural networks that relies on experience replay, a standard tool in modern deep RL. Our work shows that experience replay and KL regularization can both be viewed as approximately implementing mirror descent policy updates within a policy iteration scheme. This provides an online-learning justification for using replay buffers in policy iteration, which is different than the standard explanation for their success. Our work also suggests a new objective for storing and prioritizing replay samples, with the goal of approximating the average of value functions well. This goal has some similarities with continual learning, where experience replay has also lead to empirical successes. One interesting direction for future work would be to explore other continual learning techniques in the approximate implementation of mirror descent policy updates.

We would like to thank Mehrdad Farajtabar, Dilan Görür, Nir Levine, and Ang Li for helpful discussions.