Title: PROVABLY EFFICIENT CONTEXTUAL LINEAR QUADRATIC REGULATOR

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 O √ T regret bound in the online setting where T is the number 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.

1. INTRODUCTION:

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.

2. RELATED WORK:

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 & Szepesvá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; Abbasi-Yadkori 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.

3. PRELIMINARIES:

Notations. We begin by introducing necessary notations. We write [h] to denote the set {1, . . . , h}. We use I d ∈ R d×d to denote the d-dimensional identity matrix. We use 0 d×d to represent the allzero matrix in R d×d . If it is clear from the context, we omit the subscript d × d . Let • 2 denote the Euclidean norm of a vector in R d . 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 ⊂ R d and a closed action space U ⊂ R d . Suppose we always start from the initial state x 1 = x init ∈ X and play for H steps. Then at a state x h ∈ X , if an action u h ∈ U is played, the next state is given by

EQUATION

where A, B are matrices of proper dimension and w h+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

y h = [x h , u h ] ∈ R d+d .

Now the state transition can be rewritten as x h+1 = M y h + w h+1 . For the ease of presentation, we assume that the covariance matrix of noise vector w h+1 is E(w h+1 w h+1 ) = I d . Our analysis follows similarly if the covariance matrix is not I d (see e.g. Remark 3 of Abbasi-Yadkori & Szepesvári (2011) ). After each step, the player receives an immediate cost

x h Q h x h + u h R h u h , where Q h , R h

are positive definite (PD) matrices of proper dimensions. Throughout the paper, we assume the agent knows Q h and R h for all h. At a terminal state x H , there is no action to be played, and the player receives a terminal cost x H Q H x H , where Q H 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 {(x i , u i )} h-1 i=1 ∪ {x h } to the next action u h , such that the following objectives are minimized:

J π h (M, x) := E H-1 h =h x h Q h x h + u h R h u h + x H Q f x H x h = x h∈[H]

,

where the action u h is given by u

h = π[(x 1 , u 1 ), (x 2 , u 2 ), . . . , (x h-1 , u h-1 ), x h ],

and the expectation is over the randomness of w h 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) := K h (M )x where M = [A, B] and K h (M ) is a matrix that will be defined shortly. It is also known (see e.g. Bertsekas (1996) ) that the optimal cost function J

EQUATION

where

EQUATION

and

C h (M ) = C h+1 (M ) + E w h+1 w h+1 P h+1 (M )w h+1 ] h < H 0 h = H .

We now define K h (M ) as

EQUATION

Note that the optimal value Equation (2) satisfies Bellman equations,

∀h ∈ [H -1] : J * h (M, x) = x Q h x + π * (x) R h π * (x) + E J * h+1 (Ax + Bπ * (x) + w) and ∀h ∈ [H -1] : J * h (M, x) = x Q h x + min u E[u R h u + 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 R p×d × R p ×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 Θ * ∈ R d×(p+p ) :

EQUATION

From now on, to emphasize that the representation of LQR can be decoded from Θ * , we write

EQUATION

)

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 on K 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 , x 1 ) -J * h (M Θ * ,C,D , x 1 )] ≤ 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) :=:

EQUATION

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 .

4. MAIN ALGORITHM:

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 ∈ R d×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 ∈ R d×2p , V (1) ← I2p,2p, W (1) ← 0 ∈ R 2p×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 ← K h M Θ (k) ,C (k) ,D (k) • x (k)

h , where K h 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) := K h M (k) • x where M (k) = M Θ (k) ,C (k) ,D (k) := Θ (k) • C (k) 0 0 D (k) ,

and K h is given by Equation (4). After episode k ∈ [K], we use the following ridge regression formulation to update decoder

EQUATION

)

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:

EQUATION

The formal algorithm is presented in Algorithm 1.

4.1. ALGORITHM ANALYSIS:

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], P h (M ) 2 ≤ c q for some parameter c q > 0.

• Θ * F ≤ c Θ ; • ∀h ∈ [2, H], i ∈ [d] : w h 2 < ∞ and ∀γ > 0, E[γw h,i ] ≤ exp(γ 2 c 2 w /2); • ∀x ∈ X , u ∈ U, (C, D) ∈ supp(µ) : Cx 2 + Du 2 2 ≤ c 2 x , x 2 + u 2 2 ≤ c 2 x ; • ∀(C, D) ∈ supp(µ), x ∈ X , h ∈ [H]: K h (M Θ * ,C,D ) • x ∈ U.

where c Θ , c w , c x 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. k) ,

C (k) = Θ : tr Θ -Θ (k) V (k) Θ -Θ (k) ≤ β (

EQUATION

)

