Posterior Coreset Construction with Kernelized Stein Discrepancy for Model-Based Reinforcement Learning
Abstract
Model-based reinforcement learning (MBRL) exhibit favorable performance in practice, but their theoretical guarantees are mostly restricted to the setting when transition model is Gaussian or Lipschitz, and demands a posterior estimate whose representational complexity grows unbounded with time.
In this work, we develop a novel MBRL method (i) which relaxes the assumptions on the target transition model to belong to a generic family of mixture models; (ii) is applicable to largescale training by incorporating a compression step such that the posterior estimate consists of a Bayesian coreset of only statistically significant past state-action pairs; and (iii) exhibits a Bayesian regret of O ( dH 1+( α/ 2) T 1 − ( α/ 2) ) with coreset size of Ω( √ T 1+ α ) , where d is the aggregate dimension of state action space, H is the episode length, T is the total number of time steps experienced, and α ∈ (0 , 1] is the tuning parameter which is a novel introduction into the analysis of MBRL in this work.
To achieve these results, we adopt an approach based upon Stein’s method, which allows distributional distance to be evaluated in closed form as the kernelized Stein discrepancy (KSD).
Experimentally, we observe that this approach is competitive with several state of the art RL methodologies, and can achieve up-to 50% reduction in wall clock time in some continuous control environments.
Problem Formulation
We consider the problem of modelling an episodic finite-horizon Markov Decision Process (MDP) where the true unknown MDP is defined as M ∗ := {S , A , R ∗ , P ∗ , H, R max , ρ } , where S ⊂ R d s and A ⊂ R d a denote continuous state and action spaces, respectively.
Here, P ∗ represents the true underlying generating process for the state action transitions and R ∗ is the true rewards distribution.
After every episode of length H , the state will reset according to the initial state distribution ρ .
At time step i ∈ [1 , H ] within an episode, the agent observe s i ∈ S , selects a i ∈ A according to a policy µ , receives a reward r i ∼ R ∗ (
s i , a i ) and transitions to a new state
s i +1 ∼ P ∗ ( ·| s i , a i ) .
We consider M ∗ itself as a random process, as is the often the case in Bayesian Reinforcement
Learning, which helps us to distinguish between the true and fitted transition/reward model.
Next, we define policy µ as a mapping from state s ∈ S to action a ∈
A over an episode of length
H .
For a given MDP M , the value for time step i is the reward accumulation during the episode
V Mµ,i ( s ) =
E [Σ Hj =
i [¯ r M ( s j , µ ( s j , j ))
| s i =
s ] where actions are under policy µ ( s j , j ) ( j denotes the timestep within the episode) and ¯ r M ( s, a ) =
E r ∼ R M
( s,a )
[ r ] .
Without loss of generality, we assume the expected reward an agent receives at a single step is bounded | ¯ r M ( s, a ) | ≤ R max , ∀ s ∈ S , a ∈ A .
This further implies that | V ( s ) | ≤ HR max , ∀ s .
For a given MDP M , the optimal policy µ M is defined as µ M = arg max µ V Mµ,i ( s ) for all s and i = 1 , . . .
, H .
Next, we also define future value function U Mi ( P ) to be the expected value of the value function over all initializations and trajectories
U Mi ( P ) =
E s ′ ∼ P
( s ′ | s,a ) ,a = µ M ( s,i )
[ V Mµ M ,i ( s ′ )
| s i = s ] where P is the transition distribution under
MDP M .
According to these definitions, one would like to find the optimal policy ( ?? )
for the true model M = M ∗ .
Next, we review the PSRL algorithm, which is an adaption of Thompson sampling to RL [41] (see Algorithm 1 in Appendix B).
In PSRL, we start with a prior distribution over MDP given by ϕ .
Then at each episode, we take sample M k from the posterior given by ϕ ( ·|D k ) , where D k := { s 1 , 1 , a 1 , 1 , r 1 , 1 , · · · , s k − 1 ,H , a k − 1 ,H , r k − 1 ,H } is a data set containing past trajectory data, i.e., state-action-reward triples, which we call a dictionary .
That is, where s k,i , a k,i and r k,i indicate the state, action, and reward at time step zi in episode k .
Then, we evaluate the optimal policy µ k := µ M k .
Thereafter, information from the latest episode is appended to the dictionary as
D k
+
=
D k ∪ { s k, 1 , a k, 1 , r k, 1 , · · · , s k,H , a k,H , r k,H } .
Bayes Regret and Limitations of PSRL:
The Bayes regret of PSRL (cf.
Algorithm 1) is established to be ˜ O
( (cid:112)
d K ( R ) d E ( R ) T + E
[ L ∗ ] (cid:112) d K ( P ) d E ( P )) where d K and d E are Kolmogorov and Eluder dimensions, R and P refer to function classes of rewards and transitions, and L ∗ is a global Lipschitz constant for the future value function.
Although it is mentioned that system noise smooths the future value functions in [41], an explicit connection between H and L is absent, which leads to an exponential dependence on horizon length H in the regret [41, Corollary 2] for LQR.
A crucial assumption in deriving the best known regret bound for PSRL in continuous control is of target distribution belonging to Gaussian/symmetric class, which is often violated.
For instance, if we consider a variant of inverted pendulum with an articulated arm, the transition model has at least as many modes as there are minor-joints in the arm.
Another major challenge is related to posterior’s parameterization complexity M ( T ) :
=
|D k | , which we subsequently call the dictionary size (stepin Algorithm 1) which is used to parameterize the posterior distribution.
We note that M ( T ) =
Ω( T ) for the PSRL [41] and MPC-PSRL [19] algorithms which are state of the art approaches.
To alleviate the Gaussian restriction, we consider an alternative metric of evaluating the distributional estimation error, namely, the kerneralized Stein discrepancy (KSD).
Additionally, that KSD is easy to evaluate under appropriate conditions on the target distribution, i.e., the target distribution is smooth, one can compare its relative quality as a function of which data is included in the posterior.
Doing so allows us to judiciously choose which points to retain during the learning process in order to ensure small Bayesian regret.
To our knowledge, this work is the first to deal with the compression of posterior estimate in model-based RL settings along with provable guarantees.
These aspects are derived in detail in the following section.
Proposed Approach 2.1 Posterior Coreset Construction via KSD
The core of our algorithmic development is based upon the computation of Stein kernels and
KSD to evaluate the merit of a given transition model estimate, and determine which past samples to retain.
Doing so is based upon the consideration of Stein based integral probability metrics (IPM) rather than total variation (TV) distance which allows us to employ Stein’s identity, which under a hypothesis that the score function of the target is computable.
This approach is wellknown to yield methods to improve the sample complexity of Markov Chain Monte Carlo (MCMC) methods [13].
This turns out to be the case in model-based RL as well is a testament to its power [49, 47].
This method to approximate a target density P consists of defining an IPM [48] based on a set G consisting of test functions on X ⊂ R {+
+} d s + d a , and is defined as D G ,P ( { x i } ni =1 ) : = sup g ∈G (cid:12)(cid:12)
[Table]
Although IPMs efficiently quantify the discrepancy between an estimate and target, requires P to evaluate the integral, which may be unavailable.
To alleviate this issue, Stein’s method restricts the class of test functions g to be those such that E P [ g ( z )]
= 0 .
In this case, IPM only depends on the Dirac-delta measure ( δ ) from the stream of samples, removing dependency on the exact integration in terms of P .
Under certain smoothness assumptions on the posterior, the IPM can be evaluated in closed form through the kernelized Stein discrepancy (KSD).
The key technical upshot of employing Stein’s method is that we can now evaluate the integral probability metric of posterior ϕ D k := ϕ ( ·|D k ) parameterized by dictionary D k through the KSD, which is efficiently computable:
[Equation]
h i , h j ) which depends solely on score function of the estimated posterior and we no longer require access to the true unknown target transition model of the MDP.
This is a major merit of utilizing Stein’s method in MBRL, and allows us to improve the regret of model-based RL methods based on posterior sampling.
This the previous point is distinct from the computational burden of storing dictionary D k that parameterizes ϕ D k and evaluating the optimal value function according to the current belief model.
After this novel change (see Lemma ??
in Sec. 3), we can utilize the machinery of KSD to derive the regret rate for the proposed algorithm in this work (cf.
Algorithm 2) in lieu of concentration inequalities, as in [19].
We shift to the computational storage requirements of the posterior in continuous space next.
KSD Thinning: We develop a principled way to avoid the requirement that the dictionary D k retains all information from past episodes, and is instead parameterized by a coreset of statistically significant samples.
The issue in PSRL (or MBRL in general) is as the number of episodes experienced becomes large, the posterior representational complexity grows linearly and unbounded with episode index k .
On top of that, the posterior update in PSRL (cf.
Algorithm 1) is also parameterized by data collected in D k [-+
-]+1 (ex : Gaussian processes.[46]).
To deal with this bottleneck, we propose to sequentially remove those particles from D k +1 that contribute least in terms of KSD.
It can be interpreted as projecting posterior estimates onto “subspaces" spanned by only statistically representative past state-action-state triples.
Such nonparametric posterior representation has been shown to exhibit theoretical and numerical advantages in probability density estimation [11, 12], Gaussian Processes [32], and Monte Carlo methods [18].
Here we introduce it for the first time in model-based RL, which allows us to control the growth of the posterior complexity, which in turn permits us to obtain computationally efficient updates.
[Table]
To be more specific, suppose we are at episode k with dictionary D k associated with posterior ϕ D k , and we denote the dictionary after update as (cid:101)
D k
{++
+} =
D k + H and corresponding posterior as ϕ (cid:101)
D k +1 .
For a given dictionary D k , we can calculate the KSD of posterior ϕ D k to target.
We note that the KSD estimate goes to zero as k → ∞ due to the posterior consistency conditions [22].
At each episode k , after performing the dictionary update to obtain (cid:101) D k +1 (step 10 in
Algorithm 2), we propose to thin dictionary (cid:101) D
k +1 such that KSD ( ϕ D k +1 ) 2 <
KSD ( ϕ (cid:101) D k +1 ) 2 + ϵ k +1 where D k +1 is the dictionary following thinning and ϵ k > 0 is a scalar parameter we call the thinning budget proposed.
This means the posterior defined by compressed dictionary ϕ (cid:101)
D k +1 is at most ϵ k +1 in KSD from its uncompressed counterpart.
We will see how ϵ k permits us to trade off regret and dictionary size in practice.
KSD may be succinctly stated as
[Equation]
We summarize the proposed algorithm in Algorithm 2 with compression subroutine in Algorithm 3 in
Appendix, where KSRL is an abbreviation for Kernelized Stein Discrepancy Thinning for Model-Based
Reinforcement Learning.
Please refer to discussion Appendix B for MPC based action selection.
Bayesian Regret Analysis
In this section, we develop the Bayesian Regret Analysis for KSRL using score-based IPM that exploit’s salient structural properties of Stein’s method, and additionally provides a basis for establishing tradeoffs between regret and posterior representational complexity which is novel in this work.
In the first step, we upper bound the future value function estimation error of P k with respect to P ∗ as U ki
( P [Equation]
[Equation] ≤ [Equation] max KSD ( P k ( h i )) for all i and k .
The salient point to note here is that we do not require access to P ∗ to evaluate the RHS of the equation.
Then, we prove that E [ KSD ( ϕ D k )]
= O (cid:18)
√ k log( k ) f
( k ) (cid:19) where ϵ k = log( k ) f ( k ) 2 is the thinning budget.
For the statistical consistency of the posterior estimate, we note that it is sufficient to show that E [ KSD ( ϕ D k )]
→as k → ∞ , which imposes a lower bound on the dictionary size growth rate f ( k ) > (cid:112) k log( k ) required for convergence.
This result communicates that it is possible to achieve statistical consistency without having a linear growth in the dictionary size that is a drawback of prior art [41, 19].
Finally, with ϵ k = log( k ) f ( k ) 2 and coreset size growth condition f ( k ) =
(cid:112) k
α +1 log( k ) where α ∈ (0 , 1] , we prove total Bayes regret for KSRL is given by E [ Regret T ]
= O (cid:0) dT 1 − α 2 H 1+ α 2 (cid:1) with coreset size
M ( T ) =
˜Ω (cid:16) √ T 1+ α (cid:17) .
Observe that we match the best known prior results of PSRL with α = 1 as shown in [19] in-terms of O ( dH 32 T 12 ) , but with relaxed conditions on the posterior, allowing the approach to apply to a much broader class of problems.
Moreover, for α < 1 , we obtain a superior tradeoff in model complexity and regret.
[Figure]
Experiments
We present a detailed experimental analysis of KSRL comparing to state of the art model-based and model-free RL methods on several continuous control tasks in terms of training rewards, model complexity and KSD convergence.
For model-based we compare to MPC-PSRL method proposed in [19] since MPC-PSRL is already shown to outperform MBPO [27] and PETS [15] in this environmental settings.
For comparison to model-free approaches, we compare with Soft Actor-Critic (SAC) from [23] and Deep Deterministic Policy Gradient (DDPG) [5].
We consider continuous control environments Stochastic Pendulum , Continuous Cartpole, Reacher and Pusher with and without rewards of modified OpenAI Gym [8] & MuJoCo environments [55].
See Appendix F for additional specific details of the environments and architecture.
Discussion.
Fig.
2 compares the average reward return (top row) and model complexity (bottom row) for Cartpole, Pendulum, and Pusher, respectively.
We note that KSRL performs equally good or even better as compared to the state of the art MPC-PSRL algorithm with significant reduction in model complexity (bottom row in Fig. 2) consistently across different environments.
From the model complexity plots, we remark that KSRL is capable of automatically selecting the data points and control the dictionary growth across different environments which helps to achieve same performance in terms of average reward with fewer dictionary points to parameterize posterior distributions.
This also helps in achieving faster compute time in practice to perform the same task as detailed in Fig.in the Appendix F. We also show improvements of our algorithm over MPC with fixed buffer size as shown in Fig.
8 Appendix F with both random and sequential removal.
A Introduction
Reinforcement learning, mathematically characterized by a Markov Decision Process (MDP) [45], has gained traction for addressing sequential decision-making problems with long-term incentives and uncertainty in state transitions [52].
A persistent debate exists as to whether model-free (approximate dynamic programming [51] or policy search [57]), or model-based (model-predictive control, MPC [20, 30]) methods, are superior in principle and practice [56].
A major impediment to settling this debate is that performance certificates are presented in disparate ways, such as probably approximate correct (PAC) bounds [50, 16], frequentist regret [28, 29], Bayesian regret [1, 58, 38], and convergence in various distributional metrics [6, 2, 33].
In this work, we restrict focus to regret, as it imposes the fewest requirements on access to a generative model underlying state transitions.
In evaluating the landscape of frequentist regret bounds for RL methods, both model-based and model-free approaches have been extensively studied [28].
Value-based methods in episodic settings have been shown to achieve regret bounds ˜ O
( d p H q √ T ) [59] (with p = 1 , q = 2 ), where H is the episode length, and d is the aggregate dimension of the state and action space.
This result has been improved to p = q = 3 / 2 in [29], and further to p = 3 /and q = 1 in [60].
Recently, model-based methods have gained traction for improving upon the best known regret model-free methods with p = 1 and q = 1 / 3 [4].
A separate line of works seek to improve the dependence on T to be logarithmic through instance dependence [26, 61].
These results typically impose that the underlying MDP has a transition model that is linearly factorizable, and exhibit regret depends on the input dimension d .
This condition can be relaxed through introduction of a nonlinear feature map, whose appropriate selection is highly nontrivial and lead to large gaps between regret and practice [37], or metaprocedures [34].
Aside from the feature selection issue, these approaches require evaluation of confidence sets which is computationally costly and lead to statistical inefficiencies when approximated [42].
Thus, we prioritize Bayesian approaches to RL [21, 30] popular in robotics [17].
While many heuristics exist, performance guarantees take the form of Bayesian regret [40], and predominantly build upon posterior (Thompson) sampling [54].
In particular, beyond the tabular setting, [41] establishes a ˜ O ( σ R (cid:112)
d K ( R ) d E ( R ) T + E
[ L ∗ ] σ p (cid:112)
d K ( P ) d E ( P ))
Bayesian regret for posterior sampling RL (PSRL) combined with greedy action selections with respect to the estimated value.
Here L ∗ is a global Lipschitz constant for the future value function, d K and d E are Kolmogorov and eluder dimensions, and R and P refers to function classes of rewards and transitions.
The connection between H and L is left implicit; however, [19][Sec. 3.2] shows that L can depend exponentially on H .
Similar drawbacks manifest in the Lipschitz parameter of the Bayesian regret bound of [14], which extends the former result to continuous spaces through kernelized feature maps.
However, recently an augmentation of PSRL is proposed which employs feature embedding with Gaussian (symmetric distribution) dynamics to alleviate this issue [19], yielding the Bayesian regret of ˜ O
( H 32 d √ T ) that is polynomial in H and has no dependence on Lipschitz constant L .
These results are still restricted in the sense that it requires (i) the transition model target to be Gaussian, (ii) its representational complexity to grow unsustainably large with time.
Therefore, in this work we as the following question:
Can we achieve a trade-off between the Bayesian regret and the posterior representational complexity (aka coreset size) without oracle access to a feature map at the outset of training, in possibly continuous state-action spaces?
We provide an affirmative answer by honing in on the total variation norm used to quantify the posterior estimation error that appears in the regret analysis of [19], and identify that it can be sharpened by instead employing an integral probability metric (IPM).
Specifically, by shifting to an IPM, and then imposing structural assumptions on the target, that is, restricting it to a class of smooth densities, we can employ Stein’s identity [49, 31] to evaluate the distributional distance in closed form using the kernelized Stein discrepancy (KSD) [22, 35].
This restriction is common in Markov Chain Monte Carlo (MCMC) [3], and imposes that the target, for instance, belongs to a family of mixture models.
This modification in the metric of convergence alone leads to improved regret because we no longer require the assumption that the posterior is Gaussian [19].
However, our goal is to translate the scalability of PSRL from tabular settings to continuous spaces which requires addressing the parameterization complexity of the posterior estimate, which grows linearly unbounded [19] .
With the power to evaluate KSD in closed form, then, we sequentially remove those state-action pairs that contribute least (decided by a compression budget ϵ ) in
KSD after each episode (which is completely novel in the RL setting) from the posterior representation according to [25].
Therefore, the posterior estimate only retains statistically significant past samples from the trajectories, i.e., it is defined by a Bayesian coreset of the trajectory data [11, 12].
The budget parameter ϵ then is calibrated in terms of a rate determining factor α to yield both sublinear Bayesian regret and sublinear representational complexity of the learned posterior – see Table 2.
The resultant procedure we call Kernelized
Stein Discrepancy-based Posterior Sampling for RL ( KSRL ).
Our main contributions are, then, to: ▷ introduce Stein’s method in MBRL for the first time, and use it to develop a novel transition model estimate based upon it, which operates in tandem with a KSD compression step to remove statistically insignificant past state-action pairs, which we abbreviate as KSRL ; ▷ establish Bayesian regret bounds of the resultant procedure that is sublinear in the number of episodes experienced, without any prior access to a feature map, alleviating difficult feature selection drawbacks of prior art.
Notably, these results relax Gaussian and Lipschitz assumptions of prior related results; ▷ mathematically establish a tunable trade-off between Bayesian regret and posterior’s parameterization complexity (or dictionary size) via introducing parameter α ∈ (0 , 1] for the first time in this work; ▷ experimentally demonstrate that KSRL achieves favorable tradeoffs between sample and representational complexity relative to several strong benchmarks.
B Background
Let us write down the posterior sampling based reinforcement learning (PSRL) algorithm here in detail which is the basic building block of the research work in this paper [41].
Model Predictive Control : Model based RL planning with Model-predictive control (MPC) has achieved great success[9] in several continuous control problems especially due to its ability to efficiently incorporate uncertainty into the planning mechanism.
MPC has been an extremely useful mechanism for solving multivariate control problems with constraints [10] where it solves a finite horizon optimal control problem in a receding horizon fashion.
The integration of MPC in the Model-based RL is primarily motivated due to its implementation simplicity, where once the model is learnt, the subsequent optimization for a sequence of actions is done through MPC.
At each timepoint, the MPC applies the first action from the optimal action sequence under the estimated dynamics and reward function by solving arg max a i : i + τ (cid:80) i + τ t = i E
[ r ( s t , a t )] .
However, computing the exact arg max is nontrivial for non-convex problems and hence approximate methods like random sampling
[Table]
Tabular
[Table]
Tabular
[Table]
Tabular
[Table]
Tabular
[Table]
Tabular
[Table]
Tabular
[Table]
Continuous/Gaussian
[Table]
/ Smooth
[Table]
Algorithm
1: Input : Episode length
, Total
D episodes
Sample M ϕ
Evaluate
M do 6:
Take action
Observe r i and s i +1 action a i ∼ µ k
( s i ) 8:
Update C = C ∪ { ( s i , a i , r i , s i +1 ) } 9: end for 10:
Update D k +
=
D k ∪ C 11:
Update posterior to obtain ϕ ( · | D k +1 ) 12: end for shooting [36], cross-entropy methods [7] are commonly used.
For our specific case we use the Cross-entropy method for its effectiveness in sampling and data efficiency.
C Proof of Lemma ??
Proof.
Using the definition of the future value function with the Bellman Operator, we can write
[Equation]
where we have utilized the absolute upper bound on the value function | V ( s ) | ≤ HR max .
The norm in (3) is the total variation distance between the probability measures.
In [19,
Lemma 2], the right hand side of (44) is further upper bounded by the distance between the mean functions after assuming that the underlying distributions are Gaussian.
This is the point of departure in our analysis where we introduce the notion of Stein discrepancy to upper bound the total variation distance between the probability measures.
First, we build the analysis for d = 1 and later would extend it to multivariate scenarios.
We start the analysis by showing that the total variation distance is upper bounded by the
KSD for d = 1 .
Let us define the notion of an integral probability metric between two
( ϕ D k +1 , D k +1 ) =
KSD-Thinning ( ϕ (cid:101) D
k
+1 , (cid:101)
D k + ϵ k j as the minimial h i ∈ (cid:101) D k +1 ( ?? )
6: if KSD ( q W\{ x j } ) 2 <
α
+ ϵ then 7:
Remove the least influential point, set W = W \ { x j } 8: else 9:
Break loop 10: end if 11: end while 12: Output thinned dictionary W satisfying KSD ( q W ) 2 < α 2 + ϵ distributions p and q as
[Equation]
where F is any function space.
Now, if we restrict ourselves to a function space F ′ :
=
{ f : ∥ f ∥ ∞ ≤ 1 } , then we boils down to the definition of total variation distance between p and q given by
[Equation]
We can restrict our function class to be RKHS given by H and still write
[Equation]
where we define H ′ :
= { f : ∥ f ∥ ∞ ≤ 1 } .
We note that H ′ is uniquely determined by the kernel κ ( x, y )
(e.g., RBF kernel) .
Now, we know H ′ is a subset of RKHS H .
Now, we will try to upper bound the supremum over H ′ with supremum over a general class of functions which we call Stein class of functions as H S [35].
Note that since H ′
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Figure]
Now, for Stein class of functions H S , it holds that (cid:82) fdp = 0 .
Therefore, we can write
[Equation]
(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
fdq
[Table]
Hence, we can write
[Equation]
Utilizing the above upper bound into the right hand side of (3), we can write
[Equation]
The above result holds for d = 1 case, and assuming that the variables are independent of each other in all the dimensions, we can naively write that in d dimensions, we have
[Equation]
Hence proved.
D Proof of Lemma ??
Before starting the proof, here we discuss what is unique about it as compared to the existing literature.
The kernel Stein discrepancy based compression exists in the literature [24] but that is limited to the settings of estimating the distributions.
This is the first time we are extending the analysis to model based RL settings.
The samples in our setting are collected in the form of episodes following an optimal policy µ k (cf.
( ?? )) for each episode.
The analysis here follows a similar structure to the one performed in the proof of [24, Theorem 1] with careful adjustments to for the RL setting we are dealing with in this work.
Let us now begin the proof.
Proof.
Since we are learning for both reward and transition model denoted by R and P and they both are parameterized by the same dictionary D k , we present the KSD analysis for Λ D k :=
{ P k , R k } without loss of generality.
Further, for the proof in this section, we divide the H samples collected in k th episode into M > 1 number of batches, and then select the
KSD optimal point from each batch of H/M samples similar to SPMCMC procedure.
We will start with the one step transition at k − 1 where we have the dictionary D k − 1 and we update it to obtain (cid:101) D k which is before the thinning operation.
From the definition of KSD in ( ?? ), we can write
[Equation]
Next, for each m , since we select (SPMCMC based method in [13]) the sample h mk from
Y m :=
{ h lk } Ml =1 , and without loss of generality, we assume that H/M is an integer.
Now, from the SPMCMC based selection, we can write
[Equation]
The inequality in (8) holds because we restrict our attention to regions for which it holds that k 0 ( x , x ) ≤ S 2 k for all x ∈ Y mk for all k and m .
Utilizing the upper bound of (8) into the right hand side of (7), we get
[Equation]
Eqn (6) clearly differentiates our method from [24] highlighting the novelty of our approach is deciphering the application of KSD to our model based RL problem.
From the application of Theorem 5 [24] for our formulation with H new samples in the dictionary.
[Equation]
for any arbitrary constant r k > 0 .
We use the upper bound in (10) to the right hand side of (9), to obtain
[Equation]
Next, we divide the both sides by | (cid:101)
D k | 2 = ( |D k − 1 | + H/M ) 2 to obtain
KSD (Λ (cid:101) D k ) 2 ≤
[Equation]
Now, in this step we apply equation ( ?? ) and replace Λ (cid:101)
D k with the thinned dictionary one Λ D k and rewriting equation (12) as
[Equation]
Note that we have established a recursive relationship for the KSD of the thinned distribution in (13).
This is quite interesting because it would pave the way to establish the regret result which is the eventual goal.
After unrolling the recursion in (13)
[Equation]
we can write
[Equation]
Taking expectation on both sides of (14) and applying the log-sum exponential bound we get
[Equation]
[Equation]
where the log-sum exponential bound is used as
[Equation]
Next [Equation] we consider the inequality in (15), ignoring the exponential term, we further decompose the remaining summation term into two parts, respectively, as:
E
[Equation]
where T 1 corresponds to the sampling error and T 2 corresponds to the error due to the proposed thinning scheme.
The term T 1 represents the bias incurred at each step of the un-thinned point selection scheme.
The term T 2 is the bias incurred by the thinning operation carried out at each step.
However, for our case in the model-based reinforcement learning setting, the bias term T 1 will be different in our case as the samples are generated in a sequential manner in each episode from a Transition kernel with a decaying budget.
Let develop upper bounds on both T 1 and T 2 as follows.
Bound on T 1 : Let consider the first term on the right hand side of (17) as follow
[Equation]
[Figure]
where the equality in (18) holds by rearranging the denominators in the multiplication, and pulling ( |D k − 1 | + H/M ) 2 inside the first term.
Next, from the fact that |D j | ≤ |D j − 1 | + H/M which implies that the product will be less that 1 , we can upper bound the right hand side of (18) as follows
[Equation]
where we took expectation inside the summation and apply it to the random variable in the numerator.
From the model order growth condition, we note that |D k − 1 | + H/M ≥ |D
k | ≥ f ( k ) , which implies that 1 / ( |D k − 1 | + H/M ) 2 ≤ 1 /f ( k ) 2 , which we utilize in the right hand side of (20) to write
[Equation]
Following the similar logic mentioned in [13, Appendix A, Eqn.
(17)], we can upper bound E (cid:2) ∥
f i ∥ 2 K 0 (cid:3) as
[Equation] [Equation]
which is based on the assumption that E y ∼ P [exp( γκ 0 ( y , y ))]
= b < ∞ .
This implies that
[Equation] [Equation]
which we obtain by applying the upper bound in (23) into (22).
After simplification, we can write
[Equation]
which provide the bound on T 1 .
Next, we derive the bound on T 2 as follows.
Bound on T 2 :
Let us consider the term T 2 from (17) as follows
[Equation]
This term is extra in the analysis and appears due to the introduction of compression budget ϵ k into the algorithm.
We need control the growth of this term, and by properly designing ϵ k , we need to make sure T 2 goes to zero at least as fast at T 1 to obtain a sublinear regret analysis.
We start by observing that |D
j | ≤ ( H/M ) j which holds trivially and hence implies
[Equation]
From the algorithm construction, we know that |D
j +1 | ≤ |D j | + H/M , applying this to (27), we obtain
[Equation]
Utilize the upper bound in (28) to the right hand side of (26), to write
[Equation]
The above bound implies that we should choose the compression budget
ϵ i such that Tgoes to zero at least as fast at T[-
-]
[Equation]
To satisfy the above condition, we choose ϵ i = log( i )
f ( i ) 2 , and obtain
[Equation]
[Figure]
which satisfy the upper bound in (31), which shows that ϵ i = log( i )
f ( i ) 2 is a valid choice.
Hence,
T
[Equation]
Finally, after substituting the upper bounds for T 1 and T 2 into (17) [Equation] we obtain
E
[Equation]
Now revisiting the inequality in (14) [Equation] we write
[Equation]
where the last equality holds by the selection r j = HkM for all j .
Finally [Equation] we collect all the constants in C to write
[Equation]
Next, from applying Jensen’s inequality on the left hand side in (35) [Equation] and we can write
[Equation]
Hence proved.
E Proof of Theorem ??
Proof.
Based on the equality in (23), the first part of our regret in ( ?? ) becomes zero and we only need to derive and upper bound for ∆ IIk .
Recall that, we have
[Equation]
Following the Bellman equations representation of value functions [Equation] We can factorize ∆ IIk in-terms of the immediate rewards r and the future value function U (cf.
[41]) as follows
[Equation]
where we define
[Equation]
Now
[Equation]
for our Stein based thinning algorithm [Equation] we thin the updated dictionary (cid:101)
D k after every episode ( H steps) and obtain a compressed and efficient representation of dictionary given by D k .
Let us derive the upper bound on E [∆ k ( P )] as follows.
The Bayes regret at the k th episode can be written as
[Equation]
Utilize the upper bound from the statement of Lemma ??
to write
[Equation]
This is a an important step and point of departure from the existing state of the art regret analysis for model based RL methods [19, 41].
Instead of utilizing the naive upper bound of Total Variation distance in the right hand side of (41), we follow a different approach and bound it via the Kernel Stein Discrepancy which is a novel connection explored for the first time in this work.
We remark that the KSD upper bound in Lemma ??
is for the joint posterior where samples are h i
= ( s, a, s ′ ) .
We note here that since the score function if independent of the normalizing constant, we can write ∇ log ˜ P k ( ·| h ′ i )
=
∇ log ˜ P k
( h i ) , therefore we utilize the KSD upper bound of joint posterior in Lemma ??
on the KSD term per episode in the right hand side of (42) to obtain
[Equation]
Next, we take summation over the number of episodes which are given by [ TH ] where T is the total number of time steps in the environments, and H is the episode length.
So, after
[Figure]
summing over k = 1 to [ TH ] , we obtain
[Equation]
To derive the explicit regret rates, we assume our dictionary growth function f ( k ) =
(cid:112) k
α +1 log( k ) with range of α ∈ [0 , 1] .
Substituting this into (44), we can write
TH (cid:88) k
[Equation]
(cid:112)
[Equation]
Using (45) into (44) [Equation] we get [
[Equation]
The expression implies that
[Equation]
[Figure]
The same derivation would hold for the term E [∆ k ( r )] (similar logic to [19, Sec. 3.4]0), which would imply that (cid:80) [ TH ] k =
[Equation]
[Equation]
From (38) and (47), we can write [ TH ] (cid:88)
[Equation]
Hence proved.
F Additional Experiments and Analysis
In this section, we first provide details of the environments and the complexities added in order to validate multiple aspects of our KSRL .
Low-dimensional environments
: We consider the Continuous Cartpole ( d s = 4 , d a = 1 , H = 200 ) environment with a continuous action space which is a modified version of the discrete action classic Cartpole environment.
The continuous action space enhances the complexity of the environment and makes it hard for the agent to learn.
We also consider the Pendulum Swing Up ( d s = 3 , d a = 1 , H = 200 ) environment, a modified version of Pendulum where we limit the start state to make it harder and more challenging for exploration.
To introduce stochasticity into the dynamics, we modify the physics of the environment with independent Gaussian noises ( N (0 , 0 . 01) ).
However, these environments are primarily lower dimensional environments.
Higher-dimensional environments We consider the 7-DOF Reacher ( d s = 17 , d a = 7 , H = 150 ) and 7-DOF pusher ( d s = 20 , d a = 7 , H = 150 )
two challenging continuous control tasks as detailed in [15].
We increase the complexity from lower to higher dimensional environments and with added stochasticity to see the robustness of our algorithm and the consistency in performance.
We conduct the experiments both with and without true oracle rewards and compare the performance with other baselines.
F.1 Experimental Details
It is shown in literature [19] that a simple Bayesian Linear regression with non-linear feature representations learnt by training a Neural network works exceptionally well in the context of posterior sampling reinforcement learning.
For a fair comparison of our KSRL , we follow a similar architecture as [19] where we first train a deep neural network for both the transition and rewards model and extract the penultimate layer of the network for the
Bayesian linear regression and Posterior sampling.
As per our notation s i , a i , s i +1 , r i where h i = <
s i , a i > be the current state-action pair with the next state s i +
and reward r i and let’s assume the representation from the penultimate layer of the deep neural network of the state-action pair is denoted by z i
= NN ( h i ) ∈
R d where NN is the trained deep neural network model and d is the dimensionality of the representation.
Then the Bayesian linear regression model deals with learning the posterior distribution P post
= P ( β | D ) , with the linear model given by δ i =
β T z i + ϵ , where δ i =
s i
+
− s i as suggested in [17, 19] and ϵ ∼ N (0 , σ 2 ) , D is the size of the dictionary.
Similar to prior approaches, we choose a multivariate Gaussian prior with zero mean and Σ prior (conjugate prior) to obtain a closedform estimation of the posterior distribution which is also multivariate Gaussian.
Now, we sample β from the posterior distribution P post at the beginning of each episode and interact with the environment using MPC controller.
As described in Appendix A at each timepoint, the MPC applies the first action from the optimal action sequence under the estimated dynamics and reward function by solving arg max a i : i + τ (cid:80) i + τ t = i
E
[ r ( s t , a t )] , where τ is the horizon and is considered as a hyperparameter.
However, as described above the above method suffers from high computation complexity as the matrix multiplication step in posterior estimation is of order
O ( d 2 N ) which scales linearly with the size of the dictionary prior to that episode N .
This not only enhances the computational complexity but also makes the optimization with MPC extremely hard.
Hence, in our algorithm KSRL we construct an efficient posterior coreset with Kernelized Stein discrepancy measure from ( ?? ) and validate the average reward achieved with the compressed dictionary.
We observe in all the cases with varied complexity, our algorithm KSRL performs equally well or sometimes even better than the uncompressed current state of the art MPC-PSRL method with drastically reduced model complexity
[Figure]
F.2 Experimental Analysis
We perform a detailed multifaceted analysis comparing our algorithm KSRL with other baselines and state of the art methods in the continuous control environments.
We perform the experiments in the environments with and without oracle rewards.
The environmental setting of without oracle rewards is much more complex as here we have to model the reward function as well along with the dynamics which makes it much harder for the agent to learn with the added uncertainty in modelling.
We also enhance the complexity of the environments by adding stochasticity which makes it harder for the agent to learn even for environments with oracle rewards and we also observe the performance by varying the dimensionality of the environments.
In Figure 2 and Figure 4, we compare our KSRL with other baselines and SOTA algorithms for Cartpole, Pendulum and Pusher environments without and with oracle rewards respectively and Figure 5 for Reacher with and without rewards.
In all the cases, KSRL shows performance at-par or even better for some cases with drastically reduced model complexity where we can see a benefit of 80% improvement in the model complexity over the current SOTA with similar performance in terms of average rewards.
In Figure 3, we validate the average reward achieved by our KSRL against baselines with respect to the runtime (wallclock time) in CPU minutes and clearly observe improved performance in-terms of wallclock time where KSRL achieves optimal performance prior to the baselines and MPC-PSRL.
We also perform the convergence analysis from an empirical perspective and observe the convergence in-terms of both Kernelized Stein
Discrepancy and Posterior Variance.
In Figure 3 (d) -(f) , we study the convergence in-terms of KSD and observe the convergence of our algorithm KSRL without any bias and faster than the dense counterpart MPC-PSRL in-terms of wall clock time.
We also study the convergence from the posterior variance perspective as it is extremely important for PSRL based algorithms to accurately estimate the uncertainty.
Since, we are compressing the dictionary it is important to analyze and monitor the posterior variance over the timesteps to ensure that we are not diverging and close to the dense counterparts.
In 6, we observe the posterior variance achieved by KSRL converges to the posterior variance achieved by
MPC-PSRL [19] even though we are compressing the dictionary which highlights the efficacy of our posterior coreset.
Finally, from the above plots and analysis we conclude that our our algorithm KSRL achieves state of the art performance for continuous control environments with drastically reduced model complexity of its dense counterparts with theoretical guarantees of convergence.
F.3 Details of Hyperparamters
[Figure]
Networkarchitecture [Figure]
of
[Figure]
with 2 hidden layers of size
MLP with 4 hidden layers of size MLP with 4 hidden layers of size
[Figure]
algorithm KSRL
We keep the hyperparameters and the network architecture almost similar to [19] for a fair comparison.
For the baseline implementation of MPC-PSRL algorithm and our algorithm
KSRL we modify and leverage 1 and observe that we were able to achieve performance as described in [19].
For obtaining the results from model-free algorithms as shown in we use 2 , and could replicate the results.
Finally, for our posterior compression algorithm we modify the 3 , 4 to fit to our scenario of posterior sampling reinforcement learning.
We are thankful to all the authors for open-sourcing their repository.
Acknowledgments and Disclosure of Funding
Use unnumbered first level headings for the acknowledgments.
All acknowledgments go at the end of the paper before the list of references.
Moreover, you are required to declare funding (financial activities supporting the submitted work) and competing interests (related financial activities outside the submitted work).
More information about this disclosure can be found at: https://neurips.cc/Conferences/2022/PaperInformation/
FundingDisclosure .
Do not include this section in the anonymized submission, only in the final paper.
You can use the ack environment provided in the style file to autmoatically hide this section in the anonymized submission.