Regularization Guarantees Generalization in Bayesian Reinforcement Learning through Algorithmic Stability

How can an agent learn to quickly perform well in an unknown task? This is the basic question in reinforcement learning (RL). The most popular RL algorithms are designed in a minimax approach – seeking a procedure that will eventually learn to perform well in any task (Strehl et al. 2006; Jaksch, Ortner, and Auer 2010; Jin et al. 2018). Lacking prior information about the task, such methods must invest considerable efforts in uninformed exploration, typically requiring many samples to reach adequate performance. In contrast, when a prior distribution over possible tasks is known in advance, an agent can direct its exploration much more effectively. This is the Bayesian RL (BRL, Ghavamzadeh et al. 2016) setting. A Bayes-optimal policy – the optimal policy in BRL – can be orders of magnitude more sample efficient than a minimax approach, and indeed, recent studies demonstrated a quick solution of novel tasks, sometimes in just a handful of trials (Duan et al. 2016; Zintgraf et al. 2020; Dorfman, Shenfeld, and Tamar 2020).

For high-dimensional problems, and when the task distribution does not obey a very simple structure, solving BRL is intractable, and one must resort to approximations. A common approximation, which has been popularized under the term meta-RL (Duan et al. 2016; Finn, Abbeel, and Levine 2017), is to replace the distribution over tasks with an empirical sample of tasks, and seek an optimal policy with respect to the sample, henceforth termed the empirical risk minimization policy (ERM policy). In this paper, we investigate the performance of the ERM policy on a novel task from the task distribution, that is – we ask how well the policy generalizes.

We focus on the Probably Approximately Correct (PAC) framework, which is popular in supervised learning, and adapt it to the BRL setting. Since the space of deterministic history-dependent policies is finite (in a finite horizon setting), a trivial generalization bound for a finite hypothesis space can be formulated. However, the size of the policy space, which such a naive bound depends on, leads us to seek alternative methods for controlling generalization. In particular, regularization is a well-established method in supervised learning that can be used to trade-off training error and test error. In RL, regularized MDPs are popular in practice (Schulman et al. 2017), and have also received interest lately due to their favorable optimization properties (Shani, Efroni, and Mannor 2020; Neu, Jonsson, and Gómez 2017).

The main contribution of this work is making the connection between regularized MDPs and PAC generalization, as described above. We build on the classical analysis of Bousquet and Elisseeff (2002), which bounds generalization through algorithmic stability. Establishing algorithmic stability results for regularized MDPs, however, is not trivial, as the loss function in MDPs is not convex in the policy. Our key insight is that while not convex, regularized MDPs satisfy a certain quadratic growth criterion, which is sufficient to establish stability. To show this, we build on the recently discovered fast convergence rates for mirror descent in regularized MDPs (Shani, Efroni, and Mannor 2020). Our result, which may be of independent interest, allows us to derive generalization bounds that can be controlled by the regularization magnitude. Furthermore, we show that when the MDP prior obeys certain structural properties, our results significantly improve the trivial finite hypothesis space bound.

To our knowledge, this is the first work to formally study generalization in the BRL setting. While not explicitly mentioned as such, the BRL setting has been widely used by many empirical studies on generalization in RL (Tamar et al. 2016; Cobbe et al. 2020). In fact, whenever the MDPs come from a distribution, BRL is the relevant formulation. Our results therefore also establish a formal basis for studying generalization in RL.

This paper is structured as follows. After surveying related work in Section 2, we begin with background on MDPs, BRL, and algorithmic stability in Section 3, and then present our problem formulation and straightforward upper and lower bounds for the ERM policy in Section 4. In Section 5 we discuss fundamental properties of regularized MDPs. In Section 6 we describe a general connection between a certain rate result for mirror descent and a quadratic growth condition, and in Section 7 we apply this connection to regularized MDPs, and derive corresponding generalization bounds. We discuss our results and future directions in Section 8.