where P h is given by Equation (3) and β (k) is defined as follows,

EQUATION

)

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 • dp 2 • log 3 (dKδ -1 )

2 episodes, for some parameter c H,cq,cx,cΘ,cw depending polynomially on H, c q , c x , c Θ , c w , Then with probability at least 1δ, we have for π C,D be defined in Equation 9.

EQUATION

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 • d 1/2 p • log 3/2 (dKHc x δ -1 ) • √ KH.

where c H is a constant depending only polynomially on H, c q , c x , c M , c w .

By the definition of regret, this proposition justifies that the performance of the agent actually improves as it sees more environment.

5. EXPERIMENTS:

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 ∈ R 2 . Our goal is to exert forces u 1 , u 2 , . . . , u m ∈ R 2 on objects with different (measurable) masses to minimize the total squared distance

H h=1 z h -z * h 2 + u i 2 2 . Each state x h = [z h ; v h ] ∈ R 4

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 u h ∈ R 2 on the object, which produces an accelerations u h m ∈ R 2 . The dynamics of the system can be described as

EQUATION

)

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 vector w h 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) = 10 4 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 z h -z * h 2 + u i 2 2

. 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 = x 2 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 ∈ R 2 . 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.

6. CONCLUSION:

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) . 

A PROOF OF MAIN RESULTS:

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.

E C,D E π J π 1 M Θ * ,C,D , x 1 -E C,D J * 1 M Θ * ,C,D , x 1 = 1 K K k=1 E C,D J π k 1 M Θ * ,C,D , x 1 -E C,D J * 1 M Θ * ,C,D , x 1 = 1 K K k=1 E C,D J π k 1 M Θ * ,C,D , x 1 -J π k 1 M Θ * ,C (k) ,D (k) , x 1 + J π k 1 M Θ * ,C (k) ,D (k) , x 1 -J * 1 M Θ * ,C (k) ,D (k) , x 1 + J * 1 M Θ * ,C (k) ,D (k) , x 1 -E C,D J * 1 M Θ * ,C,D , x 1 = R 1 + R 2 + R 3 where R 1 = 1 K K k=1 E C,D J π k 1 M Θ * ,C,D , x 1 -J π k 1 M Θ * ,C (k) ,D (k) , x 1 , R 2 = 1 K K k=1 J * 1 M Θ * ,C (k) ,D (k) , x 1 -E C,D J * 1 M Θ * ,C,D , x 1 ,

and

R 3 = 1 K K k=1 J π k 1 M Θ * ,C (k) ,D (k) , x 1 -J * 1 M Θ * ,C (k) ,D (k) , x 1 .

Let F k be the filtration of fixing all randomness before episode k. We have R 1 and R 2 are Martingale difference sum. Note that the magnitude of each summand in R 1 or R 2 is upper bounded by (proved in Lemma A.3 and A.4), Hc q c x almost surely. Therefore, by Azuma's inequality (Theorem A.1), we have, with probability greater than 1δ/2,

|R 1 | + |R 2 | ≤ 2Hc q c x • 2 log(8/δ) K .

Moreover, by Proposition 4.1, we have with probability greater than 1δ/2,

|R 3 | ≤ c • d 1/2 p • log 3/2 (dKHc 2 x δ -1 ) • H K ,

where c is constant depending only polynomially on H, c q , c x , c M , and c w . 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 {X s } s≥0 is a martingale and |X s -X s-1 | ≤ c s almost surely. Then for all t > 0 and all > 0,

Pr |X t -X 0 | ≥ ≤ 2 exp -2 2 t s=1 c 2 s .

Theorem A.2 (Martingale Concentration, Theorem 16 of Abbasi-Yadkori & Szepesvári (2011) ). Let F t ; t ≥ 0 be a filtration, (z t ; t ≥ 0) be an R d -valued stochastic process adapted to (F t ). Let (η t ; t ≥ 1) be a real-valued martingale difference process adapted to F t . Assume that η t is conditionally sub-Gaussian with constant L, i.e.,

∀γ > 0 : E[γη t |F t ] ≤ exp(γ 2 L 2 /2).

Consider the following martingale

S t = t τ =1 η τ z t-1

and the matrix-valued processes

V t = I + t τ =0 z t-1 z t-1 .

Then for any δ ∈ (0, 1), with probability at least 1δ,

∀t ≥ 0, S t 2 V -1 t ≤ 2L 2 log det(V t ) 1/2 δ where S t 2 V -1 t := S t V -1 t S t .

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 + kHc 2 x /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 2 2 /p p .

By Assumption 4.1, we have z

