Langevin Thompson Sampling with Logarithmic Communication: Bandits and Reinforcement Learning

Modern machine learning often needs to balance computation and communication budgets with statistical guarantees. Existing analyses of sequential decision making have been primarily focused on the statistical aspects of the problems (Tian et al.,, 2020), less is known about their computation and communication aspects. In particular, regret minimization in multi-armed bandits (MABs) and reinforcement learning (RL) (Jaksch et al.,, 2010; Wu et al.,, 2022; Jung et al.,, 2019) is often studied under the common assumptions that computation can be performed perfectly in time, and that communication always happens in real-time (Li et al.,, 2022; Jin et al.,, 2018; Haarnoja et al.,, 2018).

A question of particular importance is whether optimal decisions can still be made under reasonable computation and communication budgets. In this work, we study the exploration-exploitation problem with low computation and communication costs using Thompson Sampling (TS) (Thompson,, 1933) (a.k.a. posterior sampling). To allow sampling from distributions that deviate from the standard restrictive assumptions, and to enable its deployment in settings where computing the exact posterior is challenging, we employ Markov Chain Monte Carlo (MCMC) methods.

TS operates by maintaining a posterior distribution over the unknown and is widely used owing to its strong empirical performance Chapelle and Li, (2011). However, the theoretical understanding of TS in bandits typically relied on the restrictive conjugacy assumptions between priors and reward distributions. Recently, approximate TS methods for general posterior distributions start to be studied Mazumdar et al., (2020); Xu et al., (2022), where MCMC algorithms are used in conjunction with TS to expand its applicability in fully-sequential settings. On the other hand, to the best of our knowledge, how to provably incorporate MCMC with posterior sampling in RL domains remains to be untackled.

Moreover, previous analyses of TS have been restricted to the fully-sequential settings, where feedback or reward is assumed to be immediately observable upon taking actions (i.e. before making the next decision). Nevertheless, it fails to account for the practical settings when delays take place in communication, or when feedback is only available in an aggregate or batched manner. Examples include clinical trials where feedback about the efficacy of medication is only available after a nontrivial amount of time, recommender systems where feedback from multiple users comes all at once and marketing campaigns Schwartz et al., (2017). This issue has been studied in literature by considering static or dynamic batching schemes and by designing algorithms that acquire feedback in batches, where the learner typically receives reward information only at the end of the batch (Karbasi et al.,, 2021; Kalkanli and Ozgur,, 2021; Vernade et al.,, 2020; Zhang et al.,, 2020). Nonetheless, the analysis of approximate TS in batched settings is unavailable for both bandits and RL.

In this paper, we tackle these challenges by incorporating TS with Langevin Monte Carlo (LMC) and batching schemes in stochastic MABs and infinite-horizon RL. Our algorithms are applicable to a broad class of distributions with only logarithmic rounds of communication between the learner and the environment, thus being robust to constraints on communication. We compare our results with other works in Table 1, and summarize our main contributions as follows:

• •

  For stochastic MABs with time horizon $T$, we present Langevin Thompson Sampling (BLTS, Algorithm 2) along with Theorem 5.2, which achieves the optimal $O(\log T)$ regret with $O(\log T)$ batches¹¹1A $T$ round game can be thought of as $T$ many batches each of size 1.. The main technical contribution here is to show that when feedback is obtained in a batched manner where the posterior concentration is weaker (Theorem 1), the convergence guarantee of SGLD continues to hold.

• •

  For large-scale infinite-horizon MDPs, we present Langevin Posterior Sampling for RL (LPSRL, Algorithm 3) along with Theorem 6.2 to show that SGLD with a static policy-switching²²2In MDP settings, the notion of a batch is more appropriately thought of as a policy-switch. scheme achieves the optimal $O(\sqrt{T})$ Bayesian regret with $O(\log T)$ policy switches. For tabular MDPs, we show that LPSRL incorporated with the Mirrored Langevin Dynamics (MLD) achieves the optimal $O(\sqrt{T})$ Bayesian regret with $O(\log T)$ policy switches. The use of approximate sampling leads to an additive error where the true model and the sampled model are no longer identically distributed. This error can be properly handled with the convergence guarantees of LMC methods.

