Meta Reinforcement Learning with Finite Training Tasks - a Density Estimation Approach
Anonymous Author(s)
AfﬁliationAddress email
Abstract
In meta reinforcement learning (meta RL), an agent learns from a set of trainingtasks how to quickly solve a new task, drawn from the same task distribution.
Theoptimal meta RL policy, a.k.a.
the Bayes-optimal behavior, is well deﬁned, andguarantees 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 toguarantee approximately optimal behavior with high probability.
Recent workprovided the ﬁrst such PAC analysis for a model-free setting, where a historydependent policy was learned from the training tasks.
In this work, we propose adifferent approach: directly learn the task distribution, using density estimationtechniques, and then train a policy on the learned task distribution.
We show thatour 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-dimensionalmanifold, we extend our analysis to use dimensionality reduction techniques andaccount for such structure, obtaining signiﬁcantly better bounds than previouswork, which strictly depend on the number of states and actions.
The key of ourapproach is the regularization implied by the kernel density estimation method.
Wefurther demonstrate that this regularization is useful in practice, when ‘plugged in’the state-of-the-art VariBAD meta RL algorithm.
Introduction
In recent years, reinforcement learning (RL) became the leading algorithmic framework for avariety of domains, including computer games [21], robotic manipulation [10] and autonomousdriving [16].
Popular RL algorithms, however, are characterized by a high sample inefﬁciency, dueto the exploration-exploitation problem – the need to balance between obtaining more informationabout the environment versus acting based on such information.
Indeed, most RL success storiesrequired a very long training process, only possible in simulation.
For an agent to learn fast, additional structure about the problem is required.
In Meta RL [3; 35; 6; 23],agents are allowed to train on a set of training tasks, sampled from the same task distribution as thetask they will eventually be tested on.
The hope is that similar structure between the tasks could beidentiﬁed during learning, and exploited to quickly solve the test task.
It has recently been observedthat the meta RL problem is related to the Bayesian RL formulation, where each task is modelled as a
Markov decision process (MDP), and the distribution over tasks is the Bayesian prior [22; 35].
Inthis formulation, the optimal meta RL policy is the Bayes-optimal policy – the policy that maximizesexpected return, where the expectation is taken with respect to the prior MDP distribution.
While signiﬁcant empirical progress in meta RL has been made, a theoretical understanding ofthe problem is still limited.
A central question, which we focus on in this work, is the probablyapproximately correct (PAC) analysis of meta RL – how many training tasks are required to guarantee performance that is approximately Bayes-optimal with high probability .
Practical motivation forstudying this question includes the lifelong learning setting [7], where training tasks are equivalentto tasks the agent had previously encountered, and we would like the agent to learn as quickly aspossible, and the ofﬂine meta RL setting [2], where training data is collected in advance, and it isimportant to estimate how much data is needed.
Recently, Tamar et al.
[29] proposed the ﬁrst PAC bounds for meta RL, using a model free approach.
In their work, a history-dependent policy was trained to optimize the return on the set of training
MDPs, where policy-regularization was added to the loss of each MDP in the data.
The bounds in[29] scales with the number of states exponentiated by the length of the history, intuitively due tothe number of possible histories that can be input to the policy.
In this work, we propose a modelbased approach for PAC meta RL.
Our idea is that instead of regularizing the policy during training,as in [29], we can learn a 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, ifwe can guarantee that the learned distribution is ‘close enough’ to the real prior, we should expect thelearned policy to be near Bayes optimal.
In this work, we pursue such guarantees using techniquesfrom the kernel density estimation (KDE) literature [26].
Building on PAC results for KDE, we derive PAC bounds for our model based meta RL approach.
Compared to the bounds in [29], our results require much less stringent assumptions on the MDPprior, and apply to both continuous and discrete state and action spaces.
Interestingly, our boundshave a linear dependence on the horizon, compared to the exponential dependence in [29], but have anexponential dependence on the dimension of the prior distribution, corresponding to the exponentialdependence on dimensionality of KDE.
We argue, however, that in many practical cases the priordistribution lies in some low dimensional subspace.
For such cases, we extend our analysis to usedimensionality reduction via principal component analysis (PCA), and obtain bounds where theexponential dependence is on the dimension of the low dimensional subspace.
To visualize the advantage of our approach, considerthe semi-circle domain in Figure 1, adapted from [2]:a 2-dimensional agent must navigate to a goal, located somewhere on the semi-circle.
A task thereforecorresponds to the goal location, and the task distribution is uniform on the 1-dimensional semi-circle.
The bounds in [29] depend on the number of statesand actions, which are continuous here, and even ifdiscretized, would result in excessive bounds.
Intuitively, however, it is the low-dimensional structureof the task distribution that should determine the difﬁculty of the problem, and our bounds can indeedaccount for such.
We remark that many benchmarkmeta RL task distributions exhibit similar low dimensional structure [3; 6; 2].
To complement our theoretical results, we demonstrate the practical potential of our approach by incorporating it in the VariBAD meta-RL method of
Zintgraf et al.
In VariBAD, a variational autoencoder (VAE) is used to learn a low-dimensionallatent representation of the task distribution, and a deep neural network is trained to approximatethe Bayes optimal policy.
We show that by estimating the task distribution using our kernel densityestimation method, applied to the VAE latent space, and training the policy on sampled tasks fromthis distribution, we improve the generalization of the policy to tasks not seen in the training data.
Background
We specify our notation, and background on MDPs, density estimation, and dimensionality reduction.
Notations: Estimators will be denoted with (cid:98)
2.1 Markov Decision Processes
We follow the formal setting of [29].
A Markov Decision Process (MDP) M is a 7-tuple M =( S , A , C , P init , C, P, H ) , where S , A and C are the state, action and cost spaces, respectively, P init isthe initial state distribution,
C : S ×A → P ( C ) is the cost function, P : S ×A → P ( S ) is the transition function and
H ∈ Z + is the horizon.
We denote by P ( c, s (cid:48) | s, a ) =
P ( s (cid:48) | s, a ) C ( c | s, a ) theprobability of transitioning from state s
given action a to state s (cid:48) with a cost of c , ∀ s, s (cid:48) ∈ S, ∀ a ∈ Aand ∀ c ∈ C .
When it is not clear from the context, we add the subscript × M for the variables in the7-tuple to 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 discretetime steps: the state at time t is s t , where s 0 ∼ P init , the agent chooses action a t , and then the statetransitions to s t +1 ∼ P ( ·| s t , a t ) and cost c t ∼ C ( ·| s t , a t ) is observed.
We denote the history at timet as h t = { s 0 , a 0 , c 0 , s 1 , a 1 , c 1 . . .
, s t } , and the space of possible histories at time t as H t .
The agentacts based on a history dependent policy π : H t → P ( A ) .
Every H steps, the episode ends and thestate is reset, meaning it is sampled again from P init .
The agent’s goal in an MDP is to minimize theexpected cumulative cost E π,M (cid:104)(cid:80) T − 1 t =
c t (cid:105) .
Note that the MDP horizon H and the horizon T of thecumulative cost might be different – this is important for our subsequent development, where theagent will face an unknown MDP, and can therefore improve its performance between episodes.
2.2 Kernel Density Estimation
Density estimation is the approximation of an unknown probability density function f ( x ) , x ∈ R d ,based on n i.i.d.
samples from it { X i } ni
=1 .
One popular approach to density estimation is kerneldensity estimation (KDE [26]).
Let K ( x ) denote a kernel function, which is a probability measureover R d .
Let H 0 denote a d × d positive deﬁnite and symmetric matrix satisfying det ( H 0 )
= 1 ,referred to as a unit bandwidth matrix.
Following the formulation of Jiang [14], we deﬁne the
KDE with bandwidth h > 0 as follows: (cid:98) f H 0 ,h ( x )
= n · h d n (cid:88) i =
K (cid:32)
H − 1 / 2( x − X i ) h (cid:33) .
Intuitively, hcontrols the bandwidth size, while H 0 allows for asymmetric bandwidth along different dimensions.
We make this clearer with a simple example:
Example 1.
[Symmetric Gaussian KDE]
For a Gaussian kernel function: K ( x ) = 1 ( √ 2 π ) d e − 12 x T x ,bandwidth matrix H
= I , and bandwidth h , we obtain the Gaussian KDE: (cid:98) f G ( x ) :
=1 n ( h √ 2 π ) d (cid:80) ni
=1 e − 12 h 2 ( x − X i ) T ( x − X i ) .
This popular instance of the KDE can be understood asplacing a Gaussian distribution with zero mean and standard deviation h over each sample
X i .
Jiang [14] provided ﬁnite sample bounds for the L ∞ norm of the KDE error under mild assumptions.
Assumption 1.
(cid:107) f (cid:107) ∞ < ∞
Assumption 2.
The kernel function is spherically symmetric and non-increasing, meaning thereexist non-increasing function k : R ≥
→ R ≥ 0 such that K ( x ) = κ ( (cid:107) x (cid:107) ) for x ∈ R d .
Further,the kernel function decays exponentially, meaning there exist ρ, C ρ , t 0 ∈ R > 0 such that κ ( t ) ≤
C ρ · exp ( − t ρ ) , ∀ t > t 0 .
Notice that the Gaussian kernel in Example 1 satisﬁes Assumption 2.
Assumption 3.
f is α -Holder continuous, meaning there exists 0 < α ≤ 1 , C α > 0 , such that| f ( x ) − f ( x (cid:48)
[Equation]
| ≤ C
[Equation]
(cid:107)
[Equation]
(cid:48) (cid:107)
[Equation]
, ∀
[Equation] [Equation]
R d .
Theorem 1.
[Theorem 2 of Jiang [14]] Under Assumptions 1 - 3, there exists a positive constant
C (cid:48) ≡ C (cid:48)
( d, (cid:107) f (cid:107) ∞ , C α , α, K ) such that the following holds with probability at least 1 − 1 /n ,uniformly in h > (log n/n ) 1 /d and valid unit bandwidths matrices H 0 :
[Equation]
where σ min is the smallest eigenvalue of H 0 .
The ﬁrst term on the right hand side of (1) 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 anoptimal bandwidth scales exponentially in the dimension d
(see Lemma 4 for a speciﬁc example).
2.3 Principal Component Analysis
Principal component analysis (PCA [32; 24]) is a popular linear dimensionality reduction technique.
For a d -dimensional random variable X , consider n i.i.d.
samples of X , denoted { X i } ni =1 .
Reducingthe dimension from d to d (cid:48) is achieved by projection – multiplying a vector in R d by a rank d (cid:48)orthogonal matrix P .
The expected projection error of P is R ( P ) :
= E (cid:2) (cid:107) X − PX (cid:107) 2 (cid:3) .
Similarly,the empirical projection error of P is R N ( P ) :
= (cid:80) ni =1 (cid:2) (cid:107) X
i − PX
i (cid:107) 2 (cid:3) .
We denote by Σ and (cid:98) Σ thecovariance and empirical covariance of X , respectively.
We further denote by λ 1 ≥ λ 2 ≥ · · · λ d > 0the eigenvalues of Σ , and by (cid:98) λ 1 ≥ (cid:98) λ 2 ≥ · · · (cid:98) λ d > 0 the eigenvalues of (cid:98) Σ .
The PCA projection (cid:98) P d (cid:48) is constructed from the eigenvectors corresponding to the d (cid:48) largest eigenvalues of (cid:98) Σ .
Its theoretical risk is: δ PCAd (cid:48) =
E (cid:104) R (cid:16) (cid:98) P d (cid:48)
(cid:17)(cid:105) − min P ∈P d (cid:48)
R ( P ) , where P d (cid:48) is the spaceof d × d orthogonal matrices with rank d (cid:48) ≤ d , and the expectation is with respect to the n samples.
Reiss and Wahl [24] bound the risk for sub-Gaussian random variables.
Assumption 4.
X is sub Gaussian, meaning that its second moment is ﬁnite and there exist a constant
C sg such that sup k ≥
[Table]
It is worth mentioning that any bounded random variable satisﬁes Assumption 4.
Theorem 2.
[Proposition 2.2 of Reiss and Wahl [24]] for a random variable X of dimension d thatsatisﬁes Assumption 4, we have that δ PCAd
[Table]
Problem Statement
In this paper we make use of a parametric representation of MDPs.
We suppose there is someparametric space Θ and a mapping g : Θ → M , which maps parameters to MDPs.
We next givetwo examples of such parametrizations, which will be referred to throughout the text to illustrateproperties of our analysis.
Example 2.
[Tabular mapping] Consider the case where S , A , T, C are ﬁxed and ﬁnite, while only
C and P differ between different MDPs.
The parametric space is deﬁned by Θ = Θ C × Θ P , whereΘ C ⊂ R |S|×|A|×|C| , and Θ P ⊂ R |S|×|A|×|S| , such that ∀ θ P ∈ Θ P , ∀ θ C ∈ Θ C , ∀
a ∈
A and∀ s ∈ S :
θ P ( a, s, · ) is on the |S| -simplex and θ C ( a, s, · ) is on the |C| -simplex.
The parametricmapping is deﬁned such that: P ( s (cid:48) | s, a ) =
θ P ( s, a, s (cid:48) ) and C ( c | s, a ) =
θ C ( s, a, c ) .
The tabular mapping in Example 2 does not assume any structure of the MDP space, and is commonin the Bayesian RL literature, e.g., in the Bayes-adaptive MDP model of Duff [4].
The next exampleconsiders a structured MDP space, and is inspired by the domain in Figure 1.
Example 3.
[Half-circle 2D navigation] Consider the following: S , A
= R 2 , and P ( a, s, s (cid:48) )
=
1for s (cid:48)
=
s + a ∀ s ∈ S , ∀ a ∈ A .
The parameter space is Θ = [0 , π ] .
The parametric mapping
C ( a, s ) = 1 , ∀ a ∈ A , ∀ s ∈ S such that (cid:107) s −
[ R cos θ, R sin θ ] (cid:107) 2 ≤ r , where r, R ∈ R .
3.1 The Learning Problem
We ﬁrst deﬁne the cumulative cost for a history-dependent policy π and an MDP M ∈ M as
L M,π = E π,M (cid:104)(cid:80) T − 1 t =
c t (cid:105) .
We follow the Bayesian RL (BRL) formulation [8], and assume a priordistribution over the MDP parameter space f ∈ P (Θ) .
Our objective is the expected cumulative cost 172 over a randomly sampled MDP from the prior:
[Equation]
where we used the parametric representation introduced above.
A policy π BO ∈ arg min π L f ( π )is termed Bayes optimal .
If the prior distribution is known, one can calculate the Bayes optimalpolicy [8].
However, in our setting we assume that f is not known in advance, motivating thefollowing learning problem.
− t 2 / ˜ C sg ) , where C sg and ˜ C sg differ by a universal constant [31].
We are given a training set of MDPs, { θ } Ni =1 , sampled independently from the prior f ( θ ) .
Our goal is to use these MDPs to calculate a policy that minimizes the regret:
[Equation]
(cid:34)
[Equation] [Equation]
(cid:35)
[Equation] [Equation]
) (cid:34) T
[Equation]
(cid:88)
[Equation]
c
[Equation]
(cid:35)(cid:35)
Note that the expectation above is with respect to the true prior, therefore, the regret can be interpretedas how well the policy calculated from the training data generalizes to an unseen test MDP.
Generalization Bounds
We next provide PAC bounds for the learning problem of Section 3.
Our general idea is to ﬁrst estimatethe prior distribution f ( θ ) from the training set using KDE, and then solve for the Bayes optimalpolicy with respect to the KDE-estimated distribution instead of the real prior.
For an estimator (cid:98) f ( θ )of the real prior f ( θ ) , we deﬁne the estimated Bayes optimal policy: π ∗ (cid:98) f ∈ arg min π L (cid:98) f
( π ) .
We start by showing that we can bound the regret of an estimated Bayes optimal policy, as a functionof the estimation error of the prior itself.
The proof, detailed in Section A.1 in the supplementary, is asimple application of norm inequalities, and exploiting the fact that the total cost is bounded.
Lemma 3.
Let (cid:98) f ∈ P
( R d ) be an estimator of the real prior f over the parametric space Θ .
We havethat:
[Equation]
1 , and for a bounded parametric space ofvolume | Θ | we have: R ( π
[Equation]
While the assumption of a ﬁnite parametric space volume is reasonable in practice, the volume size,which appears in the (cid:107)·(cid:107) ∞ bound that we’ll use in the proceeding analysis, grows exponentially withthe dimension of Θ .
In Section 4.1 we’ll discuss how, under some assumptions, we can relax this.
Lemma 3 provides a convenient framework for bounding the regret using any density estimationtechnique with a known bound.
We now consider the special case of the KDE (cf.
Section 2.2) witha Gaussian kernel function and an optimal selection of bandwidth, as calculated in the followinglemma.
We focus on the Gaussian kernel both for simplicity and because it is a popular choice inpractice.
Our method can be extended to any kernel function that satisﬁes assumption 2.
Lemma 4.
The optimal KDE bandwidth is (up to a constant independent of n ) h ∗ =
( log n / n ) 1 2 α + d .
The following result bounds the estimation error for this case.
The proof, detailed in Section A.3 ofthe supplementary material, is based on Theorem 1 from [14], where the constants are calculatedexplicitly for the Gaussian kernel case.
Lemma 5.
Under Assumptions 1 and 3, for a parametric space with ﬁnite volume | Θ | and a
KDE with a Gaussian kernel K ( u ) =
e − 12 uTu (2 π ) d 2 , H
= I , and an optimal bandwidth h ∗ , we havethat with probability at least 1 − 1 /n : sup x ∈ R d (cid:12)(cid:12) (cid:12) (cid:98)
f
G ( x ) − f ( x ) (cid:12)(cid:12) (cid:12) ≤ C d · (cid:0) log n n (cid:1)
α 2 α + d , where C d =
C
α 2 α − 1 2 + 16 d (cid:113) C
α ∆ αmax (Θ)+ 1 | Θ | √ 2(2 π ) d+ 64 d(2 π ) d 2 , and ∆ max (Θ) is the maximal L 1 distance between anytwo parameters in Θ .
Remark 1.
Usually, ∆ max (Θ) will scale polynomially (or sub-polynomially) with d .
For example, incase the parametric space is a d -dimensional hypercube with edge length B , we have: ∆ max = √ dB .
One parametric case that can be represented as a hyper cube is Example 2, where the edge lengthis B = 1 , and d = |S| 2 |A| + |S| |A| |C| .
In such cases, for large enough d , the ﬁrst term in theequation for C d will dominate, and therefore C d ≈ C
α 2 α − 1 2 .
By combining Lemma 3 and 5 we obtain a regret bound – the main result of this section.
Theorem 6.
For a prior f ( θ ) over a bounded parametric space Θ that satisﬁes Assumptions 1 and 3, and a Gaussian KDE with optimal bandwidth, we have that with probability at least 1 − 1 /n :
[Equation]
Remark 2.
While Theorem 6 assumes a bounded parametric space, the resulting KDE estimate isnot necessarily bounded (e.g., when using a Gaussian kernel).
In Section A.4 of the supplementary,we show that the result of Theorem 6 also holds when truncating the KDE estimate to a support Θ .
We next compare Theorem 6 with the the bounds of Tamar et al.
We recall that [29] considereda model-free approach that learns a history dependent policy (cid:98) π ∗ reg on the the training domains, with
L 2 policy regularization.
Tamar et al.
[29] only considered ﬁnite state, action, and cost spaces, asetting equivalent to our tabular representation in Example 2, where the parametric space dimensionis d = |S| 2 |A| + |S| |A| |C| .
Corollary 1 in [29] shows that with probability at least 1 − 1 /n ,
R
[Equation]
n 2 n
[Equation]
12 , where ρ = 2 q 2 T C 2max T 2 |A| ,and q = sup M,M (cid:48) ∈M
[Equation]
s (cid:48)
[Equation]
s (cid:48)
[Equation]
At ﬁrst sight, the exponent α 2 α + d in our bound compared to the 14 in [29] is signiﬁcantly worse.
Theintuitive reason for this is that Theorem 6 builds on KDE, where it is natural to expect an exponentialdependence on the dimension d when estimating the density f
( θ ) .
Yet, one must also considerthe constants.
Assuming that q is ﬁnite 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.
[29]claim that q may be made ﬁnite by adding small noise to every transition, but in this case q becomes
O ( |S||C| ) , leading to an exponential O (( |S||C| ) T ) term in the bound of [29], while our bound scaleslinearly with T .
The intuitive reason for this exponential dependence on T is that when learning ahistory-dependent policy, the space of possible histories grows exponentially with T .
To summarize, for a small T , and when the structure of M is such that q is ﬁnite, the model freeapproach of [29] seems preferable.
However, when the dimension d is small, and T is large, ourmodel based approach has the upper hand.
Remark 3.
Tamar et al.
[29] also considered a case where M is a ﬁnite set of size |M| , andin this case, their Corollary 2 shows that for 0 < α < 1 , with probability 1 − 1 / n α , R T (ˆ π ∗ ) ≤
[Equation]
43 n
[Equation]
3 , where P min = min M ∈M P ( M ) .
In the case of a discrete and ﬁniteparametric 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 (cid:98) P emp ( M )
= (cid:98) n ( M ) /n , ∀ M ∈ M , where(cid:98) n ( M ) is the number of occurrences of M in the training set.
We can bound the L 1 error ofthis estimator using the Bretagnolle-Huber-Carol inequality [30] (the full proof is outlined in
Section A.5 in the supplementary), and achieve that with probability at least 1 − 1 / n α we have that,
[Equation] [Equation]
log ( n log
[Equation]
) + |
[Equation]
|
[Equation]
) /n.
Observe that the dependence on T and
C max in this bound is better than in the bound from [29], and we do not have the P min term in thedenominator, which could be very small if some MDPs in the prior are rare, and is at most 1 / |M| .
Ignoring the log term, our dependence on n is also better for every α > 0 .
We conclude this section by pointing out that, as exempliﬁed in Remark 3, our approach can begeneralized to density estimation techniques beyond KDE, so long as their error can be bounded.
4.1 Bounds for Parametric Spaces with Low Dimensional Structure
The sample complexity in the general bound of Theorem 6 grows exponentially with the dimension ofthe parameter space Θ .
In many practical cases however, such as the semi-circle domain of Example 3,there may be a low dimensional representation that encodes most of the important information in thetasks, even though the dimensionality of the parametric space is higher.
In such cases, we expect thatour bounds can be improved to depend on the low dimensional representation.
In the following, weapproach this task by combining our model-based approach with the PCA dimensionality reductionmethod.
PCA is a linear method, and allows for a relatively simple analysis to demonstrate ourclaim.
In practice, a non-linear method may be preferred.
In Section 6, we verify empirically that ourapproach also works with non-linear deep neural network based dimensionality reduction.
We propose the following procedure
Reduce the training set dimensionality from d to d (cid:48) using PCA2.
Perform a Gaussian KDE estimation with optimal bandwidth in the low dimension d (cid:48)3.
Project back the estimated distribution to dimension d , to obtain the prior estimator (cid:98) f
d (cid:48)
G
The main difﬁculty in analysing this procedure, however, is calculating the error of the estimateddistribution in dimension d , that is, after the projection step .
To set the stage, we ﬁrst need to deﬁnethe projection step explicitly.
The PCA dimensionality reduction can be written as (cid:98) P d (cid:48)
=
W TL · W L , where W L ∈ R d (cid:48) × d .
Forany θ ∈ Θ , we therefore have that θ L = W L · θ is the low dimensional representation of θ , and
W TL · θ L is the projection of θ L back to the d -dimensional space.
For each θ L , we denote its inverseimage as follows:
Θ ⊥ L ( θ L ) = { θ ∈ Θ : W L · θ = θ L } .
By the law of total probability, the probabilitydistribution of a low-dimensional θ L is: f θ L ( θ L ) = (cid:82) θ ⊥ L ∈ Θ ⊥ L ( θ L ) f
θ ( θ ⊥ L ) dθ ⊥ L .
For the proceeding analysis, we require the function mapping parameters to MDPs be smooth.
Assumption 5.
In the MDP space M , only P and C can differ.
Furthermore, the parametricmapping is Lipshitz continuous with respect to P M ( · , · | s, a ) , i.e: ∃ C g such that ∀ s ∈ S , a ∈ A ,
[Equation]
We now extend the well known Simulation Lemma [15] to the case of a history-dependant policy.
This will allow us later to connect the error in the prior distribution to the error of the policy.
Lemma 7.
For any history dependant policy π and any parametric mapping g that satisﬁes Assumption 5, the following holds for any θ 1 , θ 2 ∈ Θ :
[Equation]
We note that Assumption 5 measures closeness between the joint distribution over costs and transitions.
This is slightly different than the simulation lemma for a Markov policy [15], where only the absolutedifference between the costs is required.
Intuitively, this is since a history-dependent policy candepend on observed costs, therefore even if two observed costs are very close in magnitude, theactions 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 Lipschitzcontinuity of g with respect to P and C separately.
We also note that there may be parameter spacesthat satisfy Eq. 2 without satisfying Assumption 5; in such cases our proceeding results will still hold.
In the following, we would like to formally consider MDP spaces with a low-dimensional structure.
This is captured by assuming that the MDPs essentially lie on a linear subspace of dimension d (cid:48) , suchthat the magnitude of the variability of MDPs outside this subspace is bounded by (cid:15) .
Assumption 6.
There exists d (cid:48) ≤ d such that λ d (cid:48)
+1 ≤ (cid:15) , where λ i , as deﬁned in Section 2.3, is thei ’th largest eigenvalue of the Covariance matrix of θ .
Under Assumption 6, and using the PCA bounds of Reiss and Wahl [24] (cf.
Theorem 2), we canbound 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 4 and 6, we have that:
[Equation]
There are two terms to the bound in Lemma 8.
The ﬁrst is due to the error in performing PCA with aﬁnite number of samples, and decays as n increases.
The second is due to the fact that even for aperfect PCA there is some error, as the MDPs do not lie perfectly in a low dimensional subspace.
We next assume that after the PCA projection, the distribution remains Holder continuous.
Assumption 7.
f
θ L is α (cid:48)
-Holder continuous, meaning there exists 0 < α (cid:48) ≤ 1 , C α (cid:48) > 0 , such that
[Equation]
Theorem 9.
Let (cid:98) f
d (cid:48)
G be the approximation of f as deﬁned above.
Under Assumptions 1, 4, 6, and 7we have with probability at least 1 − 1 /n : R
[Table]
| [Table] nn [Table]
+
[Table]
) ,where C d (cid:48)
= C α
[Table]
2 α (cid:48)− 1 2 + 16 d (cid:48)
[Table]
α [Table] and
[Table]
The ﬁrst term in the bound of Theorem 9 is the KDE error.
Note that, compared to the KDE error in
Theorem 6, the exponential dependence is on the low dimension d (cid:48) , and not on the higher dimensiond .
The second term in the bound is due to the PCA error, as discussed after Lemma 8.
This resultdemonstrates the potential in our approach, which accounts for structure in the parametric space.
Related Work
Meta RL has seen extensive empirical study [3; 35; 6; 23], and the connection to Bayesian RL hasbeen made in a series of recent works [18; 13; 22; 35].
Most meta RL studies assumed inﬁnite tasksduring training (effectively drawing different random MDPs from the prior at each training iteration),with few exceptions, mostly in the ofﬂine meta RL setting [2; 20].
Bayesian RL algorithms such as[11; 9] sample MDPs from the true posterior, which requires knowing the true prior.
To the best ofour knowledge, the only theoretical investigation of meta RL with ﬁnite training tasks is the modelfree approach in [29], which we compare against in our work.
Our theoretical analysis builds onideas from the study of density estimation [26; 14] and dimensionality reduction [32; 24].
To ourknowledge, this is the ﬁrst application of these ideas to the meta RL problem considered here.
We note that regularization techniques inspired by the mixup method [34] have been applied tometa learning [33], and recently also to meta RL [19], with a goal of improving generalization toout-of-distribution tasks.
In our experiments, we compared our approach with an approach inspiredby [19], and found that for in-distribution generalization, our approach worked better.
That said, webelieve 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 [1; 27; 5].
These works, whichhave a different ﬂavor from our PAC analysis, do not cover the meta RL problem considered here.
Experiments
In this section we complement our theoretical results with an empirical investigation.
Our goal is toshow that our main idea of learning a KDE over a low dimensional space of tasks is effective also forstate-of-the-art meta-RL algorithms, for which the linearity assumption of PCA clearly does not hold,and computing the optimal yet intractable π ∗ (cid:98) f is replaced with an approximate deep RL method.
Modern deep RL algorithms are known to be highly sensitive to many hyperparameters [12], andmeta RL algorithms are not different.
To demonstrate our case clearly, we chose to build on the
VariBAD algorithm of Zintgraf et al.
[35], for which we could implement our approach by replacingjust a single algorithmic component, as we describe next, and thus obtain a fair comparison.
VariBAD: We brieﬂy explain the VariBAD algorithm, to set the stage for our modiﬁed algorithm tofollow; we refer the reader to [35] for the full details.
VariBAD is composed of two main components.
The ﬁrst is a variational autoencoder (VAE [17]) with a recurrent neural network (RNN) encoder that,at each time step, encodes the history into a 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.
Using the terminology in our paper, the VAE latent state can be related to the MDP parameter θ , andthe VAE decoder learns a model of g ( θ ) .
The second component in VariBAD is a policy, mappingthe output of the VAE encoder and current state into an action.
This component, which can be seen asproducing an approximation of π ∗ (cid:98) f , is trained using a policy gradient algorithm such as PPO [28].
VariBAD Dream: Recall that our pipeline is to learn a KDE over the task parameters θ , and thentrain a policy on tasks from the estimated KDE.
Unfortunately, in our meta-RL setting, we do notassume that we directly know the θ representation for each task.
However, the VAE in VariBADallows for a convenient approximation.
We can think of the output of the VAE encoder at the end ofan episode, after a full history has been observed and the uncertainty about the task has been resolvedas best as possible, as a sample of the task parameter θ .
Thus, we propose to build our KDE estimateover these variables.
We henceforth refer to samples from the KDE as dream environments (as they
We emphasize that in the original VariBAD work [35], the algorithm was trained by sampling different tasks from the task distribution at each iteration , corresponding to an inﬁnite number of traning tasks.
As we are interested in the ﬁnite task setting, our implementation of VariBAD (and VariBAD Dream)draws a single batch of N train tasks once, at the beginning of training, and subsequently samplestasks from within this batch for training the VAE and policy.
Results: We consider the half-circle environment of Figure 1 – a popular task that requires learninga non-trivial Bayes-optimal policy [2].
In order to test generalization, we evaluated the agents on
N eval = 50 sampled environments, different from the N train training ones.
In Figure 2 we show theaffect of incorporating the dream environments on the average return for the evaluation environments.
Evidently, our approach achieved higher return, both for N train = 20 and N train = 30 .
For
N train = 40 , the original VariBad performed near optimally, and there was no advantage for VariBad
Dream to gain.
For N train = 10 , the sampling was too sparse, and both methods demonstratedcomparable failure in generalizing to the test environments.
We emphasize that while the performanceadvantage of VariBAD Dream is modest, it is remarkable that the KDE regularization demonstratesconsistent improvement, as the VAE training and PPO optimization already include signiﬁcantimplicit and explicit regularization mechanisms.
Additionally, we evaluated a method inspired by mixup [34; 33; 19], where dream environments arecreated by a random weighted average of latent vectors (instead of the KDE).
Interestingly, our KDEapproach 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 [19], where the goals weredistributed inside a rectangular area.
In Section A.10 in the supplementary we provide additionalempirical comparisons and visualizations.
Conclusion & Future Work
We propose a model-based scheme for meta-RL with a ﬁnite sample of training tasks, where weﬁrst 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 approachcan exploit low dimensional structure of the task distribution, when such exists, to obtain improvedbounds.
Finally, we showed that our approach can be “plugged-in” the VariBAD algorithm to improvegeneralization.
There are several exciting future directions for investigation.
Our theoretical analysis can be extendedto more advanced density estimation and dimensionality reduction techniques, such as VAEs [25].
Our empirical investigation hinted that regularization, such as effected by VariBAD Dream, canimprove deep meta-RL algorithms.
More sophisticated regularization can be developed based onprior knowledge about the possible tasks.
Checklist
The checklist follows the references.
Please read the checklist guidelines carefully for information onhow to answer these questions.
For each question, change the default [TODO] to [Yes] , [No] , or[N/A] .
You are strongly encouraged to include a justiﬁcation to your answer , either by referencingthe appropriate section of your paper or providing a brief inline description.
For example:
• Did you include the license to the code and datasets?
[Yes] See Section• Did you include the license to the code and datasets?
[No] The code and the data areproprietary.•
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 reﬂect the paper’s [Table] contributions and scope?
[Yes] [Table] (b) Did you describe the limitations of your work?
[Yes] [Table]
(c) Did you discuss any potential negative societal impacts of your work?
[N/A] [Table]
(d) Have you read the ethics review guidelines and ensured that your paper conforms tothem?
[Yes]2.
If you are 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 supplementarymaterials3.
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)? [Yes]
In the supplementary materials (b) Did you specify all the training details (e.g., data splits, hyperparameters, how theywere chosen)?
[Yes] In Section A.9 of the supplementary(c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[Yes] In Section 6 and Section A.10 in the supplementary(d) Did you include the total amount of compute and the type of resources used (e.g.,type of GPUs, internal cluster, or cloud provider)?
[Yes] See Section A.11 of the
supplementary
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a)
If your work uses existing assets, did you cite the creators?
[Yes] See Section A.9 ofthe supplementary
(b) Did you mention the license of the assets?
[Yes] See Section A.9 of the supplementary(c)
Did you include any new assets either in the supplemental material or as a URL?
[No]
(d) Did you discuss whether and how consent was obtained from people whose data you’reusing/curating?
[N/A]
(e) Did you discuss whether the data you are using/curating contains personally identiﬁable
information or offensive content?
[N/A]
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, ifapplicable?
[N/A](b)
Did you describe any potential participant risks, with links to Institutional ReviewBoard (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]