Meta Reinforcement Learning with Finite Training Tasks - a Density Estimation Approach
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 deﬁned, 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 ﬁrst such PAC analysis for a model-free setting, where a historydependent 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 signiﬁcantly 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.
Introduction
In recent years, reinforcement learning (RL) became a dominant algorithmic framework for a variety of domains, including computer games [24], robotic manipulation [11] and autonomous driving [18].
Popular RL algorithms, however, are characterized by a high sample inefﬁciency, due to the exploration-exploitation problem – the need to balance between obtaining more information about the environment versus acting based on such information.
Indeed, most RL success stories required a very long training process, only possible in simulation.
For an agent to learn fast, additional structure of the problem is required.
In Meta RL [4; 39; 7; 26], agents are allowed to 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 identiﬁed 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 a
Markov decision process (MDP), and the distribution over tasks is the Bayesian prior [25; 39].
In this formulation, the optimal meta RL policy is the Bayes-optimal policy – the policy that maximizes expected return, where the expectation is taken with respect to the prior MDP distribution.
While signiﬁcant empirical progress in meta RL has been made, the theoretical understanding of the problem is still limited.
A central question, which we focus on in this work, is the probably approximately correct (PAC) analysis of meta RL, namely, how many training tasks are required to guarantee performance that is approximately Bayes-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 ofﬂine 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 ﬁ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 [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 a model based approach for PAC meta RL.
Our idea is that instead of regularizing the policy during training, as in [33], 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, if 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 on 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 HalfCircle domain in Figure 1, adapted from [3]: a
2-dimensional agent
must navigate to a goal, lo cated somewhere on the half-circle.
A task therefore corresponds to the goal location, and the task distribution is uniform on
the 1-dimensional half-
circle.
The bounds in [33] depend on the number of states and actions, which are continuous here, and even if discretized, would result in excessive bounds.
Intuitively, however, it is the low-dimensional structure of the task distribution that should determine the difﬁculty of the problem, and our bounds can indeed account for such.
We remark that many benchmark meta RL task distributions exhibit similar low dimensional structure [4; 7; 3].
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-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.
Background
We specify our notation, and background on MDPs, density estimation, and dimensionality reduction.
Notations denoted explicitly.
For a set X , we denote by P ( X ) the set of distributions over X .
We denote by v d the volume of a unit ball in R d .
For a space S we denote |S| as the volume of the space.
2.1 Markov Decision Processes
We follow the formal setting of [33].
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 is the 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 ) the probability of transitioning from state s given action a to state s (cid:48) with a cost of c , ∀ s, s (cid:48) ∈ S, ∀ a ∈ A and ∀ c ∈ C .
When it is not clear from the context, we add the subscript × M for the variables in the 7-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 discrete time steps: the state at time t is s t , where s 0 ∼ P init , the agent chooses action a t , and then the state transitions 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 time t 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 agent acts based on a history-dependent policy π : H t → P ( A ) .
Every H steps, the episode ends and the state is reset, meaning it is sampled again from P init .
The agent’s goal in an MDP is to minimize the expected 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 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 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 kernel density estimation (KDE [29]).
Let K ( x ) denote a kernel function, which is a probability measure over 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 [16], 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, h controls 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 ) = [Equation] 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 as placing a Gaussian distribution with zero mean and standard deviation h over each sample X i .
Jiang [16] 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 there exists a 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 exists ρ, 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 α -Hölder continuous, meaning there exists 0 < α ≤ 1 , C α > 0 , such that
[Equation]
Theorem 1.
[Theorem 2 of Jiang [16]] 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 an optimal bandwidth scales exponentially in the dimension d
(see Lemma 4 for a speciﬁc example).
2.3 Principal Component Analysis
Principal component analysis (PCA [1; 27]) 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 .
Reducing the 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) Σ the covariance and empirical covariance of X , respectively.
We further denote by λ 1 ≥ λ 2 ≥ · · · λ d >the 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 space of d × d orthogonal matrices with rank d (cid:48) ≤ d , and the expectation is with respect to the n samples.
Reiß and Wahl [27] bound the risk for sub-Gaussian random variables.
Assumption 4.
X is sub Gaussian, meaning that its second moment is ﬁnite and there exists a constant C sg such that sup
[Equation]
It is worth mentioning that any bounded random variable satisﬁes Assumption 4.
Theorem 2.
[Proposition 2.2 of Reiß and Wahl [27]] for a random variable X of dimension d that satisﬁes Assumption 4, we have that δ PCAd
[Equation]
Problem Statement
In this paper we make use of a 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 to throughout the text to illustrate properties 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 parametric mapping 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 common in the Bayesian RL literature, e.g., in the Bayes-adaptive MDP model of Duff [5].
The next example considers 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) )
=for 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 [9], and assume a prior distribution over the MDP parameter space f ∈ P (Θ) .
Our objective is the expected cumulative cost 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 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, { θ } Ni =1 , sampled independently from the prior f ( θ ) .
Our goal is to use these MDPs to calculate a policy that minimizes the regret:
[Equation]
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.
Generalization Bounds
We next provide PAC bounds for the learning problem of Section 3.
Our general idea is to ﬁrst estimate the prior distribution f ( θ ) 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 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 function of the estimation error of the prior itself.
The proof, detailed in Section A.2 in the supplementary, is a simple 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 have that: R ( π ∗ (cid:98) f ) =
L f ( π ∗ (cid:98) f ) − L f ( π BO ) ≤
[Equation]
C max T (cid:13)(cid:13)(cid:13) f − (cid:98) f
(cid:13)(cid:13)(cid:13) 1 , and for a bounded parametric space of volume | Θ | we have: R ( π
[Equation]
=
[Equation]
∞ .
While the assumption of a ﬁnite parametric space volume is reasonable in practice, the volume size, appearing in the (cid:107)·(cid:107) ∞ bound that we shall use in the proceeding analysis, grows exponentially with the dimension of Θ .
In Section 4.1 we will discuss how, under some assumptions, we can relax this.
Lemma 3 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 in the following 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 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.4 of the supplementary material, is based on Theorem 1 from [16], where the constants are calculated explicitly 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 have that 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 nn (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 any two parameters in Θ .
Remark 1.
Usually, ∆ max (Θ) will scale polynomially (or sub-polynomially) with d .
For example, in case 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 length is B = 1 , and d = |S| 2 |A| + |S| |A| |C| .
In such cases, for large enough d , the ﬁrst term in the equation 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 is not necessarily bounded (e.g., when using a Gaussian kernel).
In Section A.5 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 [33] considered a 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.
[33] only considered ﬁnite state, action, and cost spaces, a setting equivalent to our tabular representation in Example 2, where the parametric space dimension is d = |S| 2 |A| + |S| |A| |C| .
Corollary 1 in [33] shows that with probability at least 1 − 1 /n ,
[Equation]
n (cid:17) 12 , where ρ = 2 q 2 T C 2max T 2 |A| , and q = sup M,M (cid:48)
[Equation]
[Equation] [Equation] [Equation]
,c
[Equation] [Equation]
) / P M
[Equation]
( s
[Equation]
,c | s,a ) .
At ﬁrst sight, the exponent α 2 α + d in our bound compared to the 1 4 in [33] is signiﬁcantly worse.
The intuitive reason for this is that Theorem 6 builds on KDE, where it is natural to expect an exponential dependence on the dimension d when estimating the density f
( θ ) .
Yet, one must also consider the 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.
[33] 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 [33], while our bound scales linearly with T .
The intuitive reason for this exponential dependence on T is that when learning a history-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-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 a case where M is a ﬁnite set of size |M| , and in this case, their Corollary 2 shows that for 0 < α < 1 , with probability 1 − 1 / n α , R T (ˆ π ∗ [Equation] ≤
[Equation]
3 , where P
[Equation]
= [Equation] P ( M ) .
In the case of a discrete and ﬁnite 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 (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 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 / n α , we have that
R T ( π ∗ (cid:98) P emp ) ≤ 2 C max T (cid:112) 2 ( α log ( n log 2) + |M| + 1) /n.
Observe that the dependence on T and
C max in this bound is better than in the bound of [33], and we do not have the P min term in the denominator, 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 be generalized 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 of the parameter space Θ .
In many practical cases however, such as the HalfCircle domain of Example 3, there may be a low dimensional representation that encodes most of the important information in 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 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 6, we verify empirically that our approach 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 PCA 2.
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 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 π ∗ (cid:98) f requires g ( θ ) , where θ is in dimension d .
To set the stage, we ﬁrst need to deﬁne the 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 .
For any θ ∈ Θ , 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 inverse image as follows:
Θ ⊥ L ( θ L ) = { θ ∈ Θ : W L · θ = θ L } .
By the law of total probability, the probability distribution 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 that maps parameters to 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 P M ( · , · | s, a ) , i.e: ∃ C g such that ∀ s ∈ S , a ∈ A ,
[Equation]
We now extend the well known Simulation Lemma [17] to the case of a history-dependent policy.
This will allow us later to relate the error in the prior distribution to the error of the policy.
Lemma 7.
For any history-dependent 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 [17], where only the absolute difference between the costs is required.
Intuitively, this is since a 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.
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) , such that the magnitude of the variability of MDPs outside this subspace is bounded by (cid:15) .
Assumption 6.
There exist d (cid:48) ≤ d and (cid:15) ∈ R ≥ 0 such that λ d (cid:48)
+1 ≤ (cid:15) , where λ i , as deﬁned in Section 2.3, is the i ’th largest eigenvalue of the Covariance matrix of θ .
Under Assumption 6, and using the PCA bounds of Reiß and Wahl [27] (cf.
Theorem 2), 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 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 a perfect 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 Hölder continuous.
Assumption 7.
f
θ L is α (cid:48)
-Hölder continuous, meaning there exists 0 < α (cid:48) ≤ 1 , C
[Equation]
> 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, andwe have with probability at least 1 − 1 /n :
[Equation]
where C d (cid:48)
=
C α
[Equation]
+ 16 d
[Equation]
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 dimension d .
The second term in the bound is due to the PCA error, as discussed after Lemma 8.
This result demonstrates the potential in our approach, which accounts for structure in the parametric space.
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; 25; 39].
Most meta RL studies assumed inﬁnite tasks during training (effectively drawing different random MDPs from the prior at each training iteration), with few exceptions, mostly in the ofﬂine 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, Simchowitz et al.
[32] analyse a family of posterior-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-specifed prior, and, more importantly, focuses on the sample complexity of obtaining a 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 ﬁnite 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; 16] and dimensionality reduction [1; 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 [21], 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 [2; 30; 6].
These works, which have 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 to show that our main idea of learning a KDE over a low dimensional space of tasks is effective also for state-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 [13], 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.
VariBAD: We brieﬂy explain the VariBAD algorithm, to set the stage for our modiﬁed algorithm to follow; we refer the reader to [39] for the full details.
VariBAD is composed of two main components.
The ﬁrst is a variational autoencoder (VAE [19]) 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 related to the MDP parameter θ , and the VAE decoder learns a model of g ( θ ) .
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 π ∗ (cid:98) f , is trained using a policy gradient algorithm such as PPO [31].
VariBAD Dream: Recall that our pipeline is to learn a KDE over the task parameters θ , 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 as best as possible, as a sample of the task parameter θ .
Thus, we propose to build our KDE estimate 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 on dream environments.
We refer to this method as VariBAD
Dream .
In our implementation, we train the KDE and VariBAD components simultaneously.
The full implementation details and pseudo code are in Section A.10 of the supplementary.
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 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 samples tasks from within this batch for training the VAE and policy.
Results: We consider the HalfCircle environment of Figure 1 – a popular task that requires learning a non-trivial Bayes-optimal policy [3].
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 the affect 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 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 signiﬁcant implicit and explicit regularization mechanisms.
Additionally, we evaluated a method inspired by mixup [38; 37; 21], 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 [21] were designed to evaluate out-of-distribution generalization, different from the in distribution setting considered here.
In Sections A.11 and A.12 of the supplementary we provide additional evaluations and show similar results on a MuJoCo [34] environment with a high dimensional state space.
Synthetic Experiments with a Known Parametric Space:
In our theoretical analysis, we assumed knowledge of both the parameters θ of the training MDPs, and the function g ( θ ) that maps parameters to MDPs.
We complement our empirical investigation with synthetic experiments with VariBAD
Dream, where g ( θ ) is known.
We consider again the HalfCircle environment of Figure 1.
We deﬁne the parametric space as Θ = R 2 , where g ( θ ) prescribes a corresponding MDP with a goal at location ( x, y )
= θ .
For our experiment, we sample { θ i } Ni =1 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 3 compares VariBAD Dream with the conventional VariBad (trained only on { θ i } Ni =1 ), on unseen test environments.
Evidently, the use of KDE improved the generalization performances of
VariBad by a 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 ( θ ) .
Discussion & 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 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 our 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, are 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 ﬁne-tuning a policy trained on diverse tasks display state-of-the-art performance.
There are 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 as 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 reﬂect 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.
Checklist
The checklist follows the references.
Please read the checklist guidelines carefully for information on how 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 referencing the 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 are proprietary.
• 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 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 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 supplementary materials 3.
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 they were chosen)?
[Yes] In Section A.10 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.11 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.13 of the supplementary 4.
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.10 of the supplementary (b) Did you mention the license of the assets?
[Yes] See Section A.10 of the supplementary (c) Did you include any new assets either in the supplemental material or as a URL?
[Yes] See Section A.10 of 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 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, if applicable?
[N/A] (b) Did you describe any potential participant risks, with links to 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]