Year: 2019

Abstract:
A fundamental challenge in artificially intelligence is to build an agent that generalizes and adapts to unseen environments. A common strategy is to build a decoder that takes a context of the unseen new environment and generates a policy. The current paper studies how to build a decoder for the fundamental continuous control environment, linear quadratic regulator (LQR), which can model a wide range of real world physical environments. We present a simple algorithm for this problem, which uses upper confidence bound (UCB) to refine the estimate of the decoder and balance the exploration-exploitation trade-off. Theoretically, our algorithm enjoys a Õ (√ T ) regret bound in the online setting where T is the number of environments the agent played. This also implies after playing Õ ( 1/ 2 ) environments, the agent is able to transfer the learned knowledge to obtain an -suboptimal policy for an unseen environment. To our knowledge, this is first provably efficient algorithm to build a decoder in the continuous control setting. While our main focus is theoretical, we also present experiments that demonstrate the effectiveness of our algorithm.

text: (√ T ) regret bound in the online setting where T is the num-
ber of environments the agent played. This also implies after playing Õ ( 1/ 2 ) environments, the agent is able to transfer the learned knowledge to obtain an -suboptimal policy for an unseen environment. To our knowledge, this is first provably efficient algorithm to build a decoder in the continuous control setting. While our main focus is theoretical, we also present experiments that demonstrate the effectiveness of our algorithm.

heading: 1 INTRODUCTION
text: Humans are able to solve a new task without any training based on previous experience in similar tasks. Our intelligent agent should be able do the same, learning from previous experience, adapting to the new ones and improving the performance as the agent gains more experience. This is a challenging problem as we need to design an adaptation mechanism which is fundamentally different from classical supervised learning methods.
A common approach is to build a decoder so that once the agent sees a description of new task, i.e., the context of the new task, the decoder turns the context into a succinct representation of the new task, based on which the agent is able to design a policy to solve the task. Note this procedure resembles how a human solves a new task. For example, if a human wants to push an object on a table, the human first sees the object and the table (context). Then, in his/her mind, the context becomes a representation of this task, e.g., a sense of weight of the object. Based on this representation, the human can easily reason about how much force to exert on the object.
This general approach has been applied in practice. For example, Wu et al. (2018) studied the visual navigation task and built a Bayesian model that takes the context of new environments and outputs the policy that enables the agent to navigate. Killian et al. (2016) used this approach to develop personalized medicine policies for HIV treatment.
While this is a promising approach, currently we only have limited theoretical understanding. The approach can be formulated in Contextual Markov Decision Process (CMDP) framework (Hallak et al., 2015). Recently, there is a line of work gave provable guarantees for CMDP (Abbasi-Yadkori & Neu, 2014; Hallak et al., 2015; Dann et al., 2018; Modi et al., 2018; Modi & Tewari, 2019). These work all studied tabular MDPs, and use function approximation, e.g., linear functions, generalized linear models, etc, to model the mapping from the context to the probability transition matrix. A major drawback of these work is that they are restricted to the tabular setting and thus can only deal with discrete environments. Therefore, they can hardly model real-world continuous control tasks, like the task of pushsing an object as we described above. A natural question arises:
Can we design a provably efficient decoder for continuous control problems?
In this paper, we make an important step towards answering this question. We study the fundamental setting in continuous control, linear quadratic regulator (LQR). LQR is arguably the most widely used framework in continuous control, as LQR easily models real world physical phenomena, e.g., the pushing object task we described earlier. We propose a new algorithm that builds a decoder, so that for a new LQR task, the decoder takes LQR’s context and outputs a representation based on which the agent can easily infer a near-optimal policy for new continuous control tasks. In the training phase, we build the decoder via a sequence of LQRs (in an online fashion) with unknown parameters. For each new task, we first use the current decoder to build the representation of this task, infer a policy based on this representation and use this policy to do control for this episode. There are two crucial components in our algorithm. First, after each episode, we will refine the estimate of the decoder based on the observations from this episode. Second, it is crucial to use a upper confidence bound (UCB) estimator of the decoder to build the representation so that the agent can perform a near-optimal trade-off between exploration and exploitation. In this way, we provably show the decoder improves the performance as it experiences more training tasks. Formally, we show our algorithm enjoys Õ (√ T )
regret (the difference between the cumulative rewards of our algorithm and the unknown optimal policy on every seen environment) bound in the online setting. Moreover, the algorithm is able to obtain an -suboptimal policy for an unseen LQR environment after playing Õ ( −2 )
environments. To our knowledge, this is the first provably efficient algorithm that builds a decoder for continuous control environments. Empirically, we simulate several physical environments to illustrate the effectiveness of our algorithm.
Organization This paper is organized as follows. In Section 2, we discuss related work. In Section 3, we formally describe the problem setup. In Section 4, we present our algorithm and its theoretical guarantees. In Section 5, we use simulation on physical environments to demonstrate the effectiveness of our approach. We conclude in Section 6 and defer most technical proofs to the appendix.

heading: 2 RELATED WORK
text: Recently there is a large body of literature focusing on learning for control in LQR systems. The first work we are aware of is Fiechter (1997) which studies the sample complexity of LQR in the offline setting. For the online setting, where the agent can only obtain the next state starting from the present state, the first near-optimal regret bound (Õ( √ T )) is due to Abbasi-Yadkori & Szepesvári (2011), which studies the learning problem in the infinite-horizon average-case cost setting. Later on, a sequence of papers (Tu & Recht, 2017; Dean et al., 2017; 2018; Tu & Recht, 2018; AbbasiYadkori et al., 2018; Cohen et al., 2019) studied this problem in similar settings, improved efficiency of the algorithms and characterized the gap between model-free and model-based approaches.
Building an agent that quickly adapts to new environment has received increasing interest in the machine learning community. Taylor & Stone (2009) gave a summary for the literature status before 2009. More recently, a sequence of theory papers Lehnert & Littman (2018); Spector & Belongie (2018); Abel et al. (2018); Lehnert et al. (2019) studied the transferability of reward knowledge, state-abstraction, and model features for Markov decision processes. Please also refer to references in paper cited above for more details. There are also some experimental works, e.g., Santara et al. (2019); Yu et al. (2018); Wu et al. (2018); Gamrian & Goldberg (2018), studying how to transfer knowledge from seen tasks to unseen tasks. Nevertheless, we are not aware of any study on how to provably perform continuous control with contexts.

