Improved No-Regret Algorithms for Stochastic Shortest Path with Linear MDP

Our results are built upon several technical innovations. First, as mentioned, we develop an improved analysis for the finite-horizon approximation of (Cohen et al., 2021), which might be of independent interest. The key idea is to directly bound the total approximation error with respect to the regret bound of the finite-horizon algorithm, instead of analyzing the estimation precision for each state-action pair as done in (Cohen et al., 2021).

Second, to obtain the logarithmic bound in Section 3, we note that it is not enough to simply combine the aforementioned finite-horizon approximation and the existing logarithmic regret results for the finite-horizon setting such as (He et al., 2021), since the sub-optimality gap obtained in this way is in terms of the finite-horizon counterpart instead of the original SSP and could be substantially smaller. We resolve this issue via a longer horizon in the approximation and a careful two-stage analysis.

Finally, our horizon-free result in Section 4 is obtained by a novel combination of several ideas, including the global optimization algorithm of (Zanette et al., 2020b; Wei et al., 2021b), the variance-aware confidence sets of (Zhang et al., 2021) (for a related but different setting with linear mixture MDPs), an improved analysis of the variance-aware confidence sets (Kim et al., 2021), and finally a new clipping trick and new update conditions that we propose. Our analysis does not require the recursion-based technique of (Zhang et al., 2020a) (for the tabular case), nor estimating higher order moments of value functions as in (Zhang et al., 2021) (for linear mixture MDPs), which might also be of independent interest.

Regret minimization of SSP under stochastic costs has been well studied in the tabular setting (that is, no function approximation) (Tarbouriech et al., 2020; Cohen et al., 2020, 2021; Tarbouriech et al., 2021; Chen et al., 2021a; Jafarnia-Jahromi et al., 2021). There are also several works (Rosenberg & Mansour, 2020; Chen et al., 2021b; Chen & Luo, 2021) considering the more challenging setting with adversarial costs (which is beyond the scope of this work).

Beyond linear function approximation, in the finite-horizon setting researchers also start considering theoretical guarantees for general function approximation (Wang et al., 2020; Ishfaq et al., 2021; Kong et al., 2021). The study for SSP, which again is a strict generalization of the finite-horizon problems and might be a better model for many applications, falls behind in this regard, motivating us to explore in this direction with the goal of providing a more complete picture at least for linear function approximation.

The use of variance information is crucial in obtaining optimal regret bounds in MDPs. This dates back to the work of (Lattimore & Hutter, 2012) for the discounted setting, which has been significantly extended to the finite-horizon setting (Azar et al., 2017; Jin et al., 2018; Zanette & Brunskill, 2019; Zhang et al., 2020a, b). Constructing variance-aware confidence sets for linear bandits and linear mixture MDPs has also gained recent attention (Zhou et al., 2021a; Zhang et al., 2021; Kim et al., 2021). We are among the first to do so for linear MDPs (a concurrent work (Wei et al., 2021a) also does so but for a completely different purpose of improving robustness against corruption).

Logarithmic gap-dependent bounds have been shown in different settings; see for example (Jaksch et al., 2010; Simchowitz & Jamieson, 2019; Jin et al., 2021; He et al., 2021), but to our knowledge, we are the first to show similar bounds for SSP.

The learner’s goal is to learn a policy that reaches the goal state with minimum costs. Formally, a (stationary and deterministic) policy $\pi:{\mathcal{S}}\rightarrow{\mathcal{A}}$ is a mapping that assigns an action $\pi(s)$ to each state $s\in{\mathcal{S}}$. We say $\pi$ is proper if following $\pi$ (that is, taking action $\pi(s)$ whenever in state $s$) reaches the goal state with probability $1$. Given a proper policy $\pi$, we define its value function and action-value function as follows:

