Efficient Local Planning with Linear Function Approximation

Theoretical guarantees for the random access model have been obtained for the tabular setting (Sidford et al., 2018; Agarwal et al., 2020b; Li et al., 2020; Azar et al., 2013). As for linear function approximation, different assumptions have been made for theoretical analysis. Under the linear MDP assumption, Yang and Wang (2019) derived an optimal $\mathcal{O}(d\kappa^{-2}(1-\gamma)^{-3})$ query complexity bound by a variance-reduction Q-learning type algorithm. Under the $Q_{\pi}$-realizability of all determinstic policies (a strictly weaker assumption than linear MDP (Zanette et al., 2020)), Du et al. (2020) showed a negative result for the settings with model misspecification error $\epsilon=\Omega(\sqrt{H/d})$ (see also Van Roy and Dong (2019); Lattimore et al. (2020)). When $\epsilon=o((1-\gamma)^{2}/\sqrt{d})$, assuming the access to the full feature matrix, Lattimore et al. (2020) proposed algorithms with polynomial query costs, and Du et al. (2020) proposed similar algorithm for the exact $Q_{\pi}$ realizability setting. Since these works need to find a globally optimal design or barycentric spanner, their computational costs depend polynomially on the size of the state space. Under the $V^{*}$-realizability assumption (i.e., the optimal value function is linear in some feature map), Shariff and Szepesvári (2020) proposed a planning algorithm assuming the availability of a set of core states but obtaining such core states can still be computationally inefficient. Zanette et al. (2019) proposed an algorithm that uses a similar concept named anchor points but only provided a greedy heuristic to generate these points. A notable negative result is established in Weisz et al. (2021b) that shows that with only $Q^{*}$-realizability, any agent requires $\min(\exp(\Omega(d)),\exp(\Omega(H)))$ queries to learn an optimal policy.

Many prior studies have used simulators in tree-search style algorithms (Kearns et al., 2002; Munos, 2014). Under this setting, for the online planning problem, recently Weisz et al. (2021a) established an $\mathcal{O}((dH/\kappa)^{|\mathcal{A}|})$ query cost bound to learn an $\kappa$-optimal policy by the TensorPlan algorithm assuming the $V^{*}$-realizability. Whenever the action set is small, TensorPlan is query efficient, but its computational efficiency is left as an open problem. Under $Q^{*}$-realizability and constant sub-optimality gap, for the offline planning problem, Li et al. (2021) proposed an algorithm with $\operatorname{poly}(d,H,\kappa^{-1},\Delta_{\text{gap}}^{-1})$ query and computational costs.

As mentioned, many online RL algorithms can be converted to a policy optimization algorithm under the online access model using online-to-batch conversion. There is a large body of literature on online RL with linear function approximation and here we discuss a non-exhaustive list of prior works. Under the $Q^{*}$-realizability assumption, assuming that the probability transition of the MDP is deterministic, Wen and Van Roy (2013) proposed a sample and computationally efficient algorithm via the eluder dimension (Russo and Van Roy, 2013). Assuming the MDP has low Bellman rank, Jiang et al. (2017) proposed an algorithm that is sample efficient but computationally inefficient, and similar issues arise in Zanette et al. (2020) under the low inherent Bellman error assumption. Du et al. (2021) proposed a more general MDP class named bilinear class and provided a sample efficient algorithm, but the computational efficiency is unclear.

Under $Q_{\pi}$-realizability, several algorithms, such as Politex (Abbasi-Yadkori et al., 2019; Lazic et al., 2021), AAPI (Hao et al., 2021), and MDP-EXP2 (Wei et al., 2021) achieved sublinear regret in the infinite horizon average reward setting and are also computationally efficient. However, the corresponding analysis avoids the exploration issue by imposing a feature excitation assumption which may not be satisfied in many problems. Under the linear MDP assumption, Jin et al. (2020) established a $\mathcal{O}(\sqrt{d^{3}H^{3}T})$ regret bound for an optimistic least-square value iteration algorithm. Agarwal et al. (2020a) derived a $\operatorname{poly}(d,H,\kappa^{-1})$ sample cost bound for the policy cover-policy gradient algorithm, which can also be applied in the state aggregation setting; the algorithm and sample cost were subsequently improved in Zanette et al. (2021). Under the linear mixture MDP assumption (Yang and Wang, 2020; Zhou et al., 2020), Cai et al. (2020) proved an $\mathcal{O}(\sqrt{d^{3}H^{3}T})$ regret bound for an optimistic least square policy iteration (LSPI) type algorithm. A notable negative result for the online RL setting by Wang et al. (2021) shows that an exponentially large number of samples are needed if we only assume $Q^{*}$-realizability and constant sub-optimality gap. Other related works include Ayoub et al. (2020); Jin et al. (2021); Du et al. (2019); Wang et al. (2019), and references therein.