• •

  Experiments are performed to demonstrate the effectiveness of our algorithms, which maintain the order-optimal regret with significantly lower communication costs compared to existing exact TS methods.

In this section, we introduce the problem setting with relevant background information.

Under the conjugacy assumptions on rewards, asymptotic convergence of TS was studied in stochastic MABs by Granmo, (2010) and May et al., (2012). Later, finite-time analyses with $O(\log T)$ problem-dependent regret bound were provided Agrawal and Goyal, (2012); Kaufmann et al., (2012); Agrawal and Goyal, (2013). However, in practice, exact posteriors are intractable for all but the simplest models (Riquelme et al.,, 2018), necessitating the use of approximate sampling methods with TS in complex problem domains. Recent progress has been made in understanding approximate TS in fully-sequential MABs Lu and Van Roy, (2017); Mazumdar et al., (2020); Zhang, (2022); Xu et al., (2022). On the other hand, the question of learning with TS in the presence of batched data has evolved along a separate trajectory of works (Karbasi et al.,, 2021; Kalkanli and Ozgur,, 2021; Vernade et al.,, 2020; Zhang et al.,, 2020). However, provably performing Langevin TS in batched settings remains unexplored, and in this paper, we aim at bridging these lines.

Moving to the more complex decision-making frameworks based on MDPs, TS is employed in model-based methods to learn transition models, which is known as Posterior Sampling for Reinforcement Learning (PSRL) (Strens,, 2000). When exact posteriors are intractable, MCMC methods have been empirically studied for performing Bayesian inference in policy and reward spaces in RL (Brown et al.,, 2020; Imani et al.,, 2018; Bojun,, 2020; Guez et al.,, 2014). MCMC is a family of approximate posterior inference methods that enables sampling without exact knowledge of posteriors Ma et al., (2015); Welling and Teh, (2011). However, it is unclear how to provably incorporate MCMC methods in learning transition models for RL.

Furthermore, the analysis of undiscounted infinite-horizon MDPs (Abbasi-Yadkori and Szepesvári,, 2015; Osband and Van Roy,, 2016; Ouyang et al.,, 2017; Wei et al.,, 2020, 2021) poses greater challenges compared to the well-studied episodic MDPs with finite horizon and fixed episode length (Osband et al.,, 2013). Previous works on infinite-horizon settings include model-based methods that estimate environment dynamics and switch policies when the number of visits to state-action pairs doubles (Jaksch et al.,, 2010; Tossou et al.,, 2019; Agrawal and Jia,, 2017; Bartlett and Tewari,, 2012). Nevertheless, under such dynamic schemes, the number of policy switches can be as large as $O(\sqrt{T})$, making it computationally heavy and infeasible for continuous states and actions. To enable TS with logarithmic policy switches while maintaining optimal regret, we build upon an algorithmically-independent static scheme as in Theocharous et al., 2017b , and incorporate Langevin Monte Carlo (LMC) methods to sample from inexact posteriors.

In the MAB and MDP settings, $\theta$ parameterizes the unknown reward or transition distributions respectively. TS maintains a distribution over the parameters and updates the distribution to (the new) posterior upon receiving new data. Given $p(X|\theta)$, prior $\lambda(\theta)$, and $n$ data samples $\{X_{i}\}_{i=1}^{n}$⁶⁶6Here data $\{X_{i}\}_{i=1}^{n}$ can be rewards for some arms or actual transitions of state-action pairs depending on the setting., let $\rho_{n}$ be the posterior distribution after receiving $n$ data samples which satisfies: $\rho(\theta|\{X_{i}\}_{i=1}^{n})\propto\exp(\sum_{i=1}^{n}\log p(X_{i}|\theta)+\log\lambda(\theta))$. In addition, consider the scaled posterior $\rho_{n}[\gamma]$ for some scaling parameter $\gamma$, which represents the density proportional to $\exp(\gamma(\sum_{i=1}^{n}\log p(X_{i}|\theta)+\log\lambda(\theta)))$.