where the expectation in $V^{\pi}$ is with respect to the randomness of next states $s_{i+1}\sim P_{s_{i},\pi(s_{i})}$ and the number of steps $I$ before reaching the goal $g$. Let $\Pi$ be the set of proper policies. We make the basic assumption that $\Pi$ is non-empty. Under this assumption, there exists an optimal proper policy $\pi^{\star}$, such that $V^{\pi^{\star}}(s)=\min_{a}Q^{\pi^{\star}}(s,a)$, and $V^{\pi^{\star}}(s)=\min_{\pi\in\Pi}V^{\pi}(s)$ for all $s$ (Bertsekas & Yu, 2013). We use $V^{\star}$ and $Q^{\star}$ as shorthands for $V^{\pi^{\star}}$ and $Q^{\pi^{\star}}$. The formal goal of the learner is then to minimize her regret against $\pi^{\star}$, that is, the difference between her total costs and that of the optimal proper policy, defined as

We also define $R_{K}=\infty$ if $T=\infty$.

In the so-called tabular setting, the state space is assumed to be small, and algorithms with computational complexity and regret bound depending on $S=|{\mathcal{S}}|$ are acceptable. To handle a potentially massive state space, however, we consider the same linear function approximation setting of (Vial et al., 2021), where the MDP enjoys a linear structure in both the transition and cost functions (known as linear or low-rank MDP).

We refer the reader to (Vial et al., 2021) and references therein for justification on this widely-used structural assumption (especially on the last few norm constraints). Under Assumption 1, by definition we have $Q^{\star}(s,a)=\phi(s,a)^{\top}w^{\star}$, where $w^{\star}=\theta^{\star}+\int V^{\star}(s^{\prime})d\mu(s^{\prime})\in\mathbb{R}^{d}$, that is, $Q^{\star}$ is also linear in the features.

Two extra parameters that play a key role in our analysis are: $B_{\star}=\max_{s}V^{\star}(s)$, the maximum cost of the optimal policy starting from any state, and $T_{\star}=\max_{s}T^{\pi^{\star}}(s)$, the maximum hitting time of the optimal policy starting from any state, where $T^{\pi}(s)$ is the expected number of steps before reaching the goal if one follows policy $\pi$ starting from state $s$. By definition, we have $T_{\star}\leq B_{\star}/c_{\min}$.

For simplicity, we assume that $B_{\star}$, $T_{\star}$, and $c_{\min}$ are known to the learner for most discussions, and defer to the appendix what we can achieve when some of these parameters are unknown. We also assume $B_{\star}>1$ and $c_{\min}>0$ by default (and will discuss the case $c_{\min}=0$ for specific algorithms if modifications are needed).

For $n\in\mathbb{N}_{+}$, we define $[n]=\{1,\ldots,n\}$. For any $l\leq r$, we define $[x]_{[l,r]}=\min\{\max\{x,l\},r\}$ as the projection of $x$ onto the interval $[l,r]$. The notation $\tilde{\mathcal{O}}\left(\cdot\right)$ hides all logarithmic terms including $\ln K$ and $\ln\frac{1}{\delta}$ for some confidence level $\delta\in(0,1)$.

For a parameter $B>0$ and a weight vector $w\in\mathbb{R}^{d}$, inspired by (Zhang et al., 2021) we define a variance-aware confidence set for time step $t$ as

where ${\mathcal{J}}_{B}=\{\lceil\log_{2}\epsilon\rceil,\ldots,\lceil\log_{2}(6\sqrt{d}B)\rceil\}$ with $\epsilon=\frac{c_{\min}}{150d^{3}K}$, and $\Omega^{j}_{t}(w,B)=\mathbb{B}(3\sqrt{d}B)\cap$