Let $\phi:{\mathcal{S}}\times\mathcal{A}\to\mathbb{R}^{d}$ be a feature map which assigns to each state-action pair a $d$-dimensional feature vector. For any $(s,a)\in{\mathcal{S}}\times\mathcal{A}$, the agent can obtain $\phi(s,a)$ with a computational cost of $\operatorname{poly}(d)$. Here, we emphasize that the computation of the feature vectors does not lead to a query cost. Without loss of generality, we impose the following bounded features assumption.

We consider the following two different assumptions on the linear realizability of the Q-functions:

We use a concept called core set. A core set $\mathcal{C}$ is a set of tuples $z=(s,a,\phi(s,a),q)\in{\mathcal{S}}\times\mathcal{A}\times\mathbb{R}^{d}\times(\mathbb{R}\cup\{\mathsf{none}\})$. The first three elements in the tuple denote a state, an action, and the feature vector corresponding to the state-action pair, respectively. The last element $q\in\mathbb{R}$ in the tuple denotes an estimate of $Q_{\pi}(s,a)$ for a policy $\pi$. During the algorithm, we may not always have such an estimate, in which case we write $q=\mathsf{none}$. For a tuple $z$, we use $z_{s}$, $z_{a}$, $z_{\phi}$, and $z_{q}$ to denote the $s$, $a$, $\phi$, and $q$ coordinates of $z$, respectively. We note that in prior works, the core set usually consists of the state-action pairs and their features (Lattimore et al., 2020; Du et al., 2020; Shariff and Szepesvári, 2020); whereas in this paper, for the convenience of notation, we also have the target values (Q-function estimates) in the core set elements. We denote by $\Phi_{\mathcal{C}}\in\mathbb{R}^{|\mathcal{C}|\times d}$ the feature matrix of all the elements in $\mathcal{C}$, i.e., each row of $\Phi_{\mathcal{C}}$ is the feature vector of an element in $\mathcal{C}$. Similarly, we define $q_{\mathcal{C}}\in\mathbb{R}^{|\mathcal{C}|}$ as the vector for the $Q_{\pi}$ estimate of all the tuples in $\mathcal{C}$.

It is also useful to introduce a notion of good set.

Intuitively, the good set is a set of vectors that are well-covered by the rows of $\Phi_{\mathcal{C}}$; in other words, these vectors are not closely aligned with the eigenvectors associated with the small eigenvalues of the covariance matrix of all the features in the core set.

As an overview, our algorithm Confident MC-LSPI works as follows. First, we initialize the core set using the initial state $\rho$ paired with all actions. Then, the algorithm runs least-squares policy iteration (Munos, 2003) to optimize the policy. This means that in each iteration, we estimate the Q-function value for every state-action pair in $\mathcal{C}$ using Monte Carlo rollout with the simulator, and learn a linear function to approximate the Q-function of the rollout policy, and the next policy is chosen to be greedy with respect to this linear function. Our second algorithm Confident MC-Politex works similarly, with the only difference being that instead of using the greedy policy iteration update rule, we use the mirror descent update rule with KL regularization between adjacent rollout policies (Even-Dar et al., 2009; Abbasi-Yadkori et al., 2019). Moreover, in both algorithms, whenever we observe a state-action pair whose feature is not in the good set during Monte Carlo rollout, we add the pair to the core set and restart the policy iteration process. We name the rollout subroutine ConfidentRollout. We discuss details in the following.

In our analysis, we say that we start a new loop whenever we start (or restart) the policy iteration process, i.e., going to line $(*)$ in Algorithm 2. By definition, when we start a new loop, the size of the core set $\mathcal{C}$ is increased by $1$. First, in Lemma 5.1 below, we show that the size of the core set will never exceed $C_{\max}=\widetilde{\mathcal{O}}(d)$. Therefore, the total number of loops is at most $C_{\max}$. In each loop, we run $K$ policy iterations; in each iteration, we run Algorithm 1 from at most $C_{\max}$ points from the core set; and each time when we run Algorithm 1, we query the simulator at most $\mathcal{O}(mn)$ times. Thus, for both algorithms, the total number of queries that we make is at most $C_{\max}^{2}Kmn$. Therefore, using the parameter choice in Theorems 5.1 and 5.2 and omitting logarithmic factors, we can obtain the query costs of Confident MC-LSPI and Politex in Table 2. As we can see, when $\epsilon=0$ or $\epsilon\neq 0$ but $\epsilon=o(1/\sqrt{d})$ (the regime we care about in this paper), the query cost of Confident MC-LSPI is lower than Politex. As for computational cost, since our policy improvement steps only involve matrix multiplication and matrix inversion, the computational cost is also polynomial in the aforementioned factors. One thing to notice is that during the rollout process, in each step, the agent needs to compute the features of a state paired with all actions, and thus the computational cost linearly depends on $|\mathcal{A}|$; on the contrary the query cost does not depend on $|\mathcal{A}|$ since in each step the agent only needs to query the simulator with the action sampled according to the policy.

