SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Online RL in Linearly  $q^{\pi}$ -Realizable MDPs Is as Easy as in Linear MDPs If You Learn What to Ignore
Gellert Weisz
Google DeepMind, London, UK  
University College London, London, UK
András György
Google DeepMind, London, UK
Google DeepMind, Montreal, Canada  
University of Alberta, Edmonton, Canada
We consider online reinforcement learning (RL) in episodic Markov decision processes (MDPs) under the linear  $q^{\pi}$ -realizability assumption, where it is assumed that the action-values of all policies can be expressed as linear functions of state-action features. This class is known to be more general than linear MDPs, where the transition kernel and the reward function are assumed to be linear functions of the feature vectors. As our first contribution, we show that the difference between the two classes is the presence of states in linearly  $q^{\pi}$ -realizable MDPs where for any policy, all the actions have approximately equal values, and skipping over these states by following an arbitrarily fixed policy in those states transforms the problem to a linear MDP. Based on this observation, we derive a novel (computationally inefficient) learning algorithm for linearly  $q^{\pi}$ -realizable MDPs that simultaneously learns what states should be skipped over and runs another learning algorithm on the linear MDP hidden in the problem. The method returns an  $\varepsilon$ -optimal policy after polylog $(H,d)/\varepsilon^{2}$  interactions with the MDP, where  $H$  is the time horizon and  $d$  is the dimension of the feature vectors, giving the first polynomial-sample-complexity online RL algorithm for this setting. The results are proved for the misspecified case, where the sample complexity is shown to degrade gracefully with the misspecification error.
We consider reinforcement learning where an agent interacts in an online fashion with an environment modeled as a Markov decision process: The agent, observing a state, takes an action that results in a random next state and reward, the latter of which is to be maximized over time. To tackle large, possibly infinite state spaces, additional structure needs to be introduced to this problem. One such structure is a "feature-map" that maps state-action pairs to  $d$ -dimensional vectors (for some positive integer  $d$ ) with the intention that a "good feature-map extracts important information from the state-action pairs so that learning with this extra information becomes tractable. An example is the case of linear MDPs [Jin et al., 2020a], where the assumption is that both the transition and reward functions are linearly factorizable and their left factors are given by the feature-map. In contrast, value-based approaches, such as  $q^{\pi}$ -realizability [Du et al., 2019, Lattimore et al., 2020] aim to model only the action-values with the features. In this work, we focus on the latter, a strictly more general setting than that of linear MDPs [Zanette et al., 2020, Proposition 4].
There are several sample-efficient algorithms discovering near-optimal policies in linear MDPs under various MDP access models and settings (online access: Jin et al. [2020a]; batch setting: Jin et al.
MDP class	Online RL	Planning with simulator
poly(·) sample	poly(·) compute	poly(·) sample	poly(·) compute
Linear MDP	Jin et al. [2020a]
qπ-realizable MDP	This work	Open problem	Yin et al. [2022]
Table 1: Comparison of efficiency results for linear MDPs and  $q^{\pi}$ -realizable MDPs under online RL and planing with a simulator. This work establishes that  $q^{\pi}$ -realizable MDPs are also sample efficiently solvable under online RL. The computational complexity of this problem remains open.
[2021]; reward-free setting: Wagenmaker et al. [2022]). The best known sample-complexity bound for the online access model is achieved by the computationally inefficient algorithm of Zanette et al. [2020], called ELEANOR, which serves as a starting point of our work.
In this work we consider the setting of linearly  $q^{\pi}$ -realizable MDPs. As opposed to linear MDPs, before this work, sample efficient solutions were only known for this case when the MDP is accessed through a simulator that implements some form of a state-reset function [Lattimore et al., 2020, Yin et al., 2022, Weisz et al., 2022] (Table 1). In this work we resolve an open problem by Du et al. [2019], and show that having access to a state-reset is not essential in this setting. To this end, we present SKIPPYELEANOR (Algorithm 1) and a corresponding theorem (Theorem 4.1) that shows that SKIPPYELEANOR, which uses online interactions only, is a provably sample-efficient solution to this problem. The rest of this paper is organized as follows. In Section 2 we introduce the basic definitions. In Section 3 we give an insight into the difference between linear  $q^{\pi}$ -realizability and linear MDPs, which motivates our approach. In Section 4 we describe our algorithm and the most important technical tools we discovered for its analysis. Notably, in Section 4.2 we establish a rich structure inherent in  $q^{\pi}$ -realizable MDPs, which acts as the technical foundation to this work, and may be of independent interest. Finally, Section 5 gives a summary of the proof of our main result (Theorem 4.1), before concluding with some notes on future work in Section 6.
For a linear subspace  $X$  of  $\mathbb{R}^d$ , let  $\mathrm{Proj}_X$  denote the orthogonal projection matrix onto  $X$ . Throughout we fix  $d \in \mathbb{N}^+$ . For  $L > 0$ , let  $\mathcal{B}(L) = \{x \in \mathbb{R}^d : \| x \|_2 \leq L\}$  denote the  $d$ -dimensional Euclidean ball of radius  $L$  centered at the origin, where  $\| \cdot \|_2$  denotes the Euclidean norm. Let PD denote the set of positive definite matrices in  $\mathbb{R}^{d \times d}$ . We write  $a \approx_{\varepsilon} b$  for  $a, b, \varepsilon \in \mathbb{R}$  if  $|a - b| \leq \varepsilon$ . Let  $\mathbb{I}\{B\}$  be the indicator function of a boolean-valued (possibly random)  $B$  taking value 1 if  $B$  is true and 0 if false. Let  $\mathcal{M}_1(X)$  denote the set of probability distributions supported on set  $X$ . The rest of our notation is standard, but described in Appendix A for completeness.
For the setting of episodic finite horizon RL, with horizon  $H$ , a finite-action Markov decision process (MDP) describes an environment for sequential decision-making. It is defined by a tuple  $(S, [\mathcal{A}], P, \mathcal{R})$  as follows. The state space  $S$  is split across stages:  $S = (S_{t})_{t \in [H]}$  with  $S_{1} = \{s_{1}\}$  for some designated initial state  $s_{1}$ . Without loss of generality, we assume the  $(S_{t})_{t \in [H]}$  are disjoint sets. We define the function stage:  $S \to [H]$  as stage(s)  $= t$  if  $s \in S_{t}$ . We consider finite action spaces of size  $\mathcal{A}$  for some  $\mathcal{A} \in \mathbb{N}^{+}$ , and without loss of generality, define the set of actions to be  $[\mathcal{A}] := \{1, \dots, \mathcal{A}\}$ . The transition kernel is  $P: (\bigcup_{t \in [H-1]} S_{t}) \times [\mathcal{A}] \to \mathcal{M}_{1}(\mathcal{S})$ , with the property that transitions happen between successive stages, that is, for any  $t \in [H-1]$ , state  $s_{t} \in S_{t}$ , and action  $a \in [\mathcal{A}]$ ,  $P(s_{t}, a) \in \mathcal{M}_{1}(S_{t+1})$ . The reward kernel is  $\mathcal{R}: S \times [\mathcal{A}] \to \mathcal{M}_{1}([0,1])$ . An agent interacts sequentially with this environment in an episode lasting  $H$  steps by taking some action  $a \in [\mathcal{A}]$  in the current state. The environment responds by transitioning to some next-state according to  $P$ , and giving a reward in  $[0,1]$  according to  $\mathcal{R}$ .
We describe an agent interacting with the MDP by a policy  $\pi$ , which, to each history of interaction (including states, actions and rewards) assigns a probability distribution over the actions. Policies where this distribution only depend on the last state in the history are called memoryless, and these are identified with elements of the set  $\Pi = \{\pi : S \to \mathcal{M}_1([\mathcal{A}])\}$ . Using a policy  $\pi$ , starting at some state  $s$  in an MDP induces a probability distribution over histories, which we denote by  $\mathcal{P}_{\pi,s}$ . For
<sup>1</sup>Here, the reward and next-state are independent, given the current state and last action. Independence is nonessential and is assumed only to simplify the presentation.
any  $a \in [\mathcal{A}]$ ,  $\mathcal{P}_{\pi, s, a}$  is the distribution over the histories when first action  $a$  is used in state  $s$ , after which policy  $\pi$  is followed.  $\mathbb{E}_{\bullet}$  is the expectation operator corresponding to a distribution  $\mathcal{P}_{\bullet}$  (e.g.,  $\mathbb{E}_{\pi, s}$  is the expectation with respect to  $\mathcal{P}_{\pi, s}$ ). The state- and action-value functions  $\nu^{\pi}$  and  $q^{\pi}$  are defined as the expected total reward within the first episode while  $\pi$  is used:
$$
\nu^ {\pi} (s) = \underset {\pi , s} {\mathbb {E}} \sum_ {u = \operatorname {s t a g e} (s)} ^ {H} R _ {u} \quad \text {f o r} s \in \mathcal {S} \quad \text {a n d} \quad q ^ {\pi} (s, a) = \underset {\pi , s, a} {\mathbb {E}} \sum_ {u = \operatorname {s t a g e} (s)} ^ {H} R _ {u} \quad \text {f o r} s \in \mathcal {S}, a \in [ \mathcal {A} ].
$$
Let  $\pi^{\star} \in \Pi$  be an optimal policy, satisfying  $q^{\pi^{\star}}(s,a) = \sup_{\pi \in \Pi} q^{\pi}(s,a) = \sup_{\pi \in \text{all policies}} q^{\pi}(s,a)$  for all  $(s,a) \in S \times [\mathcal{A}]$ . Let  $q^{\star}(s,a) = q^{\pi^{\star}}(s,a)$  and  $\nu^{\star}(s) = \sup_{a' \in [\mathcal{A}]} q^{\star}(s,a)$  for all  $(s,a)$ .
3 From linear  $q^{\pi}$ -realizability to linear MDPs
As described in the introduction, we endow our MDP with a feature map  $\varphi : S \times [\mathcal{A}] \to \mathcal{B}(L_1)$  for some  $L_1 > 0$ . For reference, we start with a definition of linear MDPs with a parameter norm bound  $L_2 > 0$ , formalizing that the transition kernel and the expected rewards are approximately linear functions of the features:
Definition 3.1.  $[\kappa$ -approximately linear MDP] For any  $\kappa \leq 1$ , an MDP is a  $\kappa$ -approximately linear MDP if (i) there exists  $\theta_1, \ldots, \theta_H \in \mathcal{B}(L_2)$  such that for any  $h \in [H]$  and  $(s,a) \in S_h \times [\mathcal{A}]$ ,  $\left|\mathbb{E}_{R\sim \mathcal{R}(s,a)}R - \langle \varphi(s,a),\theta_h\rangle\right| \leq \kappa$  and (ii) for any  $f: S \to [0,H]$  and  $h \in [H-1]$ , there exists  $\theta_h' \in \mathcal{B}(L_2)$  such that for all  $(s,a) \in S_h \times [\mathcal{A}]$ ,  $\left|\mathbb{E}_{S'\sim P(s,a)}f(S') - \left\langle \varphi(s,a),\theta_h'\right\rangle\right| \leq \kappa$ .
A key consequence of the linear MDP assumption is that the inherent Bellman error
$$
\sup_{\theta_{h + 1}\in \mathcal{B}(L_{2})}\inf_{\theta_{h}\in \mathcal{B}(L_{2})}\sup_{(s,a)\in \mathcal{S}_{h}\times [\mathcal{A}]} \Big|\underset {R\sim \mathcal{R}(s,a),S^{\prime}\sim P(s,a)}{\mathbb{E}}R(s,a) + \max_{a^{\prime}\in [\mathcal{A}]} \langle \varphi (S^{\prime},a^{\prime}),\theta_{h + 1}\rangle -\langle \varphi (s,a),\theta_{h}\rangle \Big|,
$$
scales with the misspecification  $\kappa$ . This property is also referred to as the closedness to the Bellman operator, and is a crucial component in the analysis of approximation errors for algorithms tackling linear MDPs.
In this work we consider a weaker linearity assumption where we only assume that the action-value functions are approximately linear:
Definition 3.2 ( $q^{\pi}$ -realizability: uniform linear function approximation error of value-functions). Given an MDP, the uniform value-function approximation error (or misspecification) induced by a feature map  $\varphi : S \times [\mathcal{A}] \to \mathcal{B}(L_1)$ , over a set of parameters in  $\mathcal{B}(L_2)$  is
$$
\eta = \sup_{\pi \in \Pi}\max_{h\in [H]}\inf_{\theta^{(h)}\in \mathcal{B}(L_{2})}\sup_{(s,a)\in \mathcal{S}_{h}\times [\mathcal{A}]}  \left|q^{\pi}(s,a) - \left\langle \varphi (s,a),\theta^{(h)}\right\rangle \right|.
$$
For the MDP and the corresponding feature map, for all  $h \in [H]$  fix any  $\theta_h : \Pi \to \mathcal{B}(L_2)$  mapping each memoryless policy  $\pi \in \Pi$  to its "parameter", such that
$$
q ^ {\pi} (s, a) \approx_ {\eta} \langle \varphi (s, a), \theta_ {h} (\pi) \rangle \quad f o r a l l \pi \in \Pi , s \in \mathcal {S} _ {h}, a n d a \in [ \mathcal {A} ]. \tag {1}
$$
The set of all parameters  $\Theta_h \subseteq \mathcal{B}(L_2)$  for a stage  $h \in [H]$  is given by  $\Theta_h = \{\theta_h(\pi) : \pi \in \Pi\}$ .
Note that  $\theta_h$  satisfying Eq. (1) always exist [Weisz et al., 2022, Appendix C]. We focus on the feasible regime where  $\eta$  is polynomially small in the relevant parameters. Specifically, we assume that  $\eta$  is bounded according to Eq. (21). The main problem of interest in this work is the following:
Problem 3.3 (informal). For any  $\varepsilon, \zeta > 0$  and any MDP with corresponding uniform value-function approximation error  $\eta$ , derive an algorithm that, with probability at least  $1 - \zeta$ , will find an  $\varepsilon$ -optimal policy (i.e., a policy  $\pi$  such that  $\nu^{\pi}(s_1) \geq \nu^{\star}(s_1) - \varepsilon$ ) by interacting with the MDP online for  $T$  steps with  $T$  bounded by a polynomial function of  $(d, H, \varepsilon^{-1}, \log \zeta^{-1}, \log L_1, \log L_2)$ . That the interaction with the MDP is online means that it is only possible to observe the features corresponding to the current state, and to take an action and subsequently observe the resulting reward and next state, which then becomes the current state. We consider the fixed horizon episodic setting, that is, the next state is reset to the initial state  $s_1$  after every  $H$  steps.
2Compared to the definition of Jin et al. [2020b], our definition does not require the existence of a vectorvalued measure to represent the transition kernel. This is a generalization that is compatible with all existing algorithms for linear MDPs.
Figure 1: Left: MDP with deterministic transitions and rewards (edges are labeled with action/reward). Right: The same MDP with the red "low-range" states "skipped" over.  $\varphi(s_1, \cdot) = (1)$ ,  $\varphi(s_3, \cdot) = (0.5)$ ,  $\varphi(\cdot, \cdot) = (0)$  otherwise. Both MDPs are  $q^{\pi}$ -realizable, but only the right MDP is linear.
Algorithms developed for linear MDPs are not directly applicable to Problem 3.3 when the MDP is only  $q^{\pi}$ -realizable: While a linear MDP is also  $q^{\pi}$ -realizable, a  $q^{\pi}$ -realizable MDP may be neither a linear MDP, nor one with a low inherent Bellman error [Zanette et al., 2020]. As an illustrative example, Fig. 1, left shows an MDP that is  $q^{\pi}$ -realizable but not linear. To see this, observe that the features for both actions in  $s_1$  are identical, but their transitions and rewards are not. As illustrated in the figure however, if we skip over the red states (with identical actions) by taking the first action on them and summing up the rewards received until we reach a black state, we arrive at a linear MDP. This serves as the main intuition behind our work: the red states have no bearing on action-values, so they can be skipped, and the resulting MDP is linear.
More generally, we can define the range of any state as the maximum possible difference in action-value that the choice of action in that state can make:
$$
\operatorname {r a n g e} (s) = \sup  _ {\theta \in \Theta_ {\text {s t a g e} (s)}} \max  _ {i, j \in [ \mathcal {A} ]} \left\langle \varphi (s, i, j), \theta \right\rangle \text {f o r a l l} h \in [ H ], s \in S _ {h}, \tag {2}
$$
where  $\varphi(s, i, j) = \varphi(s, i) - \varphi(s, j)$  is the notation for feature differences. Clearly, the choice of action in low-range states is not too important, as
$$
v ^ {\pi} (s) - q ^ {\pi} (s, a) \leq \operatorname {r a n g e} (s) + 2 \eta \quad \text {f o r a n y} \pi \in \Pi \text {a n d a l l} a \in [ \mathcal {A} ]. \tag {3}
$$
Not only are the action choices in low-range states unimportant for the task of finding a near-optimal policy for the MDP, these choices can affect transitions and rewards in a nonlinear way. Interestingly, the existence of low-range states is the reason why  $q^{\pi}$ -realizable MDPs are not necessarily linear, as shown by the next result (proved in Appendix C), which follows easily from Lemma 4.7.
Proposition 3.4. Consider an MDP with uniform value-function approximation error  $\eta \geq 0$ . If there are no states  $s \in S$  with  $\mathrm{range}(s) < \alpha$  for some  $\alpha > 0$ , then the transitions and rewards of the MDP are linear (Definition 3.1) with misspecification scaling with  $\eta$ , and parameter norms scaling inversely with  $\alpha$ .
Our approach. The above result immediately offers a strategy to learn under the (linear)  $q^{\pi}$ -realizability assumption. Assuming access to an oracle that can determine whether or not range(s) <  $\alpha$  for any state  $s$ , the MDP could be "converted" to one that has no low-range states but has near-identical state and action-value functions of any policy (compared to the original MDP), by skipping over low-range states (by executing an arbitrary action) until a state with a range at least  $\alpha$  is reached. We will call such a multi-state transition a skipped step and refer to such a policy as a skipped policy. The reward presented for a skipped step is the cumulative reward over the skipped states. When the oracle is correct, the new MDP is a linear MDP, allowing techniques such as ELEANOR to efficiently learn a near-optimal policy. This conversion argument is part of the intuition of our method, but it is not strictly part of the proof, so we defer the details to Appendix C. The only missing piece for solving the general case, Problem 3.3, is learning an oracle that can suggest when to skip over a state, and combining it with the learning algorithm for the linear MDP. This general approach leads to our algorithm, SKIPPYELEANOR, which runs a modified version of ELEANOR with guessed oracles. During the algorithm, we detect when an incorrect oracle leads to suboptimal results, and refine the oracle accordingly. The details of the algorithm are explained in the next section.
In this section we present our main results following our plan outlined above. We first give Algorithm 1, along with a high-level overview of the algorithm; the details are explained throughout the section. The parameters of the algorithm are presented in Appendix B.
Algorithm 1 SKIPPYELEANOR
1: Input: accuracy  $\varepsilon >0$  , failure probability  $\zeta >0$    
2: Initialize  $m\gets 0,m^{\prime}\gets 0,Q_{h} = L_{2}I$  for  $h\in [H],\pi^0 = (s\mapsto 1)$    
3: while  $m^{\prime}\leq m_{\mathrm{max}}^{\prime}$  do   
4:  $m\gets m + 1,m^{\prime}\gets m^{\prime} + 1$ $\triangleright m^{\prime}$  also counts iterations repeated due to Line 14   
5: Estimate optimistic problem parameters  $\hat{G},\bar{\theta}$  by solving Optimization Problem 4.10   
6: for  $k\in [H]$  do   
7: Let  $\pi^{mk}$  be the policy defined by SKIPPYPOLICY(  $\hat{G},\bar{\theta},k)$    
8: Sample  $n$  trajectories by executing  $\pi^{mk}$  from  $S_{1}$  for  $n$  episodes   
9: Record data  $(S_h^{mk_j},A_h^{mk_j},R_h^{mk_j})_{h\in [H],j\in [n]}$  and stage-mapping functions  $(p^{mk_j})_{j\in [n]}$    
10: Solve Optimization Problem 4.12 with input  $(\hat{G},\bar{\theta})$  Consistency check record its value  $x$  (maximum discrepancy), and arguments  $\nu$  (direction) and  $i$  (stage).   
11: Calculate useful component  $w\gets \mathrm{Proj}_{Z(Q,i)}\nu$  Definition 4.2   
12: if  $x>$  discrepancy_threshold then   
13:  $Q_{i}\gets (Q_{i}^{-2} + Q_{i}^{-1}ww^{\top}Q_{i}^{-1})^{-\frac{1}{2}}$  append  $Q_{i}^{-1}w$  to  $C_i$  according to Eq. (4)   
14:  $m\gets m - 1$  redo this iteration   
15: continue   
16: if average uncertaintay  $\leq$  uncertainty_threshold then   
17: return policy  $\pi^{mH}$
Algorithm 2 SKIPPPOLICY
1: Input:  $\hat{G},\bar{\theta},k$    
2: Initialize  $S_{1}\gets s_{1},j\gets 1,\pi^{0}\gets (s\mapsto 1)$  , stage mapping  $p$    
3: for  $i = 1$  to  $H$  do   
4: Compute skip probabilities  $\tau_{i}\leftarrow \tau_{\hat{G}}(S_{i})$  and non-skip action  $a^+ \leftarrow \pi_{\hat{\theta}}^+(S_i)$  from Eq. (8)   
5: Sample independently  $B_{i}\sim$  Bernoulli  $(\tau_{i})$    
6: if  $B_{i} = 0$  then  $A_{i}\gets 1$  skip (follow  $\pi^0$  ) with probability  $1 - \tau_{i}$    
7: else   
8:  $p(j)\gets i,j\gets j + 1$    
9: if  $j\leq k$  then  $A_{i}\gets a^{+}$  (Phase I) else  $A_{i}\gets 1$  (Phase II)   
10: if  $i = H$  then   
11:  $p(j^{\prime}) = H + 1$  for  $j^{\prime} = j,\dots ,H$
For every stage  $h \in [H]$ , the algorithm keeps a progressively refined estimate of the geometry of the parameter space  $\Theta_h$ , by maintaining an ever shrinking ellipsoid enclosing  $\Theta_h$ . This ellipsoid is parametrized by an 'inverse covariance matrix'-like quantity  $Q_h$ , determined by  $\tilde{O}(d)$  vectors, which guarantees  $\max_{\theta_h \in \Theta_h} \| \theta_h \|_{Q_h^{-2}} = \tilde{O}(\sqrt{d})$ . Looking at the definition of range in Eq. (2), it is clear that the smaller the ellipsoid becomes, the better estimate we can give for the ranges.
Given some data collected so far and  $(Q_{h})_{h\in [H]}$ , SKIPPYELEANOR computes optimistic estimates of the action-values by calculating an optimistic policy parameter  $\bar{\theta}$ , as well as a guess  $\hat{G}$  to a near-optimal design which is used to estimate the range for the states (due to technical reasons,  $\hat{G}$  will guess a near-optimal design for the transformed parameter space  $Q_{h}^{-1}\Theta_{h}$ ).
Data is collected by running stochastic versions of skipped policies on the MDP, where the states to be skipped over are determined based on the range estimates; when a state is skipped, an action is selected using a deterministic policy  $\pi^0$  that always chooses the first action in every state. To ensure that the estimation problem is smooth in terms of  $\hat{G}$ , we use a smoothed version of skipped policies, where states are skipped randomly, and the probability of skipping is larger for states with lower ranges, while high-range states are never skipped. Similarly to ELEANOR, we aim to estimate the action-value function of a state-action pair by adding the estimated one-step reward to the estimated value-function of the next state. However, unlike ELEANOR, we would like to do this in the reduced MDP, where the low-range states that are skipped over are removed (and the corresponding transitions are replaced by skipped steps). Since we do not know these states in advance, we run exploratory policies that skip over next states starting from any state: namely,
we run SKIPPYPOLICY( $\hat{G}, \bar{\theta}, k$ ) for all  $k \in [H]$  with a maximum number of unskipped states  $k$  (Phase I), and once this is skip budget is exhausted, all remaining states are skipped over by rolling out  $\pi^0$  (Phase II), which ensures that we collect enough data at every stage of the MDP to be able to estimate the one-skippy-step reward of any skipping mechanism. Compared to ELEANOR, this introduces an additional loop in Line 6 of SKIPPYELEANOR; see Appendix D for additional details. For any execution, SKIPPYPOLICY maintains a stage-mapping function  $p$ , which, for any stage  $h$  of the trajectory in the reduced MDP gives the stage index in the original MDP. In other words,  $p(j)$  is the stage of the landing state of the  $j^{th}$  skippy step.