Generalization to novel tasks in RL has been studied extensively, often without making the explicit connection to Bayesian RL. Empirical studies can largely be classified into three paradigms. The first increases the number of training tasks, either using procedurally generated domains (Cobbe et al. 2020), or by means such as image augmentation (Kostrikov, Yarats, and Fergus 2020) and task interpolation (Yao, Zhang, and Finn 2021). The second paradigm adds inductive bias to the neural network, such as a differentiable planning or learning computation (Tamar et al. 2016; Boutilier et al. 2020), or graph neural networks (Rivlin, Hazan, and Karpas 2020). The third is meta-RL, where an agent is explicitely trained to generalize, either using a Bayesian RL objective (Duan et al. 2016; Zintgraf et al. 2020), or through gradient based meta learning (Finn, Abbeel, and Levine 2017). We are not aware of theoretical studies of generalization in Bayesian RL.

The Bayesian RL algorithms of Guez, Silver, and Dayan (2012) and Grover, Basu, and Dimitrakakis (2020) perform online planning in the belief space by sampling MDPs from the posterior at each time step, and optimizing over this sample. The performance bounds for these algorithms require the correct posterior at each step, implicitly assuming a correct prior, while we focus on learning the prior from data. The lower bound in our Proposition 2 demonstrates how errors in the prior can severely impact performance.

Our stability-based approach to PAC learning is based on the seminal work of Bousquet and Elisseeff (2002). More recent works investigated stability of generalized learning (Shalev-Shwartz et al. 2010), multi-task learning (Zhang 2015), and stochastic gradient descent (Hardt, Recht, and Singer 2016). To our knowledge, we provide the first stability result for regularized MDPs, which, due to their non-convex nature, requires a new methodology. The stability results of Charles and Papailiopoulos (2018) build on quadratic growth, a property we use as well. However, showing that this property holds for regularized MDPs is not trivial, and is a major part of our contribution.

Finally, there is recent interest in PAC-Bayes theory for meta learning (Amit and Meir 2018; Rothfuss et al. 2021; Farid and Majumdar 2021). To our knowledge, this theory has not yet been developed for meta RL.

We give background on BRL and algorithmic stability.

We next describe our learning problem. We are given a training set of $N$ simulators for $N$ independently sampled MDPs, $\left\{M_{1},\dots,M_{N}\right\}$, where each $M_{i}\sim P(M)$; in the following, we will sometimes refer to this training set as the training data. We are allowed to interact with these simulators as we wish for an unrestricted amount of time. From this interaction, our goal is to compute a policy $\pi$ that obtains a low expected $T$-horizon cost for a test simulator $M\sim P(M)$, i.e., we wish to minimize the population risk (1). It is well known (e.g., Ghavamzadeh et al. 2016) that there exists a deterministic history dependent policy that minimizes (1), also known as the Bayes-optimal policy, and we denote it by $\pi_{\mathrm{BO}}$. Our performance measure is the $T$-horizon average regret,

$$\begin{split}\mathcal{R}_{T}(\pi)=&\mathbb{E}_{M\sim P}\Bigg{[}\mathbb{E}_{\pi;M}\bigg{[}\sum_{t=0}^{T}C_{M}(s_{t},a_{t})\bigg{]}\\ &-\!\mathbb{E}_{\pi_{\mathrm{BO}};M}\!\bigg{[}\!\sum_{t=0}^{T}\!C_{M}(s_{t},a_{t})\!\bigg{]}\!\Bigg{]}=\mathcal{L}(\pi)-\mathcal{L}(\pi_{\mathrm{BO}}).\end{split}$$

(2)

In supervised learning, a well-established method for controlling generalization is to add a regularization term, such as the $L_{2}$ norm of the parameters, to the objective function that is minimized. The works of Bousquet and Elisseeff (2002); Shalev-Shwartz et al. (2010) showed that for convex loss functions, adding a strongly convex regularizer such as the $L_{2}$ norm leads to algorithmic stability, which can be used to derive generalization bounds that are controlled by the amount of regularization. In this work, we ask whether a similar approach of adding regularization to the BRL objective (3) can be used to control generalization.

We focus on the following regularization scheme. For some $\lambda>0$, consider a regularized ERM of the form:

$$\hat{\mathcal{L}}^{\lambda}(\pi)=\frac{1}{N}\sum_{i=1}^{N}\mathbb{E}_{\pi;M_{i}}\left[\sum_{t=0}^{T}C_{M_{i}}(s_{t},a_{t})+\lambda\mathcal{R}(\pi(\cdot|h_{t}))\right],$$

