Provably Efficient Exploration in Policy Optimization

Introduction

Coupled with powerful function approximators such as neural networks, policy optimization plays a key role in the tremendous empirical successes of deep reinforcement learning (Silver et al., 2016, 2017; Duan et al., 2016; OpenAI, 2019; Wang et al., 2018). In sharp contrast, the theoretical understandings of policy optimization remain rather limited from both computational and statistical perspectives. More specifically, from the computational perspective, it remains unclear until recently whether policy optimization converges to the globally optimal policy in a finite number of iterations, even given infinite data. Meanwhile, from the statistical perspective, it still remains unclear how to attain the globally optimal policy with a finite regret or sample complexity.

A line of recent work (Fazel et al., 2018; Yang et al., 2019a; Abbasi-Yadkori et al., 2019a, b; Bhandari and Russo, 2019; Liu et al., 2019; Agarwal et al., 2019; Wang et al., 2019) answers the computational question affirmatively by proving that a wide variety of policy optimization algorithms, such as policy gradient (PG) (Williams, 1992; Baxter and Bartlett, 2000; Sutton et al., 2000), natural policy gradient (NPG) (Kakade, 2002), trust-region policy optimization (TRPO) (Schulman et al., 2015), proximal policy optimization (PPO) (Schulman et al., 2017), and actor-critic (AC) (Konda and Tsitsiklis, 2000), converge to the globally optimal policy at sublinear rates of convergence, even when they are coupled with neural networks (Liu et al., 2019; Wang et al., 2019). However, such computational efficiency guarantees rely on the regularity condition that the state space is already well explored. Such a condition is often implied by assuming either the access to a “simulator” (also known as the generative model) (Koenig and Simmons, 1993; Azar et al., 2011, 2012a, 2012b; Sidford et al., 2018a, b; Wainwright, 2019) or finite concentratability coefficients (Munos and Szepesvári, 2008; Antos et al., 2008; Farahmand et al., 2010; Tosatto et al., 2017; Yang et al., 2019b; Chen and Jiang, 2019), both of which are often unavailable in practice.

In a more practical setting, the agent sequentially explores the state space, and meanwhile, exploits the information at hand by taking the actions that lead to higher expected total rewards. Such an exploration-exploitation tradeoff is better captured by the aforementioned statistical question regarding the regret or sample complexity, which remains even more challenging to answer than the computational question. As a result, such a lack of statistical understanding hinders the development of more sample-efficient policy optimization algorithms beyond heuristics. In fact, empirically, vanilla policy gradient is known to exhibit a possibly worse sample complexity than random search (Mania et al., 2018), even in basic settings such as linear-quadratic regulators. Meanwhile, theoretically, vanilla policy gradient can be shown to suffer from exponentially large variance in the well-known “combination lock” setting (Kakade, 2003; Leffler et al., 2007; Azar et al., 2012a), which only has a finite state space.

In this paper, we aim to answer the following fundamental question:

Can we design a policy optimization algorithm that incorporates exploration and is provably sample-efficient?

To answer this question, we propose the first policy optimization algorithm that incorporates exploration in a principled manner. In detail, we develop an Optimistic variant of the PPO algorithm, namely OPPO. Our algorithm is also closely related to NPG and TRPO. At each update, OPPO solves a Kullback-Leibler (KL)-regularized policy optimization subproblem, where the linear component of the objective function is defined using the action-value function. As is shown subsequently, solving such a subproblem corresponds to one iteration of infinite-dimensional mirror descent (Nemirovsky and Yudin, 1983) or dual averaging (Xiao, 2010), where the action-value function plays the role of the gradient. To encourage exploration, we explicitly incorporate a bonus function into the action-value function, which quantifies the uncertainty that arises from only observing finite historical data. Through uncertainty quantification, such a bonus function ensures the (conservative) optimism of the updated policy. Based on NPG, TRPO, and PPO, OPPO only augments the action-value function with the bonus function in an additive manner, which makes it easily implementable in practice.

Theoretically, we establish the sample efficiency of OPPO in an episodic setting of Markov decision processes (MDPs) with full-information feedback, where the transition dynamics are linear in features (Yang and Wang, 2019b, a; Jin et al., 2019; Ayoub et al., 2020; Zhou et al., 2020). In particular, we allow the transition dynamics to be nonstationary within each episode. See also the work of Du et al. (2019a); Van Roy and Dong (2019); Lattimore and Szepesvari (2019) for a related discussion on the necessity of the linear representation. In detail, we prove that OPPO attains a $\sqrt{d^{2}H^{3}T}$-regret up to logarithmic factors, where $d$ is the feature dimension, $H$ is the episode horizon, and $T$ is the total number of steps taken by the agent. Note that such a regret does not depend on the numbers of states and actions, and therefore, allows them to be even infinite. In particular, OPPO attains such a regret without knowing the transition dynamics or accessing a “simulator”. Moreover, we prove that, even when the reward functions are adversarially chosen across the episodes, OPPO attains the same regret in terms of competing with the globally optimal policy in hindsight (Cesa-Bianchi and Lugosi, 2006; Bubeck and Cesa-Bianchi, 2012). In comparison, existing algorithms based on value iteration, e.g., optimistic least-squares value iteration (LSVI) (Jin et al., 2019), do not allow adversarially chosen reward functions. Such a notion of robustness partially justifies the empirical advantages of KL-regularized policy optimization (Neu et al., 2017; Geist et al., 2019). To the best of our knowledge, OPPO is the first provably sample-efficient policy optimization algorithm that incorporates exploration.

Our work is based on the aforementioned line of recent work (Fazel et al., 2018; Yang et al., 2019a; Abbasi-Yadkori et al., 2019a, b; Bhandari and Russo, 2019; Liu et al., 2019; Agarwal et al., 2019; Wang et al., 2019) on the computational efficiency of policy optimization, which covers PG, NPG, TRPO, PPO, and AC. In particular, OPPO is based on PPO (and similarly, NPG and TRPO), which is shown to converge to the globally optimal policy at sublinear rates in tabular and linear settings, as well as nonlinear settings involving neural networks (Liu et al., 2019; Wang et al., 2019). However, without assuming the access to a “simulator” or finite concentratability coefficients, both of which imply that the state space is already well explored, it remains unclear whether any of such algorithms is sample-efficient, that is, attains a finite regret or sample complexity. In comparison, by incorporating uncertainty quantification into the action-value function at each update, which explicitly encourages exploration, OPPO not only attains the same computational efficiency as NPG, TRPO, and PPO, but is also shown to be sample-efficient with a $\sqrt{d^{2}H^{3}T}$-regret up to logarithmic factors.

Our work is closely related to another line of work (Even-Dar et al., 2009; Yu et al., 2009; Neu et al., 2010a, b; Zimin and Neu, 2013; Neu et al., 2012; Rosenberg and Mansour, 2019a, b) on online MDPs with adversarially chosen reward functions, which mostly focuses on the tabular setting.

Assuming the transition dynamics are known and the full information of the reward functions is available, the work of Even-Dar et al. (2009) establishes a $\sqrt{\tau^{2}T\cdot\log|\mathcal{A}|}$-regret, where $\mathcal{A}$ is the action space, $|\mathcal{A}|$ is its cardinality, and $\tau$ upper bounds the mixing time of the MDP. See also the work of Yu et al. (2009), which establishes a $T^{2/3}$-regret in a similar setting.

Assuming the transition dynamics are known but only the bandit feedback of the received rewards is available, the work of Neu et al. (2010a, b); Zimin and Neu (2013) establishes an $H^{2}\sqrt{|\mathcal{A}|T}/\beta$-regret (Neu et al., 2010b), a $T^{2/3}$-regret (Neu et al., 2010a), and a $\sqrt{H|{\mathcal{S}}||\mathcal{A}|T}$-regret (Zimin and Neu, 2013), respectively, all up to logarithmic factors. Here ${\mathcal{S}}$ is the state space and $|{\mathcal{S}}|$ is its cardinality. In particular, it is assumed by Neu et al. (2010b) that, with probability at least $\beta$, any state is reachable under any policy.

Assuming the full information of the reward functions is available but the transition dynamics are unknown, the work of Neu et al. (2012); Rosenberg and Mansour (2019a) establishes an $H|{\mathcal{S}}||\mathcal{A}|\sqrt{T}$-regret (Neu et al., 2012) and an $H|{\mathcal{S}}|\sqrt{|\mathcal{A}|T}$-regret (Rosenberg and Mansour, 2019a), respectively, both up to logarithmic factors.

Assuming the transition dynamics are unknown and only the bandit feedback of the received rewards is available, the recent work of Rosenberg and Mansour (2019b) establishes an $H|{\mathcal{S}}|\sqrt{|\mathcal{A}|T}/\beta$-regret up to logarithmic factors. In particular, it is assumed by Rosenberg and Mansour (2019b) that, with probability at least $\beta$, any state is reachable under any policy. Without such an assumption, an $H^{3/2}|{\mathcal{S}}||\mathcal{A}|^{1/4}T^{3/4}$-regret is established.

In the latter two settings with unknown transition dynamics, all the existing algorithms (Neu et al., 2012; Rosenberg and Mansour, 2019a, b) follow the gradient direction with respect to the visitation measure, and thus, differ from most practical policy optimization algorithms. In comparison, OPPO is not restricted to the tabular setting and indeed follows the gradient direction with respect to the policy. OPPO is simply an optimistic variant of NPG, TRPO, and PPO, which makes it also a practical policy optimization algorithm. In particular, when specialized to the tabular setting, our setting corresponds to the third setting with $d=|{\mathcal{S}}|^{2}|\mathcal{A}|$, where OPPO attains an $H^{3/2}|{\mathcal{S}}|^{2}|\mathcal{A}|\sqrt{T}$-regret up to logarithmic factors.

Broadly speaking, our work is related to a vast body of work on value-based reinforcement learning in tabular (Jaksch et al., 2010; Osband et al., 2014; Osband and Van Roy, 2016; Azar et al., 2017; Dann et al., 2017; Strehl et al., 2006; Jin et al., 2018) and linear settings (Yang and Wang, 2019b, a; Jin et al., 2019; Ayoub et al., 2020; Zhou et al., 2020), as well as nonlinear settings involving general function approximators (Wen and Van Roy, 2017; Jiang et al., 2017; Du et al., 2019b; Dong et al., 2019). In particular, our setting is the same as the linear setting studied by Ayoub et al. (2020); Zhou et al. (2020), which generalizes the one proposed by Yang and Wang (2019a). We remark that our setting differs from the linear setting studied by Yang and Wang (2019b); Jin et al. (2019). It can be shown that the two settings are incomparable in the sense that one does not imply the other (Zhou et al., 2020). Also, our setting is related to the low-Bellman-rank setting studied by Jiang et al. (2017); Dong et al. (2019). In comparison, we focus on policy-based reinforcement learning, which is significantly less studied in theory. In particular, compared with the work of Yang and Wang (2019b, a); Jin et al. (2019); Ayoub et al. (2020); Zhou et al. (2020), which focuses on value-based reinforcement learning, OPPO attains the same $\sqrt{T}$-regret even in the presence of adversarially chosen reward functions. Compared with optimism-led iterative value-function elimination (OLIVE) (Jiang et al., 2017; Dong et al., 2019), which handles the more general low-Bellman-rank setting but is only sample-efficient, OPPO simultaneously attains computational efficiency and sample efficiency in the linear setting. Despite the differences between policy-based and value-based reinforcement learning, our work shows that the general principle of “optimism in the face of uncertainty” (Auer et al., 2002; Bubeck and Cesa-Bianchi, 2012) can be carried over from existing algorithms based on value iteration, e.g., optimistic LSVI, into policy optimization algorithms, e.g., NPG, TRPO, and PPO, to make them sample-efficient, which further leads to a new general principle of “conservative optimism in the face of uncertainty and adversary” that additionally allows adversarially chosen reward functions.