We also note that when Assumption 3.3 holds, i.e., $\epsilon\neq 0$, the sub-optimality of the output policy is $\widetilde{\mathcal{O}}(\frac{\epsilon\sqrt{d}}{(1-\gamma)^{2}})$ for LSPI and $\widetilde{\mathcal{O}}(\frac{\epsilon\sqrt{d}}{1-\gamma})$ for Politex. Therefore, in the presence of a model misspecification error, Confident MC-Politex can achieve a better final sub-optimality than Confident MC-LSPI, although it’s query cost is higher.

We now discuss our proof strategy, focusing on LSPI for simplicity.

The first step is to show that our algorithm will terminate. This is equivalent to showing that the size of the core set $\mathcal{C}$ will not exceed certain finite quantity, since whenever we receive the $\mathsf{uncertain}$ status from ConfidentRollout, we increase the size of the core set by $1$, go back to line $(*)$ in Algorithm 2, and start a new loop. The following lemma shows that the size of the core set is always bounded, and thus the algorithm will always terminate.

This result first appears in Russo and Van Roy (2013) as the eluder dimension of linear function class. We present the proof of this lemma in Appendix A for completeness.

The next step is to analyze the gap between the value of the optimal policy and the policy $\pi$ parameterized by the vector $w_{K-1}$ that the algorithm outputs in the final loop, i.e., $V^{*}(\rho)-V_{\pi_{K-1}}(\rho)$. For ease of exposition, here we only consider the case of deterministic probability transition kernel $P$. Our full proof in Appendix B considers general stochastic dynamics.

To analyze our algorithm, we note that for approximate policy iteration (API) algorithms, if in every iteration (say the $k$-th iteration), we have an approximate Q-function that is close to the true Q-function of the rollout policy (say $\pi_{k-1}$) in $\ell_{\infty}$ norm, i.e., $\|Q_{k-1}-Q_{\pi_{k-1}}\|_{\infty}\leq\eta$, then existing results (Munos, 2003; Farahmand et al., 2010) ensure that we can learn a good policy if in every iteration we choose the new policy to be greedy with respect to the approximate Q-function. However, since we only have local access to the simulator, we cannot have such $\ell_{\infty}$ guarantee. In fact, as we show in the proof, we can only ensure that when $\phi(s,a)$ is in the good set $\mathcal{H}$, our linear function approximation is accurate, i.e., $|Q_{k-1}(s,a)-Q_{\pi_{k-1}}(s,a)|\leq\eta$ where $Q_{k-1}(s,a)=w_{k}^{\top}\phi(s,a)$. To overcome the lack of $\ell_{\infty}$ guarantee, we introduce the notion of virtual policy iteration algorithm. In the virtual algorithm, we start with the same initial policy $\widetilde{\pi}_{0}=\pi_{0}$. In the $k$-th iteration of the virtual algorithm, we assume that we have access to the true Q-function of the rollout policy $\widetilde{\pi}_{k-1}$ when $\phi(s,a)\notin\mathcal{H}$, and construct

where $\widetilde{w}_{k}$ is the linear coefficient that we learn in the virtual algorithm in the same way as in Eq. (4.1). Then $\widetilde{\pi}_{k}$ is chosen to be greedy with respect to $\widetilde{Q}_{k-1}(s,a)$. In this way, we can ensure that $\widetilde{Q}_{k-1}(s,a)$ is close to the true Q-function $Q_{\widetilde{\pi}_{k-1}}(s,a)$ in $\ell_{\infty}$ norm and thus the output policy, say $\widetilde{\pi}_{K-1}$, of the virtual algorithm is good in the sense that $V^{*}(\rho)-V_{\widetilde{\pi}_{K-1}}(\rho)$ is small.

To connect the output policy of the virtual algorithm and our actual algorithm, we note that by definition, in the final loop of our algorithm, in any iteration, for any state $s$ that the agent visits in ConfidentRollout, and any action $a\in\mathcal{A}$, we have that $\phi(s,a)\in\mathcal{H}$ since the subroutine never returns $\mathsf{uncertain}$ status. Further, because the initial state, probability transition kernel, and the policies are all deterministic, we know that the rollout trajectories of the virtual algorithm and our actual algorithm are always the same in the final loop (the virtual algorithm does not get a chance to use the true Q-function $Q_{\widetilde{\pi}_{k-1}}$). With rollout length $n$, we know that when we start with state $\rho$, the output of the virtual algorithm $\widetilde{\pi}_{K-1}$ and our actual algorithm $\pi_{K-1}$ take exactly the same actions for $n$ steps, and thus $|V_{\pi_{K-1}}(\rho)-V_{\widetilde{\pi}_{K-1}}(\rho)|\leq\frac{\gamma^{n+1}}{1-\gamma}$, which implies that $V^{*}(\rho)-V_{\pi_{K-1}}(\rho)$ is small. To extend this argument to the setting with stochastic transitions, we need to use a coupling argument which we elaborate in the Appendix.