(k ) h 2 2 ≤ c 2

x . This completes the proof.

Let us then define an event E k as follows.

Definition A.1 (Good Event). We define event E k as {∀k ≤ k : Θ * ∈ C (k ) }.

We then show that the event E k happens with high probability.

Lemma A.2. For all k ∈ [K], we have Pr[E k ] ≥ 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+1 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 w (k ) h+1 z (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 k k =1 H-1 h=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 2 V := 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,j z (k ) h 2 2 = j∈[d] k k =1 H-1 h=1 w (k ) h+1,j z (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,j z (k ) h 2 (V (k) ) -1 ≤ 2c 2 w log(d det(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+1 z (k ) h 2 (V (k) ) -1 ≤ 2dc 2 w log(d det(V (k) ) 1/2 /δ).

Plugging to tr (Θ * -Θ (k) )V (k) (Θ * -Θ (k) ) , we have, with probability at least 1δ,

tr (Θ * -Θ (k) )V (k) (Θ * -Θ (k) ) ≤ c Θ + c w 2d log(d det(V (k) ) 1/2 /δ) 2 ≤ c Θ + c w 2d log d + p log(1 + kHc 2 x /p)/2 + log δ -1 2 .

This completes the proof.

We define I E K as the indicator for E K happens. We denote

EQUATION

EQUATION

] .

On E k , we have ∀k ∈ [K] :

J * 1 ( M (k) , x k 1 ) ≤ J * 1 (M (k) * , x k 1 ). We denote ∆ (k) := J π k h (M (k) * , x 1 ) -J * h (M (k) * , x 1

). We can rewrite equation 7 as

Regret(KH) = K k=1 I E k ∆ (k) + K k=1 (1 -I E k )∆ (k) ,

where the second term is non-zero with probability less than δ. For the first term, we have

I E k ∆ (k) ≤ I E k J π k 1 (M (k) * , x 1 ) -J * 1 ( M (k) , x 1 )) =: I E k • ∆ (k) 1 , where ∆ (k) h = J π k h (M (k) * , x h ) -J * h ( M (k) , x h ). Let us consider ∆ (k)

h . We denote filtration F k,h as fixing the trajectory up to time (k, h) and all {C (k ) , D (k ) } k ≤k .

We have

EQUATION

(k) h = M (k) * z (k) h P h+1 ( M (k) ) M (k) * z (k) h -M (k) z (k) h P h+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 F k,h . We can well bound the sum of them via Azuma's inequality.

Lemma A.3. For all h ∈ [H], |J π k h (M (k) * , x (k) h )| ≤ (H -h + 1) • c q • c x .

Proof. Prove by induction on h. The base case

J π k H (M (k) * , x (k) H ) = x (k) H Q H x (k)

H ≤ c q c 2 x holds straightforwardly. Consider an arbitrary h < H, we have

J π k h (M (k) * , x (k) h ) = x (k) h Q h x (k) h +u (k) h R h u (k) h +E w (k) h+1 [J π k h+1 (M (k) * , x (k) h+1 ) | F k,h ] ≤ c q c x +(H-h)•c q •c x as desired.

Lemma A.4. For all x ∈ X, we have

I E K |J * h ( M (k) , x)| ≤ c q c x .

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)qc x + c q c x 2 • log 2 δ .

And, with probability at least 1δ,

k k =1 H-1 h=1 δ (k) h ≤ 8kH • c 2 x c 2 q • 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 P h+1 ( M (k) ) 1/2 M (k) y (k) h 2 2 -P h+1 ( M (k) ) 1/2 M * y (k) h 2 2 ≤ k k =1 H-1 h=1 P h+1 ( M (k) ) 1/2 M (k) y (k) h 2 -P h+1 ( M (k) ) 1/2 M * y (k) h 2 • P h+1 ( M (k) ) 1/2 M (k) y (k) h 2 + P h+1 ( M (k) ) 1/2 M * y (k) h 2 ≤ k k =1 H-1 h=1 P h+1 ( M (k) ) 1/2 M (k) y (k) h 2 -P h+1 ( M (k) ) 1/2 M * y (k) h 2 2 1/2 • k k =1 H-1 h=1 P h+1 ( M (k) ) 1/2 M (k) y (k) h 2 + P h+1 ( M (k) ) 1/2 M * y (k) h 2 2 1/2

Notice that P h+1 ( M (k) ) 1/2 M (k) y (k) h 2 ≤ c q c x c Θ and P h+1 ( M (k) ) 1/2 M * y

(k) h 2 ≤ c q c x . Hence k k =1 H-1 h=1 P h+1 ( M (k) ) 1/2 M (k) y (k) h 2 + P h+1 ( M (k) ) 1/2 M * y (k) h 2 2 1/2 ≤ kH • (c q c x (1 + c Θ )) 2 .

Moreover, by triangle inequality, we have

P h+1 ( M (k) ) 1/2 M (k) y (k) h 2 -P h+1 ( M (k) ) 1/2 M * y (k) h 2 ≤ P h+1 ( M (k) ) 1/2 M (k) -M * y (k) h 2 ≤ c q M (k) -M * y (k) h 2 ≤ c q M (k) -M * (V (k) ) 1/2 (V (k) ) -1/2 y (k) h 2 ≤ c q M (k) -M * (V (k) ) 1/2 2 (V (k) ) -1/2 y (k) h 2 ≤ c q • β (k) • (V (k) ) -1/2 y (k) h 2 .

By Assumption 2, we also have (V (k) ) -1/2 y (k) h 2 ≤ (y

(k) h 2 ≤ √ c x . Hence, P h+1 ( M (k) ) 1/2 M (k) y (k) h 2 -P h+1 ( M (k) ) 1/2 M * y (k) h 2 ≤ c q √ c x • β (k) • min (V (k) ) -1/2 y (k) h

2 , 1 Combining the above equations, we have,

k k =1 H-1 h=1 δ (k) h ≤ kH • (c q c x (1 + c Θ )) 2 • c q √ c x • β (k) • k k =1 H-1 h=1 min (V (k) ) -1/2 y (k) h 2 2 , 1 ≤ 2c 3/2 x c 2 q c Θ • β (k) • k k =1 H-1 h=1 log 1 + (V (k) ) -1/2 y (k) h 2 2 • √ kH.

Lastly, by Lemma 8 of Yang & Wang (2019), we have k) ).