2 Notation

Throughout this paper, we denote by $C,C^{\prime},C^{\prime\prime},\ldots$ absolute constants whose values can vary from line by line.

Preliminaries

In this paper, we consider an episodic MDP $({\mathcal{S}},\mathcal{A},H,{\mathcal{P}},r)$, where ${\mathcal{S}}$ and $\mathcal{A}$ are the state and action spaces, respectively, $H$ is the length of each episode, ${\mathcal{P}}_{h}(\cdot\,|\,\cdot,\cdot)$ is the transition kernel from a state-action pair to the next state at the $h$-th step of each episode, and $r^{k}_{h}:{\mathcal{S}}\times\mathcal{A}\rightarrow$ is the reward function at the $h$-th step of the $k$-th episode. We assume that the reward function is deterministic, which is without loss of generality, as our subsequent regret analysis readily generalizes to the setting where the reward function is stochastic.

At the beginning of the $k$-th episode, the agent determines a policy $\pi^{k}=\{\pi^{k}_{h}\}_{h=1}^{H}\in\Delta(\mathcal{A}\,|\,{\mathcal{S}},H)$. We assume that the initial state $x^{k}_{1}$ is fixed to $x_{1}\in{\mathcal{S}}$ across all the episodes, which is without loss of generality, as our subsequent regret analysis readily generalizes to the setting where $x^{k}_{1}$ is sampled from a fixed distribution across all the episodes. Then the agent iteratively interacts with the environment as follows. At the $h$-th step, the agent receives a state $x^{k}_{h}$ and takes an action following $a^{k}_{h}\sim\pi^{k}_{h}(\cdot\,|\,x^{k}_{h})$. Subsequently, the agent receives a reward $r^{k}_{h}(x^{k}_{h},a^{k}_{h})$ and the next state following $x^{k}_{h+1}\sim{\mathcal{P}}_{h}(\cdot\,|\,x^{k}_{h},a^{k}_{h})$. The $k$-th episode ends after the agent receives the last reward $r^{k}_{H}(x^{k}_{H},a^{k}_{H})$.

We allow the reward function $r^{k}=\{r^{k}_{h}\}_{h=1}^{H}$ to be adversarially chosen by the environment at the beginning of the $k$-th episode, which can depend on the $(k-1)$ historical trajectories. The reward function $r^{k}_{h}$ is revealed to the agent after it takes the action $a^{k}_{h}$ at the state $x^{k}_{h}$, which together determine the received reward $r^{k}_{h}(x^{k}_{h},a^{k}_{h})$. We define the regret in terms of competing with the globally optimal policy in hindsight (Cesa-Bianchi and Lugosi, 2006; Bubeck and Cesa-Bianchi, 2012) as

By the definitions in (2.2) and (2.3), we have the following Bellman equation,

2 Linear Function Approximations

We consider the linear setting where the transition dynamics are linear in a feature map, which is formalized in the following assumption.

for any $(x,a,x^{\prime})\in{\mathcal{S}}\times\mathcal{A}\times{\mathcal{S}}$. Also, we assume that

for any $(x,a)\in{\mathcal{S}}\times\mathcal{A}$ and $V:{\mathcal{S}}\rightarrow[0,H]$.

Algorithm and Theory

We present Optimistic PPO (OPPO) in Algorithm 1, which involves a policy improvement step and a policy evaluation step.

Policy Improvement Step. In the $k$-th episode, OPPO updates $\pi^{k}$ based on $\pi^{k-1}$ (Lines 4-9 of Algorithm 1). In detail, we define the following linear function of the policy $\pi\in\Delta(\mathcal{A}\,|\,{\mathcal{S}},H)$,

which is a local linear approximation of $V_{1}^{\pi,k-1}(x^{k}_{1})$ at $\pi^{k-1}$ (Schulman et al., 2015, 2017). In particular, we have that $L_{k-1}(\pi^{k-1})=V^{\pi^{k-1},k-1}_{1}(x^{k}_{1})$. The policy improvement step is defined by

Here the KL-divergence regularizes $\pi$ to be close to $\pi^{k-1}$ so that $L_{k-1}(\pi)$ well approximates $V^{\pi,k-1}_{1}(x^{k}_{1})$, which further ensures that the updated policy $\pi^{k}$ improves the expected total reward (associated with the reward function $r^{k-1}$) upon $\pi^{k-1}$. Also, $\alpha>0$ is the stepsize, which is specified in Theorem 3.1. By executing the updated policy $\pi^{k}$, the agent receives the state-action sequence $\{(x^{k}_{h},a^{k}_{h})\}_{h=1}^{H}$ and observes the reward function $r^{k}$, which together determine the received rewards $\{r^{k}_{h}(x^{k}_{h},a^{k}_{h})\}_{h=1}^{H}$.

The policy improvement step defined in (3.2) corresponds to one iteration of NPG (Kakade, 2002), TRPO (Schulman et al., 2015), and PPO (Schulman et al., 2017). In particular, PPO solves the same KL-regularized policy optimization subproblem as in (3.2) at each iteration, while TRPO solves an equivalent KL-constrained subproblem. In the special case where the reward function $r^{k-1}_{h}$ is linear in the feature map $\phi^{k-1}_{h}$ defined subsequently, which implies that the Q-function $Q^{\pi^{k-1},k-1}_{h}$ is also linear in $\phi^{k-1}_{h}$, the updated policy $\pi^{k}$ can be equivalently obtained by one iteration of NPG when the policy is parameterized by an energy-based distribution (Agarwal et al., 2019; Wang et al., 2019). Such a policy improvement step can also be cast as one iteration of infinite-dimensional mirror descent (Nemirovsky and Yudin, 1983) or dual averaging (Xiao, 2010), where the Q-function plays the role of the gradient (Liu et al., 2019; Wang et al., 2019).

The updated policy $\pi^{k}$ obtained in (3.2) takes the following closed form,

for any $h\in[H]$ and $x\in{\mathcal{S}}$. However, the Q-function $Q^{\pi^{k-1},k-1}_{h}$ remains to be estimated through the subsequent policy evaluation step. We denote by $Q^{k-1}_{h}$ the estimated Q-function, which replaces the Q-function $Q^{\pi^{k-1},k-1}_{h}$ in (3.1)-(3.3) and is correspondingly used in Line 6 of Algorithm 1.

Policy Evaluation Step. At the end of the $k$-th episode, OPPO evaluates the policy $\pi^{k}$ based on the $(k-1)$ historical trajectories (Lines 11-18 of Algorithm 1). In detail, for any $h\in[H]$, we define the empirical mean-squared Bellman error (MSBE) (Sutton and Barto, 2018) as

while we initialize $V^{\tau}_{H+1}$ as a zero function on ${\mathcal{S}}$. The policy evaluation step is defined by iteratively updating the estimated Q-function $Q^{k}=\{Q^{k}_{h}\}_{h=1}^{H}$ associated with the reward function $r^{k}=\{r^{k}_{h}\}_{h=1}^{H}$ by

Here $\beta>0$ scales with $d$, $H$, and $K$, which is specified in Theorem 3.1.

The policy evaluation step defined in (3.1) corresponds to one iteration of least-squares temporal difference (LSTD) (Bradtke and Barto, 1996; Boyan, 2002). In particular, as we have

with high probability, which is subsequently characterized in Lemma 4.3. Here the inequality holds uniformly for any $(h,k)\in[H]\times[K]$ and $(x,a)\in{\mathcal{S}}\times\mathcal{A}$. As the fact that $r^{k}_{h}\in$ for any $h\in[H]$ implies that $Q^{\pi^{k},k}_{h}\in[0,H-h+1]$, we truncate $Q^{k}_{h}$ to the range $[0,H-h+1]$ in (3.1), which is correspondingly used in Line 17 of Algorithm 1.

2 Regret Analysis

We establish an upper bound of the regret of OPPO (Algorithm 1) in the following theorem. Recall that the regret is defined in (2.1) and $T=HK$ is the total number of steps taken by the agent, where $H$ is the length of each episode and $K$ is the total number of episodes. Also, $|\mathcal{A}|$ is the cardinality of $\mathcal{A}$ and $d$ is the dimension of the feature map $\psi$.

Let $\alpha=\sqrt{2\log{|\mathcal{A}|}/(HT)}$ in (3.2) and Line 6 of Algorithm 1, $\lambda=1$ in (3.1) and Line 12 of Algorithm 1, and $\beta=C\sqrt{dH^{2}\cdot\log(dT/\zeta)}$ in (3.1) and Line 15 of Algorithm 1, where $C>1$ is an absolute constant and $\zeta\in(0,1]$. Under Assumption 2.1 and the assumption that $\log|\mathcal{A}|=O(d^{2}\cdot[\log(dT/\zeta)]^{2})$, the regret of OPPO satisfies

with probability at least $1-\zeta$, where $C^{\prime}>0$ is an absolute constant.

See Section 4 for a proof sketch and Appendix C for a detailed proof. ∎

Theorem 3.1 proves that OPPO attains a $\sqrt{d^{2}H^{3}T}$-regret up to logarithmic factors, where the dependency on the total number of steps $T$ is optimal. In the stationary setting where the reward function and initial state are fixed across all the episodes, such a regret translates to a $d^{2}H^{4}/\varepsilon^{2}$-sample complexity (up to logarithmic factors) following the argument of Jin et al. (2018) (Section 3.1). Here $\varepsilon>0$ measures the suboptimality of the obtained policy $\pi^{k}$ in the following sense,

Discussion of Mechanisms. In the sequel, we consider the ideal setting where the transition dynamics are known, which, by the Bellman equation defined in (2.4), allows us to access the Q-function $Q_{h}^{\pi,k}$ for any policy $\pi$ and $(h,k)\in[H]\times[K]$ once given the reward function $r^{k}$. The following lemma connects the difference between two policies to the difference between their expected total rewards through the Q-function.

For any policies $\pi,\pi^{\prime}\in\Delta(\mathcal{A}\,|\,{\mathcal{S}},H)$ and $k\in[K]$, it holds that

For notational simplicity, we omit the conditioning on $x_{1}=x^{k}_{1}$, e.g., in (3.7) of Lemma 3.2 subsequently. The following lemma characterizes the policy improvement step defined in (3.2), where the updated policy $\pi^{k}$ takes the closed form in (3.3).

For any distributions $p^{*},p\in\Delta(\mathcal{A})$, state $x\in{\mathcal{S}}$, and function $Q:{\mathcal{S}}\times\mathcal{A}\rightarrow[0,H]$, it holds for $p^{\prime}\in\Delta(\mathcal{A})$ with $p^{\prime}(\cdot)\propto p(\cdot)\cdot\exp\{\alpha\cdot Q(x,\cdot)\}$ that

Corresponding to the definition of the regret in (2.1), we define the globally optimal policy in hindsight (Cesa-Bianchi and Lugosi, 2006; Bubeck and Cesa-Bianchi, 2012) as

which attains a zero-regret. In the ideal setting where the Q-function $Q^{\pi^{k},k}_{h}$ associated with the reward function $r^{k}$ is known and the updated policy $\pi^{k+1}_{h}$ takes the closed form in (3.3), Lemma 3.3 implies

for any $(h,k)\in[H]\times[K]$ and $x\in{\mathcal{S}}$. Combining (3.2) with Lemma 3.2, we obtain

Here the first inequality follows from telescoping the right-hand side of (3.2) across all the episodes and the fact that the KL-divergence is nonnegative. Also, the second inequality follows from the initialization of the policy and Q-function in Line 1 of Algorithm 1. Setting $\alpha=\sqrt{2\log|\mathcal{A}|/(HT)}$ in (3.2), we establish a $\sqrt{H^{3}T\cdot\log|\mathcal{A}|}$-regret in the ideal setting.

Such an ideal setting demonstrates the key role of the KL-divergence in the policy improvement step defined in (3.2), where $\alpha>0$ is the stepsize. Intuitively, without the KL-divergence, that is, setting $\alpha\rightarrow\infty$, the upper bound of the regret on the right-hand side of (3.2) tends to infinity. In fact, for any $\alpha<\infty$, the updated policy $\pi_{h}^{k}$ in (3.3) is “conservatively” greedy with respect to the Q-function $Q^{\pi^{k-1},k-1}_{h}$ associated with the reward function $r^{k-1}$. In particular, the regularization effect of both $\pi^{k-1}_{h}$ and $\alpha$ in (3.3) ensures that $\pi^{k}_{h}$ is not “fully” committed to perform well only with respect to $r^{k-1}$, just in case the subsequent adversarially chosen reward function $r^{k}$ significantly differs from $r^{k-1}$. In comparison, the “fully” greedy policy improvement step, which is commonly adopted by the existing work on value-based reinforcement learning (Jaksch et al., 2010; Osband et al., 2014; Osband and Van Roy, 2016; Azar et al., 2017; Dann et al., 2017; Strehl et al., 2006; Jin et al., 2018, 2019; Yang and Wang, 2019b, a), lacks such a notion of robustness. On the other hand, an intriguing question is whether being “conservatively” greedy is less sample-efficient than being “fully” greedy in the stationary setting, where the reward function is fixed across all the episodes. In fact, in the ideal setting where the Q-function $Q^{\pi^{k-1},k-1}_{h}$ associated with the reward function $r^{k-1}$ in (3.3) is known, the “fully” greedy policy improvement step with $\alpha\rightarrow\infty$ corresponds to one step of policy iteration (Sutton and Barto, 2018), which converges to the globally optimal policy $\pi^{*}$ within $K=H$ episodes and hence equivalently induces an $H^{2}$-regret. However, in the realistic setting, the Q-function $Q^{\pi^{k-1},k-1}_{h}$ in (3.1)-(3.3) is replaced by the estimated Q-function $Q^{k-1}_{h}$ in Line 6 of Algorithm 1, which is obtained by the policy evaluation step defined in (3.1). As a result of the estimation uncertainty that arises from only observing finite historical data, it is indeed impossible to do better than the $\sqrt{T}$-regret even in the tabular setting (Jin et al., 2018), which is shown to be an information-theoretic lower bound. In the linear setting, OPPO attains such a lower bound in terms of the total number of steps $T=HK$. In other words, in the stationary setting, being “conservatively” greedy suffices to achieve sample-efficiency, which complements its advantages in terms of robustness in the more challenging setting with adversarially chosen reward functions.

Proof Sketch

For the simplicity of discussion, we define the model prediction error as

For any $(k,h)\in[K]\times[H]$, we define $\mathcal{F}_{k,h,1}$ as the $\sigma$-algebra generated by the following state-action sequence and reward functions,

and $\mathcal{F}_{k,h,2}$ as the $\sigma$-algebra generated by

where, for the simplicity of discussion, we define $x^{k}_{H+1}$ as a null state for any $k\in[K]$. The $\sigma$-algebra sequence $\{\mathcal{F}_{k,h,m}\}_{(k,h,m)\in[K]\times[H]\times}$ is a filtration with respect to the timestep index

In other words, for any $t(k,h,m)\leq t(k^{\prime},h^{\prime},m^{\prime})$, it holds that $\mathcal{F}_{k,h,m}\subseteq\mathcal{F}_{k^{\prime},h^{\prime},m^{\prime}}$.

By the definition of the $\sigma$-algebra $\mathcal{F}_{k,h,m}$, for any $(k,h)\in[K]\times[H]$, the estimated value function $V^{k}_{h}$ and Q-function $Q^{k}_{h}$ are measurable to $\mathcal{F}_{k,1,1}$, as they are obtained based on the $(k-1)$ historical trajectories and the reward function $r^{k}$ adversarially chosen by the environment at the beginning of the $k$-th episode, both of which are measurable to $\mathcal{F}_{k,1,1}$.

In the following lemma, we decompose the regret defined in (2.1) into three terms. Recall that the globally optimal policy in hindsight $\pi^{*}$ is defined in (3.8) and the model prediction error $\iota^{k}_{h}$ is defined in (4.1).

which is independent of the linear setting in Assumption 2.1. Here $\{\mathcal{M}_{k,h,m}\}_{(k,h,m)\in[K]\times[H]\times}$ is a martingale adapted to the filtration $\{\mathcal{F}_{k,h,m}\}_{(k,h,m)\in[K]\times[H]\times}$, both with respect to the timestep index $t(k,h,m)$ defined in (4.2) of Definition 4.1.

