SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Lemma 5. (Proposition A.8. in [69]) Let  $z$  be a discrete random variable that takes values in  $\{1, \dots, d\}$ , distributed according to  $q$ . We write  $q$  as a vector where  $\vec{q} = [Pr(z = j)]_{j=1}^{d}$ . Assume there are  $N$  i.i.d. samples, and that the empirical estimate of  $\vec{q}$  is  $[\hat{q}]_j = \frac{1}{N} \sum_{i=1}^N \mathbf{1}[z_i = j]$ , where  $\hat{q}$  is a  $d$ -dimensional vector.
Then for any  $\epsilon > 0$ , the following result holds:
$$
P r \left(\| \hat {q} - \vec {q} \| _ {2} \geq \frac {1}{\sqrt {N}} + \epsilon\right) \leq e ^ {- N \epsilon^ {2}}, \tag {23}
$$
which implies that:
$$
P r \left(\| \hat {q} - \vec {q} \| _ {1} \geq \sqrt {d} \left(\frac {1}{\sqrt {N}} + \epsilon\right)\right) \leq e ^ {- N \epsilon^ {2}}. \tag {24}
$$
Lemma 6. (Lemma 3 in [70]) Under any initial distribution  $\eta(\cdot)$  and any transition dynamics  $P(\cdot | s, a)$ , we denote  $d_w(\cdot, \cdot)$  as the visitation measure of the visited state-action pair  $(s, a)$  with a softmax policy parameterized by parameter  $w$ . Suppose Assumption 2 holds, then for all policy parameter  $w$  and  $w'$ , we have
$$
\left\| d _ {w} (\cdot , \cdot) - d _ {w ^ {\prime}} (\cdot , \cdot) \right\| _ {\mathrm {T V}} \leq C _ {d} \| w - w ^ {\prime} \| \tag {25}
$$
where  $C_d$  is a positive constant.
Lemma 7. (Lemma 5 in [31]) Suppose Assumption 3 holds. Under the soft policy iteration defined in (13) - (14), we denote the soft  $Q$ -function under reward parameter  $\theta_{k}$  and policy  $\pi_{k+1}$  as  $Q_{k + \frac{1}{2}}$ . Moreover, recall that  $Q_{k+1}$  has been defined as the soft  $Q$ -function under the reward parameter  $\theta_{k+1}$  and policy  $\pi_{k+1}$ . Then for any  $s \in S$ ,  $a \in \mathcal{A}$  and  $k \geq 0$ , the following inequality holds:
$$
\left| Q _ {k + \frac {1}{2}} (s, a) - Q _ {k + 1} (s, a) \right| \leq L _ {q} \| \theta_ {k} - \theta_ {k + 1} \|, \tag {26}
$$
where  $L_{q} \coloneqq \frac{L_{r}}{1 - \gamma} > 0$  and  $L_{r}$  is the positive constant defined in Assumption 3.
Lemma 8. (Lemma 6 in [31]) Following the soft policy iteration defined in (13) - (14), the following holds for any iteration  $k \geq 0$ :
$$
Q _ {k} (s, a) \leq Q _ {k + \frac {1}{2}} (s, a) + \frac {2 \gamma \epsilon_ {\mathrm {a p p}}}{1 - \gamma}, \quad \forall s \in \mathcal {S}, a \in \mathcal {A}, \tag {27a}
$$
$$
\left\| Q _ {\theta_ {k}} - Q _ {k + \frac {1}{2}} \right\| _ {\infty} \leq \gamma \left\| Q _ {\theta_ {k}} - Q _ {k} \right\| _ {\infty} + \frac {2 \gamma \epsilon_ {\mathrm {a p p}}}{1 - \gamma} \tag {27b}
$$
where  $Q_{k + \frac{1}{2}}(\cdot ,\cdot)$  denotes the soft  $Q$ -function under reward parameter  $\theta_{k}$  and updated policy  $\pi_{k + 1}$ , and  $Q_{\theta_k}(\cdot ,\cdot)$  denotes the soft  $Q$ -function under reward parameter  $\theta_{k}$  and corresponding optimal policy  $\pi_{\theta_k}$ . Moreover, we denote  $\| Q_{\theta_k} - Q_{k + \frac{1}{2}}\|_{\infty} = \max_{s\in S}\max_{a\in \mathcal{A}}|Q_{\theta_k}(s,a) - Q_{k + \frac{1}{2}}(s,a)|$ .
Remark. In Lemma 7 and Lemma 8, the definitions of the soft Q-function  $Q_{k}$  and the soft value function  $V_{k}$  under a conservative MDP are given below:
$$
Q _ {k} (s, a) := r (s, a; \theta_ {k}) + U (s, a) + \gamma \mathbb {E} _ {s ^ {\prime} \sim \hat {P} (\cdot | s, a)} \left[ V _ {k} \left(s ^ {\prime}\right) \right] \tag {28a}
$$
$$
V _ {k} (s) := \mathbb {E} _ {\tau \sim (\eta , \pi_ {k}, \widehat {P})} \left[ \sum_ {t = 0} ^ {\infty} \gamma^ {t} \left(r \left(s _ {t}, a _ {t}; \theta_ {k}\right) + U \left(s _ {t}, a _ {t}\right) + \mathcal {H} \left(\pi_ {k} (\cdot | s _ {t})\right)\right) \mid s _ {0} = s \right] \tag {28b}
$$
where  $U(s, a)$  is the penalty function for the state-action pair  $(s, a)$ , which is used to quantify the uncertainty in the estimated world model  $\hat{P}$ . If we rewrite the above defined soft Q-function  $Q_{k}$  as:
$$
Q _ {k} (s, a) = \tilde {r} (s, a; \theta_ {k}) + \gamma \mathbb {E} _ {s ^ {\prime} \sim \widehat {P} (\cdot | s, a)} \big [ V _ {k} (s ^ {\prime}) \big ]
$$
where  $\tilde{r}(s, a; \theta_k) \coloneqq r(s, a; \theta_k) + U(s, a)$ , then we can directly follow the proof steps in [31] to prove Lemma 7 and Lemma 8.
Proof. According to the closed-form expressions of the optimal policy  $\pi_{\theta}$  and the optimal soft value function  $V_{\theta}$  in (4), let us decompose the objective  $L(\theta)$  defined in (2a) as below:
$$
\begin{array}{l} L (\theta) = \mathbb {E} _ {\tau^ {\mathrm {E}} \sim (\eta , \pi^ {\mathrm {E}}, P)} \left[ \sum_ {t = 0} ^ {\infty} \gamma^ {t} \log \pi_ {\theta} (a _ {t} | s _ {t}) \right] \\ \stackrel {(i)} {=} \mathbb {E} _ {\tau^ {\mathrm {E}} \sim (\eta , \pi^ {\mathrm {E}}, P)} \left[ \sum_ {t = 0} ^ {\infty} \gamma^ {t} \log \left(\frac {\exp Q _ {\theta} (s _ {t} , a _ {t})}{\sum_ {a \in \mathcal {A}} \exp Q _ {\theta} (s _ {t} , a)}\right) \right] \\ \stackrel {(i i)} {=} \mathbb {E} _ {\tau^ {\mathrm {E}} \sim (\eta , \pi^ {\mathrm {E}}, P)} \left[ \sum_ {t = 0} ^ {\infty} \gamma^ {t} \Big (Q _ {\theta} (s _ {t}, a _ {t}) - V _ {\theta} (s _ {t}) \Big) \right] \\ = \sum_ {t = 0} ^ {\infty} \gamma^ {t} \mathbb {E} _ {(s _ {t}, a _ {t}) \sim (\eta , \pi^ {\mathrm {E}}, P)} \left[ Q _ {\theta} (s _ {t}, a _ {t}) \right] - \sum_ {t = 0} ^ {\infty} \gamma^ {t} \mathbb {E} _ {s _ {t} \sim (\eta , \pi^ {\mathrm {E}}, P)} \left[ V _ {\theta} (s _ {t}) \right] \\ \stackrel {(i i i)} {=} \sum_ {t = 0} ^ {\infty} \gamma^ {t} \mathbb {E} _ {(s _ {t}, a _ {t}) \sim (\eta , \pi^ {\mathrm {E}}, P)} \left[ r (s _ {t}, a _ {t}; \theta) + U (s _ {t}, a _ {t}) + \gamma \mathbb {E} _ {s _ {t + 1} \sim \widehat {P} (\cdot | s _ {t}, a _ {t})} [ V _ {\theta} (s _ {t + 1}) ] \right] - \sum_ {t = 0} ^ {\infty} \gamma^ {t} \mathbb {E} _ {s _ {t} \sim (\eta , \pi^ {\mathrm {E}}, P)} \left[ V _ {\theta} (s _ {t}) \right] \\ = \sum_ {t = 0} ^ {\infty} \gamma^ {t} \mathbb {E} _ {(s _ {t}, a _ {t}) \sim (\eta , \pi^ {\mathrm {E}}, P)} \left[ r (s _ {t}, a _ {t}; \theta) + U (s _ {t}, a _ {t}) \right] - \mathbb {E} _ {s _ {0} \sim \eta (\cdot)} \left[ V _ {\theta} (s _ {0}) \right] \\ + \left(\sum_ {t = 0} ^ {\infty} \gamma^ {t + 1} \mathbb {E} _ {(s _ {t}, a _ {t}) \sim (\eta , \pi^ {\mathrm {E}}, P), s _ {t + 1} \sim \widehat {P} (\cdot | s _ {t}, a _ {t})} [ V _ {\theta} (s _ {t + 1}) ] - \sum_ {t = 0} ^ {\infty} \gamma^ {t + 1} \mathbb {E} _ {s _ {t + 1} \sim (\eta , \pi^ {\mathrm {E}}, P)} [ V _ {\theta} (s _ {t + 1}) ]\right) \\ = \underbrace {\left(\sum_ {t = 0} ^ {\infty} \gamma^ {t} \mathbb {E} _ {\left(s _ {t} , a _ {t}\right) \sim \left(\eta , \pi^ {\mathrm {E}} , P\right)} \left[ r \left(s _ {t} , a _ {t} ; \theta\right) + U \left(s _ {t} , a _ {t}\right) \right] - \mathbb {E} _ {s _ {0} \sim \eta (\cdot)} \left[ V _ {\theta} \left(s _ {0}\right) \right]\right)} _ {\text {T 1 : s u r r o g a t e o b j e c t i v e}} \\ + \underbrace {\left(\sum_ {t = 0} ^ {\infty} \gamma^ {t + 1} \mathbb {E} _ {\left(s _ {t} , a _ {t}\right) \sim \left(\eta , \pi^ {\mathrm {E}} , P\right) , s _ {t + 1} \sim \hat {P} \left(\cdot \mid s _ {t} , a _ {t}\right)} \left[ V _ {\theta} \left(s _ {t + 1}\right) \right] - \sum_ {t = 0} ^ {\infty} \gamma^ {t + 1} \mathbb {E} _ {\left(s _ {t} , a _ {t}\right) \sim \left(\eta , \pi^ {\mathrm {E}} , P\right) , s _ {t + 1} \sim P \left(\cdot \mid s _ {t} , a _ {t}\right)} \left[ V _ {\theta} \left(s _ {t + 1}\right) \right]\right)} _ {\text {T 2 : e r r o r t e r m d u e t o t r a n s i t i o n p r o b a b i l i t y m i s m a t c h}} \tag {29} \\ \end{array}
$$
where (i) and (ii) follows the closed-form expression of the optimal policy  $\pi_{\theta}$  and the optimal soft value function in (4). Moreover, (iii) follows the definition of the soft Q-function in (3b).
From the  $\mathrm{T}_2$  term in (29), we observe that the error term comes from the transition dynamics mismatch between the estimated world model  $\widehat{P}$  and the ground-truth dynamics model  $P$ . Then we analyze the transition dynamics mismatch:
$$
\begin{array}{l} T _ {2} = \sum_ {t = 0} ^ {\infty} \gamma^ {t + 1} \mathbb {E} _ {(s _ {t}, a _ {t}) \sim (\eta , \pi^ {\mathrm {E}}, P), s _ {t + 1} \sim \widehat {P} (\cdot | s _ {t}, a _ {t})} \left[ V _ {\theta} (s _ {t + 1}) \right] - \sum_ {t = 0} ^ {\infty} \gamma^ {t + 1} \mathbb {E} _ {(s _ {t}, a _ {t}) \sim (\eta , \pi^ {\mathrm {E}}, P), s _ {t + 1} \sim P (\cdot | s _ {t}, a _ {t})} \left[ V _ {\theta} (s _ {t + 1}) \right] \\ = \sum_ {t = 0} ^ {\infty} \gamma^ {t + 1} \mathbb {E} _ {(s _ {t}, a _ {t}) \sim (\eta , \pi^ {\mathrm {E}}, P)} \left[ \mathbb {E} _ {s _ {t + 1} \sim \widehat {P} (\cdot | s _ {t}, a _ {t})} [ V _ {\theta} (s _ {t + 1}) ] - \mathbb {E} _ {s _ {t + 1} \sim P (\cdot | s _ {t}, a _ {t})} [ V _ {\theta} (s _ {t + 1}) ] \right] \\ = \sum_ {t = 0} ^ {\infty} \gamma^ {t + 1} \mathbb {E} _ {(s _ {t}, a _ {t}) \sim (\eta , \pi^ {\mathrm {E}}, P)} \left[ \sum_ {s _ {t + 1} \in \mathcal {S}} V _ {\theta} (s _ {t + 1}) \left(\widehat {P} (s _ {t + 1} | s _ {t}, a _ {t}) - P (s _ {t + 1} | s _ {t}, a _ {t})\right) \right] \\ \end{array}
$$
Recall that  $d^{\mathrm{E}}(\cdot ,\cdot)$  is defined in (5) where  $d^{\mathrm{E}}(s,a)\coloneqq (1 - \gamma)\pi^{\mathrm{E}}(a|s)\sum_{t = 0}^{\infty}\gamma^{t}P^{\pi^{\mathrm{E}}} (s_{t} = s|s_{0}\sim \eta),$  we obtain the following result:
$$
\begin{array}{l} T _ {2} = \sum_ {t = 0} ^ {\infty} \gamma^ {t + 1} \mathbb {E} _ {(s _ {t}, a _ {t}) \sim (\eta , \pi^ {\mathrm {E}}, P)} \left[ \sum_ {s _ {t + 1} \in \mathcal {S}} V _ {\theta} (s _ {t + 1}) \left(\widehat {P} (s _ {t + 1} | s _ {t}, a _ {t}) - P (s _ {t + 1} | s _ {t}, a _ {t})\right) \right] \\ = \gamma \sum_ {t = 0} ^ {\infty} \sum_ {s \in \mathcal {S}, a \in \mathcal {A}} \gamma^ {t} P (s _ {t} = s | s _ {0} \sim \eta) \pi^ {E} (a _ {t} = a | s _ {t} = s) \Bigg (\sum_ {s _ {t + 1} \in \mathcal {S}} V _ {\theta} (s _ {t + 1}) \big (\widehat {P} (s _ {t + 1} | s _ {t} = s, a _ {t} = a) - P (s _ {t + 1} | s _ {t} = s, a _ {t} = a) \big) \Bigg) \\ = \frac {\gamma}{1 - \gamma} \cdot \mathbb {E} _ {(s, a) \sim d ^ {\mathrm {E}} (\cdot , \cdot)} \left[ \sum_ {s ^ {\prime} \in S} V _ {\theta} \left(s ^ {\prime}\right) \left(\widehat {P} \left(s ^ {\prime} \mid s, a\right) - P \left(s ^ {\prime} \mid s, a\right)\right) \right] \tag {30} \\ \end{array}
$$
Then we further denote the surrogate objective  $\widehat{L} (\cdot)$  as below:
$$
\widehat {L} (\theta) := \sum_ {t = 0} ^ {\infty} \gamma^ {t} \mathbb {E} _ {(s _ {t}, a _ {t}) \sim (\eta , \pi^ {\mathrm {E}}, P)} \left[ r (s _ {t}, a _ {t}; \theta) + U (s _ {t}, a _ {t}) \right] - \mathbb {E} _ {s _ {0} \sim \eta (\cdot)} \left[ V _ {\theta} (s _ {0}) \right]. \tag {31}
$$
By plugging (30) and (31) into (29), we obtain the decomposition of the likelihood objective  $L(\theta)$  as below:
$$
L (\theta) = \widehat {L} (\theta) + \frac {\gamma}{1 - \gamma} \cdot \mathbb {E} _ {(s, a) \sim d ^ {\mathrm {E}} (\cdot , \cdot)} \left[ \sum_ {s ^ {\prime} \in \mathcal {S}} V _ {\theta} \left(s ^ {\prime}\right) \left(\widehat {P} \left(s ^ {\prime} \mid s, a\right) - P \left(s ^ {\prime} \mid s, a\right)\right) \right]. \tag {32}
$$
Proof. According to (32), we have the following series of relations:
$$
\begin{array}{l} | L (\theta) - \widehat {L} (\theta) | = \frac {\gamma}{1 - \gamma} \cdot \left| \mathbb {E} _ {(s, a) \sim d ^ {\mathrm {E}} (\cdot , \cdot)} \left[ \sum_ {s ^ {\prime} \in \mathcal {S}} V _ {\theta} (s ^ {\prime}) \left(\widehat {P} (s ^ {\prime} | s, a) - P (s ^ {\prime} | s, a)\right) \right] \right| \\ \leq \frac {\gamma}{1 - \gamma} \cdot \mathbb {E} _ {(s, a) \sim d ^ {\mathbb {E}} (\cdot , \cdot)} \left[ \sum_ {s ^ {\prime} \in \mathcal {S}} | V _ {\theta} (s ^ {\prime}) | \cdot | \widehat {P} (s ^ {\prime} | s, a) - P (s ^ {\prime} | s, a) | \right] \\ \leq \frac {\gamma}{1 - \gamma} \cdot \max _ {\tilde {s} \in \mathcal {S}} | V _ {\theta} (\tilde {s}) | \cdot \mathbb {E} _ {(s, a) \sim d ^ {\mathrm {E}} (\cdot , \cdot)} \Big [ \sum_ {s ^ {\prime} \in \mathcal {S}} \left| \hat {P} (s ^ {\prime} | s, a) - P (s ^ {\prime} | s, a) \right| \Big ] \\ = \frac {\gamma}{1 - \gamma} \cdot R (\theta) \cdot \mathbb {E} _ {(s, a) \sim d ^ {\mathrm {E}} (\cdot , \cdot)} \left[ \left\| \widehat {P} (\cdot | s, a) - P (\cdot | s, a) \right\| _ {1} \right] \tag {33} \\ \end{array}
$$
where (33) follows the definition  $R(\theta) \coloneqq \max_{s \in S} |V_{\theta}(s)|$  and the definition of the visitation measure  $d^{\mathrm{E}}(\cdot, \cdot)$  in (5). According to the definition of  $V_{\theta}(\cdot)$  in (3a), we have
$$
R (\theta) = \max _ {s \in S} \left| V _ {\theta} (s) \right|
$$
$$
\begin{array}{l} \stackrel {(i)} {=} \max  _ {s \in \mathcal {S}} \left| \mathbb {E} _ {\tau \sim (\eta , \pi_ {\theta}, \widehat {P})} \left[ \sum_ {t = 0} ^ {\infty} \gamma^ {t} \big (r (s _ {t}, a _ {t}; \theta) + U (s _ {t}, a _ {t}) + \mathcal {H} (\pi_ {\theta} (\cdot | s _ {t})) \big) \Big | s _ {0} = s \right] \right| \\ \leq \max  _ {s \in \mathcal {S}} \mathbb {E} _ {\tau \sim (\eta , \pi_ {\theta}, \widehat {P})} \left[ \sum_ {t = 0} ^ {\infty} \gamma^ {t} \left(| r (s _ {t}, a _ {t}; \theta) | + | U (s _ {t}, a _ {t}) | + | \mathcal {H} (\pi_ {\theta} (\cdot | s _ {t})) |\right) \mid s _ {0} = s \right] \\ \end{array}
$$
$$
\begin{array}{l} \stackrel {(i i)} {\leq} \max _ {s \in \mathcal {S}} \mathbb {E} _ {\tau \sim (\eta , \pi_ {\theta}, \widehat {P})} \Big [ \sum_ {t = 0} ^ {\infty} \gamma^ {t} \big (C _ {r} + C _ {u} + | \mathcal {H} (\pi_ {\theta} (\cdot | s _ {t})) | \big) \Big | s _ {0} = s \Big ] \\ \stackrel {(i i i)} {\leq} \max  _ {s \in \mathcal {S}} \mathbb {E} _ {\tau \sim (\eta , \pi_ {\theta}, \widehat {P})} \left[ \sum_ {t = 0} ^ {\infty} \gamma^ {t} \left(C _ {r} + C _ {u} + \log | \mathcal {A} |\right) \mid s _ {0} = s \right] \\ = \frac {C _ {r} + C _ {u} + \log | \mathcal {A} |}{1 - \gamma}. \\ \end{array}
$$
where (i) follows the definition of  $V_{\theta}(\cdot)$  in (3a) and (ii) follows (8). Moreover, (iii) follows the fact that information entropy is non-negative and the maximum entropy is obtained under uniform distribution where we have  $|\mathcal{H}(\pi_{\theta}(\cdot |s))| = \mathcal{H}(\pi_{\theta}(\cdot |s))\leq -\sum_{a\in \mathcal{A}}\frac{1}{|\mathcal{A}|}\log |\frac{1}{|\mathcal{A}|}| = \log |\mathcal{A}|$ .
Denote  $C_v \coloneqq \frac{C_r + C_u + \log |\mathcal{A}|}{1 - \gamma}$ , we obtain the property:  $R(\theta) = \max_{s \in S} |V_\theta(s)| \leq C_v$ . Plugging this result into (33), we obtain the following result to finish the proof:
$$
| L (\theta) - \widehat {L} (\theta) | \leq \frac {\gamma C _ {v}}{1 - \gamma} \cdot \mathbb {E} _ {(s, a) \sim d ^ {\mathrm {E}} (\cdot , \cdot)} \left[ \left\| \widehat {P} (\cdot | s, a) - P (\cdot | s, a) \right\| _ {1} \right].
$$
F Proof of Lemma 3
Proof. To start the analysis, we first take gradient of the surrogate objective  $\widehat{L}(\theta)$  (as defined in (7)) w.r.t.  $\theta$ :
$$
\begin{array}{l} \nabla \widehat {L} (\theta) \stackrel {(i)} {=} \mathbb {E} _ {\tau^ {\mathrm {E}} \sim (\eta , \pi^ {\mathrm {E}}, P)} \left[ \sum_ {t = 0} ^ {\infty} \gamma^ {t} \Big (\nabla_ {\theta} r (s _ {t}, a _ {t}; \theta) + \nabla_ {\theta} U (s _ {t}, a _ {t}) \Big) \right] - \mathbb {E} _ {s _ {0} \sim \eta (\cdot)} \left[ \nabla_ {\theta} V _ {\theta} (s _ {0}) \right] \\ \stackrel {(i i)} {=} \mathbb {E} _ {\tau^ {\mathrm {E}} \sim (\eta , \pi^ {\mathrm {E}}, P)} \left[ \sum_ {t = 0} ^ {\infty} \gamma^ {t} \nabla_ {\theta} r (s _ {t}, a _ {t}; \theta) \right] - \mathbb {E} _ {s _ {0} \sim \eta (\cdot)} \left[ \nabla_ {\theta} \log \left(\sum_ {a \in \mathcal {A}} \exp Q _ {\theta} (s _ {0}, a)\right) \right] \\ = \mathbb {E} _ {\tau^ {\mathrm {E}} \sim (\eta , \pi^ {\mathrm {E}}, P)} \left[ \sum_ {t = 0} ^ {\infty} \gamma^ {t} \nabla_ {\theta} r (s _ {t}, a _ {t}; \theta) \right] - \mathbb {E} _ {s _ {0} \sim \eta (\cdot)} \left[ \sum_ {a \in \mathcal {A}} \left(\frac {\exp Q _ {\theta} (s _ {0} , a)}{\sum_ {\tilde {a} \in \mathcal {A}} \exp Q _ {\theta} (s _ {0} , \tilde {a})} \nabla_ {\theta} Q _ {\theta} (s _ {0}, a)\right) \right] \\ \stackrel {(i i i)} {=} \mathbb {E} _ {\tau^ {\mathrm {E}} \sim (\eta , \pi^ {\mathrm {E}}, P)} \left[ \sum_ {t = 0} ^ {\infty} \gamma^ {t} \nabla_ {\theta} r (s _ {t}, a _ {t}; \theta) \right] - \mathbb {E} _ {s _ {0} \sim \eta (\cdot)} \left[ \sum_ {a \in \mathcal {A}} \pi_ {\theta} (a | s _ {0}) \nabla_ {\theta} Q _ {\theta} (s _ {0}, a) \right] \\ = \mathbb {E} _ {\tau^ {\mathrm {E}} \sim (\eta , \pi^ {\mathrm {E}}, P)} \left[ \sum_ {t = 0} ^ {\infty} \gamma^ {t} \nabla_ {\theta} r (s _ {t}, a _ {t}; \theta) \right] - \mathbb {E} _ {s _ {0} \sim \eta (\cdot), a _ {0} \sim \pi_ {\theta} (\cdot | s _ {0})} \left[ \nabla_ {\theta} Q _ {\theta} (s _ {0}, a _ {0}) \right] \tag {34} \\ \end{array}
$$
where (i) follows the definition of the surrogate objective  $\widehat{L}(\theta)$  in (7), and (ii) follows the closed-form expression of the optimal soft value function  $V_{\theta}$  in (4) and the penalty function  $U(s, a)$  is independent of the reward parameter  $\theta$ . Moreover, (iii) follows the closed-form expression of the optimal policy  $\pi_{\theta}$  in (4). To further analyze the gradient expression in (34), we must derive the gradient of the optimal soft Q-function  $Q_{\theta}$ . Recall that  $Q_{\theta}$  is the soft Q-function under the optimal policy  $\pi_{\theta}$ , the penalty function  $U$  and the estimated world model  $\widehat{P}$ . Therefore, we have the following derivations:
$$
\begin{array}{l} \nabla_ {\theta} Q _ {\theta} (s _ {0}, a _ {0}) \\ \stackrel {(i)} {=} \nabla_ {\theta} \left(r (s _ {0}, a _ {0}; \theta) + U (s _ {0}, a _ {0}) + \gamma \mathbb {E} _ {s _ {1} \sim \widehat {P} (\cdot | s _ {0}, a _ {0})} [ V _ {\theta} (s _ {1}) ]\right) \\ \stackrel {(i i)} {=} \nabla_ {\theta} r (s _ {0}, a _ {0}; \theta) + \nabla_ {\theta} U (s _ {0}, a _ {0}) + \gamma \mathbb {E} _ {s _ {1} \sim \widehat {P} (\cdot | s _ {0}, a _ {0})} \left[ \nabla_ {\theta} \log \left(\sum_ {\tilde {a} \in \mathcal {A}} \exp Q _ {\theta} (s _ {0}, \tilde {a})\right) \right] \\ = \nabla_ {\theta} r (s _ {0}, a _ {0}; \theta) + \gamma \mathbb {E} _ {s _ {1} \sim \widehat {P} (\cdot | s _ {0}, a _ {0})} \left[ \sum_ {a \in \mathcal {A}} \frac {\exp Q _ {\theta} (s _ {1} , a)}{\sum_ {\tilde {a} \in \mathcal {A}} \exp Q _ {\theta} (s _ {1} , \tilde {a})} \nabla_ {\theta} Q _ {\theta} (s _ {1}, a) \right] \\ \stackrel {(i i i)} {=} \nabla_ {\theta} r (s _ {0}, a _ {0}; \theta) + \gamma \mathbb {E} _ {s _ {1} \sim \widehat {P} (\cdot | s _ {0}, a _ {0})} \left[ \sum_ {a \in \mathcal {A}} \pi_ {\theta} (a | s _ {1}) \nabla_ {\theta} Q _ {\theta} (s _ {1}, a) \right] \\ \stackrel {(i v)} {=} \nabla_ {\theta} r (s _ {0}, a _ {0}; \theta) + \gamma \mathbb {E} _ {s _ {1} \sim \widehat {P} (\cdot | s _ {0}, a _ {0}), a _ {1} \sim \pi_ {\theta} (\cdot | s _ {1})} \bigg [ \nabla_ {\theta} \bigg (r (s _ {1}, a _ {1}; \theta) + U (s _ {1}, a _ {1}) + \gamma \mathbb {E} _ {s _ {2} \sim \widehat {P} (\cdot | s _ {1}, a _ {1})} [ V _ {\theta} (s _ {2}) ] \bigg) \bigg ] \\ \end{array}
$$
where (i) and (iv) follows the definition of the optimal soft Q-function in (3b); (ii) follows the closed-form expression of  $V_{\theta}$  in (4); (iii) is from (4). By recursively applying the equalities (i) and (iv), we obtain the following result:
$$
\begin{array}{l} \nabla_ {\theta} Q _ {\theta} (s _ {0}, a _ {0}) \\ = \nabla_ {\theta} \left(r (s _ {0}, a _ {0}; \theta) + U (s _ {0}, a _ {0}) + \gamma \mathbb {E} _ {s _ {1} \sim \widehat {P} (\cdot | s _ {0}, a _ {0})} \left[ V _ {\theta} (s _ {1}) \right]\right) \\ = \nabla_ {\theta} r (s _ {0}, a _ {0}; \theta) + \gamma \mathbb {E} _ {s _ {1} \sim \widehat {P} (\cdot | s _ {0}, a _ {0}), a _ {1} \sim \pi_ {\theta} (\cdot | s _ {1})} \bigg [ \nabla_ {\theta} \bigg (r (s _ {1}, a _ {1}; \theta) + U (s _ {1}, a _ {1}) + \gamma \mathbb {E} _ {s _ {2} \sim \widehat {P} (\cdot | s _ {1}, a _ {1})} \big [ V _ {\theta} (s _ {2}) \big ] \bigg) \bigg ] \\ = \mathbb {E} _ {\tau^ {\mathrm {A}} \sim (\pi_ {\theta}, \widehat {P})} \left[ \sum_ {t = 0} ^ {\infty} \gamma^ {t} \nabla_ {\theta} r \left(s _ {t}, a _ {t}; \theta\right) \mid s _ {0}, a _ {0} \right]. \tag {35} \\ \end{array}
$$
Finally, by plugging (35) into (34), we obtain the gradient expression of the surrogate objective  $\widehat{L} (\theta)$  as below:
$$
\begin{array}{l} \nabla \widehat {L} (\theta) = \mathbb {E} _ {\tau^ {\mathrm {E}} \sim (\eta , \pi^ {\mathrm {E}}, P)} \left[ \sum_ {t = 0} ^ {\infty} \gamma^ {t} \nabla_ {\theta} r (s _ {t}, a _ {t}; \theta) \right] - \mathbb {E} _ {s _ {0} \sim \eta (\cdot), a _ {0} \sim \pi_ {\theta} (\cdot | s _ {0})} \left[ \nabla_ {\theta} Q _ {\theta} (s _ {0}, a _ {0}) \right] \\ = \mathbb {E} _ {\tau^ {\mathrm {E}} \sim (\eta , \pi^ {\mathrm {E}}, P)} \left[ \sum_ {t = 0} ^ {\infty} \gamma^ {t} \nabla_ {\theta} r (s _ {t}, a _ {t}; \theta) \right] - \mathbb {E} _ {s _ {0} \sim \eta (\cdot), a _ {0} \sim \pi_ {\theta} (\cdot | s _ {0})} \left[ \mathbb {E} _ {\tau^ {\mathrm {A}} \sim (\pi_ {\theta}, \widehat {P})} \left[ \sum_ {t = 0} ^ {\infty} \gamma^ {t} \nabla_ {\theta} r (s _ {t}, a _ {t}; \theta) \mid s _ {0}, a _ {0} \right] \right] \\ = \mathbb {E} _ {\tau^ {\mathrm {E}} \sim (\eta , \pi^ {\mathrm {E}}, P)} \left[ \sum_ {t = 0} ^ {\infty} \gamma^ {t} \nabla_ {\theta} r \left(s _ {t}, a _ {t}; \theta\right) \right] - \mathbb {E} _ {\tau^ {\mathrm {A}} \sim (\eta , \pi_ {\theta}, \widehat {P})} \left[ \sum_ {t = 0} ^ {\infty} \gamma^ {t} \nabla_ {\theta} r \left(s _ {t}, a _ {t}; \theta\right) \right]. \tag {36} \\ \end{array}
$$
G Proof of Lemma 4
Suppose Assumptions 2 - 3 hold, we prove the Lipschitz continuous property of the optimal soft Q-function  $Q_{\theta}$  in (19a) and the Lipschitz smooth property of the surrogate objective  $\widehat{L} (\theta)$  in (19b) respectively.
G.1 Proof of Inequality (19a)
Proof. In order to show that the optimal soft Q-function  $Q_{\theta}(s,a)$  is Lipschitz continuous w.r.t. the reward parameter  $\theta$ . for any state-action pair  $(s,a)$ , we take two steps to finish the proof. First, we show that  $Q_{\theta}$  has bounded gradient under any reward parameter  $\theta$ . Then we could use the mean value theorem to complete the proof.
According to the gradient expression of the optimal soft Q-function  $Q_{\theta}$  in (35), the following result holds:
$$
\begin{array}{l} \| \nabla_ {\theta} Q _ {\theta} (s, a) \| = \left\| \mathbb {E} _ {\tau^ {\mathrm {A}} \sim (\pi_ {\theta}, \widehat {P})} \left[ \sum_ {t = 0} ^ {\infty} \gamma^ {t} \nabla_ {\theta} r \left(s _ {t}, a _ {t}; \theta\right) \mid s _ {0} = s, a _ {0} = a \right] \right\| \\ \stackrel {(i)} {\leq} \mathbb {E} _ {\tau^ {\mathrm {A}} \sim (\pi_ {\theta}, \widehat {P})} \left[ \sum_ {t = 0} ^ {\infty} \gamma^ {t} \| \nabla_ {\theta} r (s _ {t}, a _ {t}; \theta) \| \mid s _ {0} = s, a _ {0} = a \right] \\ \stackrel {(i i)} {\leq} \mathbb {E} _ {\tau^ {\mathrm {A}} \sim (\pi_ {\theta}, \widehat {P})} \left[ \sum_ {t = 0} ^ {\infty} \gamma^ {t} L _ {r} \mid s _ {0} = s, a _ {0} = a \right] \\ = \frac {L _ {r}}{1 - \gamma} \tag {37} \\ \end{array}
$$
where (i) follows the Jensen's inequality and (ii) follows from the assumption that the reward gradient is bounded; see (18) in Assumption 3. Denote  $L_{q} \coloneqq \frac{L_{r}}{1 - \gamma}$ , we are able to show the Lipschitz continuous property of the optimal soft Q-function  $Q_{\theta}$  as below:
$$
\left| Q _ {\theta_ {1}} (s, a) - Q _ {\theta_ {2}} (s, a) \right| \stackrel {(i)} {=} \left| \langle \theta_ {1} - \theta_ {2}, \nabla_ {\theta} Q _ {\tilde {\theta}} (s, a) \rangle \right| \leq \left\| \theta_ {1} - \theta_ {2} \right\| \cdot \left\| \nabla_ {\theta} Q _ {\tilde {\theta}} (s, a) \right\| \stackrel {(i i)} {\leq} L _ {q} \| \theta_ {1} - \theta_ {2} \|
$$
where  $\tilde{\theta}$  is a convex combination between  $\theta_{1}$  and  $\theta_{2}$ . Moreover, (i) is from the mean value theorem and (ii) follows (37).
G.2 Proof of Inequality (19b)
Proof. In Lemma 3, we have shown the expression of the gradient of the surrogate objective  $\widehat{L}(\theta)$ . Then for any reward parameters  $\theta_{1}$  and  $\theta_{2}$ , we are able to obtain the following result: