# Multi-Task Imitation Learning for Linear Dynamical Systems

Thomas T. Zhang Authors contributed equally to this work. University of
Pennsylvania  Katie Kang††footnotemark:  University of California, Berkeley
Bruce D. Lee††footnotemark:  University of Pennsylvania  Claire Tomlin
University of California, Berkeley  Sergey Levine  University of California,
Berkeley  Stephen Tu  Google Research, Brain Team  Nikolai Matni  University
of Pennsylvania  Google Research, Brain Team

###### Abstract

We study representation learning for efficient imitation learning over linear
systems. In particular, we consider a setting where learning is split into two
phases: (a) a pre-training step where a shared $k$-dimensional representation
is learned from $H$ source policies, and (b) a target policy fine-tuning step
where the learned representation is used to parameterize the policy class. We
find that the imitation gap over trajectories generated by the learned target
policy is bounded by
$\tilde{O}\left(\frac{kn_{x}}{HN_{\mathrm{shared}}}+\frac{kn_{u}}{N_{\mathrm{target}}}\right)$,
where $n_{x}>k$ is the state dimension, $n_{u}$ is the input dimension,
$N_{\mathrm{shared}}$ denotes the total amount of data collected for each
policy during representation learning, and $N_{\mathrm{target}}$ is the amount
of target task data. This result formalizes the intuition that aggregating
data across related tasks to learn a representation can significantly improve
the sample efficiency of learning a target task. The trends suggested by this
bound are corroborated in simulation.

Keywords: Imitation learning, transfer learning, multi-task learning,
representation learning

##  1 Introduction

Imitation learning (IL), which learns control policies by imitating expert
demonstrations, has demonstrated success across a variety of domains including
self-driving cars (Codevilla et al., 2018) and robotics (Schaal, 1999).
However, using IL to learn a robust behavior policy may require a large amount
of training data (Ross et al., 2011), and expert demonstrations are often
expensive to collect. One remedy for this problem is multi-task learning:
using data from other tasks (source tasks) in addition to from the task of
interest (target task) to jointly learn a policy. We study the application of
multi-task learning to IL over linear systems, and demonstrate improved sample
efficiency when learning a controller via representation learning.

Our results expand on prior work that studies multi-task representation
learning for supervised learning (Du et al., 2020; Tripuraneni et al., 2021),
addressing the new challenges that arise in the imitation learning setting.
First, the data for IL is temporally dependent, as it is generated from a
dynamical system $x[t+1]=f(x[t],u[t],w[t])$. In contrast, the supervised
learning setting assumes that both the train and test data are independent and
identically distributed (i.i.d.) from the same underlying distribution.
Furthermore, we are interested in the performance of the learned controller in
closed-loop rather than its error on expert-controlled trajectories. Hence,
bounds on excess risk, which corresponds to the one-step prediction error of
the learned controller under the expert distribution, are not immediately
informative for the closed-loop performance. We instead focus our analysis on
the tracking error between the learned and expert policies, which requires us
to account for the distribution shift between the learned and expert
controllers.

We address these challenges in the setting of IL for linear systems. The
following statement captures the benefits of multi-task representation
learning on sample complexity:

######  Theorem 1.1 (main result, informal)

Suppose that the source task controllers are sufficiently related to the
target task controller. Then, the tracking error between the learned target
controller and the corresponding expert is bounded with high probability by:

$\displaystyle\textrm{tracking error}\lesssim\frac{\textrm{rep.\ dimension
}\times\textrm{state dimension}}{\textrm{\\# source task
datapoints}}+\frac{\textrm{rep.\ dimension }\times\textrm{input
dimension}}{\textrm{\\# target task datapoints}}.$

The first term in this bound corresponds to the error from learning a common
representation, and the second term the error in fitting the remaining weights
of the target task controller. The key upshot of this result is that the
numerator of the second term (rep. dimension $\times$ input dimension) is
smaller than the number of parameters (input dimension $\times$ state
dimension) in the target controller. This demonstrates an improvement in
sample complexity of multi-task IL over direct IL, where the error scales as
$\frac{\\#\mathrm{parameters}}{\\#\mathrm{datapoints}}$. Furthermore, we note
that the error in learning the representation decays along all axes of the
data: # of tasks $\times$ # of trajs $\times$ traj length for source tasks,
and # of trajs $\times$ traj length for the target task. It is non-trivial to
demonstrate that the error decays with the trajectory length, and doing so
requires tools that handle causally dependent data in our analysis.

The remainder of the paper formulates the multi-task IL problem, and the
assumptions required to prove Theorem 1.1. The main contributions may be
summarized as follows:

  * •

We provide novel interpretable notions of source task overlap with the target
task (§2 and §3).

  * •

We bound the imitation gap achieved by multi-task IL as in Theorem 1.1 (§3).

  * •

We empirically show the efficacy of multi-task IL when the assumptions are
satisfied (§4).

###  1.1 Related Work

Multi-task imitation and reinforcement learning: Multi-task RL and IL methods
seek to represent policies solving different tasks with shared parameters,
enabling the transfer of knowledge across related tasks (Teh et al., 2017;
Espeholt et al., 2018; Hessel et al., 2018; Singh et al., 2020; Deisenroth et
al., 2014), and rapid test-time adaptation to new tasks (Finn et al., 2017;
Rakelly et al., 2019; Duan et al., 2016; Yu et al., 2021; Yang and Nachum,
2021). There also exists a body of work which theoretically analyses the
sample complexity of representation learning in multi-task RL and IL (Lu et
al., 2021; Cheng et al., 2022; Lu et al., 2022; Xu et al., 2020; Maurer et
al., 2015; Arora et al., 2020). While this line of work considers a more
general MDP setting compared with the linear dynamical systems we consider,
the specific results are often stated with incompatible assumptions (such as
bounded states/cost functions and discrete action spaces), and/or do not
reflect how system-theoretic properties such as closed-loop task stability
affect the final rates.

Multi-task system identification and adaptive control: Recent work has also
considered applications of multi-task learning where the dynamics change
between tasks, and the goal is to perform adaptive control (Harrison et al.,
2018; Richards et al., 2021, 2022; Shi et al., 2021; Muthirayan et al., 2022)
or dynamics forecasting (Wang et al., 2021). Multi-task system identification
(Modi et al., 2021) and stabilization using data from related systems (Li et
al., 2022) have also been considered. Our work instead studies the problem of
learning to imitate different expert controllers while the system remains the
same, and demonstrates bounds on the tracking error between the learned
controller and its corresponding expert.

Sample complexity of multi-task learning: Numerous works have studied the
sample efficiency gains of multi-task learning for regression and
classification under various task similarity assumptions (Baxter, 1995;
Crammer et al., 2008; Maurer et al., 2016; Tripuraneni et al., 2020; Chua et
al., 2021). Most closely related to our results are Du et al. (2020) and
Tripuraneni et al. (2021), both of which show multi-task representation
learning sample complexity bounds in the linear regression setting in which
the error from learning the representation decays with the total number of
source training samples. Our work leverages these results to tackle the
setting of linear imitation learning, which has the additional challenges of
non-i.i.d. data and test time distribution shift.

##  2 Problem Formulation

###  2.1 Multi-Task Imitation Learning

Imitation learning uses state/action pairs
$(x,u)\in\mathbb{R}^{n_{x}}\times\mathbb{R}^{n_{u}}$ of expert demonstrations
to learn a controller
$\hat{\pi}:\mathbb{R}^{n_{x}}\rightarrow\mathbb{R}^{n_{u}}$, by matching the
learned controller actions to the expert actions. In particular, if
$\mathcal{D}$ is the training set of expert state/action pairs, then
$\hat{\pi}\in\operatorname*{argmin}_{\pi}\sum_{(x,u)\in\mathcal{D}}\left\|\pi(x)-u\right\|^{2}.$

We are interested in the problem of _multi-task_ imitation learning, where we
consider $H+1$ different expert controllers. We call the first $H$ controllers
source controllers and the $(H+1)^{\textrm{st}}$ controller the target
controller. We assume that we have access to $N_{1}$ trajectories for each
source task, and $N_{2}\leq N_{1}$ trajectories for the target task. For
simplicity, we assume all trajectories are of the same length $T$. In
particular, for each source task
$h\in\mathopen{}\left\\{1,\dots,H\right\\}\mathclose{}$, our source data
consists of
$\mathopen{}\left\\{\mathopen{}\left\\{(x_{i}^{(h)}[t],u_{i}^{(h)}[t])\right\\}\mathclose{}_{t=0}^{T-1}\right\\}\mathclose{}_{i=1}^{N_{1}}$,
while our target data consists of
$\mathopen{}\left\\{\mathopen{}\left\\{(x_{i}^{(H+1)}[t],u_{i}^{(H+1)}[t])\right\\}\mathclose{}_{t=0}^{T-1}\right\\}\mathclose{}_{i=1}^{N_{2}}$.
Our goal is to learn a controller which effectively imitates the target
controller. However, because we only have access to a small number ($N_{2}$)
of target expert trajectories, we leverage the $HN_{1}$ expert trajectories
from the source controllers to accelerate the learning of the target
controller.

To do so, we break our training into two stages: a pre-training stage which
learns from the combined source task data, and a target training stage which
only learns from the target task data. In the pre-training stage, we extract a
common, low dimensional representation for the source controllers, which is
used later in the target training stage. More specifically, we learn a common,
low dimensional representation mapping
$\hat{\phi}:\mathbb{R}^{n_{x}}\rightarrow\mathbb{R}^{k}$, where $k<n_{x}$ is
the dimension of the representation, and linear predictors
$\hat{F}^{(h)}\in\mathbb{R}^{n_{u}\times k}$ unique to each task:

$\displaystyle\hat{\phi},\hat{F}^{(1)},\dots,\hat{F}^{(H)}\in\operatorname*{argmin}_{\phi,F^{(1)},\dots,F^{(H)}}\sum_{h=1}^{H}\sum_{i=1}^{N_{1}}\sum_{t=0}^{T-1}\left\|F^{(h)}\phi(x_{i}^{(h)}[t])-u_{i}^{(h)}[t]\right\|^{2}.$
(1)

We do not address the details of solving the empirical risk minimization
problem, and instead perform our analysis assuming (1) can be solved to
optimality. Note however that Tripuraneni et al. (2021) demonstrate in the
linear regression setting that a method-of-moments-based algorithm can
efficiently find approximate empirical risk minimizers.

Once a common representation $\hat{\phi}$ is obtained, we move on to target
task training. During target task training, we use the common representation
mapping $\hat{\phi}$ learned from the pre-training step to map the states into
the lower dimensional representation, and learn an additional linear predictor
$\hat{F}^{(H+1)}$ unique to the target task to model the target controller:

$\displaystyle\hat{F}^{(H+1)}=\operatorname*{argmin}_{F}\sum_{i=1}^{N_{2}}\sum_{t=0}^{T-1}\left\|F\hat{\phi}\mathopen{}\left(x_{i}^{(H+1)}[t]\right)\mathclose{}-u_{i}^{(H+1)}[t]\right\|^{2}.$
(2)

Since the representation $\hat{\phi}$ is fixed from pre-training, (2) is an
ordinary least squares problem.

###  2.2 System and Data Assumptions

We focus our analysis on a linear systems setting, with state
$x[t]\in\mathbb{R}^{n_{x}}$, input $u[t]\in\mathbb{R}^{n_{u}}$, and Gaussian
process noise $w[t]\in\mathbb{R}^{n_{x}}$ obeying dynamics

$\displaystyle x[t+1]=Ax[t]+Bu[t]+w[t].$ (3)

Let each expert controller be of the form $u[t]=K^{(h)}x[t]+z[t]$, where
$z[t]\in\mathbb{R}^{n_{u}}$ is Gaussian actuator noise.111As the control
actions are the labels in IL, actuator noise corresponds to label noise in
supervised learning. In the absence of such noise, the controller is recovered
by $n_{x}$ linearly independent states and corresponding expert inputs. We
assume the system matrices $(A,B)$ remain the same between tasks, but the
process noise covariance and the controllers may change.222This could be the
case, for instance, if different controllers are designed for different levels
of noise. In particular, we have
$w^{(h)}[t]\overset{i.i.d.}{\sim}\mathcal{N}(0,\Sigma_{w}^{(h)})$ and
$z^{(h)}[t]\overset{i.i.d.}{\sim}\mathcal{N}(0,\sigma_{z}^{2}I)$ with
$\Sigma_{w}^{(h)}\succ 0$ for all $h\in[H+1]$ and $\sigma_{z}^{2}>0$.

We assume all of the expert controllers $K^{(h)}$ are stabilizing, i.e., the
spectral radii of $A+BK^{(h)}$ are less than one. Note that this implies that
$(A,B)$ is stabilizable. No other assumptions on $(A,B)$ are required. The
state distribution of the system under each expert controller will converge to
the stationary distribution $\mathcal{N}(0,\Sigma_{x}^{(h)})$, where
$\Sigma_{x}^{(h)}$ solves the following discrete Lyapunov equation:

$\Sigma_{x}^{(h)}=(A+BK^{(h)})\Sigma_{x}^{(h)}(A+BK^{(h)})^{\top}+\sigma_{z}^{2}BB^{\top}+\Sigma_{w}^{(h)}.$

For simplicity, we assume that the initial states of the expert demonstrations
in our datasets are sampled
$x_{i}^{(h)}[0]\overset{i.i.d.}{\sim}\mathcal{N}(0,\Sigma_{x}^{(h)})$. Thus at
all times, the marginal state distributions of the expert demonstrations are
equal to $\mathcal{N}(0,\Sigma_{x}^{(h)})$.333This assumption is not
restrictive, as stable systems exponentially converge to stationarity.

Finally, we assume that the expert controllers share a low dimensional
representation. Specifically, there exists some
$\Phi_{\star}\in\mathbb{R}^{k\times n_{x}}$ with $2k\leq n_{x}$444While this
assumption is more stringent than the intuitive $k<n_{x}$ assumption, it
arises from the fact that the residual of the stacked source controllers may
be of rank $2k$. and weights
$F^{(1)}_{\star},F^{(2)}_{\star},\dots,F^{(H+1)}_{\star}\in\mathbb{R}^{n_{u}\times
k}$ such that for all $h\in[H+1]$, $K^{(h)}=F_{\star}^{(h)}\Phi_{\star}$, and
the action taken at time $t$ for trajectory $i$ is:555An example of a setting
where expert controllers satisfy this assumption is when the system has high
dimensional states which exhibit low dimensional structure, e.g. when $A$ and
$B$ can be decomposed into $A=\Phi_{\star}^{\dagger}\tilde{A}\Phi_{\star}$ and
$B=\Phi_{\star}^{\dagger}\tilde{B}$, where $\tilde{A}\in\mathbb{R}^{k\times
k}$ and $\tilde{B}\in\mathbb{R}^{k\times n_{u}}$. Here, linear policies $K$
which optimize some objective in terms of the low dimensional features of the
system can be decomposed into $K=\tilde{K}\Phi_{\star}$, where
$\tilde{K}\in\mathbb{R}^{n_{u}\times k}$, mirroring the assumptions of our
expert controllers. We provide a concrete example in Section 4.

$u_{i}^{(h)}[t]=F^{(h)}_{\star}\Phi_{\star}x_{i}^{(h)}[t]+z_{i}^{(h)}[t].$

Under this assumption, the learned common representation $\hat{\phi}$ in
Section 2.1 can be restricted to linear representations, i.e.,
$\hat{\phi}(x)=\hat{\Phi}x$, where $\hat{\Phi}\in\mathbb{R}^{k\times n_{x}}$.
Note that solving Problem (2) with $\hat{\Phi}$ fixed involves solving for
only $kn_{u}$ parameters, which is smaller than the $n_{u}n_{x}$ unknown
parameters when learning from scratch. In particular, by representing the
controller as $F^{(H+1)}\Phi$, we have $k(n_{u}+n_{x})$ unknown parameters:
$kn_{x}$ of the parameters are, however, learned using the source task data,
leaving only $kn_{u}$ parameters to learn with target task data.

###  2.3 Notation

The Euclidean norm of a vector $x$ is denoted $\left\|x\right\|$. For a matrix
$A$, the spectral norm is denoted $\left\|A\right\|$, and the Frobenius norm
is denoted $\left\|A\right\|_{F}$. The spectral radius of a square matrix is
denoted $\rho(A)$. We use $\dagger$ to denote the Moore-Penrose pseudo-
inverse. For a square matrix $A$ with $\rho(A)<1$, define
$\mathcal{J}(A)=\sum_{t\geq 0}\left\|A^{t}\right\|<\infty$. A symmetric,
positive semi-definite (psd) matrix $A=A^{\top}$ is denoted $A\succeq 0$.
Similarly $A\succeq B$ denotes that $A-B$ is positive semidefinite. The
condition number of a positive definite matrix $A$ is denoted
$\kappa(A)=\frac{\lambda_{\max}(A)}{\lambda_{\min}(A)}$, where
$\lambda_{\max}$ and $\lambda_{\min}$ denote the maximum and minimum
eigenvalues, respectively. Similarly, $\sigma_{i}(A)$ denote the singular
values of $A$. We denote the normal distribution with mean $\mu$ and
covariance $\Sigma$ by $\mathcal{N}(\mu,\Sigma)$. We use standard
$\mathcal{O}(\cdot)$, $\Theta(\cdot)$ and $\Omega(\cdot)$ to omit universal
constant factors, and $\tilde{}\mathcal{O}(\cdot),\tilde{\Theta}(\cdot)$ and
$\tilde{\Omega}(\cdot)$ to also omit polylog factors. We also use $a\lesssim
b$ to denote $a=O(b)$. We use the indexing shorthand
$[K]:=\mathopen{}\left\\{1,\dots,K\right\\}\mathclose{}$. For a given task
$h\in[H+1]$, the matrix of stacked states is defined as

$\displaystyle\mathbf{X}^{(h)}$
$\displaystyle=\begin{bmatrix}x_{1}^{(h)}[0]&\dots&x_{1}^{(h)}[T-1]&\dots&x_{N_{1}}^{(h)}[0]&\dots&x_{N_{1}}^{(h)}[T-1]\end{bmatrix}^{\top}\in\mathbb{R}^{N_{1}T\times
n_{x}}.$ (4)

Lastly, let $\bar{\lambda}=\max_{1\leq h\leq
H}\lambda_{\max}(\Sigma_{x}^{(h)})$ and $\underline{\lambda}=\min_{1\leq h\leq
H}\lambda_{\min}(\Sigma_{x}^{(h)})$.

##  3 Sample Complexity of Multi-Task Imitation Learning

In order to derive any useful information from source tasks for a downstream
task, the source tasks must satisfy some notion of task diversity that
sufficiently covers the downstream task. To that end, we introduce the
following notions of source tasks covering the target task.

######  Definition 3.1 (target task covariance coverage (Du et al., 2020))

Define the constant $c$ as:

$\displaystyle
c:=\min_{h\in[H]}\lambda_{\min}((\Sigma_{x}^{(H+1)})^{-1/2}\Sigma_{x}^{(h)}(\Sigma_{x}^{(H+1)})^{-1/2}).$
(5)

Note that $c$ is well-defined and positive by our assumption that
$\Sigma_{w}^{(h)}\succ 0$ for all $h\in[H+1]$.

Definition 3.1 captures the degree to which the closed-loop distribution of
states for each source task aligns with that of the target task. We then
introduce the following notion of task similarity between the source and
target task weights, which generalizes the well-conditioning assumptions in Du
et al. (2020) and Tripuraneni et al. (2021).

######  Assumption 3.1 (diverse source controllers)

We assume the target task weights $F_{\star}^{(H+1)}$ and the source task
weights $F_{\star}^{(1)},\dots,F_{\star}^{(H)}$ satisfy

$\displaystyle\left\|F_{\star}^{(H+1)}\begin{bmatrix}F_{\star}^{(1)}\\\
\vdots\\\
F_{\star}^{(H)}\end{bmatrix}^{\dagger}\right\|^{2}\leq\mathcal{O}\mathopen{}\left(\frac{1}{H}\right)\mathclose{}.$
(6)

3.1 states that the alignment and loadings of the singular spaces between the
stacked source task weights and target task weights closely match along the
low-dimensional representation dimension. For example, if
$F_{\star}^{(h)}=F_{\star}^{(H+1)}$ for each $h\in[H]$, the RHS of (6) is
$1/H$. We note that this assumption subsumes and is more geometrically
informative than a direct bound on the ratio of singular values, e.g.

$\sigma_{\max}^{2}(F_{\star}^{(H+1)})/\sigma_{k}^{2}\mathopen{}\left(\begin{bmatrix}F_{\star}^{(1)}\\\
\vdots\\\
F_{\star}^{(H)}\end{bmatrix}\right)\mathclose{}\leq\mathcal{O}(1/H),$

which would follow by naively extending the well-conditioning assumptions in
Du et al. (2020) and Tripuraneni et al. (2021). Notably, such a condition
might not be satisfied even if $F_{\star}^{(h)}=F_{\star}^{(H+1)}$, $\forall
h\in[H]$, e.g., if $F_{\star}^{(H+1)}$ is rank-deficient.

###  3.1 Excess Risk Bound: Generalization Along Expert Target Task
Trajectories

First we show that learning controllers through multi-task representation
learning leads to favorable generalization bounds on the excess risk of the
learned controller inputs on the expert target task state distribution,
analogous to the bounds on multi-task linear regression in Du et al. (2020);
Tripuraneni et al. (2021). However, a key complicating factor in our setting
is the fact that the input noise $z^{(h)}[t]$ enters the process, and thus the
data $x^{(h)}[t]$ is causally dependent on the “label noise”. In order to
overcome this issue and preserve our statistical gains along time $T$, we
leverage the theory of self-normalized martingales, in particular generalizing
tools from Abbasi-Yadkori et al. (2011) to the matrix-valued setting. The full
argument is detailed in Appendix A and Appendix B. This culminates in the
following target task excess risk bound.

######  Theorem 3.1 (target task excess risk bound)

Given $\delta\in(0,1)$, suppose that

$\displaystyle N_{1}T$
$\displaystyle\gtrsim\max_{h\in[H]}\mathcal{J}\mathopen{}\left(A+BK^{(h)}\right)\mathclose{}^{2}\kappa\mathopen{}\left(\Sigma_{x}^{(h)}\right)\mathclose{}(n_{x}+\log(H/\delta)),$
$\displaystyle N_{2}T$
$\displaystyle\gtrsim\mathcal{J}\mathopen{}\left(A+BK^{(H+1)}\right)\mathclose{}^{2}\kappa\mathopen{}\left(\Sigma_{x}^{(H+1)}\right)\mathclose{}(k+\log(1/\delta)).$

Define $\mathcal{P}_{0:T-1}^{(H+1)}$ as the distribution over target task
trajectories $(x^{(H+1)}[0],\cdots,x^{(H+1)}[T-1])$. Then with probability at
least $1-\delta$, the excess risk of the learned representation $\hat{\Phi}$
and target task weights $\hat{F}^{(H+1)}$ is bounded by

$\displaystyle\mathrm{ER}(\hat{\Phi},\hat{F}^{(H+1)})$
$\displaystyle:=\frac{1}{2T}\mathbb{E}_{\mathcal{P}_{0:T-1}^{(H+1)}}\mathopen{}\left[\sum_{t=0}^{T-1}\left\|(F_{\star}^{(H+1)}\Phi_{\star}-\hat{F}^{(H+1)}\hat{\Phi})x^{(H+1)}[t]\right\|^{2}\right]\mathclose{}$
$\displaystyle\lesssim\sigma_{z}^{2}\mathopen{}\left(\frac{kn_{x}\log\mathopen{}\left(N_{1}T\frac{\bar{\lambda}}{\underline{\lambda}}\right)\mathclose{}}{cN_{1}TH}+\frac{kn_{u}+\log(\frac{1}{\delta})}{N_{2}T}\right)\mathclose{}.$
(7)

Note that when we are operating in the setting where we have much more source
data than target data, the second term limits the excess risk bound in (7).
The second term scales with $kn_{u}$, which is smaller than the number of
total parameters in the controller $n_{u}n_{x}$, or $k(n_{u}+n_{x})$ under the
assumption of a low rank (rank-$k$) controller. Therefore, the benefit of
multi-task learning exhibited by this bound is most clear in the setting of
underactuation, i.e., when $n_{u}\leq n_{x}$. It should also be noted that the
quantity $kn_{x}$ in the numerator of the first term will only be smaller than
the number of source controller parameters ($n_{x}n_{u}H$) if $k$ is much
smaller than $n_{u}H$. This is reasonable, because if $k\geq n_{u}H$, an
optimal representation could simply contain all of the source task
controllers.

###  3.2 Closed-Loop Guarantees: Tackling Distribution Shift

We show that using multi-task representation learning leads to favorable
generalization bounds of the performance of the learned target controller in
closed-loop. As we are studying the pure offline imitation learning
(“behavioral cloning”) setting, we do not assume that the expert controllers
are optimizing any particular objective. Therefore, to quantify the
performance of the controller, we bound the deviation of states generated by
the learned and expert target controller run in closed-loop, i.e., the
tracking error, which implies general expected-cost bounds.

In order to transfer a bound on the excess risk of the target task
$\mathrm{ER}(\hat{\Phi},\hat{F}^{(H+1)})$ into a bound on the tracking error,
we must account for the fundamental distribution shift between the expert
trajectories seen during training and the trajectories generated by running
the learned controller in closed-loop. We leverage the recent framework of
Pfrommer et al. (2022) to bound the tracking error, making the necessary
modifications to handle stochasticity. Our bound formalizes the notion that
“low training error implies low test error,” even under the aforementioned
distribution shift. A detailed exposition can be found in Appendix C.

Let us define the following coupling of the states of the expert versus
learned target task closed-loop systems: given a learned controller
$\hat{K}=\hat{F}^{(H+1)}\hat{\Phi}$ from solving the pre-training and fine-
tuning optimization problems (1) and (2), for a realization of process
randomness $x[0]\sim\mathcal{N}(0,\Sigma_{x}^{(H+1)})$ and
$z[t]\overset{\mathrm{i.i.d.}}{\sim}\mathcal{N}(0,\sigma_{z}^{2}I)$,
$w[t]\overset{\mathrm{i.i.d.}}{\sim}\mathcal{N}(0,\Sigma_{w}^{(H+1)})$ for
$t=0,\dots,T-1$, we write

$\displaystyle x_{\star}[t+1]$
$\displaystyle=(A+BK^{(H+1)})x_{\star}[t]+Bz[t]+w[t],\quad x_{\star}[0]=x[0],$
$\displaystyle\hat{x}[t+1]$
$\displaystyle=(A+B\hat{K})\hat{x}[t]+Bz[t]+w[t],\quad\hat{x}[0]=x[0].$

Thus $\hat{x}[t]$ and $x_{\star}[t]$ are the states visited by the learned and
expert target task systems with the _same_ draw of process randomness. We show
a high probability bound on the closed-loop tracking error
$\left\|x_{\star}[t]-\hat{x}[t]\right\|$ that scales with the excess risk of
the learned controller. Denote by $\mathcal{P}^{\star}_{1:T}$ and
$\hat{}\mathcal{P}_{1:T}$ the distributions of trajectories
$\mathopen{}\left\\{x_{\star}[t]\right\\}\mathclose{}_{t=1}^{T}$ and
$\mathopen{}\left\\{\hat{x}[t]\right\\}\mathclose{}_{t=1}^{T}$.

######  Theorem 3.2 (Target task tracking error bound)

Let $(\hat{\Phi},\hat{F}^{(H+1)})$ denote the learned representation and
target task weights, and $\mathrm{ER}(\hat{\Phi},\hat{F}^{(H+1)})$ denote the
corresponding excess risk. Define $A_{\mathsf{cl}}:=A+BK^{(H+1)}$. Assume that
the excess risk satisfies:

$\mathrm{ER}(\hat{\Phi},\hat{F}^{(H+1)})\lesssim\frac{\lambda_{\min}\mathopen{}\left(\Sigma_{x}^{(H+1)}\right)\mathclose{}}{\mathcal{J}\mathopen{}\left(A_{\mathsf{cl}}\right)\mathclose{}^{2}\left\|B\right\|^{2}}.$
(8)

Then with probability greater than $1-\delta$, for a new target task
trajectory sampled with process randomness
$x[0]\sim\mathcal{N}(0,\Sigma_{x}^{(H+1)})$ and
$z[t]\overset{\mathrm{i.i.d.}}{\sim}\mathcal{N}(0,\sigma_{z}^{2}I)$,
$w[t]\overset{\mathrm{i.i.d.}}{\sim}\mathcal{N}(0,\Sigma_{w}^{(H+1)})$ for
$t=0,\dots,T-1$, the tracking error satisfies

$\displaystyle\max_{1\leq t\leq T}\left\|\hat{x}[t]-x_{\star}[t]\right\|^{2}$
$\displaystyle\lesssim\mathcal{J}\mathopen{}\left(A_{\mathsf{cl}}\right)\mathclose{}^{2}\left\|B\right\|^{2}\log\mathopen{}\left(\frac{T}{\delta}\right)\mathclose{}\mathrm{ER}(\hat{\Phi},\hat{F}^{(H+1)}).$
(9)

Furthermore, for any cost function $h(\cdot)$ that is $L$-Lipschitz with
respect to the trajectory-wise metric
$d\mathopen{}\left(\bm{x}_{1:T},\bm{y}_{1:T}\right)\mathclose{}=\max_{1\leq
t\leq T}\left\|x[t]-y[t]\right\|$, we have the following bound on the expected
cost gap

$\displaystyle{\left|\mathbb{E}_{\hat{}\mathcal{P}_{1:T}}\mathopen{}\left[h(\hat{\bm{x}}_{1:T})\right]\mathclose{}-\mathbb{E}_{\mathcal{P}^{\star}_{1:T}}\mathopen{}\left[h(\bm{x}^{\star}_{1:T})\right]\mathclose{}\right|}$
$\displaystyle\lesssim
L\mathcal{J}(A_{\mathsf{cl}})\left\|B\right\|\sqrt{\log{T}}\sqrt{\mathrm{ER}(\hat{\Phi},\hat{F}^{(H+1)})}$
(10)

By invoking the bound on the excess risk from Theorem 3.1, condition (8) is
satisfied with probability at least $1-\delta^{\prime}$ if we have
sufficiently many samples $H$, $T$, $N_{1}$, $N_{2}$ such that

$\sigma_{z}^{2}\mathopen{}\left(\frac{kn_{x}\log\mathopen{}\left(N_{1}T\frac{\bar{\lambda}}{\underline{\lambda}}\right)\mathclose{}}{cN_{1}TH}+\frac{kn_{u}+\log(\frac{1}{\delta^{\prime}})}{N_{2}T}\right)\mathclose{}\lesssim\frac{\lambda_{\min}\mathopen{}\left(\Sigma_{x}^{(H+1)}\right)\mathclose{}}{\mathcal{J}\mathopen{}\left(A_{\mathsf{cl}}\right)\mathclose{}^{2}\left\|B\right\|^{2}}.$

The bound on excess risk from Theorem 3.1 may also be substituted into the
tracking error bound in (9) to find that with probability at least
$1-\delta-\delta^{\prime}$, the tracking error satisfies

$\max_{1\leq t\leq
T}\left\|\hat{x}[t]-x_{\star}[t]\right\|^{2}\lesssim\mathcal{J}\mathopen{}\left(A_{\mathsf{cl}}\right)\mathclose{}^{2}\left\|B\right\|^{2}\log\mathopen{}\left(\frac{T}{\delta}\right)\mathclose{}\sigma_{z}^{2}\mathopen{}\left(\frac{kn_{x}\log\mathopen{}\left(N_{1}T\frac{\bar{\lambda}}{\underline{\lambda}}\right)\mathclose{}}{cN_{1}TH}+\frac{kn_{u}+\log(\frac{1}{\delta^{\prime}})}{N_{2}T}\right)\mathclose{}.$

The above inequality provides the informal statement of the main result in
Theorem 1.1 by hiding $\log$ terms as well as the terms dependent on system
parameters. A bound for the expected cost gap
${\left|\mathbb{E}_{\hat{}\mathcal{P}_{1:T}}\mathopen{}\left[h(\hat{\bm{x}}_{1:T})\right]\mathclose{}-\mathbb{E}_{\mathcal{P}^{\star}_{1:T}}\mathopen{}\left[h(\bm{x}^{\star}_{1:T})\right]\mathclose{}\right|}$
can be similarly instantiated.

###### Remark 3.1

The dependence of the tracking error bound in (9) on the stability of the
target-task closed-loop system through $\mathcal{J}(A_{\mathsf{cl}})$ is tight
(see Appendix C). Intuitively, less stable systems exacerbate the input errors
from the learned controller.

###### Remark 3.2

Some immediate examples of $h(\cdot)$ include LQR state costs
$h(\bm{x}_{1:T})=\max_{t}\left\|Q^{1/2}x[t]\right\|$ and regularized tracking
costs
$h(\bm{x}_{1:T})=\max_{t}\left\|x[t]-x_{\mathrm{goal}}[t]\right\|+\lambda\left\|Rx[t]\right\|$.
Since $\frac{1}{T}\sum_{t=1}^{T}\left\|x[t]-y[t]\right\|\leq\max_{1\leq t\leq
T}\left\|x[t]-y[t]\right\|$, (10) holds with no modification for time-averaged
costs $h(\cdot)$. Bounds on the full LQR cost
$h\mathopen{}\left(\mathopen{}\left(\bm{x}_{1:T},K\right)\mathclose{}\right)\mathclose{}:=\max_{1\leq
t\leq T}\left\|\begin{bmatrix}Q^{1/2}\\\ R^{1/2}K\end{bmatrix}x[t]\right\|$
can be similarly derived, and are detailed in Appendix D.

##  4 Numerical Results

We consider a simple system with $n_{x}=4$ and $n_{u}=2$ from Hong et al.
(2021). In particular, let

$\displaystyle x[t+1]=\begin{bmatrix}.99&.03&-.02&-.32\\\ .01&.47&4.7&.00\\\
.02&-.06&.40&.00\\\
.01&-.04&.72&.99\end{bmatrix}x[t]+\begin{bmatrix}.01&.99\\\ -3.44&1.66\\\
-.83&.44\\\ -.47&.25\end{bmatrix}u[t]=:Ax[t]+Bu[t].$

We generate a collection of stabilizing controllers $K^{(1)}$,
$K^{(2)},\dots,K^{(H+1)}$ as LQR controllers with different cost matrices.
Specifically, let $R=I_{2}$, and $Q^{(h)}=\alpha^{(h)}I_{4}$ for
$\alpha^{(h)}\in\texttt{logspace}(-2,2,H+1)$, where $H=9$. The controllers
$K^{(h)}$ are then given by
$K^{(h)}=-(B^{\top}P^{(h)}B+R)^{-1}B^{\top}P^{(h)}A$, where $P^{(h)}$ solves
the following Discrete Algebraic Riccati equation:
$P^{(h)}=A^{\top}P^{(h)}A+A^{\top}P^{(h)}B(B^{\top}P^{(h)}B+R)^{-1}B^{\top}P^{(h)}A+Q^{(h)}$.

Next, assume that rather than directly observing the state, we obtain a high
dimensional observation given by an injective linear function of the state:
such an observation model can be viewed as a linear “perceptual sensor” or
camera. In particular, we suppose that $y_{t}=Gx_{t}$, where
$G\in\mathbb{R}^{50\times 4}$. For simplicity, we select the elements of $G$
i.i.d. from $\mathcal{N}(0,1)$, which ensures that $G$ is injective almost
surely. The dynamics of the observations may be written
$y[t+1]=GAx_{t}+GBu[t]=GAG^{\dagger}y[t]+GBu[t],$ with the input
$u[t]=K^{(h)}x[t]=K^{(h)}G^{\dagger}y[t]$. Define $\bar{A}=GAG^{\dagger}$ and
$\bar{B}=GB$, and $\bar{K}^{(h)}=K^{(h)}G^{\dagger}$. Consider the dynamics in
the face of process noise $w[t]\overset{i.i.d.}{\sim}\mathcal{N}(0,I_{50})$,
along with inputs corrupted by noise
$z[t]\overset{i.i.d.}{\sim}\mathcal{N}(0,I_{2})$:

$\displaystyle y[t+1]$
$\displaystyle=(\bar{A}+\bar{B}\bar{K}^{(h)})y[t]+\bar{B}z[t]+w[t],\quad
u[t]=\bar{K}^{(h)}y[t]+z[t].$ (11)

For the first $H$ controllers, we collect $N_{1}$ trajectories of length
$T=20$ to get the pairs
$\mathopen{}\left\\{\mathopen{}\left\\{\mathopen{}\left\\{(y_{i}^{h}[t],u_{i}^{h}[t])\right\\}\mathclose{}_{t=0}^{T-1}\right\\}\mathclose{}_{i=1}^{N_{1}}\right\\}\mathclose{}_{h=1}^{H}$.
For the last controller, we collect $N_{2}$ length $T=20$ trajectories to get
the dataset
$\mathopen{}\left\\{\mathopen{}\left\\{(y_{i}^{H+1}[t],u_{i}^{H+1}[t])\right\\}\mathclose{}_{t=0}^{T-1}\right\\}\mathclose{}_{i=1}^{N_{2}}$.
Our goal is to learn the controller $\bar{K}^{(H+1)}$ from the collected state
measurements and inputs. We compare the following ways of doing so:

---
(_a_) Tracking Error

---
(_b_)  Parameter Error

---
(_c_)  Percent Stable

Figure 1: We plot the tracking error between trajectories from the expert and
learned controllers, $\underset{{1\leq t\leq
T_{\textrm{test}}}}{\max}\left\|\hat{y}[t]-y_{\star}[t]\right\|^{2}$, the
parameter error,
$\mathopen{}\left(\left\|\hat{F}^{(H+1)}\hat{\Phi}-\bar{K}^{(H+1)}\right\|_{F}\right)\mathclose{}$,
and the percent of stable closed-loop systems for varying amounts of target
task data to compare multi-task IL to directly learning the controller from
target task data only. All metrics are plotted with respect to the lifted
system in Equation (11). Multi-task IL demonstrates a significant benefit over
direct IL in all metrics, especially when there is limited target task data.

  * •

Multi-task Imitation Learning: We observe that the data generating mechanism
ensures the existence of a low dimensional representation. In particular, one
possible $\Phi_{\star}$ is $G^{\dagger}$. Therefore, the stage is set for the
two step approach outlined in Section 2. In particular, we assume that the
true underlying state dimension is known, and we set the low dimensional
representation dimension to $k=4$, and jointly optimize over $\Phi$,
$F^{(1)},\dots,F^{(H)}$ in Problem (1). We approximately solve this problem
with $10000$ steps of alternating gradient descent using the adam optimizer
(Kingma and Ba, 2014) in optax (Babuschkin et al., 2020) with a learning rate
of $0.0001$. The learned representation is then fixed, and the target training
data is used to optimize $F^{(H+1)}$.

  * •

Direct Imitation Learning: We compare multi-task learning to direct learning,
which does not leverage the source data. In particular, given the target data,
direct learning solves the problem
$\operatorname*{minimize}_{F^{(H+1)}}\sum_{i=1}^{N_{1}}\sum_{t=0}^{T-1}\left\|F^{(H+1)}y_{i}^{(H+1)}[t]-u_{i}^{(H+1)}[t]\right\|^{2}$
666Note that another baseline leverages the fact that a $k$ dimensional
representation of the state exists, and learns it using target task data only
by solving
$\operatorname*{minimize}_{F^{(H+1)},\Phi}\sum_{i=1}^{N_{1}}\sum_{t=0}^{T-1}\left\|F^{(H+1)}\Phi
y^{(H+1)}_{i}[t]-u_{i}^{(H+1)}[t]\right\|^{2}$. For the current example,
however, $n_{u}<k$, so this approach is less efficient than the direct
learning approach proposed.. In this setting, we let $\hat{\Phi}=I_{50}.$

Note that a $2\times 50$ controller has $n_{u}\times n_{x}=100$ parameters to
learn from the target data. Meanwhile, multi-task imitation learning needs to
learn a total of $k\times n_{u}+k\times n_{x}=208$ parameters for the target
controller, but the $k\times n_{u}$ parameters are learned using source task
data. This leaves only $8$ parameters to learn using target task data.

Figure 1 plots three metrics that provide insight into the efficacy of these
approaches: the imitation gap given by $\max_{1\leq t\leq
T_{\textrm{test}}}\left\|\hat{y}[t]-y_{\star}[t]\right\|^{2}$ for length
$T_{\textrm{test}}=100$ observation trajectories $\hat{y}$ and $y_{\star}$
rolled out under the learned controller and expert controller, respectively,
with the same noise realizations; the parameter error,
$\left\|\hat{F}^{(H+1)}\hat{\Phi}-\bar{K}^{(H+1)}\right\|_{F}$; and the
percentage of trials where the learned controller is stabilizing. The trials
are over ten realizations of $G$, as well as ten realizations of the noise,
for a total of $100$ trials. For each trial, $N_{1}=10$, while $N_{2}$ sweeps
values in $[20]$. The medians for the imitation gap and parameter error are
shown, with the $20\%-80\%$ quantiles shaded.

---
(_a_) Tracking Error

---
(_b_)  Parameter Error

---
(_c_)  Percent Stable

Figure 2: We plot the tracking error between trajectories from the expert and
learned controllers, $\underset{{1\leq t\leq
T_{\textrm{test}}}}{\max}\left\|\hat{y}[t]-y_{\star}[t]\right\|^{2}$, the
parameter error,
$\mathopen{}\left(\left\|\hat{F}^{(H+1)}\hat{\Phi}-\bar{K}^{(H+1)}\right\|_{F}\right)\mathclose{}$,
and the percent of stable closed-loop systems for varying amounts of target
task data to compare multi-task IL to directly learning the controller from
target task data only. Similar to the setting of multi-task IL for transfer to
a new task, leveraging all source data to learn the controller for a single
source task provides a significant benefit in all three metrics over direct
IL.

In Figure 2, we additionally plot these metrics on one of the $H$ source
training tasks (arbitrarily chosen as $h=7$) for varying amounts of training
data to demonstrate the efficacy of the approach for multi-task learning.
Here, $N_{1}$ ranges from $1$ to $10$. We compare the controller
$\hat{F}^{(h)}\hat{\Phi}$ resulting from the shared training in Problem (1)
with the controller from directly training a controller on this task without
leveraging source task data. We note that our theoretical results, with mild
modification, also support the efficacy of this simultaneous training of a
representation and task weights.

##  5 Conclusion and Future Work

In this work, we study the sample complexity of multi-task imitation learning
for linear systems. We find that if the different sets of expert
demonstrations share a low dimensional representation, and the demonstrations
are sufficiently diverse, then doing multi-task representation learning will
lead to a smaller tracking error when deploying the learned policy in closed-
loop, compared to learning a policy only from target task data. Our results
are a first step towards understanding how the performance of a controller
trained on multi-task data relates to the characteristics of the multi-task
training data and the system being controlled. Some exciting directions for
future work would be to extend our analysis to nonlinear systems and nonlinear
representation functions, as well as other types of learning algorithms such
as model-based and model-free RL.

## Acknowledgements

Katie Kang is supported by a NSF Graduate Research Fellowship. Bruce D. Lee is
supported by the DoD through a National Defense Science & Engineering
Fellowship. Sergey Levine is supported in part by the DARPA Assured Autonomy
program. Nikolai Matni and Thomas Zhang are supported in part by NSF awards
CPS-2038873, CAREER award ECCS-2045834, and a Google Research Scholar award.

## References

  * Codevilla et al. (2018) Felipe Codevilla, Matthias Müller, Antonio López, Vladlen Koltun, and Alexey Dosovitskiy.  End-to-end driving via conditional imitation learning.  In _2018 IEEE international conference on robotics and automation (ICRA)_ , pages 4693–4700. IEEE, 2018.
  * Schaal (1999) Stefan Schaal.  Is imitation learning the route to humanoid robots?  _Trends in cognitive sciences_ , 3(6):233–242, 1999.
  * Ross et al. (2011) Stéphane Ross, Geoffrey Gordon, and Drew Bagnell.  A reduction of imitation learning and structured prediction to no-regret online learning.  In _Proceedings of the fourteenth international conference on artificial intelligence and statistics_ , pages 627–635. JMLR Workshop and Conference Proceedings, 2011.
  * Du et al. (2020) Simon S Du, Wei Hu, Sham M Kakade, Jason D Lee, and Qi Lei.  Few-shot learning via learning the representation, provably.  _arXiv preprint arXiv:2002.09434_ , 2020.
  * Tripuraneni et al. (2021) Nilesh Tripuraneni, Chi Jin, and Michael Jordan.  Provable meta-learning of linear representations.  In _International Conference on Machine Learning_ , pages 10434–10443. PMLR, 2021.
  * Teh et al. (2017) Yee Whye Teh, Victor Bapst, Wojciech Marian Czarnecki, John Quan, James Kirkpatrick, Raia Hadsell, Nicolas Heess, and Razvan Pascanu.  Distral: Robust multitask reinforcement learning, 2017.  URL https://arxiv.org/abs/1707.04175.
  * Espeholt et al. (2018) Lasse Espeholt, Hubert Soyer, Remi Munos, Karen Simonyan, Volodymir Mnih, Tom Ward, Yotam Doron, Vlad Firoiu, Tim Harley, Iain Dunning, Shane Legg, and Koray Kavukcuoglu.  Impala: Scalable distributed deep-rl with importance weighted actor-learner architectures, 2018.  URL https://arxiv.org/abs/1802.01561.
  * Hessel et al. (2018) Matteo Hessel, Hubert Soyer, Lasse Espeholt, Wojciech Czarnecki, Simon Schmitt, and Hado van Hasselt.  Multi-task deep reinforcement learning with popart, 2018.  URL https://arxiv.org/abs/1809.04474.
  * Singh et al. (2020) Avi Singh, Eric Jang, Alexander Irpan, Daniel Kappler, Murtaza Dalal, Sergey Levine, Mohi Khansari, and Chelsea Finn.  Scalable multi-task imitation learning with autonomous improvement, 2020\.  URL https://arxiv.org/abs/2003.02636.
  * Deisenroth et al. (2014) Marc Peter Deisenroth, Peter Englert, Jan Peters, and Dieter Fox.  Multi-task policy search for robotics.  In _2014 IEEE International Conference on Robotics and Automation (ICRA)_ , pages 3876–3881, 2014.  doi: 10.1109/ICRA.2014.6907421.
  * Finn et al. (2017) Chelsea Finn, Pieter Abbeel, and Sergey Levine.  Model-agnostic meta-learning for fast adaptation of deep networks, 2017\.  URL https://arxiv.org/abs/1703.03400.
  * Rakelly et al. (2019) Kate Rakelly, Aurick Zhou, Deirdre Quillen, Chelsea Finn, and Sergey Levine.  Efficient off-policy meta-reinforcement learning via probabilistic context variables, 2019.  URL https://arxiv.org/abs/1903.08254.
  * Duan et al. (2016) Yan Duan, John Schulman, Xi Chen, Peter L. Bartlett, Ilya Sutskever, and Pieter Abbeel.  Rl2: Fast reinforcement learning via slow reinforcement learning, 2016\.  URL https://arxiv.org/abs/1611.02779.
  * Yu et al. (2021) Tianhe Yu, Aviral Kumar, Yevgen Chebotar, Karol Hausman, Sergey Levine, and Chelsea Finn.  Conservative data sharing for multi-task offline reinforcement learning, 2021.  URL https://arxiv.org/abs/2109.08128.
  * Yang and Nachum (2021) Mengjiao Yang and Ofir Nachum.  Representation matters: Offline pretraining for sequential decision making, 2021.  URL https://arxiv.org/abs/2102.05815.
  * Lu et al. (2021) Rui Lu, Gao Huang, and Simon S. Du.  On the power of multitask representation learning in linear mdp, 2021\.  URL https://arxiv.org/abs/2106.08053.
  * Cheng et al. (2022) Yuan Cheng, Songtao Feng, Jing Yang, Hong Zhang, and Yingbin Liang.  Provable benefit of multitask representation learning in reinforcement learning, 2022.  URL https://arxiv.org/abs/2206.05900.
  * Lu et al. (2022) Rui Lu, Andrew Zhao, Simon S. Du, and Gao Huang.  Provable general function class representation learning in multitask bandits and mdps, 2022.  URL https://arxiv.org/abs/2205.15701.
  * Xu et al. (2020) Zhiyuan Xu, Kun Wu, Zhengping Che, Jian Tang, and Jieping Ye.  Knowledge transfer in multi-task deep reinforcement learning for continuous control, 2020.  URL https://arxiv.org/abs/2010.07494.
  * Maurer et al. (2015) Andreas Maurer, Massimiliano Pontil, and Bernardino Romera-Paredes.  The benefit of multitask representation learning, 2015.  URL https://arxiv.org/abs/1505.06279.
  * Arora et al. (2020) Sanjeev Arora, Simon S. Du, Sham Kakade, Yuping Luo, and Nikunj Saunshi.  Provable representation learning for imitation learning via bi-level optimization, 2020.  URL https://arxiv.org/abs/2002.10544.
  * Harrison et al. (2018) James Harrison, Apoorva Sharma, Roberto Calandra, and Marco Pavone.  Control adaptation via meta-learning dynamics.  In _Workshop on Meta-Learning at NeurIPS_ , volume 2018, 2018.
  * Richards et al. (2021) Spencer M Richards, Navid Azizan, Jean-Jacques Slotine, and Marco Pavone.  Adaptive-control-oriented meta-learning for nonlinear systems.  _arXiv preprint arXiv:2103.04490_ , 2021.
  * Richards et al. (2022) Spencer M Richards, Navid Azizan, Jean-Jacques Slotine, and Marco Pavone.  Control-oriented meta-learning.  _arXiv preprint arXiv:2204.06716_ , 2022.
  * Shi et al. (2021) Guanya Shi, Kamyar Azizzadenesheli, Michael O’Connell, Soon-Jo Chung, and Yisong Yue.  Meta-adaptive nonlinear control: Theory and algorithms.  _Advances in Neural Information Processing Systems_ , 34:10013–10025, 2021.
  * Muthirayan et al. (2022) Deepan Muthirayan, Dileep Kalathil, and Pramod P Khargonekar.  Meta-learning online control for linear dynamical systems.  _arXiv preprint arXiv:2208.10259_ , 2022.
  * Wang et al. (2021) Rui Wang, Robin Walters, and Rose Yu.  Meta-learning dynamics forecasting using task inference.  _arXiv preprint arXiv:2102.10271_ , 2021.
  * Modi et al. (2021) Aditya Modi, Mohamad Kazem Shirani Faradonbeh, Ambuj Tewari, and George Michailidis.  Joint learning of linear time-invariant dynamical systems.  _arXiv preprint arXiv:2112.10955_ , 2021.
  * Li et al. (2022) Lidong Li, Claudio De Persis, Pietro Tesi, and Nima Monshizadeh.  Data-based transfer stabilization in linear systems.  _arXiv preprint arXiv:2211.05536_ , 2022.
  * Baxter (1995) Jonathan Baxter.  Learning internal representations.  In _Proceedings of the eighth annual conference on Computational learning theory_ , pages 311–320, 1995.
  * Crammer et al. (2008) Koby Crammer, Michael Kearns, and Jennifer Wortman.  Learning from multiple sources.  _Journal of Machine Learning Research_ , 9(8), 2008.
  * Maurer et al. (2016) Andreas Maurer, Massimiliano Pontil, and Bernardino Romera-Paredes.  The benefit of multitask representation learning.  _Journal of Machine Learning Research_ , 17(81):1–32, 2016.
  * Tripuraneni et al. (2020) Nilesh Tripuraneni, Michael Jordan, and Chi Jin.  On the theory of transfer learning: The importance of task diversity.  _Advances in Neural Information Processing Systems_ , 33:7852–7862, 2020.
  * Chua et al. (2021) Kurtland Chua, Qi Lei, and Jason D Lee.  How fine-tuning allows for effective meta-learning.  _Advances in Neural Information Processing Systems_ , 34:8871–8884, 2021.
  * Abbasi-Yadkori et al. (2011) Yasin Abbasi-Yadkori, Dávid Pál, and Csaba Szepesvári.  Online least squares estimation with self-normalized processes: An application to bandit problems.  _arXiv preprint arXiv:1102.2670_ , 2011.
  * Pfrommer et al. (2022) Daniel Pfrommer, Thomas TCK Zhang, Stephen Tu, and Nikolai Matni.  Tasil: Taylor series imitation learning.  _arXiv preprint arXiv:2205.14812_ , 2022.
  * Hong et al. (2021) J Hong, N Moehle, and S Boyd.  Lecture notes in “introduction to matrix methods”, 2021.
  * Kingma and Ba (2014) Diederik P Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  _arXiv preprint arXiv:1412.6980_ , 2014.
  * Babuschkin et al. (2020) Igor Babuschkin, Kate Baumli, Alison Bell, Surya Bhupatiraju, Jake Bruce, Peter Buchlovsky, David Budden, Trevor Cai, Aidan Clark, Ivo Danihelka, Claudio Fantacci, Jonathan Godwin, Chris Jones, Ross Hemsley, Tom Hennigan, Matteo Hessel, Shaobo Hou, Steven Kapturowski, Thomas Keck, Iurii Kemaev, Michael King, Markus Kunesch, Lena Martens, Hamza Merzic, Vladimir Mikulik, Tamara Norman, John Quan, George Papamakarios, Roman Ring, Francisco Ruiz, Alvaro Sanchez, Rosalia Schneider, Eren Sezener, Stephen Spencer, Srivatsan Srinivasan, Luyu Wang, Wojciech Stokowiec, and Fabio Viola.  The DeepMind JAX Ecosystem, 2020.  URL http://github.com/deepmind.
  * Vershynin (2018) Roman Vershynin.  _High-dimensional probability: An introduction with applications in data science_ , volume 47.  Cambridge university press, 2018.
  * Laurent and Massart (2000) Beatrice Laurent and Pascal Massart.  Adaptive estimation of a quadratic functional by model selection.  _Annals of Statistics_ , pages 1302–1338, 2000.
  * Jedra and Proutiere (2020) Yassir Jedra and Alexandre Proutiere.  Finite-time identification of stable linear systems optimality of the least-squares estimator.  In _2020 59th IEEE Conference on Decision and Control (CDC)_ , pages 996–1001, 2020.  doi: 10.1109/CDC42340.2020.9304362.
  * Chen et al. (2016) Yan Mei Chen, Xiao Shan Chen, and Wen Li.  On perturbation bounds for orthogonal projections.  _Numerical Algorithms_ , 73(2):433–444, 2016\.
  * Xu (2020) Xuefeng Xu.  On the perturbation of an l2-orthogonal projection.  _Journal of Computational and Applied Mathematics_ , 368:112327, 2020.
  * Liang et al. (2015) Tengyuan Liang, Alexander Rakhlin, and Karthik Sridharan.  Learning with square loss: Localization through offset rademacher complexity.  In _Conference on Learning Theory_ , pages 1260–1285. PMLR, 2015\.
  * Khalil (1996) H.K. Khalil.  _Nonlinear Systems_.  Prentice Hall, 1996.  ISBN 9780132280242.  URL https://books.google.com/books?id=qiBuQgAACAAJ.
  * Horn and Johnson (2012) Roger A Horn and Charles R Johnson.  _Matrix analysis_.  Cambridge university press, 2012.
  * Fazel et al. (2018) Maryam Fazel, Rong Ge, Sham Kakade, and Mehran Mesbahi.  Global convergence of policy gradient methods for the linear quadratic regulator.  In _International Conference on Machine Learning_ , pages 1467–1476. PMLR, 2018.
  * Mania et al. (2019) Horia Mania, Stephen Tu, and Benjamin Recht.  Certainty equivalence is efficient for linear quadratic control.  _Advances in Neural Information Processing Systems_ , 32, 2019.
  * Villani (2021) Cédric Villani.  _Topics in optimal transportation_ , volume 58.  American Mathematical Soc., 2021.

#### Additional Notation

The set of $n\times m$ matrices $(n\geq m)$ with orthonormal columns is
denoted by $\mathcal{O}_{n,m}$. The set of $d$ dimensional unit vectors is
denoted by $\mathcal{S}^{d-1}$. We let $P_{A}=A(A^{\top}A)^{\dagger}A^{\top}$
denote the projection onto $\textbf{span}(A)$, and let $P_{A}^{\perp}=I-P_{A}$
denote the projection onto $\textbf{span}(A)^{\perp}$. The Kronecker product
of a matrix $A$ with a matrix $B$ is denoted $A\otimes B$. The vectorization
of a matrix $A$ with columns $A_{1},\dots,A_{n}$ is denoted
$\operatorname{Vec}(A)=\begin{bmatrix}A_{1}^{\top}&A_{2}^{\top}&\dots&A_{n}^{\top}\end{bmatrix}^{\top}$.
We denote the chi-squared distribution with $k$ degrees of freedom by
$\chi^{2}(k)$. Finally, we define the following $N_{1}T\times n_{u}$ data
matrices:

$\displaystyle\mathbf{U}^{(h)}$
$\displaystyle=\begin{bmatrix}u_{1}^{(h)}[0]&\dots&u_{1}^{(h)}[T-1]&\dots&u_{N_{1}}^{(h)}[0]&\dots&u_{N_{1}}^{(h)}[T-1]\end{bmatrix}^{\top},$
$\displaystyle\mathbf{Z}^{(h)}$
$\displaystyle=\begin{bmatrix}z_{1}^{(h)}[0]&\dots&z_{1}^{(h)}[T-1]&\dots&z_{N_{1}}^{(h)}[0]&\dots&z_{N_{1}}^{(h)}[T-1]\end{bmatrix}^{\top}.$

##  Appendix A Bounding the covariance concentration

We introduce the following version of the Hanson-Wright inequality for sub-
gaussian quadratic forms (Vershynin, 2018, Theorem 6.3.2), specialized to the
Gaussian case via Laurent and Massart (2000).

######  Proposition A.1 (Concentration of Gaussian Quadratic Forms)

Let $z\sim\mathcal{N}(0,I_{n})$ be a random Gaussian vector, and
$R\in\mathbb{R}^{m\times n}$ is an arbitrary fixed matrix. Then the following
concentration inequalities hold for all $\varepsilon>0$:

$\displaystyle\mathbb{P}\mathopen{}\left[\left\|Rz\right\|^{2}\geq(1+\varepsilon)\left\|R\right\|_{F}^{2}\right]\mathclose{}\leq\exp\mathopen{}\left(-\frac{1}{4}\min\mathopen{}\left\\{\frac{\varepsilon^{2}}{4},\varepsilon\right\\}\mathclose{}\frac{\left\|R\right\|_{F}^{2}}{\left\|R\right\|^{2}}\right)\mathclose{}$
(12)
$\displaystyle\mathbb{P}\mathopen{}\left[{\left|\left\|Rz\right\|^{2}-\left\|R\right\|_{F}^{2}\right|}\geq\varepsilon\left\|R\right\|_{F}^{2}\right]\mathclose{}\leq
2\exp\mathopen{}\left(-\frac{1}{4}\min\mathopen{}\left\\{\frac{\varepsilon^{2}}{4},\varepsilon\right\\}\mathclose{}\frac{\left\|R\right\|_{F}^{2}}{\left\|R\right\|^{2}}\right)\mathclose{},$
(13)

Proof  By the rotational invariance of Gaussian random vectors, we have
$z^{\top}R^{\top}Rz\overset{d}{=}\sum_{i=1}^{\min\mathopen{}\left\\{d,n\right\\}\mathclose{}}\sigma_{i}^{2}z_{i}^{2}$,
where $\sigma_{i}$ is the $i$th singular value of $R$. Noting that $z_{i}^{2}$
are i.i.d. $\chi^{2}$ random variables with $1$ degree of freedom, Laurent and
Massart (2000, Lemma 1) provides the following upper and lower tail bounds on
Gaussian quadratic forms:

$\displaystyle\mathbb{P}\mathopen{}\left[\left\|Rz\right\|^{2}\geq\left\|R\right\|_{F}^{2}+2\left\|R^{\top}R\right\|_{F}\sqrt{t}+2\left\|R\right\|^{2}t\right]\mathclose{}\leq\exp(-t)$
$\displaystyle\mathbb{P}\mathopen{}\left[\left\|Rz\right\|^{2}\leq\left\|R\right\|_{F}^{2}-2\left\|R^{\top}R\right\|_{F}\sqrt{t}\right]\mathclose{}\leq\exp(-t).$

We want to leverage these bounds to prove multiplicative concentration bounds.
Focusing on the upper tail first, we observe that

$\displaystyle\left\|R\right\|_{F}^{2}+2\left\|R^{\top}R\right\|_{F}\sqrt{t}+2\left\|R\right\|^{2}t$
$\displaystyle\leq\left\|R\right\|_{F}^{2}+2\left\|R\right\|\left\|R\right\|_{F}\sqrt{t}+2\left\|R\right\|^{2}t$
$\displaystyle\leq\left\|R\right\|_{F}^{2}+4\max\mathopen{}\left\\{\left\|R\right\|\left\|R\right\|_{F}\sqrt{t},\left\|R\right\|^{2}t\right\\}\mathclose{}$
$\displaystyle=:(1+\varepsilon)\left\|R\right\|_{F}^{2},$

such that by construction

$\mathbb{P}\mathopen{}\left[\left\|Rz\right\|^{2}\geq(1+\varepsilon)\left\|R\right\|_{F}^{2}\right]\mathclose{}\leq\exp(-t).$

Solving for the two cases of $t$ in terms of $\varepsilon$ and taking their
minimum, we find that

$\mathbb{P}\mathopen{}\left[\left\|Rz\right\|^{2}\geq(1+\varepsilon)\left\|R\right\|_{F}^{2}\right]\mathclose{}\leq\exp\mathopen{}\left(-\frac{1}{4}\min\mathopen{}\left\\{\frac{\varepsilon^{2}}{4},\varepsilon\right\\}\mathclose{}\frac{\left\|R\right\|_{F}^{2}}{\left\|R\right\|^{2}}\right)\mathclose{},$

which completes the upper tail bound. Repeating this for the lower tail bound,
we set
$t=\frac{1}{4}\frac{\left\|R\right\|_{F}^{2}}{\left\|R\right\|^{2}}\varepsilon^{2}$
to get

$\mathbb{P}\mathopen{}\left[\left\|Rz\right\|^{2}\leq\left\|R\right\|_{F}^{2}-\varepsilon\left\|R\right\|_{F}^{2}\right]\mathclose{}\leq\exp\mathopen{}\left(-\frac{1}{4}\frac{\left\|R\right\|_{F}^{2}}{\left\|R\right\|^{2}}\varepsilon^{2}\right)\mathclose{},$

which for a given $\varepsilon>0$ is always upper bounded by the upper tail
bound.

The following $\varepsilon$-net argument is standard, and a proof may be found
in Chapter 4 of Vershynin (2018).

###### Lemma A.1

Let $W$ be and $d\times d$ random matrix, and $\varepsilon\in[0,1/2)$.
Furthermore, let $\mathcal{N}$ be an $\varepsilon$-net of $\mathcal{S}^{d-1}$
with minimal cardinality. Then for all $\rho>0$,

$\operatorname{\mathbb{P}}\mathopen{}\left[\left\|W\right\|>\rho\right]\mathclose{}\leq\mathopen{}\left(\frac{2}{\varepsilon}+1\right)\mathclose{}^{d}\max_{x\in\mathcal{N}}\operatorname{\mathbb{P}}\mathopen{}\left[{\left|x^{\top}Wx\right|}>(1-2\epsilon)\rho\right]\mathclose{}.$

We now prove a generalized result on the concentration of covariances from
Jedra and Proutiere (2020) in which the covariates are pre and post multiplied
by a matrix with orthonormal rows and its transpose, respectively. This
demonstrates error bounds scaling with the lower dimension of the orthonormal
matrix.

######  Lemma A.2 (Modified Lemma 2 of Jedra and Proutiere (2020))

Let $A\in\mathbb{R}^{n\times n}$ satisfy $\rho(A)<1$ and let $\Sigma_{w}\succ
0$ with dimension $n\times n$. Let $\Sigma_{x}$ solve the discrete Lyapunov
equation:

$\Sigma_{x}=A\Sigma_{x}A^{\top}+\Sigma_{w}.$

Consider drawing $N_{1}$ trajectories of length $T$ from a linear system
$x_{i}[t+1]=Ax_{i}[t]+w_{i}[t]$, where
$w_{i}[t]\overset{i.i.d}{\sim}\mathcal{N}(0,\Sigma_{w})$ for $i=1,\dots N_{1}$
and $t=0,\dots T-1$, and
$x_{i}[0]\overset{i.i.d.}{\sim}\mathcal{N}(0,\Sigma_{x})$ for
$i=1,\dots,N_{1}$. Let $\mathbf{X}$ be the matrix of stacked state data as in
Equation (4). Let $U\in\mathcal{O}_{n,d}$ with $d\leq n$ be independent of
$\mathbf{X}$. Define
$M=\mathopen{}\left(N_{1}TU^{\top}\Sigma_{x}U\right)\mathclose{}^{-1/2}$. Then

$\left\|(\mathbf{X}UM)^{\top}\mathbf{X}UM-I\right\|>\max\mathopen{}\left\\{\varepsilon,\varepsilon^{2}\right\\}\mathclose{}$

holds with probability at most

$2\exp\mathopen{}\left(-c_{1}\varepsilon^{2}\frac{N_{1}T}{\kappa(\Sigma_{x})\mathcal{J}(A)^{2}}+c_{2}d\right)\mathclose{}$

for some absolute constants $c_{1}$ and $c_{2}$.

Proof

Note that we can write

$\mathbb{E}\mathopen{}\left[(\mathbf{X}U)^{\top}\mathbf{X}U\right]\mathclose{}=\mathbb{E}\mathopen{}\left[U^{\top}\sum_{i=1}^{N_{1}}\sum_{t=0}^{T-1}x_{i}[t]x_{i}[t]^{\top}U\right]\mathclose{}=N_{1}TU^{\top}\Sigma_{x}U.$

Then

$\displaystyle\left\|(\mathbf{X}UM)^{\top}(\mathbf{X}UM)-I\right\|$
$\displaystyle=\sup_{v\in\mathcal{S}^{d-1}}{\left|v^{\top}\mathopen{}\left((\mathbf{X}UM)^{\top}\mathbf{X}UM-I\right)\mathclose{}v\right|}$
$\displaystyle=\sup_{v\in\mathcal{S}^{d-1}}{\left|v^{\top}\mathopen{}\left((\mathbf{X}UM)^{\top}\mathbf{X}UM\right)\mathclose{}v-\mathbb{E}v^{\top}\mathopen{}\left((\mathbf{X}UM)^{\top}\mathbf{X}UM\right)\mathclose{}v\right|}$
$\displaystyle=\sup_{v\in\mathcal{S}^{d-1}}{\left|\left\|\mathbf{X}UMv\right\|^{2}-\mathbb{E}\left\|\mathbf{X}UMv\right\|^{2}\right|}$
$\displaystyle\overset{(i)}{=}\sup_{v\in\mathcal{S}^{d-1}}{\left|\left\|\sigma_{Mv}^{\top}\tilde{\Gamma}\xi\right\|^{2}-\mathbb{E}\left\|\sigma_{Mv}^{\top}\tilde{\Gamma}\xi\right\|^{2}\right|}$
$\displaystyle=\sup_{v\in\mathcal{S}^{d-1}}{\left|\left\|\sigma_{Mv}^{\top}\tilde{\Gamma}\xi\right\|^{2}-\left\|\sigma_{Mv}^{\top}\tilde{\Gamma}\right\|_{F}^{2}\right|},$

where $(i)$ follows by defining

$\Gamma=I_{N_{1}}\otimes\begin{bmatrix}\Sigma_{x}^{1/2}\\\
A\Sigma_{x}^{1/2}&\Sigma_{w}^{1/2}&&0\\\ &&\ddots\\\
A^{T-1}\Sigma_{x}^{1/2}&\ldots&A\Sigma_{w}^{1/2}&\Sigma_{w}^{1/2}\end{bmatrix},$
$\xi=\operatorname{Vec}\mathopen{}\left(\begin{bmatrix}\eta_{0}[-1]&\dots&\eta_{0}[T-1]&\dots&\eta_{N_{1}}[0]&\dots&\eta_{N_{1}}[T-2]\end{bmatrix}\right)\mathclose{},$

where $\eta_{i}[t]\overset{i.i.d.}{\sim}\mathcal{N}(0,I)$ for $t\in[-1,T-2]$,
and

$\sigma_{Mv}:=I_{N_{1}T}\otimes(Mv),\qquad\tilde{\Gamma}:=(I_{N_{1}T}\otimes
U^{\top})\Gamma,$

and recalling that
$\operatorname{Vec}(U^{\top}\mathbf{X}^{\top})=(I_{N_{1}T}\otimes
U^{\top})\Gamma\xi=\tilde{\Gamma}\xi$. Next oberserve that for any
$v\in\mathcal{S}^{d-1}$, we can apply Proposition A.1 to show that

$\operatorname{\mathbb{P}}\mathopen{}\left[{\left|\left\|\sigma_{Mv}^{\top}\tilde{\Gamma}\xi\right\|^{2}-\left\|\sigma_{Mv}\tilde{\Gamma}\right\|_{F}^{2}\right|}>\rho\left\|\sigma_{Mv}^{\top}\tilde{\Gamma}\right\|_{F}^{2}\right]\mathclose{}\leq
2\exp\mathopen{}\left(-C\min\mathopen{}\left\\{\rho^{2},\rho\right\\}\mathclose{}\frac{\left\|\sigma_{Mv}^{\top}\tilde{\Gamma}\right\|_{F}^{2}}{\left\|\sigma_{Mv}^{\top}\tilde{\Gamma}\right\|^{2}}\right)\mathclose{}$

for some positive universal constant $C$. Next observe that
$\left\|\sigma_{Mv}\tilde{\Gamma}\right\|_{F}^{2}=1$, and
$\left\|\sigma_{Mv}^{\top}\tilde{\Gamma}\right\|\leq\left\|M\right\|\left\|\tilde{\Gamma}\right\|$,
and thus the right hand side reduces to

$2\exp\mathopen{}\left(-\frac{C\min\mathopen{}\left\\{\rho^{2},\rho\right\\}\mathclose{}}{\left\|M\right\|^{2}\left\|\tilde{\Gamma}\right\|^{2}}\right)\mathclose{}.$

Applying Lemma A.1 with $\varepsilon=\frac{1}{4}$, we have that

$\operatorname{\mathbb{P}}\mathopen{}\left[\left\|(\mathbf{X}UM)^{\top}\mathbf{X}UM-I\right\|>\rho\right]\mathclose{}\leq
2\cdot
9^{d}\exp\mathopen{}\left(-\frac{C\min\mathopen{}\left\\{\rho^{2}/4,\rho/2\right\\}\mathclose{}}{\left\|M\right\|^{2}\left\|\tilde{\Gamma}\right\|^{2}}\right)\mathclose{}.$

Setting
$\rho=\max\mathopen{}\left\\{\varepsilon,\varepsilon^{2}\right\\}\mathclose{}$
and rearranging constants provides

$\operatorname{\mathbb{P}}\mathopen{}\left[\left\|(\mathbf{X}UM)^{\top}\mathbf{X}UM-I\right\|>\max\mathopen{}\left\\{\varepsilon,\varepsilon^{2}\right\\}\mathclose{}\right]\mathclose{}\leq
2\exp\mathopen{}\left(-\frac{c_{1}\varepsilon^{2}}{\left\|M\right\|^{2}\left\|\Gamma\right\|^{2}}+c_{2}d\right)\mathclose{}.$

Lastly, note that

$\displaystyle\left\|M\right\|=\frac{\left\|(U^{\top}\Sigma_{x}U)^{-1/2}\right\|}{\sqrt{N_{1}T}}=\frac{1}{\sqrt{N_{1}T}}\frac{1}{\lambda_{d}(U^{\top}\Sigma_{x}U)^{1/2}}\leq\frac{1}{\sqrt{N_{1}T}}\frac{1}{\lambda_{\min}(\Sigma_{x})^{1/2}}$

and

$\displaystyle\left\|\tilde{\Gamma}\right\|$
$\displaystyle\leq\left\|\Gamma\right\|$
$\displaystyle=\left\|\begin{bmatrix}\Sigma_{x}^{1/2}\\\
A\Sigma_{x}^{1/2}&\Sigma_{w}^{1/2}&&0\\\ &&\ddots\\\
A^{T-1}\Sigma_{x}^{1/2}&\ldots&A\Sigma_{w}^{1/2}&\Sigma_{w}^{1/2}\end{bmatrix}\right\|$
$\displaystyle\leq\left\|\begin{bmatrix}\Sigma_{x}^{1/2}\\\ \vdots\\\
A^{T-1}\Sigma_{x}^{1/2}\end{bmatrix}\right\|+\left\|\begin{bmatrix}\Sigma_{w}^{1/2}&&0\\\
&&\ddots\\\
A^{T-2}\Sigma_{w}^{1/2}&\ldots&A\Sigma_{w}^{1/2}&\Sigma_{w}^{1/2}\end{bmatrix}\right\|$
$\displaystyle\leq\sum_{s=0}^{T-1}\left\|A^{s}\right\|\left\|\Sigma_{x}\right\|^{1/2}+\left\|\begin{bmatrix}I&&0\\\
&&\ddots\\\
A^{T-2}&\ldots&A&I\end{bmatrix}\right\|\left\|\Sigma_{w}\right\|^{1/2}$
$\displaystyle\leq 2\sum_{s\geq
0}\left\|A^{s}\right\|\left\|\Sigma_{x}\right\|^{1/2}$
$\displaystyle=2\mathcal{J}(A)\left\|\Sigma_{x}\right\|^{1/2}.$

where the final inequality follows by applying Lemma 5 of Jedra and Proutiere
(2020) and the fact that $\Sigma_{w}\preceq\Sigma_{x}$. The lemma then follows
from the fact that
$\kappa(\Sigma_{x})=\frac{\left\|\Sigma_{x}\right\|}{\lambda_{\min}(\Sigma_{x})}$.

The previous lemma can now be applied to show concentration of the empirical
covariance for both the source and target task.

###### Lemma A.3

Suppose $N_{1}T\geq
C\max_{h}\mathcal{J}\mathopen{}\left(A+BK^{(h)}\right)\mathclose{}^{2}\kappa\mathopen{}\left(\Sigma_{x}^{(h)}\right)\mathclose{}(n_{x}+\log(H/\delta))$
for $\delta\in(0,1)$, where $C$ is a universal numerical constant. Then with
probability at least $1-\frac{\delta}{10}$ over the states
$\mathbf{X}^{(1)},\dots,\mathbf{X}^{(H)}$ in the source tasks, we have

$0.9\Sigma_{x}^{(h)}\preceq\frac{1}{N_{1}T}\mathopen{}\left(\mathbf{X}^{(h)}\right)\mathclose{}^{\top}\mathbf{X}^{(h)}\preceq
1.1\Sigma_{x}^{(h)},\qquad\forall h\in[H].$

Proof  Applying Lemma A.2 with $U=I$ and $\varepsilon=0.1$, as long as $C\geq
100\frac{\max\mathopen{}\left\\{c_{2},\log(1/20)\right\\}\mathclose{}}{c_{1}}$,
then for any task $h\in[H]$, we have that

$\left\|(\mathbf{X}^{(h)}M)^{\top}\mathbf{X}^{(h)}M-I\right\|\leq 0.1,$

with probability at least
$1-2\exp\mathopen{}\left(-\frac{0.01c_{1}N_{1}T}{\mathcal{J}(A+BK^{(h)})^{2}\kappa(\Sigma_{x}^{(h)})}+c_{2}n_{x}\right)\mathclose{}\geq
1-\frac{\delta}{10H}$. This may be written equivalently as

$0.9\Sigma_{x}^{(h)}=0.9\frac{M^{-2}}{N_{1}T}\preceq\frac{\mathbf{X}^{(h)\top}\mathbf{X}^{(h)}}{N_{1}T}\preceq
1.1\frac{M^{-2}}{N_{1}T}=1.1\Sigma_{x}^{(h)}.$

Taking a union bound over the $H$ tasks provides the desired result.

###### Lemma A.4

Suppose $N_{2}T\geq
C\mathcal{J}\mathopen{}\left(A+BK^{(H+1)}\right)\mathclose{}^{2}\kappa\mathopen{}\left(\Sigma_{x}^{(H+1)}\right)\mathclose{}(k+\log(1/\delta))$
for $\delta\in(0,1)$, where $C$ is a universal numerical constant. Then for
any given matrix $\Phi\in\mathbb{R}^{n_{x}\times 2k}$ independent of
$\mathbf{X}^{(H+1)}$, with probability at least $1-\frac{\delta}{10}$ over
$\mathbf{X}^{(H+1)}$ we have

$0.9\Phi^{\top}\Sigma_{x}^{(H+1)}\Phi\preceq\frac{1}{N_{2}T}\Phi^{\top}\mathopen{}\left(\mathbf{X}^{(H+1)}\right)\mathclose{}^{\top}\mathbf{X}^{(H+1)}\Phi\preceq
1.1\Phi^{\top}\Sigma_{x}^{(H+1)}\Phi.$

Proof  Let $\Phi=USV^{\top}$ be the singular value decomposition of $\Phi$
with $U\in\mathbb{R}^{n_{x}\times 2k}$. Applying Lemma A.2 with
$\varepsilon=0.1$, we have that

$\displaystyle
0.9U^{\top}\Sigma_{x}^{(H+1)}U=0.9\frac{M^{-2}}{N_{2}T}\preceq\frac{U^{\top}\mathbf{X}^{(H+1)\top}\mathbf{X}^{(H+1)}U}{N_{2}T}\preceq
1.1\frac{M^{-2}}{N_{2}T}=1.1U^{\top}\Sigma_{x}^{(H+1)}U$ (14)

with probability at least
$1-2\exp\mathopen{}\left(-0.01c_{1}\frac{N_{2}T}{\mathcal{J}(A+BK^{(H+1)})^{2}\kappa(\Sigma_{x}^{(H+1)})}+2c_{2}k\right)\mathclose{}\geq
1-\frac{\delta}{10}$ as long as $C\geq
100\frac{\max\mathopen{}\left\\{c_{2},\log(1/20)\right\\}\mathclose{}}{c_{1}}$.
Left and right multiplying Equation (14) by $V^{\top}S$ and $SV$ respectively
results in

$0.9\Phi^{\top}\Sigma_{x}^{(H+1)}\Phi\preceq\frac{\Phi^{\top}\mathbf{X}^{(H+1)\top}\mathbf{X}^{(H+1)}\Phi}{N_{2}T}\preceq
1.1\Phi^{\top}\Sigma_{x}^{(H+1)}\Phi$

with probability at least $1-\frac{\delta}{10}$.

##  Appendix B Data Guarantees

We will now use the covariance concentration results to show concentration of
the source and target controllers to the optimal controllers. We begin by
recalling Lemma A.5 of Du et al. (2020).

###### Lemma B.1

There exists a subset $\mathcal{N}\subset\mathcal{O}_{d_{1},d_{2}}$ that is an
$\varepsilon$-net of $\mathcal{O}_{d_{1},d_{2}}$ in Frobenius norm such that
${\left|\mathcal{N}\right|}\leq\mathopen{}\left(\frac{6\sqrt{d_{2}}}{\varepsilon}\right)\mathclose{}^{d_{1}d_{2}}$.

Also, note the following fact.

###### Fact B.1

For size conforming $Z,X$, with $X$ full column rank,

$\displaystyle\sup_{F}\left[4\langle
Z,XF^{\top}\rangle-\left\|XF^{\top}\right\|_{F}^{2}\right]=4\left\|(X^{\top}X)^{-1/2}X^{\top}Z\right\|_{F}^{2}=4\left\|P_{X}Z\right\|_{F}^{2}.$

The following result on perturbation of projection matrices may be found in
Chen et al. (2016); Xu (2020).

######  Lemma B.2 (Perturbation of projection matrices)

Let $A,B$ be size conforming matrices both of the same rank. We have:

$\left\|P_{A}-P_{B}\right\|\leq\min\left\\{\left\|A^{\dagger}\right\|,\left\|B^{\dagger}\right\|\right\\}\left\|A-B\right\|.$

###### Lemma B.3

Let $(x_{t})_{t\geq 1}$ be a $\mathbb{R}^{d}$-valued process adapted to a
filtration $(\mathcal{F}_{t})_{t\geq 1}$. Let $(\eta_{t})_{t\geq 1}$ be a
$\mathbb{R}^{m}$-valued process adapted to $(\mathcal{F}_{t})_{t\geq 2}$.
Suppose that $(\eta_{t})_{t\geq 1}$ is a $\sigma$-sub-Gaussian martingale
difference sequence, i.e.,:

$\displaystyle\mathbb{E}[\eta_{t}\mid\mathcal{F}_{t}]$ $\displaystyle=0,$ (15)
$\displaystyle\mathbb{E}[\exp(\lambda\langle
v,\eta_{t}\rangle)\mid\mathcal{F}_{t}]$
$\displaystyle\leq\exp\mathopen{}\left(\frac{\lambda^{2}\sigma^{2}\left\|v\right\|^{2}}{2}\right)\mathclose{}\>\>\forall\mathcal{F}_{t}\text{-measurable
}\lambda\in\mathbb{R},v\in\mathbb{R}^{m}.$ (16)

For $\Lambda\in\mathbb{R}^{m\times d}$, let $(M_{t}(\Lambda))_{t\geq 1}$ be
the $\mathbb{R}$-valued process:

$\displaystyle
M_{t}(\Lambda)=\exp\left(\frac{1}{\sigma}\sum_{i=1}^{t}\langle\Lambda
x_{i},\eta_{i}\rangle-\frac{1}{2}\sum_{i=1}^{t}\left\|\Lambda
x_{i}\right\|^{2}\right).$ (17)

The process $(M_{t}(\Lambda))_{t\geq 1}$ satisfies
$\mathbb{E}[M_{t}(\Lambda)]\leq 1$ for all $t\geq 1$.

Proof  Let $M_{0}(\Lambda):=1$. Fix $t\geq 1$. Observe:

$\displaystyle\mathbb{E}[M_{t}(\Lambda)]$
$\displaystyle=\mathbb{E}[\mathbb{E}[M_{t}(\Lambda)\mid\mathcal{F}_{t}]]=\mathbb{E}\left[M_{t-1}(\Lambda)\mathbb{E}\left[\exp\mathopen{}\left(\frac{1}{\sigma}\langle\Lambda
x_{t},\eta_{t}\rangle\right)\mathclose{}\,\bigg{|}\,\mathcal{F}_{t}\right]\exp\mathopen{}\left(-\frac{1}{2}\left\|\Lambda
x_{t}\right\|^{2}\right)\mathclose{}\right]$
$\displaystyle\leq\mathbb{E}[M_{t-1}(\Lambda)].$

The following result generalizes the self-normalized martingale inequality
from Abbasi-Yadkori et al. (2011) to handle multiple matrix valued self-
normalized martingales.

######  Proposition B.1 (Generalized self-normalized martingale inequality)

Fix $H\in\mathbb{N}_{+}$. For $h\in[H]$, let $(x_{t}^{h},\eta_{t}^{h})_{t\geq
1}$ be a $\mathbb{R}^{d}\times\mathbb{R}^{m}$-valued process and
$(\mathcal{F}_{t}^{h})_{t\geq 1}$ be a filtration such that
$(x_{t}^{h})_{t\geq 1}$ is adapted to $(\mathcal{F}_{t}^{h})_{t\geq 1}$,
$(\eta_{t}^{h})_{t\geq 1}$ is adapted to $(\mathcal{F}_{t}^{h})_{t\geq 2}$,
and $(\eta_{t}^{h})_{t\geq 1}$ is a $\sigma$-sub-Gaussian martingale
difference sequence. Suppose that for all $h_{1}\neq h_{2}$, the process
$(x_{t}^{h_{1}},\eta_{t}^{h_{1}})$ is independent of
$(x_{t}^{h_{2}},\eta_{t}^{h_{2}})$. Fix (non-random) positive definite
matrices $\\{V^{h}\\}_{h=1}^{H}$. For $t\geq 1$ and $h\in[H]$, define:

$\displaystyle\bar{V}_{t}^{h}:=V^{h}+V^{h}_{t},\>\>V^{h}_{t}:=\sum_{i=1}^{t}x_{i}x_{i}^{\top},\>\>S_{t}^{h}:=\sum_{i=1}^{t}x_{i}\eta_{i}^{\top}.$
(18)

For any fixed $T\in\mathbb{N}_{+}$, with probability at least $1-\delta$:

$\displaystyle\sum_{h=1}^{H}\left\|(\bar{V}^{h}_{T})^{-1/2}S^{h}_{T}\right\|_{F}^{2}\leq
2\sigma^{2}\left[\sum_{h=1}^{H}\log\mathrm{det}((\bar{V}_{T}^{h})^{m/2}(V^{h})^{-m/2})+\log(1/\delta)\right].$
(19)

Proof  By rescaling, we assume that $\sigma=1$ without loss of generality. For
$h\in[H]$, define:

$\displaystyle M_{T}^{h}(\Lambda):=\exp\left(\sum_{t=1}^{T}\langle\Lambda
x_{t}^{h},\eta_{t}^{h}\rangle-\frac{1}{2}\sum_{t=1}^{T}\left\|\Lambda
x_{t}^{h}\right\|^{2}\right),\>\>\Lambda\in\mathbb{R}^{m\times d}.$ (20)

By Lemma B.3, $\mathbb{E}[M_{T}^{h}(\Lambda)]\leq 1$ for any fixed $\Lambda$.
Now, we use the method of mixtures argument from Abbasi-Yadkori et al. (2011).
Let $\Gamma\in\mathbb{R}^{m\times d}$ be a random matrix with i.i.d. $N(0,1)$
entries, independent of everything else. By the tower property,
$\mathbb{E}[M_{T}^{h}(\Gamma)]=\mathbb{E}[\mathbb{E}[M_{T}^{h}(\Gamma)]\mid\Gamma]\leq
1$. On the other hand, let us compute
$\mathbb{E}[M_{T}^{h}(\Gamma)\mid\mathcal{F}^{h}_{\infty}]$. Let $p(\Lambda)$
denote the density of $\Gamma$, and also let $p(g)$ denote the density of a
$d$-dimensional isotropic normal vector. Let us momentarily drop the
superscript $h$ since the computation is identical for every $h$. First, we
note that the proof of Theorem 3 in Abbasi-Yadkori et al. (2011) shows the
following identity for any fixed $u\in\mathbb{R}^{d}$,
$V\in\mathbb{R}^{m\times m}$ positive semidefinite, and
$V_{0}\in\mathbb{R}^{m\times m}$ positive definite:

$\displaystyle\int\exp\mathopen{}\left(\langle
g,u\rangle-\frac{1}{2}\left\|g\right\|^{2}_{V}\right)\mathclose{}p(g)\,\mathrm{d}g=\mathopen{}\left(\frac{\mathrm{det}(V_{0})}{\mathrm{det}(V_{0}+V)}\right)\mathclose{}^{1/2}\exp\mathopen{}\left(\frac{1}{2}\left\|u\right\|^{2}_{(V+V_{0})^{-1}}\right)\mathclose{}.$

With this identity and independence of the rows of $\Gamma$:

$\displaystyle\mathbb{E}[M_{T}(\Gamma)\mid\mathcal{F}_{\infty}]$
$\displaystyle=\int\exp\mathopen{}\left(\langle\Lambda^{\top},S_{T}\rangle-\frac{1}{2}\left\|\Lambda^{\top}\right\|_{V_{t}}^{2}\right)\mathclose{}p(\Lambda)\,\mathrm{d}\Lambda$
$\displaystyle=\int\exp\mathopen{}\left(\sum_{i=1}^{m}\langle\Lambda^{\top}e_{i},S_{T}e_{i}\rangle-\frac{1}{2}\left\|\Lambda^{\top}e_{i}\right\|^{2}_{V_{t}}\right)\mathclose{}p(\Lambda)\,\mathrm{d}\Lambda$
$\displaystyle=\prod_{i=1}^{m}\int\exp\mathopen{}\left(\langle
g,S_{T}e_{i}\rangle-\frac{1}{2}\left\|g\right\|^{2}_{V_{t}}\right)\mathclose{}p(g)\,\mathrm{d}g$
$\displaystyle=\mathopen{}\left(\frac{\mathrm{det}(V)}{\mathrm{det}(V+V_{T})}\right)\mathclose{}^{m/2}\exp\mathopen{}\left(\frac{1}{2}\sum_{i=1}^{m}\left\|S_{T}e_{i}\right\|^{2}_{(V+V_{T})^{-1}}\right)\mathclose{}$
$\displaystyle=\mathopen{}\left(\frac{\mathrm{det}(V)}{\mathrm{det}(V+V_{T})}\right)\mathclose{}^{m/2}\exp\mathopen{}\left(\frac{1}{2}\left\|(V+V_{T})^{-1/2}S_{T}\right\|_{F}^{2}\right)\mathclose{}.$

That is, for every $h\in[H]$, we have:

$\displaystyle\mathbb{E}\left[\mathopen{}\left(\frac{\mathrm{det}(V^{h})}{\mathrm{det}(\bar{V}^{h}_{T})}\right)\mathclose{}^{m/2}\exp\mathopen{}\left(\frac{1}{2}\left\|(\bar{V}_{T}^{h})^{-1/2}S_{T}\right\|_{F}^{2}\right)\mathclose{}\right]\leq
1.$

Now by Markov’s inequality and independence of the processes across $h$:

$\displaystyle\operatorname{\mathbb{P}}\left(\sum_{h=1}^{H}\left\|(\bar{V}_{T}^{h})^{-1/2}S_{T}\right\|_{F}^{2}>2\left[\sum_{h=1}^{H}\log\mathrm{det}((\bar{V}_{T}^{h})^{m/2}(V^{h})^{-m/2})+\log(1/\delta)\right]\right)$
$\displaystyle=\operatorname{\mathbb{P}}\left(\exp\mathopen{}\left(\frac{1}{2}\sum_{h=1}^{H}\left\|(\bar{V}_{T}^{h})^{-1/2}S_{T}\right\|_{F}^{2}\right)\mathclose{}>\delta^{-1}\prod_{h=1}^{H}\mathopen{}\left(\frac{\mathrm{det}(\bar{V}_{T}^{h})}{\mathrm{det}(V^{h})}\right)\mathclose{}^{m/2}\right)$
$\displaystyle=\operatorname{\mathbb{P}}\left(\prod_{h=1}^{H}\mathopen{}\left(\frac{\mathrm{det}(V^{h})}{\mathrm{det}(\bar{V}_{T}^{h})}\right)\mathclose{}^{m/2}\exp\mathopen{}\left(\frac{1}{2}\left\|(\bar{V}_{T}^{h})^{-1/2}S_{T}\right\|_{F}^{2}\right)\mathclose{}>\delta^{-1}\right)$
$\displaystyle\leq\delta\mathbb{E}\left[\prod_{h=1}^{H}\mathopen{}\left(\frac{\mathrm{det}(V^{h})}{\mathrm{det}(\bar{V}_{T}^{h})}\right)\mathclose{}^{m/2}\exp\mathopen{}\left(\frac{1}{2}\left\|(\bar{V}_{T}^{h})^{-1/2}S_{T}\right\|_{F}^{2}\right)\mathclose{}\right]$
$\displaystyle=\delta\prod_{h=1}^{H}\mathbb{E}\left[\mathopen{}\left(\frac{\mathrm{det}(V^{h})}{\mathrm{det}(\bar{V}_{T}^{h})}\right)\mathclose{}^{m/2}\exp\mathopen{}\left(\frac{1}{2}\left\|(\bar{V}_{T}^{h})^{-1/2}S_{T}\right\|_{F}^{2}\right)\mathclose{}\right]$
$\displaystyle\leq\delta.$

###### Remark B.1

The original self-normalized martingale bound from Abbasi-Yadkori et al.
(2011) holds for an arbitrary stopping time, and consequently _uniformly_ for
all time $t\in\mathbb{N}$ by a simple reduction. In contrast, Proposition B.1
holds for a fixed $t\in\mathbb{N}$, which suffices for our purposes.

######  Lemma B.4 (source controller concentration)

Under the setting and data assumptions of Theorem 3.1, with probability at
least $1-\frac{\delta}{5}$,

$\displaystyle\sum_{h=1}^{H}\left\|\mathbf{X}^{(h)}\mathopen{}\left(\hat{F}^{(h)}\hat{\Phi}-F^{(h)}_{\star}\Phi_{\star}\right)\mathclose{}^{\top}\right\|_{F}^{2}\lesssim\sigma_{z}^{2}\mathopen{}\left(kmH+kn\log\mathopen{}\left(N_{1}T\frac{\bar{\lambda}}{\underline{\lambda}}\right)\mathclose{}+\log\mathopen{}\left(\frac{1}{\delta}\right)\mathclose{}\right)\mathclose{}.$

Proof  Before we start the proof, we first define two subsets:

$\displaystyle\Theta_{1}$ $\displaystyle:=\\{(F_{1}\Phi,\dots,F_{h}\Phi)\mid
F_{i}\in\mathbb{R}^{n_{u}\times k},\>\>\Phi\in\mathbb{R}^{k\times n_{x}}\\},$
$\displaystyle\Theta_{2}(a,b)$
$\displaystyle:=\\{(F_{1}\Phi,\dots,F_{h}\Phi)\mid
F_{i}\in\mathbb{R}^{n_{u}\times
b},\>\>\Phi^{\top}\in\mathcal{O}_{a,b}\\},\>\>a\geq b.$

It is easy to see that $\Theta_{1}=\Theta_{2}(n_{x},k)$. Furthermore, a simple
argument shows that

$\Theta_{2}(n_{x},k)-\Theta_{2}(n_{x},k)\subset\Theta_{2}(n_{x},2k),$

where the minus sign indicates the Minkowski difference of two sets.

Throughout this proof, we will condition on the following event, which holds
with probability at least $1-\frac{\delta}{10}$ by Lemma A.3:

$\displaystyle
0.9\Sigma_{x}^{(h)}\preceq\frac{1}{N_{1}T}\mathopen{}\left(\mathbf{X}^{(h)}\right)\mathclose{}^{\top}\mathbf{X}^{(h)}\preceq
1.1\Sigma_{x}^{(h)},\qquad\forall h\in[H].$ (21)

By optimality of $\hat{F}^{(1)},\hat{F}^{(2)},\dots,\hat{F}^{(H)}$ and
$\hat{\Phi}$ for Problem (1), we have that

$\displaystyle\sum_{h=1}^{H}\left\|\mathbf{U}^{(h)}-\mathbf{X}^{(h)}(\hat{F}^{(h)}\hat{\Phi})^{\top}\right\|_{F}^{2}\leq\sum_{h=1}^{H}\left\|\mathbf{U}^{(h)}-\mathbf{X}^{(h)}(F^{(h)}_{\star}\Phi_{\star})^{\top}\right\|_{F}^{2}.$

Substituting in
$\mathbf{U}^{(h)}=\mathbf{X}^{(h)}(F^{(h)}_{\star}\Phi_{\star})^{\top}+\mathbf{Z}^{(h)}$
and re-arranging yields the basic inequality:

$\displaystyle\sum_{h=1}^{H}\left\|\mathbf{X}^{(h)}(\hat{F}^{(h)}\hat{\Phi}-F^{(h)}_{\star}\Phi_{\star})^{\top}\right\|_{F}^{2}\leq
2\sum_{h=1}^{H}\operatorname{Tr}\mathopen{}\left(\mathopen{}\left(\mathbf{Z}^{(h)}\right)\mathclose{}^{\top}\mathbf{X}^{(h)}(F^{(h)}_{\star}\Phi_{\star}-\hat{F}^{(h)}\hat{\Phi})^{\top}\right)\mathclose{}.$
(22)

Let $\Delta^{(h)}:=\hat{F}^{(h)}\hat{\Phi}-F^{(h)}_{\star}\Phi_{\star}$.
Multiplying the basic inequality above by two and re-arranging again, we
obtain the _offset basic inequality_ (Liang et al., 2015):

$\displaystyle\sum_{h=1}^{H}\left\|\mathbf{X}^{(h)}(\Delta^{(h)})^{\top}\right\|^{2}_{F}$
$\displaystyle\leq\sum_{h=1}^{H}4\operatorname{Tr}\mathopen{}\left({(\mathbf{Z}^{(h)})^{\top}}{\mathbf{X}^{(h)}}{(\Delta^{(h)})^{\top}}\right)\mathclose{}-\left\|\mathbf{X}^{(h)}(\Delta^{h}{(h)})^{\top}\right\|_{F}^{2}$
$\displaystyle\leq\sup_{\\{\Delta^{(h)}\\}_{h=1}^{H}\in\Theta_{1}-\Theta_{1}}\sum_{h=1}^{H}4\operatorname{Tr}\mathopen{}\left({(\mathbf{Z}^{(h)})^{\top}}{\mathbf{X}^{(h)}}{(\Delta^{(h)})^{\top}}\right)\mathclose{}-\left\|\mathbf{X}^{(h)}(\Delta^{h}{(h)})^{\top}\right\|_{F}^{2}$
$\displaystyle=\sup_{\\{\Delta^{(h)}\\}_{h=1}^{H}\in\Theta_{2}(n_{x},k)-\Theta_{2}(n_{x},k)}\sum_{h=1}^{H}4\operatorname{Tr}\mathopen{}\left({(\mathbf{Z}^{(h)})^{\top}}{\mathbf{X}^{(h)}}{(\Delta^{(h)})^{\top}}\right)\mathclose{}-\left\|\mathbf{X}^{(h)}(\Delta^{h}{(h)})^{\top}\right\|_{F}^{2}$
$\displaystyle\leq\sup_{\\{\Delta^{(h)}\\}_{h=1}^{H}\in\Theta_{2}(n_{x},2k)}\sum_{h=1}^{H}4\operatorname{Tr}\mathopen{}\left({(\mathbf{Z}^{(h)})^{\top}}{\mathbf{X}^{(h)}}{(\Delta^{(h)})^{\top}}\right)\mathclose{}-\left\|\mathbf{X}^{(h)}(\Delta^{h}{(h)})^{\top}\right\|_{F}^{2}$
$\displaystyle=\sup_{\Phi^{\top}\in\mathcal{O}_{n_{x},2k}}\sum_{h=1}^{H}\sup_{F_{h}\in\mathbb{R}^{n_{u}\times
2k}}\left[4\operatorname{Tr}\mathopen{}\left({(\mathbf{Z}^{(h)})^{\top}}{\mathbf{X}^{(h)}}\Phi^{\top}F_{h}^{\top}\right)\mathclose{}-\left\|\mathbf{X}^{(h)}\Phi^{\top}F_{h}^{\top}\right\|_{F}^{2}\right]$
$\displaystyle=4\sup_{\Phi^{\top}\in\mathcal{O}_{n_{x},2k}}\sum_{h=1}^{H}\left\|(\Phi(\mathbf{X}^{(h)})^{\top}\mathbf{X}^{(h)}\Phi^{\top})^{-1/2}\Phi(\mathbf{X}^{(h)})^{\top}\mathbf{Z}^{(h)}\right\|_{F}^{2}.$

The last equality above follows from B.1. We now derive an upper bound on the
sum for a fixed $\Phi^{\top}\in\mathcal{O}_{n_{x},2k}$. To do this, we invoke
Proposition B.1 with the trajectories within each task concatenated in
sequence: $V^{h}\leftarrow 0.9N_{1}T\Phi\Sigma_{x}^{(h)}\Phi^{\top}$,
$x_{t}^{h}\leftarrow\Phi x_{i}^{(h)}[t]$, and $\eta_{t}^{h}\leftarrow
z_{i}^{(h)}[t]$. Note that since
$\Phi(\mathbf{X}^{(h)})^{\top}\mathbf{X}^{(h)}\Phi^{\top}\succeq V^{h}$, we
then have that:

$(\Phi(\mathbf{X}^{(h)})^{\top}\mathbf{X}^{(h)}\Phi^{\top})^{-1}\preceq
2(\Phi(\mathbf{X}^{(h)})^{\top}\mathbf{X}^{(h)}\Phi^{\top}+V^{h})^{-1}.$

Consequently, with probability at least $1-\delta^{\prime},$

$\displaystyle\sum_{h=1}^{H}$
$\displaystyle\left\|(\Phi(\mathbf{X}^{(h)})^{\top}\mathbf{X}^{(h)}\Phi^{\top})^{-1/2}\Phi(\mathbf{X}^{(h)})^{\top}\mathbf{Z}^{(h)}\right\|_{F}^{2}$
$\displaystyle\lesssim\sum_{h=1}^{H}\left\|(\Phi(\mathbf{X}^{(h)})^{\top}\mathbf{X}^{(h)}\Phi^{\top}+V^{h})^{-1/2}\Phi(\mathbf{X}^{(h)})^{\top}\mathbf{Z}^{(h)}\right\|_{F}^{2}$
$\displaystyle\lesssim\sigma^{2}\mathopen{}\left[\sum_{h=1}^{H}\log\mathrm{det}\mathopen{}\left(\mathopen{}\left(\Phi\mathopen{}\left(\mathbf{X}^{(h)}\right)\mathclose{}^{\top}\mathbf{X}^{(h)}\Phi^{\top}+V^{h}\right)\mathclose{}^{n_{u}/2}\mathopen{}\left(V^{h}\right)\mathclose{}^{-n_{u}/2}\right)\mathclose{}+\log(1/\delta^{\prime})\right]\mathclose{}$
$\displaystyle=\sigma^{2}\mathopen{}\left[\frac{n_{u}}{2}\sum_{h=1}^{H}\log\mathrm{det}\mathopen{}\left(\mathopen{}\left(\Phi\mathopen{}\left(\mathbf{X}^{(h)}\right)\mathclose{}^{\top}\mathbf{X}^{(h)}\Phi^{\top}+V^{h}\right)\mathclose{}\mathopen{}\left(V^{h}\right)\mathclose{}^{-1}\right)\mathclose{}+\log(1/\delta^{\prime})\right]\mathclose{}$
$\displaystyle=\sigma^{2}\mathopen{}\left[\frac{n_{u}}{2}\sum_{h=1}^{H}\log\mathrm{det}\mathopen{}\left(\frac{1}{0.9N_{1}T}\mathopen{}\left(\Phi\mathopen{}\left(\mathbf{X}^{(h)}\right)\mathclose{}^{\top}\mathbf{X}^{(h)}\Phi^{\top}\right)\mathclose{}\mathopen{}\left(\Phi\Sigma_{x}^{(h)}\Phi^{\top}\right)\mathclose{}^{-1}+I_{2k}\right)\mathclose{}+\log(1/\delta^{\prime})\right]\mathclose{}$
$\displaystyle\leq\sigma^{2}\mathopen{}\left[\frac{n_{u}}{2}\sum_{h=1}^{H}\log\mathrm{det}\mathopen{}\left(\frac{1.1}{0.9}+1\right)\mathclose{}I_{2k}+\log(1/\delta^{\prime})\right]\mathclose{}$
$\displaystyle\lesssim\sigma_{z}^{2}(kHn_{u}+\log(1/\delta^{\prime})).$ (23)

This holds for a fixed $\Phi^{\top}\in\mathcal{O}_{n_{x},2k}$, so it remains
to union bound over $\mathcal{O}_{n_{x},2k}$. Let
$\\{\Phi_{i}^{\top}\\}_{i=1}^{N}$ be an $\varepsilon$-net of
$\mathcal{O}_{n_{x},2k}$ in the spectral norm of resolution to be determined.
For $\Phi^{\top}\in\mathcal{O}_{n_{x},2k}$, let $\Phi_{i}^{\top}$ denote the
nearest element in the cover. By triangle inequality and $(a+b)^{2}\leq
2(a^{2}+b^{2})$:

$\displaystyle\sum_{h=1}^{H}$
$\displaystyle\left\|(\Phi(\mathbf{X}^{(h)})^{\top}\mathbf{X}^{(h)}\Phi^{\top})^{-1/2}\Phi(\mathbf{X}^{(h)})^{\top}\mathbf{Z}^{(h)}\right\|_{F}^{2}$
$\displaystyle=\sum_{h=1}^{H}\left\|P_{\mathbf{X}^{(h)}\Phi}\mathbf{Z}^{(h)}\right\|_{F}^{2}$
$\displaystyle\leq
2\sum_{h=1}^{H}\left\|(P_{\mathbf{X}^{(h)}\Phi}-P_{\mathbf{X}^{(h)}\Phi_{i}})\mathbf{Z}^{(h)}\right\|_{F}^{2}+2\sum_{h=1}^{H}\left\|P_{\mathbf{X}^{(h)}\Phi_{i}}\mathbf{Z}^{(h)}\right\|^{2}_{F}$
$\displaystyle\leq
2\sum_{h=1}^{H}\left\|P_{\mathbf{X}^{(h)}\Phi}-P_{\mathbf{X}^{(h)}\Phi_{i}}\right\|^{2}_{2}\left\|\mathbf{Z}^{(h)}\right\|_{F}^{2}+2\sum_{h=1}^{H}\left\|P_{\mathbf{X}^{(h)}\Phi_{i}}\mathbf{Z}^{(h)}\right\|^{2}_{F}.$

By Lemma B.2,

$\displaystyle\left\|P_{\mathbf{X}^{(h)}\Phi}-P_{\mathbf{X}^{(h)}\Phi_{i}}\right\|$
$\displaystyle\leq\left\|(\Phi(\mathbf{X}^{(h)})^{\top}\mathbf{X}^{(h)}\Phi^{\top})^{-1}\Phi(\mathbf{X}^{(h)})^{\top}\right\|\left\|\mathbf{X}^{(h)}(\Phi-\Phi_{i})\right\|$
$\displaystyle\leq\frac{1.1}{0.9}\frac{\lambda_{\max}(\Sigma_{x}^{(h)})}{\lambda_{\min}(\Sigma_{x}^{(h)})}\varepsilon\lesssim\frac{\overline{\lambda}}{\underline{\lambda}}\varepsilon.$

Next, we need to bound
$\sum_{h=1}^{H}\left\|\mathbf{Z}^{(h)}\right\|^{2}_{F}$. Because each
$z_{i}^{(h)}[t]$ is i.i.d. $\mathcal{N}(0,\sigma_{z}^{2}I)$, this sum is
distributed as a $\sigma_{z}^{2}\psi$, where $\psi$ is a $\chi^{2}$ random
variable with $HN_{1}Tn_{u}$ degrees of freedom. Hence by standard $\chi^{2}$
tail bounds, with probability at least $1-\delta/20$,
$\sum_{h=1}^{H}\left\|\mathbf{Z}^{(h)}\right\|_{F}^{2}\lesssim\sigma_{z}^{2}(HN_{1}Tn_{u}+\log(1/\delta))$.
This prompts us to select
$\varepsilon\lesssim\frac{k}{N_{1}T\overline{\lambda}/\underline{\lambda}}$,
from which we conclude:

$\sum_{h=1}^{H}\left\|P_{\mathbf{X}^{(h)}\Phi}-P_{\mathbf{X}^{(h)}\Phi_{i}}\right\|^{2}\left\|\mathbf{Z}^{(h)}\right\|_{F}^{2}\lesssim\sigma_{z}^{2}kHn_{u}.$

By Lemma B.1, we can then bound the size of the $\varepsilon$-covering of
$\mathcal{O}_{n_{x},2k}$ by:

$N\leq\mathopen{}\left(\frac{c\sqrt{k}}{k}N_{1}T\frac{\overline{\lambda}}{\underline{\lambda}}\right)\mathclose{}^{2n_{x}k}.$

So now we take $\delta^{\prime}=\delta/(20N)$, and conclude. In particular, by
union bounding over the elements of
$\mathopen{}\left\\{\Phi_{i}\right\\}\mathclose{}_{i=1}^{N}$, we have that
with probability at least $1-\frac{\delta}{5}$,

$\displaystyle\sup_{\Phi^{\top}\in\mathcal{O}_{n_{x},2k}}$
$\displaystyle\sum_{h=1}^{H}\left\|(\Phi(\mathbf{X}^{(h)})^{\top}\mathbf{X}^{(h)}\Phi^{\top})^{-1/2}\Phi(\mathbf{X}^{(h)})^{\top}\mathbf{Z}^{(h)}\right\|_{F}^{2}$
$\displaystyle\lesssim\sigma_{z}^{2}\mathopen{}\left(kHn_{u}+n_{x}k\log\mathopen{}\left(N_{1}T\frac{\bar{\lambda}}{\underline{\lambda}}\right)\mathclose{}+\log\mathopen{}\left(\frac{1}{\delta}\right)\mathclose{}\right)\mathclose{}.$

The probability of $1-\frac{\delta}{5}$ arises from union bounding over
Equation (21), which holds with probability at least $1-\frac{\delta}{10}$,
the bound on $\sum_{h=1}^{H}\left\|\mathbf{Z}^{(h)}\right\|_{F}^{2}$ which
holds with probability at least $1-\frac{\delta}{20}$, and the bound in
Equation (23), which holds for all
$\Phi_{i}\in\mathopen{}\left\\{\Phi_{i}\right\\}\mathclose{}_{i=1}^{N}$ with
probability at least $1-\frac{\delta}{20}$.

######  Lemma B.5 (Lemma A.7 in Du et al. (2020))

For any two matrices $A_{1}$ and $A_{2}$ with the same number of columns which
satisfy $A_{1}^{\top}A_{1}\succeq A_{2}^{\top}A_{2}$, then for any matrix $B$,
we have

$A_{1}^{\top}P_{A_{1}B}^{\perp}A_{1}\succeq
A_{2}^{\top}P_{A_{2}B}^{\perp}A_{2}.$

As a consequence, for any matrices $B$ and $B^{\prime}$, we have

$\left\|P_{A_{1}B}^{\perp}A_{1}B^{\prime}\right\|_{F}^{2}\geq\left\|P_{A_{2}B}^{\perp}A_{2}B^{\prime}\right\|_{F}^{2}.$

######  Lemma B.6 (target controller concentration)

Under the setting and data assumptions of Theorem 3.1, with probability at
least $1-\frac{2\delta}{5}$,

$\displaystyle\frac{1}{N_{2}T}$
$\displaystyle\left\|\mathbf{X}^{(H+1)}\mathopen{}\left(\hat{F}^{(H+1)}\hat{\Phi}-F^{(H+1)}_{\star}\Phi_{\star}\right)\mathclose{}^{\top}\right\|_{F}^{2}\lesssim\bar{\sigma}_{z}^{2}\mathopen{}\left(\frac{kn_{x}\log\mathopen{}\left(N_{1}T\frac{\bar{\lambda}}{\underline{\lambda}}\right)\mathclose{}}{cN_{1}TH}+\frac{kn_{u}+\log(\frac{1}{\delta})}{N_{2}T}\right)\mathclose{}.$

Proof  By the optimality of $\hat{\Phi},\hat{F}^{(1)},\dots,\hat{F}^{(H)}$ for
Problem (1), we know that

$\hat{F}^{(h)}=\mathopen{}\left(\hat{\Phi}\mathopen{}\left(\mathbf{X}^{(h)}\right)\mathclose{}^{\top}\mathbf{X}^{(h)}\hat{\Phi}^{\top}\right)\mathclose{}^{-1}\hat{\Phi}\mathopen{}\left(\mathbf{X}^{(h)}\right)\mathclose{}^{\top}\mathbf{U}^{(h)}.$

Therefore,
$\mathbf{X}^{(h)}\mathopen{}\left(\hat{F}^{(h)}\hat{\Phi}\right)\mathclose{}^{\top}=P_{\mathbf{X}^{(h)}\hat{\Phi}^{\top}}\mathbf{U}^{(h)}=P_{\mathbf{X}^{(h)}\hat{\Phi}^{\top}}\mathopen{}\left(\mathbf{X}^{(h)}\mathopen{}\left(F^{(h)}_{\star}\Phi_{\star}\right)\mathclose{}^{\top}+\mathbf{Z}^{(h)}\right)\mathclose{}$.
Then by applying Lemma B.4, we have that with probability at least
$1-\frac{\delta}{5}$,

$\displaystyle\bar{\sigma}_{z}^{2}$
$\displaystyle\mathopen{}\left(kmH+kn\log\mathopen{}\left(N_{1}T\frac{\bar{\lambda}}{\underline{\lambda}}\right)\mathclose{}+\log\mathopen{}\left(\frac{1}{\delta}\right)\mathclose{}\right)\mathclose{}\gtrsim\sum_{h=1}^{H}\left\|\mathbf{X}^{(h)}\mathopen{}\left(\hat{F}^{(h)}\hat{\Phi}-F^{(h)}_{\star}\Phi_{\star}\right)\mathclose{}^{\top}\right\|_{F}^{2}$
$\displaystyle=\sum_{h=1}^{H}\left\|P_{\mathbf{X}^{(h)}\hat{\Phi}^{\top}}\mathopen{}\left(\mathbf{X}^{(h)}\mathopen{}\left(F^{(h)}_{\star}\Phi_{\star}\right)\mathclose{}^{\top}+Z^{(h)}\right)\mathclose{}-\mathbf{X}^{(h)}\mathopen{}\left(F^{(h)}_{\star}\Phi_{\star}\right)\mathclose{}^{\top}\right\|_{F}^{2}$
$\displaystyle=\sum_{h=1}^{H}\left\|(P_{\mathbf{X}^{(h)}\hat{\Phi}^{\top}}-I)\mathbf{X}^{(h)}\mathopen{}\left(F^{(h)}_{\star}\Phi_{\star}\right)\mathclose{}^{\top}+P_{\mathbf{X}^{(h)}\hat{\Phi}^{\top}}\mathbf{Z}^{(h)}\right\|_{F}^{2}$
$\displaystyle=\sum_{h=1}^{H}\left\|P_{\mathbf{X}^{(h)}\hat{\Phi}^{\top}}^{\perp}\mathbf{X}^{(h)}\mathopen{}\left(F^{(h)}_{\star}\Phi_{\star}\right)\mathclose{}^{\top}\right\|_{F}^{2}+\left\|P_{\mathbf{X}^{(h)}\hat{\Phi}^{\top}}\mathbf{Z}^{(h)}\right\|_{F}^{2}$
$\displaystyle\geq\sum_{h=1}^{H}\left\|P_{\mathbf{X}^{(h)}\hat{\Phi}^{\top}}^{\perp}\mathbf{X}^{(h)}\mathopen{}\left(F^{(h)}_{\star}\Phi_{\star}\right)\mathclose{}^{\top}\right\|_{F}^{2}$
$\displaystyle\overset{(i)}{\geq}0.9N_{1}T\sum_{h=1}^{H}\left\|P_{\mathopen{}\left(\Sigma^{(h)}\right)\mathclose{}^{1/2}\hat{\Phi}^{\top}}^{\perp}\mathopen{}\left(\Sigma^{(h)}\right)\mathclose{}^{1/2}\mathopen{}\left(F^{(h)}_{\star}\Phi_{\star}\right)\mathclose{}^{\top}\right\|_{F}^{2}$
$\displaystyle\overset{(ii)}{\geq}0.9cN_{1}T\sum_{h=1}^{H}\left\|P_{\mathopen{}\left(\Sigma^{(H+1)}\right)\mathclose{}^{1/2}\hat{\Phi}^{\top}}^{\perp}\mathopen{}\left(\Sigma^{(H+1)}\right)\mathclose{}^{1/2}\mathopen{}\left(F^{(h)}_{\star}\Phi_{\star}\right)\mathclose{}^{\top}\right\|_{F}^{2}$

where $(i)$ follow from the inequality in Lemma A.3 in addition to Lemma B.5,
while $(ii)$ follows from the inequality in (5) (note that we have already
conditioned on this event through Lemma B.4) in addition to Lemma B.5. If we
define

$\displaystyle\mathbf{F}_{\star}=\begin{bmatrix}F_{\star}^{(1)}\\\ \vdots\\\
F_{\star}^{(H)}\end{bmatrix},$

we can write the above result concisely as

$\displaystyle\bar{\sigma}_{z}^{2}$
$\displaystyle\mathopen{}\left(kmH+kn\log\mathopen{}\left(N_{1}T\frac{\bar{\lambda}}{\underline{\lambda}}\right)\mathclose{}+\log\mathopen{}\left(\frac{1}{\delta}\right)\mathclose{}\right)\mathclose{}\gtrsim
0.9cN_{1}T\left\|P_{\mathopen{}\left(\Sigma^{(H+1)}\right)\mathclose{}^{1/2}\hat{\Phi}^{\top}}^{\perp}\mathopen{}\left(\Sigma^{(H+1)}\right)\mathclose{}^{1/2}\Phi_{\star}^{\top}\mathbf{F}_{\star}^{\top}\right\|_{F}^{2}$
(24)

Now, letting
$\Phi=\begin{bmatrix}\hat{\Phi}^{\top}&\hat{\Phi}_{\star}^{\top}\end{bmatrix}$,
Lemma A.4 gives us
$\frac{1}{N_{2}T}\Phi^{\top}\mathopen{}\left(\mathbf{X}^{(H+1)}\right)\mathclose{}^{\top}\mathbf{X}^{(H+1)}\Phi\preceq
1.1\Phi^{\top}\Sigma_{x}^{(H+1)}\Phi$. Combining with Lemma B.5, we have that
with probability at least $1-\frac{\delta}{10}$,

$\displaystyle
1.1\left\|P_{\mathopen{}\left(\Sigma^{(H+1)}\right)\mathclose{}^{1/2}\hat{\Phi}^{\top}}^{\perp}\mathopen{}\left(\Sigma^{(H+1)}\right)\mathclose{}^{1/2}\Phi_{\star}^{\top}\mathbf{F}_{\star}^{\top}\right\|_{F}^{2}\geq\frac{1}{N_{2}T}\left\|P_{\mathbf{X}^{(H+1)}\hat{\Phi}^{\top}}^{\perp}\mathbf{X}^{(H+1)}\Phi_{\star}^{\top}\mathbf{F}_{\star}^{\top}\right\|_{F}^{2}.$

Plugging this in above, we have that with probability at least
$1-\frac{2\delta}{5}$,

$\displaystyle\bar{\sigma}_{z}^{2}$
$\displaystyle\mathopen{}\left(kmH+kn\log\mathopen{}\left(N_{1}T\frac{\bar{\lambda}}{\underline{\lambda}}\right)\mathclose{}+\log\mathopen{}\left(\frac{1}{\delta}\right)\mathclose{}\right)\mathclose{}\gtrsim\frac{cN_{1}}{N_{2}}\left\|P_{X^{(H+1)}\hat{\Phi}^{\top}}^{\perp}\mathbf{X}^{(H+1)}\Phi_{\star}^{\top}\mathbf{F}_{\star}^{\top}\right\|_{F}^{2}.$

To conclude the proof, note that
$\mathbf{X}^{(H+1)}\hat{\Phi}^{\top}\mathopen{}\left(\hat{F}^{(H+1)}\right)\mathclose{}^{\top}=P_{\mathbf{X}^{(H+1)}\hat{\Phi}^{\top}}\mathbf{U}^{(h)}$
and
$\mathbf{F}_{\star}^{\dagger}:=(\mathbf{F}_{\star}^{\top}\mathbf{F}_{\star})^{-1}\mathbf{F}_{\star}^{\top}$
such that $\mathbf{F}_{\star}^{\dagger}\mathbf{F}_{\star}=I_{k}$. Thus

$\displaystyle\frac{1}{N_{2}T}{\left\|\mathbf{X}^{(H+1)}\mathopen{}\left(\hat{F}^{(H+1)}\hat{\Phi}-F^{(H+1)}_{\star}\Phi_{\star}\right)\mathclose{}^{\top}\right\|_{F}^{2}}$
$\displaystyle=\frac{1}{N_{2}T}{\left\|P_{\mathbf{X}^{(H+1)}\hat{\Phi}^{\top}}^{\perp}\mathbf{X}^{(H+1)}\mathopen{}\left(F^{(H+1)}_{\star}\Phi_{\star}\right)\mathclose{}^{\top}+P_{\mathbf{X}^{(H+1)}\hat{\Phi}^{\top}}\mathbf{Z}^{(H+1)}\right\|_{F}^{2}}$
$\displaystyle\leq\frac{1}{N_{2}T}{\left\|P_{\mathbf{X}^{(H+1)}\hat{\Phi}^{\top}}^{\perp}\mathbf{X}^{(H+1)}\mathopen{}\left(F^{(H+1)}_{\star}\Phi_{\star}\right)\mathclose{}^{\top}\right\|_{F}^{2}}+\frac{1}{N_{2}T}\left\|P_{\mathbf{X}^{(H+1)}\hat{\Phi}^{\top}}\mathbf{Z}^{(H+1)}\right\|_{F}^{2}$
$\displaystyle=\frac{1}{N_{2}T}{\operatorname{Tr}\mathopen{}\left(P_{\mathbf{X}^{(H+1)}\hat{\Phi}^{\top}}^{\perp}\mathbf{X}^{(H+1)}\Phi_{\star}^{\top}\mathopen{}\left(F^{(H+1)}_{\star}\right)\mathclose{}^{\top}F^{(H+1)}_{\star}\Phi_{\star}\mathopen{}\left(\mathbf{X}^{(H+1)}\right)\mathclose{}^{\top}P_{\mathbf{X}^{(H+1)}\hat{\Phi}^{\top}}^{\perp}\right)\mathclose{}}$
$\displaystyle\qquad\qquad+\frac{1}{N_{2}T}\left\|P_{X^{(H+1)}\hat{\Phi}^{\top}}\mathbf{Z}^{(H+1)}\right\|_{F}^{2}$
$\displaystyle=\frac{1}{N_{2}T}{\operatorname{Tr}\mathopen{}\left(P_{\mathbf{X}^{(H+1)}\hat{\Phi}^{\top}}^{\perp}\mathbf{X}^{(H+1)}\Phi_{\star}^{\top}\mathbf{F}_{\star}^{\top}\mathopen{}\left(\mathbf{F}_{\star}^{\dagger}\right)\mathclose{}^{\top}\mathopen{}\left(F^{(H+1)}_{\star}\right)\mathclose{}^{\top}F^{(H+1)}_{\star}\mathbf{F}_{\star}^{\dagger}\mathbf{F}_{\star}\Phi_{\star}\mathopen{}\left(\mathbf{X}^{(H+1)}\right)\mathclose{}^{\top}P_{\mathbf{X}^{(H+1)}\hat{\Phi}^{\top}}^{\perp}\right)\mathclose{}}$
$\displaystyle\qquad\qquad+\frac{1}{N_{2}T}\left\|P_{X^{(H+1)}\hat{\Phi}^{\top}}\mathbf{Z}^{(H+1)}\right\|_{F}^{2}$
$\displaystyle\leq\frac{1}{N_{2}T}\left\|\mathopen{}\left(\mathbf{F}^{\dagger}_{\star}\right)\mathclose{}^{\top}{\mathopen{}\left(F^{(H+1)}_{\star}\right)\mathclose{}^{\top}F^{(H+1)}_{\star}}\mathbf{F}^{\dagger}_{\star}\right\|\operatorname{Tr}\mathopen{}\left(P_{\mathbf{X}^{(H+1)}\hat{\Phi}^{\top}}^{\perp}\mathbf{X}^{(H+1)}\Phi_{\star}^{\top}\mathbf{F}_{\star}^{\top}\mathbf{F}_{\star}\Phi_{\star}\mathopen{}\left(\mathbf{X}^{(H+1)}\right)\mathclose{}^{\top}P_{\mathbf{X}^{(H+1)}\hat{\Phi}^{\top}}^{\perp}\right)\mathclose{}$
$\displaystyle\qquad\qquad+\frac{1}{N_{2}T}\left\|P_{\mathbf{X}^{(H+1)}\hat{\Phi}^{\top}}\mathbf{Z}^{(H+1)}\right\|_{F}^{2}$
$\displaystyle\lesssim\frac{1}{N_{2}T}\left\|F^{(H+1)}\mathbf{F}_{\star}^{\dagger}\right\|^{2}\left\|P_{\mathbf{X}^{(H+1)}\hat{\Phi}^{\top}}^{\perp}\mathbf{X}^{(H+1)}\Phi_{\star}^{\top}\mathbf{F}_{\star}^{\top}\right\|_{F}^{2}+\frac{1}{N_{2}T}\left\|P_{\mathbf{X}^{(H+1)}\hat{\Phi}^{\top}}\mathbf{Z}^{(H+1)}\right\|_{F}^{2}$
$\displaystyle\lesssim\frac{1}{cN_{1}TH}\bar{\sigma}_{z}^{2}\mathopen{}\left(kn_{u}H+kn_{x}\log\mathopen{}\left(N_{1}T\frac{\bar{\lambda}}{\underline{\lambda}}\right)\mathclose{}+\log\mathopen{}\left(\frac{1}{\delta}\right)\mathclose{}\right)\mathclose{}+\frac{1}{N_{2}T}\left\|P_{\mathbf{X}^{(H+1)}\hat{\Phi}^{\top}}\mathbf{Z}^{(H+1)}\right\|_{F}^{2}$

where the last line follows from the inequality in (24) and applying
Assumption 3.1 on $\left\|F^{(H+1)}\mathbf{F}_{\star}^{\dagger}\right\|^{2}$.
As $\hat{\Phi}$ is independent from $\mathbf{X}^{(H+1)}$ and
$\mathbf{Z}^{(H+1)}$, the second term may be bounded by applying Proposition
B.1 as in Equation (23). In particular, with probability at least
$1-\frac{\delta}{5}$,

$\left\|P_{\mathbf{X}^{(H+1)}\hat{\Phi}^{\top}}\mathbf{Z}^{(H+1)}\right\|_{F}^{2}\leq\sigma_{z}^{2}(kn_{u}+\log(1/\delta)).$

Therefore,

$\displaystyle\frac{1}{N_{2}T}\left\|\mathbf{X}^{(H+1)}\mathopen{}\left(\hat{F}^{(H+1)}\hat{\Phi}-F^{(H+1)}_{\star}\Phi_{\star}\right)\mathclose{}^{\top}\right\|_{F}^{2}$
$\displaystyle\lesssim\frac{1}{cN_{1}TH}\sigma_{z}^{2}\mathopen{}\left(kn_{u}H+kn_{x}\log\mathopen{}\left(N_{1}T\frac{\bar{\lambda}}{\underline{\lambda}}\right)\mathclose{}+\log\mathopen{}\left(\frac{1}{\delta}\right)\mathclose{}\right)\mathclose{}+\frac{\sigma_{z}^{2}kn_{u}+\sigma_{z}^{2}\log(\frac{1}{\delta})}{N_{2}T}$
$\displaystyle=\sigma_{z}^{2}\mathopen{}\left(\frac{kn_{u}H}{cN_{1}TH}+\frac{kn_{x}\log\mathopen{}\left(N_{1}T\frac{\bar{\lambda}}{\underline{\lambda}}\right)\mathclose{}}{cN_{1}TH}+\frac{\log(1/\delta)}{cN_{1}TH}+\frac{kn_{u}+\log(\frac{1}{\delta})}{N_{2}T}\right)\mathclose{}$
$\displaystyle\lesssim\sigma_{z}^{2}\mathopen{}\left(\frac{kn_{x}\log\mathopen{}\left(N_{1}T\frac{\bar{\lambda}}{\underline{\lambda}}\right)\mathclose{}}{cN_{1}TH}+\frac{kn_{u}+\log(\frac{1}{\delta})}{N_{2}T}\right)\mathclose{}.$

See 3.1

Proof  The excess risk may be written as follows:

$\displaystyle\mathrm{ER}(\hat{\Phi},\hat{F}^{(H+1)})$
$\displaystyle=\frac{1}{2T}\mathbb{E}_{x[0],w[0],\dots,w[T-1]}\mathopen{}\left[\sum_{t=0}^{T-1}\left\|(F^{(H+1)}\Phi_{\star}-\hat{F}^{(H+1)}\hat{\Phi})x[t]\right\|_{2}^{2}\right]\mathclose{}$
$\displaystyle=\frac{1}{2T}\mathbb{E}_{x[0],w[0],\dots,w[T-1]}\mathopen{}\left[\sum_{t=0}^{T-1}\operatorname{Tr}\mathopen{}\left(\mathopen{}\left(F^{(H+1)}\Phi_{\star}-\hat{F}^{(H+1)}\hat{\Phi}\right)\mathclose{}x[t]x[t]^{\top}\mathopen{}\left(F^{(H+1)}\Phi_{\star}-\hat{F}^{(H+1)}\hat{\Phi}\right)\mathclose{}\right)\mathclose{}\right]\mathclose{}$
$\displaystyle=\frac{1}{2T}\sum_{t=0}^{T-1}\operatorname{Tr}\mathopen{}\left(\mathopen{}\left(F^{(H+1)}\Phi_{\star}-\hat{F}^{(H+1)}\hat{\Phi}\right)\mathclose{}\Sigma_{x}^{(H+1)}\mathopen{}\left(F^{(H+1)}\Phi_{\star}-\hat{F}^{(H+1)}\hat{\Phi}\right)\mathclose{}\right)\mathclose{}$
$\displaystyle=\frac{1}{2}\operatorname{Tr}\mathopen{}\left(\mathopen{}\left(F^{(H+1)}\Phi_{\star}-\hat{F}^{(H+1)}\hat{\Phi}\right)\mathclose{}\Sigma_{x}^{(H+1)}\mathopen{}\left(F^{(H+1)}\Phi_{\star}-\hat{F}^{(H+1)}\hat{\Phi}\right)\mathclose{}\right)\mathclose{}.$

Then by applying the inequality in Lemma A.4 with
$\Phi=\begin{bmatrix}\hat{\Phi}&\Phi_{\star}\end{bmatrix}$, we have

$\displaystyle
0.9\Phi^{\top}\Sigma_{x}^{(H+1)}\Phi\preceq\frac{1}{N_{2}T}\Phi^{\top}\mathopen{}\left(\mathbf{X}^{(H+1)}\right)\mathclose{}^{\top}\mathbf{X}^{(H+1)}\Phi.$

Using this result this above, we have that

$\displaystyle\mathrm{ER}(\hat{\Phi},\hat{F}^{(H+1)})$
$\displaystyle\leq\frac{1}{1.8N_{2}T}\operatorname{Tr}\mathopen{}\left(\mathopen{}\left(F^{(H+1)}\Phi_{\star}-\hat{F}^{(H+1)}\hat{\Phi}\right)\mathclose{}\mathopen{}\left(\mathbf{X}^{(H+1)}\right)\mathclose{}^{\top}\mathbf{X}^{(H+1)}\mathopen{}\left(F^{(H+1)}\Phi_{\star}-\hat{F}^{(H+1)}\hat{\Phi}\right)\mathclose{}\right)\mathclose{}$
$\displaystyle\leq\frac{1}{1.8N_{2}T}\left\|\mathopen{}\left(F^{(H+1)}\Phi_{\star}-\hat{F}^{(H+1)}\hat{\Phi}\right)\mathclose{}\mathbf{X}^{(H+1)}\right\|_{F}^{2}.$

The claim then follows by application of Lemma B.6.

##  Appendix C Bounding the imitation gap

We recall that the issue with translating a bound on the excess risk of the
target task $\mathrm{ER}(\hat{\Phi},\hat{F}^{(H+1)})$ into a bound on the
tracking error between the closed-loop learned and expert trajectories comes
from the fundamental distribution shift between the expert trajectories seen
during training and the trajectories generated by running the learned
controller in closed-loop. This has traditionally been a difficult problem to
analyze due to the circular nature of requiring the feedback controller error
to be small to guarantee small state deviation, which depends on the state
deviation being small to begin with.

Toward addressing this issue, recent work by Pfrommer et al. (2022) proposes a
theoretical framework for non-linear systems to bound the tracking error by
the imitation error via matching the higher-order input/state derivatives
$D^{p}_{x}\pi(x)$ of the expert controller. For linear systems, matching the
Jacobian is sufficient, where we note the Jacobian of a linear controller
$u[t]=Kx[t]$ is simply $K$. Furthermore, a generalization bound on the
empirical risk of a learned controller such as Theorem 3.1 implicitly bounds
the controller error $\left\|\hat{K}-K^{(H+1)}\right\|$.

However, the framework described in Pfrommer et al. (2022) crucially relies on
assuming a noiseless system, whereas the problem considered in this work is
only non-trivial in the presence of process noise. Furthermore, due to the
lack of excitatory noise and the generality of non-linear systems, the
tracking error bounds in Pfrommer et al. (2022) scale with
$\tilde{}\mathcal{O}\mathopen{}\left(\frac{\mathrm{\\#param}}{N_{2}\delta^{\prime}}\right)\mathclose{}$,
where $\delta^{\prime}$ is the failure probability over a new trajectory. To
that end, the main goal of this section is to introduce bounds on the tracking
error that scale multiplicatively with the favorable generalization bounds on
the excess risk, which improve with the trajectory length $T$, and also
improving the system-agnostic Markov scaling $1/\delta^{\prime}$ to
$\log(1/\delta^{\prime})$ in our linear systems setting.

Toward proving a bound on the imitation gap, we first introduce the notion of
general stochastic incremental stability ($\delta$-SISS). We recall
definitions of standard comparison functions (Khalil, 1996): a function
$\gamma(x)$ is class $\mathcal{K}$ if it is continuous, strictly increasing,
and satisfies $\gamma(0)=0$. A function $\beta(x,t)$ is class
$\mathcal{K}\mathcal{L}$ if it is continuous, $\beta(\cdot,t)$ is class
$\mathcal{K}$ for each fixed $t$, and $\beta(x,\cdot)$ is decreasing for each
fixed $x$.

######  Definition C.1 ($\delta$SISS)

Consider the discrete-time control system $x[t+1]=f(x[t],u[t],w[t])$ subject
to input perturbations $\Delta[t]$ and additive zero-mean process noise
$\mathopen{}\left\\{w[t]\right\\}\mathclose{}_{t\geq 0}$. The system
$f(x[t],u[t],w[t])$ is _incremental stochastic input-to-state stable_
($\delta$-SISS) if for all initial conditions $x[0],y[0]\in\mathcal{X}$,
perturbation sequences
$\mathopen{}\left\\{\Delta[0]\right\\}\mathclose{}_{t\geq 0}$ and noise
realizations $\mathopen{}\left\\{w[0]\right\\}\mathclose{}_{t\geq 0}$ there
exists a class $\mathcal{K}\mathcal{L}$ function $\beta$ and a class
$\mathcal{K}$ function $\gamma$ such that the trajectories
$x[t+1]=f(x[t],u[t],w[t])$ and $y[t+1]=f(y[t],u[t]+\Delta[t],w[t])$ satisfy
for all $t\in\mathbb{N}$

$\displaystyle\left\|x[t]-y[t]\right\|$
$\displaystyle\leq\beta\mathopen{}\left(\left\|x[0]-y[0]\right\|,t\right)\mathclose{}+\gamma\mathopen{}\left(\max_{0\leq
k\leq t-1}\left\|\Delta[k]\right\|\right)\mathclose{}.$ (25)

######  Lemma C.1 (Stable linear dynamical systems are $\delta$SISS)

Let $(A,B)$ describe a linear time-invariant system $x[t+1]=Ax[t]+Bu[t]$. Let
$A$ be Schur-stable, i.e. $\rho(A)<1$, where $\rho(\cdot)$ denotes the
spectral radius. Recall the constant $\mathcal{J}(A):=\sum_{t\geq
0}\left\|A^{t}\right\|$. Furthermore, for any given $\nu$ such that
$\rho(A)<\nu<1$, define the corresponding constant

$\displaystyle\tau(A,\nu):=\sup_{k\in\mathbb{N}}\frac{\left\|A^{k}\right\|}{\nu^{k}},$

which is guaranteed to be finite by Gelfand’s formula (Horn and Johnson, 2012,
Corollary 5.6.13). The linear system $(A,B)$ is incrementally stochastic
input-to-state stable, where we may set

$\displaystyle\beta(x,t)$ $\displaystyle:=\tau(A,\nu)\nu^{t}x$
$\displaystyle\gamma(x)$ $\displaystyle:=\mathcal{J}(A)\left\|B\right\|x.$

Proof  Let us define the nominal and perturbed states $x_{t}$ and $y_{t}$
defined by

$\displaystyle x[t+1]$ $\displaystyle=Ax[t]+Bu[t]+w[t],\quad x[0]=\xi_{1}$
$\displaystyle y[t+1]$ $\displaystyle=Ay[t]+B(u[t]+\Delta[t])+w[t],\quad
y[0]=\xi_{2}.$

We observe that by linearity, we may write

$\displaystyle x[t]$
$\displaystyle=A^{t}x[0]+\sum_{k=0}^{t-1}(A^{t-k-1}Bu[k]+A^{t-k-1}w[k])$
$\displaystyle y[t]$
$\displaystyle=A^{t}y[0]+\sum_{k=0}^{t-1}(A^{t-k-1}Bu[k]+A^{t-k-1}B\Delta[k]+A^{t-k-1}w[k]).$

Therefore, their difference can be written as

$\displaystyle x[t]-y[t]$
$\displaystyle=A^{t}(x[0]-y[0])-\sum_{k=0}^{t-1}A^{t-k-1}B\Delta[k]$
$\displaystyle\implies\left\|x[t]-y[t]\right\|$
$\displaystyle\leq\left\|A^{t}\right\|\left\|x[0]-y[0]\right\|+\mathopen{}\left(\sum_{k=0}^{t-1}\left\|A^{t-k-1}\right\|\right)\mathclose{}\left\|B\right\|\max_{0\leq
k\leq t-1}\left\|\Delta[k]\right\|$
$\displaystyle\leq\tau(A,\nu)\nu^{t}\left\|x[0]-y[0]\right\|+\mathcal{J}(A)\left\|B\right\|\max_{0\leq
k\leq t-1}\left\|\Delta[k]\right\|$

where we can match the functions $\beta(\cdot,t)$ and $\gamma(\cdot)$ to their
respective quantities above. Crucially, we observe that due to linearity, the
noise terms in $x[t]$ and $y[t]$ cancel out, and therefore $\delta$SISS is
effectively the same as the standard $\delta$ISS.

By showing stable linear systems are $\delta$SISS, we may adapt Corollary A.1
from Pfrommer et al. (2022) to yield the following guarantee with respect to
the expert closed-loop system of the target task $A+BK^{(H+1)}$, which allows
us to bound the imitation gap in terms of the maximal deviation between the
learned and target controller inputs.

###### Proposition C.1

Let the target task closed-loop system $(A+BK^{(H+1)},B)$ be $\delta$SISS for
$\beta(\cdot,t)$ and $\gamma(\cdot)$ defined in Lemma C.1. For a given test
controller $K$, and realizations of randomness
$x[0]\sim\mathcal{N}(0,\Sigma_{x}^{(H+1)})$,
$z[t]\sim\mathcal{N}(0,\sigma_{z}^{2}I)$, $w[t]\sim
N(0,\Sigma_{w}^{(H+1)}),t=0,\dots,T-1$, we write

$\displaystyle\begin{split}x_{\star}[t+1]&=(A+BK^{(H+1)})x_{\star}[t]+Bz[t]+w[t],\quad
x_{\star}[0]=x[0]\\\
\hat{x}[t+1]&=(A+BK)\hat{x}[t]+Bz[t]+w[t],\quad\hat{x}[0]=x[0].\end{split}$
(26)

In other words, $x_{\star}[t]$ is the state from rolling out the expert target
task controller $K^{(H+1)}$, and $\hat{x}[t]$ is the state from rolling out
the test controller $K$ under the same conditions. If $K$ satisfies

$\displaystyle\left\|K-K^{(H+1)}\right\|\leq\frac{1}{2\mathcal{J}(A+BK^{(H+1)})\left\|B\right\|},$

then the tracking error satisfies for any $T\geq 1$,

$\displaystyle\max_{1\leq t\leq
T}\left\|\hat{x}[t]-x_{\star}[t]\right\|\leq\max_{0\leq t\leq
T-1}2\mathcal{J}(A+BK^{(H+1)})\left\|B\right\|\left\|\mathopen{}\left(K-K^{(H+1)}\right)\mathclose{}x_{\star}[t]\right\|.$

In other words, Proposition C.1 allows us to bound the imitation gap by the
excess risk induced by controller $K$, as long as $K$ is sufficiently close to
the expert controller $K^{(H+1)}$ in the spectral norm.

Proof  We borrow Proposition 3.1 from Pfrommer et al. (2022), which states a
general result for non-linear controllers and non-linear systems, but in the
noiseless setting, and specify it to the linear systems setting with process
noise. In particular, the result states that for a given state-feedback
controller $\pi(x)$ and an expert state-feedback controller $\pi_{\star}(x)$
that induces a $\delta$-SISS closed-loop system
$x_{\star}[t+1]=f(x_{\star}[t],\pi_{\star}(x_{\star}[t]))+w[t]$ (the result
can be extended from $\delta$-ISS with no modification to the proof), then for
a given $\varepsilon>0$, if $\pi$ satisfies the following on an expert
trajectory generated by realizations of randomness $x[0],w[0],\dots,w[T-1]$:

$\displaystyle\max_{0\leq t\leq
T-1}\sup_{\left\|\delta\right\|\leq\varepsilon}\left\|\pi_{\star}(x_{\star}[t]+\delta)-\pi(x_{\star}[t]+\delta)\right\|$
$\displaystyle\leq\gamma^{-1}(\varepsilon),$ (27)

where $\gamma^{-1}$ is the (generalized) inverse of the class-$\mathcal{K}$
function $\gamma(\cdot)$, then the imitation gap is bounded by

$\max_{0\leq t\leq T-1}\left\|\hat{x}[t]-x_{\star}[t]\right\|\leq\varepsilon,$

where $\hat{x}[t]$ are the states generated by running controller $\pi$ in
closed-loop on the same noise realizations as those generating $x_{\star}[t]$.
Now, substituting $\pi(x):=Kx$, $\pi_{\star}(x):=K^{(H+1)}x$, and
$\gamma(x):=\mathcal{J}(A+BK^{(H+1)})\left\|B\right\|x$, such that
$\gamma^{-1}(x)=\mathopen{}\left(\mathcal{J}(A+BK^{(H+1)})\left\|B\right\|\right)\mathclose{}^{-1}x$,
we observe

$\displaystyle\max_{0\leq t\leq
T-1}\sup_{\left\|\delta\right\|\leq\varepsilon}\left\|\pi_{\star}(x_{\star}[t]+\delta)-\pi(x_{\star}[t]+\delta)\right\|$
$\displaystyle=$ $\displaystyle\max_{0\leq t\leq
T-1}\sup_{\left\|\delta\right\|\leq\varepsilon}\left\|(K-K^{(H+1)})x_{\star}[t]+(K-K^{(H+1)})\delta\right\|$
$\displaystyle\leq$ $\displaystyle\max_{0\leq t\leq
T-1}\left\|(K-K^{(H+1)})x_{\star}[t]\right\|+\left\|K-K^{(H+1)}\right\|\varepsilon.$

Therefore, in order to satisfy (27), it suffices to satisfy

$\displaystyle\max_{0\leq t\leq
T-1}\left\|(K-K^{(H+1)})x_{\star}[t]\right\|+\left\|K-K^{(H+1)}\right\|\varepsilon\leq\mathopen{}\left(\mathcal{J}(A+BK^{(H+1)})\left\|B\right\|\right)\mathclose{}^{-1}\varepsilon.$

Since $\varepsilon>0$ is arbitrary, setting
$\varepsilon:=\max_{t}2\mathcal{J}(A+BK^{(H+1)})\left\|B\right\|\left\|(K-K^{(H+1)})x_{\star}[t]\right\|$,
it is sufficient for

$\left\|K-K^{(H+1)}\right\|\leq\frac{1}{2\mathcal{J}(A+BK^{(H+1)})\left\|B\right\|},$

to satisfy the above inequality, which leads to the following bound on the
imitation gap

$\displaystyle\max_{0\leq t\leq T-1}\left\|\hat{x}[t]-x_{\star}[t]\right\|$
$\displaystyle\leq\varepsilon=\max_{0\leq t\leq
T-1}2\mathcal{J}(A+BK^{(H+1)})\left\|B\right\|\left\|Kx_{\star}[t]-K^{(H+1)}x_{\star}[t]\right\|.$

This completes the proof.

Before moving on, we discuss a few qualitative properties of Proposition C.1.
First off, a high-probability bound on the tracking error is somewhat higher
resolution than the LQR-type bounds in Fazel et al. (2018); Mania et al.
(2019), where the quantity of interest is the difference in the expected
infinite-horizon costs of the two controllers, which does not directly imply a
bound on the deviation between the states induced by the two controllers at a
given time. Furthermore, we note Proposition C.1 is meaningful precisely
because there is driving process noise. If the target task is noiseless, and
test controller $K$ is stabilizing, then
$\left\|\hat{x}[t]-x_{\star}[t]\right\|$ is trivially upper bounded by an
exponentially decaying quantity, which for a sufficiently long horizon $T$
will beat out any generalization bound scaling with
$\mathrm{poly}(N_{1},N_{2},H,T)$–however, this noiseless setting is
correspondingly uninteresting to study as a statistical problem.

We propose the following key result in the form of a high probability bound on
the tracking error.

######  Theorem C.1 (Full version of Theorem 3.2, (9))

Let the target task closed-loop system $(A+BK^{(H+1)},B)$ be $\delta$SISS for
$\beta(\cdot,t)$ and $\gamma(\cdot)$ as defined in Lemma C.1. Given a
generalization bound on the excess risk of the learned representation and
target task weights $(\hat{\Phi},\hat{F}^{(H+1)})$ (such as Theorem 3.1) of
the form: with probability greater than $1-\delta$

$\displaystyle\mathrm{ER}\mathopen{}\left(\hat{\Phi},\hat{F}^{(H+1)}\right)\mathclose{}$
$\displaystyle\leq f(N_{1},T,H,N_{2},\delta),$

we have the following bound on the tracking error. Assuming we have enough
samples such that

$f(N_{1},T,H,N_{2},\delta)\leq\frac{\lambda_{\min}(\Sigma_{x}^{(H+1)})}{8\mathcal{J}(A+BK^{(H+1)})^{2}\left\|B\right\|^{2}},$

then with probability greater than $1-\delta-\delta^{\prime}$, for a new
target task trajectory sampled with i.i.d. process randomness
$x[0]\sim\Sigma_{x}^{(H+1)}$, $w[t]\sim\Sigma_{w}^{(H+1)}$, $t=0,\dots,T-1$,
the tracking error satisfies

$\displaystyle\max_{1\leq t\leq T}\left\|\hat{x}[t]-x_{\star}[t]\right\|^{2}$
$\displaystyle\leq
4\mathcal{J}(A+BK^{(H+1)})^{2}\left\|B\right\|^{2}\mathopen{}\left(1+4\log\mathopen{}\left(\frac{T}{\delta^{\prime}}\right)\mathclose{}\right)\mathclose{}\mathrm{ER}(\hat{\Phi},\hat{W}^{H+1})$
(28)
$\displaystyle\lesssim\mathcal{J}(A+BK^{(H+1)})^{2}\left\|B\right\|^{2}\log\mathopen{}\left(\frac{T}{\delta^{\prime}}\right)\mathclose{}f(N_{1},T,H,N_{2},\delta),$
(29)

where the expert and learned trajectory states $\hat{x}[t]$ and $x_{\star}[t]$
are as defined in Proposition C.1.

Proof  By Proposition C.1, we have that if the learned controller
$\hat{K}=\hat{F}^{(H+1)}\hat{\Phi}$ satisfies

$\displaystyle\left\|\hat{K}-K^{(H+1)}\right\|\leq\frac{1}{2\mathcal{J}(A+BK^{(H+1)})\left\|B\right\|},$

where we plugged in the $\beta(\cdot,t)$, $\gamma(\cdot)$ functions derived
from Lemma C.1 on the target task expert closed-loop system
$(A+BK^{(H+1)},B)$, then we have for learned and expert trajectories generated
by independent instances of randomness $x[0],w[0],\dots,w[T-1]$ the following
imitation gap:

$\max_{1\leq t\leq T}\left\|\hat{x}[t]-x_{\star}[t]\right\|\leq\max_{0\leq
t\leq
T-1}2\mathcal{J}(A+BK^{(H+1)})\left\|B\right\|\left\|\mathopen{}\left(\hat{K}-K^{(H+1)}\right)\mathclose{}x_{\star}[t]\right\|.$

In order derive sample complexity bounds from this, we observe a
generalization bound on the excess risk

$\displaystyle\mathrm{ER}(\hat{\Phi},\hat{F}^{(H+1)})$
$\displaystyle:=\frac{1}{2}\mathbb{E}_{\mathopen{}\left\\{x_{\star}[t]\right\\}\mathclose{}\sim\mathcal{P}^{(H+1)}}\mathopen{}\left[\frac{1}{T}\sum_{t=0}^{T-1}\left\|\mathopen{}\left(\hat{K}-K^{(H+1)}\right)\mathclose{}x_{\star}[t]\right\|^{2}\right]\mathclose{}$
$\displaystyle=\frac{1}{2T}\sum_{t=0}^{T-1}\mathbb{E}\mathopen{}\left[\operatorname{Tr}\mathopen{}\left((\hat{K}-K^{(H+1)})x[t]x[t]^{\top}(\hat{K}-K^{(H+1)})^{\top}\right)\mathclose{}\right]\mathclose{}$
$\displaystyle=\frac{1}{2}\operatorname{Tr}\mathopen{}\left((\hat{K}-K^{(H+1)})\Sigma_{x}^{(H+1)}(\hat{K}-K^{(H+1)})^{\top}\right)\mathclose{}$
$\displaystyle\leq f(N_{1},T,H,N_{2},\delta)\quad\text{w.p. }\geq 1-\delta$

directly implies a generalization bound on the Frobenius norm deviation
between the learned and expert controllers:

$\displaystyle\frac{1}{2}\operatorname{Tr}\mathopen{}\left((\hat{K}-K^{(H+1)})\Sigma_{x}^{(H+1)}(\hat{K}-K^{(H+1)})^{\top}\right)\mathclose{}\geq\frac{1}{2}\lambda_{\min}(\Sigma_{x}^{(H+1)})\left\|\hat{K}^{(H+1)}-K^{(H+1)}\right\|_{F}^{2}$
$\displaystyle\implies$
$\displaystyle\left\|\hat{K}^{(H+1)}-K^{(H+1)}\right\|_{F}^{2}\leq\frac{2}{\lambda_{\min}(\Sigma_{x}^{(H+1)})}f(N_{1},T,H,N_{2},\delta)\text{
w.p. }\geq 1-\delta.$

Since the Frobenius norm upper bounds the spectral norm, it suffices to have
enough samples $N_{1},N_{2},T,H$ such that

$f(N_{1},T,H,N_{2},\delta)\leq\frac{\lambda_{\min}(\Sigma_{x}^{(H+1)})}{8\mathcal{J}(A+BK^{(H+1)})^{2}\left\|B\right\|^{2}}$

If the sample complexity requirement is satisfied, then we have with
probability greater than $1-\delta$ that the spectral norm gap is satisfied

$\displaystyle\left\|\hat{K}-K^{(H+1)}\right\|\leq\frac{1}{4\mathcal{J}(A+BK^{(H+1)})^{2}\left\|B\right\|^{2}},$

which in turn implies the bound on the tracking error

$\displaystyle\max_{1\leq t\leq
T}\left\|\hat{x}[t]-x_{\star}[t]\right\|^{2}\leq\max_{0\leq t\leq
T-1}4\mathcal{J}(A+BK^{(H+1)})^{2}\left\|B\right\|^{2}\left\|\mathopen{}\left(\hat{K}-K^{(H+1)}\right)\mathclose{}x_{\star}[t]\right\|^{2}.$
(30)

In order to convert the RHS of the above expression, which involves a maximum
over time of the imitation error between $\hat{K}$ and $K^{(H+1)}$, into
something involving the excess risk of the controller $\hat{K}$, which is the
expected value of the imitation error, we again appeal to the Hanson-Wright
inequality, adapted specifically for Gaussian quadratic forms in Proposition
A.1.

We recall from (30) the tracking error is with probability greater than
$1-\delta$ bounded by

$\max_{1\leq t\leq T}\left\|\hat{x}[t]-x_{\star}[t]\right\|^{2}\leq\max_{0\leq
t\leq
T-1}4\mathcal{J}\mathopen{}\left(A+BK^{(H+1)}\right)\mathclose{}^{2}\left\|B\right\|^{2}\left\|\mathopen{}\left(\hat{K}-K^{(H+1)}\right)\mathclose{}x_{\star}[t]\right\|^{2}.$

In order to bound the RHS of the above inequality by a multiplicative factor
of the excess risk $\mathrm{ER}(\hat{\Phi},\hat{F}^{(H+1)})$, we apply
Proposition A.1, setting

$\displaystyle
R:=\mathopen{}\left(\hat{K}-K^{(H+1)}\right)\mathclose{}\mathopen{}\left(\Sigma_{x}^{(H+1)}\right)\mathclose{}^{1/2},\quad
z_{t}:=\mathopen{}\left(\Sigma_{x}^{(H+1)}\right)\mathclose{}^{-1/2}x_{\star}[t].$
(31)

By an application of the union bound, we have

$\displaystyle\mathbb{P}\mathopen{}\left[\max_{t\leq
T-1}z_{t}^{\top}R^{\top}Rz_{t}\geq(1+\varepsilon)\left\|R\right\|_{F}^{2}\right]\mathclose{}$
$\displaystyle\leq\sum_{t=0}^{T-1}\mathbb{P}\mathopen{}\left[z_{t}^{\top}R^{\top}Rz_{t}\geq(1+\varepsilon)\left\|R\right\|_{F}^{2}\right]\mathclose{}$
$\displaystyle\leq
T\exp\mathopen{}\left(-\frac{1}{4}\min\mathopen{}\left\\{\frac{\varepsilon^{2}}{4},\varepsilon\right\\}\mathclose{}\frac{\left\|R\right\|_{F}^{2}}{\left\|R\right\|^{2}}\right)\mathclose{}.$

Setting the last line less than the desired failure probability
$\delta^{\prime}\in(0,1)$, we get

$\displaystyle\min\mathopen{}\left\\{\frac{\varepsilon^{2}}{4},\varepsilon\right\\}\mathclose{}$
$\displaystyle\geq
4\frac{\left\|R\right\|^{2}}{\left\|R\right\|_{F}^{2}}\log\mathopen{}\left(\frac{T}{\delta^{\prime}}\right)\mathclose{}.$

Since $\left\|R\right\|\leq\left\|R\right\|_{F}$, it suffices to choose

$\displaystyle\varepsilon$
$\displaystyle=4\log\mathopen{}\left(\frac{T}{\delta^{\prime}}\right)\mathclose{},$

such that with probability greater than $1-\delta^{\prime}$

$\displaystyle\max_{0\leq t\leq
T-1}\left\|\mathopen{}\left(\hat{K}-K^{(H+1)}\right)\mathclose{}x_{\star}[t]\right\|^{2}$
$\displaystyle\leq\mathopen{}\left(1+4\log\mathopen{}\left(\frac{T}{\delta^{\prime}}\right)\mathclose{}\right)\mathclose{}\left\|\mathopen{}\left(\hat{K}-K^{(H+1)}\right)\mathclose{}\mathopen{}\left(\Sigma_{x}^{(H+1)}\right)\mathclose{}^{1/2}\right\|_{F}^{2}$
$\displaystyle=\mathopen{}\left(1+4\log\mathopen{}\left(\frac{T}{\delta^{\prime}}\right)\mathclose{}\right)\mathclose{}\mathrm{ER}\mathopen{}\left(\hat{\Phi},\hat{F}^{(H+1)}\right)\mathclose{}.$

Plugging this back into the tracking error bound and union bounding over the
generalization bound event on
$\mathrm{ER}\mathopen{}\left(\hat{\phi},\hat{F}^{(H+1)}\right)\mathclose{}$ of
probability $1-\delta$ and the concentration event of probability
$1-\delta^{\prime}$ yields the desired result.

The high probability bound introduced in Theorem C.1 provides a very high
resolution control over the deviation between closed-loop learned and expert
states, where we control the maximum deviation of states over time by the
expected time-averaged excess risk of the learned controller, accruing only a
$\log(T)$ factor in the process. Furthermore, our high-probability bounds are
multiplicative with respect to only $\log(1/\delta)$ rather than the naive
Markov’s inequality factor $1/\delta$, which is immediately higher-resolution
than the corresponding in-expectation bound (for example derived from
Proposition C.2), on which only Markov’s inequality can be applied without
more detailed analysis. However, this being said, in much of existing
literature in control and reinforcement learning, the performance of a policy
is evaluated as an expected cost/reward over the trajectories it generates,
for example LQG cost (Fazel et al., 2018), or expected reward in online
imitation learning (Ross et al., 2011). This therefore motivates understanding
whether the tools we develop directly imply bounds in-expectation on general
cost functions evaluated on the learned trajectory distribution versus the
expert trajectory distribution. Intuitively, if we treat the tracking error
between two trajectories as a metric, if the cost function varies continuously
(e.g. Lipschitz) with respect to this metric, our generalization bounds should
imply some sort of distributional distance between the closed-loop learned and
expert trajectory distributions. This is made explicit in the following
result.

######  Theorem C.2 (Full version of Theorem 3.2, (10))

Let us denote stacked trajectory vectors
$\bm{x}_{1:T}=\begin{bmatrix}x[1]^{\top}&\cdots&x[T]^{\top}\end{bmatrix}^{\top}\in\mathbb{R}^{n_{x}T}$,
and denote $\bm{x}^{\star}_{1:T}\sim\mathcal{P}^{\star}_{1:T}$ and
$\hat{\bm{x}}_{1:T}\sim\hat{}\mathcal{P}_{1:T}$ as the distributions of
closed-loop expert and learned trajectories generated by $K^{(H+1)}$ and
$\hat{K}$, respectively. Let $h$ be any cost function that is $L$-Lipschitz
with respect to the metric between trajectories
$d\mathopen{}\left(\bm{x}_{1:T},\bm{y}_{1:T}\right)\mathclose{}=\max_{1\leq
t\leq T}\left\|x[t]-y[t]\right\|$, i.e.,

$\displaystyle{\left|h(x)-h(y)\right|}\leq Ld(x,y),\;\forall
x,y\in\mathcal{X}$

Assume the sample requirements of Theorem C.1 are satisfied for given
$\delta$. Then with probability greater than $1-\delta$, the following bound
holds on the gap between the expected costs of expert and learned
trajectories:

$\displaystyle{\left|\mathbb{E}_{\hat{}\mathcal{P}_{1:T}}\mathopen{}\left[h(\hat{\bm{x}}_{1:T})\right]\mathclose{}-\mathbb{E}_{\mathcal{P}^{\star}_{1:T}}\mathopen{}\left[h(\bm{x}^{\star}_{1:T})\right]\mathclose{}\right|}$
$\displaystyle\leq
2\sqrt{3}L\mathcal{J}(A_{\mathsf{cl}})\left\|B\right\|\sqrt{1+\log(T)}\sqrt{\mathrm{ER}(\hat{\Phi},\hat{F}^{(H+1)})}$

where $A_{\mathsf{cl}}:=A+BK^{(H+1)}$.

###### Remark C.1

We note that since
$\frac{1}{T}\sum_{t=1}^{T}\left\|x[t]-y[t]\right\|\leq\max_{1\leq t\leq
T}\left\|x[t]-y[t]\right\|$, the above result holds with minimal modification
for
$d(\bm{x}_{1:T},\bm{y}_{1:T})=\frac{1}{T}\sum_{t=1}^{T}\left\|x[t]-y[t]\right\|$.

Before proceeding with the proof of Theorem C.2, the following are valid
examples of $h(\cdot)$:

  * •

$h(\bm{x}_{1:T})=\max_{1\leq t\leq T}\left\|Q^{1/2}x[t]\right\|$,

  * •

$h(\bm{x}_{1:T})=\max_{1\leq t\leq
T}\left\|x[t]-x_{\mathrm{goal}}[t]\right\|+\lambda\left\|Rx[t]\right\|$.

Proof  We appeal to an application of Kantorovich-Rubinstein duality to
establish general in-expectation guarantees. Let us define the metric

$\displaystyle d\mathopen{}\left(\bm{x}_{1:T},\bm{y}_{1:T}\right)\mathclose{}$
$\displaystyle=\max_{1\leq t\leq T}\left\|x[t]-y[t]\right\|.$

Define $\Gamma(\mathcal{P}^{\star}_{1:T},\hat{}\mathcal{P}_{1:T})$ as the set
of all couplings between the distributions $\mathcal{P}^{\star}_{1:T}$ and
$\hat{}\mathcal{P}_{1:T}$. In particular, trajectories
$\bm{x}^{\star}_{1:T},\hat{\bm{x}}_{1:T}$ following the same instances of
$w[t]$ and $z[t]$ considered in (26) is one such coupling. The Wasserstein
$1$-distance between $\mathcal{P}^{\star}_{1:T}$ and $\hat{}\mathcal{P}_{1:T}$
is defined as

$\displaystyle\mathcal{W}_{1}\mathopen{}\left(\mathcal{P}^{\star}_{1:T},\hat{}\mathcal{P}_{1:T}\right)\mathclose{}$
$\displaystyle=\inf_{\gamma\in\Gamma(\mathcal{P}^{\star}_{1:T},\hat{}\mathcal{P}_{1:T})}\mathbb{E}_{\mathopen{}\left(\bm{x}^{\star}_{1:T},\hat{\bm{x}}_{1:T}\right)\mathclose{}\sim\gamma}\mathopen{}\left[d\mathopen{}\left(\bm{x}^{\star}_{1:T},\hat{\bm{x}}_{1:T}\right)\mathclose{}\right]\mathclose{}$
$\displaystyle=\inf_{\gamma\in\Gamma(\mathcal{P}^{\star}_{1:T},\hat{}\mathcal{P}_{1:T})}\mathbb{E}_{\mathopen{}\left(\bm{x}^{\star}_{1:T},\hat{\bm{x}}_{1:T}\right)\mathclose{}\sim\gamma}\mathopen{}\left[\max_{1\leq
t\leq T}\left\|x_{\star}[t]-\hat{x}[t]\right\|\right]\mathclose{}.$

In particular, defining the coupling (26) as $\overline{\gamma}$, we
immediately have

$\displaystyle\mathcal{W}_{1}\mathopen{}\left(\mathcal{P}^{\star}_{1:T},\hat{}\mathcal{P}_{1:T}\right)\mathclose{}$
$\displaystyle\leq\mathbb{E}_{\mathopen{}\left(\bm{x}^{\star}_{1:T},\hat{\bm{x}}_{1:T}\right)\mathclose{}\sim\overline{\gamma}}\mathopen{}\left[\max_{1\leq
t\leq T}\left\|x_{\star}[t]-\hat{x}[t]\right\|\right]\mathclose{}.$

Kantorovich-Rubinstein duality (cf. Villani (2021)) provides the following
dual characterization of the Wasserstein distance:

$\displaystyle\frac{1}{L}\sup_{\left\|h\right\|_{\mathrm{Lip}}\leq
L}\mathbb{E}_{\hat{}\mathcal{P}_{1:T}}\mathopen{}\left[h(\hat{\bm{x}}_{1:T})\right]\mathclose{}-\mathbb{E}_{\mathcal{P}^{\star}_{1:T}}\mathopen{}\left[h(\bm{x}^{\star}_{1:T})\right]\mathclose{}$
$\displaystyle=\mathcal{W}_{1}\mathopen{}\left(\mathcal{P}^{\star}_{1:T},\hat{}\mathcal{P}_{1:T}\right)\mathclose{},$

where

$\left\|h\right\|_{\mathrm{Lip}}=\sup_{x,y\in\mathcal{X}}\frac{{\left|h(x)-h(y)\right|}}{d(x,y)}.$

Therefore, given any cost function $h$ that is $L$-Lipschitz continuous with
respect to $d(x,y)$, the gap between the expected closed-loop costs of the
learned and expert controllers, we can chain inequalities to yield

$\displaystyle\mathbb{E}_{\hat{}\mathcal{P}_{1:T}}\mathopen{}\left[h(\hat{\bm{x}}_{1:T})\right]\mathclose{}-\mathbb{E}_{\mathcal{P}^{\star}_{1:T}}\mathopen{}\left[h(\bm{x}^{\star}_{1:T})\right]\mathclose{}\leq
L\mathbb{E}_{\mathopen{}\left(\bm{x}^{\star}_{1:T},\hat{\bm{x}}_{1:T}\right)\mathclose{}\sim\overline{\gamma}}\mathopen{}\left[\max_{1\leq
t\leq T}\left\|x_{\star}[t]-\hat{x}[t]\right\|\right]\mathclose{}.$

It remains to provide a bound on
$\mathbb{E}_{\mathopen{}\left(\bm{x}^{\star}_{1:T},\hat{\bm{x}}_{1:T}\right)\mathclose{}\sim\overline{\gamma}}\mathopen{}\left[\max_{1\leq
t\leq T}\left\|x_{\star}[t]-\hat{x}[t]\right\|\right]\mathclose{}$. Recalling
the bound (30), this reduces to bounding the imitation error along expert
trajectories $x_{\star}[t]$. In order to convert a bound on $\max_{1\leq t\leq
T}\left\|x_{\star}[t]-\hat{x}[t]\right\|$ to $\max_{1\leq t\leq
T}\left\|x_{\star}[t]-\hat{x}[t]\right\|$, we apply Jensen’s inequality to
yield

$\displaystyle\mathbb{E}\mathopen{}\left[\max_{1\leq t\leq
T}\left\|x_{\star}[t]-\hat{x}[t]\right\|\right]\mathclose{}^{2}\leq\mathbb{E}\mathopen{}\left[\max_{1\leq
t\leq T}\left\|x_{\star}[t]-\hat{x}[t]\right\|^{2}\right]\mathclose{}$

The rest follows from an application of the maximal inequality.

###### Proposition C.2

Given a test controller $K$, we have the following maximal-type inequality
between the excess maximal risk and the excess average risk over the expert
target task data:

$\displaystyle\mathbb{E}\mathopen{}\left[\max_{0\leq t\leq
T-1}\left\|(K-K^{(H+1)})x_{\star}[t]\right\|^{2}\right]\mathclose{}$
$\displaystyle\leq
3(1+\log(T))\mathbb{E}\mathopen{}\left[\frac{1}{T}\sum_{t=0}^{T-1}\left\|(K-K^{(H+1)})x_{\star}[t]\right\|^{2}\right]\mathclose{}.$

Proof  We recall that the population distribution of $x_{\star}[t]$ is
$\mathcal{N}(0,\Sigma_{x}^{(H+1)})$. Recalling the definitions in (31), we can
re-write

$\displaystyle\left\|(K-K^{(H+1)})x_{\star}[t]\right\|^{2}$
$\displaystyle=z_{t}^{\top}R^{\top}Rz_{t},$

where $z_{t}\sim\mathcal{N}(0,I)$. Therefore, what we need to show is a
maximal inequality on

$\displaystyle\mathbb{E}\mathopen{}\left[\max_{0\leq t\leq
T-1}z_{t}^{\top}R^{\top}Rz_{t}\right]\mathclose{}.$

Toward establishing such an inequality, we prove the following bound on the
moment-generating function of $z_{t}^{\top}R^{\top}Rz_{t}$.

###### Lemma C.2

Let $z\sim\mathcal{N}(0,I_{n})$ be a standard Gaussian random vector, and
$R\in\mathbb{R}^{d\times n}$ is an arbitrary fixed matrix. Then, we have the
following bound on the moment-generating function of $z^{\top}R^{\top}Rz$

$\displaystyle\mathbb{E}\mathopen{}\left[\exp\mathopen{}\left(\lambda
z^{\top}R^{\top}Rz\right)\mathclose{}\right]\mathclose{}$
$\displaystyle\leq\exp\mathopen{}\left(3\lambda\left\|R\right\|_{F}^{2}\right)\mathclose{},\quad\text{for
}{\left|\lambda\right|}\leq\frac{1}{3\left\|R\right\|^{2}}.$

Proof  The proof largely follows standard results, for example from Vershynin
(2018, Chapter 6.2), but we derive the result from scratch to preserve
explicit numerical constants. By the rotational invariance of Gaussian random
vectors, we may write

$\displaystyle\mathbb{E}\mathopen{}\left[\exp\mathopen{}\left(\lambda
z^{\top}R^{\top}Rz\right)\mathclose{}\right]\mathclose{}$
$\displaystyle=\mathbb{E}\mathopen{}\left[\exp\mathopen{}\left(\sum_{i=1}^{\min\mathopen{}\left\\{n,d\right\\}\mathclose{}}\lambda\sigma_{i}^{2}g_{i}^{2}\right)\mathclose{}\right]\mathclose{}$
$\displaystyle=\prod_{i=1}^{\min\mathopen{}\left\\{n,d\right\\}\mathclose{}}\mathbb{E}\mathopen{}\left[\exp\mathopen{}\left(\lambda\sigma_{i}^{2}g_{i}^{2}\right)\mathclose{}\right]\mathclose{},$

where $\sigma_{i}$ is the $i$th singular value of $R$ and
$g_{i}\sim\mathcal{N}(0,1)$ are i.i.d. standard Gaussian random variables.
Thus we have reduced the problem to bounding the MGF of a $\chi^{2}(1)$ random
variable. We recall that the MGF of a $\chi^{2}(1)$ random variable is given
by

$\mathbb{E}\mathopen{}\left[\exp\mathopen{}\left(tg_{i}^{2}\right)\mathclose{}\right]\mathclose{}=\frac{1}{\sqrt{1-2t}},\quad
t<\frac{1}{2}.$

We claim that there exists constants $C,c>0$ such that

$\mathbb{E}\mathopen{}\left[\exp\mathopen{}\left(tg_{i}^{2}\right)\mathclose{}\right]\mathclose{}\leq\exp\mathopen{}\left(Ct\right)\mathclose{},\quad
t\leq c<\frac{1}{2}.$

We can derive candidates for $C,c$ by solving

$\displaystyle\frac{1}{\sqrt{1-2t}}$
$\displaystyle\leq\exp\mathopen{}\left(Ct\right)\mathclose{}\;\land\;t<1/2$
$\displaystyle\iff-\frac{1}{2}\log(1-2t)$ $\displaystyle\leq
Ct\;\land\;t<1/2.$

We observe the function $f(t):=-\frac{1}{2}\log(1-2t)-Ct$ satisfies $f(0)=0$
and $f^{\prime}(t)=\frac{1}{1-2t}-C$ is monotonically increasing for $t<1/2$,
and $f^{\prime}(0)<0$ for $C>1$. Therefore, for a given $C>1$ it suffices to
find $0<c<1/2$ such that $f^{\prime}(c)=0$, which implies $f(c)\leq 0$.
Plugging in $C=3$, $c=1/3$ satisfies this. Therefore, we get the MGF bound

$\mathbb{E}\mathopen{}\left[\exp\mathopen{}\left(tg_{i}^{2}\right)\mathclose{}\right]\mathclose{}\leq\exp\mathopen{}\left(3t\right)\mathclose{},\quad
t\leq\frac{1}{3}.$

Applying this back on the MGF of $z^{\top}R^{\top}Rz$, we have

$\displaystyle\mathbb{E}\mathopen{}\left[\exp\mathopen{}\left(\lambda
z^{\top}R^{\top}Rz\right)\mathclose{}\right]\mathclose{}$