Lemma 4.2 allows us to characterize the regret by upper bounding terms (i), (ii), and (iii) in (4.2), respectively. In detail, term (i) corresponds to the right-hand side of (3.2) in Lemma 3.2 with the Q-function $Q^{\pi^{k},k}_{h}$ replaced by the estimated Q-function $Q^{k}_{h}$, which is obtained by the policy evaluation step defined in (3.1). In particular, as the updated policy $\pi^{k+1}_{h}$ is obtained by the policy improvement step in Line 6 of Algorithm 1 using $\pi^{k}_{h}$ and $Q^{k}_{h}$, term (i) can be upper bounded following a similar analysis to the discussion in Section 3.2, which is based on Lemmas 3.2 and 3.3 as well as (3.2). Also, by the Azuma-Hoeffding inequality, term (ii) is a martingale that scales as $O(B_{\mathcal{M}}\sqrt{T_{\mathcal{M}}})$ with high probability, where $T_{\mathcal{M}}$ is the total number of timesteps and $B_{\mathcal{M}}$ is an upper bound of the martingale differences. More specifically, we prove that $T_{\mathcal{M}}=2HK=2T$ and $B_{\mathcal{M}}=2H$ in Appendix C, which implies that term (ii) is $O(\sqrt{H^{2}T})$ with high probability. Meanwhile, term (iii) corresponds to the model prediction error, which is characterized subsequently in Section 4.2. Note that the regret decomposition in (4.2) of Lemma 4.2 is independent of the linear setting in Assumption 2.1, and therefore, applies to any forms of estimated Q-functions $Q^{k}_{h}$ in more general settings. In particular, as long as we can upper bound term (iii) in (4.2), our regret analysis can be carried over even beyond the linear setting.

2 Model Prediction Error

To upper bound term (iii) in (4.2) of Lemma 4.2, we characterize the model prediction error $\iota^{k}_{h}$ defined in (4.1) in the following lemma. Recall that the bonus function $\Gamma^{k}_{h}$ is defined in (3.1).

Let $\lambda=1$ in (3.1) and Line 12 of Algorithm 1, and $\beta=C\sqrt{dH^{2}\cdot\log(dT/\zeta)}$ in (3.1) and Line 15 of Algorithm 1, where $C>1$ is an absolute constant and $\zeta\in(0,1]$. Under Assumption 2.1, it holds with probability at least $1-\zeta/2$ that

for any $(k,h)\in[K]\times[H]$ and $(x,a)\in{\mathcal{S}}\times\mathcal{A}$.

Lemma 4.3 demonstrates the key role of uncertainty quantification in achieving sample-efficiency. More specifically, due to the uncertainty that arises from only observing finite historical data, the model prediction error $\iota^{k}_{h}(x,a)$ can be possibly large for the state-action pairs $(x,a)$ that are less visited or even unseen. However, as is shown in Lemma 4.3, explicitly incorporating the bonus function $\Gamma^{k}_{h}$ into the estimated Q-function $Q^{k}_{h}$ ensures that $\iota^{k}_{h}(x,a)\leq 0$ with high probability for any $(k,h)\in[K]\times[H]$ and $(x,a)\in{\mathcal{S}}\times\mathcal{A}$. In other words, the estimated Q-function $Q^{k}_{h}$ is “optimistic in the face of uncertainty”, as $\iota^{k}_{h}(x,a)\leq 0$ or equivalently

To illustrate the intuition behind the model prediction error $\iota^{k}_{h}$ defined in (4.1), we define the implicitly estimated transition dynamics as

where $\Lambda^{k}_{h}$ is defined in (3.1). Correspondingly, the policy evaluation step defined in (3.1) takes the following equivalent form (ignoring the truncation step for the simplicity of discussion),

Conclusion

We study the sample efficiency of policy-based reinforcement learning in the episodic setting of linear MDPs with full-information feedback. We proposed an optimistic variant of the proximal policy optimization algorithm, dubbed as OPPO, which incorporates the principle of “optimism in the face of uncertainty” into policy optimization. When applied to the episodic MDP with unknown transition and adversarial reward, OPPO provably achieves a near-optimal $\sqrt{d^{2}H^{3}T}$-regret. To the best of our knowledge, OPPO is the first provably efficient policy optimization algorithm that explicitly incorporates exploration.

Acknowledgements

The authors would like to thank Lingxiao Wang, Wen Sun, and Sham Kakade for pointing out a technical issue in the first version regarding the covering number of value functions in the linear setting. This version has fixed the technical issue with a definition of the linear MDP different from the one in the first version. The authors would also like to thank Csaba Szepesvári, Lin F. Yang, Yining Wang, and Simon S. Du for helpful discussions.

References

Appendix A Proofs of Lemmas in Section 3

Meanwhile, by (A.2) we have that, on the right-hand side of (A.1),

where the last equality follows from (2.4). Combining (A.1), (A.1), (A.5), and the linearity of the Bellman evaluation operator defined in (A.1), we obtain

which concludes the proof of Lemma 3.2. ∎

A.2 Proof of Lemma 3.3

which implies that $\langle z,p^{*}-p^{\prime}\rangle=0$ in (A.2) as $p^{\prime},p^{*}\in\Delta(\mathcal{A})$. Moreover, by (A.2) we have

for any state $x\in{\mathcal{S}}$. Meanwhile, by Pinsker’s inequality, it holds that

Combining (A.2), (A.8), and the fact that $\|Q(x,\cdot)\|_{\infty}\leq H$ for any state $x\in{\mathcal{S}}$, we obtain

which concludes the proof of Lemma 3.3. ∎

Appendix B Proofs of Lemmas in Section 4

for any $(k,h)\in[K]\times[H]$ and state $x\in{\mathcal{S}}$.

We decompose the instantaneous regret at the $k$-th episode into the following two terms,

for any $(k,h)\in[K]\times[H]$. For any $k\in[K]$, recursively expanding (B.6) across $h\in[H]$ yields

where $V^{\pi^{*},k}_{H+1}=V^{k}_{H+1}={\bm{0}}$. Therefore, we obtain

for any $(k,h)\in[K]\times[H]$. By the definition of the model prediction error $\iota^{k}_{h}$ in (4.1), we have

