Posterior Coreset Construction with Kernelized Stein Discrepancy for Model-Based Reinforcement Learning

Abstract

Recent progress on improving theoretical sample efficiency of model-based reinforcement learning (RL), which exhibits superior sample complexity in practice, requires Gaussian and Lipschitz assumptions on the transition model, and additionally defines a posterior representation that grows unbounded with time. In this work, we propose a novel Kernelized Stein Discrepancy-based Posterior Sampling for RL algorithm (named KSRL) which extends model-based RL based upon posterior sampling (PSRL) in several ways: we (i) relax the need for any smoothness or Gaussian assumptions, allowing for complex mixture models; (ii) ensure it is applicable to large-scale training by incorporating a compression step such that the posterior consists of a Bayesian coreset of only statistically significant past state-action pairs; and (iii) develop a novel regret analysis of PSRL based upon integral probability metrics, which, under a smoothness condition on the constructed posterior, can be evaluated in closed form as the kernelized Stein discrepancy (KSD). Consequently, we are able to improve the O(H 3/2 d √ T ) regret of PSRL to O(H 3/2 √ T ), where d is the input dimension, H is the episode length, and T is the total number of episodes experienced, alleviating a linear dependence on d . Moreover, we theoretically establish a trade-off between regret rate with posterior representational complexity via introducing a compression budget parameter based on KSD, and establish a lower bound on the required complexity for consistency of the model. Experimentally, we observe that this approach is competitive with several state of the art RL methodologies, with substantive improvements in computation time. 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.

STEERING: Stein Information Directed Exploration for Model-Based Reinforcement Learning

“…We start by deriving an upper bound for E[K π k (M ; H k,H )] with an SPMCMC style local optimization method proposed originally in (Chen et al, 2019) and later used in sequential decision-making scenarios by (Chakraborty et al, 2022b;Hawkins et al, 2022). We start with the definition in (8) to write…”

Posterior Coreset Construction with Kernelized Stein Discrepancy for Model-Based Reinforcement Learning

“…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. More specifically, observe that in step 10 and 11 in PSRL (see Algorithm in Appendix (Chakraborty et al 2022)), the dictionary at each episode k retains H additional points, i.e., |D k+1 | = |D k | + H. Hence, 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 step 11 in PSRL (cf. Algorithm ??)…”

“…We summarize the proposed algorithm in Algorithm 1 with compression subroutine in Algorithm 2, where KSRL is an abbreviation for Kernelized Stein Discrepancy Thinning for Model-Based Reinforcement Learning. Please refer to discussion Appendix (Chakraborty et al 2022) for MPC-based action selection.…”

“…See Appendix (Chakraborty et al 2022) for proof. Observe that Lemma 4.1 is unique to this work, and is a key point of departure from (Osband and Van Roy 2014;Fan and Ming 2021).…”

“…See Appendix (Chakraborty et al 2022) for proof. The inequality established in Lemma 4.2 relates the KSD to the number of episodes experienced by a model-based RL method, and may be interpreted as an adaption of related rates of posterior contraction in terms of KSD that have appeared in MCMC (Chen et al 2019;?).…”