P ROVABLY EFFICIENT MULTI - TASK R EINFORCEMENT L EARNING IN LARGE STATE SPACES
Anonymous authors
Paper under double-blind review
A BSTRACT
We study multi-task Reinforcement Learning where shared knowledge among different environments is distilled to enable scalable generalization to a variety of problem instances.
In the context of general function approximation, Markov
Decision Process (MDP) with low Bilinear rank encapsulates a wide range of structural conditions that permit polynomial sample complexity in large state spaces, where the Bellman errors are related to bilinear forms of features wi ‘th low intrinsic dimensions.
To achieve multi-task learning in MDPs, we propose online representation learning algorithms to capture the shared features in the different task-speciﬁc bilinear forms.
We show that in the presence of low-rank structures in the features of the bilinear forms, the algorithms beneﬁt from sample complexity improvements compared to single-task learning.
Therefore, we achieve the ﬁrst sample efﬁcient multi-task reinforcement learning algorithm with general function approximation
I NTRODUCTION
The ability to capture informative representations that generalize among multiple tasks has become signiﬁcant in various machine learning applications Li et al.
; Tsiakas et al.
; Baevski et al.
; D’Eramo et al.
; Kubota et al.
; Liu et al. (2019b).
In the context of multi-task learning Caruana (1997); Baxter (2000); Yu et al.
(2005), this ability is highly desirable and becomes vital to learn with fewer amount of samples than learning each single task individually.
Representation learning Bengio et al.
(2013) is a powerful approach for achieving such sample efﬁciency improvement.
This paper considers representation learning in Multitask Reinforcement Learning, an important class of meta Reinforcement Learning (meta-RL)
Wang et al.
Finn et al.
Ritter et al.
Reinforcement learning (RL) is a sequential decision-making problem where an agent aims to learn the optimal decisions by interacting with an unknown environment Sutton & Barto (2018).
Empowered by representation learning with deep neural networks LeCun et al.
; Goodfellow et al.
, RL has achieved tremendous success in various real-world applications, such as Go
Silver et al.
(2016), Atari Mnih et al.
Dota2 Berner et al.
(2019), Texas Holdém poker
Moravˇcík et al.
(2017), and autonomous driving Shalev-Shwartz et al.
Therefore, the beneﬁt of using representation learning to extract joint feature embedding from different but related tasks emerged as an essential problem to investigate.
Speciﬁcally, this paper studies the problem of learning multiple RL problems jointly with the help of representation learning.
Although multi-task learning in online-decision making problems has received increasing research interest Lazaric & Ghavamzadeh (2010); Mutti et al.
; Maurer et al.
; Qin et al. (2021); Yang et al. (2021); Hu et al.
(2021), most existing works focus on tabular or linear models.
Indeed, how general function approximations extrapolate across huge state spaces remains largely an open problem itself.
Recently, Bilinear class Du et al.
(2021) proposes a promising structural framework of generalization in reinforcement learning through the use of function approximation.
Bilinear class postulates that the Bellman error can be related to a bilinear form depending on the hypothesis and captures nearly all existing function approximation models, e.g.
Jin et al.
(2020a)
; Zanette et al.
; Yang & Wang (2020); Jiang et al.
(2017); Sun et al.
Kakade et al.
Agarwal et al.
However, in the presence of shared information in the bilinear forms across multiple tasks, the Bilin-UCB proposed in Du et al.
(2021) is not able to adapt to such knowledge and challenges abound in adopting representation learning to ﬁnd nearly-optimal policies with limited data.
In this paper, we give the ﬁrst sample efﬁcient algorithm of multi-task RL with general function approximation through the usage of representation learning in Bilinear class.
In particular, to study representation learning, we propose Low-rank Multi-task Bilinear class—a structural framework that permits generalization both within and across tasks in multi-task RL.
Speciﬁcally, such a model class speciﬁes M MDP instances, where M > 0 is a ﬁxed integer and each one belongs to the
Bilinear class Du et al.
(2021), i.e., the Bellman error admits a low-rank factorization in R d .
Since our multi-task setting have M MDP instances, there are M different features maps speciﬁed by the deﬁnition of Bilinear class, each corresponding to one MDP task and taking values in R d .
We additionally assume that these M features maps, when packed together as a matrix-valued mapping in R d × M , has rank k (cid:28) d .
In other words, in Low-rank Multi-task Bilinear class, the bilinear form of each task possesses a low-dimensional task-speciﬁc feature and a shared representation.
Under this setting, it is desirable that the RL algorithm utilize the intrinsic low-dimensional structure to achieve an improved sample efﬁciency compared to solving each task separately.
To this end, under the online setting where the agent learn from its past experiences without knowing the model, we design a sample efﬁcient algorithm that provably ﬁnds nearly-optimal policies for all tasks.
Our algorithm is based on the principle of Optimism in the Face of Uncertainty (OFU) which constructs an conﬁdence region that contains the true hypothesis based on the historical data across the M tasks, and then update the policy according to the most optimistic hypothesis within the conﬁdence region.
In particular, here the hypothesis can denote the true transition models or optimal value functions of these M tasks.
When constructing the conﬁdence region, we explicitly utilize the low-dimensional structure by joint learning the task-speciﬁc
features and shared representation via Empirical Risk Minimization (ERM) with multi-task data.
Moreover, as for planning, we ﬁnd the hypothesis in the conﬁdence region which leads to the highest aggregated value in these M tasks.
In the analysis, we show a concentration result where the estimation noise can be embedded into low dimension space and thus prove that our algorithm is able to ﬁnd nearly-optimal policies within limited samples.
Concretely, compared to learning each task separately using Bilin-UCB Du et al.
(2021), an algorithm designed for Bilinear class without utilizing the shared representation, our algorithm enjoys a ( d/k ) -time improvement in the sample efﬁciency whenever the feature classes are small.
To our best knowledge, our work seems to propose the ﬁrst provably sample efﬁcient multi-task RL algorithm with general function approximation.
Notations Let R d denote the d -dimensional space and R d × k denote the space of d -by- k matrices in R .
The inner product of two vectors x, y ∈ R
d is denoted as (cid:104) x, y (cid:105) .
For sets A 1 , . . .
, A n , deﬁne ⊗ k ∈ [ n ] A k
=
A 1 ⊗ · · · ⊗ A k
= { ( a 1 , . . . , a n ) : a k ∈ A k , k ∈ [ n ] } .
Given scalars a 1 , · · · , a n , let a 1: n denote the vector ( a 1 , · · · , a n ) .
Also let ( a ω ) ω ∈ Ω denote the tuple consisting of a ω where ω comes from a countable set Ω .
For variables v 1 , . . .
, v k , we denote by v 1: k the k -tuple ( v 1 , . . .
, v k ) .
Roadmap
In Section 2 we introduce the basic problem setup and notations.
In Section 3 we introduce Low-rank Multi-task Bilinear class—a framework that captures shared information in bilinear class.
Next, we display the main algorithm of learning Low-rank Multi-task Bilinear class models, empowered by representation learning and optimism principle in decision making, in Section 4.
We show the main theoretical result in Section 5 and the overview of techniques in
Section 6.
A couples of examples are given in Section A.
We conclude with discussions of further directions.
1. R ELATED W ORK
General function approximation in Reinforcement Learning Theoretical understanding of the sample complexity of RL with general function approximation seems relatively scarce.
In recent years, there has been a surge of theoretical insights on linear function approximation and non-linear function approximations Jin et al.
(2020b;c); Wang et al.
; Zanette et al.
(2020); Agarwal et al.
(2020); Kakade et al.
; Wen & Van Roy (2017); Dann et al.
Du et al.
Dong et al.
; Liu et al. (2019a); Wang et al.
(2020a)
; Dong et al.
; Zhou et al.
; Yang et al.
(2020); Jin et al.
(2021a);
Du et al.
Among them, Bilinear class Du et al.
(2021) is one of the most general framework.
Multi-task learning and Meta learning Meta Learning Thrun & Pratt (2012) and multi-task representation learning Caruana (1997) are important tools for capturing shared knowledge across tasks and achieving generalization to a new task.
Theoretical analysis dated back to Baxter (2000);
Maurer (2005); Yu et al.
Maurer et al.
; Amit & Meir (2018); Konobeev et al.
(2020) consider generalization errors averaged over the meta distribution under the assumption of shared distribution for sampling the source tasks.
Recently, Du et al.
Tripuraneni et al.
(2020b;a);
Chua et al.
(2021) consider the diversity of task distributions and establish risk bounds based on learning shared representation across tasks and transferring to new tasks.
Multi-task bandit and RL In multi-task bandit and RL problems, theoretical analysis often postulates the existence of low-rank of sparsity structures in the representation Lazaric & Ghavamzadeh (2010); Brunskill & Li (2013); Calandriello et al.
; Mutti et al.
; Maurer et al.
D’Eramo et al.
Arora et al.
Qin et al. (2021); Yang et al.
; Hu et al.
The most related works are Hu et al.
; Yang et al.
(2021) where the beneﬁts of representation learning of linear bandits and linear MDPs are studied.
However, it remains open whether multi-task
RL can beneﬁt from representation learning via general function approximation.
P RELIMINARIES
We consider learning a set of M problem instances
[Equation]
where each M m ∈ P denotes an episodic Markov Decision Process (MDP) in which S m denotes the state space, A m denotes the action set, H denotes the number of time steps in each episode, { P m,h } Hh =1 denotes the transition kernel, r m denotes the reward function, and s m, 1 denotes the ﬁxed initial state.
We assume r m ∈ (0 , 1) without loss of generality.
For MDP M m , let E m denote the expectation under { P m,h } Hh =1 .
A (deterministic) policy π m is a length-
H sequence of functions π m = { π m,h : S m (cid:55)→
A m } Hh =1 .
To interact with M m , the agent starts at a ﬁxed initial state s m, 1 and at each time step h ∈
[ H ] , it takes action a m,h ∼ π m,h , receives reward r h ( s m,h , a m,h ) and transits to s m,h +1 ∼ P ( ·| s m,h , a m,h ) .
Let E mπ m denote the expectation under MDP M m and taking policy π m .
We use o m,h
= ( s m,h , a m,h , s m,h +1 ) to denote the history at h -th time step in MDP M m .
Given a policy π m , we deﬁne the value function V π m m,h ( s ) as the expected sum of reward under policy π m starting from s m,h
= s at time step h :
[Equation]
Similarly, we deﬁne the Q-function Q π m m,h ( s, a ) as the the expected sum of reward taking action a in state s m,h = s and then following π m,h :
[Equation]
The Bellman operator T m,h applied to Q-function Q : S m × A m (cid:55)→ R is deﬁned via
[Equation]
There exists an optimal policy π ∗ m that gives the optimal value function for all states, i.e. , V π ∗ m m,h ( s ) =
sup π V πm,h ( s ) holds for all h ∈ [ H ] and s ∈ S m .
Therefore Q π ∗ m m,h satisﬁes the following Bellman optimality equations for all s ∈ S m , a ∈
A m and h ∈
[ H ] :
[Equation]
The agent aims at using fewer samples to ﬁnd a set of policies { π m } M m =1 such that
[Equation]
holds with probability at least 1 − δ .
In the following, we deﬁne the ﬁltration H t to be the σ -ﬁeld induced by all the random variables up to round t .
M ULTI - TASK B ILINEAR CLASS
In general function approximation, a hypothesis class denoted by G is applied to address the set P of
M problem instances.
Here, each hypothesis g ∈ G is a function approximation that captures the information of both task-speciﬁc information of all M problem instances and the shared knowledge across tasks.
For example, in model-based RL, g may denote the transition models of the M tasks, that is, { P m,h } M,Hm =1 ,h =1 .
Whereas in model-free RL, g can be the optimal value functions of the M tasks, i.e., { V π ∗ m m,h } M,Hm =1 ,h =1 .
Using the notion of hypothesis g ∈ G , we aim to cover a large class of
MDP models studied under the function approximation setting.
We will introduce concrete examples in Section A.
Under our setting, each hypothesis g ∈ G is associated with Q-functions { Q m,h,g ( · , · ) } m ∈
[ M ] ,h ∈ [ H ] and value functions { V m,h,g ( · ) } m ∈
[ M ] ,h ∈ [ H ] in the M tasks such that V m,h,g ( · ) = max a Q m,h,g ( · , a ) holds for all m ∈
[ M ] and h ∈ [ H ] .
Bilinear class Du et al.
(2021) is a general framework that allows generalization in RL for a wide range of function approximators.
In the following, we develop the low-rank multi-task bilinear framework so that it permits generalization in meta-RL across different tasks.
The key intuition behind this framework is that it captures the common representation via low-rank structures in the features of task-speciﬁc bilinear forms.
Deﬁnition 3.1 (Low-rank Multi-task Bilinear class) .
We say the following tuple
[Equation]
is a multi-task Bilinear class with rank k if there exists g ∗ ∈ G such that Q m,h,g ∗ =
Q π ∗ m m,h and
V m,h,g ∗ =
V π ∗ m m,h hold for all m ∈ [ M ] , h ∈
[ H ] , and there exist functions W ∗ m,h : G (cid:55)→ R d , v ∗ m,h : G (cid:55)→ R k , B ∗ h : G (cid:55)→ R d × k , and X ∗ m,h : G (cid:55)→ R d
( d (cid:29) k ) such that for each g ∈ G , m ∈
[ M ] and h ∈ [ H ] :
The features W ∗ m,h possess low rank structures:
[Equation]
We can upper bound the expected Bellman error as follows:
[Equation]
Here π m,g,h ( s ) =
arg max a Q m,h,g ( s, a ) is the optimal policy in MDP M m under hypothesis g .
For any f ∈ G there exist policy π est ,m ( f ) and discrepancy measure l m,h,f ( o m,h , g ) that can be used for estimating (cid:104) v m,h ( g ) B h ( g ) − v m,h ( g ∗ )
B h ( g ∗ ) , X m,h ( f ) (cid:105) for any g ∈ G , such that:
[Equation]
where o m,h
= ( s m,h , a m,h , s m,h +1 ) .
Remark 3.2 .
A few remarks are in order.
First, when d = k and M = 1 , we recover the Bilinear class introduced in Du et al.
Thus, our model can be viewed as an multi-task extension of
Bilinear class with a shared low-dimensional representation.
Second, here the hypothesis class G can either be model-based function approximation or value-based function approximation, as in either case the Q-functions Q m,h,g and the value functions V m,h,g are well deﬁned.
Third, the common representation B ∗ h ( · ) captures the knowledge shared in all problem instances and enables multi-task learning with fewer samples.
Finally, the discrepancy measure l m,h,f ( o m,h , · ) can be computed for all hypothesis g , an important property that facilitates data reusage as explained below.
Examples of
Low-rank Multi-task Bilinear class can be found in below and in Appendix A.
To understand this framework, we ﬁrst notice that the expected Bellman errors of the following forms
[Equation]
serve as upper bounds for the sub-optimality of policies π m,g .
Thus the Eq.
(2) indicates that the sub-optimality of policies π m,g can be controlled if the feature W ∗ m,h ( g ) is close to the feature
W ∗ m,h ( g ∗ ) corresponding to the best-in-class function approximator g ∗ .
Eq.
(3) further establishes the connection between the error W ∗ m,h ( g ) − W ∗ m,h ( g ∗ ) and the discrepancy measure l m,h,f ( o m,h , g ) .
One important observation is information sharing among the function class through the feature
X ∗ m,h ( f ) : given samples from a single function approximator f , the quantity
[Equation]
can be estimated for all g ∈ G .
This thus allows for data reusage for each on-policy samples in our algorithm.
Eq.
(1) postulates that the features W ∗ m,h ( g ) share a low-dimension structure.
Note that this structure is not ﬁxed for all function approximator g — in fact, the common representation B h ( g ) may also be a function of g .
This allows for much generality in multi-task RL models.
We will show that the low-rank multi-task bilinear class cover many existing function approximation models in multitask setting.
As compared with a great number of multi-task/meta RL algorithms that use a single representation to solve multiple tasks, our function approximators allow for more generality and therefore handles more structural conditions of multi-task RL such as those studied in Yang et al.
(2021) and Hu et al.
Moreover, our algorithm reduces to using a single representation for the feature W in many special cases, for example when the feature v ∗ m,h is trivial.
In this sense, our algorithm can be seen as encompassing the method of using a single feature for multiple tasks.
Now we give an example of low-rank multi-task bilinear framework.
This example shows that the proposed framework captures the linear mixture model Modi et al.
(2020) in multi-task RL where the mixing coefﬁcients of different tasks lie in the same low-dimension space.
First, we recall the deﬁnition of Linear Mixture Model in the following.
Deﬁnition 3.3 (Linear Mixture Model) .
MDP M = ( S , A , { P h } Hh =1 , r, s 1 ) is a linear mixture model if there exist known feature maps φ : S ×
A × S (cid:55)→ R d and ψ : S × A (cid:55)→ R d and unknown θ ∗ h ∈ R d ,
h ∈
[ H ] such that
[Equation]
It is known that linear mixture model is Bilinear class with hypothesis g =
( θ 1 , . .
. , θ H ) and
[Equation]
The discrepancy measure can chosen as
[Equation]
and the estimation policies can be chosen as π est ( f ) =
π f .
We consider learning the special case of
Low-rank Multi-task Bilinear class where each MDP M m is a linear mixture model with
[Equation]
and there exist ν ∗ m,h ∈ R k , B ∗ h ∈ R d × k such that θ ∗ m,h
= B ∗ h ν ∗ m,h for all h ∈
[ H ] , m ∈
[ M ] .
Then we let g =
( ν m,h , B h ) h ∈
[ H ] ,m ∈
[ M ] and use the (ﬁxed) features v m,h ( g ) =
ν m,h , B h
( g ) =
B h .
Notice that each ( ν m,h , B h ) h ∈ [ H ] ,m ∈
[ M ] will deﬁne the expectation E mπ m,g (via P m,h ( s, a, s (cid:48) )
= (cid:104) θ m,h , φ m ( s, a, s (cid:48) ) (cid:105) with θ m,
h = B h ν m,h ).
Thus we can set the feature class X induced by G in which the feature X m,h ( g ) can be computed for each g ∈ G as follows:
[Equation]
Input: V , B , X for t ← 1 , . .
, T do
Find g ( t ) as the solution of the optimization problem in Eq.
For ∀ h ∈ [ H ] , m ∈
[ M ] , sample n 0 times from a m, 1: h − 1 ∼ π m,g ( t ) , a m,h ∼ π m, est and create datasets D ( t ) m,h . end for
Let t 0 ← max t ∈
[ T ] (cid:80)
Mm =
1 V π m,g ( t ) m,( s 1 )
Return π g ( t 0)
R EPRESENTATION L EARNING IN L OW - RANK M ULTI - TASK B ILINEAR CLASS
In this section we present the main algorithm to learn Low-rank Multi-task Bilinear class.
First, we deﬁne in the following the feature classes to capture v ∗ m,h , B ∗ h , X ∗ m,h for m ∈
[ M ] , h ∈ [ H ] .
Deﬁnition 4.1 (Function approximations) .
We deﬁne feature class V
= V 1 ⊗ · · · ⊗ V H where
V h ⊂ { v : G (cid:55)→ R k } ⊗ M , ∀ h ∈
[ H ] , representation class B = B 1 ⊗ · · · ⊗ B H where B h ⊂ { B : G (cid:55)→ R d × k } , and feature class X
=
X 1 ⊗ · · · ⊗ X H where X h ⊂ { X : G (cid:55)→ R d } ⊗ M , ∀ h ∈ [ H ] .
Here we assume that function classes V , B , and X completely captures the mappings speciﬁed in the multi-task Bilinear class.
Therefore, given expressive V , B , X as inputs, we can learn the representations v 1: M, 1: H ∈ V , B 1: H ∈ B , X 1: M, 1:
H ∈ X to approximate v ∗ 1: M, 1: H , B ∗ 1: H , X ∗ 1: M, 1: H by minimizing some proper loss function.
Based on this assumption, we propose an algorithm based on the OFU principle while leveraging representation learning to improve sample efﬁciency.
The procedure is shown in Algorithm 1.
In particular, in the t -th iteration, an (optimistic) hypothesis g ( t ) ∈ G is computed using Upper
Conﬁdence Bound (UCB) by ﬁnding a hypothesis that achieves the highest total value in these M tasks.
Speciﬁcally, consider the following constrained optimization problem:
[Equation]
which maximizes the sum of candidate value functions V m, 1 ,g i ( s m, 1 ) subject to the constraint that the hypothesis g ( t ) belongs to the conﬁdence set G ( t )
As we will show in the proof, G ( t ) contains the true hypothesis g ∗ for all t ∈
[ T ] with high probability.
Thus, by solving (4), the sum of value functions of g ( t ) serves as an upper bound of that of g ∗ .
The key issue now is how the conﬁdence set is chosen.
In BiLin-UCB, the conﬁdence set is chosen to contain all hypotheses that possess small values in the average of discrepancy measures across the available batch data.
Since the discrepancy measures serve as unbiased estimates for the bilinear forms which upper bound the Bellman error, this conﬁdence set essentially ﬁnds all hypotheses with low Bellman error.
However, this approach fails to exploit the shared information among tasks.
Instead, we make use of the feature classes V , B , X and learn the common representation { B h } h ∈
[ H ] by Empirical Risk Minimization (ERM).
For each hypothesis g ∈ G and h ∈ [ H ] , let ( v ( g ) 1: m,h , B ( g ) h , X ( g ) 1: m,h , (cid:101) g ) be the solution of the following optimization problem:
[Equation]
Notice that (cid:101) g ( g ) depends on g and in the rest of this paper we use (cid:101) g for simplicity of notation.
With features ( v ( g ) 1: m,h , B ( g ) h , X ( g ) 1: m,h ) , the conﬁdence set G ( t ) is then given as the collection of all possible hypothesis g such that the sum of squares of the bilinear forms across the available batch data is not greater than a pre-deﬁned parameter R 2 :
[Equation]
Due to the deﬁnition of ( v ( g ) 1: m,h , B ( g ) h , X ( g ) 1: m,h , (cid:101) g ) , the bilinear form
[Equation]
approximates (cid:104) W ∗ m,h ( g ) − W ∗ m,h ( g ∗ ) , X ∗ m,h ( g ( τ ) )
(cid:105) for all g ∈ G .
Therefore the conﬁdence set G ( t ) contains all the hypothesis in which the sum of (cid:0) (cid:104) W ∗ m,h ( g ) − W ∗ m,h
( g ∗ ) , X ∗ m,h
( g ( τ ) )
(cid:105) (cid:1) 2 over all history data and tasks is upper bounded by R 2 .
We will show that this can be used to quantify
the
Bellman error of the candidate action-value functions under the corresponding greedy policy, which then can be related to the sub-optimality of the corresponding greedy policy in the true environment.
The width parameter R 2 can be enlarged to handle the cases where realizability assumption is not strictly satisﬁed.
Since ERM is robust to misspeciﬁcation, optimism is still guaranteed and the analysis follows similarly.
The resulting algorithm would thus have sub-optimality depending additively on the misspeciﬁcation error.
This conﬁdence set captures the shared information B h ( · ) across tasks and contains all hypothesis with low Bellman errors within a smaller R 2 .
Algorithm 1 then samples trajectories according to the greedy policies of the chosen hypothesis g ( t ) and the estimation policy π est and augments the available batch data with D ( t ) m,h , h ∈
[ H ] , m ∈
[ M ] .
We note that the main computation workload of Algorithm 1 is the ERM step in Line 3.
This means that Algorithm 1 is oracle-efﬁcient with access to an ERM oracle.
In general, Line 3 pertains to a difﬁcult optimization problem over the set of candidate function approximators G ( t ) .
Although no efﬁcient algorithms are currently known to solve this problem, we note that for certain instances of the bilinear class (e.g., Yang et al. (2021)) where G is parameterized by variables in a Euclidean space, computationally efﬁcient gradient-based algorithms may exist to solve Line 3.
M AIN T HEORY
This section presents the theoretical analysis of Algorithm 1.
Without loss of generality, we assume that the feature classes are all bounded.
Deﬁnition 5.1.
Assume (cid:107) v m,h ( g ) (cid:107) 2 , (cid:107) B
h
( g ) (cid:107) F ≤ C W and (cid:107) X m,h ( g ) (cid:107) 2 ≤ C X hold for all v m,h ∈
V h , B h ∈ B h , X m,h ∈ X h and g ∈ G , h ∈
[ H ] , m ∈
[ M ] .
Here C W , C X ∈ R .
Next, it is important to consider the expressiveness of function classes V , B , X .
The following assumption is common in the theory of reinforcement learning Jin et al.
(2020b)
; Wang et al.
(2020b);
Jin et al.
(2021b);
Du et al.
Assumption 5.2 (Realizability) .
Assume g ∗ ∈ G and v ∗ 1: M, 1: H ∈ V , B ∗ 1: H ∈ B , X ∗ 1: M, 1: H ∈ X .
Now we present the theory for Algorithm 1.
Theorem 5.3.
Set T = 8 HMd log(1 + MC 2 X λ ) and
[Equation]
With probability 1 − δ the algorithm outputs a set of policies π g ( t 0) such that
[Equation]
The total number of trajectories used in Algorithm 1 is upper bounded by O ( MHTn 0 ) .
Therefore, with probability at least 1 − δ , Algorithm 1 is able to use
[Equation]
trajectories to ﬁnd a set of policies { π m } m ∈
[ M ] such that
[Equation]
Notice that if we use Bilin-UCB to learn each single task individually, then the total number of trajectories to achieve the above guarantee is
[Equation]
[Equation]
Indeed, to achieve Eq.
(6), the average sub-optimality of each task should be O ( (cid:15)/M ) .
Using Bilin-UCB, achieving V π ∗ m m, 1 ( s 1 ) − V π m m, 1 ( s 1 ) ≤ (cid:15)/M costs O
( H 6 M 2 d 2 log( |G| /δ ) /(cid:15) 2 ) samples for each task.
Thus the total number of trajectories is
[Equation]
via learning each task individually.
Therefore, Theorem 5.3 improves the sample complexity of learning Low-rank Multi-task Bilinear class given small sizes of expressive feature classes V , B , X , for example, log( |X||V||B| ) ≤ Md .
In general, the dependence on d can not be reduced since we estimate a d-by-k matrix for shared representation.
Furthermore, we believe that without further assumption, a polynomial improvement w.r.t M is impossible because the learner has to learn the task-speciﬁc features v ∗ m,h ∈ R k for all m ∈
[ M ] . For an important special case known as Linear Mixture Model (Deﬁnition 3.3), we have the following result via plugging |V| = |B| = 1 , |X| = |G| into Theorem 5.3.
Corollary 5.4.
Consider Low-rank Multi-task Bilinear class where each MDP M m is a linear mixture model with P m,h ( s, a, s (cid:48) )
= (cid:104) θ ∗ m,h , φ m ( s, a, s (cid:48) ) (cid:105) , r m,h ( s, a )
= (cid:104) θ ∗ m,h , ψ m ( s, a ) (cid:105) and there exist ν ∗ m,h ∈ R k , B ∗ h ∈ R d × k such that θ ∗ m,h
= B ∗ h ν ∗ m,h for all h ∈
[ H ] , m ∈
[ M ] .
Let
X
[Equation]
G
[Equation]
N
[Equation] k , (cid:15) [Equation] ) where N ( R k , (cid:15) ) denotes the (cid:15) -covering of R k and N ( R d × k , (cid:15) ) denotes the (cid:15) -covering of R d × k .
Under Assumption 5.2, there exists an algorithm that with probability at least 1 − δ ﬁnds a set of policies { π m } m ∈
[ M ] such that
[Equation]
using
[Equation]
trajectories.
Remark 5.5 .
Using Bilin-UCB to learn each task individually, it takes O (cid:0) H 6 M 3 d 3 log(1 /δ ) · (cid:15) − 2 (cid:1) trajectories to learn a set of policies { π m } m ∈
[ M ] such that
[Equation]
Therefore, Algorithm 1 achieves sample complexity improvement comparing to single-task learning.
Without further assumption, a polynomial improvement w.r.t M appears impossible because the learner has to learn the task-speciﬁc features v m,h ∈ R k for all m ∈
[ M ] with an Ω( M 3 k ) sample complexity in total.
The main beneﬁt of multi-task learning in this case is to reduce the dependence on the ambient dimension d to k .
T ECHNIQUE O VERVIEW
This section gives an overview of the analysis and the main techniques.
Owing to optimism principle and the construction of conﬁdence set via representation learning, the proof of Theorem 5.3 will crucially depend on the following three observations:
Risk bounds for representation learning.
Since the discrepancy measures serve as unbiased estimations of bilinear forms, i.e.
[Equation]
the solutions ( v ( g ) 1: m,h , B ( g ) h , X ( g ) 1: m,h , (cid:101) g ( g ) ) of ERM will be able to concentrate to the population mean, i.e.
[Equation]
This means that Algorithm 1 has approximately captured the correct bilinear forms.
Thus G ( t ) essentially ﬁnds all hypothesis g such that (cid:80) τ
( (cid:104) B ∗ h ( g
[Equation]
X
[Equation]
g
[Equation]
is small.
Regret decomposition associated to the bilinear forms.
One key property of Bilinear class is the upper bound of Bellman error as follow:
[Equation]
Furthermore, we know that the sub-optimality of the greedy policy of candidate hypothesis can be decomposed into the sum of Bellman errors across time steps.
Therefore, we show the following upper bound
[Equation]
Coverage condition of features X m,h ( · ) .
Owing to the previous two observations, the issue left is to upper bound |(cid:104)
B ∗ h ( g ) v ∗ m,h ( g ) − B ∗ h ( g ∗ ) v ∗ m,h ( g ∗ ) , X ∗ m,h ( g ) (cid:105)| for some g
= g ( t ) by the quantity (cid:80) τ ( (cid:104) B ∗ h ( g ) v ∗ m,h ( g ) − B ∗ h ( g ∗ ) v ∗ m,h ( g ∗ ) , X [Equation]
m,h ( g ( τ ) )
(cid:105) )
Using Hölder’s inequality, it sufﬁces that (cid:80)
[Equation] ( g [Equation] X
[Equation] (cid:62) possess sufﬁcient coverage condition of X ∗ m, h ( g ) for all m ∈
[ M ] .
Setting T
= Ω( Md ) , we will conﬁrm this via elliptical potential lemma.
The proof of Theorem 5.3 is then established based on the above three observations.
Since the second and third steps are natural extensions of the analysis of Bilin-UCB, the technical bulk is then to build risk bounds for representation learning.
Speciﬁcally, we deﬁne the following failure event:
Deﬁnition 6.1.
Deﬁne E as the event that there exist t ∈
[ T ] and h ∈
[ H ] such that
[Equation]
We will show that P [ E ] ≤ 1 − δ with the choice of R 2 in Algorithm 1.
The key step is embedding the estimation noise into low dimensional space R k and achieve improved concentration.
C ONCLUSION
This paper considers learning multiple RL problems jointly with representation learning.
A structural framework is proposed that permits generalization across tasks via general function approximation.
A sample efﬁcient algorithm is designed based on representation with ERM and optimistic principle where the conﬁdence sets are constructed based on learned features.
Theoretical analysis is displayed that the algorithm ﬁnds nearly-optimal policies within limited samples.
Several examples are discussed and sample complexity improvements are illustrated.
Given the success of representation learning in multi-task RL, it is an interesting future direction to study transfer learning for quickly adapting prior knowledge to a new, unseen task with limited data and computational power.