where $\mathcal{R}$ is some regularization function applied to the policy. In particular, we will be interested in $L_{2}$ regularization, where $\mathcal{R}(\pi(\cdot|h_{t}))=\|\pi(\cdot|h_{t})\|_{2}$. We also define the regularized population risk,

$$\mathcal{L}^{\lambda}(\pi)=\mathbb{E}_{M\sim P}\mathbb{E}_{\pi;M}\left[\sum_{t=0}^{T}C_{M}(s_{t},a_{t})+\lambda\mathcal{R}(\pi(\cdot|h_{t}))\right].$$

In standard (non-Bayesian) RL, regularized MDPs have been studied extensively (Neu, Jonsson, and Gómez 2017). A popular motivation has been to use the regularization to induce exploration (Fox, Pakman, and Tishby 2015; Schulman et al. 2017). Recently, Shani, Efroni, and Mannor (2020) showed that for optimizing a policy using $k$ iterations of mirror descent (equivalent to trust region policy optimization Schulman et al. 2015) with $L_{2}$ or entropy regularization enables a fast $O(1/k)$ convergence rate, similarly to convergence rates for strongly convex functions, although the MDP objective is not convex. In our work, we build on these results to show a stability property for regularization in the BRL setting described above. We begin by adapting a central result in Shani, Efroni, and Mannor (2020), which was proved for discounted MDPs, to our finite horizon and history dependent policy setting.

The BRL objectives in Eq. (1) (similarly, Eq. (3)) can be interpreted as follows: we first choose a history dependent policy $\pi(h_{t})$, and then nature draws an MDP $M\sim P(M)$ (similarly, $M\sim\hat{P}_{N}$), and we then evaluate $\pi(h_{t})$ on $M$. The expected cost (over the draws of $M$), is the BRL performance. In the following discussion, for simplicity, the results are given for the prior $P(M)$, but they hold for $\hat{P}_{N}$ as well.

Let $P(M|h_{t};\pi)$ denote the posterior probability of nature having drawn the MDP $M$, given that we have seen the history $h_{t}$ under policy $\pi$. From Bayes rule, we have that

$$P(M|h_{t};\pi)\propto P(h_{t}|M;\pi)P(M).$$

Let us define the regularized expected cost,

$$C_{\lambda}(h_{t},a_{t};\pi)=\mathbb{E}_{M|h_{t};\pi}C_{M}(s_{t},a_{t})+\lambda\mathcal{R}(\pi_{t}(\cdot|h_{t})),$$

and the value function,

$$V_{t}^{\pi}(h_{t})=\mathbb{E}_{\pi;M|h_{t}}\left[\left.\sum_{t^{\prime}=t}^{T}C_{\lambda}(h_{t^{\prime}},a_{t^{\prime}};\pi)\right|h_{t}\right].$$

The value function satisfies Bellman’s equation. Let

$$\begin{split}&P(c_{t},s_{t+1}|h_{t},a_{t})=\\ &\sum_{M}P(M|h_{t})P(c_{t}|M,s_{t},a_{t})P(s_{t+1}|M,s_{t},a_{t})\end{split}$$

denote the posterior probability of observing $c_{t},s_{t+1}$ at time $t$. Then

$$V_{T}^{\pi}(h_{T})=\sum_{a_{T}}\pi(a_{T}|h_{T})C_{\lambda}(h_{T},a_{T};\pi),$$

and, letting $h_{t+1}=\left\{h_{t},a_{t},c_{t},s_{t+1}\right\}$,

$$\begin{split}V_{t}^{\pi}(h_{t})=&\sum_{a_{t}}\pi(a_{t}|h_{t})\bigg{(}C_{\lambda}(h_{t},a_{t};\pi)\\ &+\sum_{c_{t},s_{t+1}}P(c_{t},s_{t+1}|h_{t},a_{t})V_{t+1}^{\pi}(\left\{h_{t},a_{t},c_{t},s_{t+1}\right\})\bigg{)}.\end{split}$$

Consider two histories $h_{t},\bar{h}_{\bar{t}}\in\mathcal{H}$, and let