where the last equality follows from (2.4). Plugging (B.1) into (B.8), we obtain

for any $(k,h)\in[K]\times[H]$. For any $k\in[K]$, recursively expanding (B.11) across $h\in[H]$ yields

where $V^{k}_{H+1}(x^{k}_{H+1})=V^{\pi^{k},k}_{H+1}(x^{k}_{H+1})=0$. Therefore, we obtain

By Definition 4.1 and the definitions of $D_{k,h,1}$ and $D_{k,h,2}$ in (B.11), we have

for any $(k,h)\in[K]\times[H]$. Here we have that $\mathcal{F}_{k,0,2}=\mathcal{F}_{k-1,H,2}$ for any $k\geq 2$, as (4.2) of Definition 4.1 implies

Also, we define $\mathcal{F}_{1,0,2}$ to be empty. Thus, (B.13) allows us to define the martingale

with respect to the timestep index $t(k,h,m)$ defined in (4.2) of Definition 4.1. Such a martingale is adapted to the filtration $\{\mathcal{F}_{k,h,m}\}_{(k,h,m)\in[K]\times[H]\times}$. In particular, we have that, on the right-hand side of (B.12),

Combining (B.3), (B.7), (B.12), and (B.15), we obtain

which concludes the proof of Lemma 4.2. ∎

B.2 Proof of Lemma 4.3

Recall that $\phi_{h}^{k}$ defined in (3.4) takes the following form,

for any $(k,h)\in[K]\times[H]$ and $(x,a)\in{\mathcal{S}}\times\mathcal{A}$. Also, recall that the estimated Q-function $Q^{k}_{h}$ obtained by the policy evaluation step defined in (3.1) takes the following form,

for any $(k,h)\in[K]\times[H]$ and $(x,a)\in{\mathcal{S}}\times\mathcal{A}$. Here $\Gamma^{k}_{h}$ and $\Lambda_{h}^{k}$ are defined in (3.1). Meanwhile, by Assumption 2.1 we have

for any $(k,h)\in[K]\times[H]$ and $(x,a)\in{\mathcal{S}}\times\mathcal{A}$. Plugging the definition of $\Lambda^{k}_{h}$ in (3.1) into (B.2), we obtain

for any $(k,h)\in[K]\times[H]$ and $(x,a)\in{\mathcal{S}}\times\mathcal{A}$. Here the second equality follows from (B.2) with $V^{k}_{h+1}$ replaced by $V_{h+1}^{\tau}$ for any $\tau\in[k-1]$. Combining (B.16) and (B.2), we obtain

for any $(k,h)\in[K]\times[H]$ and $(x,a)\in{\mathcal{S}}\times\mathcal{A}$.

Term (i). As is defined in (3.1), $(\Lambda^{k}_{h})^{-1}$ is a positive-definite matrix. By the Cauchy-Schwarz inequality, the absolute value of term (i) is upper bounded as

for any $(k,h)\in[K]\times[H]$ and $(x,a)\in{\mathcal{S}}\times\mathcal{A}$. Under the event $\mathcal{E}$ defined in (D.1) of Lemma D.1, which happens with probability at least $1-\zeta/2$, it holds that

for any $(k,h)\in[K]\times[H]$ and $(x,a)\in{\mathcal{S}}\times\mathcal{A}$. Here $C^{\prime\prime}>0$ is an absolute constant and $\zeta\in(0,1]$.

Term (ii). Similar to (B.20), the absolute value of term (ii) is upper bounded as

for any $(k,h)\in[K]\times[H]$ and $(x,a)\in{\mathcal{S}}\times\mathcal{A}$. Here the first inequality follows from the Cauchy-Schwarz inequality, the second inequality follows from the fact that $\Lambda^{k}_{h}\succeq\lambda\cdot{I}$, and the last inequality follows from Assumption 2.1, which assumes that $\|\theta_{h}\|_{2}\leq\sqrt{d}$.

Combining (B.19), (B.21), (B.2), and the fact that $\lambda=1$, we obtain

for any $(k,h)\in[K]\times[H]$ and $(x,a)\in{\mathcal{S}}\times\mathcal{A}$ under the event $\mathcal{E}$ defined in (D.1) of Lemma D.1. Here $C>1$ is an absolute constant. Setting

in the bonus function $\Gamma^{k}_{h}$ defined in (3.1), by (B.2) we obtain

for any $(k,h)\in[K]\times[H]$ and $(x,a)\in{\mathcal{S}}\times\mathcal{A}$ under $\mathcal{E}$. Hence, for the model prediction error $\iota^{k}_{h}$ defined in (4.1), by (B.16), (B.24), and (B.25) we have

for any $(k,h)\in[K]\times[H]$ and $(x,a)\in{\mathcal{S}}\times\mathcal{A}$ under $\mathcal{E}$. Thus, combining (B.2), (B.2), and Lemma D.1, which implies that $\mathcal{E}$ happens with probability at least $1-\zeta/2$, we conclude the proof of Lemma 4.3. ∎

Appendix C Proof of Theorem 3.1

We upper bound terms (i)-(iii) in (4.2) of Lemma 4.2 respectively, that is,

Term (i). By Lemma 3.3 and the policy improvement step in Line 6 of Algorithm 1, we have

Here the second last inequality follows from the fact that the KL-divergence is nonnegative. Also, the last inequality follows from the initialization of the policy and Q-function in Line 1 of Algorithm 1, which implies that $\pi^{1}_{h}(\cdot\,|\,x_{h})$ is a uniform distribution on $\mathcal{A}$ and hence

Here the inequality follows from the fact that the entropy of $\pi^{*}_{h}(\cdot\,|\,x_{h})$ is nonnegative. Thus, setting $\alpha=\sqrt{2\log{|\mathcal{A}|}/(HT)}$ in Line 6 of Algorithm 1, by (C) we obtain

