Provable General Function Class Representation Learning in Multitask Bandits and MDP
Authors: Rui Lu
Andrew Zhao
Simon Shaolei Du
Gao Huang
Venue: NeurIPS 2022 Accept (2022)
Keywords: reinforcement learning
multi-task
representation learning
theory

Abstract:  While multitask representation learning has become a popular approach in reinforcement learning (RL) to boost the sample efficiency, the theoretical understanding of why and how it works is still limited. Most previous analytical works could only assume that the representation function is already known to the agent or from linear function class, since analyzing general function class representation encounters non-trivial technical obstacles such as generalization guarantee, formulation of confidence bound in abstract function space, etc. However, linear-case analysis heavily relies on the particularity of linear function class, while real-world practice usually adopts general non-linear representation functions like neural networks. This significantly reduces its applicability. In this work, we extend the analysis to general function class representations. Specifically, we consider an agent playing $M$ contextual bandits (or MDPs) concurrently and extracting a shared representation function $\phi$ from a specific function class $\Phi$ using our proposed Generalized Functional Upper Confidence Bound algorithm (GFUCB). We theoretically validate the benefit of multitask representation learning within general function class for bandits and linear MDP for the first time. Lastly, we conduct experiments to demonstrate the effectiveness of our algorithm with neural net representation.

