Provable Benefit of Multitask Representation Learning in Reinforcement Learning
Anonymous Author(s)
AffiliationAddress
Abstract
As representation learning becomes a powerful technique to reduce sample complexity in reinforcement learning (RL) in practice, theoretical understanding of itsadvantage is still limited.
In this paper, we theoretically characterize the benefitof representation learning under the low-rank Markov decision process (MDP)model.
We first study multitask low-rank RL (as upstream training), where all tasksshare a common representation, and propose a new multitask reward-free algorithmcalled REFUEL.
REFUEL learns both the transition kernel and the near-optimalpolicy for each task, and outputs a well-learned representation for downstreamtasks.
Our result demonstrates that multitask representation learning is provablymore sample-efficient than learning each task individually, as long as the totalnumber of tasks is above a certain threshold.
We then study the downstream offline
RL, where the agent is given a new task sharing the same representation as theupstream tasks and an offline dataset, and aims to find a near-optimal policy.
Wedevelop a sample-efficient algorithm with the suboptimality gap bounded by theestimation error of the learned representation in the upstream plus a vanishingterm that decreases as the number of offline samples becomes large.
Our resultfurther captures the benefit of employing the learned representation from upstreamtraining as opposed to learning the representation of the low-rank model directly.
To the best of our knowledge, this is the first theoretical study that characterizes thebenefit of representation learning in exploration-based reward-free multitask RL.
Introduction
Multitask representation learning has arisen as a popular and powerful framework to deal withdata-scarce environment in supervised learning, bandits and reinforcement learning (
RL) recently(Maurer et al., 2016;
Du et al., 2020; Tripuraneni et al., 2021;
Yang et al., 2021; Cella et al., 2022;
Lu et al., 2021; Arora et al., 2020; Hu et al., 2021).
The general multitask representation learningframework consists of both upstream multiple source tasks learning and downstream target tasklearning.
In RL applications, in upstream, the agent is assigned with T source tasks sharing acommon representation and the goal is to learn a near-optimal policy for each task via learningthe representation.
In downstream, with the help of the learned representation, the agent aims tofind a near-optimal policy of a new task that shares the same representation as the source tasks.
While representation learning has achieved great success in supervised learning (Du et al., 2020;
Tripuraneni et al., 2021; Maurer et al., 2016; Kong et al., 2020) and multi-armed bandits (MAB)problems (Yang et al., 2021;
Qin et al., 2022; Cella et al., 2022), most works in multitask RL mainlyfocus on empirical algorithms (Sodhani et al., 2021; Arulkumaran et al., 2022; Teh et al., 2017) withlimited theoretical works (Arora et al., 2020;
Hu et al., 2021; Brunskill and Li, 2013; Müller and
Pacchiano, 2022; Calandriello et al., 2015; Lu et al., 2021; D’Eramo et al., 2020).
Generally speaking,there are two main challenges for multitask representation learning in RL.
First, differently from supervised learning with static datasets, in MAB and RL, representation learning and data collectionare intertwined with each other.
The agent has to choose an exploration policy in each iteration basedon the representation learned in the past.
Such iterative process introduces temporal dependence tothe collected data.
Note that none of the aforementioned theoretical multitask RL studies took theinfluence of sequential exploration and temporal dependence in data into consideration.
Second, datadistribution mismatch between the target task and the source tasks may occur.
Additional assumptionis usually required in order to ensure the representation learned from sources task will benefit thetarget task (Lu et al., 2021).
In this work, we aim to unveil the fundamental impact of multitask representation learning on RL byanswering the following two open questions:
• First, in the upstream, is it possible to design exploration-based multitask representation learningalgorithms for RL to achieve efficiency gain compared with single-task RL after taking theinteractive data collection process and temporal data dependency into account?• Second, can the learned representation from upstream bring efficiency gain to the downstreamlearning?
Main contributions: We focus on multitask representation learning under low-rank MDPs (Agarwalet al., 2020), where the transition kernel admits a low-rank decomposition that naturally involvesrepresentation learning.
It answers the aforementioned questions affirmatively, as summarized below
• First, we propose a new multitask representation learning algorithm called REFUEL under lowrank MDPs in upstream.
REFUEL features joint modeling learning, pseudo cumulative valuefunction construction for exploration, joint exploration policy learning and a unique terminationcriterion.
These novel designs enable REFUEL to leverage the shared representation acrosssource tasks to improve the learning efficiency.
• We show that REFUEL can 1) find a near-optimal policy and 2) learn a near-accurate modelfor each of T source tasks to certain average accuracy ϵ u , with a sample complexity of
[Equation]
O
[Equation]
, where H, K, d denote the horizon steps,the cardinality of action space and the low-rank dimension, respectively.
Compared with stateof-the-art single-task representation in low-rank MDPs (Uehara et al., 2022), REFUEL achievesreduced sample complexity as long as T = Ω (cid:0) Kd 5 (cid:1) , which shows the benefit of multitask learningprovably.
To the best of our knowledge, this is the first theoretical result demonstrating thebenefit of representation learning in exploration-based reward-free multitask RL.• Finally, when the learned representation from REFUEL is leveraged for downstream offlinelearning, we propose a new sample efficient algorithm with suboptimality gap of (cid:101) O
( H √ dξ down +Hd/ √ N down ) , where N down is the number of samples of offline dataset.
The first term capturesthe approximation error ξ down that measures the quality of the learned representation fromupstream, and is negligible when the upstream learning is sufficiently accurate, and the secondterm improves the dependency on H and d compared with the state-of-the-art (Uehara et al. ,2022) without knowledge of the upstream representation, showing the benefit of employing thelearned representation from upstream.
Notations: Throughout the paper, we use [ N ] to denote set { 1 , . . .
, N } for N ∈ N , use ∥ x ∥ 2 todenote the ℓ 2 norm of vector x , use ∥ X ∥ F to denote the Frobenius norm of matrix X , use △ ( A ) todenote the probability simplex over set A , and use U ( A ) to denote uniform sampling over A , given|A| < ∞ .
Furthermore, for any symmetric positive definite matrix Σ , we let ∥ x ∥ Σ :=
√ x ⊤ Σ x .
Related work
Multitask bandits, multitask RL and meta-RL: The benefit of multitask learning in linear banditshas been investigated in Yang et al.
; Qin et al. (2022); Cella et al. (2022); Azizi et al.
Deshmukh et al.
; Cella and Pontil (2021); Hu et al.
For multitask RL, Arora et al.(2020) showed that representation learning can reduce sample complexity for imitation learning.
Calandriello et al.
(2015) studied multitask learning under linear MDPs with known representation.
Hu et al.
(2021) studied multi-task RL with low inherent Bellman error (Zanette et al., 2020) withknown representation.
Assuming all tasks are sampled from an unknown distribution, Brunskilland Li (2013) studied the benefit of multitask RL when each task is sampled independently from a distribution over a finite set of MDPs while Müller and Pacchiano (2022) focused on meta-RL forlinear mixture model
(Yang and Wang, 2020).
Considering all tasks share a common representation,
D’Eramo et al.
(2020) showed the benefit in the convergence rate on value iteration and Lu et al.(2021) proved the sample efficiency gain of multitask RL under low-rank MDPs.
None of the aboveworks took the influence of sequential exploration and temporal dependence in data into consideration,where Lu et al.
(2021) assumed a generative model is accessible and D’Eramo et al.
(2020) assumedthat the dataset is given in prior.
In our study, we design reward-free exploration for multitask RL.
Single task RL under low-rank MDPs: Low-rank MDPs are first studied byAgarwal et al.
(2020),and then Uehara et al.
(2022) studied low-rank MDPs for both online and offline RL.
Modi et al.(2021) proposed a model-free algorithm MOFFLE for reward-free RL under low-nonnegative-rank
MDPs.
Du et al.
; Misra et al.
; Zhang et al.
(2022) studied block MDP which is a specialcase of low-rank MDPs.
Offine RL under linear and low-rank MDPs: Most of the existing works on offline RL focus onsingle-task learning, among which Jin et al.
(2021) studied offline RL under linear MDPs and proposedan algorithm called PEVI with provable suboptimality gap, Uehara et al.
(2022) studied offline RLunder low-rank MDPs and proposed an algorithm called REP-LCB that leverages pessimism tolearn under a partial coverage condition, and Yin et al.
(2022) studied offline RL leveraging varianceinformation.
Problem formulation
3.1 Episodic MDP
Consider an episodic non-stationary Markov Decision Process (MDP) with finite horizon H , denotedby M =
( S , A , H, P, r ) , where S is the state space (possibly infinite), A is the finite action spacewith cardinality K , H is the number of steps in each episode, P = { P h } h ∈
[ H ] is the collection oftransition measures with P h : S × A (cid:55)→ △ ( S ) , r = { r h } h ∈
[ H ] is the collection of deterministicreward functions with r h : S ×A (cid:55)→ [0 , 1] .
The initial state s 1 is assumed to be fixed for each episodeand the sum of the reward is normalized with (cid:80) Hh =1 r h ∈ [0 , 1] .
Denote policy π = { π h } h ∈ [ H ] as a sequence of mappings where π h : S (cid:55)→ △ ( A ) .
We useπ h ( a | s ) to denote the probability of selecting action a in state s at step h .
Given a starting states h , we use s h ′ ∼ ( P, π ) to denote a state sampled by executing policy π under the transitionmodel P for h ′ − h steps and E ( s h ,a h ) ∼ ( P,π ) [ · ] to denote the expectation over states s h ∼ ( P, π )and actions a h ∼ π .
Given policy π and a state s at step h , the (state) value function is definedas V πh,P,r ( s ) :
=
E ( s h ′ ,a h ′ )
∼ ( P,π ) (cid:104)(cid:80) Hh ′ =
h r h ′
( s h ′ , a h ′ )
(cid:12)(cid:12)(cid:12) s h
= s (cid:105) .
Similarly, the (state-)action valuefunction for a given state-action pair ( s, a ) under policy π at step h is defined as Q πh,P,r ( s, a ) =
r h ( s, a )+
( P h V πh +1 ,P,r )( s, a ) , where ( P h f )( s, a ) =
E s ′ ∼ P h ( ·| s,a ) [ f ( s ′ )] for any function f : S (cid:55)→ R .
For simplicity, we omit subscript h for h = 1 .
Since the MDP begins with the same initial state s 1 , tosimplify the notation, we use V πP,r to denote V π 1 ,P,r ( s 1 ) without causing ambiguity in the following.
Let P ∗ denote the transition kernel of the underlying environment.
Given a reward function r ,there always exists an optimal policy π ∗ that yields the optimal value function V π ∗ P ∗ ,r = sup π V πP ∗ ,r ,abbreviated as V ∗ P ∗ ,r .
This paper focuses on low-rank MDPs (Jiang et al., 2017; Agarwal et al., 2020) defined as follows.
Definition 1 (Low-rank MDPs) .
A transition kernel P ∗ h : S ×
A → △ ( A ) admits a low-rankdecomposition with dimension d ∈ N if there exists two embedding functions ϕ ∗ h : S
×
A → R d andµ ∗ h : S → R d such that
[Equation]
Without loss of generality, we assume ∥ ϕ ∗ h ( s, a ) ∥ 2 ≤ 1 for all ( s, a ) ∈ S × A and for any function g :
S (cid:55)→ [0 , 1] , (cid:13)(cid:13)(cid:82) µ ∗ h ( s ) g ( s ) ds
(cid:13)(cid:13) 2 ≤ √ d .
An MDP is a low-rank MDP with dimension d if for h ∈ [ H ] ,its transition kernel P ∗ h admits a low rank decomposition with dimension d .
Let ϕ ∗ =
{ ϕ ∗ h } h ∈ [ H ] andµ ∗ = { µ ∗ h } h ∈
[ H ] be the embeddings for P ∗ , where ϕ ∗ is also called representation in RL literature.
3.2 Reward-free multitask RL and its downstream learning
We consider a multitask RL consisting of reward-
In reward-free multitask upstream learning, the agent is expected to conduct reward-free exploration over T source tasks, where all reward functions { r t } t ∈
[ T ] are unknown.
Each taskt ∈ [ T ] is associated with a low-rank MDP M t =
( S , A , H, P t , r t ) .
All T tasks are identicalexcept for (a) their transition model P t , which admits a low-rank decomposition with dimension d :
P t ( s ′ | s, a ) =
⟨ ϕ ∗ ( s ′ ) , µ ( ∗ ,t ) ( s, a ) ⟩ for all t ∈
[ T ] and (b) their reward r t .
Namely, all tasks sharethe same feature function ϕ ∗ , but may differ in µ ( ∗ ,t ) as well as reward r t .
The goal of upstreamlearning is to find a near-optimal policy and a near-accurate model for any task t ∈
[ T ] and any rewardfunction r t , t ∈ [ T ] via sufficient reward-free exploration, and output a well-learned representationfor downstream task.
For each t ∈ [ T ] , we use P ( ∗ ,t ) to denote the true transition kernel of task t .
In downstream learning, the agent is assigned with a new target task T + 1 associated with a low-rank
MDP M T +1 =
( S , A , H, P T +1 , r T +1 ) .
The task shares the same S , A , H and ϕ ∗ with the Tupstream tasks, but has a task-specific µ ( ∗ ,T +1) .
We also assume that the reward for the new taskr T +1 is given to the agent.
Here, the agent is assumed to take the learned representation function (cid:98) ϕduring the upstream training, and then interacts with the new task environment M T +1 in an offlinemanner to find a near-optimal policy for the new task.
By utilizing the representation learned in theupstream, the agent is expected to achieve better sample efficiency in the downstream.
Finally, weuse P ( ∗ ,T +1) to denote the true transition kernel of task T + 1 .
Upstream reward-free multitask RL
In this section, we present our proposed algorithm for upstream reward-free multitask RL in low-rank
MDPs and characterize its theoretical performance.
Since it is impossible to learn a model in polynomial time if there is no assumption on features ϕ hand µ h , we first adopt the following conventional assumption (Agarwal et al., 2020).
Assumption 1.
(Realizability).
A learning agent can access to a model class { (Φ , Ψ) } that containsthe true model, namely, for any h ∈
[ H ] , t ∈
[ T ] , the embeddings ϕ ∗ h ∈
Φ , µ ( ∗ ,t ) h ∈ Ψ .
While we assume the cardinality of function classes to be finite for simplicity, extensions to infiniteclasses with bounded statistical complexity are not difficult (Sun et al., 2019; Agarwal et al., 2020).
4.1 Algorithm design
We first describe our proposed algorithm REFUEL depicted in Algorithm 1.
In each iteration n , for each task t and time step h , the agent executes the exploration policy π n − 1 tfollowed by two uniformly chosen actions to collect trajectories for H episodes.
Note that π n − 1 tis designed at the previous iteration.
Then, to utilize the common feature of data collected fromdifferent source tasks, the agent processes the data jointly as following steps.
Joint MLE-based Modeling Learning: The agent passes all previously collected data toestimate the low-rank components (cid:98) ϕ ( n ) h , (cid:98) µ ( n, 1) h , . . .
, (cid:98) µ ( n,t ) h simultaneously via the MLE oracle
MLE (cid:16)(cid:83) t ∈
[ T ] D ( n,t ) h (cid:17) (line 9) on the joint distribution defined as follows:
[Equation]
The MLE oracle above generalizes the MLE oracle commonly adopted for single-task RL in the 174 literature (Agarwal et al., 2020; Uehara et al., 2022).
In practice, MLE oracle can be efficientlyimplemented if the model classes Φ , Ψ are properly parameterized such as by neural networks.
Next, for each task t , the agent uses the learned embeddings (cid:98) ϕ ( n ) h , (cid:98) µ ( n,t ) h to update the estimatedtransition kernel (cid:98) P ( n,t ) at each step h as
[Equation]
Pseudo Cumulative Value Function Construction for Exploration: The agent first uses therepresentation estimator (cid:98) ϕ ( n ) h to update the empirical covariance matrix (cid:98) U ( n,t ) h as
[Equation]
to
− uniformly choose t ) h − h and enter and collect data
), a , s ( n,t,
) h , a ( n,t,
) h
Add s dataset h =
D ( n − 1 ,t ) h ∪
Learn (cid:16)(cid:98) ϕ
( n ) h , (cid:98) µ ( n
1) h , (cid:98) µ ( n T h (cid:17) (cid:16)(cid:83) t t h (cid:17)
, T do 11:
Update estimated transition kernels (cid:98) P ( n,t ) h ( ·|· , · ) as Equation (2), empirical covariance matrix (cid:98) U ( n,t ) h as Equation (3) and exploration-driven reward (cid:98) b ( n,t ) ( · , · ) as Equation (4).
12: end for 13: end for 14:
Set Pseudo Cumulative Value Function PCV (cid:16) (cid:98) P ( n,t ) , (cid:98) b ( n,t ) h , π t ; T (cid:17) as Equation (5).
Get exploration policy π n 1 , . .
, π nT = arg max π 1 ,...,π T PCV (cid:16) (cid:98) P ( n,t ) , (cid:98) b ( n,t ) h , π t ; T (cid:17) .
16: if 2 PCV (cid:16) (cid:98) P ( n,t ) , (cid:98) b ( n,t ) h , π t ; T (cid:17) + 2 √ KTζ n ≤ Tϵ u then 17:
Terminate and set n u = n and output (cid:98) ϕ
= (cid:98) ϕ
( n u ) , (cid:98) P (1) = (cid:98) P ( n u , 1) , . . .
, (cid:98) P ( T ) =
(cid:98) P ( n u ,T ) .
18: end if 19: end for 20: Output: (cid:98) ϕ
= (cid:98) ϕ
( n u ) , (cid:98) P (1) = (cid:98) P ( n u , 1) , . . .
, (cid:98) P ( T ) =
(cid:98) P ( n u ,T ) .
where { s ( τ,t,h +1) h , a ( τ,t,h +1) h , s ( τ,t,h +1) h + } is the tuple collected at the τ -th iteration, ( h +1) -th episode,and step h .
Then, the agent uses both (cid:98) ϕ ( n ) h and (cid:98) U ( n,t ) h to provide an exploration-driven reward as
[Equation]
where B and (cid:101) α n are pre-determined parameters.
To integrate information of learned models andexploration-driven rewards for T tasks, we define a Pseudo Cumulative Value Function (PCV) as
[Equation]
Joint Policy Learning for Exploration: The agent then learns T policies via optimization overthe PCV, and uses these policies as the exploration policy for each task in the next iteration.
Since
PCV measures the cumulative model estimation error over all tasks, these policies will explore thestate-action space where the model estimation has large uncertainty on average of T tasks.
Here weadopt an optimization oracle for line 15 in Algorithm 1.
Termination: REFUEL is equipped with a termination condition and outputs the current estimatedmodel and representation during iterations if the PCV plus certain minor term is below a threshold,which suffices to guarantee that for each task t , the returned model (cid:98) P ( t ) from REFUEL reaches anaverage proximity to the true model P ( ∗ ,t ) .
The joint learning and decision making among all tasks involved in the four main components differentiate REFUEL from single-task RL algorithms, and enable it to exploit the shared representationacross the source tasks to improve the sample efficiency.
4.2 Theoretical results on sample complexity
In this section, we first establish the sample complexity upper bound for REFUEL and then compareour multitask result with that of single-task RL.
Theorem 1.
For any fixed δ ∈ (0 , | Ψ | min { T, − Kd 2 } ) , set λ n =
O ( d log( | Φ | nTH/δ )) , ζ n =
O (log (cid:0) | Φ || Ψ | T nH/δ (cid:1) /n ) , (cid:101)
α n = O ( (cid:112) K log( | Φ || Ψ | T nH/δ ) + λ n dT ) and B = O ( (cid:112) T + K/d 2 ) .
Let (cid:98) ϕ, (cid:98) P (1) , . . .
, (cid:98) P ( T ) be the output of Algorithm 1.
Then, under Assumption 1, with probability atleast 1 − δ , for any policy π t of task t ∈ [ T ] and any step h ∈
[ H ] , we have
[Equation]
≤ [Equation] u
[Equation]
Further, for any task t , given any reward r t , let (cid:98) π t = arg max π V π (cid:98) P ( t ) ,r t .
Then with probability atleast 1 − δ , 1 T (cid:80) Tt =
[ V ∗ P ( ∗ ,t ) ,r t − V (cid:98) π t P ( ∗ ,t ) ,r t ] ≤ ϵ u .
Meanwhile, the number of trajectories collected by each task is at most
[Equation]
Equation (6) in Theorem 1 indicates that the estimated transition kernels of the T MDPs meet therequired accuracy on average.
Theorem 1 further indicates that with the data collected during suchreward-free exploration, if every task is further given any reward function, near-optimal policies canbe found for all tasks on average as well.
For the sample complexity bound in Equation (7), the first term is inversely proportional to T , whichcaptures sample-benefit for each task due to learning the T models jointly.
The second term is relatedto the number of samples to guarantee the concentration of empirical covariance matrix (cid:98) U ( n,t ) h inorder to identify desirable exploration policies.
Such concentration needs to be satisfied for eachtask t independently and can not be implemented jointly, which makes the second term independentwith T .
The third term originally arises from the model estimation error shift among T tasks whenimplementing joint MLE-based model estimation.
In the algorithm, this error shift is then transferredto B defined in Equation (4), which provides an explicit upper bound for the exploration-drivenreward (cid:98) b ( n,t ) h .
Note that (cid:98) b ( n,t ) h consists the PCV to guide exploration, which finally contributes toconvergence of REFUEL and affects its sample complexity.
Compared with the state-of-the-art sample complexity that scales in (cid:101) O (cid:16) H 5 d 4 K 2 ϵ 2 u (cid:17) 1 when performing
T single-task RL independently (Uehara et al., 2022), REFUEL achieves lower sample complexitywhen T
= Ω (cid:0)
Kd 5 (cid:1) , and it reduces the single-task sample complexity in (Uehara et al., 2022) by anorder of O (cid:16) min { H 2 d 2 T, H 2 K, d 3 , dK, d 5 TK } (cid:17) .
Technically, the joint learning feature of REFUEL requires new analytical techniques.
We outline themain steps of the proof of Theorem 1 here and defer the detailed proof to Appendix A. Step 1: Wedevelop a new upper bound on model estimation error for each task, which captures the advantageof joint MLE model estimation over single-task learning, as shown in Proposition 1.
Step 2: Weestablish the PCV as an uncertainty measure that captures the difference between the estimated andground truth models.
This justifies using such a PCV as a guidance for further exploration, as shownin Proposition 2.
Step 3: We show that the summation of the PCVs over all iterations is sublinearwith respect to the total number N u of iterations, as shown in Proposition 3, which further impliespolynomial efficiency of REFUEL in learning the models.
Downstream offline RL
In downstream RL, the agent is given a new target task T + 1 under MDP M T +1 and aims to finda near-optimal policy for it.
The downstream task is assumed to share the same representation asthe upstream tasks, and hence the agent can adopt the representation learned from the upstream toexpedite its learning.
Since the agent does not know the task-specific embedding function µ T +1 , sheis allowed to have access to an offline dataset D down =
( s τ h , a τ h , r τ h , s τ h +1 )
} N down ,H τ,h = , which is rolledout from some behavior policy ρ .
In this section, we provide an algorithm for downstream offline RL, introduce the connectionsbetween upstream and downstream MDPs, and characterize the suboptimality gap of the algorithm.
5.1 Algorithm for downstream offline RL
We present our pessimistic value iteration algorithm for downstream offline RL in Algorithm 2.
Although our algorithm shares similar design principles as traditional algorithms for linear MDPs (Jinet al., 2021), it differs from them significantly as described in the following, due to the misspecificationof representation from upstream and the general rather than linear reward function adopted.
For easeof exposition, we define Bellman operator B h as ( B h f )( s, a ) =
r h ( s, a ) + ( P ( ∗ ,T +1) h f )( s, a ) forany f : S × A (cid:55)→ R .
The main body of the algorithm consists of a backward iteration over steps.
Ineach iteration h , the agent executes the following main steps.
6: (cid:98) V h ( · ) = (cid:98) Q h ( · , (cid:98) π h ( · )) , where (cid:98) π h ( · ) =
arg max π h (cid:98) Q h ( · , (cid:98) π h ( · )) .
7: end for 8: Output: { (cid:98) π h } Hh =1 .
Approximate Bellman update with general reward function.
We construct an estimated Bellmanupdate (cid:98) B h (cid:98) V h +1 based on the dataset D down to approximate B h (cid:98) V h
+1 .
Here (cid:98) V h
+1 is an estimated valuefunction constructed in the previous iteration based on D down .
Since our reward function is moregeneral and not necessarily linear in the feature, we choose to parameterized P ( ∗ ,T +1)
h
V ∗
h
+1 insteadof Q ∗ h ( s, a ) =
r h ( s, a ) + ( P ( ∗ ,T +1)
h
V ∗ h +1 )
( s, a ) due to the linear structure of P ( ∗ ,T +1) h .
Althoughthe ground truth representation is unknown, the agent is able to use the estimated representationlearned from the upstream learning to parameterize the linear term ( P ( ∗ ,T +1)
h
V ∗ h +1 )
( s, a ) .
Then, thecoefficient w ∗ h can be estimated by solving the following regularized least squares problem:
[Equation]
which has the closed-from solution (cid:98) w h = Λ − 1 h (cid:80) N
[Equation] [Equation]
ϕ h
[Equation]
s τh
[Equation] a τh ) [Equation] V h
+
( s τh +1 ) with Λ h = λI + 259 (cid:80) N down τ = (cid:98)
ϕ
h ( s τh , a τh ) (cid:98) ϕ
h ( s τh , a τh ) ⊤ .
We obtain an estimate ( (cid:98) B h (cid:98) V h +1 )( · , · ) =
r h ( · , · ) + (cid:98) w ⊤ h (cid:98) ϕ
h ( · , · ) .
Design pessimism for incorporating upstream misspecification of representation.
Pessimismis realized by subtracting an uncertainty metric Γ h from (cid:98) Q h (line 5).
The uncertainty metric Γ hquantifies the Bellman update error | ( (cid:98) B h (cid:98) V h
+1 − B h (cid:98) V h +1 )
( s, a ) | at step h .
Our selection of Γ h isin the form of ξ down + β ( (cid:98) ϕ
h Λ − 1 h (cid:98)
ϕ h ) 1 / 2 , which captures both the approximation error (first term)and the estimation error in the Bellman update (second term).
Both of the above errors are affectedby the misspecification of representation from upstream learning.
Intuitively, m :=
( (cid:98) ϕ
h Λ − 1 h (cid:98) ϕ h ) − 1represents the effective number of samples the agent has explored along the (cid:98) ϕ direction, and thepenalty term β/ √ m represents the uncertainty along the (cid:98)
ϕ direction.
By choosing a proper valuefor β , we can prove that with high probability, the uncertainty metric Γ h always upper bounds the
Bellman update error.
Select greedy policy.
Line 6 executes a greedy policy (cid:98) π h to maximize (cid:98) Q h , which will be the outputof the algorithm after the backward iteration over steps completes.
5.2 Connections between upstream and downstream MDPs
In downstream RL tasks, the agent will adopt the feature (cid:98) ϕ learned from upstream Algorithm 1.
Toensure that the feature learned in upstream can still work well in downstream tasks, upstream anddownstream MDPs should have certain connections.
We next elaborate some reasonable assumptionson transition kernels to build such connections.
The following reachability assumption is common in relevant literature (Modi et al., 2021).
Assumption 2 (Reachability) .
For each source task t in upstream, whose transition kernel is P ( ∗ ,t ) ,there exists a policy π 0 t such that min s ∈S P ( π 0 t ,t ) h ( s ) ≥ κ u , where P ( π 0 t ,t ) h ( · ) : S → R is the densityfunction over S using policy π 0 t to roll into state s at timestep h .
Assumption 3.
The state space S is compact. Combining Assumption 3 with Assumption 2, it can be shown that there exists a uniform distributionon S has the density function f ( s ) =
υ , where 1 /υ is the measure of S .
Then, we denote U ( S ) asthe uniform distribution over S and U ( S , A ) as the uniform distribution over S × A .
Assumption 4.
For any two different models in the model class { (Φ , Ψ) } , say P 1 ( s ′ | s, a ) = 286 (cid:10) ϕ[-
-]
1 ( s, a ) , µ 1 ( s ′ ) (cid:11) and P 2 ( s ′ | s, a ) =
(cid:10) ϕ
( s, a ) , µ 2 ( s ′ ) (cid:11) , there exists a constant C R such that for all( s, a ) ∈ S × A and h ∈ [ H ] ,
[Equation]
[Equation]
For normalization, we assume that for any ϕ ∈ Φ , ∥ ϕ ( s, a ) ∥ 2 ≤ 1 and for any µ ∈ Ψ and anyfunction g : S → [0 , 1] , (cid:13)(cid:13)(cid:82) [Equation]
d . Assumption 4 ensures that for each source task t any ( s, a ) ∈ S × A , the total variation distancebetween the learned (cid:98) P ( t ) ( ·| s, a ) and true transition kernels (cid:98) P ( ∗ ,t ) ( ·| s, a ) will not be too large whenthe expectation of the total variation distance over the entire state action space is small.
Assumption 5.
For the underlying MDP of task T + 1 , we assume that the reward function r T +1 ischosen from a function class R with ε -covering number N R ( ε ) for any given ε , and N R ( ε ) =
(cid:101) O (cid:0) 1 ε p (cid:1) .
The true transition kernel P ( ∗ ,T +1) can be ξ -approximated by a linear combination of T source tasks,i.e. there exist T (unknown) coefficients c 1 , . . .
, c T ∈ [0 , C L ] such that
[Equation]
Assumption 5 establishes the underlying connection between upstream source tasks and the downstream target task.
The assumption on N ε ( R ) is more general than the common assumption onreward under linear MDP and can be easily achieved in practice.
For example, if the reward islinear with respect to an unknown feature (cid:101) ϕ : S ×
A → R p , i.e., there exists θ ∈ R p and ∥ θ ∥ ≤ R ,r ( s, a ) = ⟨ (cid:101) ϕ
( s, a ) , θ ⟩ , then N ε ( R ) =
(cid:101) O (cid:16)(cid:0)
+ 2 Rε (cid:1) p (cid:17) (see Lemma D.5 in Jin et al.
, whichsatisfies the assumption.
We remark here (cid:101) ϕ is not necessarily the same as ϕ ∗ , even its dimension pcan be different from d .
Finally, we make the following coverage assumption for the estimated feature, which can be easilychecked in practice.
Assumption 6 (Feature coverage) .
There exists absolute constant κ ρ such that
[Equation]
5.3 Sample complexity for downstream offline RL
Since the downstream task applies the estimated feature from the upstream, the downstream performance is highly affected by the upstream estimation accuracy of the feature.
Hence, we first introducethe definition of ϵ -approximate linear MDP, and provide a guarantee for the estimated feature.
Definition 2 ( ϵ -Approximate linear MDP) .
For any ϵ > 0 , MDP M = ( S , A , H, P, r ) is an ϵapproximate linear MDP with a time-dependent feature map ϕ : S ×
A → R d if, for any h ∈ [ H ] ,there exist a time-dependent unknown (signed) measures µ over S such that for any ( s, a ) ∈ S × A ,we have
[Equation]
Any ϕ satisfying Equation (11) is called an ϵ -approximate feature map for M .
Let ξ down = ξ + C L C R
Tυϵ u κ u .
The following lemma shows that the feature (cid:98) ϕ learned in upstream is aξ down -approximate feature map and can approximate the true feature in the new task.
Lemma 1.
Under Assumptions 2 to 5, the output of Algorithm 1 (cid:98) ϕ is a ξ down -approximate featurefor MDP M T +1 =
( S , A , H, P ( ∗ ,T +1) , r, s 1 ) , i.e. there exist a time-dependent unknown (signed)measures (cid:98) µ ∗ over S such that for any ( s, a ) ∈ S × A , we have
[Equation]
Furthermore, for any g : S → [0 , 1] , (cid:13)(cid:13)(cid:82) (cid:98) µ ∗ h ( s ) g ( s ) ds
(cid:13)(cid:13)
[Equation]
d .
Now we are ready to provide our main result in the following theorem.
Theorem 2.
Under Assumptions 2 to 5, for any δ ∈ (0 , 1) , if we set λ = 1 and β =
O (cid:0) d √ ι + √ dN down ξ down + √ p log N down (cid:1) , where ι = O (log ( pdHN down max { ξ down , 1 } /δ )) ,then with probability at least 1 − δ , the suboptimality gap of Algorithm 2 is at most
[Equation]
If additionally Assumption 6 holds, sample size satisfies N down ≥ 40 κ ρ · log(4 dH/δ ) , then the 327 suboptimality gap is at most
O (cid:18)
[Equation]
To the best of our knowledge, Theorem 2 provides the first suboptimality gap with approximationerror in representation.
Our first result characterizes the suboptimality gap if offline data is rolled outfrom a behavior policy 2 .
If the offline data satisfies the feature coverage assumption, we are able tofurther upper bound the suboptimality gap by three terms as stated in Theorem 2.
We note that Equation (14) can be simplified as (cid:101) O (cid:16) Hd 1 / 2 ξ down + N − 1 / 2 down H max { d, √ p } (cid:17) , wherethe first term is a constant error that arises from the approximation error in Bellman update.
Thesecond term diminishes as N down grows to infinity, and it arises from the estimation error of (cid:98) w hin regularized linear regression.
The dependence on d and √ p arises from the model complexity oflinear function approximation class when we seek a concentration of the estimated value function.
Such an suboptimality also demonstrates that our downstream offline RL benefits from the learned representation from upstream.
Namely, as long as the estimated representation in upstream approximatesthe ground truth well (i.e., ξ down is small), then the remaining term (cid:101) O (cid:16) N − 1 / 2 down H max { d, √ p } (cid:17) ofthe suboptimality gap improves (cid:101)
O
( H 2 d 2 ) 3 of REP-LCB (Uehara et al., 2022) under low-rank MDPwith unknown representation, with better dependence on H and d due to the benefit of upstreamrepresentation learning.
Further, compared to standard linear MDP with the knowledge of ground truth representation Jin et al.(2021), our suboptimality gap contains an additional upstream misspecification error (i.e., ξ down ).
The remaining term yields an order of (cid:101) O ( Hd ) which matches that in Jin et al.
(2021) if we specializeour reward to be linear with p =
O ( d ) .
Conclusion
In this paper, we show a shared representation of low rank MDPs can be learned efficiently andtransferred.
We provide upper bound on both the sample complexity of upstream reward-freemultitask RL and the suboptimality gap of downstream offline RL with the help of the learnedrepresentation from upstream.
Our results show that multitask representation learning is provablymore sample efficient than learning each task individually and capture the benefit of employingthe learned representation from upstream to learn a near-optimal policy of a new similar task indownstream.
As future work, we note that our analysis of the downstream offline RL depends onsome assumptions that connect between the upstream and downstream MDPs.
It is interesting toexplore whether these assumptions can be further relaxed.