heading: 3 PRELIMINARIES
text: Notations. We begin by introducing necessary notations. We write [h] to denote the set {1, . . . , h}. We use Id ∈ Rd×d to denote the d-dimensional identity matrix. We use 0d×d′ to represent the allzero matrix in Rd×d′ . If it is clear from the context, we omit the subscript d × d′. Let ‖·‖2 denote the Euclidean norm of a vector in Rd. For a symmetric matrix A, let ‖A‖op denote its operator
norm and λi (A) denote its i-th eigenvalue. Throughout the paper, all sets are multisets, i.e., a single element can appear multiple times.
Finite Horizon Linear Quadratic Regulator. We now formally define the finite horizon Linear Quadratic Regulator (LQR) problem. In the LQR problem, there is a state space X ⊂ Rd and a closed action space U ⊂ Rd′ . Suppose we always start from the initial state x1 = xinit ∈ X and play for H steps. Then at a state xh ∈ X , if an action uh ∈ U is played, the next state is given by
xh+1 = Axh +Buh + wh+1, (1)
where A,B are matrices of proper dimension and wh+1 is a zero-mean random vector. Here A,B can be viewed as the succinct representation of this LQR because as will be explained below, given A,B, we can easily infer the optimal policy for this LQR. For simplicity, we additionally denote
M = [A,B], and yh = [x>h , u > h ] > ∈ Rd+d′ .
Now the state transition can be rewritten as xh+1 = Myh + wh+1. For the ease of presentation, we assume that the covariance matrix of noise vectorwh+1 is E(wh+1w>h+1) = Id. Our analysis follows similarly if the covariance matrix is not Id (see e.g. Remark 3 of Abbasi-Yadkori & Szepesvári (2011)). After each step, the player receives an immediate cost x>hQhxh +u > hRhuh, where Qh, Rh are positive definite (PD) matrices of proper dimensions. Throughout the paper, we assume the agent knows Qh and Rh for all h. At a terminal state xH , there is no action to be played, and the player receives a terminal cost x>HQHxH , where QH is a PD matrix of proper dimension. The goal of the player is to find a policy π : (X ×U)∗×X → U , which is a function that maps the trajectory {(xi, ui)}h−1i=1 ∪ {xh} to the next action uh, such that the following objectives are minimized:{
Jπh (M,x) := E [( H−1∑
h′=h
x>hQhxh + u > hRhuh ) + x>HQfxH ∣∣∣∣ xh = x]} h∈[H] ,
where the action uh is given by uh = π[(x1, u1), (x2, u2), . . . , (xh−1, uh−1), xh], and the expectation is over the randomness of wh and π.
It is well-known that the optimal policy π∗ is Markovian Puterman (2014), i.e., it only depends on the present state. For an unconstrained action space U , we have
∀x ∈ X , h ∈ [H − 1] : π∗h(M,x) := Kh(M)x where M = [A,B] and Kh(M) is a matrix that will be defined shortly. It is also known (see e.g. Bertsekas (1996)) that the optimal cost function J∗h(x) := J π∗ h (x) is given by
J∗h(M,x) := x >Ph(M)x+ Ch(M) = inf π Jπh (M,x) (2)
where
Ph(M) =
{ Qh +A
>Ph+1(M)A−A>Ph+1B(Rh +B>Ph+1(M)B)−1B>Ph+1(M)A h < H QH h = H
(3)
and
Ch(M) =
{ Ch+1(M) + Ewh+1 [ w>h+1Ph+1(M)wh+1] h < H
0 h = H .
We now define Kh(M) as
Kh(M) := −(Rh +B>Ph+1(M)B)−1B>Ph+1(M)A. (4) Note that the optimal value Equation (2) satisfies Bellman equations,
∀h ∈ [H − 1] : J∗h(M,x) = x>Qhx+ π∗(x)>Rhπ∗(x) + E [ J∗h+1(Ax+Bπ ∗(x) + w) ]
and
∀h ∈ [H − 1] : J∗h(M,x) = x>Qhx+ min u E[u>Rhu+ J∗h+1(Ax+Bu+ w)].
Now we have shown that if we are given A and B, then we can obtain the optimal policy directly. In this paper, we deal with setting where A and B are unknown and we need to use decoder to decode A and B from the contexts of the current LQR, as specified below.
Learning to Control LQR with Contexts In the continuous control with contexts setting, in each episode we observe a context (C,D) ∼ µ, where µ is a distribution on Rp×d × Rp′×d′ . The context [C,D] encodes the information of the environment. Formally, the representation ([A,B]) of this environment can be decoded from the context via a decoding matrix Θ∗ ∈ Rd×(p+p ′):
[A,B] = Θ∗ · [
C 0p×d′ 0p′×d D
] . (5)
From now on, to emphasize that the representation of LQR can be decoded from Θ∗, we write MΘ∗,C,D := Θ∗ · [
C 0p×d′ 0p′×d D
] = [A,B]. (6)
If it is clear from the context, we ignore [C,D] for notational simplicity. Note the optimal decoder Θ∗ is unknown to the agent and the goal is to learn Θ∗ from contexts and interactions with the environment. Below we formally define the problem that we study. Definition 3.1 (Contextual Transfer Learning Problem). Build an agent that plays onK LQR games (one trajectory per game) with context pairs {(C(1), D(1)), (C(2), D(2)), . . . , (C(K), D(K))} ∼ µ, for some integer K ≥ 0 such that for another new context pair (C,D) ∼ µ, the agent outputs a policy π based on (C,D) which satisfies
E[Jπh (MΘ∗,C,D, x1)− J∗h(MΘ∗,C,D, x1)] ≤ for some given target accuracy > 0.
Here K is the sample complexity which ideally scales polynomially with and problem-dependent parameters. The performance of the agent can also be measured by regret, as defined below.
Regret(KH) := K∑ k=1 J π̃ (k) 1 ( MΘ∗,C(k),D(k) , x1 ) − J∗1 ( MΘ∗,C(k),D(k) , x1 ) , (7)
where π̃(k) is the policy played at episode k by the agent. This quantity measurse the sub-optimality of policies the agent played in the first K episodes. Remark 3.1. We consider matrix-type linear maps from context to the representation only for sake of presentation. Our algorithm and analysis can be readily extended to other linear maps, e.g., [A∗(C), B∗(D)] := f(C,D) for some unknown linear function f .

heading: 4 MAIN ALGORITHM
text: In this section, we first describe the algorithm and then present its sample complexity guarantees.
Since the decoder is linear, a straightforward algorithm is first to estimate (A,B) using the trajectory from the single episode with system identification techniques (Mehra, 1974), and then to use linear regression to estimate the mapping from (C,D) to (A,B). However, this naive algorithm has two drawbacks. First, estimating (A,B) accurately requires a long horizon. However, in our setup, we do not have any restrictions on the horizon. To fix this, we stack contexts and observations to construct a more direct estimate on the decoder (cf. Equation equation 8). Second, in order to achieve√ T -type regret guarantee, one needs to balance exploration and exploitation carefully, but the naive algorithm does not have such an component. Our algorithm uses UCB to construct a confidence set which helps balance exploration and exploitation.
Algorithm We describe the high-level idea of the algorithm below. The agent maintains a decoder that maps the context (C,D) to the representation (A,B). We denote Θ(k) the decoder at the k-th episode. Initially, we know nothing about Θ∗, so we initialize our decoder by setting Θ(1) = 0 ∈ Rd×p. At the k-th episode, the agent plays policy π(k) and in each time step h ∈ [H−1], it collects data
x (k) h , u (k) h , x (k) h+1, z (k) h ←
[ C(k)x
(k) h
D(k)u (k) h
] ,
Algorithm 1 Linear Continuous Control with Contexts 1: Input Total number of episodes K; 2: Initialize Θ(1) ← 0 ∈ Rd×2p, V (1) ← I2p,2p, W (1) ← 0 ∈ R2p×d; 3: for episode k = 1, 2, . . . ,K do 4: Let x(k)1 ← xinit, V (k+1) ← V (k), W (k+1) ←W (k); 5: Obtain context [C(k), D(k)] ∼ µ; 6: Solve for the present policy:
Θ̃(k) = arg min Θ∈C(k) J∗1
( MΘ,C(k),D(k) , x (k) 1 ) (10)
7: where J∗1 is given by Equation 2, and C(k) is defined in Equation 11; 8: for stage h = 1, 2, . . . , H − 1 do 9: Let the current state be x(k)h ;
10: Play action u(k)h ← Kh ( MΘ(k),C(k),D(k) ) · x(k)h , where Kh is defined in Equation 4; 11: Obtain the next state x(k)h+1;
12: Let z(k)h ← [ C(k)x (k) h
D(k)u (k) h
] ;
13: Update: V (k+1) ← V (k+1) + z(k)h z (k)> h ; 14: Update: W (k+1) ←W (k+1) + z(k)h ( x (k) h+1 )> ;
15: Compute Θ(k+1)> ← ( V (k+1) )−1 W (k+1);
16: output Θ̃(k) where k is chosen from [K] uniformly at random.
where z(k)h can be viewed as the context regularized observation. We now describe how to obtain policy π(k). We first solve the following optimization problem,
Θ̃(k) = arg min Θ∈C(k) J∗1
( MΘ,C(k),D(k) , x (k) 1 ) where J∗1 is given by Equation (2), and the confidence set C(k) will be defined shortly. C(k) represents our confidence region on Θ∗. Since we choose the one that minimizes the cost, this represents the principle “optimism in the face of uncertainty” and it is the key to balance exploration and exploitation which will be more clear in the proof. Notice that the above optimization problem is a polynomial optimization problem. Then the policy is given by
π (k) h (x) := Kh
( M (k) ) · x where M (k) = MΘ(k),C(k),D(k) := Θ(k) · [ C(k) 0
0 D(k)
] ,
and Kh is given by Equation (4). After episode k ∈ [K], we use the following ridge regression formulation to update decoder
Θ(k) = (( V (k+1) )−1 W (k+1) )> . (8)
where
V (k+1) = I + k∑ k′=1 H−1∑ h=1 z (k′) h z (k′)> h and W (k+1) = k∑ k′=1 H−1∑ h=1 z (k′) h x (k′)> h+1 .
After playing K episodes, the algorithm outputs a Θ̃ by picking one from {Θ̃(k)}k∈[K] uniformly at random. Now for a new task with its context, our learned policy map is given by:
∀C,D ∼ µ, x ∈ X , h ∈ [H − 1] : π̃C,D,h(x) = Kh ( Θ̃ · [ C 0 0 D ]) · x. (9)
The formal algorithm is presented in Algorithm 1.

heading: 4.1 ALGORITHM ANALYSIS
text: To present the analysis of the algorithm, we first introduce some assumptions. Assumption 4.1. The contexts and LQR satisfy the following properties.
• ∀h ∈ [H], ‖Ph(M)‖2 ≤ cq for some parameter cq > 0. • ‖Θ∗‖F ≤ cΘ; • ∀h ∈ [2, H], i ∈ [d] : ‖wh‖2 <∞ and ∀γ > 0, E[γwh,i] ≤ exp(γ2c2w/2); • ∀x ∈ X , u ∈ U , (C,D) ∈ supp(µ) : ‖Cx‖2 + ‖Du‖22 ≤ c2x, ‖x‖2 + ‖u‖22 ≤ c2x; • ∀(C,D) ∈ supp(µ), x ∈ X , h ∈ [H]: Kh(MΘ∗,C,D) · x ∈ U . where cΘ, cw, cx are some positive parameters.
The first assumption is standard to ensure controllability. The second is a regularity condition on the optimal decoder Θ∗. The third assumption imposes almost sure boundedness of the noise w. The fourth assumption is a regularity condition on the observation. The last assumption guarantees the optimal controller for the unconstrained LQR problem is realizable in our control set U . Given these assumptions, We are now ready to define confidence set C(k) as follows.
C(k) = { Θ : tr [( Θ−Θ(k) ) V (k) ( Θ−Θ(k) )>] ≤ β(k), and ∀h ∈ [H], (C,D) ∈ supp(µ),
∥∥Ph(MΘ,C,D)∥∥2 ≤ cq}, (11) where Ph is given by Equation (3) and β(k) is defined as follows,
β(k) = ( cΘ + cw √ 2d ( log d+ p log(1 + kHc2x/p)/2 + log δ −1 ))2 . (12)
We remark that C(k) is changing at every episode because we update Θ(k) and V (k) at every episode. The size of C(k) is decreasing because V (k) is strictly increasing at every episode. With the above assumptions, the guarantee of Algorithm 1 is formally presented in the next theorem. Theorem 4.1. Suppose we run Algorithm 1 for
K ≥ c′H,cq,cx,cΘ,cw · dp2 · log 3(dKδ−1)
2
episodes, for some parameter c′H,cq,cx,cΘ,cw depending polynomially onH , cq, cx, cΘ, cw, Then with probability at least 1− δ, we have for π̃C,D be defined in Equation 9.
E [C,D]∼µ
[ E
π̃C,D
( J π̃C,D 1 ([Θ∗C,Θ∗D], x1) ) − J∗1 ([Θ∗C,Θ∗D], x1) ] ≤ . (13)
Theorem 13 states after playing polynomial number of episodes, our agent can learn a decoder Θ̃ such that given a new LQR with contexts (C,D), this decoder can turns the contexts into a nearoptimal policy π̃C,D without any training on the new LQR. Note this is the desired agent we want to build as described in the introduction. We emphasize again that this is the first provably efficient algorithm that builds a decoder for continuous control environments.
Remark 4.1. Via similar analysis, it is easy to show that if the output Θ̃ is picked uniformly at random from {Θ(k)}k∈[K], the policy achieves similar accuracy.
In fact, Theorem 4.1 is implies by the following regret bound of our algorithm. Proposition 4.1. With probability at least 1− δ,
Regret(KH) ≤ c′H · d1/2p · log3/2(dKHcxδ−1) · √ KH.
where c′H is a constant depending only polynomially on H , cq, cx, cM , cw.
By the definition of regret, this proposition justifies that the performance of the agent actually improves as it sees more environment.

heading: 5 EXPERIMENTS
text: In this section, we validate the effectiveness of our algorithm via numerical simulations.
We perform experiments on a path-following task. In this task, we are given a trajectory z∗1 , z ∗ 2 , . . . , z ∗ H ∈ R2. Our goal is to exert forces u1, u2, . . . , um ∈ R2 on objects with different
(measurable) masses to minimize the total squared distance ∑H h=1 ‖zh − z∗h‖2 + ‖ui‖22. Each state xh = [zh; vh] ∈ R4 is a vector whose first two dimensions represent the current position and the last two dimension represent the current velocity. In each stage h, we may exert a force uh ∈ R2 on the object, which produces an accelerations uhm ∈ R2. The dynamics of the system can be described as{
zh+1 = zh + vh vh+1 = k · vh + uh/m
(14)
where 0 < k ≤ 1 is the decay rate of velocity induced by resistance. In our setting, the decay rate of velocity k is fixed (encoded in Θ∗), where the mass of the object m is drawn from the uniform distribution over [0.1, 10]. In our experiments, we set the noise vectorwh in the dynamics of the LQR system (cf. Equation 1) to be a Gaussian random vector with zero mean and covariance 10−4 · I . In each episode, we receive an object with mass m where m is draw from the uniform distribution over [0.1, 10], train one trajectory using that object, and the goal is to recover the physical law described in Equation 14 so that our model can deal with objects with unseen mass m. Please see Appendix B for the concrete value of Θ∗, Q and R and the distribution of C and D.
In our experiments, we use 100 different masses as training masses (fixed among all experiments), and use 100 different masses as test masses (again fixed among all experiments). All the training masses and test masses are drawn from the uniform distribution over [0.1, 10]. We implement a practical version of Algorithm 1. In particular, instead of solving the optimization problem in Equation 10 exactly, we sample 100 different Θ from C(k) uniformly at random, and choose the Θ which minimizes the objective function. Moreover, instead of using the theoretical bound for β(k) in Equation 12, we treat β(k) as a tunable parameter and set β(k) = 104 in our experiments to encourage exploration at early stage of the algorithm. We use two different metrics to measure the accuracy of the learned model. First, we use ‖Θk − Θ∗‖F where Θk is calculated in Line 15 to measure the accuracy of the learned Θ. Moreover, using the learned Θ, we test on 100 objects whose masses are the 100 test masses to calculate the control cost ∑H h=1 ‖zh− z∗h‖2 + ‖ui‖22. We compare the control cost of the learned Θ and the optimal control cost, and use the mean value of the differences (named mean control error) to measure the accuracy.
In all experiments we fix H = 20. We use three different types of trajectories: unit circle, parabola y = x2 with x ∈ [0, 1] and Lemniscate of Bernoulli with a = 11. For all three types of trajectories we use their parametric equation x = x(t) and y = y(t), divide the interval [0, 1] evenly into H parts, and set t to be the endpoints of these parts. We use these t values to define the trajectory z∗1 , z ∗ 2 , . . . , z ∗ H ∈ R2. We set the decay ratio k to be k = 1 or k = 0.7 in our experiments.
We plot the accuracy of the learned model in Figure 1. Here we vary the number of training episodes (the number of training masses) and observe its effect on the accuracy. It can be observed that our algorithm achieves an satisfactory accuracy using only 5 episodes. We also illustrate trajectories obtained by our resulting controllers in Figure 2. From Figure 2, it is clear that as the agent plays more environments, it can enjoy better performance.

heading: 6 CONCLUSION
text: In this paper, we give a provably efficient algorithm for learning LQR with contexts. Our result bridges two major fields, learning with contexts and continuous control from a theoreticallyprincipled view. For future work, it is interesting to study more complex settings, include non-linear control. Another interesting direction is to design provable algorithm in our setting with safety guarantees (Dann et al., 2018).
1https://en.wikipedia.org/wiki/Lemniscate_of_Bernoulli.

heading: A PROOF OF MAIN RESULTS
text: This sections devotes to proving the main results. Before we prove Proposition 4.1, let us use it to prove Theorem 4.1.
Proof of Theorem 4.1. We rewrite the Equation equation 13 as follows. EC,DEπ̃ [ J π̃1 ( MΘ∗,C,D, x1 )] − EC,D [ J∗1 ( MΘ∗,C,D, x1 )] = 1
K K∑ k=1 EC,D [ Jπ k 1 ( MΘ∗,C,D, x1 )] − EC,D [ J∗1 ( MΘ∗,C,D, x1 )] = 1
K K∑ k=1 ( EC,D [ Jπ k 1 ( MΘ∗,C,D, x1 )] − Jπk1 ( MΘ∗,C(k),D(k) , x1 ) + Jπ k 1 ( MΘ∗,C(k),D(k) , x1 ) − J∗1 ( MΘ∗,C(k),D(k) , x1 ) + J∗1 ( MΘ∗,C(k),D(k) , x1 ) − EC,D [ J∗1 ( MΘ∗,C,D, x1
)]) = R1 +R2 +R3
where
R1 = 1
K K∑ k=1 ( EC,D [ Jπ k 1 ( MΘ∗,C,D, x1 )] − Jπk1 ( MΘ∗,C(k),D(k) , x1 )) ,
R2 = 1
K K∑ k=1 ( J∗1 ( MΘ∗,C(k),D(k) , x1 ) − EC,D [ J∗1 ( MΘ∗,C,D, x1 )]) ,
and
R3 = 1
K K∑ k=1 ( Jπ k 1 ( MΘ∗,C(k),D(k) , x1 ) − J∗1 ( MΘ∗,C(k),D(k) , x1 )) .
LetFk be the filtration of fixing all randomness before episode k. We haveR1 andR2 are Martingale difference sum. Note that the magnitude of each summand inR1 orR2 is upper bounded by (proved in Lemma A.3 and A.4),
Hcqcx
almost surely. Therefore, by Azuma’s inequality (Theorem A.1), we have, with probability greater than 1− δ/2,
|R1|+ |R2| ≤ 2Hcqcx · √ 2 log(8/δ)
K .
Moreover, by Proposition 4.1, we have with probability greater than 1− δ/2,
|R3| ≤ c · d1/2p · log3/2(dKHc2xδ−1) · √ H
K ,
where c is constant depending only polynomially on H , cq, cx, cM , and cw. Combining the above two inequalities, and setting K appropriately, we complete the proof of Theorem 4.1.
A.1 USEFUL CONCENTRATION BOUNDS
Before we prove the main proposition, we first recall some useful concentration bounds.
Theorem A.1 (Azuma’s inequality). Assume that {Xs}s≥0 is a martingale and |Xs −Xs−1| ≤ cs almost surely. Then for all t > 0 and all > 0,
Pr [ |Xt −X0| ≥ ] ≤ 2 exp
( − 2
2 ∑t s=1 c 2 s
) .
Theorem A.2 (Martingale Concentration, Theorem 16 of Abbasi-Yadkori & Szepesvári (2011)). Let Ft; t ≥ 0 be a filtration, (zt; t ≥ 0) be an Rd-valued stochastic process adapted to (Ft). Let (ηt; t ≥ 1) be a real-valued martingale difference process adapted to Ft. Assume that ηt is conditionally sub-Gaussian with constant L, i.e.,
∀γ > 0 : E[γηt|Ft] ≤ exp(γ2L2/2). Consider the following martingale
St = t∑ τ=1 ητzt−1
and the matrix-valued processes
Vt = I + t∑ τ=0 zt−1z > t−1.
Then for any δ ∈ (0, 1), with probability at least 1− δ,
∀t ≥ 0, ‖St‖2V −1t ≤ 2L 2 log
( det(Vt) 1/2
δ ) where ‖St‖2V −1t := S > t V −1 t St.
A.2 PROOF OF PROPOSITION 4.1
In this section, we prove the main proposition. We first bound det(V (k)) for any k.
Lemma A.1. For all k ∈ [K],
det(V (k)) ≤ ( 1 + kHc2x/p )p .
Proof. Since V (k) is PD, we have,
det(V (k)) ≤ ( tr(V (k))/p )p ≤ ( 1 + k∑ k′=1 H−1∑ h=1 ∥∥z(k′)h ∥∥22/p)p. By Assumption 4.1, we have ‖z(k
′) h ‖22 ≤ c2x. This completes the proof.
Let us then define an event Ek as follows. Definition A.1 (Good Event). We define event Ek as {∀k′ ≤ k : Θ∗ ∈ C(k ′)}.
We then show that the event Ek happens with high probability.
Lemma A.2. For all k ∈ [K], we have Pr[Ek] ≥ 1− δ.
Proof. Now we consider Θ∗ −Θ(k). We immediately have
Θ>∗ −Θ(k)> = Θ>∗ − (V (k))−1 ( k∑ k′=1 H−1∑ h=1 z (k′) h (Θ∗z (k′) h + w (k′) h+1) > )
= ( I − (V (k))−1 k∑ k′=1 H−1∑ h=1 z (k′) h z (k′)> h ) Θ>∗ + (V (k))−1 k∑ k′=1 H−1∑ h=1 z (k′) h w (k′)> h+1 .
Next, we have
(Θ∗ −Θ(k))V (k)(Θ∗ −Θ(k))>
=Θ∗
( I − (V (k))−1 k∑ k′=1 H−1∑ h=1 z (k′) h z (k′)> h )> V (k) ( I − (V (k))−1 k∑ k′=1 H−1∑ h=1 z (k′) h z (k′)> h ) Θ>∗
+ Θ∗
( I − (V (k))−1 k∑ k′=1 H−1∑ h=1 z (k′) h z (k′)> h )> k∑ k′=1 H−1∑ h=1 z (k′) h w (k′)> h+1
+ k∑ k′=1 H−1∑ h=1 w (k′) h+1z (k′)> h ( I − (V (k))−1 k∑ k′=1 H−1∑ h=1 z (k′) h z (k′)> h ) Θ>∗
+ k∑ k′=1 H−1∑ h=1 w (k′) h+1z (k′)> h (V (k))−1 k∑ k′=1 H−1∑ h=1 z (k′) h w (k′)> h+1
Note that ∑k k′=1 ∑H−1 h=1 z (k′) h z (k′)> h = V (k) − I and thus (V (k))−1∑kk′=1∑H−1h=1 z(k′)h z(k′)>h = I − (V (k))−1. Hence we have tr [ (Θ∗ −Θ(k))V (k)(Θ∗ −Θ(k))>
] = ‖Θ∗‖2(V (k))−1 + 2tr ( Θ∗(V (k))−1 k∑
k′=1 H−1∑ h=1 z (k′) h w (k′)> h+1 ) + ∥∥∥ k∑ k′=1 H−1∑ h=1 z (k′) h w (k′)> h+1 ∥∥∥2 (V (k))−1
≤ ‖Θ∗‖2(V (k))−1 + 2 ∥∥∥Θ∗∥∥∥
(V (k))−1 ∥∥∥ k∑ k′=1 H−1∑ h=1 z (k′) h w (k′)> h+1 ∥∥∥ (V (k))−1 + ∥∥∥ k∑ k′=1 H−1∑ h=1 z (k′) h w (k′)> h+1 ∥∥∥2 (V (k))−1
where ‖X‖2V := tr ( X>V X ) and the last inequality uses Cauchy-Schwartz inequality. Notice that
‖Θ∗‖(V (k))−1 ≤ ‖Θ∗‖F . Moreover, we have∥∥∥ k∑
k′=1 H−1∑ h=1 z (k′) h w (k′)> h+1 ∥∥∥2 (V (k))−1 = ∥∥∥∥(V (k))−1/2 k∑ k′=1 H−1∑ h=1 z (k′) h w (k′)> h+1 ∥∥∥∥2 F
= ∑ j∈[d] ∥∥∥∥(V (k))−1/2 k∑ k′=1 H−1∑ h=1 w (k′) h+1,jz (k′) h ∥∥∥∥2 2 = ∑ j∈[d] ∥∥∥ k∑ k′=1 H−1∑ h=1 w (k′) h+1,jz (k′) h ∥∥∥2 (V (k))−1
By Theorem A.2, we have, for every j ∈ [d], with probability at least 1− δ/d, we have,∥∥∥ k∑ k′=1 H−1∑ h=1 w (k′) h+1,jz (k′) h ∥∥∥2 (V (k))−1 ≤ 2c2w log(ddet(V (k))1/2/δ).
By an union bound, we have, with probability at least 1− δ,∥∥∥ k∑ k′=1 H−1∑ h=1 w (k′) h+1z (k′) h ∥∥∥2 (V (k))−1 ≤ 2dc2w log(ddet(V (k))1/2/δ).
Plugging to tr [ (Θ∗ −Θ(k))V (k)(Θ∗ −Θ(k))> ] , we have, with probability at least 1− δ, tr [ (Θ∗ −Θ(k))V (k)(Θ∗ −Θ(k))> ] ≤ ( cΘ + cw √ 2d log(ddet(V (k))1/2/δ)
)2 ≤ ( cΘ + cw √ 2d ( log d+ p log(1 + kHc2x/p)/2 + log δ −1 ))2 .
This completes the proof.
We define IEK as the indicator for EK happens. We denote
M (k) ∗ = MΘ∗,C(k),D(k) , M (k) = MΘ(k),C(k),D(k) , and y (k) h = [x (k)> h , u (k)> h ] >.
On Ek, we have ∀k ∈ [K] : J∗1 (M̃ (k), xk1) ≤ J∗1 (M (k)∗ , xk1). We denote ∆(k) := Jπ k
h (M (k) ∗ , x1)− J∗h(M (k) ∗ , x1). We can rewrite equation 7 as
Regret(KH) = K∑ k=1 IEk∆(k) + K∑ k=1 (1− IEk)∆(k),
where the second term is non-zero with probability less than δ. For the first term, we have IEk∆(k) ≤ IEk [ Jπ k 1 (M (k) ∗ , x1)− J∗1 (M̃ (k), x1)) ] =: IEk · ∆̃ (k) 1 ,
where ∆̃
(k) h = J πk h (M (k) ∗ , xh)− J∗h(M̃ (k), xh).
Let us consider ∆̃(k)h . We denote filtration Fk,h as fixing the trajectory up to time (k, h) and all {C(k′), D(k′)}k′≤k. We have
∆̃ (k) h =x (k)> h Qhx (k) h + u (k)> h Rhu (k) h + Ew(k)h+1 [J πk h+1(M (k) ∗ , x (k) h+1) | Fk,h]
− x(k)>h Qhx (k) h − u (k)> h Rhu (k) h
− E w
(k) h+1
[ (M̃ (k)z
(k) h + w (k) h+1) >Ph+1(M̃ (k))(M̃ (k)z (k) h + w (k) h+1) | Fk,h ] − Ch+1(M̃ (k))
=E w
(k) h+1
[Jπ k h+1(M (k) ∗ , x (k) h+1) | Fk,h]
− E w
(k) h+1
[ (M̃ (k)z
(k) h + w (k) h+1) >Ph+1(M̃ (k))(M̃ (k)z (k) h + w (k) h+1) | Fk,h ] − Ch+1(M̃ (k))
=E w
(k) h+1
[Jπ k h+1(M (k) ∗ , x (k) h+1) | Fk,h]− Jπ k h+1(M (k) ∗ , x (k) h+1) + J πk h+1(M (k) ∗ , x (k) h+1)
− ( M̃ (k)z
(k) h )> Ph+1(M̃ (k)) ( M̃ (k)z (k) h ) − Ch+1(M̃ (k))
− E w
(k) h+1
[( w
(k) h+1 )> Ph+1(M̃ (k))w (k) h+1 ∣∣ Fk,h] =δ
(k) h + J πk h+1(M (k) ∗ , x (k) h+1)−
( M̃ (k)z
(k) h )> Ph+1(M̃ (k)) ( M̃ (k)z (k) h ) − Ch+1(M̃ (k))
− E w
(k) h+1
[( x
(k) h+1 −M (k) ∗ z (k) h )> Ph+1(M̃ (k)) ( x (k) h+1 −M (k) ∗ z (k) h ) ∣∣ Fk,h] =δ
(k) h + J πk h+1(M (k) ∗ , x (k) h+1)−
( M̃ (k)z
(k) h )> Ph+1(M̃ (k)) ( M̃ (k)z (k) h ) − Ch+1(M̃ (k))
− E w
(k) h+1
[( x
(k) h+1 )> Ph+1(M̃ (k)) ( x (k) h+1 ) ∣∣ Fk,h]+ (M (k)∗ z(k)h )>Ph+1(M̃ (k))(M (k)∗ z(k)h ) =δ
(k) h + δ ′(k) h + δ ′′(k) h + J πk h+1(M∗, x (k) h+1)− J∗h+1(M̃ (k), x (k) h+1)
where
δ (k) h = Ew(k)h+1 [J πk h+1(M (k) ∗ , x (k) h+1) | Fk,h]− Jπ k h+1(M (k) ∗ , x (k) h+1) (15)
δ ′(k) h =
( x
(k) h+1 )> Ph+1(M̃ (k)) ( x (k) h+1 ) − E
w (k) h+1
[( x
(k) h+1 )> Ph+1(M̃ (k)) ( x (k) h+1 ) ∣∣ Fk,h] (16) δ ′′(k) h = ( M (k) ∗ z (k) h )> Ph+1(M̃ (k)) ( M (k) ∗ z (k) h ) − ( M̃ (k)z (k) h )> Ph+1(M̃ (k)) ( M̃ (k)z (k) h ) . (17)
By induction, we have k∑
k′=1
∆̃ (k) 1 ≤ k∑ k′=1 H−1∑ h=1 ( δ (k) h + δ ′(k) h + δ ′′(k) h ) .
Notice that δ(k)h and δ (k) h are Martingale difference adapted to Fk,h. We can well bound the sum of them via Azuma’s inequality.
Lemma A.3. For all h ∈ [H], |Jπkh (M (k) ∗ , x (k) h )| ≤ (H − h+ 1) · cq · cx.
Proof. Prove by induction on h. The base case Jπ k H (M (k) ∗ , x (k) H ) = x (k)> H QHx (k) H ≤ cqc2x holds straightforwardly. Consider an arbitrary h < H , we have
Jπ k h (M (k) ∗ , x (k) h ) = x (k)> h Qhx (k) h +u (k)> h Rhu (k) h +Ew(k)h+1 [J πk h+1(M (k) ∗ , x (k) h+1) | Fk,h] ≤ cqcx+(H−h)·cq·cx
as desired.
Lemma A.4. For all x ∈ X , we have IEK |J∗h(M̃ (k), x)| ≤ cqcx.
Proof. Follows from Assumption 4.1.
We are now ready to prove Proposition 4.1.
Proof of Proposition 4.1. Thus by Azuma’s inequality, we have, with probability at least 1− δ,∣∣∣ k∑ k′=1 H−1∑ h=1 δ (k) h ∣∣∣ ≤√2kH · [(H − h+ 1)qcx + cqcx]2 · log 2 δ .
And, with probability at least 1− δ,∣∣∣ k∑ k′=1 H−1∑ h=1 δ ′(k) h ∣∣∣ ≤√8kH · c2xc2q · log 2δ . For
∑ δ ′′(k) h , we bound it here.∣∣∣ k∑
k′=1 H−1∑ h=1 δ ′′(k) h ∣∣∣ ≤ k∑ k′=1 H−1∑ h=1 ∣∣∣δ′′(k)h ∣∣∣ = k∑ k′=1 H−1∑ h=1 ∣∣∣‖Ph+1(M̃ (k))1/2(M̃ (k)y(k)h )‖22 − ‖Ph+1(M̃ (k))1/2(M∗y(k)h )‖22∣∣∣ ≤
k∑ k′=1 H−1∑ h=1 ∣∣∣(‖Ph+1(M̃ (k))1/2(M̃ (k)y(k)h )‖2 − ‖Ph+1(M̃ (k))1/2(M∗y(k)h )‖2) · ( ‖Ph+1(M̃ (k))1/2 ( M̃ (k)y (k) h ) ‖2 + ‖Ph+1(M̃ (k))1/2 ( M∗y (k) h ) ‖2 )∣∣∣
≤ [ k∑ k′=1 H−1∑ h=1 ( ‖Ph+1(M̃ (k))1/2 ( M̃ (k)y (k) h ) ‖2 − ‖Ph+1(M̃ (k))1/2 ( M∗y (k) h ) ‖2 )2]1/2
· [ k∑ k′=1 H−1∑ h=1 ( ‖Ph+1(M̃ (k))1/2 ( M̃ (k)y (k) h ) ‖2 + ‖Ph+1(M̃ (k))1/2 ( M∗y (k) h ) ‖2 )2]1/2
Notice that ‖Ph+1(M̃ (k))1/2 ( M̃ (k)y (k) h ) ‖2 ≤ cqcxcΘ and ‖Ph+1(M̃ (k))1/2 ( M∗y (k) h ) ‖2 ≤ cqcx. Hence [ k∑ k′=1 H−1∑ h=1
∣∣∣(‖Ph+1(M̃ (k))1/2(M̃ (k)y(k)h )‖2 + ‖Ph+1(M̃ (k))1/2(M∗y(k)h )‖2)2]1/2 ≤ √ kH · (cqcx(1 + cΘ))2.
Moreover, by triangle inequality, we have∣∣∣‖Ph+1(M̃ (k))1/2(M̃ (k)y(k)h )‖2 − ‖Ph+1(M̃ (k))1/2(M∗y(k)h )‖2∣∣∣ ≤ ‖Ph+1(M̃ (k))1/2 ( M̃ (k) −M∗ ) y (k) h ‖2
≤ cq‖ ( M̃ (k) −M∗ ) y (k) h ‖2
≤ cq‖ ( M̃ (k) −M∗ ) (V (k))1/2(V (k))−1/2y (k) h ‖2
≤ cq‖ ( M̃ (k) −M∗ ) (V (k))1/2‖2‖(V (k))−1/2y(k)h ‖2
≤ cq · √ β(k) · ‖(V (k))−1/2y(k)h ‖2.
By Assumption 2, we also have ‖(V (k))−1/2y(k)h ‖2 ≤ ‖(y (k) h ‖2 ≤ √ cx. Hence,∣∣∣‖Ph+1(M̃ (k))1/2(M̃ (k)y(k)h )‖2 − ‖Ph+1(M̃ (k))1/2(M∗y(k)h )‖2∣∣∣
≤ cq √ cx · √ β(k) ·min ( ‖(V (k))−1/2y(k)h ‖2, 1 ) Combining the above equations, we have,∣∣∣ k∑ k′=1 H−1∑ h=1 δ ′′(k) h ∣∣∣ ≤√kH · (cqcx(1 + cΘ))2 · cq√cx ·√β(k) · √√√√ k∑ k′=1 H−1∑ h=1 min ( ‖(V (k))−1/2y(k)h ‖22, 1 )
≤ 2c3/2x c2qcΘ · √ β(k) · √√√√ k∑ k′=1 H−1∑ h=1 log ( 1 + ‖(V (k))−1/2y(k)h ‖22 ) · √ kH.
Lastly, by Lemma 8 of Yang & Wang (2019), we have k∑
k′=1 H−1∑ h=1 log ( 1 + ‖(V (k))−1/2y(k)h ‖22 ) ≤ 2H log det(V (k)).
Together with Lemma A.1, we have 2H log det(V (k)) ≤ 2Hp · log ( 1 + kHc2x/p ) .
Overall, we have,∣∣∣ k∑ k′=1 H−1∑ h=1 δ ′′(k) h ∣∣∣ ≤ 2c3/2x c2qcΘ ·√2Hp · log (1 + kHc2x/p) · (β(k)) · kH. Putting everything together, with probability at least 1− 2δ, we have
reg(KH) ≤ K∑ k′=1 H−1∑ h=1 ( δ (k) h + δ ′(k) h + δ ′′(k) h ) ≤ √ 2KH · [ (H − h+ 1)cqcx + cqcx ]2 · log 2 δ + √ 8KH · c2xc2q · log 2 δ
+ 2c3/2x c 2 qcΘ ·
√ 2Hp · log ( 1 +KHc2x/p ) · (β(K)) ·KH
≤ cH · d1/2p · log3/2(dKHc2xδ−1) · √ KH,
where cH is a constant depending on H, cq, cx, cΘ and cw.

heading: B CONCRETE CHOICE OF THE PARAMETERS
text: We further augment the state so that the first coordinate is a constant with value 1. More specifically, we set the state xh = [1; zh; vh] ∈ R5. We set
Qh =  ‖z∗h‖22 −z∗h 0−z∗h I 0 0 0 0  so that for any state xh, xThQhxH = ‖zh − z∗h‖22. We set Rh = I with size 2× 2. We set
Θ∗ =  1 0 0 0 0 0 0 0 1 0 1 0 0 0 0 0 1 0 1 0 0 0 0 0 k 0 1 0 0 0 0 0 k 0 1  , C to be the 5 × 5 identity matrix and D to be I/m with size 2 × 2 where m is sampled from the uniform distribution over [0.1, 10], to represent the physical law in Equation 14.

Review summaries:
This paper considers the problem of changing environments for LQR. The authors model this through the use of a decoder that maps an incoming context (C,D) to the LQR matrices (A,B). They provide an algorithm for this setting based on a UCB strategy, prove sample complexity and regret bounds, and experimental results.

In order to generalize the RL agent to unseen environment, in this work the authors studied the theoretical learning problem of building a decoder on top of linear continuous control using linear quadratic regulator (LQR). They presented a simple, UCB-based algorithm that refines the estimates of the encoder while doing LQR and balances  the exploration-exploitation trade-off. In the online setting, the proposed algorithm has a O(\sqrt{T}) regret bound, where T is the number of environments the agent played. This also implies after certain exploration, the agent is able to transfer the learned knowledge to obtain a near-optimal policy to an unseen environment. To justify their theoretical bounds the authors also present experiments that demonstrate the effectiveness of the algorithm.