Meta Reinforcement Learning with Finite Training Tasks - a Density Estimation Approach 
Authors: Zohar Rimon
Aviv Tamar
Gilad Adler
Venue: NeurIPS 2022 Accept (2022)
Keywords: RL
Bayesian RL
Meta RL

Abstract: In meta reinforcement learning (meta RL), an agent learns from a set of training tasks how to quickly solve a new task, drawn from the same task distribution. The optimal meta RL policy, a.k.a.~the Bayes-optimal behavior, is well defined, and guarantees optimal reward in expectation, taken with respect to the task distribution. The question we explore in this work is how many training tasks are required to guarantee approximately optimal behavior with high probability. Recent work provided the first such PAC analysis for a model-free setting, where a history-dependent policy was learned from the training tasks. In this work, we propose a different approach: directly learn the task distribution, using density estimation techniques, and then train a policy on the learned task distribution. We show that our approach leads to bounds that depend on the dimension of the task distribution. In particular, in settings where the task distribution lies in a low-dimensional manifold, we extend our analysis to use dimensionality reduction techniques and account for such structure, obtaining significantly better bounds than previous work, which strictly depend on the number of states and actions. The key of our approach is the regularization implied by the kernel density estimation method. We further demonstrate that this regularization is useful in practice, when `plugged in' the state-of-the-art VariBAD meta RL algorithm.

Meta Reinforcement Learning with Finite Training Tasks - a Density Estimation Approach
Zohar Rimon
Aviv Tamar
Technion Israel Institute of Technology zohar, rimon@campus technion.
Technion Israel Institute of Technology
ac. ac.onn..........................
avivt@technion.c ac
Gilad Adler Ford Research Center Israel gadler38ford.cm com
Abstract
In meta reinforcement learning (meta RL), an agent learns from a of training tasks how to quickly solve new task, drawn from the same task distribution. The optimal meta RL policy, a.k.a. the Bayes-optimal behavior, is well defined, and guarantees optimal reward in expectation, taken with respect to the task distribution. The question we explore this work is how many training tasks are required guarantee approximately optimal behavior with high probability. Recent work provided the first such PAC analysis for model-free setting, where history- dependent policy was learned from the training tasks. In this work, we propose a different approach: directly learn the task distribution, using density estimation techniques, and then train a policy on the learned task distribution. We show that our approach leads to bounds that depend on the dimension In particular, in settings where the task distribution lies in www-dimnsioonl manifold, we extend our analysis to use dimensionality reduction techniques and account for such structure, obtaining significantly better bounds than previous work, which strictly depend on the number of states and actions. The key of our approach is the regularization implied by the kernel density estimation method. We further demonstrate that this regularization useful in practice, when 'plugged in' the state-of-the-art VariBAD meta RL algorithm.
task distribution.
Introduction
In recent years, reinforcement learning (RL) became dominant algorithmic framework for variety of domains, including computer games [24), robotic manipulation [I], and autonomous driving [18]. Popular RL algorithms, however, characterized by a high sample inefficiency, due to the peploration-explottaaie about the environment versus acting based on such information. Indeed, most RL success stories required a very long training process, only possible in simulation.
problem the need to balance between obtaining more information
For an agent to learn fast, additional structure of the problem is required. In Meta RL 39 26 agents are allowed train on a set of training tasks, sampled from the same task distribution as the task they will eventually be tested on. The hope is that similar structure between the tasks could be identified during learning, and exploited to quickly solve the test task. It has recently been observed that the meta RL problem is related to the Bayesian RL formulation, where each task is modelled as Markov decision process (MDP), and the distribution over tasks is the Bayesian prior [5] 39]. In this formulation, the optimal meta RL policy is the Bayes-optinal policy the policy that maximizes expected return, where the expectation is taken with respect to the prior MDP distribution.
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

While significant empirical progress in meta RL has been made, the theoretical understanding of the problem still limited. A central question, which we focus on in this work, the probably approximately correct (PAC) analysis of meta RI, namely, how many training tasks are required to guarantee performance that is approximately Hayes-optimal with high probability. Practical motivation for studying this question includes the lifelong learning setting [8]. where training tasks are equivalent to tasks that the agent had previously encountered, and we seek agents that learn as quickly as possible, and the offline meta RL setting [3], where training data is collected in advance, and therefore estimating how much data is needed is important.
Recently, Tamar et al. [33] proposed the first PAC bounds for meta RL, using a model free approach. In their work, a history- policy was trained to optimize the return on the set of training MDPs, where policy-regularizatinn was added to the loss of each MDP in the data. The bounds in 33 scale with the number of states exponentiated by the length of the history, intuitively due to the number of possible histories that can be input to the policy. In this work, we propose model based approach for PAC meta RL. Our idea is that instead of regularizing the policy during training, as in [33] we can learn regularized version of the distribution of MDPs from the training data. Subsequently, we can train an agent to be Bayes optimal on this learned distribution. Intuitively, we can guarantee that the learned distribution is 'close enough' to the real prior, we should expect the learned policy to be near Bayes optimal. Here, we derive such guarantees using techniques from the kernel density estimation (KDE) literature [29]
Building on PAC results for KDE, we derive PAC bounds for our model based meta RL approach. Compared to the bounds in 33 our results require much less stringent assumptions on the MDP prior; and apply to both continuous and discrete state and action spaces. Interestingly, our bounds have a linear dependence on the horizon, compared to the exponential dependence in [33] but have an exponential dependence on the dimension of the prior distribution, corresponding to the exponential dependence dimensionality of KDE. We argue, however, that in many practical cases the prior distribution lies in some low dimensional subspace. For such cases, we extend our analysis to use dimensionality reduction via principal component analysis (PCA), and obtain bounds where the exponential dependence is on the dimension of the low dimensional subspace.
To visualize the advantage of our approach. consider the HalfCirele domain in Figure [ adapted from 33:
Episode 1
2-dimensional agent must navigate to a goal, lo-
Episode 2
cated somewhere on the half-cirele. A task therefore corresponds to the goal location, and the task distri- bution is uniform on the I-dimensional half-circle. The bounds in [3 depend on the number of states and actions, which are continuous here, and even if discretized, would result in excessive bounds. Intu- itively. however, it is the low-dimensional structure of the task distribution that should determine the dif- ficulty of the problem, and our bounds can indeed account for such. We remark that many benchmark meta RL task distributions exhibit similar low dimen- sional structure [4 33.
Figure 1: The HalfCirele domain (taken from B): the task is navigate a goal position that can be anywhere on the half-circle (light blue). A Bayes-optimal agent first searches along the half circle for the goal, and once found, moves directly towards it.
To complement our theoretical results, we demon- strate the practical potential of our approach by in- corporating it in the VariBAD meta-RL method of Zintgraf et al. [39] In VariBAD, variational autoencoder (VAE) is used to learn now-dimensional latent representation of the task distribution, and a deep neural network is trained to approximate the Bayes optimal policy. We show that by estimating the task distribution using our kernel density estimation method, applied to the VAE latent space, and training the policy on sampled tasks from this distribution, we improve the generalization of the policy to tasks not seen in the training data.
2 Background
We specify our notation, and background on MDPs, density estimation, and dimensionality reduction. Notations: Estimators will be denoted with X. ||-|| denotes the Euclidean norm, other norms will be

---

denoted explicitly. For set X, we denote by P(X) the set of distributions over X. We denote by va the volume of a unit ball in TRd For a space S we denote S as the volume of the space.
2.1 Markoy Decision Processes
We follow the formal setting of [33]. A Markov Decision Process (MDP) M is a 7-tuple M (S,A, A. C, Puit: C, P. H), where S, and C are the state, action and cost spaces, respectively, Pmnt is the initial state distribution, C:: S A + P(C) is the cost function, P S x A P(8) is the transi- tion function and H E Z+ the horizon. We denote by (c. 8, 5) P (s' | 8, a) C (c 8, a) the probability of transitioning from state given action to state s' with cost of c, Vs, E S. Va E A and Vc E C. When it is not clear from the context, we add the subscript XM for the variables in the 7-tuple indicate that they correspond to the particular MDP M. We denote the space of MDPs as M, and use the term MDP and task interchangeably. An agent interacts with an MDP at discrete time steps: the state at time t is St, where So P'uitt the agent chooses action ag, and then the state transitions to S1+1 ~ P(-[st, az) and cost C2 ~ C(-[st, a.) is observed. We denote the history at time t as h4 = { ag Co. $1, d1 C1 ....................... St), and the space of possible histories time as H4. The agent acts based on a history cestory-dependent policy π : H4 + P(A). Every steps, the episode ends and the state is reset, meaning it is sampled again from Pinit- The agent's goal in MDP is minimize the
expected cumulative cost Ex.M Σ ca Note that the MDP horizon H and the horizon T of the
cumulative cost might be different this is important for our subsequent development, where the agent will face an unknown MDP, and can therefore improve its performance between episodes.
2.2 Kernel Density Estimation
Density estimation is the approximation of unknown probability density function f(x), xr Rd based on n i.i.d. samples from it {X, [X,)=_-- 1-1 One popular approach to density estimation is kernel density estimation (KDE (2)). ). Let K (zr) denote kernel function, which is a probability measure over Rdd Let Ho denote d x d positive definite and symmetric matrix satisfying det (Ho) = 1, referred to as a unit bandwidth matrix. Following the formulation of Jiang [16]. we define the
H 1/2 (x x.) Intuitively,
KDE with bandwidth h > 0 as follows: A
ni hd Σ
controls the bandwidth size, while Ho allows for asymmetric bandwidth along different dimensions. We make this clearer with simple example:
Example 1. [Symmetric Gaussian KDE] For a Gaussian kernel function: K(x) (V27)² 1
bandwidth matrix Ho I, and bandwidth h, we obtain the Gaussian KDE: Fa(x)
n 2 Ei= e zha (x.xx"(xxX)) This popular instance of the KDE can be understood as placing Gaussian distribution with zero mean and standard deviation h over each sample X,.
Jiang 16 provided finite sample bounds for the Loo norm of the KDE error under mild assumptions.
Assumption 1.
|| < 30
Assumption 2. The kernel function is spherically symmetric and non-increasing, meaning there exists inn-iicceesing function R2zo -> R220 such that K(x) = =([[=]]) [[]]=]]])o for x E Rdd. Further, the kernel function decays exponentially, meaning there exists p. Cp: to R > 0 such that k(t) < Cp exp (-t),, (2))yt Vtt > to-
Notice that the Gaussian kernel in Example satissies Assumption Assumption 3. ff is Höödde continuous, meaning there exists 0 < a C. > 0, such that [f()) - f (a') S C. || 11xxx2 2'11", Vx, E Rdd.
Theorem 1. [Theorem 2 of Jiang [16]] Under Assumptions 7 there exists positive constant
C = C² (d, II Ca, a, K) such that the following holds with probability at least 1/n,
uniformly in h > (logn) noggn/))1 n)dd and valid unit bundwidths matrices Ho:
h
log
-(()())-(()((__)) <C'
+ V
(1)
a/2
n hd
0 min
where min is the smallest eigenvalue of Ho.
3

---

The first term on the right hand side of ① is a bias term, which can be reduced by reducing h. However, the second term will grow when h is reduced. In general, the sample complexity under an optimal bandwidth scales exponentially in the dimension d (see Lemma 4 for specific example).
2.3 Principal Component Analysis
Principal component analysis (PCA [ (7l) is a popular linear dimensionality reduction technique. For a d-dimensional random variable X, consider nt i..d. samples of X, denoted {Xi [x+)== Reducing the dimension from to is achieved by projection multiplying a vector in IR' by a rank d'
orthogonal matrix P. The expected projection error of is R R(P) E [[[X- PX]]']] Similarly, the empirical projection error of P is R == covariance and empirical covariance of X, respectively. We further denote by A 2 A2 Ad
[||; P We denote by Σ and the
the eigenvalues of Σ, and by X, 2 Ag 2 ... Xa > 0 the eigenvalues of Σ.
The PCA projection Pa is constructed from the cigenvectors corresponding to the d' largest eigenval- of Σ. Its theoretical risk is: 8. E R Pv minpeP R(P), where Pa is the space of d x d orthogonal matrices with rank d' S d, and the expectation is with respect to the n samples. ReiB and Wahl (27] bound the risk for uub-Gaussian random variables.
Assumption 4. X is sub Gaussian, meaning that its second moment is finite and there exists constant Cag such that supx21 1 [|XX ul[]] 1/k S C.gg [[X u)²]+/2 Yu E g
is worth mentioning that any bounded random variable satisfies Assumption
Theorem 2. [Proposition 2.2 of Reiß and Wahl (27]] for a random variable X of dimension d that 640;, (13))
satisfies Assumption 4 we have that 8% &PCA & min 6C', (T
Vn
( -
3 Problem Statement
In this paper make use of parametric representation of MDPs. We suppose there is some parametric space @ and a mapping g θ - .M, which maps parameters to MDPs. We next give two examples of such parametrizations, which will be referred throughout the text illustrate properties of our analysis.
Example 2. ['Tabular mapping] Consider the case where S,A, T, are fixed and finite, while only C' and P differ between different MDPs. The parametric space is defined by θ = Oc x Op, where Oc C R|S|x|Aix CI, and Op C R|S)x|A(x(s [()) such that 10p E Op, Vlc E Oc. Va A and Vs S: 0p(a,8,)) is on the (S) -implex and (c(a, 8, is the -ii-mipplex The parametric mapping is defined such that: P(s'|s, a) 0p(s, a, 8') and C(cls, a) = 0c(8,a,c).
The tabular mapping in Example 2 does not assume any structure of the MDP space, and is common in the Bayesian RL literature, e.g.. in the Bayes-adaptin MDP model of Duff [5]. The next example considers a structured MDP space, and is inspired by the domain in Figure [
Example 3. [Half-circle 2D navigation] Consider the following: S,A = RR,,,nd and P(a, 8, 8)) = 1
for s' s Vs E S, Va E A. The parameter space is θ [0. The parametric mapping C(a,s) 8) = 1, Va E A, Vs E such that ||s - R cos 8, Rin 0] lg I r, where r, R E R.
3.1 The Learning Problem
We first define the cumulative cost for a history-dependent policy π and an MDP M E M as LM, E..M ΣT= Ct We follow the Bayesian RL (BRL) formulation 9 and assume prior distribution over the MDP parameter space f E P()). Our objective is the expected cumulative cost over randomly sampled MDP from the prior:
T-1
( E [L0(0)=]] Ee-s Ea,
Σ
co
1=
'This definition is equivalent to the standard sub- Gaussian definition, P([X) > 1) S exp(1 exp(1-1ppp)) P/(.,)) where Cop and Con differ by universal constant (36).

---

hyper cube is Example 2 where the edge length where we used the parametric representation introduced above. A policy тво E arg min, L (x) termed Bayes optimal. If the prior distribution known, one can calculate the Bayes optimal policy 9 However, in our setting we assume that f is not known in advance, motivating the following learning problem.
We are given a training set of MDPs, {0);- sampled independently from the prior f(0). Our goal is to use these MDPs to calculate policy that minimizes the regret:
R(x) - (T()) - Ly (FBO) = Be-s Br Br,
Σ
ci
ETRO M=g(())
Σ
cE
Note that the expectation above is with respect to the true prior, therefore. the regret can be interpreted as how well the policy calculated from the training data generalizes to an unseen test MDP.
4 Generalization Bounds
We next provide PAC bounds for the learning problem of Section Our general idea first estimate the prior distribution [()) from the training set using KDE, and then solve for the Bayes optimal policy with respect to the KDE-estimated distribution instead of the real prior. For estimator f(0) of the real prior (()), we define the estimated Bayes optimal policy: " E arg min, fz(TT).
We start by showing that we can bound the regret of an estimated Bayes optimal policy, function of the estimation error of the prior itself. The proof, detailed Section A.. in the supplementary, is a simple application of norm inequalities, and exploiting the fact that the total cost is bounded.
Lemma 3. Let J E P(RR)) be an estimator of the real prior J over the parametric space O. We have that: R(x))= (x( (TBO) < 20man T |s 5i, and for a bounded parametric space of volume Q we have: R(TT)) = Ls (T²) L (TBO) S 100ma
=
While the assumption of finite parametric space volume is reasonable in practice, the volume size, appearing in the ||-1log bound that we shall use in the proceeding analysis, grows exponentially with the dimension of O. In Section 4.1 we will discuss how, under some assumptions, we can relax this.
Lemma D provides a convenient framework for bounding the regret using any density estimation technique with a known bound. We now consider the special case of the KDE (cf. Section 2.2) with a Gaussian kernel function and an optimal selection of bandwidth, as calculated lemma. We focus on the Gaussian kernel both for simplicity and because it is a popular choice in practice. Our method can be extended to any kernel function that satisfies assumption 2
the following
Lemma 4. The optimal KDE bandwidth is (up to a constant independent of n) h (log 30क
The following result bounds the estimation error for this case. The proof, detailed in Section A.4 of the supplementary material, is based on Theorem from (16), where the constants are calculated explicitly for the Gaussian kernel case.
Lemma 5. Under Assumptions and 3 for a parametric space with finite volume |0] and KDE with a Gaussian kernel K(u)
Juu Ho = I, and an optimal bandwidth h", we have
(2m)
that with probability at least 1 1/nz supeerd [F((s) s(x)) s Ca (logn) za-a where Ca
C.2 18d, Co. (9))0(9))
6ad
and Amar (0) is the maximal Lı distance between any
/2(2x)
(2x)
twO parameters in A.
Remark 1. Usually, Aman (0) will scale polynomially (or sub- (ob-polynomially) with d. For example, in case the parametric space is a d-dimensional hypercube with edge length B, we have: Amaz = VdB. One parametric case that can be represented [C) In such cases, for large enough the first term in the a-1
B 1, and d | |A| equation for Ca will dominate, and therefore Ca se Co2
By combining Lemma 3 Jand 5 we obtain regret bound - the main result of this section.
5

---

Theorem 6. For prior f() over bounded parametric space Q that satisfies Assumptions 7 and and a Gaussian KDE with optimal bandwidth, we have that with probability at least 1 n:
(log
RI RI To )= S 2a | [O) Ca
Remark 2. While Theorem aasumes a bounded parametric space, the resulting KDE estimate is not necessarily bounded (e.g., when using a Gaussian kernel). In Section[A.3 Af3 the supplementary, we show that the result of Theorem Calso holds when truncating the KDE estimate to a support O.
We next compare Theorem 6 with the the bounds of Tamar et al. [33 We recall that 33 considered model-free approach that learns a
policy Rieg the the training domains, with Le policy regularization. Tamar et al. [33] only considered finite state, action, and cost spaces, a where the parametric space dimension |A| [C). Corollary 1 in [33 shows that with probability at least 1 1/n,
setting equivalent to our tabular representation in Example 2
is d = |s? |A| |S|
Rr (f*) S p n 'T 2/Pm + AVE + +30maxx)
loga 2m
where p 22 222
m T [A],
---
and q = sup ".M, At first sight, the exponent 204a intuitive reason for this is that Theorem dependence on the dimension d when estimating the density f(0). Yet, one must also consider the constants. Assuming that is finite places a severe constraint on the space of possible MDPs essentially, each cost and transition must be possible in all MDPs in the prior! Tamar et al. B3) claim that q may be made finite by adding small noise to every transition, but in this case becomes (((S](C), leading to exponential O (((()C))))) term in the bound of [33, while our bound scales linearly with T. The intuitive reason for this exponential dependence on T is that when learning a history-dependee policy, the space of possible histories grows exponentially with T.
Pax (s.ss)/ Prre/sesff (slc)s.a)
a in our bound compared to the á in 33 is significantly worse. The
on KDE, where is natural to expect an exponential
To summarize, for a small T, and when the structure of M is such that q is finite, the model-free approach of (33] seems preferable. However, when the dimension d is small, and T is large, our model based approach has the upper hand.
Remark 3. Tamar et al. [33] also considered case where M finite set of size [M], and in this case, their Corollary 2 shows that for 0 < < 1, with probability 1 1/1n,, RR (+") S
4803
IAI
Tin 338 where Ponin MinNEM P(M) In the case of a discrete and finite
PPeain
parametric space there is no need for the KDE, but we can still use our model based approach, and estimate the prior using the empirical distribution Pemp (M) = n(M)/n, VM E M,, where i(M) is the number of occurrences of M in the training set. We can bound the Lı error of this estimator using the Bretagnolle-Huber Carol inequality (35) (the full proof is outlined in Section A.6 in the supplementary), and obtain that with probability at least 1 - 1/ne,, we have that R Pamp
S 2Cmaas /2 (ax log (n 20 + [M] + /n. Observe that the dependence on T and Omar in this bound is better than in the bound of (33), and we do not have the Pmai term in the denominator, which could be very small if some MDPs in the prior are rare, and is at most M// Ignoring the log term, our dependence n is also better for every >>> 0.
We conclude this section by pointing out that, as exemplified in Remark 3 our approach can be generalized to density estimation techniques beyond KDE, so long their error can be bounded.
4.1 Bounds for Parametric Spaces with Low Dimensional Structure
The sample complexity in the general bound of Theorem O grows exponentially with the dimension of the parameter space O. In many practical cases however, such as the HalfCircle domain of Example 3 there may be low dimensional representation that encodes most the important information the tasks, even though the dimensionality of the parametric space is higher. In such cases, we expect that our bounds can be improved to depend on the low dimensional representation. In the following, we
'hee bound in 33 also has a term X for the weight of the L2 regularization term. Here, we present the result for the optimal X, which can be calculated similarly to the optimal bandwidth in Lemma
We present the result of (33 for the optimal regularization coefficient, and ignore subleading n term.
t-e

---

approach this task by combining our model-based approach with the PCA dimensionality reduction method. PCA is a linear method, and allows for a relatively simple analysis to demonstrate our claim. In practice, a non-linear method may be preferred. In Section O we verify empirically that our approach also works with non-linear deep neural network based dimensionality reduction.
We propose the following procedure.
1. Reduce the training set dimensionality from d to d' using PCA
2. Perform a Gaussian KDE estimation with optimal bandwidth in the low dimension d'
3. Project back the estimated distribution to dimension d, to obtain the prior estimator Pd'
The main difficulty in analysing this procedure, however, is calculating the error of the estimated distribution in dimension d, that is, after the projection step. We remark that projecting back is necessary, as calculating the estimated Bayes optimal policy z requires g(0), where 0 is in dimension d. To set the stage, we first need to define the projection step explicitly.
The PCA dimensionality reduction can be written Pa WT. WI, where WT Rd''dd For any 0 E θ, we therefore have that pg = WT- - 0 is the low dimensional representation of 8, and WT 0r is the projection of o,, back to the d-dimensional space. For each H,, denote its inverse image follows: 0 (8L) {0 θ WL 0 82 }. By the law of total probability, the probability distribution of low-dimensional 0r is: fer (0z) = Eo (0L) fe(0))d0
For the proceeding analysis, require the function that maps parameters MDPs to be smooth. Assumption 5. In the MDP space M, only P and C can differ. Furthermore, the parametric mapping is Lipshitz continuous with respect to Pas (s,- 8, 8,a), ie:: 3C, such that Vs E &,a E A. [PM=-(00)) (v.- s, a)- PMs=g(22) (.,- s, a) 1 S Cg 1 02l1
We now extend the well known Simulation Lemmna [17] to the case of a distory-dependenn policy. This will allow us later to relate the error in the prior distribution to the error of the policy.
Lemma 7. For any eestorydeppenden policy and any parametric mapping g that satisfies Assump- tion the following holds for any 01, 02 E O:
[.M=g(0r).a == S Cm ||0 02111 TT2.
We note that Assumption 5 measures closeness between the joint distribution costs and transitions. This is slightly different than the simulation lemma for a Markov policy [17]. where only the absolute difference between the costs is required. Intuitively, this is since history-dependent policy can depend on observed costs, therefore even if two observed costs are very close in magnitude, the actions resulting from observing them could be very different. We note that in the case where P' and C' are bounded (which is always true when S, A and C' are discrete), it is enough to assume Lipschitz continuity of g with respect to P and C' separately. We also note that there may be parameter spaces that satisfy Eq. 2 without satisfying Assumption 5) in such cases our proceeding results will still hold.
would like to formally consider MDP spaces with a low-dimesioonal structure. This is captured by assuming that the MDPs essentially lie on a linear subspace of dimension d', such that the magnitude of the variability of MDPs outside this subspace is bounded by e.
In the following,
Assumption 6. There exist d' < d and e R20 such that A S 6, where Xi as defined in Section 2.3 is the i'th largest eigenvalue of the Covariance matrix of 0.
Under Assumption 6 and using the PCA bounds of Reiß and Wahl (27] (ef. Theorem we can bound the dimensionality reduction error due to performing PCA on the sampled MDPs in our data, as given by the following lemma. Lemma 8. Under Assumptions
and we have that:
E
Pa
S min
8C22 (etr(E)
64C ( "Ayy -
+ e - (d-d').
Vn
There are two terms to the bound in Lemma 8 finite number of samples, and decays n increases. The second is due to the fact that even for perfect PCA there is some error, as the MDPs do not lie perfectly in a low dimensional subspace.
The first is due to the error in performing PCA with a
We next assume that after the PCA projection, the distribution remains Hölder continuous.
7

---

Assumption 7. for is a' -Höler continuous, meaning there exists 0 < a' S 1, Co > 0, such that (f(Or) - (0,)) S Co' 100,,000,, - 01111 vo,, 0/, E E.,.
Theorem 9. Let fd be the approximation of f as defined above. Under Assumptions we have with probability at least 1 1/n:
and
Rr (miØe)
S
20man' 16 Ca' (logn) za + 20 C co/ min
Cyy vn
Cht c(d
d(),
where Ca' = C 2 10d /o, Ameer
84d ToLl
(2m) 4 C2g = tr (E)), and Ax, Ar -
(2(2m)
The first term in the bound of Theorem 2 is the KDE error. Note that, compared to the KDE error in Theorem 6. the exponential dependence is on the low dimension d', and not on the higher dimension d. The second term in the bound is due to the PCA error, as discussed after Lemma This result demonstrates the potential in our approach, which accounts for structure in the parametric space.
5 Related Work
Meta RL has seen extensive empirical study [4] 39 7 26], and the connection to Bayesian RL has been made in a series of recent works [20 14 during training (effectively drawing different random MDPs from the prior at each training iteration),
25 39 Most meta RL studies assumed infinite tasks
with few exceptions, mostly in the offline meta RL setting 3 22).
Bayesian RL algorithms such as [12 10 sample MDPs from the true posterior, which requires knowing the true prior. In their comprehensive study, Simehowitz et al. [32] analyse family of pnsterior-sampling based algorithms, termed n-Monte Carlo methods, under mis-specifed priors, and bound the corresponding difference in accumulated regret. In contrast, our work considers a zero-shot setting, comparing the Bayes-optimal policies with respect to the true and the mis siiisspeeieed prior, and, more importantly. focuses on the sample complexity of obtaining prior that is accurate enough an issue that is not addressed in (32 for MDPs. To the best of our knowledge, the only theoretical investigation of meta RL with finite training tasks in a similar setting to ours is the model free approach in [33), which we compare against in our work. Our theoretical analysis builds on ideas from the study of density estimation (29: [6] and dimensionality reduction [ 27].
We note that regularization techniques inspired by the mixup method (38 have been applied to meta learning (37,, and recently also to meta RL (21) with a goal of improving generalization to out-of-distribution tasks. In our experiments, we compared our approach with an approach inspired by (2T) and found that for in-distribution generalization, our approach worked better. That said, we believe there is much more to explore in developing effective regularization methods for meta RL. We conclude with studies on PAC-Bayes theory for meta learning 22, 30 6). These works, which have a different flavor from our PAC analysis, do not cover the meta RL problem considered here.
Experiments
In this section we complement our theoretical results with empirical investigation. Our goal is to show that our main idea of learning a KDE over a low dimensional space of tasks is effective also for state-of-the-arn meta-RL algorithms, for which the linearity assumption of PCA clearly does not hold, and computing the optimal yet intractable r is replaced with an approximate deep RL method.
Modern deep RL algorithms are known to be highly sensitive many hyperparameters [3], and meta RL algorithms are not different. To demonstrate our case clearly, we chose to build on the VariBAD algorithm of Zintgraf et al. [39], for which we could implement our approach by replacing just a single algorithmic component, as we describe next, and thus obtain a fair comparison.
We briefly explain the VariBAD algorithm, to set the stage for our modified algorithm to follow; we refer the reader to 39 for the full details. VariBAD composed of two main components. The first is variational autoencoder (VAE [9] with recurrent neural network (RNN) encoder that, at each time step, encodes the history into Gaussian distribution over low dimensional latent vectors. The decoder part of the VAE is trained to predict the next state and reward from the encoded history.
VariBAD:
8

---

HalfCircle 20 Real Environments
HalfCircle 30 Real Environments
192
290
343
340
Average Eval Retur
Average Eval Retur
022
220
323
300
WETTOOD SNOTT A77O
INTERIA PTEERO--ER CTVO
Frames
Frames
187
Figure 2: Average return on HalfCirele with KDE and mixup dream environments and without dream environments. The average is shown in dashed lines, with the 95% confidence intervals (15 random seeds). We do not show the intervals for the mixup run for visualization clarity; mixup obtained similar intervals as without dream. The full comparison is in Section A. of the supplementary.
Using the terminology in our paper, the VAE latent state can related to the MDP parameter 0, and the VAE decoder learns a model of g(0). The second component in VariBAD is a policy, mapping the output of the VAE encoder and current state into an action. This component, which can be seen as producing an approximation of T is trained using policy gradient algorithm such PPO 31
VariBAD Dream:
Recall that our pipeline is to learn a KDE over the task parameters 0, and then train a policy on tasks from the estimated KDE. Unfortunately, in our meta-RL. setting, we do not assume that we directly know the representation for each task. However, the VAE in VariBAD allows for a convenient approximation. We can think of the output of the VAE encoder at the end of an episode, after a full history has been observed and the uncertainty about the task has been resolved sample of the task parameter O. Thus,
as best as possible, over these variables. We henceforth refer to samples from the KDE as dream environments (as they are not present in the real data), and we note that the VAE decoder can be used to sample rewards and state transitions from these environments. Thus, we can train the VariBAD policy on both the sampled training environments, and also dream environments. We refer to this method as VariBAD Dream. In our implementation. we train the KDE and VariBAD components simultancously. The full implementation details and pseudo code are in Section A.10 of the supplementary.
propose to build our KDE estimate
We emphasize that in the original VariBAD work (39,, the algorithm was trained by sampling different tasks from the task distribution at each iteration, corresponding to an infinite number of traning tasks. As we are interested in the finite task setting, our implementation of VariBAD (and VariBAD Dream) draws a single batch of Nerain tasks once, at the beginning of training, and subsequently samples tasks from within this batch for training the VAE and policy.
We consider the HalfCirele environment of Figure popular task that requires learning a non-trivial Bayes-optimal policy (3). In order to test generalization, we evaluated the agents on Nevat = 50 sampled environments, different from the Ntrain training ones. In Figure 2e show the affect of incorporating the dream environments on the average return for the evaluation environments. Evidently. our approach achieved higher return, both for Ntrain = 20 and Ntrain 30. For Ntrain 40, the original VariBad performed near optimally, and there was no advantage for VariBad Dream to gain. For Ntrain = 10, the sampling was too sparse, and both methods demonstrated comparable failure in generalizing to the test environments. We emphasize that while the performance advantage of VariBAD Dream is modest, it is remarkable that the KDE regularization demonstrates consistent improvement, as the VAE training and PPO optimization already include significant implicit and explicit regularization mechanisms.
Results:
where dream environments are created by a random weighted average of latent vectors (instead of the KDE). Interestingly, our KDE approach outperformed this method, which may be the result of the geometry of the task distribution there are no goals within the half circle, in contrast with the experiments in [21]. where the goals were distributed inside a rectangular area. We remark that the experiments in [2 were designed to
Additionally, we evaluated a method inspired by mixup 38
evaluate out-of-d nut-of-distriuutiin generalization, different from the in distribution setting considered here. In Sections A. and A.12 of the supplementary we provide additional evaluations and show similar
results on MuJoCo 34] environment with high dimensional state space.

---

Synthetic Experiments with a Known Parametric Space: knowledge of both the parameters the training MDPs, and the function g(0) that maps parameters to MDPs. We complement our empirical investigation with synthetic experiments with VariBAD Dream, where g(0) is known.
In our theoretical analysis, we assumed
We consider again the HalfCircle environment of Figure I We define the parametric space as è = I2, where g(0) prescribes a corresponding MDP with a goal at location (x, y) O. For our experiment, we sample {0 i.i.d. from the training environment distribution, and perform KDE directly on these samples. Our variant of VariBad Dream samples parameters from the KDE and generates dream MDPs using the known g
Figure compares VariBAD Dream with the conventional VariBad (trained only {0.) on unseen test environments. Evidently, the use of KDE improved the generalization performances of VariBad by large margin, especially when the number of training environments is small. Note that the improvement is starker in this synthetic setting compared to Figure 2 as VariBAD Dream here does not suffer from inaccuracies due to learning a model of g(0).
HaitCirde- Real Eentron
HellCircle Real Envircrements
HellCrrcee 20 Real Ervironments
Soveage
Promes
Fremees
Framee
Figure 3: Average return on half-eircle of the original VariBad and the VariBAD Dream variant with access to the MDP parametric space. The average is shown in dashed lines, with the 95% confidence intervals (using 6 random seeds).
7 Discussion & Future Work
We propose model-based scheme for meta-RL with finite sample of training tasks, where we first estimate the prior distribution of tasks, and train a Bayes optimal policy on the estimated prior. Using KDE for density estimation, we obtain state-of-the-art PAC bounds. Further, our approach can exploit low dimensional structure of the task distribution, when such exists, to obtain improved bounds. Finally, we showed that our approach can be "plugged-in" the VariBAD algorithm to improve generalization.
A key takeaway from analysis is that the dimensionality of the task distribution determines, with exponential dependency, the amount of training samples required to act well. This insight provides a rule-of-thumb of when meta RL approaches based on task inference, such as VariBAD, expected to work well. Indeed, recent empirical work by Mandi et al. (23) claimed that in benchmarks such as RLBench [15] where tasks are very diverse, simpler meta RL methods based on fine-tuning a policy trained on diverse tasks display state-of-the-art performance.
There several exciting future directions for investigation. Our theoretical analysis can be extended to more advanced density estimation and dimensionality reduction techniques, such as VAEs (28). Our empirical investigation hinted that regularization, such affected by VariBAD Dream, can improve deep meta-RL algorithms. More sophisticated regularization can be developed based on prior knowledge about the possible tasks.
Acknowledgements
This work received funding from Ford Inc., and from the European Union (ERC, Bayes-RL, Project Number 101041250). Views and opinions expressed are however those of the author(s) only and do not necessarily reflect those of the European Union or the European Research Council Executive Agency. Neither the European Union nor the granting authority can be held responsible for them.
10

---

References
[1] Hervé Abdi and Lynne J Williams. Principal component analysis. Wiley interdisciplinary reviews: computational statistics, 2(4))433-459) 2010.
[2] Ron Amit and Ron Meir. Meta-leaming by adjusting priors based extended pac-bayes theory.
In International Conference on Machine Learning, pages 205-214. PMLR, 2018
[3] Ron Dorfman, Idan Shenfeld, and Aviv Tamar. Offline meta learning of exploration. In Advances in Neural Information Processing Systems, 2021.
[4] Yan Duan, John John Schulman, Xi Chen, Peter L Bartlett, Ilya Sutskever, and Picter Abbeel. RL?: Fast reinforcement learning via slow reinforcement learning. arXiv preprint arXiv:161..0277 2016.
[5] Michael O' O'Gordon Duff. Optimal Learning: Computational procedures for Bayes-adaptive Markov decision processes. PhD thesis, University of Massachusetts at Amherst, 2002.
[6] Alec Farid and Anirudha Majumdar. Pac-bus: Meta-learning bounds via pac- bac-s and uniform stability. In Advances in Neural Information Processing Systems, 2021.
[7] Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adap- tation of deep networks. In Proceedings of the 34th International Conference on Machine Learning-Volue 70, pages 1135. JMLR. org, 2017.
[8] Francisco Garcia and Philip S Thomas. A meta-mdp approach to exploration for lifelong reinforcement learning. Advances in Neural Information Processing Systems, 32, 2019.
[9] Mohammad Ghavamzadeh, Shie Mannor, Joelle Pineau, and Aviv Tamar. Bayesian reinforce- ment learning: A survey. arXiv preprint trXivv:609 04436, 2016.
160..046999,,,
[10] Divya Grover, Debabrota Basu, and Christos Dimitrakakis. Bayesian reinforcement learning via deep, sparse sampling. In International Conference on Artificial Intelligence and Statistics, pages 3036-3045. PMLR, 2020.
[11] Shixiang Gu, Ethan Holly, Timothy Lillicrap, and Sergey Levine. Deep reinforcement learning for robotic manipulation with asynchronous off-policy updates. In 2017 IEEE international conference on robotics and automation (ICRA), pages 3389-3396. IEEE, 2017.
[12] Arthur Guez, David Silver, and Peter Dayan, Efficient bayes-adaptive reinforcement learning using sample-based search. In Advances in neural information processing systems, pages 1025-1033, 2012.
[13] Peter Henderson. Riashat Islam, Philip Bachman, Joelle Pineau, Doina Precup, and David Meger. Deep reinforcement learning that matters. In Proceedings of the AAAI conference on artificial intelligence, volume 32, 2018.
[14] Jan Humplik, Alexandre Galashov, Leonard Hasenelever, Pedro A Ortega, Yee Whye Teh, and Nicolas Heess. Meta reinforcement learning as task inference. arXiv preprint 195444190,.06224, 2019.
[15] Stephen James, Zicong Ma, David Rovick Arrojo, and Andrew J. Davison. RLBench: The robot learning benchmark & learning environment. IEEE Robotics and Automation Letters, 2020.
[16] Heinrich Jiang. Uniform convergence rates for kemne density estimation. In Doina Precup and Yee Whye Teh, editors, Proceedings of the 34th International Conference on Machine Learning. volume 70 of Proceedings of Machine Learning Research, pages (694-1703. PMLR, 06-11 Aug 2017.
[17] Michael Kearns and Satinder Singh. Near-optimal reinforcement learning in polynomial time. Machine learning. 49(2):209-22) 232, 2002.
[18] Alex Kendall, Jeffrey Hawke, David Janz, Przemyslaw Mazur, Daniele Reda, John-Mark Allen, Vinh-Dieu I Lam, Alex Bewley, and Amar Shah. Learning to drive in a day. 2019 International Conference Robotics and Automation (ICRA), pages 22544254. IEEE, 2019.
11

---

[19] Diederik P Kingma and Max Welling, Auto-encoding variational bayes. In Advances in Neural Information Processing Systems, 2013.
[20] Gilwoo Lee, Brian Hou, Aditya Mandalika, Jeongseok Lee, Sanjiban Choudhury, and Sid- dhartha S Srinivasa. Bayesian policy optimization for model uncertainty. In International Conference on Learning Representation, 2019.
[21] Suyoung Lee and Sae- Young Chung. Improving generalization in meta-rl with imaginary tasks from latent dynamics mixture. In Advances in Neural Information Processing Systems, 2021.
[22] Jiachen Li, Quan Vuong, Shuang Liu, Minghua Liu, Kamil Ciosek, Henrik Christensen, and Hao Su. Multi-task batch reinforcement learning with metric learning. In Advances in Neural Information Processing Systems, 2020.
[23] Zhao Mandi, Pieter Abbeel, and Stephen James. On the effectiveness of fine-tuning versus meta-reinforcement learning, 2022.
[24] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wierstra, and Martin A. Riedmiller. Playing atari with deep reinforcement learning. In Advances in Neural Information Processing Systems, 2013.
[25] Pedro A Ortega, Jane X Wang, Mark Rowland, Tim Genewein, Zeb Kurth-Nelson, Razvan Pascanu, Nicolas Heess, Joel Veness, Alex Pritzel, Pablo Sprechmann, et al. Meta-learning of sequential strategies. arXiv preprint 10X.x::9950003000 2019.
[26] Kate Rakelly, Aurick Zhou, Deirdre Quillen, Chelsea Finn, and Sergey Levine. Efficient off-policy renforeeiforre learning via probabilistic context variables. In International Conference Machine Learning, 2019.
[27] Markus ReiB and Martin Wahl. Nonasymptotic upper bounds for the reconstruction error of pca. The Annals of Statistics, 1022):1398-11,,, 2020.
[28] Michal Rolinek, Dominik Zietlow, and Georg Martius. Variational autoencoders pursue pea directions (by accident). In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattem Recognition, pages 12406 10616,1215,,100 2019.
[29] Murray Rosenblatt. Remarks on Some Nonparametric Estimates of Density Function. The Annals of Mathematical Statistics, 27(3):832 837, 1956. doi: 10. 10.1214/a0mm/117777221099
[30] Jonas Rothfuss, Vincent Fortuin, Martin Josifoski, and Andreas Krause. Pacoh: Bayes-optimal meta-lcarning with pac-guarantees. In International Conference on Machine Learning, pages 9116-9126 PMLR, 2021.
[31] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint IrXiv:1707.06347 2017.
[32] Max Simchowitz, Christopher Tosh, Akshay Krishnamurthy, Daniel J Hsu, Thodoris Lykouris, Miro Dudik, and Robert E Schapire. Bayesian decision decision-making under misspecified priors with applications to meta-learning. In Advances in Neural Information Processing Systems, 2021.
[33] Aviv Tamar, Daniel Soudry, and Ey Zisselman. Regularization guarantees generalization in bayesian reinforcement learning through algorithmic stability. Proceedings of the AAAI Conference on Artificial Intelligence, 36(8):8423-844 2022.
[34] Emanuel Todorov. Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In 2012 IEEE/RSJ International Conference on Intelligent Robots and Systems, pages 5026-5033. IEEE, 2012.
[35] Aad W Vaart and Jon Wellner. Weak convergence. In Weak convergence and empirical processes, pages 16-28. Springer, 1996.
[36] Roman Vershynin. Introduction to the non-asymptotic analysis of random matrices. arXiv preprint 2010.
12

---

[37] Huaxiu Yao, Linjun Zhang, and Chelsea Finn. Meta-learning with fewer tasks through task interpolation. In International Conference on Learning Representation, 2022.
[38] Hongyi Zhang, Moustapha Cisse, Yann Dauphin, and David Lopez-Paz. mixup: Beyond empirical risk minimization. In International Conference Learning Representation, 2018.
[39] Luisa Zintgraf, Kyriacos Shiarlis, Maximilian Igl, Sebastian Schulze, Yarin Gal, Katja Hofmann, and Shimon Whiteson. Varibad: A very good method for bayes-adaptive deep via meta- learning. In International Conference on Learning Representation, 2020.
Checklist
The checklist follows the references. Please read the checklist guidelines carefully for information how to answer these questions. For each question, change the default [TODO] to [Yes], [No], , or [N/A] You are strongly encouraged include justification your answer, either by referencing the appropriate section of your paper or providing a brief inline description. For example:
Did you include the license the code and datasets? [Yes] See Section
Did you include the license the code and datasets? [No] The code and the data proprictary.
Did you include the license to the code and datasets? [N/A]
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's
contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes]
(c) Did you discuss any potential negative societal impacts of your work? [N/A]
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [Yes]
(b) Did you include complete proofs of all theoretical results? [Yes] In the supplementary materials
3. If you ran experiments...
(a) Did you include the code, data, and instructions needed reproduce the main perimental results (either in the supplemental material or as URL)? [Yes] In the supplementary materials
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes] In Bection[A.10jo the supplementary
(c) Did you report error bars (c.g., with respect the random seed after running experi-
ments multiple times)? [Yes] In Section and Section A. A.11 in the supplementary (d) Did you include the total amount compute and the type of resources used (e.g.. type of GPUs, internal cluster, cloud provider)? [Yes] See Section A.13 of the supplementary
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) your work uses existing assets, did you cite the creators? [Yes] See Section A. the supplementary
(b) Did you mention the license of the assets? [Yes] See Section of the supplementary
(c) Did you include any assets either in the supplemental material URL? [Yes] See Section o0 the supplementary
(d) Did you discuss whether and how consent was obtained from people whose data you're using/curating? [N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]
13

---

5. If you used crowdsourcing conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, applicable? [N/A]
(b) Did you describe any potential participant risks, with links Institutional Review Board (IRB) approvals, if applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]
14


Reviews:
[{"review_id": "tkSV3QlBoN3", "reviewer": "Reviewer_6rr7", "summary": "In this paper, PAC bounds for meta reinforcement learning with finite training tasks are derived. The derived bounds is applicable to both continuous and discrete state and action spaces. The derived bounds have exponential dependence on the dimensionality of the task space. To address this, the analysis is extended to the case with dimensionality reduction using principal component analysis. To support the theoretical results, the technique to improve the existing meta-RL method was evaluated in a simple environment.", "questions": "- In practice, how should we determine the optimal bandwidth of the Gaussian kernel for KDE? How did the authors determined the optimal bandwidth for the experiment?\n- I did not clearly understand how practical the problem setting of meta RL with finite training tasks is. I would appreciate it if the authors provide examples in the real world where meta RL with finite training tasks should be considered.", "limitations": "I think that the authors clearly described the limitation of their theoretical analysis. For example, it is clearly described that analysis regarding the dimensionality reduction is limited to PCA, which is a linear method. ", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strength:\n- Novel and practical PAC bounds for meta-RL are derived\n- Unlike the bound proposed by Tamar et al. (2021), the derived bound can be applied to both continuous and discrete state action spaces\n- In addition to theoretical results, experimental results are reported\n\nWeakness\n- The work is largely based on the prior work by Tamar et al. (2021), and one might think that the contribution is incremental\n- The experiments seem simple, while I think that the theoretical results are strong enough\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": "l8wFJ4yGdXV", "reviewer": "Reviewer_ST9x", "summary": "This paper can be separated into two quite different sections. \n\nThe theoretical part derives complexity bounds for an algorithm that would have access to the MDP’s parameters (noted theta), that would run a PCA in order to reduce the dimension of that parameter’s space, and learn the density as modeled by an independent gaussian distribution (presented as an example of KDE). \nThen, the Bayes-optimal policy can be computed from this learned distribution under better PAC criteria than by using the naive density approximation in the high-dimensional space.\n\nThe experimental part considers data augmentation of the training set of MDPs by using the latent space learned by VariBAD. One can link this approach to the theoretical part by comparing the space reduction performed by the VAE of VariBAD with a PCA, \nwhere the high-dimensional parameter is the history itself (through enough episodes to reduce uncertainty), and the low-dimensional projection is the latent state at the end of an episode.\n", "questions": "L 263: why projecting back? It should be sufficient and even easier to train an agent using the reduced representation of priors. ", "limitations": "Not applicable.", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The theoretical part looks correctly driven, and brings up an interesting proven result: it is worth reducing the dimension to concentrate on the low-dimensional manifold that structures the set of MDPs. \n\nHowever, the experimental part is rather disappointing, as it does not directly verify the first part’s claims, and treats a very different setting (where theta is not observed) with quite a different algorithm, despite the fact that this separated part reveals another interesting result: one can use the VAE’s latent space to augment the data and improve generalisation with less examples. Also, this part is missing higher dimensional MuJoCo environments in order to fairly compare with VariBAD.\n\nIn my opinion, this paper could even be divided into two papers that would be worth separated publication: \n\n1) the first including this theoretical part with a good experimental study: where theta is observed, using both discrete and continuous environments (including the half-circle with function approximation), \nAnd comparing to naive baselines (ppo/sac with concatenated state + theta observation, KDE without PCA dimension reduction)\n\n2) the second including the experimental part with a proper method / algorithm description, and adding larger envs (mujoco). Eventually, a theoretical explanation that would rather look at the improvement as the result of a data augmentation.\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": "JdnWxn1TAl8", "reviewer": "Reviewer_ZFq8", "summary": "The paper focuses on designing PAC meta RL algorithms. While prior works took a model-free approach of learning regularized history conditioned policies, this work first estimates the underlying task distribution from the training task and then trains a meta RL algorithm on the learned task distribution. The resulting approach obtains state-of-the-art PAC bounds and shows that the approach can be used to improve practical meta RL algorithms (such as VariBAD).", "questions": "N/A", "limitations": "The paper doesn't explicitly discuss its limitations", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths:** \n  1. The paper uses KDE techniques to improve the PAC bounds\n  2. The resulting approach is simple which can be used to improve existing meta RL algorithms such as VariBAD\n\n**Weaknesses:**\n  1. While the approach obtains better PAC bounds, my concern is that KDE estimation will only work on tasks which have some low-dimensional embedding (such as the reaching task). I don’t think this approach of estimating task distribution and then learning from the task distribution will scale up to more complex settings (say, meta-world ML10 benchmark, Yu et al., 2021, https://arxiv.org/pdf/1910.10897.pdf).\n\nOverall, I do think that the strength outweighs the weakness of the paper. However, I should say that PAC bounds for meta RL isn’t exactly my area of expertise and hence my opinion should be taken with a grain of salt.\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": "0s8dkmggk-", "reviewer": "Reviewer_RVe9", "summary": "This work proposes an imaginary task generation method for meta-RL tasks with a limited number of training tasks using the prior estimated by kernel density estimation (KDE). The authors provide a theoretical bound for the regret of the estimation and empirically evaluate the method in conjunction with variBAD [1] on the HalfCircle environment. \n\n[1] Zintgraf et al., Varibad: A very good method for bayes-adaptive deep RL via meta-learning.", "questions": "1. Is there any additional methods the authors tried to improve the generalization of the generated tasks? Because it seems that the policy may overfit the KDE-estimated prior. E.g., generating contexts with perturbation on the prior. \n\n2. It would be great if the authors provide a table of all notations used in this paper in Appendix, as some of the notations are confusing and similar. \n\n3. It would be great if the authors review the reference and update them to the recent conference proceedings. \ne.g., [18] arXiv->ICLR2019, [19] CorR->NeurIPS 2021, [20] arXiv->NeurIPS 2020...\n\n", "limitations": "This work is based on in-distribution parametric meta-RL with finite training tasks.\nAlthough it is not the focus of this work in the current state, it is hard to generalize to non-parametric, non-stationary, and out-of-distribution meta-RL tasks [3].\n\n[3] Yu et al., Meta-World: A Benchmark and Evaluation for Multi-Task and Meta Reinforcement Learning", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "***Strengths***\n\nOverall, I think the strength of this study lies in its novelty and theoretical analysis.\n\n1. Novelty\n\nThere are some previous works that generate imaginary tasks using the latent model as LDM [2]. However, this is the first work that directly estimates the density using KDE and uses the density to sample new tasks. This novel idea sounds reasonable and effective for parametric meta-RL tasks with finite tasks.\n\n2. Solid theoretical bound\n\nI appreciate the authors for the solid derivations for the regret bounds in section 2.2. Although the conclusion that the bound is proportional to the dimension of the parameter space is somewhat natural, it has a great meaning that it was shown through a solid proof. Especially the use of PCA to reduce dimension is well supported by this theory.\n\n***Weaknesses***\n\nOn the other hand, I think there are weaknesses in practical application and experimentation. I think the contribution of this work can be more significant when the empirical evaluation and analysis are conducted as follows. \n\n1. Ambiguous link between theory and algorithm\n\nI felt the transition from section 4 to section 6 to be somewhat abrupt. Maybe providing a pseudo-code of variBAD+KDE would prevent confusion. It is unclear whether the proposed method would work on the latent space generated by the encoder of VAE, which may deviate from the assumptions made. \n\nI know that the VAE of variBAD [1] also does some implicit density estimation based on multivariate Gaussian distribution initialized to standard normal distribution. Due to the nature of the VAE, the latent space generated by variBAD may already be regular (continuous and complete). Maybe applying the KDE on top of this latent space makes the latent space more regular?  One possibile way to show this is an empirical comparison on the latent space for unseen tasks as the t-SNE plot in [2], Figure 9. \n\nAlso, I'm curious how regular and general the generated tasks will be given unseen context or states without regularizations as a dropout in [2]. I appreciate the visualization in Figure 4, but I'm not sure if the reward decoder would generate high rewards for states that were not rewarding during training for any training tasks. Maybe 20~30 tasks are enough to contain all states in a light blue circle in Figure 1?\n\nAs the authors claim, it could have been because of the special geometry of the HalfCircle's task distribution. Therefore I consider testing only one environment as a weakness of this work.\n\n2. Limited experiment domains\n\nThe HalfCircle environment in this work is a very meaningful and intuitive domain. However, it is limited to a one-dimensional task space, therefore it is hard to conclude that the proposed method would outperform generally. Maybe the proposed method works well by overfitting the policy for the learned task prior. To emphasize the contribution of this work for finite training tasks, I strongly recommend evaluating at least tasks evaluated by variBAD [1] or out-of-distribution tasks evaluated by LDM [2].\n\n\n[1] Zintgraf et al., Varibad: A very good method for bayes-adaptive deep RL via meta-learning.\n\n[2] Lee et al., Improving Generalization in Meta-RL with Imaginary Tasks from Latent Dynamics Mixture.", "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