with $\mathbb{B}(r)=\{x\in\mathbb{R}^{d}:\left\|{x}\right\|_{2}\leq r\}$ being the $d$-dimensional $L_{2}$-ball of radius $r$, $\mathcal{G}_{\xi}(r)=\{\xi n,n\in\mathbb{Z}\}^{d}\cap\mathbb{B}(r)$ being the $\xi$-net of $\mathbb{B}(r)$, $\textsf{clip}_{j}(x)=[x]_{[-2^{j},2^{j}]}$ (recall $[x]_{[l,r]}=\min\{\max\{x,l\},r\}$), $\phi_{i}$ being a shorthand of $\phi(s_{i},a_{i})$, $V_{w,B}(s)=\min_{a}[\phi(s,a)^{\top}w]_{[0,2B]}$ (and $V_{w,B}(g)=0$), $\epsilon^{i}_{V}(w^{\prime})=\phi_{i}^{\top}w^{\prime}-c_{i}-V(s^{\prime}_{i})$, $\eta^{i}_{V}(w^{\prime})=\epsilon^{i}_{V}(w^{\prime})^{2}$, and finally $\iota_{B,t}=2^{11}d\ln\frac{48dBt}{\epsilon\delta}$ for some failure probability $\delta$. The key difference between our confidence set and that of (Zhang et al., 2021) is in the definition of $\epsilon^{i}_{V}(w^{\prime})$ and $\eta^{i}_{V}(w^{\prime})$ due to the different structures between linear MDPs and linear mixture MDPs. In particular, we note that the value function $V$ (more formally $V_{w,B}$) in our definitions is itself defined with respect to another weight vector $w$.

With this confidence set, when VA-GOPO decides to update $w_{t}$, it searches over all $w$ such that $w\in\Omega_{t}(w,B_{t})$ and finds the one that minimizes the value at the current state $V_{w,B_{t}}(s_{t})$ (Line 3). Here, $B_{t}$ is a running estimate of $B_{\star}$. VA-GOPO maintains the inequality $V_{t}(s_{t})\leq B_{t}$ during the update by doubling the value of $B_{t}$ and repeating Line 3 whenever this is violated (Line 3). Note that the constraint $w\in\Omega_{t}(w,B_{t})$ is in a sense self-referential — we consider $w$ within a confidence set defined in terms of $w$ itself, which is an important distinction compared to (Zhang et al., 2021) and is critical for linear MDPs.

To provide some intuition on our confidence set, denote $V_{w_{t},B_{t}}$ by $V_{t}$ and $\Omega_{t}(w_{t},B_{t})$ by $\Omega_{t}$. Note that if we ignore the dependency between $V_{t}$ and $\{\phi_{i}\}_{i<t}$ (an issue that will eventually be addressed by some covering arguments), then $\{\epsilon^{i}_{V_{t}}(w^{\prime})\}_{i<t}$ forms a martingale sequence when $w^{\prime}=\widetilde{w}_{t}\triangleq\theta^{\star}+\int V_{t}(s^{\prime})d\mu(s^{\prime})$, and thus the inequality in Eq. (3) holds with high probability by some Bernstein-style concentration inequality (Lemma 36). Formally, this allows us to show the following.

Since $w_{t}$ is also in $\Omega_{t}$, the difference between $\phi(s,a)^{\top}w_{t}$ and $c(s,a)+\mathbb{E}_{s^{\prime}\sim P_{s,a}}[V_{t}(s^{\prime})]=\phi(s,a)^{\top}\widetilde{w}_{t}$ is controlled by the size of the confidence set $\Omega_{t}$, which is overall shrinking and thus making sure that $w_{t}$ is getting closer and closer to $w^{\star}$. In addition, we also show that $V_{t}$ is optimistic at state $s_{t}$ whenever an update is performed and that $B_{t}$ never overestimates $B_{\star}$ significantly.

VA-GOPO updates $w_{t}$ whenever one of the three conditions in Line 3 is triggered. The first condition $s^{\prime}_{t-1}=g$ simply indicates that the current time step is the start of a new episode. The second condition is

where $t^{\prime}$ is the most recent update time step (Line 3) and $\Phi^{j}_{t}(\nu)=\sum_{i<t}f_{j}(\phi_{i}^{\top}\nu)+2^{j}\left\|{\nu}\right\|_{2}^{2}$ with $f_{j}(x)=\textsf{clip}_{j}(x)x$. This lazy update condition makes sure that the algorithm does not update $w_{t}$ too often (see Lemma 27) while still enjoying a small enough estimation error. The last condition $V_{t^{\prime}}(s_{t})=2B_{t}$ (we call it overestimate condition) tests whether the current state has an overestimated value (note that $2B_{t}$ is the maximum value of $V_{t^{\prime}}$ due to the truncation in its definition). This condition helps remove a factor of $d^{1/4}$ in the regret bound without using some complicated ideas as in previous works; see Appendix C.5 for more explanation.