k k =1 H-1 h=1 log 1 + (V (k) ) -1/2 y (k) h 2 2 ≤ 2H log det(V (

Together with Lemma A.1, we have 2H log det(V (k) ) ≤ 2Hp • log 1 + kHc 2 x /p . Overall, we have,

k k =1 H-1 h=1 δ (k) h ≤ 2c 3/2 x c 2 q c Θ • 2Hp • log 1 + kHc 2 x /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)c q c x + c q c x 2 • log 2 δ + 8KH • c 2 x c 2 q • log 2 δ + 2c 3/2 x c 2 q c Θ • 2Hp • log 1 + KHc 2 x /p • (β (K) ) • KH ≤ c H • d 1/2 p • log 3/2 (dKHc 2 x δ -1 ) • √ KH,

where c H is a constant depending on H, c q , c x , c Θ and c w .

B CONCRETE CHOICE OF THE PARAMETERS:

We further augment the state so that the first coordinate is a constant with value 1. More specifically, we set the state x h = [1; z h ; v h ] ∈ R 5 . We set

Q h =   z * h 2 2 -z * h 0 -z * h I 0 0 0 0   so that for any state x h , x T h Q h x H = z h -z * h 2 2

. We set R h = 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.

:

https://en.wikipedia.org/wiki/Lemniscate_of_Bernoulli.

Figures and Tables:

Figure 1: Θ -Θ * F and Mean Control Error.

Peer review (iclr2020):
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. Overall the paper was well written but I had several concerns. 1. The results of this paper were not contrasted with other papers in this area. For example, if C,D are constant, and Theta_* is a Block diagonal matrix with A,B on the diagonal - then the contextual case reduces to the standard LQR problem. It*s unclear how the results given compare to past results in this setting, for example those of Abbasi-Yadkori/Szepesvari 2011. 2. I did not fully understand the UCB nature of the algorithm. In each round Theta^(k) (the least squares estimator) seems to be used to compute the optimal policy (line 10 of the algorithm) instead 	ilde{Theta}^(k) -the optimistic estimate. The optimistic estimate is only used in line 16 - a randomized procedure that is unmotivated. 3. Building on (1), it is hard to understand the results as given since there are no lower bounds given nor is there a discussion of the problem dependent parameters that arise. For example, in Theorem 1, is dp^2 suspected to be tight? Since the number of parameters in Theta^{ast} is d(p+p*), perhaps this is off by a factor of p? 3. I struggled to understand the setup of the experiments - as described the algorithm given was not used at all, rather Theta^(k) was approximated and beta^k was set to be a constant. This does not seem like a fair evaluation of the method. In summary, I would reject this submission unless the authors couch it better in past work, explain their results better, and improve the experiment setup. Finally, a typo: I think the indexing variable in the equation on the top of page 4 is h* not h.