SlowGuess/NeurIPS_2023_OCR huggingface dataset:
In this section, we present a generic guarantee for PSDP (Algorithm 3). We show that given any reward functions  $r_{1:h} : \mathcal{X} \times \mathcal{A} \to \mathbb{R}$  and a function class  $\mathcal{G} \subseteq \{g : \mathcal{X} \times \mathcal{A} \to \mathbb{R}\}$  that "realizes" these reward functions (we formalize this in the next definition), if  $\Psi^{(1:h)}$  are  $\alpha$ -policy covers for layers 1 through  $h$ , then for sufficiently large  $n \geq 1$  and with high probability, the output  $\hat{\pi} = \mathrm{PSDP}(h, r_{1:h}, \mathcal{G}, \Psi^{(1:h)}, n)$  is an approximate maximizer of the objective
$$
\max  _ {\pi \in \Pi_ {M}} \mathbb {E} ^ {\pi} \left[ \sum_ {t = 1} ^ {h} r _ {t} \left(\boldsymbol {x} _ {t}, \boldsymbol {a} _ {t}\right) \right].
$$
To formalize this result, we define the notion of realizability we require for the function class  $\mathcal{G}$ .
Definition D.1. We say that the function class  $\mathcal{G} \subseteq \{g : \mathcal{X} \times \mathcal{A} \to \mathbb{R}\}$  realizes the reward functions  $r_{1:h} : \mathcal{X} \times \mathcal{A} \to \mathbb{R}$  if for all  $t \in [h]$  and all  $\pi \in \Pi_{\mathbb{M}}^{t:h}$ ,
$$
Q _ {t} ^ {\pi} \in \mathcal {G}, \quad \text {w h e r e} \quad Q _ {t} ^ {\pi} (x, a) := r _ {t} (x, a) + \mathbb {E} ^ {\pi} \left[ \sum_ {\ell = t + 1} ^ {h} r _ {\ell} \left(\boldsymbol {x} _ {\ell}, \boldsymbol {a} _ {\ell}\right) \mid \boldsymbol {x} _ {t} = x, \boldsymbol {a} _ {t} = a \right]. \tag {14}
$$
Note that  $Q_{t}^{\pi}$  in (14) represents the state-action value function ( $Q$ -function) at layer  $t \in [h]$  with respect to the rewards  $r_{1:h}$  and partial policy  $\pi$ .
In what follows, given a function class  $\mathcal{G} \subseteq \{g : \mathcal{X} \times \mathcal{A} \to \mathbb{R}\}$ , we use  $\mathcal{N}_{\mathcal{G}}(\varepsilon)$  to denote the  $\varepsilon$ -covering number of  $\mathcal{G}$  in  $\ell_{\infty}$  distance. With this, we now state a guarantee for PSDP.
Theorem D.2. Let  $\varepsilon, \delta \in (0,1)$ ,  $B > 0$ , and  $h \in [H]$ . Suppose reward functions  $r_{1:h}: \mathcal{X} \times \mathcal{A} \to \mathbb{R}$ , function class  $\mathcal{G}$ , a collection of policies  $\Psi^{(1:h)}$ , and a parameter  $n \geq 1$  satisfy the following:
- The function class  $\mathcal{G}$  realizes the reward functions  $r_{1:h}$  (in the sense of Definition D.1), functions in  $\mathcal{G}$  are uniformly bounded by  $B$ , and  $\lim_{n\to \infty}n^{-1}\cdot \log \mathcal{N}_{\mathcal{G}}(1 / n) = 0$
- For some  $0 < \alpha \leq 1$ , for each  $1 \leq t \leq h$ , it holds that  $\Psi^{(t)}$  is an  $\alpha$ -policy cover for layer  $t$  and moreover  $|\Psi^{(t)}| \leq d$ .
- The parameter  $n$  is chosen such that  $cdH\alpha^{-1} \cdot \varepsilon_{\mathrm{stat}}(n, \delta / H) \leq \varepsilon$ , where  $\varepsilon_{\mathrm{stat}}(n, \delta') \coloneqq \sqrt{B^2n^{-1} \cdot (\log\mathcal{N}_{\mathcal{G}}(1 / n) + \log(n / \delta))}$  and  $c > 0$  is a large enough absolute constant.
Then, with probability at least  $1 - \delta$ , the policy  $\hat{\pi} = \mathsf{PSDP}\big(h,r_{1:h},\mathcal{G},\Psi^{(1:h)},n\big)$  coming from Algorithm 3, satisfies the following guarantee:
$$
\operatorname * {m a x} _ {\pi \in \Pi_ {\mathfrak {M}}} \mathbb {E} ^ {\pi} \left[ \sum_ {t = 1} ^ {h} r _ {t} (\pmb {x} _ {t}, \pmb {a} _ {t}) \right] \leq \mathbb {E} ^ {\hat {\pi}} \left[ \sum_ {t = 1} ^ {h} r _ {t} (\pmb {x} _ {t}, \pmb {a} _ {t}) \right] + \varepsilon .
$$
To prove the theorem, we need two intermediate results. The first shows that the  $Q$  function is the Bayes-optimal predictor of the sum of rewards when rolling out with policy  $\pi$ .
Lemma D.1. Let  $t \in [H]$ ,  $r_{1:h} : \mathcal{X} \times \mathcal{A} \to \mathbb{R}$  be reward functions,  $\pi \in \Pi_{\mathbb{M}}$ , and let  $g_{\mathrm{bayes}}^{\pi}$  denote the Bayes-optimal predictor for the sum of rewards under policy  $\pi$ , i.e.,
$$
g _ {\text {b a y e s}} ^ {\pi} \in \underset {g: \mathcal {X} _ {t} \times \mathcal {A} \rightarrow \mathbb {R}} {\arg \min } \mathbb {E} ^ {\pi} \left[\left(g \left(\boldsymbol {x} _ {t}, \boldsymbol {a} _ {t}\right) - \sum_ {\ell = t} ^ {h} r _ {\ell} \left(\boldsymbol {x} _ {\ell}, \boldsymbol {a} _ {\ell}\right)\right) ^ {2} \right], \tag {15}
$$
Then,  $g_{\mathrm{bayes}}^{\pi} = Q_t^{\pi}$ , where  $Q_{t}^{\pi}$  is the  $Q$ -function at layer  $t \in [h]$  with respect to the policy  $\pi$  and rewards  $r_{1:h}$  defined in (14).
Proof of Lemma D.1. The least-squares solution  $g_{\mathrm{bayes}}^{\pi}$  of the problem in (15) satisfies, for all  $a \in \mathcal{A}$  and  $x \in \mathcal{X}_t$ ,
$$
\begin{array}{l} g _ {\mathrm {b a y e s}} ^ {\pi} (x, a) = \mathbb {E} ^ {\pi} \left[ \sum_ {\ell = t} ^ {h} r _ {\ell} (\boldsymbol {x} _ {\ell}, \boldsymbol {a} _ {\ell}) \Bigg | \boldsymbol {x} _ {t} = x, \boldsymbol {a} _ {t} = a \right], \\ = \mathbb {E} \left[ r _ {t} (\boldsymbol {x} _ {t}, \boldsymbol {a} _ {t}) \mid \boldsymbol {x} _ {t} = x, \boldsymbol {a} _ {t} = a \right] + \mathbb {E} ^ {\pi} \left[ \sum_ {\ell = t + 1} ^ {h} r _ {\ell} (\boldsymbol {x} _ {\ell}, \boldsymbol {a} _ {\ell}) \Bigg | \boldsymbol {x} _ {t} = x, \boldsymbol {a} _ {t} = a \right], \\ = r _ {t} (x, a) + \mathbb {E} ^ {\pi} \left[ \sum_ {\ell = t + 1} ^ {h} r _ {\ell} (\boldsymbol {x} _ {\ell}, \boldsymbol {a} _ {\ell}) \Bigg | \boldsymbol {x} _ {t} = x, \boldsymbol {a} _ {t} = a \right], \\ = Q _ {t} ^ {\pi} (s, a), \\ \end{array}
$$
where the last step follows by the definition of the  $Q$ -function in (14).
We now show that the solution  $\hat{g}^{(t)}$  of the least-squares problem in (6) of Algorithm 3 is close to the  $Q$ -function in the appropriate sense.
Lemma D.2. Let  $\varepsilon, \delta \in (0,1)$ ,  $B > 0$ , and  $1 \leq t \leq h \leq H$ . Further, let  $(\varepsilon_{\mathrm{stat}}, r_{1:h}, \mathcal{G}, \Psi^{(1:h)}, n)$  be as in Theorem D.2. Then, the solution  $\hat{g}^{(t)}$  of the least-squares problem in (6) in Algorithm 3 when calling  $\mathsf{PSDP}(h, r_{1:h}, \mathcal{G}, \Psi^{(1:h)}, n)$  satisfies with probability at least  $1 - \delta$ ,
$$
\mathbb {E} ^ {\text {u n i f} (\Psi^ {(t)})} \left[ \max  _ {a \in \mathcal {A}} \left(\hat {g} ^ {(t)} \left(\boldsymbol {x} _ {t}, a\right) - Q _ {t} ^ {\tilde {\pi} ^ {(t + 1)}} \left(\boldsymbol {x} _ {t}, a\right)\right) ^ {2} \right] \leq c ^ {2} \cdot \varepsilon_ {\text {s t a t}} ^ {2} (n, \delta),
$$
where  $\hat{\pi}^{(t + 1)}\in \Pi_{\mathbb{M}}^{t + 1:h}$  is as in Algorithm 3, and  $c > 0$  is an absolute constant.
Proof of Lemma D.2. Fix  $t \in [h]$  and let  $g_{\mathrm{bayes}}^{(t)} = g_{\mathrm{bayes}}^{\pi}$  be as in Lemma D.1 with
$$
\pi = \operatorname {u n i f} \left(\Psi^ {(t)}\right) \circ_ {t} \pi_ {\text {u n i f}} \circ_ {t + 1} \hat {\pi} ^ {(t + 1)},
$$
and reward functions  $r_{1:h}$  as in the lemma's statement. By Lemma D.1,  $g_{\mathrm{bayes}}^{(t)}$  is the Bayes-optimal solution of the least-squares problem in (6) of Algorithm 3. Thus, since  $\mathcal{G}$  realizes the reward functions  $r_{1:h}$ , a standard uniform-convergence guarantee for least-square regression (see e.g. Mhammedi et al. [29, Proposition B.1] with  $e = 0$  almost surely) implies that there exists an absolute constant  $c > 0$  (independent of  $t, h$ , and any other problem parameters) such that with probability at least  $1 - \delta$ ,
$$
\mathbb {E} ^ {\mathrm {u n i f} (\Psi^ {(t)})} \left[ \max _ {a \in \mathcal {A}} \left(\hat {g} ^ {(t)} (\boldsymbol {x} _ {t}, a) - g _ {\mathrm {b a y e s}} ^ {(t)} (\boldsymbol {x} _ {t}, a)\right) ^ {2} \right] \leq c ^ {2} B ^ {2} \cdot \frac {\log \mathcal {N} _ {\mathcal {G}} (1 / n) + \log (n / \delta)}{n}.
$$
The desired result follows by the fact that  $g_{\mathrm{bayes}}^{(t)} \equiv Q_t^{\hat{\pi}^{(t + 1)}}$ ; see Lemma D.1.
We also require the classical performance difference lemma from Kakade [25].
<sup>9</sup>Observe that because the loss is strongly convex, this predictor is unique up to sets of measure zero, justifying our usage of "the" Bayes optimal reward.
Lemma D.3 (Performance Difference Lemma). Let  $\pi_{\star}, \pi \in \Pi_{\mathsf{M}}$  be arbitrary, and  $Q_{t}^{\pi}$  be as defined in (14). Then, for any  $h \geq 1$
$$
\mathbb {E} ^ {\pi_ {\star}} \left[ \sum_ {t = 1} ^ {h} r _ {t} (\pmb {x} _ {t}, \pmb {a} _ {t}) \right] - \mathbb {E} ^ {\pi} \left[ \sum_ {t = 1} ^ {h} r _ {t} (\pmb {x} _ {t}, \pmb {a} _ {t}) \right] = \sum_ {t = 1} ^ {h} \mathbb {E} ^ {\pi_ {\star}} \left[ Q _ {t} ^ {\pi} (\pmb {x} _ {t}, \pi_ {\star} (\pmb {x} _ {t})) - Q _ {t} ^ {\pi} (\pmb {x} _ {t}, \pi (\pmb {x} _ {t})) \right].
$$
Using these results, we now prove Theorem D.2.
Proof of Theorem D.2. We first show that for any  $t \in [h]$ , there is an event  $\mathcal{E}_t$  of probability at least  $1 - \delta / H$  under which the learned partial policies  $\hat{\pi}^{(t)}, \hat{\pi}^{(t + 1)}$  are such that
$$
\mathbb {E} ^ {\pi_ {\star}} \left[ Q _ {t} ^ {\hat {\pi} ^ {(t + 1)}} \left(\boldsymbol {x} _ {t}, \pi_ {\star} \left(\boldsymbol {x} _ {t}\right)\right) - Q _ {t} ^ {\hat {\pi} ^ {(t + 1)}} \left(\boldsymbol {x} _ {t}, \hat {\pi} ^ {(t)} \left(\boldsymbol {x} _ {t}\right)\right) \right] \leq \frac {\varepsilon}{H}, \tag {16}
$$
where  $\pi_{\star} \in \arg \max_{\pi \in \Pi_{\mathbb{M}}} \mathbb{E}^{\pi}\left[\sum_{t=1}^{h} r_{t}(\boldsymbol{x}_{t}, \boldsymbol{a}_{t})\right]$  is the optimal policy and  $Q_{t}^{\pi}$  is the  $Q$ -function defined in (14). Once we establish (16) for all  $t \in [h]$ , we will apply the performance difference lemma (Lemma D.3) and the union bound to obtain the desired result.
For any  $t\in [h]$  , we have
$$
\begin{array}{l} \mathbb {E} ^ {\pi_ {\star}} \left[ Q _ {t} ^ {\hat {\pi} ^ {(t + 1)}} \left(\boldsymbol {x} _ {t}, \pi_ {\star} \left(\boldsymbol {x} _ {t}\right)\right) - Q _ {t} ^ {\hat {\pi} ^ {(t + 1)}} \left(\boldsymbol {x} _ {t}, \hat {\pi} ^ {(t)} \left(\boldsymbol {x} _ {t}\right)\right) \right] \\ = \mathbb {E} ^ {\pi_ {\star}} \left[ Q _ {t} ^ {\hat {\pi} ^ {(t + 1)}} \left(\boldsymbol {x} _ {t}, \pi_ {\star} \left(\boldsymbol {x} _ {t}\right)\right) - \hat {g} ^ {(t)} \left(\boldsymbol {x} _ {t}, \pi_ {\star} \left(\boldsymbol {x} _ {t}\right)\right) + \hat {g} ^ {(t)} \left(\boldsymbol {x} _ {t}, \pi_ {\star} \left(\boldsymbol {x} _ {t}\right)\right) - Q _ {t} ^ {\hat {\pi} ^ {(t + 1)}} \left(\boldsymbol {x} _ {t}, \hat {\pi} ^ {(t)} \left(\boldsymbol {x} _ {t}\right)\right) \right], \\ \leq \mathbb {E} ^ {\pi_ {\star}} \left[ Q _ {t} ^ {\hat {\pi} ^ {(t + 1)}} (\boldsymbol {x} _ {t}, \pi_ {\star} (\boldsymbol {x} _ {t})) - \hat {g} ^ {(t)} (\boldsymbol {x} _ {t}, \pi_ {\star} (\boldsymbol {x} _ {t})) + \hat {g} ^ {(t)} (\boldsymbol {x} _ {t}, \hat {\pi} ^ {(t)} (\boldsymbol {x} _ {t})) - Q _ {t} ^ {\hat {\pi} ^ {(t + 1)}} (\boldsymbol {x} _ {t}, \hat {\pi} ^ {(t)} (\boldsymbol {x} _ {t})) \right], \\ \leq 2 \cdot \mathbb {E} ^ {\pi_ {*}} \left[ \max  _ {a \in \mathcal {A}} \left| Q _ {t} ^ {\hat {\pi} ^ {(t + 1)}} \left(\boldsymbol {x} _ {t}, a\right) - \hat {g} ^ {(t)} \left(\boldsymbol {x} _ {t}, a\right) \right| \right], \tag {17} \\ \end{array}
$$
where the penultimate inequality follows by the fact that  $\hat{\pi}^{(t)}(x)\in \arg \max_{a\in \mathcal{A}}\hat{g}^{(t)}(x,a)$ , for all  $(x,a)\in \mathcal{X}_t\times \mathcal{A}$  by its construction in (8). Now, using the assumption on  $\Psi^{(t)}$ , we have for any  $t\in [t]$
$$
\begin{array}{l} \mathbb {E} ^ {\pi_ {\star}} \left[ Q _ {t} ^ {\hat {\pi} ^ {(t + 1)}} \left(\boldsymbol {x} _ {t}, \pi_ {\star} (\boldsymbol {x} _ {t})\right) - Q _ {t} ^ {\hat {\pi} ^ {(t + 1)}} \left(\boldsymbol {x} _ {t}, \hat {\pi} ^ {(t)} (\boldsymbol {x} _ {t})\right) \right] \\ \leq 2 \int_ {\mathcal {X} _ {t}} \left(\max  _ {a \in \mathcal {A}} \left| Q _ {t} ^ {\hat {\pi} ^ {(t + 1)}} (x, a) - \hat {g} ^ {(t)} (x, a) \right|\right) d ^ {\pi_ {*}} (x) \mathrm {d} \nu (x), \\ \leq 2 \alpha^ {- 1} \int_ {\mathcal {X} _ {t}} \left(\max  _ {a \in \mathcal {A}} \left| Q _ {t} ^ {\hat {\pi} ^ {(t + 1)}} (x, a) - \hat {g} ^ {(t)} (x, a) \right|\right) \max  _ {\pi \in \Psi^ {(t)}} d ^ {\pi} (x) \mathrm {d} \nu (x), \\ \leq 2 \alpha^ {- 1} d \int_ {\mathcal {X} _ {t}} \left(\max  _ {a \in \mathcal {A}} \left| Q _ {t} ^ {\pi^ {(t + 1)}} (x, a) - \hat {g} ^ {(t)} (x, a) \right|\right) d ^ {\text {u n i f} (\Psi^ {(t)})} (x) d \nu (x), \quad (| \Psi^ {(t)} | \leq d) \\ = 2 \alpha^ {- 1} d \cdot \mathbb {E} ^ {\text {u n i f} (\Psi^ {(t)})} \left[ \max  _ {a \in A} \left| Q _ {t} ^ {\hat {\pi} ^ {(t + 1)}} \left(\boldsymbol {x} _ {t}, a\right) - \hat {g} ^ {(t)} \left(\boldsymbol {x} _ {t}, a\right) \right| \right], \tag {18} \\ \end{array}
$$
where the first inequality is by (17), the second inequality follows from the fact that  $\Psi^{(t)}$  is an  $\alpha$ -policy cover, the third inequality follows from the fact that  $|\Psi^{(t)}| \leq d$ , and the equality is by definition. Now, by Lemma D.2, we have that for any  $t \in [h]$ , there is an absolute constant  $c > 0$  (independent of  $t$  and other problem parameters) and an event  $\mathcal{E}_t$  of probability at least  $1 - \delta / H$  under which the solution  $\hat{g}^{(t)}$  of the least-squares regression problem on (6) of Algorithm 3 satisfies,
$$
\mathbb {E} ^ {\text {u n i f} (\Psi^ {(t)})} \left[ \max  _ {a \in \mathcal {A}} \left| \hat {g} ^ {(t)} \left(\boldsymbol {x} _ {t}, a\right) - Q _ {t} ^ {\hat {\pi} ^ {(t + 1)}} \left(\boldsymbol {x} _ {t}, a\right) \right| \right] \leq c \cdot \varepsilon_ {\text {s t a t}} (n, \frac {\delta}{H}) \leq \frac {\alpha \varepsilon}{2 d H}, \tag {19}
$$
where the last inequality follows by the choice of  $n$  in the theorem's statement. Combining (19) with (18) establishes (16) under the event  $\mathcal{E}_t$ .
Now, by the performance difference lemma (Lemma D.3), we have
$$
\begin{array}{l} \mathbb {E} ^ {\pi \star} \left[ \sum_ {t = 1} ^ {h} r _ {t} \left(\boldsymbol {x} _ {t}, \boldsymbol {a} _ {t}\right) \right] - \mathbb {E} ^ {\hat {\pi}} \left[ \sum_ {t = 1} ^ {h} r _ {t} \left(\boldsymbol {x} _ {t}, \boldsymbol {a} _ {t}\right) \right] \\ = \sum_ {t = 1} ^ {h} \mathbb {E} ^ {\pi_ {\star}} \left[ Q _ {t} ^ {\hat {\pi} ^ {(t + 1)}} (\boldsymbol {x} _ {t}, \pi_ {\star} (\boldsymbol {x} _ {t})) - Q _ {t} ^ {\hat {\pi} ^ {(t + 1)}} (\boldsymbol {x} _ {t}, \hat {\pi} ^ {(t)} (\boldsymbol {x} _ {t})) \right]. \\ \end{array}
$$
Algorithm 6 RepLearn  $(h,\mathcal{F},\Phi ,P,n)$  : Representation Learning for Low-Rank MDPs [35]
Require:
- Target layer  $h \in [H]$ .
- Discriminator class  $\mathcal{F}$ .
- Feature class  $\Phi$ .
- Policy distribution  $P \in \Delta(\Pi_{\mathbb{M}})$ .
- Number of samples  $n \in \mathbb{N}$ .
1: Set  $\varepsilon_{\mathrm{stat}} = O\left(\sqrt{cd^2n^{-1}\log(|\Phi| / \delta)}\right)$  for sufficiently absolute constant  $c > 0$  (see Appendix E).
2: Let  $\phi^{(1)} \in \Phi$  be arbitrary.
3: Set  $\mathcal{D} \gets \emptyset$ .
4: for  $n$  times  $\mathbf{do}$ 
5: Sample  $\pi \sim P$ .
6: Sample  $(x_h, a_h, x_{h+1}) \sim \pi \circ h \pi_{\mathrm{unif}}$ .
7: Update dataset:  $\mathcal{D} \gets \mathcal{D} \cup \{(x_h, a_h, x_{h+1})\}$ .
8: Define  $\mathcal{L}_{\mathcal{D}}(\phi, w, f) = \sum_{(x,a,x') \in \mathcal{D}} (\phi(x, a)^{\top} w - f(x'))^2$ .
9: for  $t = 1, 2, \ldots$  do /* Discriminator selection */
10: Solve
[ f^{(t)} \in \arg \max_{f \in \mathcal{F}} \widehat{\Delta}(f), \text{ where } \widehat{\Delta}(f) := \max_{\tilde{\phi} \in \Phi} \left\{ \min_{w \in B(3d^{3/2})} \mathcal{L}_{\mathcal{D}}(\phi^{(t)}, w, f) - \min_{\tilde{w} \in B(2\sqrt{d})} \mathcal{L}_{\mathcal{D}}(\tilde{\phi}, \tilde{w}, f) \right\}. ]
11: if  $\widehat{\Delta}(f^{(t)}) \leq 16dt\varepsilon_{\mathrm{stat}}^2$  then
12: Return  $\phi^{(t)}$ .
/* Feature selection via least-squares minimization */
13: Solve
[ \phi^{(t+1)} \in \arg \min_{\phi \in \Phi} \min_{(w_1, \ldots, w_t) \in B(2\sqrt{d})^t} \sum_{\ell=1}^{t} \mathcal{L}_{\mathcal{D}}(\phi, w_\ell, f^{(\ell)}). ]
Thus, under the event  $\mathcal{E} \coloneqq \bigcup_{t=1}^{h} \mathcal{E}_{t}$ , we have that
$$
\mathbb {E} ^ {\pi_ {\star}} \left[ \sum_ {t = 1} ^ {h} r _ {t} \left(\boldsymbol {x} _ {t}, \boldsymbol {a} _ {t}\right) \right] - \mathbb {E} ^ {\hat {\pi}} \left[ \sum_ {t = 1} ^ {h} r _ {t} \left(\boldsymbol {x} _ {t}, \boldsymbol {a} _ {t}\right) \right] \leq \varepsilon .
$$
The desired result follows by the fact that a union bound implies  $\mathbb{P}[\mathcal{E}] \geq 1 - \delta$ .
E Guarantee for RepLearn
In this section, we give a generic guarantee for RepLearn (Algorithm 6). Compared to previous guarantees in Modi et al. [35], Zhang et al. [49], we prove a fast  $1 / n$ -type rate of convergence for RepLearn, and show that the algorithm succeeds even when the norm of the weight  $w$  in Line 10 does not grow with the number of iterations. We also use the slightly simpler discriminator class:
$$
\mathcal {F} := \left\{f: x \mapsto \max  _ {a \in \mathcal {A}} \theta^ {\top} \phi (x, a) \mid \theta \in \mathcal {B} (1), \phi \in \Phi \right\}. \tag {20}
$$
The main guarantee for RepLearn is as follows.
Theorem E.1. Let  $h \in [H]$ ,  $\delta \in (0, e^{-1})$ , and  $n \in \mathbb{N}$  be given, and suppose that  $\mu_{h+1}^{\star}$  satisfies the normalization assumption in Eq. (4). For any function  $f \in \mathcal{F}$ , define
$$
w _ {f} = \int_ {\mathcal {X} _ {h + 1}} f (x) \mu_ {h + 1} ^ {\star} (x) \mathrm {d} \nu (x).
$$
Let  $P \in \Delta(\Pi_{\mathbb{M}})$  be a distribution over policies,  $\mathcal{F}$  be as (20), and  $\Phi$  be a feature class satisfying Assumption 2.2. With probability at least  $1 - \delta$ , RepLearn with input  $(h, \mathcal{F}, \Phi, P, n)$  terminates after  $t \leq T := \left[ d \log_{3/2}(2nd^{-1/2}) \right]$  iterations, and its output  $\phi^{(t)}$  satisfies
$$
\sup  _ {f \in \mathcal {F}} \inf  _ {w \in \mathcal {B} \left(3 d ^ {3 / 2}\right)} \mathbb {E} _ {\pi \sim P} \mathbb {E} ^ {\pi \circ_ {h} \pi_ {\text {u n i f}}} \left[ \left(w ^ {\top} \phi^ {(t)} \left(\boldsymbol {x} _ {h}, \boldsymbol {a} _ {h}\right) - w _ {f} ^ {\top} \phi_ {h} ^ {\star} \left(\boldsymbol {x} _ {h}, \boldsymbol {a} _ {h}\right)\right) ^ {2} \right] \leq \varepsilon_ {\text {R e p L e a r n}} ^ {2} (n, \delta), \tag {21}
$$
where  $\varepsilon_{\mathrm{RepLearn}}^2 (n,\delta)\coloneqq cTd^3 n^{-1}\log (|\Phi | / \delta)$ , for some sufficiently large absolute constant  $c > 0$
To prove the theorem, we need a technical lemma, which follows from Modi et al. [35, Lemma 14].
Lemma E.1. Consider a call to RepLearn  $(h, \mathcal{F}, \Phi, P, n)$  (Algorithm 6) in the setting of Theorem E.1. Further, let  $\mathcal{L}_{\mathcal{D}}$  be as in Algorithm 6 and define
$$
\left(\phi^ {(t)}, \widehat {w} _ {1} ^ {(t)}, \dots , \widehat {w} _ {t - 1} ^ {(t)}\right) \in \underset {\phi \in \Phi , (w _ {1}, \dots , w _ {t - 1}) \in \mathcal {B} (2 \sqrt {d}) ^ {t - 1}} {\arg \min } \sum_ {\ell = 1} ^ {t - 1} \mathcal {L} _ {\mathcal {D}} \left(\phi , w _ {\ell}, f ^ {(\ell)}\right). \tag {22}
$$
For any  $\delta \in (0,1)$ , there is an event  $\mathcal{E}^{(t)}(\delta)$  of probability at least  $1 - \delta$  such that under  $\mathcal{E}^{(t)}(\delta)$ , if Algorithm 6 does not terminate at iteration  $t \geq 1$ , then for  $w^{(\ell)} \coloneqq w_{f(\ell)}$ :
$$
\sum_ {\ell = 1} ^ {t - 1} \mathbb {E} _ {\pi \sim P} \mathbb {E} ^ {\pi \circ_ {h} \pi_ {\text {u n i f}}} \left[ \left(\phi^ {(t)} \left(\boldsymbol {x} _ {h}, \boldsymbol {a} _ {h}\right) ^ {\top} \widehat {w} _ {\ell} ^ {(t)} - \phi_ {h} ^ {\star} \left(\boldsymbol {x} _ {h}, \boldsymbol {a} _ {h}\right) ^ {\top} w ^ {(\ell)}\right) ^ {2} \right] \leq t \varepsilon_ {\text {s t a t}} ^ {2} (n, \delta), \tag {23}
$$
$$
\inf _ {w \in \frac {3}{2} \mathcal {B} (d ^ {3 / 2})} \mathbb {E} _ {\pi \sim P} \mathbb {E} ^ {\pi \circ_ {h} \pi_ {\mathrm {u n i f}}} \left[ \left(\phi^ {(t)} (\pmb {x} _ {h}, \pmb {a} _ {h}) ^ {\top} w - \phi_ {h} ^ {\star} (\pmb {x} _ {h}, \pmb {a} _ {h}) ^ {\top} w ^ {(t)}\right) ^ {2} \right] > 8 d t \varepsilon_ {\mathrm {s t a t}} ^ {2} (n, \delta),
$$
where  $\varepsilon_{\mathrm{stat}}^2 (n,\delta)\coloneqq cd^2 n^{-1}\log (|\Phi | / \delta)$  and  $c\geq 1$  is a sufficiently large absolute constant.
With this, we prove Theorem E.1.
Proof of Theorem E.1. Let us abbreviate  $\varepsilon \coloneqq \varepsilon_{\mathrm{stat}}(n,\delta)$ , with  $\varepsilon_{\mathrm{stat}}(n,\delta)$  defined as in Lemma E.1. Further, let  $N\coloneqq 1 + \left[d\log_{3 / 2}(2d^{3 / 2} / \varepsilon)\right]$ ,  $\delta^{\prime}\coloneqq \frac{\delta}{2N}$ , and define
$$
\tilde {\varepsilon} _ {\text {s t a t}} := \varepsilon_ {\text {s t a t}} (n, \delta^ {\prime}). \tag {24}
$$
Note that  $\varepsilon \leq \tilde{\varepsilon}_{\mathrm{stat}}$  and  $N - 1 \leq T$ , where  $T$  is the number of iterations in the theorem statement; the latter inequality follows by the facts that the absolute constant  $c$  in Lemma E.1 is at least 1 and  $\log (|\Phi| / \delta) \geq 1$ . We define an event  $\mathcal{E} \coloneqq \mathcal{E}^{(1)}(\delta') \cap \dots \cap \mathcal{E}^{(N)}(\delta')$ , where  $(\mathcal{E}^t(\cdot))_t$  are the success events in Lemma E.1. Note that  $\mathbb{P}[\mathcal{E}] \geq 1 - \delta / 2$  by the union bound. Throughout this proof, we condition on the event  $\mathcal{E}$ .
To begin the proof, we define a sequence of vectors  $(v_{1:d}^{(\ell)})_{\ell \geq 0}$  in an inductive fashion, with  $v_{i}^{(\ell)} \in \mathbb{R}^{d}$  for all  $i \in [d]$  and  $\ell \geq 0$ . For  $\ell = 0$ , we let  $v_{i}^{(0)} = \varepsilon e_{i} / d$ , for all  $i \in [d]$ . For  $\ell \geq 1$ , we consider two cases:
Case I: If
$$
\mathcal {J} ^ {(\ell)} := \left\{j \in [ d ] \left| | \det  (V _ {- j} ^ {(\ell - 1)}, w ^ {(\ell)}) | > (1 + C) \cdot | \det  (V ^ {(\ell - 1)}) | \right. \right\} \neq \varnothing ,
$$
where  $V^{(\ell -1)}\coloneqq (v_1^{(\ell -1)},\ldots ,v_d^{(\ell -1)})\in \mathbb{R}^{d\times d}$  and  $w^{(\ell)}\coloneqq w_{f^{(\ell)}}$ , then we let  $j\coloneqq \arg \min_{j'\in \mathcal{J}^{(\ell)}}j'$  and define
$$
v _ {i} ^ {(\ell)} := \left\{ \begin{array}{l l} w ^ {(\ell)}, & \text {i f} i = j, \\ v _ {i} ^ {(\ell - 1)}, & \text {o t h e r w i s e}. \end{array} \right.
$$
- Case II: If  $\mathcal{J}^{(\ell)} = \varnothing$ , we let  $v_{i}^{(\ell)} = v_{i}^{(\ell -1)}$ , for all  $i\in [d]$ .
We first show that  $\mathcal{J}^{(t)}\neq \emptyset$  at any iteration  $t\in [N]$  where RepLearn does not terminate. Let  $t\in [N]$  be an iteration where the algorithm does not terminate, and suppose that  $\mathcal{J}^{(t)} = \emptyset$ . This means that
$$
\forall j \in [ d ], \quad | \det  \left(V _ {- j} ^ {(t - 1)}, w ^ {(t)}\right) | \leq (1 + C) \cdot | \det  \left(V ^ {(t - 1)}\right) |. \tag {25}
$$
Now, since  $\operatorname{det}(V^{(t-1)}) \neq 0$  (note that  $|\operatorname{det}(V^{(t)})|$  is non-decreasing with  $t$ ), we have that  $\operatorname{span}(V^{(t-1)}) = \mathbb{R}^d$ . Thus, there exist  $\beta_1, \ldots, \beta_d \in \mathbb{R}$  such that  $w^{(t)} = \sum_{i=1}^{d} \beta_i v_i^{(t-1)}$ . By the linearity of the determinant and (25), we have
$$
\begin{array}{l} \forall j \in [ d ], \quad (1 + C) | \cdot \det  (V ^ {(t - 1)}) | \geq | \det  (V _ {- j} ^ {(t - 1)}, w ^ {(t)}) |, \\ = \left| \det  \left(V _ {- j} ^ {(t - 1)}, \sum_ {i = 1} ^ {d} \beta_ {i} v _ {i} ^ {(t - 1)}\right) \right|, \\ = \left| \sum_ {i \in [ d ]} \beta_ {i} \cdot \det  \left(V _ {- j} ^ {(t - 1)}, v _ {i} ^ {(t - 1)}\right) \right|, \\ = \left| \beta_ {j} \right| \cdot \left| \det  \left(V ^ {(t - 1)}\right) \right|. \\ \end{array}
$$
This implies that  $|\beta_j| \leq (1 + C)$  for all  $j \in [d]$ . Now, note that by the definition of  $(v_i^{(t-1)})$ , we have that for any  $i \in [d]$  such that  $v_i^{(t-1)} \neq \varepsilon e_i / d$ , there exists  $\ell \in [t-1]$  such that  $w^{(\ell)} = v_i^{(t-1)}$ . Let
$$
\mathcal {I} ^ {(t)} := \left\{i \in [ d ] \mid v _ {i} ^ {(t - 1)} \neq \varepsilon e _ {i} / d \right\},
$$
and for any  $i \in \mathcal{I}^{(t)}$ , let  $\ell_i \in [t - 1]$  be such that  $w^{(\ell_i)} = v_i^{(t - 1)}$ . Further, define