$$\mathbf{P}^{\pi}(\bar{h}_{\bar{t}}|h_{t})\!=\!\left\{\begin{array}[]{lr}\!\!\!\sum_{a_{t}}\!\!\pi(a_{t}|h_{t})P(\bar{c}_{t},\bar{s}_{t+1}|h_{t},a_{t}),&\text{if }{\bar{t}}=t+1\\ \!\!\!0,&\text{else}\end{array}\right.$$

denote the transition probability between histories. Also, define

$$C^{\pi}(h_{t})=\sum_{a_{t}}\pi(a_{t}|h_{t})C_{\lambda}(h_{t},a_{t};\pi).$$

We can write the Bellman equation in matrix form as follows

$$V^{\pi}=C^{\pi}+\mathbf{P}^{\pi}V^{\pi},$$

(4)

where $V^{\pi}$ and $C^{\pi}$ are vectors in $\mathbb{R}^{|\mathcal{H}|}$, and $\mathbf{P}^{\pi}$ is a matrix in $\mathbb{R}^{|\mathcal{H}|\times|\mathcal{H}|}$.

The uniform trust region policy optimization algorithm of Shani, Efroni, and Mannor (2020) is a type of mirror descent algorithm applied to the policy in regularized MDPs. An adaptation of this algorithm for our setting is given in Sec. B.4 of the supplementary material. The next result provides a fundamental inequality that the policy updates of this algorithm satisfy, in the spirit of an inequality that is used to establish convergence rates for mirror descent (cf. Lemma 8.11 in Beck 2017). The proof follows Lemma 10 in Shani, Efroni, and Mannor (2020), with technical differences due to the finite horizon setting; it is given in Sec. B.4.

In the following, we shall show that Proposition 3 can be used to derive stability bounds in the regularized BRL setting. To simplify our presentation, we first present a key technique that our approach builds on in a general optimization setting, and only then come back to MDPs.

Standard stability results, such as in Bousquet and Elisseeff (2002); Shalev-Shwartz et al. (2010), depend on convexity of the loss function, and strong convexity of the regularizing function (Bousquet and Elisseeff 2002). While our $L_{2}$ regularization is strongly convex, the MDP objective is not convex in the policy.²²2The linear programming formulation is not suitable for establishing stability in our BRL setting, as changing the prior would change the constraints in the linear program. In this work, we show that nevertheless, algorithmic stability can be established. To simplify our presentation, we first present the essence of our technique in a general form, without the complexity of MDPs. In the next section, we adapt the technique to the BRL setting.

Our key insight is that the fundamental inequality of mirror descent (cf. Prop. 3), actually prescribes a quadratic growth condition. The next lemma shows this for a general iterative algorithm, but it may be useful to think about mirror descent when reading it. In the sequel, we will show that similar conditions hold for regularized MDPs.

We now present a stability result for a regularized ERM objective. The proof resembles Shalev-Shwartz et al. (2010), but replaces strong convexity with quadratic growth.

We are now ready to present stability bounds for the regularized Baysian RL setting. Let $\mu\in\mathbb{R}^{|\mathcal{H}|}$ denote the distribution over $h_{0}$, the initial history (we assume that all elements in the vector that correspond to histories of length greater than $0$ are zero). Recall the regularized ERM loss $\hat{\mathcal{L}}^{\lambda}(\pi)$, and let $\pi^{*}$ denote its minimizer. Define the leave-one-out ERM loss,

$$\hat{\mathcal{L}}^{\lambda,\backslash j}(\pi)=\frac{1}{N}\sum_{\begin{subarray}{c}i=1\\ i\neq j\end{subarray}}^{N}\mathbb{E}_{\pi;M_{i}}\left[\sum_{t=0}^{T}C(s_{t},a_{t})+\lambda\mathcal{R}(\pi(\cdot|h_{t}))\right],$$

and let $\pi^{\backslash j,*}$ its minimizer. Recall the definition of $\mathbf{P}^{\pi}$ – the transition probability between histories under policy $\pi$, which depends on the prior. In the following, we use the following notation: $\mathbf{P}^{\pi}$ refers to the empirical prior $\hat{P}_{N}$, while $\mathbf{P}_{M_{j}}^{\pi}$ refers to a prior that has all its mass on a single MDP $M_{j}$. The following theorem will be used to derive our stability results. The proof is in Sec. C.

Following the proof of Proposition 4, we would like to use the two expressions in Theorem 3 to bound $\left\|\pi^{\backslash j,*}-\pi^{*}\right\|_{2}$, which would directly lead to a stability result. This is complicated by the fact that different factors $(I-\mathbf{P}^{\pi^{\backslash j,*}})^{-1}$ and $(I-\mathbf{P}_{M_{j}}^{\pi^{\backslash j,*}})^{-1}$ appear in the two expressions. Our first result assumes that these expressions cannot be too different; a discussion of this assumption follows.

Let us define the regularized loss for MDP $M$, $\mathcal{L}_{M}^{\lambda}(\pi)=\mathbb{E}_{\pi;M}\left[\sum_{t=0}^{T}C_{M}(s_{t},a_{t})+\lambda\mathcal{R}(\pi(\cdot|h_{t}))\right].$ We have the following result.

Note that each element that corresponds to history $h_{t}$ in the vector $\mu^{\top}(I-\mathbf{P}^{\pi}_{M})^{-1}$ is equivalent to $P(h_{t}|M;\pi)$, the probability of observing $h_{t}$ under policy $\pi$ and MDP $M$ (see Sec. B.1 for formal proof). Thus, Assumption 1 essentially states that two different MDPs under the prior cannot visit completely different histories given the same policy. With our regularization scheme, such an assumption is required for uniform stability: if the test MDP can reach completely different states than possible during training, it is impossible to guarantee anything about the performance of the policy in those states. Unfortunately, the constant $D$ can be very large. Let

$$q=\sup_{M,M^{\prime}\in\mathcal{M},s,s^{\prime}\in S,a\in A,c\in\mathcal{C}}\frac{P_{M}(s^{\prime},c|s,a)}{P_{M^{\prime}}(s^{\prime},c|s,a)},$$

where we assume that $0/0=1$. Then, $P(h_{t}|M;\pi)/P(h_{t}|M^{\prime};\pi)=\Pi_{t}\frac{P_{M}(s_{t+1},c_{t}|s_{t},a_{t})}{P_{M^{\prime}}(s_{t+1},c_{t}|s_{t},a_{t})}$ is at most $q^{T}$, and therefore $D$ can be in the order of $q^{T}$. One way to guarantee that $D$ is finite, is to add a small amount of noise to any state transition. The following example estimates $q$ is such a case.

Let us now compare Corollary 1 with the trivial bound in Proposition 1. First, Corollary 1 allows to control generalization by increasing the regularization $\lambda$. The term $2\lambda T$ is a bias, incurred by adding the regularization to the objective, and can be reduced by decreasing $\lambda$. Comparing the constants of the $\mathcal{O}(1/\sqrt{N})$ term, the dominant terms are $D^{2}$ vs. $(|S||\mathcal{C}||A|)^{T}$. Since $D$ does not depend on $|A|$, the bound in Corollary 1 is important for problems with large $|A|$. The example above shows that in the worst case, $D^{2}$ can be $\mathcal{O}((|S||\mathcal{C}|)^{2T})$. There are, of course, more favorable case, where the structure of $\mathcal{M}$ is such that $D$ is better behaved.

Another case is where the set $\mathcal{M}$ is finite. In this case, we can show that the pointwise hypothesis stability does not depend on $D$, and we obtain a bound that does not depend exponentially on $T$, as we now show.

In the generalization bounds of both Corollary 1 and Corollary 2, reducing $\lambda$ and increasing $N$ at a rate such that the stability is $o(1/{\sqrt{N}})$ will guarantee learnability, i.e., convergence to $\pi_{\mathrm{BO}}$ as $N\to\infty$.

For a finite $N$, and when there is structure the hypothesis space $\mathcal{M}$, as displayed for example in $D$, the bounds in Corollaries 1 and 2 allow to set $\lambda$ to obtain bounds that are more optimistic than the trivial bound in Proposition 1. In these cases, our results show that regularization allows for improved generalization.

In this work, we analyzed generalization in Bayesian RL, focusing on algorithmic stability and a specific form of policy regularization. The bounds we derived can be controlled by the amount of regularization, and under some structural assumptions on the space of possible MDPs, compare favorably to a trivial bound based on the finite policy space. We next outline several future directions.

Aviv Tamar thanks Kfir Levy for insightful discussions on optimization and stability, and Shie Mannor for advising on the presentation of this work. This research was partly funded by the Israel Science Foundation (ISF-759/19).