SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} \left. \left| \mathbb {E} _ {z \sim P _ {1}} [ m (z) ] - \mathbb {E} _ {z \sim P _ {2}} [ m (z) ] \right| \right. \\ = \left| \int m (z) \left(p _ {1} (z) - p _ {2} (z)\right) \mathrm {d} z \right| \\ \leq \left| \int_ {\| z \| _ {2} \leq R} \| z \| _ {2} (p _ {1} (z) - p _ {2} (z)) \mathrm {d} z \right| + \int_ {\| z \| _ {2} > R} \| z \| _ {2} p _ {1} (z) \mathrm {d} z + \int_ {\| z \| _ {2} > R} \| z \| _ {2} p _ {2} (z) \mathrm {d} z \\ \leq R \mathrm {d} _ {\mathrm {T V}} \left(P _ {1}, P _ {2}\right) + \sqrt {\int_ {\| z \| _ {2} > R} \| z \| _ {2} ^ {2} p _ {1} (z) \mathrm {d} z} + \sqrt {\int_ {\| z \| _ {2} > R} \| z \| _ {2} ^ {2} p _ {2} (z) \mathrm {d} z}, \\ \end{array}
$$
where  $\int_{\|z\|_2 > R} \|z\|_2^2 p_1(z) \, \mathrm{d}z$  and  $\int_{\|z\|_2 > R} \|z\|_2^2 p_2(z) \, \mathrm{d}z$  are bounded by some constant  $C_1$  when  $R^2 \geq C_2 \max \{\mathbb{E}_{z \sim P_1}[\|z\|_2^2], \mathbb{E}_{z \sim P_2}[\|z\|_2^2]\}$  as suggested by Lemma 16 in [8].
$$
\begin{array}{l} | \mathbb {E} _ {z \sim P _ {1}} [ \| z \| _ {2} ] - \mathbb {E} _ {z \sim P _ {2}} [ \| z \| _ {2} ] | \leq \sqrt {C _ {2} \max  \{\mathbb {E} _ {P _ {1}} [ \| z \| _ {2} ^ {2} ] , \mathbb {E} _ {P _ {2}} [ \| z \| _ {2} ^ {2} ] \}} \cdot \mathrm {d} _ {\mathrm {T V}} (P _ {1}, P _ {2}) + 2 C _ {1} \\ \leq \left(\sqrt {C _ {2} \mathbb {E} _ {z \sim P _ {1}} [ \| z \| _ {2} ^ {2} ]} + \sqrt {C _ {2} \mathbb {E} _ {z \sim P _ {2}} [ \| z \| _ {2} ^ {2} ]}\right) \cdot \mathrm {d} _ {\mathrm {T V}} \left(P _ {1}, P _ {2}\right) + 2 C _ {1} \\ = \mathcal {O} \left(\left(\sqrt {\mathbb {E} _ {z \sim P _ {1}} [ \| z \| _ {2} ^ {2} ]} + \sqrt {\mathbb {E} _ {z \sim P _ {2}} [ \| z \| _ {2} ^ {2} ]}\right) \cdot \mathrm {d} _ {\mathrm {T V}} \left(P _ {1}, P _ {2}\right)\right). \\ \end{array}
$$
E.8 Proof of Lemma E.3
Proof. Recall from (E.12) that
$$
P ^ {L D} (a) = P _ {z} (z \mid \widehat {f} (A z) = a) = \mathsf {N} (\mu (\mathsf {a}), \Gamma)
$$
with  $\mu(a) := \Sigma \widehat{\beta} \left( \widehat{\beta}^{\top} \Sigma \widehat{\beta} + \nu^{2} \right)^{-1} a, \Gamma := \Sigma - \Sigma \widehat{\beta} \left( \widehat{\beta}^{\top} \Sigma \widehat{\beta} + \nu^{2} \right)^{-1} \widehat{\beta}^{\top} \Sigma$ .
$$
\begin{array}{l} \mathbb {E} _ {z \sim P ^ {L D} (a)} \left[ \| z \| _ {2} ^ {2} \right] = \mu (a) ^ {\top} \mu (a) + \operatorname {t r a c e} (\Gamma) \\ = \frac {\widehat {\beta} ^ {\top} \Sigma^ {2} \widehat {\beta}}{\left(\| \widehat {\beta} \| _ {\Sigma} ^ {2} + \nu^ {2}\right) ^ {2}} a ^ {2} + \operatorname {t r a c e} (\Sigma - \Sigma \widehat {\beta} \left(\widehat {\beta} ^ {\top} \Sigma \widehat {\beta} + \nu^ {2}\right) ^ {- 1} \widehat {\beta} ^ {\top} \Sigma) \\ =: M (a). \\ \end{array}
$$
$$
\begin{array}{l} M (a) = \mathcal {O} \left(\frac {\widehat {\beta} ^ {\top} \Sigma^ {2} \widehat {\beta}}{\left(\| \widehat {\beta} \| _ {\Sigma} ^ {2}\right) ^ {2}} a ^ {2} + \operatorname {t r a c e} (\Sigma)\right) \\ = \mathcal {O} \left(\frac {a ^ {2}}{\| \widehat {\beta} \| _ {\Sigma}} + d\right), \\ \end{array}
$$
and by Lemma D.6
$$
\| \widehat {\beta} \| _ {\Sigma} \leq \frac {1}{2} \| \beta^ {*} \| _ {\Sigma}.
$$
Thus  $\mathbb{E}_{z\sim P^{LD}(a)}\left[\| z\| _2^2\right] = M(a),M(a) = \mathcal{O}\left(\frac{a^2}{\|\beta^*\|_\Sigma} +d\right).$
Thus after adding diffusion noise at  $t_0$ , we have for  $\alpha(t) = e^{-t/2}$  and  $h(t) = 1 - e^{-t}$ :
$$
\begin{array}{l} \mathbb {E} _ {z \sim P _ {t _ {0}} ^ {L D} (a)} \left[ \| z \| _ {2} ^ {2} \right] = \mathbb {E} _ {z _ {0} \sim P ^ {L D} (a)} \mathbb {E} _ {z \sim \mathsf {N} (\alpha (\mathbf {t} _ {0}) \cdot \mathbf {z} _ {0}, \mathsf {h} (\mathbf {t} _ {0}) \cdot \mathsf {l} _ {\mathrm {d}})} \left[ \| z \| _ {2} ^ {2} \right] \\ = \mathbb {E} _ {z _ {0} \sim P ^ {L D} (a)} \left[ \alpha^ {2} (t _ {0}) \| z _ {0} \| _ {2} ^ {2} + d \cdot h (t _ {0}) \right] \\ = \alpha^ {2} \left(t _ {0}\right) \cdot \mathbb {E} _ {z _ {0} \sim P ^ {L D} (a)} \left[ \| z _ {0} \| _ {2} ^ {2} \right] + d \cdot h (t _ {0}) \\ = e ^ {- t _ {0}} \cdot \mathbb {E} _ {z _ {0} \sim P ^ {L D} (a)} \left[ \| z _ {0} \| _ {2} ^ {2} \right] + \left(1 - e ^ {- t _ {0}}\right) \cdot d. \\ \end{array}
$$
Thus  $\mathbb{E}_{z\sim P_{t_0}^{LD}(a)}\left[\| z\| _2^2\right]\leq M(a) + t_0d.$
By orthogonal decomposition we have
$$
\begin{array}{l} \mathbb {E} _ {x \sim \widehat {P} _ {a}} \left[ \| x \| _ {2} ^ {2} \right] \leq \mathbb {E} _ {x \sim \widehat {P} _ {a}} \left[ \| \left(I _ {D} - V V ^ {\top}\right) x \| _ {2} ^ {2} \right] + \mathbb {E} _ {x \sim \widehat {P} _ {a}} \left[ \| V V ^ {\top} x \| _ {2} ^ {2} \right] \\ = \mathbb {E} _ {x \sim \widehat {P} _ {a}} \left[ \| (I _ {D} - V V ^ {\top}) x \| _ {2} ^ {2} \right] + \mathbb {E} _ {x \sim \widehat {P} _ {a}} \left[ \| U ^ {\top} V ^ {\top} x \| _ {2} ^ {2} \right], \\ \end{array}
$$
where  $\mathbb{E}_{x\sim \widehat{P}_a}\left[\| (I_D - VV^\top)x\| _2^2\right]$  is bounded by (D.5) and the distribution of  $U^{\top}V^{\top}x$  , which is  $(U^{\top}V^{\top})_{\#}\widehat{P}_a$  , is close to  $\mathbb{P}_{t_0}^{LD}(a)$  up to  $TV(\widehat{P}_a)$  , which is defined in Definition D.2. Then by Lemma E.2, we have
$$
\mathbb {E} _ {x \sim \widehat {P} _ {a}} \left[ \| U ^ {\top} V ^ {\top} x \| _ {2} ^ {2} \right] = \mathcal {O} \left(\mathbb {E} _ {z \sim P _ {t _ {0}} ^ {L D} (a)} \left[ \| z \| _ {2} ^ {2} \right] \left(1 + T V (\widehat {P} _ {a})\right)\right).
$$
Thus  $\mathbb{E}_{x\sim \widehat{P}_a}\left[\| x\| _2^2\right] = \mathcal{O}\left(ct_0D + (M(a) + t_0d)\cdot (1 + TV(\widehat{P}_a))\right).$
E.9 Proof of Lemma E.4
Proof. Firstly, one can verify the following two equations by direct calculation:
$$
\left(\lambda I _ {D} + A \Sigma_ {1} A ^ {\top}\right) ^ {- 1} = \frac {1}{\lambda} \left(I _ {D} - A \left(\lambda I _ {d} + \Sigma_ {1}\right) ^ {- 1} \Sigma_ {1} A ^ {\top}\right),
$$
$$
\left(\lambda I _ {d} + \Sigma_ {1}\right) ^ {- 1} = \frac {1}{\lambda} \left(I _ {d} - \left(\lambda I _ {d} + \Sigma_ {1}\right) ^ {- 1} \Sigma_ {1}\right).
$$
$$
\begin{array}{l} \left(\lambda I _ {D} + A \Sigma_ {1} A ^ {\top}\right) ^ {- 1} A \Sigma_ {2} A ^ {\top} = \frac {1}{\lambda} \left(I _ {D} - A \left(\lambda I _ {d} + \Sigma_ {1}\right) ^ {- 1} \Sigma_ {1} A ^ {\top}\right) A \Sigma_ {2} A ^ {\top} \\ = \frac {1}{\lambda} \left(A \Sigma_ {2} A ^ {\top} - A \left(\lambda I _ {d} + \Sigma_ {1}\right) ^ {- 1} \Sigma_ {1} \Sigma_ {2} A ^ {\top}\right). \\ \end{array}
$$
$$
\begin{array}{l} \operatorname {T r} \left(\left(\lambda I _ {D} + A \Sigma_ {1} A ^ {\top}\right) ^ {- 1} A \Sigma_ {2} A ^ {\top}\right) = \operatorname {T r} \left(\frac {1}{\lambda} \left(A \Sigma_ {2} A ^ {\top} - A \left(\lambda I _ {d} + \Sigma_ {1}\right) ^ {- 1} \Sigma_ {1} \Sigma_ {2} A ^ {\top}\right)\right) \\ = \operatorname {T r} \left(\frac {1}{\lambda} \left(\Sigma_ {2} - (\lambda I _ {d} + \Sigma_ {1}) ^ {- 1} \Sigma_ {1} \Sigma_ {2}\right)\right) \\ = \operatorname {T r} \left(\frac {1}{\lambda} \left(I _ {d} - \left(\lambda I _ {d} + \Sigma_ {1}\right) ^ {- 1} \Sigma_ {1}\right) \Sigma_ {2}\right) \\ = \operatorname {T r} \left(\left(\lambda I _ {d} + \Sigma_ {1}\right) ^ {- 1} \Sigma_ {2}\right), \\ \end{array}
$$
which has finished the proof.
F Theory in Nonparametric Setting
Built upon the insights from Section 4, we provide analysis to the nonparametric reward and general data sampling setting. We generalize Assumption 4.2 to the following.
Assumption F.1. The ground truth reward  $f^{*}$  is decomposed as
$$
f ^ {*} (x) = g ^ {*} (x _ {\parallel}) - h ^ {*} (x _ {\perp}),
$$
where  $g^{*}(x_{\parallel})$  is  $\alpha$ -Hölder continuous for  $\alpha \geq 1$  and  $h^*(x_\perp)$  is nondecreasing in terms of  $\| x_\perp \|_2$  with  $h^*(0) = 0$ . Moreover,  $g^*$  has a bounded Hölder norm, i.e.,  $\| g^* \|_{\mathcal{H}^\alpha} \leq 1$ .
Holder continuity is widely studied in nonparametric statistics literature [15, 49].  $h^*$  here penalizes off-support extrapolation.
Under Assumption F.1, we use nonparametric regression for estimating  $f^{*}$ . Specifically, we specialize (3.1) in Algorithm 1 to
$$
\widehat {f} _ {\theta} \in \operatorname * {a r g m i n} _ {f _ {\theta} \in \mathcal {F}} \frac {1}{2 n} \sum_ {i = 1} ^ {n _ {1}} \left(f _ {\theta} \left(x _ {i}\right) - y _ {i}\right) ^ {2},
$$
where  $\mathcal{F} = \mathrm{NN}(L,M,J,K,\kappa)$  is chosen to be a class of neural networks. Hyperparameters in  $\mathcal{F}$  will be chosen properly in Theorem F.4.
Our theory also considers generic sampling distributions on  $x$ . Since  $x$  lies in a low-dimensional subspace, this translates to a sampling distribution assumption on latent variable  $z$ .
Assumption F.2. The latent variable  $z$  follows distribution  $P_{z}$  with density  $p_{z}$ , such that there exists constants  $B, C_{1}, C_{2}$  verifying  $p_{z}(z) \leq (2\pi)^{-(d+1)/2} C_{1} \exp \left(-C_{2} \| z \|_{2}^{2}/2\right)$  whenever  $\| z \|_{2} > B$ . And  $c_{0} I_{d} \preceq \mathbb{E}_{z \sim P_{z}}[zz^{\top}]$ .
Assumption F.2 says  $P_{z}$  has a light tail, which is standard in high-dimensional statistics [51, 53]. Assumption F.2 also encodes distributions with a compact support. Furthermore, we assume that the curated data  $(x,\widehat{y})$  induces Lipschitz conditional scores. Motivated by Chen et al. [8], we show that the linear subspace structure in  $x$  leads to a similar conditional score decomposition  $\nabla \log p_t(x|\widehat{y}) = s_{\parallel}(x,\widehat{y},t) + s_{\perp}(x,\widehat{y},t)$ , where  $s_{\parallel}$  is the on-support score and  $s_{\perp}$  is the orthogonal score. The decomposition for conditional score is as (H.1), which applies to both parametric and non-parametric cases. The following assumption is imposed on  $s_{\parallel}$ .
Assumption F.3. The on-support conditional score function  $s_{\parallel}(x,\widehat{y},t)$  is Lipschitz with respect to  $x,\widehat{y}$  for any  $t\in (0,T]$ , i.e., there exists a constant  $C_\mathrm{lip}$ , such that for any  $x,\widehat{y}$  and  $x^{\prime},\widehat{y}^{\prime}$ , it holds
$$
\left\| s _ {\parallel} \left(x, \widehat {y}, t\right) - s _ {\parallel} \left(x ^ {\prime}, \widehat {y} ^ {\prime}, t\right) \right\| _ {2} \leq C _ {\mathrm {l i p}} \| x - x ^ {\prime} \| _ {2} + C _ {\mathrm {l i p}} | \widehat {y} - \widehat {y} ^ {\prime} | _ {2}.
$$
Lipschitz score is commonly adopted in existing works [9, 25]. Yet Assumption F.3 only requires the Lipschitz continuity of the on-support score, which matches the weak regularity conditions in Lee et al. [25], Chen et al. [8]. We then choose the score network architecture similar to that in the linear reward setting, except we replace  $m$  by a nonlinear network. Recall the linear encoder and decoder estimate the representation matrix  $A$ .
We consider feedforward networks with ReLU activation functions as concept classes  $\mathcal{F}$  and  $\mathcal{S}$  for nonparametric regression and conditional score matching. Generalization to different network architectures poses no real difficulty. Given an input  $x$ , neural networks compute
$$
f _ {\mathrm {N N}} (x) = W _ {L} \sigma \left(\dots \sigma \left(W _ {1} x + b _ {1}\right) \dots\right) + b _ {L}, \tag {F.1}
$$
where  $W_{i}$  and  $b_{i}$  are weight matrices and intercepts, respectively. We then define a class of neural networks as
$\mathrm{NN}(L,M,J,K,\kappa) = \left\{f:f\text{in the form of (F.1) with} L\text{ layers and width bounded by} M,\right.$
$$
\sup  _ {x} \| f (x) \| _ {2} \leq K, \max  \left\{\| b _ {i} \| _ {\infty}, \| W _ {i} \| _ {\infty} \right\} \leq \kappa \text {f o r} i = 1, \dots , L, \text {a n d} \sum_ {i = 1} ^ {L} \left(\| W _ {i} \| _ {0} + \| b _ {i} \| _ {0}\right) \leq J \Big \}.
$$
For the conditional score network, we will additionally impose some Lipschitz continuity requirement, i.e.,  $\| f(x) - f(y)\| _2\leq c_{\mathrm{lip}}\| x - y\| _2$  for some Lipschitz coefficient  $c_{\mathrm{lip}}$ .
Recall the distribution shift defined in Definition D.3 that
$$
\mathcal {T} (P _ {1}, P _ {2}; \mathcal {L}) = \sup _ {l \in \mathcal {L}} \mathbb {E} _ {x \sim P _ {1}} [ l (x) ] / \mathbb {E} _ {x \sim P _ {2}} [ l (x) ]
$$
for arbitrary two distributions  $P_{1}, P_{2}$  and function class  $\mathcal{L}$ . Similar to the parametric case, use notation  $\widehat{P}_a \coloneqq \widehat{P}(\cdot|\widehat{y} = a)$  and  $P_a \coloneqq P(\cdot|\widehat{y} = a)$ . Then we can bound  $\operatorname{SubOpt}(\widehat{P}_a; y^* = a)$  in Theorem F.4 in terms of non-parametric regression error, score matching error and distribution shifts in both regression and score matching.
Theorem F.4. Suppose Assumption 4.1, F.1, F.2 and F.3 hold. Let  $\delta(n) = \frac{d\log\log n}{\log n}$ . Properly chosen  $\mathcal{F}$  and  $\mathcal{S}$ , with high probability, running Algorithm 1 with a target reward value  $a$  and stopping at  $t_0 = \left(n_1^{-\frac{2 - 2\delta(n_1)}{d + 6}} + Dn_1^{-\frac{d + 4}{d + 6}}\right)^{\frac{1}{3}}$  gives rise to  $\angle(V,A) \leq \widetilde{\mathcal{O}}\left(\frac{1}{c_0}\left(n_1^{-\frac{2 - 2\delta(n_1)}{d + 6}} + Dn_1^{-\frac{d + 4}{d + 6}}\right)\right)$  and
SubOpt  $(\widehat{P}_a;y^* = a)$
$$
\begin{array}{l} \leq \underbrace {\sqrt {\mathcal {T} (P (x | \widehat {y} = a) , P _ {x} ; \bar {\mathcal {F}})} \cdot \widetilde {\mathcal {O}} \left(n _ {2} ^ {- \frac {\alpha - \delta (n _ {2})}{2 \alpha + d}} + D / n _ {2}\right)} _ {\mathcal {E} _ {1}} \\ + \underbrace {\left(\sqrt {\frac {\mathcal {T} (P (x , \widehat {y} = a) , P _ {x \widehat {y}} ; \bar {\mathcal {S}})}{c _ {0}}} \cdot \| g ^ {*} \| _ {\infty} + \sqrt {M (a)}\right) \cdot \widetilde {\mathcal {O}} \left(\left(n _ {1} ^ {- \frac {2 - 2 \delta (n _ {1})}{d + 6}} + D n _ {1} ^ {- \frac {d + 4}{d + 6}}\right) ^ {\frac {1}{3}}\right)} _ {\mathcal {E} _ {2}} \\ + \underbrace {\mathbb {E} _ {x \sim \widehat {P} _ {a}} [ h ^ {*} (x _ {\perp}) ]} _ {\mathcal {E} _ {3}}, \\ \end{array}
$$
where  $M(a)\coloneqq \mathbb{E}_{z\sim \mathbb{P}(a)}[\| z\| _2^2 ]$  and
$$
\bar {\mathcal {F}} := \left\{| f ^ {*} (x) - f (x) | ^ {2}: f \in \mathcal {F} \right\}, \quad \bar {\mathcal {S}} = \left\{\frac {1}{T - t _ {0}} \int_ {t _ {0}} ^ {T} \mathbb {E} _ {x _ {t} | x} \| \nabla \log p _ {t} (x _ {t} \mid y) - s (x _ {t}, y, t) \| _ {2} ^ {2} \mathrm {d} t: s \in \mathcal {S} \right\},
$$
$\mathcal{E}_3$  penalizes the component in  $\widehat{P}_a$  that is off the truth subspace. The function classes  $\mathcal{F}$  and  $\mathcal{S}$  are chosen as  $\mathcal{F} = \mathrm{NN}(L_f, M_f, J_f, K_f, \kappa_f)$  with
$$
\begin{array}{l} L _ {f} = \mathcal {O} (\log n _ {2}), M _ {f} = \mathcal {O} \left(n _ {2} ^ {- \frac {d}{d + 2 \alpha}} (\log n _ {2}) ^ {d / 2}\right), J _ {f} = \mathcal {O} \left(n _ {2} ^ {- \frac {d}{d + 2 \alpha}} (\log n _ {2}) ^ {d / 2 + 1}\right) \\ K _ {f} = 1, \kappa_ {f} = \mathcal {O} \left(\sqrt {\log n _ {2}}\right) \\ \end{array}
$$
and  $\mathcal{S} = \mathrm{NN}(L_s,M_s,J_s,K_s,\kappa_s)$  with
$$
\begin{array}{l} L _ {s} = \mathcal {O} (\log n _ {1} + d), M _ {s} = \mathcal {O} \left(d ^ {d / 2} n _ {1} ^ {- \frac {d + 2}{d + 6}} (\log n _ {1}) ^ {d / 2}\right), J _ {s} = \mathcal {O} \left(d ^ {d / 2} n _ {1} ^ {- \frac {d + 2}{d + 6}} (\log n _ {1}) ^ {d / 2 + 1}\right) \\ K _ {s} = \mathcal {O} \left(d \log (d n _ {1})\right), \kappa_ {s} = \mathcal {O} \left(\sqrt {d \log (n _ {1} d)}\right). \\ \end{array}
$$
Moreover,  $S$  is also Lipschitz with respect to  $(x,y)$  and the Lipschitz coefficient is  $c_{\mathrm{lip}} = \mathcal{O}(10dC_{\mathrm{lip}})$ .
Remark. The proof is provided in Appendix G.2. Here we correct a typo on  $n_1^{-\frac{1}{3(d + 5)}}$  in main paper Section 4.3. Quantities  $\mathcal{T}(P(x|\widehat{y} = a), P_x; \bar{\mathcal{F}})$  and  $\mathcal{T}(P(x, \widehat{y} = a), P_{x\widehat{y}}; \bar{\mathcal{S}})$  depend on  $a$  characterizing the distribution shift. The  $\delta(n)$  terms account for the unbounded domain of  $x$ , which is negligible when  $n$  is large. In the main paper, we omit  $\delta(n)$  in the regret bound.
G Omitted Proofs in Section F
G.1 Conditional Score Decomposition and Score Matching Error
Lemma G.1. Under Assumption 4.1, F.2 and F.3, with high probability
$$
\frac {1}{T - t _ {0}} \int_ {t _ {0}} ^ {T} \| \widehat {s} (\cdot , t) - \nabla \log p _ {t} (\cdot) \| _ {L ^ {2} (P _ {t})} ^ {2} \mathrm {d} t \leq \epsilon_ {\text {d i f f}} ^ {2} (n _ {1})
$$
with  $\epsilon_{diff}^2 (n_1) = \widetilde{\mathcal{O}}\left(\frac{1}{t_0}\left(n_1^{-\frac{2 - 2\delta(n_1)}{d + 6}} + Dn_1^{-\frac{d + 4}{d + 6}}\right)\right)$  for  $\delta (n_{1}) = \frac{d\log\log n_{1}}{\log n_{1}}$
Proof. [8, Theorem 1] is easily adapted here to prove Lemma G.1 with the input dimension  $d + 1$  and the Lipschitzness in Assumption F.3. Network size of  $S$  is implied by [8, Theorem 1] with  $\epsilon = n_1^{-\frac{1}{d + 6}}$  accounting for the additional dimension of reward  $\widehat{y}$  and then the score matching error follows.
G.2 Proof of Theorem F.4
Additional Notations: Similar as before, use  $P_{t}^{LD}(z)$  to denote the low-dimensional distribution on  $z$  corrupted by diffusion noise. Formally,  $p_{t}^{LD}(z) = \int \phi_{t}(z'|z)p_{z}(z)\mathrm{d}z$  with  $\phi_t(\cdot |z)$  being the density of  $\mathsf{N}(\alpha (t)z,h(t)I_d)$ .  $P_{t_0}^{LD}(z\mid \widehat{f} (Az) = a)$  the corresponding conditional distribution on  $\widehat{f} (Az) = a$  at  $t_0$ , with shorthand as  $P_{t_0}^{LD}(a)$ . Also give  $P_{z}(z\mid \widehat{f} (Az) = a)$  a shorthand as  $P^{LD}(a)$ .
G.2.1 SubOpt  $(\widehat{P}_a; y^* = a)$  Decomposition
By the same argument as in §D.3.1, we have
$$
\begin{array}{l} \operatorname {S u b 0 p t} \left(\widehat {P} _ {a}; y ^ {*} = a\right) \leq \underbrace {\mathbb {E} _ {x \sim P _ {a}} \left[ \left| f ^ {*} (x) - \widehat {f} (x) \right| \right]} _ {\mathcal {E} _ {1}} + \underbrace {\left| \mathbb {E} _ {x \sim P _ {a}} \left[ g ^ {*} \left(x _ {\parallel}\right) \right] - \mathbb {E} _ {x \sim \widehat {P} _ {a}} \left[ g ^ {*} \left(x _ {\parallel}\right) \right] \right|} _ {\mathcal {E} _ {2}} \\ + \underbrace {\mathbb {E} _ {x \sim \widehat {P} _ {a}} [ h ^ {*} (x _ {\perp}) ]} _ {\mathcal {E} _ {3}}. \\ \end{array}
$$
G.2.2  $\mathcal{E}_1$ : Nonparametric Regression Induced Error
Nonparametric Regression Error of  $\widehat{f}$  Since  $P_{z}$  has a light tail due to Assumption F.2, by union bound and [8, Lemma 16], we have
$$
\mathbb {P} \left(\exists x _ {i} \text {w i t h} \| x _ {i} \| _ {2} > R \text {f o r} i = 1, \dots , n _ {2}\right) \leq n _ {2} \frac {C _ {1} d 2 ^ {- d / 2 + 1}}{C _ {2} \Gamma (d / 2 + 1)} R ^ {d - 2} \exp \left(- C _ {2} R ^ {2} / 2\right),
$$
where  $C_1, C_2$  are constants and  $\Gamma(\cdot)$  is the Gamma function. Choosing  $R = \mathcal{O}(\sqrt{d\log d + \log\frac{n}{\delta}})$  ensures  $\mathbb{P}(\exists x_i$  with  $\|x_i\|_2 > R$  for  $i = 1, \ldots, n_2) < \delta$ . On the event  $\mathcal{E} = \{\|x_i\|_2 \leq R$  for all  $i = 1, \ldots, n_2\}$ , denoting  $\delta(n_2) = \frac{d\log\log n_2}{\log n_2}$ , we have
$$
\left\| f ^ {*} - \widehat {f} \right\| _ {L ^ {2}} ^ {2} = \widetilde {\mathcal {O}} \left(n _ {2} ^ {- \frac {2 (\alpha - \delta (n _ {2}))}{d + 2 \alpha}}\right)
$$
by [31, Theorem 7] with a new covering number of  $S$ , when  $n_2$  is sufficiently large. The corresponding network architecture follows from Theorem 2 in "Nonparametric Regression on Low-Dimensional Manifolds using Deep ReLU Networks: Function Approximation and Statistical Recovery".
We remark that linear subspace is a special case of low Minkowski dimension. Moreover,  $\delta (n_{2})$  is asymptotically negligible and accounts for the truncation radius  $R$  of  $x_{i}$ 's (see also [8, Theorem 2 and 3]). The covering number of  $\mathcal{S}$  is  $\widetilde{\mathcal{O}}\left(d^{d / 2}n^{-\frac{d}{\alpha}}(\log n_2)^{d / 2} + Dd\right)$  as appear in [8, Proof of Theorem
2]. Therefore
$$
\begin{array}{l} \mathbb {E} _ {x \sim P _ {a}} \left[ \left| f ^ {*} (x) - \widehat {f} (x) \right| \right] \leq \sqrt {\mathbb {E} _ {x \sim P _ {a}} \left[ \left| f ^ {*} (x) - \widehat {f} (x) \right| ^ {2} \right]} \\ \leq \sqrt {\mathcal {T} (P (x | \widehat {y} = a) , P _ {x} ; \bar {\mathcal {F}}) \cdot \| f ^ {*} - \widehat {f} \| _ {L ^ {2}} ^ {2}} \\ = \sqrt {\mathcal {T} (P (x | \widehat {y} = a) , P _ {x} ; \bar {\mathcal {F}})} \cdot \widetilde {\mathcal {O}} \left(n _ {2} ^ {- \frac {\alpha - \delta (n _ {2})}{2 \alpha + d}} + D / n _ {2}\right). \\ \end{array}
$$
G.2.3  $\varepsilon_{2}$ : Diffusion Induced On-support Error
Suppose  $L_{2}$  score matching error is  $\epsilon_{diff}^{2}(n_{1})$ , i.e.
$$
\frac {1}{T - t _ {0}} \int_ {t _ {0}} ^ {T} \mathbb {E} _ {x, \widehat {f}} \| \nabla_ {x} \log p _ {t} (x, \widehat {f}) - s _ {\widehat {w}} (x, \widehat {f}, t) \| _ {2} ^ {2} \mathrm {d} t \leq \epsilon_ {d i f f} ^ {2} (n _ {1}),
$$
We revoke Definition D.2 measuring the distance between  $\widehat{P}_a$  to  $P_a$  that
$$
T V (\widehat {P} _ {a}) := \mathrm {d} _ {\mathrm {T V}} \left(P _ {t _ {0}} ^ {L D} (z \mid \widehat {f} (A z) = a), (U ^ {\top} V ^ {\top}) _ {\#} \widehat {P} _ {a}\right).
$$
Lemma D.4 applies to nonparametric setting, so we have
$$
(I _ {D} - V V ^ {\top}) x \sim \mathbb {N} (0, \Lambda), \quad \Lambda \prec c t _ {0} I _ {D}, \tag {G.1}
$$
$$
\angle (V, A) = \widetilde {\mathcal {O}} \left(\frac {t _ {0}}{c _ {0}} \cdot \epsilon_ {d i f f} ^ {2} \left(n _ {1}\right)\right). \tag {G.2}
$$
$$
T V \left(\widehat {P} _ {a}\right) = \widetilde {\mathcal {O}} \left(\sqrt {\frac {\mathcal {T} \left(P \left(x , \widehat {y} = a\right) , P _ {x \widehat {y}} ; \bar {\mathcal {S}}\right)}{c _ {0}}} \cdot \epsilon_ {d i f f} \left(n _ {1}\right)\right). \tag {G.3}
$$
$\mathcal{E}_2$  will be bounded by
$$
\begin{array}{l} \mathcal {E} _ {2} = \left| \mathbb {E} _ {x \sim P _ {a}} [ g ^ {*} (x) ] - \mathbb {E} _ {x \sim \widehat {P} _ {a}} [ g ^ {*} (x) ] \right| \\ \leq \left| \mathbb {E} _ {x \sim P _ {a}} \left[ g ^ {*} \left(A A ^ {\top} x\right) \right] - \mathbb {E} _ {x \sim \widehat {P} _ {a}} \left[ g ^ {*} \left(V V ^ {\top} x\right) \right] \right| + \left| \mathbb {E} _ {x \sim \widehat {P} _ {a}} \left[ g ^ {*} \left(V V ^ {\top} x\right) - g ^ {*} \left(A A ^ {\top} x\right) \right] \right|, \\ \end{array}
$$
where for  $\left|\mathbb{E}_{x\sim \widehat{P}_a}\big[g^* (VV^\top x) - g^* (AA^\top x)\big]\right|$ , we have
$$
\left| \mathbb {E} _ {x \sim \widehat {P} _ {a}} \left[ g ^ {*} (V V ^ {\top} x) - g ^ {*} (A A ^ {\top} x) \right] \right| \leq \mathbb {E} _ {x \sim \widehat {P} _ {a}} \left[ \| V V ^ {\top} x - A A ^ {\top} x \| _ {2} \right] \leq \| V V ^ {\top} - A A ^ {\top} \| _ {F} \cdot \mathbb {E} _ {x \sim \widehat {P} _ {a}} [ \| x \| _ {2} ]. \tag {G.4}
$$
For the other term  $\left|\mathbb{E}_{x\sim P_a}[g^* (AA^\top x)] - \mathbb{E}_{x\sim \widehat{P}_a}[g^* (VV^\top x)]\right|$ , we will bound it with  $TV(\widehat{P}_a)$ .
$$
\begin{array}{l} \left| \mathbb {E} _ {x \sim P _ {a}} \left[ g ^ {*} \left(A A ^ {\top} x\right) \right] - \mathbb {E} _ {x \sim \hat {P} _ {a}} \left[ g ^ {*} \left(V V ^ {\top} x\right) \right] \right| \\ \leq \left| \mathbb {E} _ {z \sim \mathbb {P} _ {t _ {0}} (a)} [ g ^ {*} (A z) ] - \mathbb {E} _ {z \sim (V ^ {\top}) _ {\#} \widehat {P} _ {a}} [ g ^ {*} (V z) ] \right| + \left| \mathbb {E} _ {z \sim \mathbb {P} (a)} [ g ^ {*} (A z) ] - \mathbb {E} _ {z \sim \mathbb {P} _ {t _ {0}} (a)} [ g ^ {*} (A z) ] \right| \\ \end{array}
$$
Since any  $z\sim \mathbb{P}_{t_0}(a)$  can be represented by  $\alpha (t_0)z + \sqrt{h(t_0)} u$  , where  $z\sim \mathbb{P}(a),u\sim \mathsf{N}(0,I_d)$  , then