Term (ii). Recall that the martingale differences $D_{k,h,1}$ and $D_{k,h,2}$ defined in (B.11) take the following forms,

By the truncation of $Q^{k}_{h}$ to the range $[0,H-h+1]$ in (3.1), we have

which implies that $|D_{k,h,1}|\leq 2H$ and $|D_{k,h,2}|\leq 2H$ for any $(k,h)\in[K]\times[H]$. Therefore, applying the Azuma-Hoeffding inequality to the martingale defined in (B.1), we obtain

for any $t>0$. Setting $t=\sqrt{16H^{2}T\cdot\log(4/\zeta)}$ with $\zeta\in(0,1]$, we obtain

with probability at least $1-\zeta/2$, where $T=HK$.

Term (iii). As is shown in Lemma 4.3, it holds with probability at least $1-\zeta/2$ that

for any $(k,h)\in[K]\times[H]$ and $(x,a)\in{\mathcal{S}}\times\mathcal{A}$. Meanwhile, by the definitions of $\iota_{h}^{k}$ and $Q_{h}^{k}$ in (4.1) and (3.1), respectively, we have that $|\iota_{h}^{k}(x,a)|\leq 2H$, which together with (C.4) implies

for any $(k,h)\in[K]\times[H]$ and $(x,a)\in{\mathcal{S}}\times\mathcal{A}$ with probability at least $1-\zeta/2$. Hence, we obtain

with probability at least $1-\zeta/2$. By the definition of $\Gamma^{k}_{h}$ in (3.1), we have

with $C>1$ being an absolute constant, which implies that $H\leq\beta$. Thus, (C.6) implies

By Lemma D.3 and the definition of $\Lambda^{k}_{h}$ in (3.1), we obtain

for any $h\in[H]$, where $\Lambda^{1}_{h}=\lambda\cdot I$ and $\Lambda^{K+1}_{h}\in\mathcal{F}_{K,H,2}$ by Definition 4.1. Moreover, Assumption 2.1 implies

for any $(k,h)\in[K]\times[H]$ and $(x,a)\in{\mathcal{S}}\times\mathcal{A}$, which further implies

for any $h\in[H]$. As we set $\lambda=1$, it holds for any $h\in[H]$ that

Combining (C.8)-(C.10) and the Cauchy-Schwarz inequality, we obtain

By (C.5), (C.7), and (C), it holds with probability at least $1-\zeta/2$ that

where $C>1$ is an absolute constant, $\zeta\in(0,1]$, and $T=HK$.

Plugging the upper bounds of terms (i)-(iii) in (C.2), (C.3), and (C), respectively, into (4.2) of Lemma 4.2, we obtain

with probability at least $1-\zeta$, where $C^{\prime}>0$ is an absolute constant. Here we use the fact that $\log|\mathcal{A}|=O(d^{2}\cdot[\log(dT/\zeta)]^{2})$ in (C.2) and (C). Therefore, we conclude the proof of Theorem 3.1. ∎

Appendix D Supporting Lemmas

In this section, we present the supporting lemmas.

Let $\lambda=1$ in (3.1) and Line 12 of Algorithm 1. For any $\zeta\in(0,1]$, the event $\mathcal{E}$ that, for any $(k,h)\in[K]\times[H]$,

happens with probability at least $1-\zeta/2$, where $C^{\prime\prime}>0$ is an absolute constant.

By the definition of the filtration $\{\mathcal{F}_{k,h,m}\}_{(k,h,m)\in[K]\times[H]\times}$ in Definition 4.1 and the Markov property, we have

Conditioning on $\mathcal{F}_{\tau,h,1}$, the only randomness comes from $x^{\tau}_{h+1}$, while $V^{\tau}_{h+1}$ is a deterministic function. To see this, note that $V^{\tau}_{h+1}$ is determined by $Q^{\tau}_{h+1}$ and $\pi^{\tau}_{h+1}$, which are further determined by the historical data in $\mathcal{F}_{\tau,h,1}$. We define

By (D.2), conditioning on $\mathcal{F}_{\tau,h,1}$, $\eta_{\tau,h}$ is a zero-mean random variable. Moreover, as $V^{\tau}_{h+1}\in[0,H]$, conditioning on $\mathcal{F}_{\tau,h,1}$, $\eta_{\tau,h}$ is an $H/2$-sub-Gaussian random variable, which is defined in (D.5) of Lemma D.2. Also, $\eta_{\tau,h}$ is $\mathcal{F}_{k,h,2}$-measurable, as $\mathcal{F}_{\tau,h,1}\subseteq\mathcal{F}_{k,h,2}$ for any $\tau\in[k-1]$. Hence, for any fixed $h\in[H]$, by Lemma D.2, it holds with probability at least $1-\zeta/(2H)$ that

for any $k\in[K]$. To upper bound $\det(\Lambda_{h}^{k})$ in (D), recall that $\Lambda^{k}_{h}$ is defined by

By the triangle inequality, the spectral norm of $\Lambda^{k}_{h}$ is upper bounded as

Here the last inequality follows from Assumption 2.1, which implies

for any $V\colon{\mathcal{S}}\rightarrow[0,H]$. Hence, $\det(\Lambda^{k}_{h})$ in (D) is upper bounded as

Moreover, setting $\lambda=1$, combining (D) and (D.4), and applying the union bound for any $h\in[H]$, we obtain that, with probability at least $1-\zeta/2$,

for any $(k,h)\in[K]\times[H]$, where $C^{\prime\prime}>0$ is an absolute constant. Thus, we conclude the proof of Lemma D.1. ∎

For any $\delta>0$, it holds with probability at least $1-\delta$ that

See Theorem 1 of Abbasi-Yadkori et al. (2011) for a detailed proof. ∎

See Lemma 11 of Abbasi-Yadkori et al. (2011) for a detailed proof. ∎