SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} \left| \left\langle \widehat {w}, V ^ {\pi} \right\rangle - \left\langle \widehat {w} ^ {(\delta)}, V ^ {\pi} \right\rangle \right| \leq \underbrace {\left| \sum_ {i = 1} ^ {d _ {\delta}} a _ {i} \left(\left\langle \widehat {w} , \widehat {A} _ {i} \right\rangle - \left\langle w ^ {(\delta)} , \widehat {A} _ {i} \right\rangle\right) \right|} _ {(a)} + \underbrace {\left| \sum_ {i = d _ {\delta} + 1} ^ {d} a _ {i} \left\langle \widehat {w} , \psi_ {i} \right\rangle \right|} _ {(b)} \\ + \underbrace {\left| \sum_ {i = d _ {\delta} + 1} ^ {d} a _ {i} \left\langle \widehat {w} ^ {(\delta)} , \psi_ {i} \right\rangle \right|} _ {(c)}. \\ \end{array}
$$
Since  $\widehat{A}^{(\delta)}\widehat{w} = \widehat{A}^{(\delta)}\widehat{w}^{(\delta)} = \mathbf{e}_1$ , we have term  $(a) = 0$
We move on to bound term (c).
Note that the vectors  $\{\psi_i | i = d_\delta + 1, \ldots, d\}$  are orthogonal to  $\operatorname{span}(V^{\pi_1}, \ldots, V^{\pi_{d_\delta}})$  and that together with  $V^{\pi_1}, \ldots, V^{\pi_{d_\delta}}$  they form a basis for  $\operatorname{span}(\{V^\pi | \pi \in \Pi\})$ .
Thus, we can let  $b_{i}$  in the proof of Lemma 2 be  $\psi_{i + d_{\delta}}$ .
In addition, all the properties of  $\{b_i|i\in [d - d_\delta ]\}$  also apply to  $\{\psi_i|i = d_\delta +1,\dots ,d\}$  as well.
Hence, similarly to Eq (9),
$$
\sqrt {\sum_ {i = d \delta + 1} ^ {d} a _ {i} ^ {2}} \leq \sqrt {k} \delta .
$$
Consequently, we can bound term  $(c)$  is by
$$
(c) \leq \sqrt {k} \delta \left\| \widehat {w} ^ {(\delta)} \right\| _ {2} = \frac {3}{2} \sqrt {k} \delta \| w ^ {\prime} \| _ {2}
$$
since  $\| \widehat{w}^{(\delta)}\| _2\leq \frac{3}{2}\| w'\| _2$  when  $C_V\sqrt{d_\delta}\| w'\| _2\eta_{\epsilon_\alpha ,\delta}\leq \frac{1}{3}$  as discussed in the proof of Lemma 6.
Now all is left is to bound term (b).
We cannot bound term (b) in the same way as that of term (c) because  $\| \widehat{w} \|_2$  is not guaranteed to be bounded by  $\| w' \|_2$ .
For  $i = d_{\delta} + 1,\dots ,d$  we define
$$
\epsilon_ {i} := \left| \left\langle \psi_ {i}, \widehat {A} _ {i} \right\rangle \right|.
$$
For any  $i,j = d_{\delta} + 1,\ldots ,d$ $\psi_{i}$  is perpendicular to  $V^{\pi_1}$ , thus  $\left|\left\langle \psi_i,\widehat{A}_j\right\rangle \right| = |\langle \psi_i,\widehat{\alpha}_{j - 1}V^{\pi_1} - V^{\pi_j}\rangle | = |\langle \psi_i,V^{\pi_j}\rangle |.$  Especially, we have
$$
\epsilon_ {i} = \left| \left\langle \psi_ {i}, \widehat {A} _ {i} \right\rangle \right| = | \langle \psi_ {i}, V ^ {\pi_ {i}} \rangle |.
$$
Let  $\widehat{A}_i^{\parallel} := \widehat{A}_i - \sum_{j=d_\delta+1}^{d} \left\langle \widehat{A}_i, \psi_j \right\rangle \psi_j$  denote  $\widehat{A}_i$ 's projection into  $\operatorname{span}(\widehat{A}_1, \ldots, \widehat{A}_{d_\delta})$ .
Since  $\widehat{A}_i$  has zero component in direction  $\psi_j$  for  $j > i$ , we have  $\widehat{A}_i^{\parallel} = \widehat{A}_i - \sum_{j = d_\delta +1}^i\left\langle \widehat{A}_i,\psi_j\right\rangle \psi_j$ .
$$
0 = \left\langle \widehat {w}, \widehat {A} _ {i} \right\rangle = \widehat {w} \cdot \widehat {A} _ {i} ^ {\parallel} + \widehat {w} \cdot \sum_ {j = d _ {\delta} + 1} ^ {i} \left\langle \widehat {A} _ {i}, \psi_ {j} \right\rangle \psi_ {j} = \widehat {w} \cdot \widehat {A} _ {i} ^ {\parallel} - \sum_ {j = d _ {\delta} + 1} ^ {i} \left\langle V ^ {\pi_ {i}}, \psi_ {j} \right\rangle \left\langle \widehat {w}, \psi_ {j} \right\rangle ,
$$
where the first equation holds due to  $\widehat{A}\widehat{w} = \mathbf{e}_1$
$$
\langle V ^ {\pi_ {i}}, \psi_ {i} \rangle \langle \widehat {w}, \psi_ {i} \rangle = \widehat {w} \cdot \widehat {A} _ {i} ^ {\parallel} - \sum_ {j = d _ {\delta} + 1} ^ {i - 1} \langle V ^ {\pi_ {i}}, \psi_ {j} \rangle \langle \widehat {w}, \psi_ {j} \rangle . \tag {17}
$$
Recall that at iteration  $j$  of Algorithm 1, in line 8 we pick an orthonormal basis  $\rho_{1},\ldots ,\rho_{k + 1 - j}$  of  $\mathrm{span}(V^{\pi_1},\ldots ,V^{\pi_{j - 1}})^\perp$ . Since  $\psi_{j}$  is in  $\mathrm{span}(V^{\pi_1},\ldots ,V^{\pi_{j - 1}})^\perp$  according to the definition of  $\psi_{j}$ ,  $|\langle V^{\pi_i},\psi_j\rangle |$  is no greater than the norm of  $V^{\pi_i}$ 's projection into  $\mathrm{span}(V^{\pi_1},\ldots ,V^{\pi_{j - 1}})^\perp$ .
$$
\left| \langle V ^ {\pi_ {i}}, \psi_ {j} \rangle \right| \leq \sqrt {k} \max  _ {l \in [ k + 1 - j ]} \left| \langle V ^ {\pi_ {i}}, \rho_ {l} \rangle \right| \stackrel {(d)} {\leq} \sqrt {k} \max  _ {l \in [ k + 1 - j ]} \max  (\left| \left\langle V ^ {\pi^ {\rho_ {l}}}, \rho_ {l} \right\rangle \right|, \left| \left\langle V ^ {\pi^ {- \rho_ {l}}}, - \rho_ {l} \right\rangle \right|)
$$
$$
\stackrel {(e)} {=} \sqrt {k} | \langle V ^ {\pi_ {j}}, u _ {j} \rangle | \stackrel {(f)} {\leq} \sqrt {k} | \langle V ^ {\pi_ {j}}, \psi_ {j} \rangle | = \sqrt {k} \epsilon_ {j}, \tag {18}
$$
where inequality (d) holds because  $\pi^{\rho_l}$  is the optimal personalized policy with respect to the preference vector  $\rho_{l}$ , and Equation (e) holds due to the definition of  $u_{j}$  (line 9 of Algorithm 1). Inequality (f) holds since  $\langle V^{\pi_j},\psi_j\rangle$  is the norm of  $V^{\pi_j}$ 's projection in  $\operatorname {span}(V^{\pi_1},\ldots ,V^{\pi_{j - 1}})^\perp$  and  $u_{j}$  belongs to  $\operatorname {span}(V^{\pi_1},\dots ,V^{\pi_{j - 1}})^\perp$
By taking absolute value on both sides of Eq (17), we have
$$
\epsilon_ {i} \left| \langle \widehat {w}, \psi_ {i} \rangle \right| = \left| \widehat {w} \cdot \widehat {A} _ {i} ^ {\parallel} - \sum_ {j = d _ {\delta} + 1} ^ {i - 1} \left\langle V ^ {\pi_ {i}}, \psi_ {j} \right\rangle \left\langle \widehat {w}, \psi_ {j} \right\rangle \right| \leq \left| \widehat {w} \cdot \widehat {A} _ {i} ^ {\parallel} \right| + \sqrt {k} \sum_ {j = d _ {\delta} + 1} ^ {i - 1} \epsilon_ {j} | \langle \widehat {w}, \psi_ {j} \rangle |. \tag {19}
$$
We can now bound  $\left|\widehat{w} \cdot \widehat{A}_i^{\parallel}\right|$  as follows.
$$
\begin{array}{l} \left| \widehat {w} \cdot \widehat {A} _ {i} ^ {\parallel} \right| = \left| \widehat {w} ^ {(\delta)} \cdot \widehat {A} _ {i} ^ {\parallel} \right| (20) \\ = \left| \widehat {w} ^ {(\delta)} \cdot \left(\widehat {A} _ {i} - \sum_ {j = d _ {\delta} + 1} ^ {d} \left\langle \widehat {A} _ {i}, \psi_ {j} \right\rangle \psi_ {j}\right) \right| = \left| \widehat {w} ^ {(\delta)} \cdot \widehat {A} _ {i} \right| (21) \\ \leq \left| \widehat {w} ^ {(\delta)} \cdot A _ {i} \right| + \left| \widehat {w} ^ {(\delta)} \cdot (\widehat {A} _ {i} - A _ {i}) \right| \\ \leq \left| w ^ {\prime} \cdot A _ {i} \right| + \left| \left(\widehat {w} ^ {(\delta)} - w ^ {\prime}\right) \cdot A _ {i} \right| + \left| \widehat {w} ^ {(\delta)} \cdot \left(\widehat {A} _ {i} - A _ {i}\right) \right| \\ \leq 0 + \left(C _ {\alpha} + 1\right) \sup  _ {\pi} \left| \left(\widehat {w} ^ {(\delta)} - w ^ {\prime}\right) \cdot V ^ {\pi} \right| + C _ {V} \left\| \widehat {w} ^ {(\delta)} \right\| _ {2} \epsilon_ {\alpha} \\ \leq C ^ {\prime} C _ {\alpha} \epsilon^ {(\delta)}, \\ \end{array}
$$
for some constant  $C' > 0$ .
Eq (20) holds because  $\widehat{A}^{(\delta)}\widehat{w} = \widehat{A}^{(\delta)}\widehat{w}^{(\delta)} = \mathbf{e}_1$  and  $\widehat{A}_i^{\parallel}$  belongs to  $\operatorname {span}(\widehat{A}^{(\delta)})$ . Eq (21) holds because  $\widehat{w}^{(\delta)}$  is the minimum norm solution to  $\widehat{A}^{(\delta)}x = \mathbf{e}_1$ , which implies that  $\widehat{w}^{(\delta)}\cdot \psi_i = 0$ . The last inequality follows by applying Lemma 2.
We will bound  $\epsilon_{i}|\langle \widehat{w},\psi_{i}\rangle |$  by induction on  $i = d_{\delta} + 1,\ldots ,d$
In the base case of  $i = d_{\delta} + 1$
$$
\epsilon_ {d _ {\delta} + 1} | \langle \widehat {w}, \psi_ {d _ {\delta} + 1} \rangle | \leq \left| \widehat {w} \cdot \widehat {A} _ {d _ {\delta} + 1} ^ {\parallel} \right| \leq C ^ {\prime} C _ {\alpha} \epsilon^ {(\delta)}.
$$
Then, by induction through Eq (19), we have for  $i = d_{\delta} + 2,\dots ,d$
$$
\epsilon_ {i} | \langle \widehat {w}, \psi_ {i} \rangle | \leq (\sqrt {k} + 1) ^ {i - d _ {\delta} - 1} C ^ {\prime} C _ {\alpha} \epsilon^ {(\delta)}.
$$
Similar to the deviation of Eq (18), we pick an orthonormal basis  $\rho_1, \ldots, \rho_{k+1-i}$  of  $\operatorname{span}(V^{\pi_1}, \ldots, V^{\pi_{i-1}})^\perp$  at line 8 of Algorithm 1, then we have that, for any policy  $\pi$ ,
$$
| \langle V ^ {\pi}, \psi_ {i} \rangle | \leq \sqrt {k} \max _ {l \in [ k + 1 - i ]} | \langle V ^ {\pi}, \rho_ {l} \rangle | \leq \sqrt {k} | \langle V ^ {\pi_ {i}}, u _ {i} \rangle | \leq \sqrt {k} | \langle V ^ {\pi_ {i}}, \psi_ {i} \rangle | = \sqrt {k} \epsilon_ {i}.
$$
Then we have that term (b) is bounded by
$$
(b) = \left| \sum_ {i = d _ {\delta} + 1} ^ {d} \langle V ^ {\pi}, \psi_ {i} \rangle \langle \widehat {w}, \psi_ {i} \rangle \right| \leq \sum_ {i = d _ {\delta} + 1} ^ {d} | \langle V ^ {\pi}, \psi_ {i} \rangle | \cdot | \langle \widehat {w}, \psi_ {i} \rangle |
$$
$$
\leq \sqrt {k} \sum_ {i = d _ {\delta} + 1} ^ {d} \epsilon_ {i} | \langle \widehat {w}, \psi_ {i} \rangle | \leq (\sqrt {k} + 1) ^ {d - d _ {\delta}} C ^ {\prime} C _ {\alpha} \epsilon^ {(\delta)}.
$$
Hence we have that for any policy  $\pi$
$$
\left| \langle \widehat {w}, V ^ {\pi} \rangle - \left\langle \widehat {w} ^ {(\delta)}, V ^ {\pi} \right\rangle \right| \leq (\sqrt {k} + 1) ^ {d - d _ {\delta}} C ^ {\prime} C _ {\alpha} \epsilon^ {(\delta)} + \frac {3}{2} \sqrt {k} \delta \| w ^ {\prime} \| _ {2}.
$$
H Proof of Theorem 2
Theorem 2. Consider the algorithm of computing  $\widehat{A}$  and any solution  $\widehat{w}^{(\delta)}$  to  $\widehat{A}^{(\delta)}x = \mathbf{e}_1$  for  $\delta = k^{\frac{5}{3}}\epsilon^{\frac{1}{3}}$  and outputting the policy  $\pi^{\widehat{w}^{(\delta)}} = \arg \max_{\pi \in \Pi}\left\langle \widehat{w}^{(\delta)},V^{\pi}\right\rangle$ . Then the policy  $\pi^{\widehat{w}^{(\delta)}}$  satisfies that  $v^{*} - \left\langle w^{*},V^{\pi^{\widehat{w}^{(\delta)}}}\right\rangle \leq \mathcal{O}\left(k^{\frac{13}{6}\epsilon^{\frac{1}{3}}}\right)$ .
Proof of Theorem 2. As shown in Lemma 1, we set  $C_{\alpha} = 2k$  and have  $\epsilon_{\alpha} = \frac{4k^2\epsilon}{v^*}$ . We have  $\| w' \| = \frac{\|w^{*}\|}{\langle w^{*}, V^{\pi_{1}} \rangle}$  and showed that  $\langle w^{*}, V^{\pi_{1}} \rangle \geq \frac{v^{*}}{2k}$  in the proof of Lemma 1.
By applying Lemma 2 and setting  $\delta = \left(\frac{C_V^4k^5\|w^*\|_2\epsilon}{v^{*2}}\right)^{\frac{1}{3}}$ , we have
$$
\begin{array}{l} v ^ {*} - \left\langle w ^ {*}, V ^ {\pi^ {\hat {w} ^ {(\delta)}}} \right\rangle = \left\langle w ^ {*}, V ^ {\pi_ {1}} \right\rangle \left(\left\langle w ^ {\prime}, V ^ {\pi^ {*}} \right\rangle - \left\langle w ^ {\prime}, V ^ {\pi^ {\hat {w} ^ {(\delta)}}} \right\rangle\right) \\ \leq \langle w ^ {*}, V ^ {\pi_ {1}} \rangle \left(\left\langle \widehat {w} ^ {(\delta)}, V ^ {\pi^ {*}} \right\rangle - \left\langle \widehat {w} ^ {(\delta)}, V ^ {\pi^ {\widehat {w} ^ {(\delta)}}} \right\rangle + \mathcal {O} (\sqrt {k} \| w ^ {\prime} \| _ {2} \left(\frac {C _ {V} ^ {4} k ^ {5} \| w ^ {*} \| _ {2} \epsilon}{v ^ {* 2}}\right) ^ {\frac {1}{3}})\right) \\ = \mathcal {O} (\sqrt {k} \| w ^ {*} \| _ {2} (\frac {C _ {V} ^ {4} k ^ {5} \| w ^ {*} \| _ {2} \epsilon}{v ^ {* 2}}) ^ {\frac {1}{3}}). \\ \end{array}
$$
I Dependency on  $\epsilon$
In this section, we would like to discuss a potential way of improving the dependency on  $\epsilon$  in Theorems 1 and 2.
Consider a toy example where the returned three basis policies are  $\pi_1$  with  $V^{\pi_1} = (1,0,0)$ ,  $\pi_2$  with  $V^{\pi_2} = (1,1,1)$  and  $\pi_3$  with  $V^{\pi_3} = (1,\eta, -\eta)$  for some  $\eta > 0$  and  $w^* = (1,w_2,w_3)$  for some  $w_2, w_3$ .
The estimated ratio  $\widehat{\alpha}_{1}$  lies in  $[1 + w_{2} + w_{3} - \epsilon, 1 + w_{2} + w_{3} + \epsilon]$ , and  $\widehat{\alpha}_{2}$  lies in  $[1 + \eta w_{2} - \eta w_{3} - \epsilon, 1 + \eta w_{2} - \eta w_{3} + \epsilon]$ . Suppose that  $\widehat{\alpha}_{1} = 1 + w_{2} + w_{3} + \epsilon$  and  $\widehat{\alpha}_{2} = 1 + \eta w_{2} - \eta w_{3} + \epsilon$ .
By solving
$$
\left( \begin{array}{c c c} 1 & 0 & 0 \\ w _ {2} + w _ {3} + \epsilon & - 1 & - 1 \\ \eta w _ {2} - \eta w _ {3} + \epsilon & - \eta & \eta \end{array} \right) \widehat {w} = \left( \begin{array}{c} 1 \\ 0 \\ 0 \end{array} \right)
$$
we can derive  $\widehat{w}_2 = w_2 + \frac{\epsilon}{2} (1 + \frac{1}{\eta})$  and  $\widehat{w}_3 = w_3 + \frac{\epsilon}{2} (1 - \frac{1}{\eta})$ .
The quantity measuring sub-optimality we care about is  $\sup_{\pi}|\langle \widehat{w},V^{\pi}\rangle -\langle w^{*},V^{\pi}\rangle |$  , which is upper bounded by  $C_V\| \widehat{w} -w^*\| _2$  . But the  $\ell_2$  distance between  $\widehat{w}$  and  $w^{*}$  depends on the condition number of  $\widehat{A}$  , which is large when  $\eta$  is small. To obtain a non-vacuous upper bound in Section 3, we introduce another estimate  $\widehat{w}^{(\delta)}$  based on the truncated version of  $\widehat{A}$  and then upper bound  $\left\| \widehat{w}^{(\delta)} - w^{*}\right\|_{2}$  and  $\sup_{\pi}\big|\langle \widehat{w},V^{\pi}\rangle -\big{\langle}\widehat{w}^{(\delta)},V^{\pi}\big{\rangle}\big|$  separately.
However, it is unclear if  $\sup_{\pi}|\langle \widehat{w},V^{\pi}\rangle -\langle w^{*},V^{\pi}\rangle |$  depends on the condition number of  $\widehat{A}$ . Due to the construction of Algorithm 1, we can obtain some extra information about the set of all policy values.
First, since we find  $\pi_2$  before  $\pi_3$ ,  $\eta$  must be no greater than 1. According to the algorithm,  $V^{\pi_2}$  is the optimal policy when the preference vector is  $u_2$  (see line 11 of Algorithm 1 for the definition of  $u_2$ ) and  $V^{\pi_3}$  is the optimal policy when the preference vector is  $u_3 = (0,1,-1)$ . Note that the angle between  $u_2$  and  $V^{\pi_2}$  is no greater than 45 degrees according to the definition of  $u_2$ . Then the values of all policies can only lie in the small box  $B = \{x \in \mathbb{R}^3 | |u_2^\top x| \leq |\langle u_2, V^{\pi_2} \rangle|, |u_3^\top x| \leq |\langle u_3, V^{\pi_3} \rangle|\}$ . It is direct to check that for any  $x \in B$ ,  $|\langle \widehat{w}, x \rangle - \langle w^*, x \rangle| < (1 + \sqrt{2})\epsilon$ . This example illustrates that even when the condition number of  $\widehat{A}$  is large,  $\sup_{\pi} |\langle \widehat{w}, V^{\pi} \rangle - \langle w^*, V^{\pi} \rangle|$  can be small. It is unclear if this holds in general. Applying this additional information to upper bound  $\sup_{\pi} |\langle \widehat{w}, V^{\pi} \rangle - \langle w^*, V^{\pi} \rangle|$  directly instead of through bounding  $C_V \| \widehat{w} - w^* \|_2$  is a possible way of improving the term  $\epsilon^{\frac{1}{3}}$ .
J Description of C4 and Proof of Lemma 4
Algorithm 4 C4: Compress Convex Combination using Carathéodory's theorem
1: input a set of  $k$ -dimensional vectors  $M \subset \mathbb{R}^k$  and a distribution  $p \in \mathrm{Simplex}^M$   
2: while  $|M| > k + 1$  do  
3: arbitrarily pick  $k + 2$  vectors  $\mu_1, \ldots, \mu_{k+2}$  from  $M$   
4: solve for  $x \in \mathbb{R}^{k+2}$  s.t.  $\sum_{i=1}^{k+2} x_i (\mu_i \circ 1) = 0$ , where  $\mu \circ 1$  denote the vector of appending 1 to  $\mu$   
5:  $i_0 \gets \arg \max_{i \in [k+2]} \frac{|x_i|}{p(\mu_i)}$   
6: if  $x_{i_0} < 0$  then  $x \gets -x$   
7:  $\gamma \gets \frac{p(\mu_{i_0})}{x_{i_0}}$  and  $\forall i \in [k+2]$ ,  $p(\mu_i) \gets p(\mu_i) - \gamma x_i$   
8: remove  $\mu_i$  with  $p(\mu_i) = 0$  from  $M$   
9: end while  
10: output  $M$  and  $p$
Proof. The proof is similar to the proof of Carathéodory's theorem. Given the vectors  $\mu_1, \ldots, \mu_{k+2}$  picked in line 3 of Algorithm 4 and their probability masses  $p(\mu_i)$ , we solve  $x \in \mathbb{R}^{k+2}$  s.t.  $\sum_{i=1}^{k+2} x_i (\mu_i \circ 1) = 0$  in the algorithm.
Note that there exists a non-zero solution of  $x$  because  $\{\mu_i \circ 1 | i \in [k + 2]\}$  are linearly dependent. Besides,  $x$  satisfies  $\sum_{i=1}^{d+2} x_i = 0$ .
$$
\sum_ {i = 1} ^ {d + 2} (p (\mu_ {i}) - \gamma x _ {i}) = \sum_ {i = 1} ^ {d + 2} p (\mu_ {i}).
$$
For all  $i$ , if  $x_{i} < 0$ ,  $p(\mu_i) - \gamma x_i \geq 0$  as  $\gamma > 0$ ; if  $x_{i} > 0$ , then  $\frac{x_i}{p(\mu_i)} \leq \frac{x_{i_0}}{p(\mu_{i_0})} = \frac{1}{\gamma}$  and thus  $p(\mu_i) - \gamma x_i \geq 0$ .
Hence, after one iteration, the updated  $p$  is still a probability over  $M$  (i.e.,  $p(\mu) \geq 0$  for all  $\mu \in M$  and  $\sum_{\mu \in M} p(M) = 1$ ). Besides,  $\sum_{i=1}^{d+2} (p(\mu_i) - \gamma x_i) \mu_i = \sum_{i=1}^{d+2} p(\mu_i) \mu_i - \gamma \sum_{i=1}^{d+2} x_i \mu_i = \sum_{i=1}^{d+2} p(\mu_i) \mu_i$ .
Therefore, after one iteration, the expected value  $\mathbb{E}_{\mu \sim p}[\mu ]$  is unchanged.
When we finally output  $(M', q)$ , we have that  $q$  is a distribution over  $M$  and that  $\mathbb{E}_{\mu \sim q}[\mu] = \mathbb{E}_{\mu \sim p}[\mu]$ .
Due to line 6 of the algorithm, we know that  $x_{i_0} > 0$ . Hence  $p(\mu_{i_0}) - \gamma x_{i_0} = p(\mu_{i_0}) - \frac{p(\mu_{i_0})}{x_{i_0}} x_{i_0} = 0$ .
We remove at least one vector  $\mu_{i_0}$  from  $M$  and we will run for at most  $|M|$  iterations.
K Flow Decomposition Based Approach
We first introduce an algorithm based on the idea of flow decomposition.
For that, we construct a layer graph  $G = ((L^{(0)} \cup \ldots \cup L^{(H + 1)}), E)$  with  $H + 2$  pairwise disjoint layers  $L^{(0)}, \ldots, L^{(H + 1)}$ , where every layer  $t \leq H$  contains a set of vertices labeled by the (possibly duplicated) states reachable at the corresponding time step  $t$ , i.e.,  $\{s \in S | \operatorname*{Pr}(S_t = s | S_0 = s_0) > 0\}$ .
Let us denote by  $x_{s}^{(t)}$  the vertex in  $L^{(t)}$  labeled by state  $s$ .
Layer  $L^{(H + 1)} = \{x_{*}^{(H + 1)}\}$  contains only an artificial vertex,  $x_{*}^{(H + 1)}$ , labeled by an artificial state *.
For  $t = 0, \ldots, H - 1$ , for every  $x_{s}^{(t)} \in L^{(t)}, x_{s'}^{(t + 1)} \in L^{(t + 1)}$ , we connect  $x_{s}^{(t)}$  and  $x_{s'}^{(t + 1)}$  by an edge labeled by  $(s, s')$  if  $P(s'|s, \pi(s)) > 0$ . Every vertex  $x_{s}^{(H)}$  in layer  $H$  is connected to  $x_{*}^{(H + 1)}$  by one edge, which is labeled by  $(s, *)$ .
We denote by  $E^{(t)}$  the edges between  $L^{(t)}$  and  $L^{(t + 1)}$ . Note that every trajectory  $\tau = (s_0, s_1, \ldots, s_H)$  corresponds to a single path  $(x_{s_0}^{(0)}, x_{s_1}^{(1)}, \ldots, x_{s_H}^{(H)}, x_*^{(H + 1)})$  of length  $H + 2$  from  $x_{s_0}^{(0)}$  to  $x_*^{(H + 1)}$ .
This is a one-to-one mapping and in the following, we use path and trajectory interchangeably.
The policy corresponds to a  $(x_{s_0}^{(0)}, x_*^{(H + 1)})$ -flow with flow value 1 in the graph  $G$ . In particular, the flow is defined as follows.
When the layer  $t$  is clear from the context, we actually refer to vertex  $x_{s}^{(t)}$  by saying vertex  $s$ .
For  $t = 0,\dots ,H - 1$  , for any edge  $(s,s^{\prime})\in E^{(t)}$  , let  $f:E\to \mathbb{R}^{+}$  be defined as
$$
f (s, s ^ {\prime}) = \sum_ {\tau : \left(s _ {t} ^ {\tau}, s _ {t + 1} ^ {\tau}\right) = \left(s, s ^ {\prime}\right)} q ^ {\pi} (\tau), \tag {22}
$$
where  $q^{\pi}(\tau)$  is the probability of  $\tau$  being sampled.
For any edge  $(s, *) \in E^{(H)}$ , let  $f(s, *) = \sum_{(s', s) \in E^{(H-1)}} f(s', s)$ . It is direct to check that the function  $f$  is a well-defined flow. We can therefore compute  $f$  by dynamic programming.
For all  $(s_0,s)\in E^{(0)}$ , we have  $f(s_{0},s) = P(s|s_{0},\pi (s_{0}))$  and for  $(s,s^{\prime})\in E^{(t)}$
$$
f (s, s ^ {\prime}) = P \left(s ^ {\prime} \mid s, \pi (s)\right) \sum_ {s ^ {\prime \prime}: \left(s ^ {\prime \prime}, s\right) \in E ^ {(t - 1)}} f \left(s ^ {\prime \prime}, s\right). \tag {23}
$$
Now we are ready to present our algorithm by decomposing  $f$  in Algorithm 5.
Each iteration in Algorithm 5 will zero out at least one edge and thus, the algorithm will stop within  $|E|$  rounds.
Algorithm 5 Flow decomposition based approach
1: initialize  $Q \gets \emptyset$ .
2: calculate  $f(e)$  for all edge  $e \in E$  by dynamic programming according to Eq (23)
3: while  $\exists e\in E$  s.t.  $f(e) > 0$  do
4: pick a path  $\tau = (s_0, s_1, \ldots, s_H, *) \in L^{(0)} \times L^{(1)} \times \ldots \times L^{(H+1)}$  s.t.  $f(s_i, s_{i+1}) > 0 \forall i \geq 0$
5:  $f_{\tau}\gets \min_{e\text{in}\tau}f(e)$
6:  $Q \gets Q \cup \{(\tau, f_{\tau})\}$ ,  $f(e) \gets f(e) - f_{\tau}$  for  $e$  in  $\tau$
7: end while
8: output  $Q$
Theorem 4. Algorithm 5 outputs  $Q$  satisfying that  $\sum_{(\tau, f_{\tau}) \in Q} f_{\tau} \Phi(\tau) = V^{\pi}$  in time  $\mathcal{O}(H^{2} |S|^{2})$ .
Finally, solving  $x$  takes  $\mathcal{O}(k^3)$  time and thus, Algorithm 4 takes  $\mathcal{O}(|M| k^3)$  time in total.
The core idea of the proof is that for any edge  $(s, s') \in E^{(t)}$ , the flow on  $(s, s')$  captures the probability of  $S_{t} = s \wedge S_{t+1} = s'$  and thus, the value of the policy  $V^{\pi}$  is linear in  $\{f(e) | e \in E\}$ .
The output  $Q$  has at most  $|E|$  number of weighted paths (trajectories). We can further compress the representation through C4, which takes  $O(|Q|k^3)$  time.
Corollary 2. Executing Algorithm 5 with the output  $Q$  first and then running C4 over  $\{(\Phi(\tau), f_{\tau}) | (\tau, f_{\tau}) \in Q\}$  returns a  $(k + 1)$ -sized weighted trajectory representation in time  $\mathcal{O}(H^2 |S|^2 + k^3 H |S|^2)$ .
We remark that the running time of this flow decomposition approach underperforms that of the expanding and compressing approach (see Theorem 3) whenever  $|S|H + |S|k^3 = \omega(k^4 + k|S|)$ .
L Proof of Theorem 4
Theorem 4. Algorithm 5 outputs  $Q$  satisfying that  $\sum_{(\tau, f_{\tau}) \in Q} f_{\tau} \Phi(\tau) = V^{\pi}$  in time  $\mathcal{O}(H^{2}|\mathcal{S}|^{2})$ .
Proof. Correctness: The function  $f$  defined by Eq (22) is a well-defined flow since for all  $t = 1, \ldots, H$ , for all  $s \in L^{(t)}$ , we have that
$$
\begin{array}{l} \sum_ {s ^ {\prime} \in L ^ {(t + 1)}: (s, s ^ {\prime}) \in E ^ {(t)}} f (s, s ^ {\prime}) = \sum_ {s ^ {\prime} \in L ^ {(t + 1)}: (s, s ^ {\prime}) \in E ^ {(t)}} \sum_ {\tau : (s _ {t} ^ {\tau}, s _ {t + 1} ^ {\tau}) = (s, s ^ {\prime})} q ^ {\pi} (\tau) = \sum_ {\tau : s _ {t} ^ {\tau} = s} q ^ {\pi} (\tau) \\ = \sum_ {s ^ {\prime \prime} \in L ^ {(t - 1)}: (s ^ {\prime \prime}, s) \in E ^ {(t - 1)}} f \left(s ^ {\prime \prime}, s\right). \\ \end{array}
$$
In the following, we first show that Algorithm 5 will terminate with  $f(e) = 0$  for all  $e \in E$ .
First, after each iteration,  $f$  is still a feasible  $(x^{(0)}, x^{(H + 1)})$ -flow feasible flow with the total flow out-of  $x^{(0)}$  reduced by  $f_{\tau}$ . Besides, for edge  $e$  with  $f(e) > 0$  at the beginning, we have  $f(e) \geq 0$  throughout the algorithm because we never reduce  $f(e)$  by an amount greater than  $f(e)$ .
Then, since  $f$  is a  $(x^{(0)}, x^{(H + 1)})$ -flow and  $f(e) \geq 0$  for all  $e \in E$ , we can always find a path  $\tau$  in line 4 of Algorithm 5.
Otherwise, the set of vertices reachable from  $x^{(0)}$  through edges with positive flow does not contain  $x^{(H + 1)}$  and the flow out of this set equals the total flow out-of  $x^{(0)}$ . But since other vertices are not reachable, there is no flow out of this set, which is a contradiction.
In line 6, there exists at least one edge  $e$  such that  $f(e) > 0$  is reduced to 0. Hence, the algorithm will run for at most  $|E|$  iterations and terminate with  $f(e) = 0$  for all  $e \in E$ .
Thus we have that for any  $(s,s^{\prime})\in E^{(t)}$ $f(s,s^{\prime}) = \sum_{(\tau ,f_{\tau})\in Q:(s_{t}^{\tau},s_{t + 1}^{\tau}) = (s,s^{\prime})}f_{\tau}$