SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Laixi Shi, Gen Li, Yuting Wei, Yuxin Chen, and Yuejie Chi. Pessimistic Q-learning for offline reinforcement learning: Towards optimal sample complexity. arXiv preprint arXiv:2202.13890, 2022.
Csaba Szepesvári and Rémi Munos. Finite time bounds for sampling based fitted value iteration. In Proceedings of the 22nd international conference on Machine learning, pages 880-887, 2005.
Masatoshi Uehara and Wen Sun. Pessimistic model-based offline reinforcement learning under partial coverage. In International Conference on Learning Representations, 2021.
Masatoshi Uehara, Jiawei Huang, and Nan Jiang. Minimax weight and Q-function learning for off-policy evaluation. In International Conference on Machine Learning, pages 9659-9668. PMLR, 2020.
Sara Van de Geer. Empirical Processes in M-estimation, volume 6. Cambridge university press, 2000.
Lingxiao Wang, Qi Cai, Zhuoran Yang, and Zhaoran Wang. Neural policy gradient methods: Global optimality and rates of convergence. In International Conference on Learning Representations, 2019.
Tengyang Xie and Nan Jiang. Batch value-function approximation with only realizability. In International Conference on Machine Learning, pages 11404-11413. PMLR, 2021a.
Tengyang Xie and Nan Jiang. Batch value-function approximation with only realizability. In International Conference on Machine Learning, pages 11404-11413. PMLR, 2021b.
Tengyang Xie, Ching-An Cheng, Nan Jiang, Paul Mineiro, and Alekh Agarwal. Bellman-consistent pessimism for offline reinforcement learning. Advances in neural information processing systems, 34:6683-6694, 2021.
Rui Yang, Yiming Lu, Wenzhe Li, Hao Sun, Meng Fang, Yali Du, Xiu Li, Lei Han, and Chongjie Zhang. Rethinking goal-conditioned supervised learning and its connection to offline rl. arXiv preprint arXiv:2202.04478, 2022.
Ming Yin and Yu-Xiang Wang. Towards instance-optimal offline reinforcement learning with pessimism. Advances in neural information processing systems, 34:4065-4078, 2021.
Ming Yin, Yu Bai, and Yu-Xiang Wang. Near-optimal offline reinforcement learning via double variance reduction. arXiv preprint arXiv:2102.01748, 2021.
Ming Yin, Wenjing Chen, Mengdi Wang, and Yu-Xiang Wang. Offline stochastic shortest path: Learning, evaluation and towards optimality. In Uncertainty in Artificial Intelligence, pages 2278-2288. PMLR, 2022.
Wenhao Zhan, Baihe Huang, Audrey Huang, Nan Jiang, and Jason Lee. Offline reinforcement learning with realizability and single-policy concentrability. In Conference on Learning Theory, pages 2730-2775. PMLR, 2022.
Junyu Zhang, Alec Koppel, Amrit Singh Bedi, Csaba Szepesvari, and Mengdi Wang. Variational policy gradient method for reinforcement learning with general utilities. arXiv preprint arXiv:2007.02151, 2020.
A Missing Proofs for Propositions in Section 3
Proof. Since both  $d_{\alpha}^{*}$  and  $d^{*}$  satisfy the Bellman flow constraint (1), due to the optimality of  $d_{\alpha}^{*}$  in the regularized program (5), we have
$$
\begin{array}{l} \mathbb {E} _ {(s, g) \sim d ^ {*} (s, g)} [ r (s; g) ] + \alpha D _ {f} (d ^ {*} \| \mu) \leq \mathbb {E} _ {(s, g) \sim d _ {\alpha} ^ {*} (s, g)} [ r (s; g) ] + \alpha D _ {f} (d _ {\alpha} ^ {*} \| \mu) \\ \Longrightarrow \mathbb {E} _ {(s, g) \sim d ^ {*} (s, g)} [ r (s; g) ] - \mathbb {E} _ {(s, g) \sim d _ {\alpha} ^ {*} (s, g)} [ r (s; g) ] \leq \alpha D _ {f} \left(d _ {\alpha} ^ {*} \| \mu\right) \leq \alpha \left(C _ {\alpha} ^ {*}\right) ^ {2} / 2. \\ \end{array}
$$
$$
J (\pi^ {*}) - J (\pi_ {\alpha} ^ {*}) = \mathbb {E} _ {(s, g) \sim d ^ {*} (s, g)} [ r (s; g) ] - \mathbb {E} _ {(s, g) \sim d _ {\alpha} ^ {*} (s, g)} [ r (s; g) ] \leq O \left(\alpha (C _ {\alpha} ^ {*}) ^ {2}\right).
$$
Proof. By Lagrangian duality, we can obtain that the dual problem is
$$
\begin{array}{l} \min  _ {V (s; g) \geq 0} \max  _ {d (s, a; g) \geq 0} \mathbb {E} _ {(s, g) \sim d (s, g)} [ r (s; g) ] - \alpha D _ {f} (d (s, a; g) \| \mu (s, a; g)) \\ + \sum_ {s, g} p (g) V (s; g) \left((1 - \gamma) \rho (s) + \gamma \sum_ {s ^ {\prime}, a ^ {\prime}} P \left(s \mid s ^ {\prime}, a ^ {\prime}\right) d \left(s ^ {\prime}, a ^ {\prime}; g\right) - \sum_ {a} d (s, a; g)\right) \tag {15} \\ \end{array}
$$
where  $p(g)V(s;g)$  is the Lagrangian vector. Similar to Ma et al. [2022c],
$$
\begin{array}{l} \sum_ {s, g} p (g) V (s; g) \sum_ {s ^ {\prime}, a ^ {\prime}} P (s | s ^ {\prime}, a ^ {\prime}) d (s ^ {\prime}, a ^ {\prime}; g) \\ = \sum_ {s ^ {\prime}, a ^ {\prime}, g} p (g) d \left(s ^ {\prime}, a ^ {\prime}; g\right) \sum_ {s} P \left(s \mid s ^ {\prime}, a ^ {\prime}\right) V (s; g) \\ = \sum_ {s ^ {\prime}, a ^ {\prime}, g} p (g) d \left(s ^ {\prime}, a ^ {\prime}; g\right) \mathcal {T} V \left(s ^ {\prime}, a ^ {\prime}; g\right). \\ \end{array}
$$
Then (15) is equivalent to
$$
\begin{array}{l} \min  _ {V (s; g) \geq 0} \max  _ {d (s, a; g) \geq 0} (1 - \gamma) \mathbb {E} _ {(s, g) \sim (\rho , p (g))} [ V (s; g) ] + \mathbb {E} _ {(s, a, g) \sim d} [ r (s; g) + \gamma \mathcal {T} V (s, a; g) - V (s; g) ] \\ - \alpha D _ {f} (d (s, a; g) \| \mu (s, a; g)), \\ \end{array}
$$
which can be further represented as
$$
\begin{array}{l} \min  _ {V (s; g) \geq 0} (1 - \gamma) \mathbb {E} _ {(s, g) \sim (\rho , p (g))} [ V (s; g) ] + \max  _ {d (s, a; g) \geq 0} \mathbb {E} _ {(s, a, g) \sim d} [ r (s; g) + \gamma \mathcal {T} V (s, a; g) - V (s; g) ] \\ - D _ {g} (d (s, a; g) \| \mu (s, a; g)), \\ \end{array}
$$
where  $g = \alpha \cdot f$ . Combining the constraint that  $d(s, a; g) \geq 0$  and the proof of Proposition 4.3 (Proposition B.2) of Ma et al. [2022c], the above program is equivalent to
$$
\min _ {V (s; g) \geq 0} (1 - \gamma) \mathbb {E} _ {(s, g) \sim (\rho , p (g))} [ V (s; g) ] + \mathbb {E} _ {(s, a, g) \sim \mu} [ \mathbb {1} \{g _ {*} ^ {\prime} (A _ {V} (s, a; g)) \geq 0 \} \bar {g} _ {*} (A _ {V} (s, a; g)) ],
$$
$$
d _ {\alpha} ^ {*} (s, a; g) = \mu (s, a; g) g _ {*} ^ {\prime} (r (s; g) + \gamma \mathcal {T} V _ {\alpha} ^ {*} (s, a; g) - V _ {\alpha} ^ {*} (s; g)) _ {+}.
$$
B Proof for  $V$ -Learning
In this section, we provide theoretical analysis for the guarantees of  $V$ -learning in different settings (deterministic dynamics in Appendix B.1 and stochastic dynamics in Appendix B.2).
Recall that we choose  $f(x) = \frac{(x - 1)^2}{2}$ , and thus  $f_{*}(x) = \frac{(x + 1)^{2}}{2} - \frac{1}{2}$ . Therefore, we have  $g(x) = \alpha f(x) = \frac{\alpha(x - 1)^2}{2}$  and  $g_{*}(x) = \frac{\alpha(x / \alpha + 1)^{2}}{2} - \frac{\alpha}{2}$ . Also, we define  $g_{*,\max} = \max_{v \in [-V_{\max}, V_{\max} + 1]} g_{*}(v)$ ,  $g_{*,\min} = \min_{v \in [-V_{\max}, V_{\max} + 1]} g_{*}(v)$  and  $g_{*,\Delta} = g_{*,\max} - g_{*,\min}$ . We can easily obtain that  $g_{*,\Delta} \leq \alpha (1 + \frac{V_{\max}}{\alpha})^{2} = O(V_{\max}^{2} / \alpha)$ .
For the purpose of theoretical analysis, we further define
$$
L _ {1} (V) = (1 - \gamma) \mathbb {E} _ {(s, g) \sim (\rho , p (g))} [ \alpha \cdot V (s; g) ],
$$
$$
L _ {2} (V) = \mathbb {E} _ {(s, a, g) \sim \mu} [ \alpha \cdot g _ {* +} (r (s; g) + \gamma \mathcal {T} V (s, a; g) - V (s; g)) ],
$$
$$
\hat {L} _ {1} ^ {(d)} (V) = \frac {1 - \gamma}{N _ {0}} \sum_ {i = 1} ^ {N _ {0}} \alpha \cdot V (s _ {0, i}; g _ {0, i}),
$$
$$
\hat {L} _ {2} ^ {(d)} (V) = \frac {1}{N} \sum_ {i = 1} ^ {N} \alpha \cdot g _ {* +} (r (s _ {i}; g _ {i}) + \gamma V (s _ {i} ^ {\prime}; g _ {i}) - V (s _ {i}; g _ {i})).
$$
Note that  $L_{\alpha}(V) = L_1(V) + L_2(V)$  and  $\hat{L}^{(d)}(V) = \hat{L}_1^{(d)}(V) + \hat{L}_2^{(d)}(V)$ .
B.1 Proof for Deterministic Dynamics
We first show that with high probability, the estimator  $\hat{L}^{(d)}(V)$  in Algorithm 2 concentrates well on  $L_{\alpha}(V)$  for all  $V \in \mathcal{V}$ .
Lemma 3 (Concentration of  $\hat{L}_1^{(d)},\hat{L}_2^{(d)}$ ). Under Assumptions 1 and 2, when the dynamic of the environment is deterministic, with probability at least  $1 - \delta$ ,  $\forall V\in \mathcal{V}$ , it simultaneously holds that
$$
\left| \hat {L} _ {1} ^ {(d)} (V) - L _ {1} (V) \right| \leq O \left(\alpha V _ {\max } \sqrt {\frac {\log (| \mathcal {V} | / \delta)}{N _ {0}}}\right),
$$
$$
\left| \hat {L} _ {2} ^ {(d)} (V) - L _ {2} (V) \right| \leq O \left(\alpha \cdot g _ {*, \Delta} \sqrt {\frac {\log (| \mathcal {V} | / \delta)}{N}}\right) = O \left(V _ {\max} ^ {2} \sqrt {\frac {\log (| \mathcal {V} | / \delta)}{N}}\right),
$$
which immediately implies
$$
\left| \hat {L} ^ {(d)} (V) - L _ {\alpha} (V) \right| \leq O \left(\alpha V _ {\max } \sqrt {\frac {\log (| \mathcal {V} | / \delta)}{N _ {0}}} + V _ {\max } ^ {2} \sqrt {\frac {\log (| \mathcal {V} | / \delta)}{N}}\right) \triangleq \epsilon_ {s t a t}.
$$
Proof. First, fix any  $V \in \mathcal{V}$ . The expectation of  $\hat{L}_1^{(d)}(V)$  is
$$
\begin{array}{l} \mathbb {E} [ \hat {L} _ {1} ^ {(d)} (V) ] = \frac {1 - \gamma}{N _ {0}} \sum_ {i = 1} ^ {N _ {0}} \mathbb {E} _ {(s _ {0, i}, g _ {0, i}) \sim (\rho , p (g))} [ \alpha \cdot V (s _ {0, i}; g _ {0, i}) ] \\ = (1 - \gamma) \mathbb {E} _ {(s, g) \sim (\rho , p (g))} [ \alpha \cdot V (s; g) ] \\ = L _ {1} (V). \\ \end{array}
$$
$$
\begin{array}{l} \mathbb {E} [ \hat {L} _ {2} ^ {(d)} (V) ] = \frac {1}{N} \sum_ {i = 1} ^ {N} \mathbb {E} _ {(s _ {i}, a _ {i}, g _ {i}) \sim \mu} [ \alpha \cdot g _ {* +} (r (s _ {i}; g _ {i}) + \gamma V (s _ {i} ^ {\prime}; g _ {i}) - V (s _ {i}; g _ {i})) ] \\ = \frac {1}{N} \sum_ {i = 1} ^ {N} \mathbb {E} _ {\left(s _ {i}, a _ {i}, g _ {i}\right) \sim \mu} \left[ \alpha \cdot g _ {* +} \left(r \left(s _ {i}; g _ {i}\right) + \gamma \mathcal {T} V \left(s _ {i}, a _ {i}; g _ {i}\right) - V \left(s _ {i}; g _ {i}\right)\right) \right] \\ = L _ {2} (V). \\ \end{array}
$$
Also note that  $V(s_{0,i};g_{0,i}) \in [0,V_{\max}]$  and  $g_{* + }(r(s_i;g_i) + \gamma TV(s_i,a_i;g_i) - V(s_i;g_i)) \in [0,g_{*,\Delta}]$ . By Hoeffding's inequality and a union bound, we have with probability at least  $1 - \delta$ ,
$$
\left| \hat {L} _ {1} ^ {(d)} (V) - L _ {1} (V) \right| \leq O \left(\alpha V _ {\max } \sqrt {\frac {\log (1 / \delta)}{N _ {0}}}\right),
$$
$$
\left| \hat {L} _ {2} ^ {(d)} (V) - L _ {2} (V) \right| \leq O \left(\alpha \cdot g _ {*, \Delta} \sqrt {\frac {\log (1 / \delta)}{N}}\right).
$$
Applying a union bound over all  $V \in \mathcal{V}$  concludes the result.
Next, we show that the value of  $L_{\alpha}$  at the regularized optimal  $V$ -function  $V_{\alpha}^{*}$  and the learned function  $\hat{V}$  is close.
Lemma 4 (Closeness of the population objective of  $\hat{V}$  and  $V_{\alpha}^{*}$ ). Under Assumptions 1 and 2, with probability at least  $1 - \delta$ , we have
$$
L _ {\alpha} (\hat {V}) - L _ {\alpha} (V _ {\alpha} ^ {*}) \leq O (\epsilon_ {s t a t}).
$$
Proof. We condition on the high probability event in Lemma 3. Note that
$$
L _ {\alpha} (\hat {V}) - L _ {\alpha} (V _ {\alpha} ^ {*}) = \underbrace {L _ {\alpha} (\hat {V}) - \hat {L} ^ {(d)} (\hat {V})} _ {(1)} + \underbrace {\hat {L} ^ {(d)} (\hat {V}) - \hat {L} ^ {(d)} (V _ {\alpha} ^ {*})} _ {(2)} + \underbrace {\hat {L} ^ {(d)} (V _ {\alpha} ^ {*}) - L _ {\alpha} (V _ {\alpha} ^ {*})} _ {(3)}.
$$
(1), (3)  $\leq O(\epsilon_{\mathrm{stat}})$  by Lemma 3 and (2)  $\leq 0$  by the definition of  $\hat{V}$ , which completes the proof.
Now we define "semi-strong" convexity, which is used to prove Lemma 1 and might be of independent interests.
Definition 2 (Positive semi-strong convexity). For an arbitrary set  $\mathcal{X}$ , assume  $\mathcal{F} \subset \{f \mid f : \mathcal{X} \to \mathbb{R}\}$  is a convex set. Let  $\rho \in \Delta(\mathcal{X})$  be a probability distribution over  $\mathcal{X}$ . The  $\| \cdot \|_{2,\rho}$ -norm of  $f \in \mathcal{F}$  is defined as  $\| f \|_{2,\rho} = \sqrt{\mathbb{E}_{x \sim \rho}[f^2(x)]}$ . Define function  $l: \mathbb{R} \to \mathbb{R}$  and define functional  $\mathcal{L}: \mathcal{F} \to \mathbb{R}$  as
$$
\mathcal {L} (f) = \mathbb {E} _ {x \sim \rho} [ l (f (x)) ].
$$
We say  $\mathcal{L}$  is  $\sigma$ -positive-semi-strongly convex w.r.t.  $f$  in  $\| \cdot \|_{2,\rho}$ -norm for some  $\sigma > 0$  if
$$
\mathcal {L} (f) - \frac {\sigma}{2} \| f _ {+} \| _ {2, \rho} ^ {2}
$$
is convex w.r.t.  $f$  in  $\| \cdot \|_{2,\rho}$ -norm.
Note that the common  $\sigma$ -strongly convex definition requires  $\mathcal{L}(f) - \frac{\sigma}{2}\| f\|_{2,\rho}^2$  to be convex. Now we prove the following property for  $\sigma$ -positive-semi-strong convexity.
Proposition B.1. Under the same setting of Definition 2, if  $\mathcal{L}$  is  $\sigma$ -positive-semi-strongly convex w.r.t.  $f$  in  $\| \cdot \|_{2,\rho}$ -norm, it holds that for any  $f,g\in \mathcal{F}$ ,
$$
\mathcal {L} (f) - \mathcal {L} (g) \geq \nabla \mathcal {L} (g) ^ {\top} (f - g) + \frac {\sigma}{2} \| f _ {+} - g _ {+} \| _ {2, \rho} ^ {2}.
$$
Proof. By Definition 2,  $\tilde{\mathcal{L}}(f) \triangleq \mathcal{L}(f) - \frac{\sigma}{2} \| f_+ \|_{2,\rho}^2$  is convex, which, by the definition of convex, implies that
$$
\tilde {\mathcal {L}} (f) - \tilde {\mathcal {L}} (g) \geq \nabla \tilde {\mathcal {L}} (g) ^ {\intercal} (f - g).
$$
Plugging in the definition of  $\tilde{\mathcal{L}} (\cdot)$ , we have
$$
\mathcal {L} (f) - \frac {\sigma}{2} \| f _ {+} \| _ {2, \rho} ^ {2} - \mathcal {L} (g) + \frac {\sigma}{2} \| g _ {+} \| _ {2, \rho} ^ {2} \geq \nabla \mathcal {L} (g) ^ {\top} (f - g) - \sigma \mathbb {E} _ {x \sim \rho} [ g (x) _ {+} (f (x) - g (x)) ].
$$
Rearranging, we can obtain that
$$
\begin{array}{l} \mathcal {L} (f) - \mathcal {L} (g) \geq \nabla \mathcal {L} (g) ^ {\top} (f - g) - \sigma \mathbb {E} _ {x \sim \rho} [ g (x) _ {+} (f (x) - g (x)) ] + \frac {\sigma}{2} \| f _ {+} \| _ {2, \rho} ^ {2} - \frac {\sigma}{2} \| g _ {+} \| _ {2, \rho} ^ {2} \\ = \nabla \mathcal {L} (g) ^ {\top} (f - g) + \frac {\sigma}{2} \mathbb {E} _ {x \sim \rho} [ f _ {+} ^ {2} (x) - g _ {+} ^ {2} (x) - 2 g (x) _ {+} (f (x) - g (x)) ] \\ = \nabla \mathcal {L} (g) ^ {\top} (f - g) + \frac {\sigma}{2} \mathbb {E} _ {x \sim \rho} [ f _ {+} ^ {2} (x) - g _ {+} ^ {2} (x) - 2 g (x) _ {+} (f (x) - g (x)) ] \\ = \nabla \mathcal {L} (g) ^ {\top} (f - g) + \frac {\sigma}{2} \| f _ {+} - g _ {+} \| _ {2, \rho} ^ {2} \\ + \sigma \mathbb {E} _ {x \sim \rho} [ g _ {+} (x) (f _ {+} (x) - f (x)) ] + \sigma \mathbb {E} _ {x \sim \rho} [ g _ {+} (x) (g (x) - g _ {+} (x)) ]. \\ \end{array}
$$
$$
\mathbb {E} _ {x \sim \rho} [ g _ {+} (x) (f _ {+} (x) - f (x)) ] \geq 0, \quad \mathbb {E} _ {x \sim \rho} [ g _ {+} (x) (g (x) - g _ {+} (x)) ] \geq 0.
$$
Note that for any  $x \in \mathcal{X}$ ,  $g_{+}(x) \geq 0$  and  $f_{+}(x) \geq f(x)$ , which implies  $g_{+}(x)(f_{+}(x) - f(x)) \geq 0$ . Also, note that when  $g(x) \geq 0$ , we have  $g(x) - g_{+}(x) = 0$ , and when  $g(x) < 0$ , we have  $g_{+}(x) = 0$ , which means  $g_{+}(x)(g(x) - g_{+}(x)) \equiv 0$ . Therefore, we have
$$
\mathbb {E} _ {x \sim \rho} [ g _ {+} (x) (f _ {+} (x) - f (x)) ] \geq 0, \quad \mathbb {E} _ {x \sim \rho} [ g _ {+} (x) (g (x) - g _ {+} (x)) ] = 0,
$$
which concludes.
Finally, by semi-strong convexity of  $L_{\alpha}$  w.r.t.  $U_{V}$  in  $\| \cdot \|_{2,\mu}$ -norm, we can show that  $\hat{U}$  and  $U_{\alpha}^{*}$  are also close.
Proof of Lemma 1. We condition on the high probability event in Lemma 3. Recall that  $U_V(s, a; g) = r(s; g) + \gamma \mathcal{T}V(s, a; g) - V(s; g) + \alpha$  and in deterministic dynamics,  $\hat{U} = U_{\hat{V}}$ . Let  $\mathcal{U} = \{U_V : V \in \mathbb{R}_+^{|\mathcal{S}| \times |\mathcal{A}| \times |\mathcal{G}|}\} \subseteq \mathbb{R}^{|\mathcal{S}| \times |\mathcal{A}| \times |\mathcal{G}|}$  which is a convex set by definition. Also, since  $U_V$  is linear in  $V$ , we can also obtain that  $V$  can be linearly represented by  $U$ . Therefore, we can define that  $\tilde{L}_{\alpha}(U_V) = L_{\alpha}(V)$  and  $\tilde{L}_{\alpha}(U_V) - \frac{1}{2}\mathbb{E}_{(s, a, g) \sim \mu}[U_V^2(s, a, g)_+]$  is linear and thus convex in  $U_V$ , which implies that  $\tilde{L}_{\alpha}(U_V)$  is 1-positive-semi-strongly convex w.r.t.  $U_V$  and  $\| \cdot \|_{2,\mu}$ . Then, by Proposition B.1, we have
$$
\tilde {L} _ {\alpha} (\hat {U}) - \tilde {L} _ {\alpha} (U _ {\alpha} ^ {*}) \geq \nabla \tilde {L} _ {\alpha} (U _ {\alpha} ^ {*}) ^ {\top} (\hat {U} - U _ {\alpha} ^ {*}) + \frac {1}{2} \| \hat {U} _ {+} - U _ {\alpha +} ^ {*} \| _ {2, \mu} ^ {2}.
$$
Since  $U_{\alpha}^{*} = \arg \min_{U \in \mathcal{U}} \tilde{L}_{\alpha}(U)$ , by the first order optimality condition, it holds that  $\nabla \tilde{L}_{\alpha}(U_{\alpha}^{*})^{\top}(\hat{U} - U_{\alpha}^{*}) \geq 0$ . Combining Lemma 4, we have
$$
\left\| \hat {U} _ {+} - U _ {\alpha +} ^ {*} \right\| _ {2, \mu} \leq \sqrt {2 (\tilde {L} _ {\alpha} (\hat {U}) - \tilde {L} _ {\alpha} (U _ {\alpha} ^ {*}))} \leq O \left(\sqrt {\epsilon_ {\mathrm {s t a t}}}\right).
$$
B.2 Proof for Stochastic Dynamics
In stochastic dynamic settings, we first learn the ground-truth transition model and then calculate  $\hat{U}$ . The following lemma provides a theoretical guarantee for maximum likelihood estimation (MLE) that  $\hat{P}$  and  $P^{\star}$  are close.
Lemma 5 (Convergence rate of MLE, Van de Geer [2000]). For any fixed  $\delta > 0$ , with probability at least  $1 - \delta$ , we have
$$
\mathbb {E} _ {(s, a, g) \sim \mu} \left[ \| \hat {P} (\cdot | s, a) - P ^ {\star} (\cdot | s, a) \| _ {\mathsf {T V}} ^ {2} \right] \lesssim \frac {\log (| \mathcal {P} | / \delta)}{N},
$$
where  $\hat{P}$  is defined as in (11). This immediately implies that
$$
\mathbb {E} _ {(s, a, g) \sim \mu} \left[ \| \hat {P} (\cdot | s, a) - P ^ {\star} (\cdot | s, a) \| _ {\mathsf {T V}} \right] \lesssim \sqrt {\frac {\log (| \mathcal {P} | / \delta)}{N}}.
$$
Equipped with Lemma 5, we can guarantee that for any  $V$ , the population objective  $L_{\alpha}(V)$  is close to the empirical objective  $\hat{L}^{(s)}(V)$ , which is presented in Lemma 6.
Lemma 6 (Concentration of  $\hat{L}^{(s)}$ ). Under Assumptions 1 to 3, with probability at least  $1 - \delta$ ,  $\forall V \in \mathcal{V}$ , it holds that
$$
\begin{array}{l} \left| \hat {L} ^ {(s)} (V) - L _ {\alpha} (V) \right| \\ \leq O \left(\alpha V _ {\max } \sqrt {\frac {\log (| \mathcal {V} | / \delta)}{N _ {0}}} + V _ {\max } ^ {2} \sqrt {\frac {\log (| \mathcal {V} | / \delta)}{N}} + \gamma V _ {\max } ^ {2} \sqrt {\frac {\log (| \mathcal {P} | / \delta)}{N}}\right) \triangleq \epsilon_ {s t a t} ^ {s t o c h a s t i c}. \\ \end{array}
$$
Proof. For convenience, define
$$
\begin{array}{l} \hat {L} _ {1} ^ {(s)} (V) = \frac {1 - \gamma}{N _ {0}} \sum_ {i = 1} ^ {N _ {0}} \alpha \cdot V (s _ {0, i}; g _ {0, i}), \\ \hat {L} _ {2} ^ {(s)} (V) = \frac {1}{N} \sum_ {i = 1} ^ {N} \alpha \cdot g _ {* +} (r (s _ {i}; g _ {i}) + \gamma \hat {\mathcal {T}} V (s _ {i}, a _ {i}; g _ {i}) - V (s _ {i}; g _ {i})). \\ \end{array}
$$
Note that  $\hat{L}_1^{(s)}(V) = \hat{L}_1^{(d)}(V)$  and thus by Lemma 3 we have that with probability at least  $1 - \delta$ , for all  $V \in \mathcal{V}$ ,
$$
\left| \hat {L} _ {1} ^ {(s)} (V) - L _ {1} (V) \right| \leq O \left(\alpha V _ {\max } \sqrt {\frac {\log (| \mathcal {V} | / \delta)}{N _ {0}}}\right). \tag {16}
$$
$$
\tilde {L} _ {2} ^ {(s)} (V) = \mathbb {E} _ {(s, a, g) \sim \mu} [ \alpha \cdot g _ {* +} (r (s; g) + \gamma \hat {\mathcal {T}} V (s, a; g) - V (s; g)) ]
$$
Since  $\mathbb{E}_{(s,a,g)\sim \mu}[\hat{L}_2^{(s)}(V)] = \tilde{L}_2^{(s)}(V)$ , by Lemma 3, we have that with probability  $1 - \delta$ , for any  $V\in \mathcal{V}$ , it holds that
$$
\left| \hat {L} _ {2} ^ {(s)} (V) - \tilde {L} _ {2} ^ {(s)} (V) \right| \leq O \left(V _ {\max } ^ {2} \sqrt {\frac {\log \left(| \mathcal {V} | / \delta\right)}{N}}\right). \tag {17}
$$
$$
\begin{array}{l} \left| L _ {2} (V) - \tilde {L} _ {2} ^ {(s)} (V) \right| \\ = \left| \alpha \mathbb {E} _ {(s, a, g) \sim \mu} [ g _ {* +} (r (s; g) + \gamma \mathcal {T} V (s, a; g) - V (s; g)) - g _ {* +} (r (s; g) + \gamma \hat {\mathcal {T}} V (s, a; g) - V (s; g)) ] \right| \\ = \frac {1}{2} \left| \mathbb {E} _ {(s, a, g) \sim \mu} \left[ (r (s; g) + \gamma \mathcal {T} V (s, a; g) - V (s; g) + \alpha) _ {+} ^ {2} - (r (s; g) + \gamma \hat {\mathcal {T}} V (s, a; g) - V (s; g) + \alpha) _ {+} ^ {2} \right] \right| \\ \lesssim V _ {\max } \mathbb {E} _ {(s, a, g) \sim \mu} [ \gamma \cdot | (\mathcal {T} - \hat {\mathcal {T}}) V (s, a; g) | ] \\ = V _ {\max } \mathbb {E} _ {(s, a, g) \sim \mu} [ \gamma \cdot | \mathbb {E} _ {s ^ {\prime} \sim P ^ {*} (\cdot | s, a)} [ V (s ^ {\prime}; g) ] - \mathbb {E} _ {s ^ {\prime} \sim \hat {P} (\cdot | s, a)} [ V (s ^ {\prime}; g) ] | ] \\ \lesssim \gamma V _ {\max } ^ {2} \mathbb {E} _ {(s, a, g) \sim \mu} [ \| P ^ {\star} (\cdot | s, a) - \hat {P} (\cdot | s, a) \| _ {\mathsf {T V}} ]. \\ \end{array}
$$
By Lemma 5, with probability at least  $1 - \delta$ , for all  $V \in \mathcal{V}$
$$
\left| L _ {2} (V) - \tilde {L} _ {2} ^ {(s)} (V) \right| \lesssim \gamma V _ {\max } ^ {2} \mathbb {E} _ {(s, a, g) \sim \mu} [ \| P ^ {\star} (\cdot | s, a) - \hat {P} (\cdot | s, a) \| _ {\mathrm {T V}} ] \lesssim \gamma V _ {\max } ^ {2} \sqrt {\frac {\log (| \mathcal {P} | / \delta)}{N}}. \tag {18}
$$
By rescaling  $\delta$  and applying a union bound, we can obtain that with probability at least  $1 - \delta$ , (16), (17), (18) hold simultaneously. The conclusion holds by applying the triangle inequality.
Finally, we show that  $\hat{U}_{+}$  and  $U_{\alpha +}^{*}$  are close.
Lemma 7 (Closeness of  $\hat{U}_{+}$  and  $U_{\alpha +}^{*}$  in stochastic dynamics). Under Assumptions 1 to 3, with probability at least  $1 - \delta$
$$
\left\| \hat {U} _ {+} - U _ {\alpha +} ^ {*} \right\| _ {2, \mu} \leq O \left(\sqrt {\epsilon_ {s t a t} ^ {s t o c h a s t i c}}\right).
$$
where  $\hat{U}$  is the output of Algorithm 3,  $\alpha = \Omega \left(\frac{\log(|\mathcal{V}||\mathcal{P}| / \delta)}{\sqrt{N}}\right)$  and
$$
\epsilon_ {s t a t} ^ {s t o c h a s t i c} \asymp V _ {\max } ^ {2} \sqrt {\frac {\log (| \mathcal {V} | | \mathcal {P} | / \delta)}{N}}.
$$
Proof. For convenience, let  $\tilde{U} = U_{\hat{V}}$ . Following the same analysis as in the deterministic case, we have  $\| U_{\alpha +}^{*} - \tilde{U}_{+}\|_{2,\mu}\leq O\left(\sqrt{\epsilon_{\mathrm{stat}}^{\mathrm{stochastic}}}\right)$  by Lemma 6. Also,
$$
\begin{array}{l} \| \tilde {U} _ {+} - \hat {U} _ {+} \| _ {2, \mu} ^ {2} \leq \| \tilde {U} - \hat {U} \| _ {2, \mu} ^ {2} \\ = \mathbb {E} _ {(s, a, g) \sim \mu} \left[ \left(\gamma \mathbb {E} _ {s ^ {\prime} \sim P ^ {\star} (\cdot | s, a)} [ \hat {V} (s; g) ] - \gamma \mathbb {E} _ {s ^ {\prime} \sim \hat {P} (\cdot | s, a)} [ \hat {V} (s; g) ]\right) ^ {2} \right] \\ \lesssim \gamma^ {2} V _ {\max } ^ {2} \mathbb {E} _ {(s, a, g) \sim \mu} \left[ \| \hat {P} (\cdot | s, a) - P ^ {\star} (\cdot | s, a) \| _ {\mathsf {T V}} ^ {2} \right] \\ \lesssim \gamma^ {2} V _ {\max } ^ {2} \frac {\log (| \mathcal {P} | / \delta)}{N}, \\ \end{array}
$$
where the last inequality holds by Lemma 5. Therefore,
$$
\begin{array}{l} \| \hat {U} _ {+} - U _ {\alpha +} ^ {*} \| _ {2, \mu} \leq \| \tilde {U} _ {+} - U _ {\alpha +} ^ {*} \| _ {2, \mu} + \| \tilde {U} _ {+} - \hat {U} _ {+} \| _ {2, \mu} \\ \lesssim \sqrt {\epsilon_ {\mathrm {s t a t}} ^ {\mathrm {s t o c h a s t i c}}} + \gamma V _ {\max } \sqrt {\frac {\log (| \mathcal {P} | / \delta)}{N}} \\ \lesssim \sqrt {\epsilon_ {\mathrm {s t a t}} ^ {\mathrm {s t o c h a s t i c}}}. \\ \end{array}
$$
C Proof for Policy Learning
We provide a theoretical analysis of policy learning in this section. We first show two key lemmas in Appendix C.1, and then provide missing proofs of the main text in Appendix C.2 and Appendix C.3.
C.1 Key Lemmas
Lemma 8 (Statistical error of the weighted MLE objective). Under Assumption 5, with probability at least  $1 - \delta$ , for any policy  $\pi \in \Pi$ , it holds that
$$
\left| L _ {\alpha} ^ {M L E} (\pi) - \hat {L} ^ {M L E} (\pi) \right| \leq \left(\epsilon_ {U} + \epsilon_ {\Pi}\right) \log (1 / \tau) \triangleq \epsilon_ {s t a t} ^ {M L E}.
$$
where  $\epsilon_U \triangleq O\left(\sqrt{\epsilon_{stat} / \alpha^2}\right)$  and  $\epsilon_{\Pi} \triangleq O\left(C_{\alpha}^{*}\sqrt{\frac{\log(|\Pi| / \delta)}{N}} + \frac{V_{\max}}{\alpha N}\log (|\Pi | / \delta)\right)$ .