SlowGuess/NeurIPS_2023_OCR huggingface dataset:
We then choose  $r_{\pmb{\theta}}(A_{\star})$  to be defined such that, for all  $A \in \mathcal{B}_{\mathrm{F}}(A_{\star}; r_{\pmb{\theta}}(A_{\star}))$ , we have  $\lambda_{\min}(\nabla_{\pmb{\theta}}^{2}\mathcal{J}(\pmb{\theta}; A)|_{\pmb{\theta} = \pmb{\theta}_{\star}(A)}) \geq \frac{1}{2}\lambda_{\min}(\nabla_{\pmb{\theta}}^{2}\mathcal{J}(\pmb{\theta}; A_{\star})|_{\pmb{\theta} = \pmb{\theta}_{\star}(A_{\star}))$ . By Lemma D.1, we know that  $\nabla_{\pmb{\theta}}^{2}\mathcal{J}(\pmb{\theta}; A)$  is continuous and furthermore we know that eigenvalues are continuous. Using the gradient bounds given in Lemma D.1 to bound the Lipschitz constant of  $\nabla_{\pmb{\theta}}^{2}\mathcal{J}(\pmb{\theta}; A)$ , it follows that we can take
$$
r _ {\pmb {\theta}} (A _ {\star}) = \mathrm {p o l y} \left(\frac {1}{\lambda_ {\mathrm {m i n}} (\nabla_ {\pmb {\theta}} ^ {2} \mathcal {J} (\pmb {\theta} ; A _ {\star}) | _ {\pmb {\theta} = \pmb {\theta} _ {\star} (A _ {\star}))}}, \| A _ {\star} \| _ {\mathrm {o p}}, B _ {\phi}, L _ {\phi}, L _ {\pmb {\theta}}, L _ {\mathrm {c o s t}}, \sigma_ {\pmb {w}} ^ {- 1}, H, d _ {\pmb {x}}\right) ^ {- 1}.
$$
Bounding Norm of Gradients. Fix  $A \in \mathcal{B}_{\mathrm{F}}(A_{\star}; r_{\theta}(A_{\star}))$ . We know that  $\theta_{\star}(A)$  satisfies
$$
\nabla_ {\boldsymbol {\theta}} \mathcal {J} (\boldsymbol {\theta}; A) | _ {\boldsymbol {\theta} = \boldsymbol {\theta} _ {\star} (A)} = 0.
$$
We wish to differentiate  $\theta_{\star}(A)$  with respect to  $A$ , and bound the magnitude of up to the third derivative. Similar to the proof of Lemma D.1, we let  $A_{t} \coloneqq A + t_{1}\Delta_{1}^{A} + t_{2}\Delta_{2}^{A} + t_{3}\Delta_{3}^{A}$  for some  $\Delta_i^A$  satisfying  $\| \Delta_i^A\|_{\mathrm{op}} \leq 1$ . As noted in the proof of Lemma D.1, we have
$$
\frac {\mathrm {d}}{\mathrm {d} t _ {i}} \boldsymbol {\theta} _ {\star} (A _ {t}) | _ {t = 0} = \nabla_ {A} \boldsymbol {\theta} _ {\star} (A) [ \Delta_ {i} ^ {A} ]
$$
(and similarly for higher-order derivatives). Thus, to show the result, it suffices to show that  $\theta_{\star}(A_t)$  is differentiable in  $t_1, t_2, t_3$  for all  $\Delta_i^A$ , and to bound the magnitude of this derivative for all  $\Delta_i^A$  with  $\| \Delta_i^A \|_{\mathrm{op}} \leq 1$ . We have
$$
\begin{array}{l} \frac {\mathrm {d}}{\mathrm {d} t _ {1}} \nabla_ {\boldsymbol {\theta}} \mathcal {J} (\boldsymbol {\theta}; A _ {t}) | _ {\boldsymbol {\theta} = \boldsymbol {\theta} _ {*} (A _ {t})} | _ {t = 0} = 0 \\ \Longrightarrow \underbrace {\nabla_ {A ^ {\prime}} \nabla_ {\boldsymbol {\theta}} \mathcal {J} \left(\boldsymbol {\theta} ; A ^ {\prime}\right) \mid_ {\boldsymbol {\theta} = \boldsymbol {\theta} _ {\star} (A) , A ^ {\prime} = A} \left[ \Delta_ {1} ^ {A} \right]} _ {=: G _ {1} \left(A, \Delta_ {1} ^ {A}\right)} + \nabla_ {\boldsymbol {\theta}} ^ {2} \mathcal {J} (\boldsymbol {\theta}; A) \mid_ {\boldsymbol {\theta} = \boldsymbol {\theta} _ {\star} (A)} \cdot \nabla_ {A} \boldsymbol {\theta} _ {\star} (A) \left[ \Delta_ {1} ^ {A} \right] = 0 \quad (D. 1) \\ \end{array}
$$
$$
\nabla_ {A} \boldsymbol {\theta} _ {\star} (A) [ \Delta_ {1} ^ {A} ] = - \left(\nabla_ {\boldsymbol {\theta}} ^ {2} \mathcal {J} (\boldsymbol {\theta}; A) | _ {\boldsymbol {\theta} = \boldsymbol {\theta} _ {\star} (A)}\right) ^ {- 1} \cdot G _ {1} (A, \Delta_ {1} ^ {A})
$$
which is well-defined since we have assumed that  $\nabla_{\pmb{\theta}}^{2}\mathcal{I}(\pmb {\theta};A)|_{\pmb {\theta} = \pmb{\theta}_{\star}(A)}$  is full-rank, and  $\mathcal{J}$  is differentiable in both its arguments by Lemma D.1. To compute the second derivative of  $\pmb{\theta}_{\star}$ , we differentiate through (D.1) which gives
$$
\begin{array}{l} \frac {\mathrm {d}}{\mathrm {d} t _ {2}} \left(G _ {1} (A _ {t}, \Delta_ {1} ^ {A}) + \nabla_ {\pmb {\theta}} ^ {2} \mathcal {J} (\pmb {\theta}; A _ {t}) | _ {\pmb {\theta} = \pmb {\theta} _ {\star} (A _ {t})} \cdot \nabla_ {A} \pmb {\theta} _ {\star} (A _ {t}) [ \Delta_ {1} ^ {A} ]\right) | _ {t = 0} = 0 \\ \Longrightarrow \underbrace {\frac {\mathrm {d}}{\mathrm {d} t _ {2}} \left(G _ {1} (A _ {t} , \Delta_ {1} ^ {A}) + \nabla_ {\boldsymbol {\theta}} ^ {2} \mathcal {I} (\boldsymbol {\theta} ; A _ {t}) | _ {\boldsymbol {\theta} = \boldsymbol {\theta} _ {\star} (A _ {t})} \cdot \nabla_ {A} \boldsymbol {\theta} _ {\star} (A) [ \Delta_ {1} ^ {A} ]\right) | _ {\boldsymbol {t} = 0}} _ {=: G _ {2} (A, \Delta_ {1} ^ {A}, \Delta_ {2} ^ {A})} \\ + \nabla_ {\pmb {\theta}} ^ {2} \mathcal {J} (\pmb {\theta}; A) | _ {\pmb {\theta} = \pmb {\theta} _ {\star} (A)} \cdot \nabla_ {A} ^ {2} \pmb {\theta} _ {\star} (A) [ \Delta_ {1} ^ {A}, \Delta_ {2} ^ {A} ] = 0. \\ \end{array}
$$
Note that  $G_{2}(A, \Delta_{1}^{A}, \Delta_{2}^{A})$  involves at most a third-order derivative of  $\mathcal{J}(\theta; A)$  and first-order derivative of  $\theta_{\star}(A)$ , both of which we know exist by Lemma D.1 and what we showed above. This then further implies
$$
\nabla_ {A} ^ {2} \boldsymbol {\theta} _ {\star} (A) [ \Delta_ {1} ^ {A}, \Delta_ {2} ^ {A} ] = - \left(\nabla_ {\boldsymbol {\theta}} ^ {2} \mathcal {J} (\boldsymbol {\theta}; A) | _ {\boldsymbol {\theta} = \boldsymbol {\theta} _ {\star} (A)}\right) ^ {- 1} \cdot G _ {2} (A, \Delta_ {1} ^ {A}, \Delta_ {2} ^ {A}),
$$
which is well-defined since we have assumed that  $\nabla_{\pmb{\theta}}^{2}\mathcal{J}(\pmb {\theta};A)|_{\pmb {\theta} = \pmb{\theta}_{\star}(A)}$  is full-rank. Finally, we compute
$$
\begin{array}{l} \frac {\mathrm {d}}{\mathrm {d} t _ {3}} \left(G _ {2} \left(A _ {t}, \Delta_ {1} ^ {A}, \Delta_ {2} ^ {A}\right) + \nabla_ {\boldsymbol {\theta}} ^ {2} \mathcal {J} (\boldsymbol {\theta}; A _ {t}) | _ {\boldsymbol {\theta} = \boldsymbol {\theta} _ {\star} (A _ {t})} \cdot \nabla_ {A} ^ {2} \boldsymbol {\theta} _ {\star} (A _ {t}) [ \Delta_ {1} ^ {A}, \Delta_ {2} ^ {A} ]\right) | _ {t = 0} = 0 \\ \Longrightarrow \underbrace {\frac {\mathrm {d}}{\mathrm {d} t _ {3}} \left(G _ {2} \left(A _ {t} , \Delta_ {1} ^ {A} , \Delta_ {2} ^ {A}\right) + \nabla_ {\boldsymbol {\theta}} ^ {2} \mathcal {I} (\boldsymbol {\theta} ; A _ {t}) | _ {\boldsymbol {\theta} = \boldsymbol {\theta} _ {\star} (A _ {t})} \cdot \nabla_ {A} ^ {2} \boldsymbol {\theta} _ {\star} (A) \left[ \Delta_ {1} ^ {A} , \Delta_ {2} ^ {A} \right]\right) \mid_ {t = 0}} _ {=: G _ {3} \left(A, \Delta_ {1} ^ {A}, \Delta_ {2} ^ {A}, \Delta_ {3} ^ {A}\right)} \\ + \nabla_ {\boldsymbol {\theta}} ^ {2} \mathcal {J} (\boldsymbol {\theta}; A) | _ {\boldsymbol {\theta} = \boldsymbol {\theta} _ {\star} (A)} \cdot \nabla_ {A} ^ {3} \boldsymbol {\theta} _ {\star} (A) [ \Delta_ {1} ^ {A}, \Delta_ {2} ^ {A}, \Delta_ {3} ^ {A} ] = 0. \\ \end{array}
$$
Note that  $G_{3}(A, \Delta_{1}^{A}, \Delta_{2}^{A}, \Delta_{3}^{A})$  involves at most a fourth-order derivative of  $\mathcal{J}(\pmb{\theta}; A)$  and second-order derivative of  $\pmb{\theta}_{\star}(A)$ , both of which we know exist by Lemma D.1 and what we showed above. We therefore have
$$
\nabla_ {A} ^ {3} \boldsymbol {\theta} _ {\star} (A) [ \Delta_ {1} ^ {A}, \Delta_ {2} ^ {A}, \Delta_ {3} ^ {A} ] = - \left(\nabla_ {\boldsymbol {\theta}} ^ {2} \mathcal {J} (\boldsymbol {\theta}; A) | _ {\boldsymbol {\theta} = \boldsymbol {\theta} _ {\star} (A)}\right) ^ {- 1} \cdot G _ {3} (A, \Delta_ {1} ^ {A}, \Delta_ {2} ^ {A}, \Delta_ {3} ^ {A})
$$
which is well-defined since we have assumed that  $\nabla_{\pmb{\theta}}^{2}\mathcal{J}(\pmb{\theta};A)|_{\pmb{\theta} = \pmb{\theta}_{\star}(A)}$  is full-rank. As each of these expressions is defined for all choice of  $\Delta_i^A$ , the differentiability of  $\pmb{\theta}_{\star}(A)$  follows.
Note that the above expressions for  $\nabla_A\pmb {\theta}_\star (A)[\Delta_1^A ],\nabla_A^2\pmb {\theta}_\star (A)[\Delta_1^A,\Delta_2^A ]$  and  $\nabla_A^3\pmb {\theta}_\star (A)[\Delta_1^A,\Delta_2^A,\Delta_3^A ]$  all depend on at most a fourth derivative of  $\mathcal{J}(\pmb {\theta};A)$  , as well as  $(\nabla_{\pmb{\theta}}^{2}\mathcal{J}(\pmb {\theta};A)|_{\pmb {\theta} = \pmb{\theta}_{\star}(A)})^{-1}$  . The norm bounds are then a direct consequence of Lemma D.1.
Proof of Proposition 6. By Lemma D.1 we have that  $\mathcal{J}(\pmb{\theta}; A)$  is four-times differentiable in its arguments. Since we have assumed  $\nabla_{\pmb{\theta}}^2 \mathcal{J}(\pmb{\theta}; A_\star) |_{\pmb{\theta} = \pmb{\theta}_\star(A_\star)} \succ 0$ , by the Implicit Function Theorem [60], it follows that there exists some  $r_\theta' > 0$  and mapping  $\widetilde{\pmb{\theta}}(A)$  such that, for all  $A \in \mathcal{B}_{\mathrm{F}}(A_\star; r_\theta')$ ,  $\nabla_{\pmb{\theta}} \mathcal{J}(\pmb{\theta}; A) |_{\pmb{\theta} = \widetilde{\pmb{\theta}}(A)} = 0$ , and  $\widetilde{\pmb{\theta}}(A)$  is three-times differentiable.
Our goal is now to show that  $\widetilde{\theta}(A) = \theta_{\star}(A)$  for  $A$  close enough to  $A_{\star}$ . By the continuity of eigenvalues,  $\mathcal{J}(\theta; A)$ , and  $\widetilde{\theta}(A)$ , we have that there exists some  $r$  and  $r_{\theta}''$  such that, for all  $\theta \in \mathcal{B}_{\mathrm{F}}(\theta_{\star}(A_{\star}); r)$  and  $A \in \mathcal{B}_{\mathrm{F}}(A_{\star}; r_{\theta}^{\prime \prime})$ , we have  $\nabla_{\theta}^{2}\mathcal{J}(\theta; A) > 0$  and, furthermore,  $\widetilde{\theta}(A) \in \mathcal{B}_{2}(\theta_{\star}(A_{\star}); r/2)$  for all  $A \in \mathcal{B}_{\mathrm{F}}(A_{\star}; r_{\theta}^{\prime \prime})$ . This implies that  $\widetilde{\theta}(A)$  is strict local minimum of  $\mathcal{J}(\theta; A)$  and, in particular, that
$$
\mathcal {J} (\boldsymbol {\theta}; A) > \mathcal {J} (\widetilde {\boldsymbol {\theta}} (A); A), \quad \forall \boldsymbol {\theta} \in \mathcal {B} _ {2} (\boldsymbol {\theta} _ {\star} (A _ {\star}); r), \boldsymbol {\theta} \neq \widetilde {\boldsymbol {\theta}} (A).
$$
Let  $\Delta^{\star} \coloneqq \min_{\pmb{\theta} \notin \mathcal{B}_2(\pmb{\theta}_{\star}(A_{\star}); r)} \mathcal{J}(\pmb{\theta}; A_{\star}) - \mathcal{J}(\pmb{\theta}_{\star}(A_{\star}); A_{\star})$  and note that, since we have assumed the global minimum of  $\mathcal{J}(\pmb{\theta}; A_{\star})$  is unique, we have  $\Delta^{\star} > 0$ .
Fix some  $A \in \mathcal{B}_{\mathrm{F}}(A_{\star}; r_{\theta}^{\prime \prime})$  and assume that  $\widetilde{\pmb{\theta}}(A)$  is not the global minimum of  $\mathcal{J}(\pmb{\theta}; A)$ . This implies that  $\pmb{\theta}_{\star}(A)$ , the global minimum of  $\mathcal{J}(\pmb{\theta}; A)$ , is outside of  $\mathcal{B}_2(\pmb{\theta}_{\star}(A_{\star}); r)$ . Furthermore, by the continuity of  $\mathcal{J}(\pmb{\theta}; A)$ , we have, for some  $L, L' > 0$ ,
$$
\begin{array}{l} \mathcal {J} \left(\boldsymbol {\theta} _ {\star} (A); A _ {\star}\right) \leq \mathcal {J} \left(\boldsymbol {\theta} _ {\star} (A); A\right) + L \| A - A _ {\star} \| _ {\mathrm {F}} \\ \leq \mathcal {J} (\tilde {\boldsymbol {\theta}} (A); A) + L \| A - A _ {\star} \| _ {\mathrm {F}} \\ \leq \mathcal {J} (\tilde {\boldsymbol {\theta}} (A); A _ {\star}) + 2 L \| A - A _ {\star} \| _ {\mathrm {F}} \\ \leq \mathcal {J} \left(\boldsymbol {\theta} _ {\star} \left(A _ {\star}\right); A _ {\star}\right) + 2 L \| A - A _ {\star} \| _ {\mathrm {F}} + L \| \tilde {\boldsymbol {\theta}} (A) - \boldsymbol {\theta} _ {\star} \left(A _ {\star}\right) \| _ {2} \\ \leq \mathcal {J} \left(\boldsymbol {\theta} _ {\star} \left(A _ {\star}\right); A _ {\star}\right) + L ^ {\prime} \| A - A _ {\star} \| _ {\mathrm {F}}. \\ \end{array}
$$
$$
L ^ {\prime} r _ {\theta} ^ {\prime \prime} \geq L ^ {\prime} \| A - A _ {\star} \| _ {F} \geq \mathcal {J} \left(\boldsymbol {\theta} _ {\star} (A); A _ {\star}\right) - \mathcal {J} \left(\boldsymbol {\theta} _ {\star} \left(A _ {\star}\right); A _ {\star}\right) \geq \Delta^ {\star}.
$$
However, for  $r_{\theta}''$  small enough, this is a contradiction. Thus, it follows that  $\widetilde{\pmb{\theta}}(A)$  is the global minimum of  $\mathcal{J}(\pmb{\theta}; A)$ , so  $\widetilde{\pmb{\theta}}(A) = \pmb{\theta}_{\star}(A)$ .
The result then follows since we already have that  $\widetilde{\pmb{\theta}}(A)$  is three-times differentiable and satisfies  $\nabla_{\pmb{\theta}}\mathcal{J}(\pmb{\theta};A)|_{\pmb{\theta} = \widetilde{\pmb{\theta}}(A)} = 0$ , and by taking  $r_{\pmb{\theta}}(A_{\star})$  to be the minimum of  $r_{\pmb{\theta}}^{\prime}$  and  $r_{\pmb{\theta}}^{\prime \prime}$ . The boundedness of  $L_{\pi_{\star}}$  follows as in the proof of Proposition 5.
Proof of Proposition 7. Note that, by convexity and the KKT conditions, the solutions to  $\arg \min_{\pmb {\theta}\in \mathbb{R}^{d_{\pmb{\theta}}}}\mathcal{J}(\pmb {\theta};A)$  are described by
Thus, an equivalent definition for  $\theta_{\star}(A)$  is that it satisfies  $\nabla_{\pmb{\theta}}\mathcal{J}(\pmb {\theta};A)|_{\pmb {\theta} = \pmb{\theta}_{\star}(A)} = 0$  . Assumption 13 can then be shown to hold by an argument analogous to Proposition 5.
Proof of Lemma D.2. Let  $A(t) = t\widehat{A} + (1 - t)A_{\star}$  and  $g(t) \coloneqq \mathcal{J}(\pmb{\theta}_{\star}(A(t));A_{\star})$ . By Lemma D.1 and under Assumption 5, we have that both  $\mathcal{J}(\pmb{\theta};A)$  and  $\pmb{\theta}_{\star}(A)$  are three-times differentiable for all  $A = t\widehat{A} + (1 - t)A_{\star}, t \in [0,1]$ , so it follows that  $g(t)$  is three-times differentiable in  $t$ . We can therefore apply Taylor's Theorem to expand  $g(1)$  about the point  $t = 0$  to get:
$$
\begin{array}{l} g (1) = g (0) + \nabla_ {\pmb {\theta}} \mathcal {J} (\pmb {\theta}; A _ {\star}) | _ {\pmb {\theta} = \pmb {\theta} _ {\star} (A _ {\star})} \cdot \nabla_ {A} \pmb {\theta} _ {\star} (A) | _ {A = A _ {\star}} [ \widehat {A} - A _ {\star} ] \\ + \nabla_ {A} \theta_ {\star} (A) | _ {A = A _ {\star}} ^ {\top} \nabla_ {\boldsymbol {\theta}} ^ {2} \mathcal {J} (\boldsymbol {\theta}; A _ {\star}) | _ {\boldsymbol {\theta} = \theta_ {\star} (A _ {\star})} \nabla_ {A} \theta_ {\star} (A) | _ {A = A _ {\star}} [ \widehat {A} - A _ {\star}, \widehat {A} - A _ {\star} ] \\ + \nabla_ {\pmb {\theta}} \mathcal {J} (\pmb {\theta}; A _ {\star}) | _ {\pmb {\theta} = \pmb {\theta} _ {\star} (A _ {\star})} \cdot \nabla_ {A} ^ {2} \pmb {\theta} _ {\star} (A) | _ {A = A _ {\star}} [ \widehat {A} - A _ {\star}, \widehat {A} - A _ {\star} ] \\ + \nabla_ {A} ^ {3} \mathcal {J} \left(\boldsymbol {\theta} _ {\star} (A); A _ {\star}\right) | _ {A = A ^ {\prime}} [ \widehat {A} - A _ {\star}, \widehat {A} - A _ {\star}, \widehat {A} - A _ {\star} ] \\ \end{array}
$$
where  $A' = A(t')$  for some  $t' \in [0,1]$ . Under Assumption 13, we have that  $\nabla_{\pmb{\theta}} \mathcal{J}(\pmb{\theta}; A_{\star})|_{\pmb{\theta} = \pmb{\theta}_{\star}(A_{\star})} = 0$ , which implies that, plugging in the definition of  $g(1)$  and  $g(0)$ ,
$$
\begin{array}{l} \mathcal {J} \left(\boldsymbol {\theta} _ {\star} (\widehat {A}); A _ {\star}\right) = \mathcal {J} \left(\boldsymbol {\theta} _ {\star} \left(A _ {\star}\right); A _ {\star}\right) + \nabla_ {A} \boldsymbol {\theta} _ {\star} (A) | _ {A = A _ {\star}} ^ {\top} \nabla_ {\boldsymbol {\theta}} ^ {2} \mathcal {J} \left(\boldsymbol {\theta}; A _ {\star}\right) | _ {\boldsymbol {\theta} = \boldsymbol {\theta} _ {\star} \left(A _ {\star}\right)} \nabla_ {A} \boldsymbol {\theta} _ {\star} (A) | _ {A = A _ {\star}} \left[ \widehat {A} - A _ {\star}, \widehat {A} - A _ {\star} \right] \\ + \nabla_ {A} ^ {3} \mathcal {J} \left(\boldsymbol {\theta} _ {\star} (A); A _ {\star}\right) | _ {A = A ^ {\prime}} [ \widehat {A} - A _ {\star}, \widehat {A} - A _ {\star}, \widehat {A} - A _ {\star} ]. \\ \end{array}
$$
We can bound
$$
| \nabla_ {A} ^ {3} \mathcal {J} (\pmb {\theta} _ {\star} (A); A _ {\star}) | _ {A = A ^ {\prime}} [ \widehat {A} - A _ {\star}, \widehat {A} - A _ {\star}, \widehat {A} - A _ {\star} ] | \leq \| \nabla_ {A} ^ {3} \mathcal {J} (\pmb {\theta} _ {\star} (A); A _ {\star}) | _ {A = A ^ {\prime}} \| _ {\mathrm {o p}} \cdot \| \widehat {A} - A _ {\star} \| _ {\mathrm {o p}} ^ {3}.
$$
The expression for  $\nabla_A^3\mathcal{J}(\pmb {\theta}_\star (A);A_\star)$  contains up to the third derivative of both  $\mathcal{J}(\pmb {\theta};A_{\star})$  and  $\pmb{\theta}_{\star}(A)$ . By Lemma D.1 and under Assumption 13, since  $A^{\prime}\in \mathcal{B}_{\mathrm{F}}(A_{\star};\min \{r_{\mathrm{cost}}(A_{\star}),r_{\pmb{\theta}}(A_{\star})\})$  by construction, we can then bound
$$
\left\| \nabla_ {A} ^ {3} \mathcal {J} \left(\boldsymbol {\theta} _ {\star} (A); A _ {\star}\right) \right| _ {A = A ^ {\prime}} \| _ {\mathrm {o p}} \leq \operatorname {p o l y} \left(L _ {\pi_ {\star}}, \| A _ {\star} \| _ {\mathrm {o p}}, B _ {\phi}, L _ {\phi}, L _ {\boldsymbol {\theta}}, L _ {\text {c o s t}}, \sigma_ {\boldsymbol {w}} ^ {- 1}, H, d _ {\boldsymbol {x}}\right).
$$
The result follows by the definition of  $\mathcal{H}(A_{\star})$
Proof of Lemma D.3. Recall that  $\mathcal{H}(\widehat{A}) = \nabla_A^2\mathcal{J}(\theta_\star(A);\widehat{A})|_{A=\widehat{A}}$ . To prove this, we will use that this is differentiable by Lemma D.1, and will apply Taylor's Theorem.
First, note that by Taylor's Theorem we have
$$
\left. \nabla_ {A} ^ {2} \mathcal {J} (\pmb {\theta} _ {\star} (A); \widehat {A}) \right| _ {A = \widehat {A}} = \left. \nabla_ {A} ^ {2} \mathcal {J} (\pmb {\theta} _ {\star} (A); \widehat {A}) \right| _ {A = A _ {\star}} + \left. \nabla_ {A} ^ {3} \mathcal {J} (\pmb {\theta} _ {\star} (A); \widehat {A}) \right| _ {A = A ^ {\prime}} [ \widehat {A} - A _ {\star} ]
$$
for  $A' = t\widehat{A} + (1 - t)A_{\star}$  for some  $t \in [0,1]$ . The third derivative of  $\mathcal{J}(\theta_{\star}(A); \widehat{A})$  will involve up to the third derivative of both  $\mathcal{J}(\theta; \widehat{A})$  and  $\theta_{\star}(A)$ , so using Lemma D.1 and Assumption 13, since  $A' \in \mathcal{B}_{\mathrm{F}}(A_{\star}; \min\{r_{\mathrm{cost}}(A_{\star}), r_{\theta}(A_{\star})\})$  by assumption, we can bound
$$
\| \nabla_ {A} ^ {3} \mathcal {J} (\pmb {\theta} _ {\star} (A); \widehat {A}) | _ {A = A ^ {\prime}} [ \widehat {A} - A _ {\star} ] \| _ {\mathrm {o p}} \leq \mathrm {p o l y} (L _ {\pi_ {\star}}, \| A _ {\star} \| _ {\mathrm {o p}}, B _ {\phi}, L _ {\phi}, L _ {\pmb {\theta}}, L _ {\mathrm {c o s t}}, \sigma_ {\pmb {w}} ^ {- 1}, H, d _ {\pmb {x}}) \cdot \| \widehat {A} - A _ {\star} \| _ {\mathrm {o p}}.
$$
Next, we wish to relate  $\nabla_A^2\mathcal{J}(\pmb {\theta}_\star (A);\widehat{A})|_{A = A_\star}$  to  $\nabla_A^2\mathcal{J}(\pmb {\theta}_\star (A);A_\star)|_{A = A_\star} = \mathcal{H}(A_\star)$ . Again applying Taylor's Theorem, we have
$$
\nabla_ {A} ^ {2} \mathcal {J} (\boldsymbol {\theta} _ {\star} (A); \widehat {A}) | _ {A = A _ {\star}} = \nabla_ {A} ^ {2} \mathcal {J} (\boldsymbol {\theta} _ {\star} (A); A _ {\star}) | _ {A = A _ {\star}} + \nabla_ {A ^ {\prime}} \nabla_ {A} ^ {2} \mathcal {J} (\boldsymbol {\theta} _ {\star} (A); A ^ {\prime}) | _ {A = A _ {\star}, A ^ {\prime} = A ^ {\prime \prime}} [ \widehat {A} - A _ {\star} ]
$$
for  $A'' = t\widehat{A} + (1 - t)A_{\star}$  for some  $t \in [0,1]$ . By Lemma D.1 and Assumption 13, we can bound
$$
\begin{array}{l} \left\| \nabla_ {A ^ {\prime}} \nabla_ {A} ^ {2} \mathcal {J} \left(\boldsymbol {\theta} _ {\star} (A); A ^ {\prime}\right) \right| _ {A = A _ {\star}, A ^ {\prime} = A ^ {\prime \prime}} \left[ \widehat {A} - A _ {\star} \right] \| _ {\mathrm {o p}} \\ \leq \operatorname {p o l y} \left(L _ {\pi_ {\star}}, \| A _ {\star} \| _ {\text {o p}}, L _ {\phi}, L _ {\boldsymbol {\theta}}, L _ {\text {c o s t}}, \sigma_ {\boldsymbol {w}} ^ {- 1}, H, d _ {\boldsymbol {x}}\right) \cdot \| \widehat {A} - A _ {\star} \| _ {\text {o p}}. \\ \end{array}
$$
Lemma D.4. Under Assumptions 1, 2, 4, 5 and 13, for all  $A \in \mathcal{B}_{\mathrm{F}}(A_{\star}; \min \{r_{\mathrm{cost}}(A_{\star}), r_{\theta}(A_{\star})\})$ , we can bound
$$
\| \mathcal {H} (A) \| _ {\mathrm {o p}} \leq \mathrm {p o l y} (\| A _ {\star} \| _ {\mathrm {o p}}, B _ {\phi}, L _ {\phi}, L _ {\theta}, L _ {\mathrm {c o s t}}, L _ {\pi_ {\star}}, \sigma_ {\boldsymbol {w}} ^ {- 1}, H, d _ {\boldsymbol {x}})
$$
Proof. Recall that  $\mathcal{H}(\widehat{A}) = \nabla_A^2\mathcal{J}(\pmb {\theta}_\star (A);\widehat{A})|_{A = \widehat{A}}$ . The bound then follows from Lemma D.1 and Assumption 13.
E High-Probability Regret Bounds in Nonlinear Systems
In this section, we modify the proof the main result of [17] slightly to show a high probability regret bound for  $\mathrm{LC^3}$ . For the sake of brevity, we omit details that are identical to the proof given in [17]. We will need the following assumption.
Assumption 14 (Bounded Cost). We assume that, for all trajectories  $\pmb{\tau}$ , we have  $\mathrm{cost}(\pmb{\tau}) \leq c_{\max}$ .
We adopt the notation used in this work, modifying somewhat the notation from [17]. In particular, we let  $\mathcal{J}(\pi; A)$  denote the expected cost of playing policy  $\pi$  under system  $A$ , and we set
$$
\boldsymbol {\Sigma} _ {t} = \sum_ {s = 1} ^ {t} \sum_ {h = 1} ^ {H} \phi \left(\boldsymbol {x} _ {h} ^ {t}, \boldsymbol {u} _ {h} ^ {t}\right) \phi \left(\boldsymbol {x} _ {h} ^ {t}, \boldsymbol {u} _ {h} ^ {t}\right) ^ {\top} + \lambda I
$$
denote the covariates obtained by the first  $t$  episodes of  $\mathrm{LC}^3$  (plus a regularizer). We let  $\pi^t$  denote the policy played at episode  $t$  of  $\mathrm{LC}^3$ . For a policy set  $\Pi$ , we define regret as
$$
\mathcal {R} _ {T} (\Pi) := \sum_ {t = 1} ^ {T} \mathcal {J} \left(\pi^ {t}; A _ {\star}\right) - T \cdot \min  _ {\pi \in \Pi} \mathcal {J} \left(\pi ; A _ {\star}\right).
$$
We will also denote  $\pi_{\star} := \arg \min_{\pi \in \Pi} \mathcal{J}(\pi; A_{\star})$ .
In addition to these notational changes, we modify  $\mathrm{LC}^3$  slightly to use the parameter
$$
\beta^ {t} := \sqrt {\lambda} B _ {A} + \sqrt {8 d _ {\boldsymbol {x}} \log 5 + 8 \log (T \det (\boldsymbol {\Sigma} _ {t}) \det (\boldsymbol {\Sigma} _ {0}) ^ {- 1} / \delta)}
$$
in the construction of the confidence set,  $\mathrm{BALL}^t$ .
Besides the aforementioned changes, in the following proofs we adopt the same notation as [17]. We have the following result.
Theorem 6. Under Assumptions 1 and 14 and with any policy class  $\Pi$ , with probability at least  $1 - \delta$ ,  $\mathrm{LC}^3$  has regret bounded as
$$
\mathcal {R} _ {T} (\Pi) \leq C \cdot c _ {\max } H \sqrt {d _ {\phi} \cdot (d _ {\phi} + d _ {\boldsymbol {x}} + B _ {A} + \log \frac {1}{\delta}) \cdot T} \cdot \log (1 + B _ {\phi} H T / \sigma_ {\boldsymbol {w}})
$$
Proof of Theorem 6. By Lemma E.2, we have that the event  $\mathcal{E}_1$  holds with probability at least  $1 - \delta$ . We therefore assume  $\mathcal{E}_1$  holds for the remainder of the proof.
By the definition of the confidence set in  $\mathrm{LC}^3$ , on  $\mathcal{E}_1$  we have that  $A_{\star}$  is the in confidence set for all  $t \leq T$ . It follows that on  $\mathcal{E}_1$ ,
$$
\begin{array}{l} \mathcal {R} _ {T} = \sum_ {t = 1} ^ {T} \left[ \mathcal {J} \left(\pi^ {t}; A _ {\star}\right) - \mathcal {J} \left(\pi_ {\star}; A _ {\star}\right) \right] \\ \stackrel {(a)} {\leq} \sum_ {t = 1} ^ {T} \left[ \mathcal {J} \left(\pi^ {t}; A _ {\star}\right) - \mathcal {J} \left(\pi_ {\star}; \widehat {A} ^ {t}\right) \right] \\ \stackrel {(b)} {\leq} \sum_ {t = 1} ^ {T} c _ {\max } \cdot \mathbb {E} _ {A _ {\star}, \pi^ {t}} \left[ \sum_ {h = 1} ^ {H} \min  \left\{\frac {1}{\sigma_ {\boldsymbol {w}}} \| \left(A _ {\star} - \widehat {A} ^ {t}\right) \cdot \phi \left(\boldsymbol {x} _ {h}, \boldsymbol {u} _ {h}\right) \| _ {2}, 1 \right\} \right] \tag {E.1} \\ \end{array}
$$
where  $(a)$  follows from the optimistic property of  $\mathrm{LC}^3$  when  $A_{\star}\in \mathrm{BALL}^{t}$ , and  $(b)$  follows from Lemma E.1. On  $\mathcal{E}_1$ , we have