The introduction of MCMC methods arises from the need for sampling from intractable posteriors in the absence of conjugacy assumptions. We resort to a gradient-based MCMC method that performs noisy updates based on Langevin dynamics: Stochastic Gradient Langevin Dynamics (SGLD). Algorithm 1 presents SGLD with bached data to generate samples from an approximation of the true posterior. For a detailed exposition, please refer to Welling and Teh, (2011); Ma et al., (2015) and Appendix A. Algorithm 1 takes all available data $\{X_{s}\}_{s=1}^{n}$ at the start of a batch $b$ as input, subsamples data, performs gradient updates by computing $\nabla\widehat{U}(\theta)=-\frac{n}{|D|}\sum_{X_{s}\in D}\nabla\log p(X_{s}|\theta)-\nabla\log\lambda(\theta)$, and outputs the posterior for batch $b$.

In the batched setting, new data is received at the end of a batch, or when making the decision to perform a new policy switch. Due to the way that the learner receives new data and the fact that the batch data size may increase exponentially⁷⁷7Data received in batch $k$ can be doubled compared to the previous batch., the posterior concentrates slower. This differs from the fully-sequential problem where the distribution shift of successive true posteriors is small owing to data being received in an iterative manner. We show that in batched settings, with only constant computational complexity in terms of iterations, SGLD is able to provide strong convergence guarantee as in the fully-sequential setting Mazumdar et al., (2020). Theorem 1 shows the convergence of SGLD in the Wasserstein-$p$ distance can be achieved with a constant number of iterations and data.

$\rho_{n}$ denotes the true posterior corresponding to $n$ data samples and $\tilde{\rho}_{n}$ is the approximate posterior outputted by Algorithm 1. We also note that similar concentration bounds can be achieved by using the Unadjusted Langevin Algorithm (ULA) for batched data, which adopts full-batch gradient evaluations and therefore leads to a growing iteration complexity. The proofs of Theorem 1 are adapted to the batched setting, which differs from Mazumdar et al., (2020).

In this section, we introduce Langevin Thompson Sampling for batched stochastic MAB setting in Algorithm 2, namely, BLTS. It leverages SGLD and batching schemes to learn a wide class of unknown reward distributions while reducing communication and computation costs. We have previously discussed the results of SGLD in Section 4 for both MABs and MDPs. Here, we focus on the batching strategy in Algorithm 2 for bandits, and discuss the resulting regret guarantee.

In RL frameworks, posterior sampling is commonly used in model-based methods to learn unknown transition dynamics and is known as PSRL⁸⁸8We also depart from using TS for the RL setting and stick to the more popular posterior sampling terminology for RL.. In infinite-horizon settings, PSRL operates by sampling a model and solving for an optimal policy based on the sampled MDP at the beginning of each policy switch. The learner then follows the same policy until the next policy switch. In this context, the concept of a batch corresponds to a policy switch.

Previous analyses of PSRL have primarily focused on transition distributions that conform to well-behaved conjugate families. Handling transitions that deviate from these families and computing the corresponding posteriors has been heuristically left to MCMC methods. Here, we provably extend PSRL with LMC and introduce Langevin Posterior Sampling for RL (LPSRL, Algorithm 3) using a static doubling policy-switch scheme. Analyses of PSRL have crucially relied on the true transition dynamics $\theta^{*}$ and the sampled MDPs being identically distributed (Osband et al.,, 2013; Osband and Van Roy,, 2016; Russo and Van Roy,, 2014; Ouyang et al.,, 2017; Theocharous et al., 2017b, ). However, when the dynamics are sampled from an approximation of the true posterior, this fails to hold. To address the issue, we introduce the Langevin posterior sampling lemma (Lemma 6), which shows approximate sampling yields an additive error in the Wasserstein-$1$ distance.

lemmaLangevinPS (Langevin Posterior Sampling). Let $t_{k}$ be the beginning time of policy-switch $k$, $\mathcal{H}_{t_{k}}:=\{s_{\tau},a_{\tau}\}_{\tau=1}^{t_{k}}$ be the history of observed states and actions till time $t_{k}$, and $\theta^{k}\sim\tilde{\rho}_{t_{k}}$ be the sampled model from the approximate posterior $\tilde{\rho}_{t_{k}}$ at time $t_{k}$. Then, for any $\sigma(\mathcal{H}_{t_{k}})$-measurable function $f$ that is $1$-Lipschitz, it holds that:

By the tower rule, $\Big{|}\mathbb{E}[f(\theta^{*})]-\mathbb{E}[f(\theta^{k})]\Big{|}\leq W_{1}(\tilde{\rho}_{t_{k}},\rho_{t_{k}})$.