Provable General Function Class Representation Learning in Multitask Bandits and MDPs
Rui Lu', Andrew Zhao', Simon S. Du?, Gao Huang'
Deparrment of Automation, BNRist, Tsinghua University
Paul G. Allen School of Computer Science and Engineering, University of Washington
{r-1u21, , 2 @mails.... tsinghua. edu. ssduêcs washington. com, gaohuangOtsinghua.edi cn
Abstract
While multitask representation learning has become a popular approach in reinforce- ment learning (RL) to boost the sample efficiency, theoretical understanding why and how it works is still limited. Most previous analytical works could only assume that the representation function is already known the agent or from linear function class, since analyzing general function class representation encounters non-trivial technical obstacles such as generalization guarantee, formulation confidence bound in abstract function space, etc. However, linear-case analysis heavily relies on the particularity of linear function class, while real-world practice usually adopts general non-linear representation functions like neural networks. This significantly reduces its applicability. In this work, we extend the analysis to general function class representations. Specifically, consider an agent playing M contextual bandits (or MDPs) concurrently and extracting a shared representa- tion function é from specific function class & using our proposed Generalized Functional Upper Confidence Bound algorithm (GFUCB). We theoretically vali- date the benefit of multitask representation learning within general function class for bandits and linear MDP for the first time. Lastly, we conduct experiments to demonstrate the effectiveness of our algorithm with neural net representation.
Introduction
Recently, reinforcement learning (RL) has achieved many successful applications in games 6 34 robotics (23), and many other fields. However, due to the large cardinality of state space or action space in real-world problems, the large sample complexity has been major problem for employing these RL algorithms in reality. A popular method called multitask representation learning tries to tackle this problem by extracting shared low-dimensional representation function among multiple related tasks, then using simple function (e.g., linear) on top of this common representation to solve each task 24
Despite the empirical success for multitask representation learning, particularly in reinforcement learning because of its effectiveness in reducing sample complexity, the theoretical understanding about it still limited. A march of works (37 36 22 on function approximation in bandits and RI, which permits representation. In these frameworks, an agent is considered playing M related tasks concurrently. Each task is a distinct contextual bandit linear MDP problem and all these M tasks share common representation
33 25
16
30 give results
where = (Ø: : S A + R$} is representation function class extracting a k-dimensional representation
Although the name of linear MDP contains term "linear", actually has infinite degrees of freedom because the representation function & could be general non lonerr function.
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

vector from state-action pair. Such representation function can reduce the complexity of problem from a huge space S A a simple regression problem in k-dimensional space. The value approximation function class is defined by F = L a h, here o means composition and L means linear function, which means the value of any state-action pair (s, a) is linear in its representation ((s,a)).
However, previous analyses either assume & is linear [40], or assume that the agent already knows the concrete function ó [1,21]] which equivalently reduces learning linear weight parameters. This limits their applicability, since general non-linear value estimation is ubiquitous and is the essence for the success of multitask representation learning. For instance, DQN (2 achieves great success by employing deep network to approximate Q-value function. Also, assuming the agent already knows a good representation function unrealistic in practice. Therefore, we aim extend the analysis to unknown general non-linear representation functions. This would not only reveal the more essential benefit of multitask representation learning, but also inspire and facilitate future practice.
1.1 Our Contribution
The focus of previous works on linear analyses has its own reasons. The particularity of linear function could circumvent many non-trivial obstacles analysis, which hinders previous work from from being applicable to real world scenarios. For instance, the formulation of confidence set in linear parameter space is simply an ellipsoid, and its update is straightforward via covariance matrix. More importantly, linear function class generically ensures generalization. The analysis [88 43 26 21) only requires the samples to span the whole input space to let the covariance matrix converge, then is able to derive uniform prediction error guarantee for the whole input space. However, generalization issue is much more complicated for general non-linear scenarios.
In summary, our work embraces following contributions, which solves the challenges for previous works and extends the analysis for the role of representation function in more general setting.
Eliminate the Dependency on Linearity. Towards general function class analysis, we adopt the idea of confidence set [32] 17). The algorithm extends the idea of upper confidence bound and maintains a confidence set for all the possible value estimation functions. The confidence set contains all the functions whose total empirical error at step t is less than a predetermined bound 3r As more seen data reveals more information about the environment, the confidence set will gradually shrink until converge. Therefore, our algorithm and analysis framework is applicable to general function class.
Note that designing 3e to achieve low regret for general function class c is non-trivial We firstly determine the concrete UCB form for general function class 3, ( and propose straightforward algorithm called Generalized Functional Upper Confidence Bound (or GFUCB in abbreviation) for general non-linear function class approximation. We use Eluder dimension 32 measure the complexity of the function class T to give an efficient sample complexity that ensures generalization.
Multihead Function Class. To derive sharp regret bound for our algorithm and theoretically demon- strate the benefit of multitask representation learning, we firstly introduce multihead function class FOM, which is the key technical contribution of our work. The efficacy of multitask representation learning essentially originates from the shared knowledge and structure among tasks. Hence it is vital and necessary to characterize such relation between multiples tasks that the agent simultaneously learns. However, such structure is absent in previous single task work [39 32), and it calls for special techniques to analyze the efficiency for learning these correlated functions.
To this end, we introduce multihead function class, namely FOM in section 4. This abstract function space captures the relation between different task functions, which concatenate the val- ues of (s, a) for all M tasks together as the output. Being more compact by sharing a common backbone function in F requires much fewer samples to learn compared to M indepen- dent tasks space JMM All the tasks contribute to shape a good representation, then feedback to each task for faster convergence. We formally prove that our algorithm enjoys regret bound as δ VMT dim + log lo+)) where T is the number of steps, M is the number of tasks and (())) means the covering number of function space t. We also extend the algorithm and analysis to multitask episodic RL with general value approxination under low inherent Bellman error. By simultaneously solving M different but correlated MDP tasks, our method is sample-efficient with regret Ô (VMTH H dimg(F)) (MM + log ( + MTHI2) where T is the number of episodes. H is planning horizon and I denotes the inherent Bellman error.

---

To the best of our knowledge, this is the first provably sample efficient algorithm for general representation function bandits and linear MDP. It is comparable to the most optimal regret bound when Φ is specialized to linear representation, and is better than the bounds which solve each task independently. This also theoretically explains how multitask representation learning reduces sample complexity. Essentially, the joint training for the shared representation function helps accelerate the convergence of the common backbone by having more samples from all the tasks.
Empirical Value. Finally, we conduct experiments to verify our theoretical result. We design a neural network based bandit environment and implement the GFUCB algorithm. Experimental results corroborate the effect of multitask representation learning in boosting sample efficiency in non-linear bandits. For the first time, the efficacy of the general representation algorithm proposed in theoretical analysis is validated in a proofoff-concee experiment.
2 Related Work
In the supervised learning setting, line of works have been done multitask learning and repre-
sentation learning with various assumptions 15 27 28 14 38 'Teee results assumed
that all tasks share joint representation function. It is also worth mentioning that (38) gave the method-of- ofooomntsssmssmm estimator and built the confidence ball for the feature extractor, which inspired our algorithm for the infinite-action setting.
The benefit of representation learning has been studied in sequential decision-making problems, especially in RL domains. Arora et al. (3] proved that representation learning could reduce the sample complexity of imitation learning. D' D'eramo et al. [I1] showed that representation learning could improve the convergence rate of the value iteration algorithm. Both require a probabilistic assumption similar to that in (28). and the statistical rates are of similar forms as those in (28). Following these works, we study special class of MDP called Linear MDP. Linear MDP 42 21 is popular model in RI,, which uses linear function approximation to generalize large state-action space. [44) extends the definition low inherent Bellman error (or IBE in short) MDPs. This model assumes that both the transition and the reward are near-linear in given features.
Recently, Yang et al. 40 showed multitask representation learning reduces the regret in linear bandits, using the framework developed by Du et al. [4]. Moreover, some works [17 26 21 proved results the benefit of multitask representation learning RL with generative model or linear representation function. However, these works either restrict the representation function class to be linear, or the representation function is known to agent. This is unrealistic in real world practice, which limits these works' meaning.
The most relevant works that need to be mentioned is general function class value approximation for bandits and MDPs. Russo et al. [32 first proposed the concept of cluder dimension to measure the complexity of function class and gave a regret bound for general function bandits using this dimension. Wang et al. (39 further proved that it can also be adopted in MDP problems. Dong et al. [12] extended the analysis with sequential Rademacher complexity. Inspired by these works, we adopt eluder dimension and develop our own analysis. But it should be pointed out that all those works focus single task setting, which give a provable bound for just one single MDP or bandit problem. They lack the insight for why simultancously dealing with multiple distinct but correlated tasks is more sample efficient. Our work aim to establish framework to explain this. By considering locating the ground truth value function in multihead function space J.M M (see detailed definition in section 4), we are able theoretically explain the main reason for the boost of sample efficiency. Informally speaking, the shared feature extraction backbone ò receives samples from all the tasks, therefore accelerating the convergence for every single task compare with solving them separately.
3 Preliminaries
3.1 Notations
We use [n]] to denote the set (1.2..... .n) and (..-) to denote the inner product between two vectors. We use f(x) = (g(()) to represent fxx) I C g(x) holds for any x > 2o with some C > 0 and
To > 0. Ignoring the logarithm term, we use f(x) Ö(g(z))
3

---

3.2 Multitask Contextual Bandits
We first study multitask representation learning in contextual bandits. Each task i [M] is associated with an unknown function f(i) E F from certain function class J. At each step E [T]] the agent is given a context vector a from certain context space C' and a set of actions Ar, selected from certain action space for each task i. The agent needs to choose one action Ar.; E Ar, and then receives a reward as Re = l (Cr,s, Ar,)) + 1]t,s: where i is the random noise sampled from some i.i.d.
distribution. The agent's 's goal is to understand function f(() and maximize the cumulative reward, or equivalently, minimize the total regret from all M tasks in T steps defined as below.
T
M
(eg(T)) def
ΣΣ t=
(CO) - (C,,s,
where A7,; A (Co A) is the optimal action with respect to context Ct, in task i.
3.3 Multitask MDP
Going beyond contextual bandits, we also study how this shared low-dimensional representation could benefit the sequential decision making problem like Markov Decision Process (MDP). In this work, we study undiscounted episodic finite horizon MDP problem. Consider an MDP . = (S, A, H), where S is the state space, A is the action space, P is the transition dynamics, r(::)) is the reward function and H is the planning horizon. The agent starts from an initial state $1 which can be either fixed or sampled from a certain distribution, then interacts with environment for H rounds. In the single task framework, at each round (also called level) h, the agent needs perform an action 0y according to a policy function 04 = Ty (*())) Then the agent will receive a reward R4 (sh, an) r(5h, an) + 77/. where ?h again is the noise term. The environment then transits the state from to 8441 according to distribution Pl.[sp, ax). The estimation for action value function
given following action policy т is defined as Q7, ($h, ax)) = r(5h, an) + Σ Rg(se: The (+:))) and state value function is defined as V (s%) = 0% (sh, (s())) Note that there always exists deterministic optimal policy x for which V (s) = max, Vi and Q (s, a) = max, Q7 a).
we will denote them VV, (s) and Q7, (s, a) for simplicity.
In the multitask setting, the agent gets a batch of states (s(0 simultaneously from M different MDP tasks (.M(0) LM at each round h in episode t, then performs a batch of actions (T)(() M Ji=1 for each task E [M]. Every H rounds form an episode, and the agent will interact with the environment for totally T episodes. The goal for the agent is minimizing the regret defined as
T
M
Reg( (T)
(1) $1.3
Vi
(2) $1.2
where ((i)+ is the optimal value of task i and $1, is the initial state for task i at episode f.
To let representation function play a role, it is assumed that all tasks share the same state space S and action space A. Moreover, there exists representation function Φ x A + Ø² such that action and state value function of all tasks ( is always (approximately) linear in this representation. For example, given a representation function o, the action value approximation function at level
parametrized by a vector 0 R. as Qnlo, 0,] def (Ø(s, a). 0p), similar for Vnlø, (s)(s) def (é(s, a), On)- We denote all such action value functions as Qn = fQu R*), also value function approximation space as Vn- = (V, (b) 0x. : & E + O, E R*). J. Each task M is a linear MDP, which means 2n is always approximately close under Bellman operator To (Qh+1) a) def (s, + E.//Pa(().))) max Qu+1(s',
o 0x] @ E $ 0, e
Linear MDP Definition. A finite horizon MDP M = (S, (S,A,P,r H) is a linear MDP, if there exists a representation function Φ S x A + R² and its induced value approximation function class O h E [,, such that the inherent Bellman error
IA
inf sup Qa-1EQu+1 QuER sES
sup
1Q Ta (Q (1) (s,

---

is always smaller than some small constant I.
The definition essentially assumes that for any Q-value approximation function Q++1 Q++1 level h + 1, the Q-value function Qn at level induced by it can always be closely approximated in class 2h, which assures the accuracy through sequential levels.
3.4 Eluder Dimension
To measure the complexity of a general function class f, we adopt the concept of eluder dimension [3]. First, define c-dependence and independence.
Definition 1 (e-dependent). An input x is e-dependent on set X {xi, D2, Inf with respect
to function class F, if any pair of functions f F satisfying (J(a))( f(x)))² S also satisfies ff(x) f(x) S c. Otherwise, call action be e-independent of data set X.
Intuitively, e-dependence captures the exhaustion of interpolation flexibility for function class F. Given an unknown function f's value on set X { on some particular input with only c-scale prediction error.
we are able to pin down its value
Definition (c-eluder dimension). The c-eluder dimension dim (F,,) c) is the maximum length for a sequence of inputs 21, 22,... U. X, such that for some e' every element is e'-independent of its predecessors.
This definition is similar to the definition of the dimensionality of a linear space, which is the maximum length of a sequence of vectors such that each one is linear independent to its predecessors. O( 1/e) since
For instance, if J {f(z): R + IR, f(x) 0 x}. have cimg(F,e) any d linear independent input's 's estimated value can fully describe a linear mapping function. We also omit the and use dim (F) when it only has logarithm dependent term
4 Main Results for Contextual Bandits
In this section, will present theoretical analysis the proposed GFUCB algorithm for contextual bandits.
4.1 Assumptions
This section will list the assumptions that we make for our analysis. The main assumption is the existence of shared feature extraction function from class = (6 C A + R$} that any task's value function linear in this Φ)
Assumption 1.1 (Shared Space and Representation) All the tasks share the same context space C and action space A. Also, there exists shared representation function E T and a set of k-dimensional parameters M such that each (6) has the form (.)) (d(-.-),
Following standard regularization assumptions for bandits distribution and function parameters.
40
we make assumptions on noise
Assumption 1.2 (Conditional Sub-Gaussian Noise) Denote He, = oCCi,, ALs...., Ct. A be the o-field summarizinn the history information available before reward RG, is observed for every task i E [M]. We have Wt. is sampled from a I-Sub-Gaussian distribution, namely E (exp(>n..) à) H4,3] < exp
x2)) for
Assumption 1.3 (Bounded-Norm Feature and Parameter) We assume that the parameter O, and the feature vector for any centext-action pair (C, A) E C x A is constant bounded for each task i [M], namely ||0||2 S Vk for vi E [M] and [), A)ll2 S 1 for VC C, A E A.
Apart from these assumptions, we add assumption to measure and constrain the complexity of value approximation function class J L +
Assumption 1.4 (Bounded Eluder Dimension). We assume that function class F has bounded Eluder dimension d, which means for any dimg (F,e) Öd).
5

---

4.2 Algorithm Details
Algorithm 1 Generalized Functional UCB Algorithm
1: for step :: 1 T do
2: Compute Fr according to
3: Receive contexts Ct, and action sets Ar,,, i E [M]
4: Ar, = argmax A,E.,. EM f() (Ca,s, A,)
5: Play Ar,; for task i, and get reward R for E [M].
6: end for
The details of the algorithm is in Algorithm At each stept, the algorithm first solves the optimization problem below to get the empirically optimal solution ft that best predicts the rewards for context- input pairs seen so far.
M
f + argmin ΣΣ
(CK,,, An.)) Rk.t)
FEFOM ==
Here we abuse the notation of FaM as FOM = to denote the M-head prediction version of F, parametrized by a shared representation function (.)) and weight matrix W [wi...., WM] REXM We use f(0) to denote the ish head of function which specially serves for iask
(.) )T) E F}
After obtaining ft. we maintain a functional confidence set Fi S CFMM for possible value approxi- mation functions
Fz def
SE | 11 2,E,
S $3,,15 S Cx
[
Here, for the sake of simplicity, we use | | 2,Fr
=E Σ (1() (2x..) f(²) (I
to denote the empirical 2-norm of function f. - f- -(F -() ..... f(.
f(M) Basically,
(+) contains all the functions in JSM whose value estimation difference all collected context- action pairs æy.; (Ck,i, Ak.i) compared with empirical loss minimizer fe does not exceed preset parameter Br. We show that with high probability, the real value function fa is always contained in F, when 3r is carefully chosen as Ó MMk- + log (N(#, a, || llx))). where N/F, a. l1xo) is the a-covering number of function class & in the sup-norm || max A (((a))))) and a is set to be a small number as EMT (see detailed definition and proof in Lemma 1).
For the action choice, our algorithm follows OFUL, which estimates each action value with the most optimistic function value in our confidence set F,, and chooses the action whose optimistic value estimation is the highest. In the multitask setting, choose action from each task to form an action tuple (A1, Az,..., AM) ) such that the summation of the optimistic value estimation EM (Cr,,, A:) is maximized by some function f E Ft.
Intractability. Some may have concerns the intractability of building the confidence set (*) and solving the optimization problem to get ft, At, The solution comes as two folds. From the theoretical perspective, since the focus of problem is sample complexity rather than computational complexity, a computational oracle can simply be assumed to give the solution of the optimization.
This is the common practice for theoretical works 220 35
19 in order to focus on the sample complexity analysis. From empirical perspective, there are great chances to optimize it with gradient methods. For example, solving tively solved with gradient methods [3]. As for fi and Ar,i note that it is not necessary to explicitly build the confidence set F, by listing all the candidates. The approximation algorithm just need to search within the confidence set via gradient method to optimize objective ΣM f()) (Cr,s, (0,,,,A,) A.). The start point is fi, and the algorithm knows that it approaches the border of F, when Il.ft f11., approaches Br. The details of implementation are in section 6.
a standard empirical risk minimization problem, and can effec-

---

Mechanism. GFUCB algorithm solves the exploration problem in an implicit way. For context- action pair æ (C, in task which has not been fully understood and explored yet, the possible value estimation fl( (z) will vary in large range with regard to constraint || ft f.llã, S B. This is because within Fe there many possible function value this z while agreeing all past context-action pairs' value. Therefore, the optimistic value (() will become high by getting a significant implicit bonus, encouraging the agent to try such action under context C, which achieves natural exploration.
The reduction of sample complexity is achieved through joint training for function Φ. If we solve these tasks independently, the confidence set width 3r is at scale M log (W( a, || ll=o)) because it needs to cover M representation function space respectively. By involving e in the prediction for all tasks, our algorithm reduces the size of confidence set by M times, since now the samples from all the tasks can contribute learn the representation Φb Usually log (N(Q, a, [lx))) is much greater
than & and M, hence our confidence set shrinks at much faster speed. This explains how GFUCB achieves lower regret, since the sub-optimality at each step is proportional the confidence set width 3r when real value function fo E Fe.
4.3 Regret Bound
Based on the assumptions above, we have the regret guarantee as below.
Theorem 1. Based on assumption 1.1 to 1.4, denote the cumulative regret in T steps as Reg(T), with
probability at least 8 we have Reg(T)- = Ō (VMdT(MK log or, || llco))
Here, d dim a.r) is the Eluder dimension for value approximation function class F L ob and Crr is discretization scale which only appears in logarithm term thus omitted. The detailed proof left in appendix.
To the best of knowledge, this is the first regret bound for general function class representation learning in contextual bandits. To get sense of its sharpness, note that when 4 is specialized as linear function class as Ø {Ø(z)= Bæ, B E Rexd,, J. we have log (Q, oT, Il-llso) Ödkk) and dimg((F) d, then our bound is reduced Ó(M the current best provable regret bound for linear representation class bandits in [17
dv dTMT),, which is the same optimal as
5
Main Results for MDP
5.1
Assumptions
For multitask Linear MDP setting, we adopt Assumption 3 from 17 which generalizes the inherent Bellman error (44) multitask setting.
Assumption 2.1 (Low IBE for multitask) Define multi-takk IBE is defined as
Tmall def
sup J.m, eQa-I
sup S.me EQ, (EEMM
inf M
(00 -Th²" (ol (s, (s,a)/-
We have I def sup, ITmill is small for all Qn. h E [H].
Assumption 2.1 generalize low IBE to multitask setting. It assumes that for every task i [M], its Q-value function space is always close under Bellman operator.
Assumption 2.2 (Parameter Regularization) We assume that
|ø(s, a))l < 1, 0 < (x(s,a) a) & for V(s a) E A, (H) Vr.
There exists constant D such that for any h E [H] and g it holds that 110% I D.
For any fixed { { } M
E Qutr. the random noise = del
R (s. a) +
hhts max, o ,a) The) (Q a) is bounded in [-1, and is always independent
to all other random variables for (ss a) S A.h E [H] [M].
7

---

These assumptions are widely adopted in linear MDP analytical works A4 17 26], which regularizes the parameter, feature, and noise scale. Again we add bounded Eluder dimension constraint for the Q-value estimation class.
Assumption 2.3 (Bounded Eluder Dimension). We assume that function class 2n has bounded Eluder dimension d for any h E [H].
5.2 Algorithm Details
Algorithm 2 multitask Linear MDP Algorithm
1: for episode
do
2:
H+1 =
(c)
3: for h H (6)
do
4:
Φn,t 1.t + solving (1)
0h2 (.--) dn.( T ()) max, (.
5:
6: end for
7:
for h + H do
8:
Compute Fi.e according to Lemma 4 M
(6) hh
Sh,e a = jegna Ei a11)) ((t
(5()
9: Receive states
i=1" R for task E [M].
Play ane and get reward end for 12: end for
The algorithm for multitask linear MDP is similar to contextual bandits as above. The optimization problem line 4 of of[olgori 2 is finding the empirically best solution for Q-value estimation at level h in episode as below
argmii
L(o, 9)
(1)
$ 6E0,0=(01 (0(L))...,0(N)
s.t.
| < D, [
0 S (ss T 0, s L. A, [M],
where L(d, Q) is the empirical loss function defined as
M ΣΣ
2
O R
(h+1.))
The framework of work resembles LSVI 21 and (26 which learns the Q-value estimation in reverse order, at each level h, the algorithm uses just-learned value estimation function V to build the regression target value as R.(0) + v7(6)
(s(9) (8)+1s) and find empirically best estimation f(())
h+1
TT. h.t for each task i [M]. The optimistic value estimation of cach action is again searched within confidence set FR. which centered at fú,e and shrinks as the constraint ||. fl..113.. S 3 becomes increasingly tighter. Note that the contextual bandit problem can be regarded MDP problem without transition dynamics, and our framework at each level h is indeed a copy of procedures in Algorithm 1.
1-horizon
5.3 Regret Bound
Based on assumptions 2.1 to 2.3, we prove that our algorithm enjoys a regret bound guaranteed by the following theorem. Detailed proof is left in appendix.
Theorem Based on assumption 2.1 to 2.3, denote the cumulative regret in T episodes as Reg(T"), we have the following regret bound for Reg(T) holds with probability at least - 8 for Algorithm
ò (MHVTdk- + H W/MMdllogN(00.)) +MHTIVd)

---

where a is discretization scale smaller than 1 KMT
Remark. Compared with naively executing single task general value function approximation algorithm (39 for M tasks, whose regret bound is O
to achieve same average regret, our algorithm outperforms this naive algorithm with a boost of sample efficiency by Ö(Md). This benefit mainly attributes to learning in function space F&M L b instead of J.M (C o () &)M,
M, the former is more compact and requires much less samples to learn.
Experiments
To validate our theoretical findings, we conduct experiments on non-linear neural network bandits. Note that it is proof-of-concept experiment. Our main purpose is to realize the GFUCB algorithm and check its efficacy but not to beat sophisticated real-world algorithms. The point to demonstrate is that sample efficiency of GFUCB is scalable to the number of tasks and better than naive exploration.
6.1 Task Design
To test the efficacy of our algorithm, we use the MNIST dataset [10] to build bandit problem that involves non-linear value approximation. The reward function of the bandit environment maps the same digit into the same base reward Ty. which ranges from 0 to I, plus a noise sampled from a zero-mean Gaussian with standard deviation of 0.01. At every round, each task will present the agent a context C' consists of K different digit images and ask the agent to take action an integer j E [K] meaning which image to choose, then return the reward according the agent's choice.
For the multitask setting, we construct M different tasks using different digit-to-reward mappings o; {O..... 9} + [0, 1] i [M]. where (k) will give a unique reward for all images of digit k in task i. Different tasks have different reward mapping function (.). By designing the environment this way, requires learn a common representation recognize digits for different tasks.
6.2 Implementation Details
We use a simple CNN as our feature extraction function o, which takes a digit image as input and outputs 10-dimensional normalized vector representation. It consists of two 3x3 convolution layers and two fully-comnected layers, followed by ReLU activation and a normalization procedure.
The biggest challenge for implementation is how to solve complex optimization problem in general functional space. principle, finding parameters for neural network to achieve the (near) minimal empirical error is an NP-Hard problem. To solve this issue, we use a gradient-based method to approximately find local-optimal solution. For finding the empirically best ft. we use Adam with Ir = le - 3 to train for sufficiently long steps; in our setting, it is set to be 200 epochs at every step t, to ensure that the training loss is sufficiently low.
The next major challenge is estimating the optimistic value for each action within the abstract function set F,. To tackle this problem, we enumerate all possible action tuples (A:).M, and then solve the equivalent optimization below to compute its optimistic estimated value
(C s.t.
| fo
2.E. S Bc.
JEF
Still, this a complicated optimization problem within an abstract function set. Inspired by the Lagrangian operator, transform it into unconstrained optimization problem minimizing loss function {(J) Σ (A.) X max (0, 11.. -Ill B.), where À is hyperparameter to be determined, in our algorithm we set be 30 by empirical search. Also B, a log(b- 6 t + e) is an approximation for 3, since 3e includes (4 a) which is intractable to be exactly computed, we found (a, c) = (0.4. 0.5. be a good parameter of UCB in single task. We use SGD with small learning rate (5e 4) to finctune the model fo for 200 iterations to optimize E(f).
The basic intuition is that, through optimizing é(f), the algorithm will try to maximize function E, < Br, such constraint will not appear
value ΣM, (A.). And as long as ff satisfies Ilfe fll,.,, in the loss term, thus has no effect optimization. When comes to the border of F,, where

---

Ilft - f approaches B,, the second term adds regularization term to the loss as punishment, preserving Ilft fl2,r, at a near-constant level around Bt. So we can approximately simulate the optimistic value estimating procedure via searching in the neighborhood of ft.
6.3 Connection to Algorithm
The main difference between our practical version algorithm and the theoretical one is that we did not list out all the functions in the whole confidence set JF, explicitly, but just use gradient-based method to implicitly search within very small fraction of Fr with heuristics. Getting a candidate within the confidence set much easier and tractable than rigorously exhausting all functions J+ to optimize. We can start from the parameter of ft and use gradient method to approximately find ft and Ar,i.
Another difference is we do not rigorous compute Be which involves (()), but directly determine parametrized function form. Rigorously speaking, our tuned value 3r is much smaller than the the- oretical guaranteed ones, so all the candidate functions that we search along the trajectory of gradient method still satisfy the theoretical requirement (but it may omit many other potential candidates). Therefore, our practical version algorithm should be regarded as an inaccurate approximation to the theoretical algorithm. Moreover, also plays a role as regularization to enable the convergence of Fe since we only consider regular ones the neighborhood of ft.
6.4 Results
We test the performance of our algorithm against naive eps- eps-greedy baseline that solves each task independently by training the same CNN value prediction module. We show
50
Epsiion Greedy 1 task
Annelled Epailon Grewdy 1 last
GFUCI 1 task
GFUCB 5 tasks
Cumuletivv Regret
GFUCB 12 tasks
1,5, 10 in Figure our results with number of tasks different value [ Firstly, randomly generate 10 mapping functions ()) 1,.... 10. The total 10 tasks divided into 10/M groups; each group forms M-task problem and is solved by an individual copy of some al- gorithm. At each step I, the cumulative regret from all 10 tasks is averaged to estimate the method's performance. Our result in Figure verified that the multitask training does accelcrate learning, which empirically validates our theoretical analysis. The multitask training utilizes the samples from all M tasks to jointly learn a good represen- tation Φ which significantly accelerates the learning procedure of the CNN backbone. Also, the improvement in GFUCB algorithm's 's performance with M = 1 validates the effect of our finetune procedure for getting a bonus. Detailed dissection and discussion are left in appendix.
30
20
10
100
500
Figure 1: Cumulative regret over steps
for = 1,5, 10.
7 Conclusion
In this work, we extend the analysis of the benefit of multitask representation learning from linear representation class to general function class. We propose a straightforward algorithm that can utilize samples from all the tasks to jointly train a representation function, which is demonstrated theoretically and empirically to accelerate the sample efficiency and outperform naively single-task learning. Also, we extend the analysis to the MDP setting and show that the benefit of multitask representation learning is similar. Furthermore, our experimental result reveals that our proposed algorithm is also effective in practice even for highly non-linear neural network representations.
Acknowledgments and Disclosure of Funding
This work is supported part by the National Science and Technology Major Project of the Ministry of Science and Technology of China under Grants 2018AAA0101604 the National Natural Science Foundation of China under Grants 62022048 and the State Key Lab of Autonomous Intelligent Unmanned Systems.
10

---

References
[1] Alekh Agarwal, Daniel Hsu, Satyen Kale, John Langford, Lihong Li, and Robert Schapire. Taming the monster: A fast and simple algorithm for contextual bandits. In International Conference on Muchine Learning, pages 1638-1646. PMLR, 2014.
[2] Rie Kubota Ando and Tong Zhang. framework for learning predictive structures from multiple tasks and unlabeled data. Journal of Machine Learning Research, 6(Nov):1817-1853, 2005.
[3] Sanjeev Arora, Simon S Du, Sham Kakade, Yuping Luo, and Nikunj Saunshi. Provable represen- tation learning for imitation learning via bi-level optimization. arXiv preprint arXiv:2002.10544 2020.
[4] Jonathan Baxter. A model of inductive bias learning. Journal of artificial intelligence research, 19119-198, 2000.
[5] Shai Ben-David and Reba Schuller. Exploiting task relatedness for multiple task learning. In Learning Theory and Kernel Machines, pages 567-580. Springer, 2003.
[6] Christopher Berner, Greg Brockman, Brooke Chan, Vicki Cheung, Przemyslaw Debiak, Christy Dennison, David Farhi, Quirin Fischer, Shariq Hashme, Chris Hesse, et al. Dota 2 with large scale deep reinforcement learning. arXiv preprini 19tXiv:://2.0000 2019.
[7] Rich Caruana. Multitask learning. Machine learning. 28(1):41-75, 1997.
[8] Giovanni Cavallanti, Nicolo Cesa-Bianchi, and Claudio Gentile. Linear algorithms for online multitask classification. Journal of Machine Learning Research, 1 11(Oct):2901-22 2934, 2010.
[9] Yuan Cheng, Songtao Feng, Jing Yang, Hong Zhang, and Yingbin Liang. Provable benefit of multitask representation learning in reinforcement learning. arXiv preprint (XKi::2206.059000 2022.
[10] Li Deng. The mnist database of handwritten digit images for machine learning research [best of 2012.
the web]. IEEE Signal Processing Magazine,
[11] Carlo D' D'Eramo, Davide Tateo, Andrea Bonarini, Marcello Restelli, and Jan Peters. Sharing knowledge in multi-task deep reinforcement learning. In International Conference on Learning Representations, 2020.
[12] Kefan Dong, Jiaqi Yang. and Tengyu Ma. Provable model-based nonlinear bandit and reinforce- ment learning: Shelve optimism, embrace virtual curvature. arXiv preprint (et::2002.41110 2021.
[13] Simon Du, Jason Lee, Haochuan Li, Liwei Wang, and Xiyu Zhai. Gradient descent finds global minima of deep neural networks. In International conference on machine learning, pages 1675-1685. PMLR, 2019.
[14] Simon S Du, Wei Hu, Sham M Kakade, Jason D Lee, and Qi Lei. Few-shot learning via learning the representation, provably. arXiv preprint arXiv:2002..0 09434, 2020.
[15] Simon S Du, Jayanth Koushik, Aarti Singh, and Barnabás Póczos. Hypothesis transfer learning via transformation functions. In Proceedings of the 3 st International Conference on Neural Information Processing Systems, pages 574-584, 2017.
[16] Matteo Hessel, Hubert Soyer, Lasse Espeholt, Wojeiech Czamecki, Simon Schmitt, and Hado van Hasselt. Multi-task deep reinforcement learning with popart. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pages 3796-3803, 2019.
[17] Jiachen Hu, Xiaoyu Chen, Chi Jin, Lihong Li, and Liwei Wang. Near-optimal representation learning for linear bandits and linear rl. In International Conference on Machine Learning, pages 4349 4358. PMLR, 2021.
[18] Yao Hu, Debing Zhang, Jieping Ye, Xuelong Li, and Xiaofei He. Fast and accurate matrix completion via truncated nuclear norm regularization. Pattern Analysis and Machine Intelligence. IEEE Transactions on, 3509)22117-110, 2013.
[19] Nan Jiang, Akshay Krishnamurthy, Alekh Agarwal, John Langford, and Robert E Schapire Contextual decision processes with low bellman rank are PAC-learnable. In Proceedings of the 34th International Conference on Machine Learning- Volume 70, pages 1704-1713. JMLR, org,
2017.
11

---

[20] Chi Jin, Qinghua Liu, and Sobhan Miryoosefi. Bellman eluder dimension: New rich classes of rl problems, and sample-efiiciet algorithms. Advances in neural information processing systems, 34:13406-13418, 2021.
[21] Chi Jin, Zhuoran Yang. Zhaoran Wang, and Michael Jordan. Provably efficient reinforcement learning with linear function approximation. arXiv preprint (rXXiv/907.055 0003880 2019.
[22] Alessandro Lazaric and Marcello Restelli. Transfer from multiple mdps. Advances in Neural Information Processing Systems, pages 1746-1754, 2011.
[23] Sergey Levine, Chelsea Finn, Trevor Darrell, and Pieter Abbeel. End-to-end training of deep visuomotor policies. The Journal of Machine Learning Research, 17(1):1334-137, 2016.
[24] Lihong Li, Wei Chu, John Langford, and Robert E Schapire. A contextual-bande approach to personalized news article recommendation. In Proceedings of the 19th international conference on World wide web, pages 661-670, 2010.
[25] Lydia T Liu, Urun Dogan, and Katja Hofmann. Decoding multitask dqn in the world of minecraft. In The 13th European Workshop on Reinforcement Learning (EWRL) 2016, 2016.
[26] Rui Lu, Gao Huang, and Simon S Du. On the power of multitask representation learning in linear mdp. arXiv preprint arXiv:2 106.0803, 2021.
[27] Andreas Maurer. Bounds for linear multi-task learning. Journal of Machine Learning Research, 17731): 2006.
[28] Andreas Maurer, Massimiliano Pontil, and Bernardino Romera-Paredes. The benefit of multitask representation learning. The Journal of Machine Learning Research, 17(1):2853-288 28884, 2016.
[29] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wierstra, and Martin Riedmiller. Playing atari with deep reinforcement learning. arXiv preprint arXiv: /3/2.56 2013.
[30] Matteo Papini, Andrea Tirinzoni, Aldo Pacchiano, Marcello Restelli, Alessandro Lazaric, and Matteo Pirotta. Reinforcement learning in linear mdps: Constant regret and representation selection. Advances in Neural Information Processing Systems, 34:16371-1638; 2021.
[31] Emilio Parisotto, Jimmy Lei Ba, and Ruslan Salakhutdinov. Actor-mimic: Deep multitask and transfer reinforcement learning. arXiv preprint 15X101551 2015.
[32] Daniel Russo and Benjamin Van Roy. Eluder dimension and the sample complexity of optimistic exploration. In NIPS, pages 2256-2264. Citeseer, 2013.
[33] Andrei A Rusu, Sergio Gomez Colmenarejo, Caglar Gulcehre, Guillaume Desjardins, James Kirkpatrick, Razvan Pascanu, Volodymyr Mnih, Koray Kavukcuoglu, and Raia Hadsell. Policy distillation. arXiv preprint wrxiv:1511.06299 2015.
[34] David Silver, Julian Schritwwiseer, Karen Simonyan, Ioannis Antonoglou, Aja Huang, Arthur Guez, Thomas Hubert, Lucas Baker, Matthew Lai, Adrian Bolton, et al. Mastering the game of go without human knowledge. nature, (55(7666)3354-33333 2017.
[35] Wen Sun, Nan Jiang, Akshay Krishnamurthy, Alekh Agarwal, and John Langford. Model-based reinforcement learning in contextual decision processes. arXiv preprint arXiv: /81 1eee:811100555 1.08540, 2018.
[36] Matthew E Taylor and Peter Stone. Transfer learning for reinforcement learning domains: survey, Journal of Machine Learning Research, 10(30):1633-1665, 2009.
[37] Yee Teh, Victor Bapst, Wojeiech M Czarnecki, John Quan, James Kirkpatrick, Raia Hadsell, Nicolas Heess, and Razvan Pascanu. Distral: Robust multitask reinforcement learning. In Advances in Neural Information Processing Systems, pages 4496-4506, 2017.
[38] Nilesh Tripuraneni, Chi Jin, and Michael I Jordan. Provable meta-learning of linear representa- tions. arXiv preprint arXiv:200..11684, 2020.
[39] Ruosong Wang, Ruslan Salakhutdinov, and Lin Yang. Reinforcement learning with general value function approximation: Provably efficient approach via bounded eluder dimension. arXiv preprint arXiv:2005. 10804, 2020.
[40] Jiaqi Yang, Wei Hu, Jason D. Lee, and Simon Shaolei Du Impact of representation learning in linear bandits. In International Conference on Learning Representations, 2021.
12

---

[41] Jiaqi Yang, Qi Lei, Jason D Lee, and Simon S Du. Nearly minimax algorithms for linear bandits with shared representation. arXiv preprint arXiv:2203. 15664. 2022.
[42] Lin Yang and Mengdi Wang. Sample-optimal parametric q-learning using linearly additive features. In International Conference on Machine Learning, pages 6995-7004. PMLR, 2019.
[43] Lin Yang, Chengzhuo Ni, and Mengdi Wang. Learning to control in metric space with optimal regret. arXiv preprint arXiv: 1rXXv:1905.01566 2019.
[44] Andrea Zanette, Alessandro Lazaric, Mykel Kochenderfer, and Emma Brunskill. Learning near optimal policies with low inherent bellman error. In International Conference on Machine Learning, pages (0978-10989. PMLR, 2020.
13


Reviews:
[{"review_id": "uF6anSQA4jM", "reviewer": "Reviewer_rCEX", "summary": "The paper extends regret analysis for multi-task linear MDPs and bandit problems when a shared representation that is unknown to the agent is used across tasks. A general representation function class $\\Phi$ and value function approximation class $\\mathcal{L} \\circ \\Phi$ that is linear in $\\Phi$ is used. To enable analysis in this setting, the eluder dimension of Russo et al. (2013) is adopted and $\\mathcal{L} \\circ \\Phi$ is assumed to have bounded eluder dimension. In this setting, the paper provides novel regret bounds for a UCB algorithm for multi-task bandits and linear MDPs. In the bandit case for linear representation class, the paper states that the bound recovers the best known bound as a special case. For linear MDPs, the bound extends known bounds for the single task setting and achieves improved sample complexity over a naive baseline that does not share representations across tasks. An empirical analysis is provided as a sanity check to demonstrate reduced cumulative regret over time with increasing number of tasks.", "questions": "Minor suggestions (typos, etc.):\n- L318: \"boarder\" should be \"border\"\n- Maybe capitalize section headers 3.2 and 3.3.\n", "limitations": "I was unable to find discussion of limitations of the algorithm/analysis in the text.", "rating": 7, "confidence": 2, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- To the best of my knowledge, the paper is the first to characterize how sample complexity is improved in the multi-task MDP setting when shared structure is assumed for MDPs with low inherent Bellman error (IBE).\n- The fact that the previous known bound of [14] is recovered as a special case for multi-task bandits is convincing regarding the sharpness of the bounds.\n\nWeaknesses:\n- Low IBE (Assumption 2.1) for all tasks strikes me as a potentially strong assumption. It would have been interesting to see how the algorithm behaves when a small subset of tasks do not have low IBE. How robust would the algorithm be in that case? If cumulative regret indeed grows linearly in the worst-case IBE as predicted by Thm. 2, this might be a shortcoming of the algorithm.\n- The mathematical extension of regret bounds from single task to the multi-task setting seems somewhat straightforward. Nevertheless, shared representation structure is a key differentiator.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "tYYtXyn5WIT", "reviewer": "Reviewer_8wSD", "summary": "The authors consider multitasks reinforcement learning where generalized representations function among the tasks and task-specific parameters will be estimated sequentially in order to minimize cumulative regret. Previous works only focus on either linear function class or single-task which limits their applicability. Inspired by them, the authors employ eluder dimension to handle the concept of general function class and treat the problem as linear MDP (as well as linear contextual bandit). In the end, they provide proof-of-concept experiments to evaluate the benefit of shared representation among tasks ", "questions": "Firstly, I suggest the authors put a table to highlight the distinction from a single task to multiple tasks as well as from linear class to general class. As I mentioned in Weakness 1, I don't see the significant contribution made by this work. If the authors can convince me that challenges from previous works to this one are nontrivial, I will consider increasing my score. \n\nSecondly, I hope the author can connect the experiment to your theorems. For example, N(\\Phi, \\alpha) apparently is the main feature distinct from [14], but how can I see its influence shown in the experiment? Can the authors please point out in the rebuttal?", "limitations": "I don't see any potential negative social impact from their work and the authors address limitations properly.", "rating": 3, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strength:\nThis work extends the results in [14] from linear representation class to the general class with the tools provided from [25]. Their upper bound (Thm 1 and 2) explain the cost for learning representations function, O(\\sqrt{MdTlog N(\\Phi)}), will be sublinear in M, hence this huge cost (compared to learning \\theta_i for each task) will be mitigated by jointly learning multiple tasks.\n\nWeakness:\n1. It is a not surprising combination of previous works, where most relevant techniques can be found in [14] and [25]. i.e., [14] designed and analysised  the same style algorithms for learning linear representation for multitasks in linear bandits and RL; [25] (and [31]) demostrated the construction of confidence bound for general class. Although I have quickly checked the proofs and believe the correcness of the main theorems (there are many inequalities and stataments that I have to guess the reason behind them, hope the authors can improve their writing), it is unclear what this work contributes.\n\n2. Provable but intractable. The main difference between this work and [14] should be the representation extractor and confidence set (Line 202, 208 and 262 in this work and definition 1 in [14]). As mentioned by the authors (line 39), it is relatively easy to construct in linear class but hard in general. I expected the authors show how they construct the confidence set (*), but it is disappointing they used CNN as the representation extractor rather than solving line 202 and 262. It is mainly because the algorithms presented in this work suffer intractable computational cost, hence this makes us hard to derive any useful insight from their theorems to the practice.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "jxfCGH8JVNA", "reviewer": "Reviewer_1Kyf", "summary": "The paper proves the first near-optimal regret analysis for multitask bandits and MDPs with a general function class representation learning. Different from previous literature, this paper consider general representation function that is not known to the algorithm. The algorithms follow the well-known optimism-in-face-of-uncertainty design. The effectiveness of the algorithm for bandits is validated on MINIST dataset.", "questions": "Line 70: it would be better to point it out that this result is for contextual bandit setting.\nAssumption 2.1: there is a typo.\nLine 280: How is the performance boosted by $Md$? Should it be $\\sqrt{Md}$?", "limitations": "I believe the main limitation of the proposed algorithms is that they are not practical. Discussing a practical algorithm for MDP case would be interesting.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\nThe paper adds to the line of works that try to understand the benefits of multitask representation learning on sequential decision-making problems. As far as I know, it is the first paper to consider the unknown representation function beyond linear class. The paper is well-written and easy to follow. The proofs are complete.\n\nWeakness:\n\nMy concern is that the contribution of this paper is highly incremental. Benefits of the similar form have been extensively discussed in the literature though it is a relatively new setting in this paper. I briefly went through the proofs. The paper makes very limited technical contributions. If there is any significant technical contribution, I suggest the authors point it out in the main text. The paper discussed both bandit and MDP settings, while the experiments are only conducted for bandits problems. It would be good to discuss the practice algorithm for MDP as well. \n\nIt would be interesting to discuss the worst-case generalization performance for a learned representation, which seems to be a more interesting problem for multitask RL.\n\n\n\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "edv32qiLFlb", "reviewer": "Reviewer_jNGD", "summary": "This paper studies multitask representation learning with *general* function class in contextual bandits and MDP setting. In this setting, the goal is to learn a shared representation function that maps the state (or context) and action to a $k$-dimensional vector that allows the value function to be expressed as a linear function of this representation. They proposed an algorithm, Generalized Functional Upper Confidence Bound (GFUCB) which adapts UCB for a general function class in multitask setting and prove a regret bound for both the contextual bandits setting and multitask episodic RL setting with M related tasks (sharing common state and action space). Experiments were conducted in the non-linear neural network bandits setting based on the MNIST dataset to demonstrate that their GFUCB algorithm indeed outperforms naive epsilon greedy on a single task and scales to a larger number of tasks (i.e. lower cumulative regret with more tasks). ", "questions": "1. While the focus is on the theoretical analysis in this paper, one possible improvement that can further strengthen the paper is providing also a proof of concept experiment for a linear MDP. Perhaps it can be built off the MNIST dataset as well (such as translating certain digits to their designated spot in each of the multitask instantiation), or a grid world fed as an image into a CNN. \n2. A suggestion: it will be helpful if you summarize the practical approximations to Algorithm 1 in the Experiment section 6.2 as a separate “Algorithm 3” to show the correspondence between the original GFUCB algorithm and the practical instantiation (with running gradient descent for the optimization step, etc.). \n", "limitations": "The authors did not explicitly discuss the limitations of their work beyond stating the assumptions for the analysis. One observation noted in the appendix was that while the eluder dimensions of neural networks can be exponentially large, in practice (as shown in their experiment), the UCB bonus term still shrinks as the number of training samples increases. Thus specific interplay of using neural networks as the function class and its generalization / interpolation behavior is up for future study for this problem setting. \n\nAs a mainly theory paper, the authors specified “N/A” for potential negative societal impact. Perhaps we can still speculate on potential harmful uses of highly efficient multitask RL training algorithms leading to an agent capable of very large number of tasks. \n", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Originality**: The paper presents the first provably sample efficient (with regret bound) algorithm for general function class representation learning in contextual bandits and linear MDPs. The authors clearly positioned their contributions in the background of prior works, which either studied multitask representation learning in linear representation functions or assumed prior known representation functions, or only studied single task setting when tackling general function class value approximation for bandits and MDPs. \n\n**Quality**: The paper appears to be technically sound. The definitions and assumptions for the main results are given leading up to the algorithm and regret bound. The empirical experiments however only tackle the bandits setting, but not an MDP setting. \n\n**Clarity**: While the paper is fairly dense in mathematical notations, overall the paper is adequately well-written. The intuitive interpretation of several theoretical results was helpful. \n\n**Significance**: Given the abundant interest in multitask deep RL settings, this theoretical result is a step towards a deeper (pun intended) understanding in the community. The analysis in this paper and the resulting algorithm can give further inspiration for principled approaches in designing more practical algorithms with deep neural networks as well. ", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept