SlowGuess/NeurIPS_2023_OCR huggingface dataset:
The quantity  $\mathcal{H}(A_{\star}) := \nabla_A^2 \mathcal{J}(\pi_{\star}(A); A_{\star})|_{A = A_{\star}}$ , referred to as the model-task Hessian in [2], corresponds to the curvature of the loss of the certainty-equivalence controller  $\pi_{\star}(A)$  around  $A \gets A_{\star}$ . It precisely quantifies how estimation error in each coordinate of  $A_{\star}$  translates into suboptimality of the controller—providing an answer to our question of which parameters are most relevant to learning a good controller—and reduces the problem of minimizing the controller loss to estimating  $A_{\star}$  in a particular norm. The following result gives a bound on this estimation error,  $\| \mathrm{vec}(A_{\star} - \widehat{A}) \|_{\mathcal{H}(A_{\star})}^2$ .
Proposition 2 (Informal). Consider interacting with (1.1) for  $T$  episodes, and let  $\Lambda_T = \sum_{t=1}^{T} \sum_{h=1}^{H} \phi(\pmb{x}_h^t, \pmb{u}_h^t) \phi(\pmb{x}_h^t, \pmb{u}_h^t)^\top$  denote the observed covariates and
$$
\widehat {A} = \arg \min _ {A} \sum_ {t = 1} ^ {T} \sum_ {h = 1} ^ {H} \| \pmb {x} _ {h + 1} ^ {t} - A \phi (\pmb {x} _ {h} ^ {t}, \pmb {u} _ {h} ^ {t}) \| _ {2} ^ {2}
$$
the least-squares estimate of  $A_{\star}$ . Recalling that  $\tilde{\Lambda}_T = I_{d_\infty} \otimes \Lambda_T$ , we have, with high probability:
$$
\left\| \operatorname {v e c} \left(A _ {\star} - \widehat {A}\right) \right\| _ {\mathcal {H} \left(A _ {\star}\right)} ^ {2} \lesssim \sigma_ {\boldsymbol {w}} ^ {2} \cdot \operatorname {t r} \left(\mathcal {H} \left(A _ {\star}\right) \check {\boldsymbol {\Lambda}} _ {T} ^ {- 1}\right).
$$
4.1 Algorithm and Upper Bound
Proposition 2 motivates our algorithmic approach: explore to collect covariates  $\Lambda_T$  minimizing  $\mathrm{tr}(\mathcal{H}(A_\star)\check{\Lambda}_T^{-1})$ . There are two primary challenges to achieving this: we do not know  $\mathcal{H}(A_\star)$ , as it depends on the (unknown) parameter  $A_\star$  and, even if we did know  $\mathcal{H}(A_\star)$ , it is not clear how to explore so as to collect data minimizing  $\mathrm{tr}(\mathcal{H}(A_\star)\check{\Lambda}_T^{-1})$ . We address both of these challenges in Algorithm 1.
Algorithm 1 Optimal Exploration in Nonlinear Systems (informal)
1: inputs: episodes  $T$ ,  $(\mathrm{cost}_h)_h = 1$  confidence  $\delta$ , control policies  $\Pi^{\star}$ , exploration policies  $\Pi_{\mathrm{exp}}$   
2:  $\widehat{A}^{1} \gets 0, \ell_{T} \gets \lceil \log_2 T / 8 \rceil, T_{\ell} \gets 2^{\ell}$   
3: for  $\ell = 1, 2, 3, \ldots, \ell_{T}$  do  
4: Compute estimate of model-task Hessian:  $\mathcal{H}_{\ell} \gets \mathcal{H}(\widehat{A}^{\ell})$   
5: Run DYNAMICOED on  $\Phi_{\ell}(\Lambda) \gets \operatorname{tr}(\mathcal{H}_{\ell} \cdot \Lambda^{-1})$  to learn exploration policies  $\Pi_{\ell} \subseteq \Pi_{\mathrm{exp}}$   
6: Rerun each policy in  $\Pi_{\ell} N_{\ell} = \lceil T_{\ell} / |\Pi_{\ell}| \rceil$  times, denote collected data  $\mathfrak{D}_{\ell}$   
7: Estimate  $A_{\star}$ :  $\widehat{A}^{\ell + 1} = \arg \min_A \sum_{h = 1}^H \sum_{(\boldsymbol{x}_{h + 1}, \boldsymbol{u}_h, \boldsymbol{x}_h) \in \mathfrak{D}_{\ell}} \| \boldsymbol{x}_{h + 1} - A\phi(\boldsymbol{x}_h, \boldsymbol{u}_h) \|_2^2$   
8: return  $\widehat{\pi}_{T} \gets \pi_{\star}(\widehat{A}^{\ell_{T} + 1}) \in \Pi^{\star}$
Algorithm 1 proceeds in epochs of exponentially increasing length. At each epoch it first approximates  $\mathcal{H}(A_{\star})$  by computing the model-task Hessian of the estimated system,  $\widehat{A}^{\ell}$ . Using this approximation of  $\mathcal{H}(A_{\star})$ , it seeks to explore to minimize  $\mathrm{tr}(\mathcal{H}(\widehat{A}^{\ell})\check{\Lambda}_{T}^{-1})$ . This exploration routine is encapsulated in the DYNAMICOED (dynamic optimal experiment design) function, an adaptive experiment-design routine inspired by recent work in reinforcement learning [56] and described in more detail in Section 5. DYNAMICOED returns a set of exploration policies,  $\Pi_{\ell}$ , which we run to collect data  $\mathfrak{D}_{\ell}$ . As we will show, the collected covariates,  $\pmb{\Lambda}_{\ell} := \sum_{h=1}^{H} \sum_{(\pmb{u}_h, \pmb{x}_h) \in \mathfrak{D}_{\ell}} \phi(\pmb{x}_h, \pmb{u}_h) \phi(\pmb{x}_h, \pmb{u}_h)^{\top}$ , satisfy
$$
\operatorname {t r} \left(\mathcal {H} \left(\widehat {A} ^ {\ell}\right) \check {\Lambda} _ {\ell} ^ {- 1}\right) \lesssim T _ {\ell} ^ {- 1} \cdot \min  _ {\boldsymbol {\Lambda} \in \Omega} \operatorname {t r} \left(\mathcal {H} \left(\widehat {A} ^ {\ell}\right) \check {\boldsymbol {\Lambda}} ^ {- 1}\right),
$$
which implies that DYNAMICOED collects data minimizing  $\mathrm{tr}(\mathcal{H}(\widehat{A}^{\ell})\tilde{\Lambda}_{\ell}^{-1})$  at a near-optimal rate. Given the data  $\mathfrak{D}_{\ell}$ , we form the least-squares estimate of  $A_{\star}$ ,  $\widehat{A}^{\ell +1}$ , and the process repeats. After running for  $T$  episodes, the certainty-equivalence controller on the last estimate obtained,  $\widehat{\pi}_T = \pi_\star (\widehat{A}^{\ell_T + 1})$ , is returned. The following result bounds the suboptimality of  $\widehat{\pi}_T$  as compared to  $\pi_{\star}(A_{\star})$ .
Theorem 1. Under Assumptions 1 to 6, if  $T \geq C_{\mathrm{poly}} \cdot \max \{1, r_{\mathrm{cost}}(A_\star)^{-2}, r_\theta (A_\star)^{-2}\}$ , then with probability at least  $1 - \delta$ , Algorithm 1 explores with policies in  $\Pi_{\mathrm{exp}}$  at every episode, runs for at most  $T$  episodes, and returns  $\widehat{\pi}_T \in \Pi^\star$  satisfying:
$$
\mathcal {J} (\widehat {\pi} _ {T}; A _ {\star}) - \mathcal {J} (\pi_ {\star} (A _ {\star}); A _ {\star}) \leq \frac {\sigma_ {\boldsymbol {w}} ^ {2}}{T} \cdot \min _ {\boldsymbol {\Lambda} \in \Omega} \mathrm {t r} \left(\mathcal {H} (A _ {\star}) \breve {\boldsymbol {\Lambda}} ^ {- 1}\right) \cdot C \log \frac {6 d _ {\boldsymbol {x}} d _ {\phi}}{\delta} + \frac {C _ {\mathrm {p o l y}}}{T ^ {3 / 2}}
$$
where we recall  $\Omega$  is the set of possible expected covariates on (1.1),  $C$  is a universal constant, and
$$
C _ {\mathrm {p o l y}} = \operatorname {p o l y} \left(d _ {\phi}, d _ {\boldsymbol {x}}, H, B _ {A}, B _ {\phi}, L _ {\phi}, L _ {\boldsymbol {\theta}}, L _ {\mathrm {c o s t}}, L _ {\pi_ {\star}}, \sigma_ {\boldsymbol {w}}, \sigma_ {\boldsymbol {w}} ^ {- 1}, \frac {1}{\lambda_ {\min} ^ {\star}}, \log \frac {T}{\delta}\right).
$$
Theorem 1 shows that Algorithm 1 is able to explore so as to optimally minimize the exploration loss  $\mathrm{tr}(\mathcal{H}(A_{\star})\check{\Lambda}_T^{-1})$ , up to a lower-order term scaling as  $T^{-3/2}$  and polynomially in system parameters. While Propositions 1 and 2 show that collecting data which minimizes  $\mathrm{tr}(\mathcal{H}(A_{\star})\check{\Lambda}_T^{-1})$  is in some sense fundamental, it is not clear it is necessary. We next show that it is indeed necessary.
4.2 Lower Bounds on Learning Controllers
Our goal is to show that, up to constants and lower-order terms, the bound given in Theorem 1 is not improvable, regardless of which controller estimate we use. To obtain such lower bounds, we need several additional assumptions. In particular, we require that the loss  $\mathcal{J}(\pi^{\theta};A)$  grows quadratically in the distance  $\pmb{\theta}$  is from  $\theta_{\star}(A)$ , and strengthen Assumption 3 to ensure (1.1) is sufficiently easy to excite. Formal statements of these conditions are given in Appendix F. Our lower bound is as follows.
Theorem 2 (Informal). Under Assumptions 1 to 6 and the additional regularity assumptions mentioned above, as long as  $T \geq C_{\mathrm{lb}}$ , for any  $\omega_{\mathrm{exp}} \in \triangle_{\Pi_{\mathrm{exp}}}$ , we have
$$
\min  _ {\widehat {\pi}} \max  _ {A \in \mathcal {B} _ {T}} \mathbb {E} _ {\mathfrak {D} _ {T} \sim A, \omega_ {\exp}} [ \mathcal {J} (\widehat {\pi} (\mathfrak {D} _ {T}); A) - \mathcal {J} (\pi_ {\star} (A); A) ] \geq \frac {\sigma_ {\boldsymbol {w}} ^ {2}}{3 T} \cdot \min  _ {\boldsymbol {\Lambda} \in \boldsymbol {\Omega}} \operatorname {t r} (\mathcal {H} (A _ {\star}) \breve {\boldsymbol {\Lambda}} ^ {- 1}) - \frac {C _ {\mathrm {l b}}}{T ^ {5 / 4}}
$$
for  $\mathcal{B}_T\coloneqq \mathcal{B}_{\mathrm{F}}(A_\star ;\mathcal{O}(T^{-5 / 6}))$ $\mathbb{E}_{\mathfrak{D}_T\sim A,\omega_{\mathrm{exp}}}[\cdot ] = \mathbb{E}_{\pi_{\mathrm{exp}}\sim \omega_{\mathrm{exp}}}\bigl [\mathbb{E}_{\mathfrak{D}_T\sim A,\pi_{\mathrm{exp}}}[\cdot ]\bigr ]$  the expectation over trajectories generated by running policies  $\pi \sim \omega_{\mathrm{exp}}$  on system  $A$  for  $T$  episodes,  $\hat{\pi}$  any mapping from observations to policies in  $\Pi^{\star}$ , and  $C_\mathrm{lb}$  some value scaling polynomially in problem parameters.
Note that this lower bound holds for any  $A_{\star}$  and mapping  $\phi$ , as long as our assumptions are met. Up to constants and lower-order terms, the scaling of Theorem 2 matches that of Theorem 1—both scale with  $\min_{\Lambda \in \Omega} \operatorname{tr}(\mathcal{H}(A_{\star}) \tilde{\Lambda}^{-1})$ —which implies that Algorithm 1 is indeed optimal (under certain additional regularity conditions). To the best of our knowledge, this is the first result characterizing the optimal statistical rates for learning in nonlinear dynamical systems. We emphasize that Theorem 2 holds for any decision rule  $\widehat{\pi}$ —it does not require that we use the certainty equivalence decision rule. As Algorithm 1 does rely on certainty equivalence, this result also implies that the certainty equivalence decision rule is optimal for (certain classes of) nonlinear dynamical systems.
The proof of Theorem 2 builds on the work [2], which shows a similar result for linear dynamical systems. It critically relies on our quadratic decomposition of the controller loss in Proposition 1, which reduces the problem of obtaining a lower bound on controller loss to a lower bound on estimating  $A_{\star}$  in the  $\mathcal{H}(A_{\star})$  norm. Given this, the result can be obtained by applying lower bounds on regression in general norms.
5 Optimal Experiment Design in Arbitrary Dynamical Systems
We turn now to the DYNAMICOED routine, which is the key algorithmic tool we use to prove Theorem 1. DYNAMICOED is a general reduction from policy optimization to optimal experiment design in arbitrary dynamical systems, and is an extension of a recently proposed approach for experiment design in linear MDPs [56]. This section may be of independent interest.
To illustrate the generality of this reduction, in this section we consider the following system:
$$
\boldsymbol {x} _ {h + 1} = f _ {h} \left(\boldsymbol {x} _ {h}, \boldsymbol {u} _ {h}, \boldsymbol {w} _ {h}\right), \quad h = 1, 2, \dots , H, \tag {5.1}
$$
where  $\boldsymbol{x}_h \in \mathcal{X} \subseteq \mathbb{R}^{d_\infty}$  denotes the state,  $\boldsymbol{u}_h \in \mathcal{U} \subseteq \mathbb{R}^{d_u}$  the input, and  $\boldsymbol{w}_h \in \mathbb{R}^{d_w}$  the noise. We take the dynamics  $(f_h)_{h=1}^H$  to be unknown and arbitrary. We assume there is some known featurization of
our system that is of interest,  $\phi (\pmb {x},\pmb {u})\to \mathbb{R}^{d_{\phi}}$  , and an experiment design object on this featurization,  $\Phi :\mathbb{R}^{d_{\phi}\times d_{\phi}}\rightarrow \mathbb{R}$  . Our goal is to collect some set of trajectories  $\{\pmb {\tau}_t\}_{t = 1}^T$  which minimizes  $\Phi$  ..
$$
\Phi \left(\frac {1}{T H} \cdot \sum_ {t = 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}\right).
$$
As an example, if  $\Phi(\Lambda) = \log \det(\Lambda)$ , this reduces to  $D$ -optimal design, and if  $\Phi(\Lambda) = \operatorname{tr}(\mathcal{H} \cdot \Lambda^{-1})$ , the setting considered in Section 4, this reduces to weighted  $A$ -optimal design. As before, we assume we have access to some set of exploration policies  $\Pi_{\mathrm{exp}}$ , and define  $\Lambda_{\pi}$  and  $\Omega$  as in Section 3, but with respect to this new feature map  $\phi$  and system (5.1). We also define  $\widehat{\Omega}$  to be the space of all possible covariance matrices:
$$
\widehat {\boldsymbol {\Omega}} := \left\{\sum_ {h = 1} ^ {H} \phi (\boldsymbol {x} _ {h}, \boldsymbol {u} _ {h}) \phi (\boldsymbol {x} _ {h}, \boldsymbol {u} _ {h}) ^ {\top}: \boldsymbol {x} _ {h} \in \mathcal {X}, \boldsymbol {u} _ {h} \in \mathcal {U}, \forall h \in [ H ] \right\}.
$$
To facilitate efficient experiment design in this setting, we will make the following assumption on  $\Phi$ . Assumption 7 (Regularity of  $\Phi$ ).  $\Phi$  is regular in the following sense:
1.  $\Phi$  is convex, differentiable, and  $\beta$ -smooth in the norm  $\|\cdot\|$  (with dual-norm  $\|\cdot\|_*$ ):
$$
\left\| \nabla_ {\boldsymbol {\Lambda}} \Phi (\boldsymbol {\Lambda}) - \nabla_ {\boldsymbol {\Lambda} ^ {\prime}} \Phi (\boldsymbol {\Lambda} ^ {\prime}) \right\| _ {*} \leq \beta \cdot \| \boldsymbol {\Lambda} - \boldsymbol {\Lambda} ^ {\prime} \|, \quad \forall \boldsymbol {\Lambda}, \boldsymbol {\Lambda} ^ {\prime} \in \widehat {\Omega}.
$$
2. There exists some  $M < \infty$  satisfying  $\sup_{\Lambda \in \widehat{\Omega}} \sup_{\boldsymbol{x} \in \mathcal{X}, \boldsymbol{u} \in \mathcal{U}} |\phi(\boldsymbol{x}, \boldsymbol{u})^\top \nabla_\Lambda \Phi(\Lambda) \phi(\boldsymbol{x}, \boldsymbol{u})| \leq M$ .
The key algorithmic assumption we make is access to a regret minimization oracle on (5.1).
Assumption 8 (Regret Minimization Oracle). Let  $\mathrm{cost}_h(\pmb {x},\pmb {u}) = \phi (\pmb {x},\pmb {u})^\top Q_h\phi (\pmb {x},\pmb {u})$  for some  $Q_{h}\in \mathbb{R}^{d_{\phi}\times d_{\phi}}$  such that  $|\sum_{h}\mathrm{cost}_{h}(\pmb{x}_{h},\pmb{u}_{h})|\leq 1$  for all  $\pmb {x}_h\in \mathcal{X},\pmb {u}_h\in \mathcal{U}$  . We assume we have access to some learner  $\mathbb{A}_{\mathcal{R}}$  which is able to achieve low regret on costs  $\{\mathrm{cost}_h(\cdot ,\cdot)\}_{h = 1}^H$  with respect to policy class  $\Pi_{\mathrm{exp}}$  . That is, with probability at least  $1 - \delta$  ..
$$
\sum_ {t = 1} ^ {T} \mathbb {E} _ {f, \pi_ {t}} [ \sum_ {h = 1} ^ {H} \operatorname {c o s t} _ {h} (\boldsymbol {x} _ {h} ^ {t}, \boldsymbol {u} _ {h} ^ {t}) ] - T \cdot \min  _ {\pi \in \Pi_ {\exp}} \mathbb {E} _ {f, \pi} [ \sum_ {h = 1} ^ {H} \operatorname {c o s t} _ {h} (\boldsymbol {x} _ {h}, \boldsymbol {u} _ {h}) ] \leq C _ {\mathcal {R}} \cdot \log^ {p _ {\mathcal {R}}} \frac {T}{\delta} \cdot T ^ {\alpha}
$$
for some  $C_{\mathcal{R}} > 0, p_{\mathcal{R}} > 0$ , and  $\alpha \in (0,1)$ , and where  $\pi_t$  is the policy  $\mathbb{A}_{\mathcal{R}}$  plays at episode  $t$ .
Note that the regret minimization algorithm satisfying Assumption 8 may be arbitrary. For example, for linear systems, we could apply provably efficient algorithms for the Linear Quadratic Regulator [25, 20]; for nonlinear systems of the form (1.1) we could apply the  $\mathrm{LC^3}$  algorithm of [17]; for more general settings of reinforcement learning with function approximation, algorithms such as BILIN-UCB [30] or E2D [31] could be applied. In practice, though they may not formally satisfy the guarantee of Assumption 8, deep RL approaches could be used. We have the following result.
Theorem 3. Fix  $T > 0$  and denote  $R \coloneqq \sup_{\Lambda, \Lambda' \in \widehat{\Omega}} \| \Lambda - \Lambda' \|$ . Under Assumption 7, and assuming we have access to a learner  $\mathbb{A}_{\mathcal{R}}$  satisfying Assumption 8 with  $\alpha = 1/2$ , DYNAMICOED runs for  $T$  episodes on (5.1), and with probability at least  $1 - \delta$  collects data  $\{(\pmb{x}_h^t, \pmb{u}_h^t)\}_{h \in [H], t \in [T]}$  satisfying
$$
\Phi \bigg (\frac {1}{T} \cdot \sum_ {t = 1} ^ {T} \sum_ {h = 1} ^ {H} \phi_ {h} ^ {t} (\phi_ {h} ^ {t}) ^ {\top} \bigg) - \min _ {\Lambda \in \Omega} \Phi (\Lambda) \leq \frac {\beta R ^ {2} \log T + H M (C _ {\mathcal {R}} \log^ {p _ {\mathcal {R}}} \frac {2 T}{\delta} + 3 \log^ {1 / 2} \frac {4 T}{\delta})}{T ^ {1 / 3}}
$$
where  $\phi_h^t \coloneqq \phi(\pmb{x}_h^t, \pmb{u}_h^t)$ , and we recall  $\Omega$  is the set of possible expected covariates on (5.1).
Theorem 3 shows that, given access only to a regret minimization oracle, it is possible to solve experiment design problems on arbitrary dynamical systems. The requirement that  $\alpha = 1 / 2$  is for expositional purposes only—we generalize this result to arbitrary  $\alpha$  (and more general feature maps) in Appendix C. Under certain conditions, it can be shown that, if the exploration policies DYNAMICOED runs to collect  $\mathfrak{D}$  are rerun, the newly collected data satisfies a similar guarantee as Theorem 3. This lets us run DYNAMICOED to learn an approximate solution of  $\min_{\Lambda \in \Omega} \Phi(\Lambda)$ , and then rerun the learned policies as many times as desired to collect additional data approximately minimizing  $\Phi$ .
5.1 Overview of DYNAMICOED Algorithm
DYNAMICOED is inspired by recent work on experiment design in reinforcement learning [57, 58, 56, 59], and can be seen as an extension of the FWREGRET algorithm of [56] to arbitrary systems. We
refer the reader to [56] for a more in-depth discussion of the FWREGRET algorithm, and briefly sketch its extension to arbitrary systems here (see Appendix C and Algorithm 4 for precise definitions).
Algorithm 2 Dynamic Optimal Experiment Design (DYNAMICOED, Informal)
1: input: objective  $\Phi$ , episodes  $T$ , confidence  $\delta$ , regret algorithm  $\mathbb{A}_{\mathcal{R}}$ , exploration policies  $\Pi_{\mathrm{exp}}$   
2: Set  $K \gets \mathcal{O}(T^{2/3})$ ,  $N \gets \mathcal{O}(T^{1/3})$ ,  $\gamma_n \gets \frac{1}{n+1}$   
//  $\phi_h^{k,n} := \phi(\pmb{x}_h^{k,n}, \pmb{u}_h^{k,n})$  for  $(\pmb{x}_h^{k,n}, \pmb{u}_h^{k,n})$  the state-input at step  $h$  of episode  $k$  of iteration  $n$   
3: Play any  $\pi_{\mathrm{exp}} \in \Pi_{\mathrm{exp}}$  for  $K$  episodes, set  $\Lambda_0 \gets \frac{1}{K} \sum_{k=1}^{K} \sum_{h=1}^{H} \phi_h^{k,0} (\phi_h^{k,0})^\top$   
4: for  $n = 1, 2, \ldots, N$  do  
5: Compute derivative of  $\Phi(\Lambda_{n-1})$ ,  $\Xi_n \gets \nabla_\Lambda \Phi(\Lambda)|_{\Lambda=\Lambda_{n-1}}$   
6: Run  $\mathbb{A}_{\mathcal{R}}$  on cost  $\mathrm{cost}_h^n(\pmb{x}, \pmb{u}) \gets \frac{1}{M} \cdot \phi(\pmb{x}, \pmb{u})^\top(\Xi_n) \phi(\pmb{x}, \pmb{u})$  for  $K$  episodes  
7:  $\Lambda_n \gets (1 - \gamma_n) \Lambda_{n-1} + \frac{\gamma_n}{K} \cdot \sum_{k=1}^{K} \sum_{h=1}^{H} \phi_h^{k,n} (\phi_h^{k,n})^\top$   
8: return  $\frac{1}{T} \sum_{n=0}^{N} \sum_{k=1}^{K} \sum_{h=1}^{H} \phi_h^{k,n} (\phi_h^{k,n})^\top$
Conceptually, DYNAMICOED runs a variant of conditional gradient descent on the objective  $\Phi(\Lambda)$ . At each iteration,  $n$ , it computes the gradient of the loss at the current iterate,  $\Xi_n \gets \nabla_{\Lambda} \Phi(\Lambda)|_{\Lambda = \Lambda_{n-1}}$ . To run a standard gradient descent algorithm on this objective, we would simply update  $\Lambda_{n-1}$  by taking a step in the direction  $\Xi_n$ . However, our objective is to minimize  $\Phi$  over the constraint set,  $\Omega$ . Thus, rather than taking a step in the direction  $\Xi_n$ , we wish to take a step in the direction of steepest descent within the constraint set.
The challenge is that the constraint set in our setting,  $\Omega$ , is unknown, as it depends on the expectation over trajectories induced on the unknown dynamics  $(f_h)_{h=1}^H$ , and therefore we cannot directly compute this steepest descent direction. The key observation is that the computation of this steepest descent direction is equivalent to solving:
$$
\arg \min _ {\pi_ {\exp} \in \Pi_ {\exp}} \mathbb {E} _ {f, \pi_ {\exp}} \left[ \sum_ {h = 1} ^ {H} \phi (\boldsymbol {x} _ {h}, \boldsymbol {u} _ {h}) ^ {\top} (\Xi_ {n}) \phi (\boldsymbol {x} _ {h}, \boldsymbol {u} _ {h}) \right].
$$
This is simply a policy optimization problem, however, and can be solved approximately by  $\mathbb{A}_{\mathcal{R}}$  under Assumption 8. Thus, in the call to  $\mathbb{A}_{\mathcal{R}}$  on Line 6, we approximate the steepest descent direction, and on Line 7 update  $\Lambda_{n-1}$  in this direction. Convergence of this procedure to the optimal value,  $\min_{\Lambda \in \Omega} \Phi(\Lambda)$ , can then be shown by the standard analysis of conditional gradient descent. We remark that, under Assumption 7 and Assumption 8, this argument is completely generic and does not require that our system, (5.1), exhibit any additional properties.
5.2 From Theorem 3 to Theorem 1
In Algorithm 1, our goal is to collect covariates,  $\tilde{\Lambda}_{T_{\ell}}^{-1}$ , on (1.1) such that  $\mathrm{tr}(\mathcal{H}(\widehat{A}^{\ell})\check{\Lambda}_{T_{\ell}}^{-1})$  is as small as possible. To achieve this, we apply DYNAMICOED to the dynamics (1.1) and objective  $\Phi_{\ell}(\Lambda) = \mathrm{tr}(\mathcal{H}(\widehat{A}^{\ell})\check{\Lambda}^{-1})$ , with Assumption 8 instantiated by the LC $^3$  algorithm of [17]. By the guarantee given in Theorem 3, after running for a number of episodes  $N$  which scales polynomially in problem parameters, DYNAMICOED will collect covariates  $\Lambda_N$  such that  $\Phi_{\ell}(\frac{1}{N}\Lambda_N)\leq 2\cdot \min_{\Lambda \in \Omega}\Phi_{\ell}(\Lambda)$ , which implies  $\mathrm{tr}(\mathcal{H}(\widehat{A}^{\ell})\check{\Lambda}_N^{-1})\leq \frac{2}{N}\cdot \min_{\Lambda \in \Omega}\mathrm{tr}(\mathcal{H}(\widehat{A}^{\ell})\check{\Lambda}^{-1})$ . By rerunning the policies DYNAMICOED used to collect this  $\Lambda_N$  for  $T_{\ell} / N$  additional times, we can ensure  $\mathrm{tr}(\mathcal{H}(\widehat{A}^{\ell})\check{\Lambda}_{T_{\ell}}^{-1})\lesssim \frac{1}{T_{\ell}}\cdot \min_{\Lambda \in \Omega}\mathrm{tr}(\mathcal{H}(\widehat{A}^{\ell})\check{\Lambda}^{-1})$  as desired.
We remark that the  $\mathrm{LC}^3$  algorithm requires access to a computation oracle. As the focus of this work is primarily statistical, we leave addressing this computational challenge for future work. Furthermore, as we show in the following section, computationally efficient, sampling-based implementations of our approach are very effective in practice. We remark as well that the objective we ultimately care about minimizing is  $\operatorname{tr}(\mathcal{H}(A_{\star})\check{\Lambda}^{-1})$ . As we show, by including a small amount of uniform exploration, we can ensure that the suboptimality incurred optimizing  $\operatorname{tr}(\mathcal{H}(\widehat{A}^\ell)\check{\Lambda}^{-1})$  instead of  $\operatorname{tr}(\mathcal{H}(A_{\star})\check{\Lambda}^{-1})$  only contributes to the lower-order terms of the final guarantee in Theorem 1.
Finally, we demonstrate the effectiveness of our proposed approach (Algorithm 1, the Task-Driven Exploration method in Figures 1 to 3) on several systems motivated by robotic applications. We
Figure 2: Performance on Drone
Figure 3: Performance on Car
compare Algorithm 1 with an approach that plays  $\pmb{u}_h \sim \mathcal{N}(0, \sigma_{\pmb{u}}^2 \cdot I)$  (Gaussian Exploration), and an approach inspired by [1] (Uniform Exploration), which seeks to estimate  $A_{\star}$  uniformly well, playing inputs that reduce  $\| \widehat{A} - A_{\star} \|_{\mathrm{op}}$ .