As demonstrated later, this error term can be effectively controlled and does not impact the overall regret (Theorems 6.2 and 6.3). It only requires the average reward function $J^{\pi}(\theta)$ to be 1-Lipschitz, as specified in Assumption B⁹⁹9Mathematical statement is in Appendix B.. Let us consider the parameterization of the transition dynamics $p$ with $\theta\in\mathbb{R}^{d}$, where $\theta^{*}\in\mathbb{R}^{d}$ denotes the true (unknown) parameter governing the dynamics. We explore two distinct settings based on these parameterizations:

• •

  General Parameterization (Section 6.2): In this setting, we consider modeling the full transition dynamics using $\theta^{*}\in\mathbb{R}^{d}$, where $d\ll|\mathcal{S}||\mathcal{A}|$. This parameterization can be particularly useful for tackling large-scale MDPs with large (or even continuous) state and action spaces. Towards this end, we consider $\mathcal{S}\cong\mathbb{R}$. Examples of General Parameterization include linear MDPs with feature mappings Jin et al., (2020), RL with general function approximation Yang et al., (2020), and the low-dimensional structures that govern the transition (Gopalan and Mannor,, 2015; Yang and Wang,, 2020). We provide a real-world example that adopts such parameterization in Appendix E.3.

  Despite Theocharous et al., 2017b studies a similar setting, their work confines the parameter space to $\mathbb{R}$. To accommodate a broader class of MDPs, we generalize the parameter space to $\mathbb{R}^{d}$. As suggested by Theorem 6.2, our algorithm retains the optimal $O(\sqrt{T})$ regret with $O(\log T)$ policy switches, making it applicable to a wide range of general transition dynamics.

• •

  Simplex Parameterization (Section 6.3): Here, we consider the classical tabular MDPs with finite states and actions. For each state-action pair, there exists a probability simplex $\Delta^{|\mathcal{S}|}$ that encodes the likelihood of transitioning into each state. Hence, in this case, $\theta^{*}\in\mathbb{R}^{d}$ with $d=|\mathcal{S}|^{2}|\mathcal{A}|$. This structure necessitates sampling transition dynamics from constrained distributions, which naturally leads us to instantiate LPSRL with the Mirrored Langevin Dynamics (Hsieh et al.,, 2018) (See Appendix  A for more discussions). As proven in Theorem 6.3, LPSRL with MLD achieves the optimal $O(\sqrt{T})$ regret with $O(\log T)$ policy switches for general transition dynamics subject to the probability simplex constraints.

In this section, we perform empirical studies in simulated environments for bandit and RL to corroborate our theoretical findings. By comparing the actual regret (average rewards) and the number of batches for interaction (maximum policy switches), we show Langevin TS algorithms empowered by LMC methods achieve appealing statistical accuracy with low communication cost. For additional experimental details, please refer to Appendix F.

In this paper, we jointly address two challenges in the design and analysis of Thompson sampling (TS) methods. Firstly, when dealing with posteriors that do not belong to conjugate families, it is necessary to generate approximate samples within a reasonable computational budget. Secondly, when interacting with the environment in a batched manner, it is important to limit the amount of communication required. These challenges are critical in real-world deployments of TS, as closed-form posteriors and fully-sequential interactions are rare. In stochastic MABs, approximate TS and batched interactions are studied independently. We bridge the two lines of work by providing a Langevin TS algorithm that works for a wide class of reward distributions with only logarithmic communication. In the case of undiscounted infinite-horizon MDP settings, to the best of our knowledge, we are the first to provably incorporate approximate sampling with the TS paradigm. This enhances the applicability of TS for RL problems with low communication costs. Finally, we conclude with experiments to demonstrate the appealing empirical performance of the Langevin TS algorithms.

This work is supported in part by the National Science Foundation Grants NSF-SCALE MoDL(2134209) and NSF-CCF-2112665 (TILOS), the U.S. Department Of Energy, Office of Science, and the Facebook Research award. Amin Karbasi acknowledges funding in direct support of this work from NSF (IIS-1845032), ONR (N00014- 19-1-2406), and the AI Institute for Learning-Enabled Optimization at Scale (TILOS).