Conceptio › Archive › arXiv CS

arXiv CSopen access

Optimal Posterior Sampling for Policy Identification in Tabular Markov Decision Processes

2026 · arxiv_cs

arXiv CS · Papers · License: Open Access · 2026

Open Source ↗Direct PDF ↓

neuralnetworks

machine learning, deep learning, neural networks

Optimal Posterior Sampling for Policy Identification in Tabular Markov Decision Processes

arXiv:2605.03921v1 [cs.LG] 5 May 2026

Cyrille Kone Univ. Lille, CNRS, Inria, Centrale Lille, UMR 9189-CRIStAL, F-59000 Lille, France

Abstract We study the (ε, δ)-PAC policy identification problem in finite-horizon episodic Markov Decision Processes. Existing approaches provide finite-time guarantees for approximate settings (ε > 0) but suffer from high computational cost, rendering them hard to implement, and also suffer from suboptimal dependence on log(1/δ). We propose a randomized and computationally eﬀicient algorithm for best policy identification that combines posterior sampling with an online learning algorithm to guide exploration in the MDP. Our method achieves asymptotic optimality in sample complexity, also in terms of posterior contraction rate, and runs in O(S 2 AH) per episode, matching standard model-based approaches. Unlike prior algorithms such as MOCA and PEDEL, our guarantees remain meaningful in the asymptotic regime and avoid sub-optimal polynomial dependence on log(1/δ). Our results provide both theoretical insights and practical tools for eﬀicient policy identification in tabular MDPs.

INTRODUCTION

In online reinforcement learning (RL), an agent sequentially interacts with an unknown Markov decision process (MDP) by observing trajectories composed of transitions and rewards generated by the environment. We are interested in the setting where interactions occur in episodes of finite length H: an episodic finitehorizon setting (Puterman 1994). Put formally, the MDP is described by M ≜ (S, A, H, p, µ, sinit ) where Proceedings of the 29th International Conference on Artificial Intelligence and Statistics (AISTATS) 2026, Tangier, Morocco. PMLR: Volume 300. Copyright 2026 by the author(s).

Kevin Jamieson University of Washington, Seattle, USA S is the state space of size S, A is the action space of size A, p and µ denote the transition and reward kernels, and sinit is the initial state. Each episode starts in the initial state sinit . At stage h ∈ [H], when in state sh ∈ S, the agent selects an action ah ∈ A according to its strategy, collects a (random) reward rh , and transitions to the (random) next state sh+1 according to the environment dynamics. After H steps, the environment resets to sinit . The agent must learn to eﬀiciently explore the environment in order either to maximize cumulative rewards across episodes (Jaksch et al. 2010) or, after an exploration phase, to identify a “good” policy that maximizes the expected sum of rewards per episode (Fiechter 1994; Kearns & Singh 1998). In this work, we focus on the latter objective, studied in the celebrated (ε, δ)-PAC (Probably Approximately Correct) policy identification setting, where the agent explores adaptively and eventually stops to recommend a policy it believes to be ε-optimal. The goal is to keep the exploration phase as short as possible while guaranteeing the correctness of the recommendation. In this framework, the agent follows adaptive exploratory policies and stops at a (random) time τ , when it recommends a candidate policy π̂τ believed to be optimal. The objective is to minimize the exploration time while ensuring PAC correctness. This setting is a pure exploration problem, since performance is evaluated only through the final recommended policy, not during the exploration phase. An algorithm for (ε, δ)-PAC policy identification typically has three components : (i) an adaptive exploration policy ρt to gather new observations at episode t, (ii) a recommendation rule π̂t representing the guess of the learning agent for the best policy and (iii) a stopping rule τ specifying when to terminate interaction with the environment. Early work on (ε, δ)-PAC policy identification adapted algorithms from the regret-minimization setting by incorporating adaptive stopping rules, and established

Optimal Posterior Sampling for Policy Identification in Tabular Markov Decision Processes

worst-case guarantees on the sample complexity (Azar, Munos, et al. 2012; Kaufmann, Ménard, et al. 2021a; Menard et al. 2021). However, subsequent results (A. J. Wagenmaker, Simchowitz, et al. 2022; Tirinzoni, Al-Marjani, et al. 2023) proved such rates to be inherently suboptimal from a problem-dependent perspective. Recent research has therefore shifted toward identifying the exact complexity term governing policy identification. In this regime, racing-style elimination algorithms have been proposed (Tirinzoni, Al Marjani, et al. 2022; A. Wagenmaker & Jamieson 2022; A. J. Wagenmaker, Simchowitz, et al. 2022; Narang et al. 2024). These approaches often rely on sophisticated algorithmic designs or exhaustive enumeration of deterministic policies, which leads to significant computational challenges. Another line of work has introduced tracking-based algorithms (Al Marjani et al. 2021; Marjani & Proutiere 2021); while promising, these methods are also computationally costly. Importantly, none of these approaches achieves instancedependent optimality when δ is small. In this work, we revisit the problem of (ε, δ)-PAC policy identification in the tabular episodic setting with a focus on the instance-dependent complexity when δ is small. We propose an algorithm that relies on posterior sampling and avoids costly policy enumeration and optimization steps required by many prior methods. 1.1

Related Work

The earliest results on (ε, δ)-PAC policy identification focused on establishing worst-case sample complexity bounds, i.e., the dependence on S, A, H, ε, and δ (Fiechter 1994; Kearns & Singh 1998; Azar, Munos, et al. 2012; Azar, Munos, et al. 2013; Dann & Brunskill 2015; Sidford, Wang, Wu, L. Yang, et al. 2018; Sidford, Wang, Wu & Ye 2018; Agarwal et al. 2020; Menard et al. 2021). Notably, the most recent of these 3

results matches the lower bound O SAH log(1/δ) ε2 proven by Domingues et al. 2021. The corresponding algorithms are often derived from optimistic methods originally designed for regret minimization (Azar, Osband, et al. 2017) or from online-to-batch and regretto-PAC reductions (Jin, Allen-Zhu, et al. 2018; Tirinzoni, Al-Marjani, et al. 2023). While such worstcase guarantees provide a complete characterization of performance in the hardest environments, they do not capture the true, instance-dependent complexity of PAC policy identification. In particular, many MDP instances can be solved with far fewer samples than the worst-case bounds would suggest. A series of works has recently attempted to characterize the true instance-dependent complexity of

(ε, δ)-PAC policy identification. The seminal work of A. J. Wagenmaker, Simchowitz, et al. 2022 notably observed that algorithms relying on the optimism principle, and, more generally, any algorithm satisfying good regret properties, cannot be instance-optimal in the (ε, δ)-PAC setting for all MDPs. The authors then introduced MOCA, a racing-type algorithm with action elimination, and showed that, with high probability, its sample complexity is upper-bounded by CMOCA (M, ε) log(1/δ) + poly(log(1/δ), log(1/ε), SAH)/ε, where CMOCA (M, ε) PH 1 scales with H 2 h=1 minρ maxs,a (ε∨∆h (s,a)) 2 w ρ (s,a) , h ρ and wh (s, a) is the probability that (s, a) is visited at step h when ρ is played during an episode. CMOCA (M, ε) is therefore identified as the leading complexity term when ε is small or when the suboptimality gaps are small. However, due to second-order terms, it may cease to be the dominating term when δ is small (A. J. Wagenmaker, Simchowitz, et al. 2022). A. Wagenmaker & Jamieson 2022 proposed PEDEL, a policy-elimination algorithm that estimates policy values using a sophisticated G-optimal design. For tabular MDPs the sample complexity of PEDEL scales with CPEDEL (M, ε) ≜ π PH P wh (s,a)2 H 4 h=1 minρ maxπ∈Πdet s,a wρ (s,a)(ε∨∆(π)) where 2 h Πdet is the set of deterministic (Markovian) policies and ∆(π) is the policy gap. See Tirinzoni, Al-Marjani, et al. 2023 for further discussion of the different gaps in (ε, δ)-PAC policy identification. By design, PEDEL requires enumerating all deterministic policies, which for finite-horizon MDPs is of size ASH . The identified complexities for PEDEL and MOCA are generally not comparable. Al-Marjani et al. 2023a introduced an algorithm whose identified complexity ranges between MOCA and PEDEL but suffers from a highly suboptimal polynomial dependence on log(1/δ). Narang et al. 2024 improved upon PEDEL by estimating not the individual policy values but the differences in policy values with respect to a reference policy. However, like PEDEL, their algorithm still requires enumerating the set of deterministic policies. Tirinzoni, Al Marjani, et al. 2022 introduced EPRL, a fully adaptive action-elimination algorithm for MDPs with deterministic transitions. The sample complexity of EPRL also scales with the sum of the inverse squared suboptimality gaps of suboptimal reachable actions. The authors proved a lower bound on the sample complexity of a (ε, δ)-PAC algorithm for deterministic MDPs, showing that EPRL is optimal up to O(H 2 ) terms. Importantly, this optimality guarantee holds only in the deterministic setting and does not extend to stochastic MDPs.

Cyrille Kone, Kevin Jamieson

Marjani & Proutiere 2021 studied the instancedependent complexity of (ε, δ)-PAC policy identification with a generative model for ε = 0 and stated the optimal complexity as a solution to an intensive min– max program, similar to the techniques popularized in Garivier & Kaufmann 2016. The authors proposed an algorithm that solves a proxy of this problem. Applying the same technique with a forward model, Al Marjani et al. 2021 proposed MDP-NaS and proved a sample complexity asymptotically (as δ → 0) equivalent to C(M ) log(1/δ), where C(M ) is worse than CMOCA (M, 0). To date, all existing algorithms for policy identification either suffer from computational intractability, exhibit a suboptimal dependence on the leading term involving log(1/δ) as δ → 0, or use regret minimization algorithms for their sampling rules. Posterior Sampling and Top-Two Methods. Posterior sampling has emerged as a powerful tool for exploration in bandits and MDPs. Algorithms such as Top-Two Thompson Sampling (Russo 2016; Jourdan et al. 2022; Z. Li et al. 2024) achieve optimal exploration in bandits by maintaining uncertainty over the true model. They appear to be statistically optimal, computationally eﬀicient, and empirically competitive with other approaches. In RL, particularly for regret minimization, posterior sampling algorithms have become popular for designing eﬀicient and tractable methods (Osband et al. 2013; Agrawal & Jia 2017; Russo 2019; Tiapkin et al. 2022). Yet, to the best of our knowledge, no posterior sampling algorithm has been analyzed for best policy identification in RL. We aim to close this gap by proposing an optimal and tractable algorithm for the episodic tabular setting. However, we emphasize that our goal is not merely to establish a sample complexity bound for an existing regret-minimization algorithm, as such algorithms are known to be suboptimal from an instance-dependent perspective (see above). Instead, we aim to design an instance-optimal algorithm that retains the computational eﬀiciency of posterior sampling. 1.2

Main Contributions

In this paper, we focus on the theoretical analysis of the high-confidence regime (small δ) and the exact identification case (ε = 0): • Algorithmic framework. We introduce PIPS, a top-two, posterior-sampling-based algorithm for PAC policy identification, which breaks the policy-enumeration bottleneck by using posteriordriven exploration coupled with an online learner. • Theoretical guarantees.

We establish optimal

instance-dependent sample complexity bounds in the high-confidence and exact identification regimes. In addition, we prove sharp posterior contraction results, providing a refined characterization of how uncertainty about the optimal policy decreases with the number of episodes. We further discuss the extension of PIPS to the ε > 0 case for which a simple modification of PIPS is proposed. • Empirical validation. We complement our theory with experiments showing that our algorithm outperforms existing baselines, both in terms of sample eﬀiciency and computational scalability.

PROBLEM FORMULATION

We consider an episodic, finite-horizon Markov decision process (MDP) with finite state space S of size S and action space A of size A. The MDP is deH fined by the tuple M ≜ (S, A, {ph }H h=1 , {Rh }h=1 , sinit ) where ph and Rh denote the transition and reward kernels at step h, and sinit is the initial state. A (Markovian) policy is a sequence π = {πh }H h=1 with πh (·|s) a distribution over actions given state s ∈ S. Under policy π, at each stage h ∈ [H] the agent draws ah ∼ πh (sh ), receives reward rh ∼ Rh (sh , ah ) with mean µh (sh , ah ) = µ(s, a, h), and transitions to sh+1 ∼ ph (·|sh , ah ) ≡ p(· | s, a, h). A policy is called deterministic if each πh (·|s) assigns all probability mass to a single action denoted πh (s) or π(s, h). Learning Problem. Given the MDP M , for any policy π = {πh }H is defined at each h=1 , the Q-function PH π π state-action as Q (s, a) ≜ E M h k=h rk |sh = s, ah = a , i.e., the expected accumulated reward when starting from (s, a) at step h and following π thereafter. The corresponding state-value function is Vhπ (s) ≜ PH π π EπM k=h rk |sh = s and we write V0 ≜ V1 (sinit ) for the expected return when starting from the initial state. Let Π denote the set of all policies (deterministic and stochastic). The learner’s goal is to identify an optimal policy π ⋆ ∈ argmaxπ∈Π V0π using as few episodes as possible. For an MDP M, we let Π⋆ (M ) (resp. Π⋆det (M )) denote the set of optimal (resp. deterministic optimal) policies, and drop the dependence on M when clear from context. Finally, we denote by {Ht }t⩾1 the filtration of the learning process, i.e., the information collected up to the end of the t-th episode. Assumption 1. The MDP M admits a unique optimal policy π ⋆ . We assume normally distributed rewards, i.e. Rh (s, a) ∼ N (µ(s, a, h), 1), and we further assume

Optimal Posterior Sampling for Policy Identification in Tabular Markov Decision Processes

without loss of generality that the expected reward µ(s, a, h) ∈ (0, 1). Any tabular MDP in this parametric set is uniquely identified by its transition probabilities p ≜ {ph }H h=1 and expected reward functions . Maintaining a distribution over these µ ≜ {µh }H h=1 parameters is therefore equivalent to maintaining a distribution over MDP models. In the sequel, when the context is clear, we identify an MDP by its transition probabilities p and expected rewards µ. We denote by M the set of such parametric MDPs. ′ Empirical MDP. For t ⩾ 1, let nt (s | s, a, h) ≜ Pt i i i ′ be the number of i=1 1 sh = s, ah = a, sh+1 = s transitions to state s′ after taking action a in state s at step h during the P first t′ episodes. We also define nt (s, a, h) ≜ | s, a, h) = s′ nt (s Pt−1 k k 1 s = s, a = a , the total number of visits to h h k=1 (s, a, h) up to the end of the t-th episode. The empirical transition probability at (s, a, h) is then given by  n (s′ | s, a, h)   t , if nt (s, a, h) > 0, nt (s, a, h) p̂t (s′ |s, a, h) ≜  1, otherwise. S We denote by µ̂t (s, a, h) the empirical mean reward at (s, a, h). The empirical reward and transition kernels (p̂t , µ̂t ) define a unique MDP in M, which we denote π by M̂t . For any policy π ∈ Π, we write (V̂t,h ) and π (Q̂t,h ) for the value and Q-functions of π in M̂t .

Prior and Posterior over MDPs. Our distribution over MDPs is defined over the parameters that uniquely determine the model. For transitions, independent Dirichlet priors are considered, which, by conjugacy, yield Dirichlet posteriors. For rewards, we assign independent uniform priors on (0, 1), leading to truncated normal posteriors. After observing an episode {(sth , ath , rht , sth+1 )}H h=1 , the posterior is updated via Bayes’ rule. Direct computation gives the closed form Y Pr(p, µ|Ht−1 ) = fs,a,h (µh (s, a))gs,a,h (ph (·|s, a)), s,a,h

(1) where fs,a,h is the density of density of the trun1 , and gs,a,h is cated normal N[0,1] µ̂t (s, a, h), nt (s,a,h) the density of Dir u+nt (· | s, a, h) , the Dirichlet with parameter +nt (· | s, a, h), for ∈ Rd+ , that we will specify latter. Our goal is not to develop a full Bayesian model; the posterior primarily serves as an algorithmic tool to guide exploration. Information-theoretic Lower Bounds. For two probability distributions P, Q supported on a dis-

crete space S, the Kullback–Leibler (KL) divergence P P (s) is KL(P k Q) ≜ s∈S P (s) log Q(s) . For two MDPs M, M ′ , we say that M is absolutely continuous with respect to M ′ (denoted M M ′ ) if for every (s, a, h), ph (·|s, a) p′h (·|s, a) and Rh (s, a) Rh′ (s, a). Given an MDP M , define Alt(M ) ≜ M ′ ∈ M : M M ′ and Π⋆ (M ) ∩ Π⋆ (M ′ ) = ∅ , the set of alternative MDPs M ′ such that M is absolutely continuous with respect to M ′ but where no policy optimal in M remains optimal in M ′ . We denote by (KL(M k M ′ ))s,a,h the KL divergence between the transition and reward kernels of M and M ′ . For each (s, a, h), introducing Ms,a,h ≜ Rh (s, a) ⊗ ph (·|s, a), ′ KL(M kM′ )s,a,h = KL Ms,a,h kMs,a,h ′ = KL Rs,a,h kRs,a,h + KL ps,a,h kp′s,a,h , where ps,a,h ≜ ph (·|s, a) and Rs,a,h ≜ Rh (s, a). An algorithm is said to be δ-PAC for best policy identification (BPI) if its stopping time τ and recommendation rule π̂τ satisfy PrM τ < ∞, π̂τ 6= π ⋆ ⩽ δ, ∀ M ∈ M . Let ΩM denote the set of visitation probabilities induced by Markovian policies on M , ΩM ≜ {{whπ (s, a)}s,a,h : π ∈ Π}. From the information-contraction principle and change-of-distribution techniques (as popularized in Garivier & Kaufmann 2016), the following holds. Theorem 1. Any δ-PAC algorithm with stopping time τ satisfies, 1 , where EM [τ ] ⩾ Γ−1 M log 2.4δ " # X f)(s,a,h) . inf wh (s, a) KL(M kM ΓM ≜ sup f∈Alt(M ) w∈ΩM M

We can express ΓM as

s,a,h

H X π exp f)s ,a ,h ΓM = max inf E KL(M kM M h h π expM f∈Alt(M ) h=1

# , (2)

where the expectation is taken over trajectories generated by policy ρ with dynamics governed by M and the KL terms act as deterministic rewards. The problem in Eq.(2) can be interpreted as the value of a twoplayer zero-sum game: the learner chooses an exploration policy π exp to maximize the accumulated information (measured by KL terms), while an adversary f (with different optimal selects an alternative MDP M policies than in M ) to minimize it. The quantity in Eq.(2) not only characterizes the sample complexity of BPI but also governs the contraction

Cyrille Kone, Kevin Jamieson

rate of the posterior distribution of any adaptive algorithm for BPI. Theorem 2. Fix an MDP M ∈ M with optimal policy π ⋆ and let νt denote the posterior distribution defined in Eq. (1). For any adaptive algorithm for best policy identification, with probability one, ⋆ f) ⩽ ΓM . lim sup − 1t log PrM / Π⋆ ( M f∼νt | Ht−1 π ∈

ft serves as an informative alternachallenger MDP M tive to the current estimate, concentrating exploration on the state-action pairs that are most relevant for distinguishing optimal from suboptimal policies. Once the challenger is obtained, an online RL learner directs exploration toward the regions where the empirical model and the challenger differ the most, as measured by their KL divergence.

t→∞

In words, for any adaptive algorithm, the posterior probability of misidentifying the optimal policy decays at most as e−tΓM asymptotically in t. While analogous results are known in the multi-armed bandit setting (Russo 2016), our result is novel in the context of tabular MDPs (see Appendix H for a full proof).

POLICY IDENTIFICATION VIA POSTERIOR SAMPLING

Our algorithm uses the posterior distribution to guide exploration in the environment. At episode t, the learner computes the optimal policy of the empirical MDP M̂t , denoted by π̂t , based on all data observed up to episode t. The core idea is then to construct a ft , sampled from an inflated postechallenger MDP M ft . rior distribution, such that π̂t is not optimal in M This challenger highlights uncertainty about the optimal policy and directs further exploration. Formally, let 1/ηt be the inflation parameter. We define the density of the inflated posterior distribution over MDPs as Y t t νtηt (p, µ) ≜ fs,a,h (µ(s, a, h)) gs,a,h (p(·|s, a, h)) , s,a,h t where fs,a,h is the density of inflated 1 t N[0,1] (µ̂t (s, a, h), ηt nt (s,a,h) ) and gs,a,h is the density of the inflated Dirichlet Dir(u + ηt nt (· | s, a, h)), and we set the prior hyperparameter to u = (1, . . . , 1). ft is generated via conditional The challenger MDP M posterior sampling: we draw candidate MDPs f f Mt,1 , Mt,2 , . . . i.i.d from νtηt and select the first one for which π̂t is not optimal. Concretely, for ft,k , we compute its optimal each sampled MDP M k,⋆ e Q-function Q t,h by backward induction (Puterman 1994) and then check whether π̂t is optimal for this Q-function. As soon as a sample is found where π̂t is suboptimal, this instance is designated as the ft . challenger MDP M

This sampling procedure naturally favors randomness in insuﬀiciently explored state-action pairs, where deviations from the empirical model are most likely to produce a different optimal policy. Consequently, the

Exploration Policy. The learner maintains an online RL reward-maximization strategy over the space of exploration policies to ensure that suﬀicient evidence is collected to distinguish between the empirical f. An initial exploMDP M̂ and the challenger MDP M ration policy π1exp is given. At the end of each episode, the exploration policy is updated via a policy improvement step performed by a mirror descent learner. By ft ≜ (θt , Φt ), we define the exploration reward letting M kernel as 2 µ̂t (s, a, h) − θt (s, a, h) + 2 ′ X p̂t (s | s, a, h) p̂t(s′ | s, a, h) log . max (Φt (s′ | s, a, h), e−tα ) s′ (3) which corresponds to the KL divergence between the ft (s, a, h) and M̂t (s, a, h) with transition marginals M α f in Mt clipped to e−t to avoid singularities and control the magnitude of the reward. An optimistic policy evaluation step estimates the Q-function of πtexp in the MDP with reward kernel rtexp and (unknown) transition kernel p (hence the use of optimism). Concretely, using Bernstein type confidence bonuses, and initializing V̄t,H+1 (·) = 0, the optimisitic Q-function is recursively evaluated as rtexp (s, a, h) ≜

Q̄t,h (s, a) ≜ rtexp (s, a, h) + p̂t,h V̄t,h+1 (s, a) + 2gt,h β p (nt (s, a, h), 1/t3 ) + 3 max (1, nt (s, a, h)) (4) s 2 Varp̂t,h V̄t,h+1 (s, a) · β p (nt (s, a, h), 1/t3 ) , max (1, nt (s, a, h)) with the value function Passociated V̄t,h+1 (s) ≜ ≜ a Q̄t,h (s, a)ρh (a|s) and gt,h maxs V̄t,h+1 (s). We introduce p̂ V̄ (s, a) ≜ t,h t,h+1 Es′ ∼p̂t,h (·|s,a) V̄t,h+1 (s′ ) , and Varp̂t,h V̄t,h+1 (s, a) = Vars′ ∼p̂t,h (·|s,a) (V̄t,h+1 (s′ )). As proven in Appendix D, for some properly tuned threshold β p , the Q-function computed in Eq.(4) is, with probability at least πtexp 1 − 1/t3 , an optimistic estimate of Qt,h (s, a), the Q-function of πtexp in (rtexp , p). Next, one step of policy improvement via an online learner on the

Optimal Posterior Sampling for Policy Identification in Tabular Markov Decision Processes

A-simplex. We maintain SH independent online learners (Ls,h )s,h on the simplex 4A (distributions over actions).

Algorithm 1: PIPS: Policy Identification via Posterior Sampling in Tabular MDPs Require: Initial exploration π1exp = {π1exp (·|s, h)}s,h ; initial forced exploration c1 ; mixing parameter γ ∈ (0, 1); clipping parameter α ∈ (0, 12 ); posterior inflation (ηt )t⩾1 ; initial MDP M̂0 ∈ M (arbitrary)

The policy improvement step updates each online learner to improve the exploration policy for each (s, h) via the estimated exploration Q-function. Notable choices include AdaHedge and other instances of mirror descent learners. This is analyzed in Appendix D. To correct for potential empirical bias in initial estimates of some state-action pairs, the exploration policy πtexp is mixed with a forced exploration policy ct , which ensures that any reachable state will be explored suﬀiciently. Given some parameter γ ∈ (0, 1), at episode t, with probability t−γ , the behaviour policy πtexp is replaced a forced exploration policy ct (explained below). This ensures broad coverage while keeping the forced exploration negligible: the total number of episodes allocated is sublinear in t, and for γ close to 1, only O(log t) episodes are allocated to forced exploration, letting the exploration mainly be led by ComputeExplorationPolicy.

for episode t = 1, 2, . . . do // Empirical best policy

π Compute π̂ t ∈ arg maxπ V̂t,0

// Computing a challenger MDP 3 4 5

6 7 8 9

The forced exploration policy performs a single iter- 10 ation update of UCBVI (Azar, Osband, et al. 2017) 11 with a reward function designed as an indicator of a randomly picked state-action pair. A triplet (s̃t , ãt , h̃t ) is selected randomly and a reward kernel is defined as f rt : (s, a, h) 7→ 1 s = s̃t , a = ãt , h = h̃t and ct+1 is defined as the greedy policy wrt an optimistic estimate of the Q-function for the reward kernel rtf . The procedure is described and analyzed in Appendix E. 12 In Appendix L, we describe a simple modification to 13 construct the challenger MDP when the goal is to iden- 14 15 tify an ε-optimal policy (with ε > 0). 16

for k = 1, 2, . . . do ft,k ∼ νtηt Sample candidate challenger M e k,⋆ Compute the optimal Q-function Q of t ft,k by backward induction M e k,⋆ if π̂t is suboptimal for Q then t ft ← M ft,k and break Set challenger M end end Compute ct as in Algorithm 4 Draw Zt ∼ Bernoulli(t−γ ) and define ∀(s, h) ∈ S × [H] ( ct (·|s, h), if Zt = 1, exp π̃t (·|s, h) ← πtexp (·|s, h), if Zt = 0, // Rollout and data collection

Set st1 ← sinit for h = 1 to H do Sample action ath ∼ π̃texp (·|sth , h) Observe reward rht and transition to sth+1 end

// KL-based bonuses for the online learner Computational Cost. The computational cost of Construct bonus kernel rtexp (s, a, h) from M̂t Algorithm 1 is dominated by the sampling and plan- 17 ft as in Eq.(3) ning steps performed at each episode. Concretely, the and M algorithm generates SAH Dirichlet samples of dimen// Behaviour policy improvement exp sion S and SAH samples from a normal distribution. 18 Update main exploration policy πt+1 ← It then verifies whether the sampled MDP and the emComputeExplorationPolicy t, πtexp , rtexp , p̂t , nt pirical MDP share the same optimal policy, which can be done via backward induction (Puterman 1994) in 19 end time O(S 2 AH). Altogether, the per-episode computa2 tional complexity is O(S AH), with a memory requirement of O(S 2 AH), consistent with standard modelbased approaches. Our first result characterizes the posterior contraction rate. If we were to sample an MDP from the poste⋆ rior the ”Bayesian” error would be PrM / f∼νt |Ht−1 (π ∈ 4 MAIN THEORETICAL RESULTS R ⋆ f Π (M )) = dνt ((p, µ)) which by Theorem 2 is Alt(M )

In this section, we present the guarantees attained by PIPS both in terms of posterior contraction rate and sample complexity upper bound.

asymptotically larger than e−tΓM for any adaptive algorithm. The result below shows that PIPS reaches this optimal rate.

Cyrille Kone, Kevin Jamieson

Algorithm 2: Exploration Policy Helper Function 1 Function ComputeExplorationPolicy(t, πtexp , rtexp , p̂t , nt ) Require: Independent online learning instances (Ls,h )s,h

mally defined as n o ft,k ) . (5) τδ ≜ inf t ⩾ 1 : ∀ 1 ⩽ k ⩽ B(t, δ),π̂t ∈ Π⋆ (M

Set V̄t,H+1 (s) ← 0 ∀ s ∈ S for episode stage h = H, H − 1, . . . , 1 do foreach (s, a) ∈ S × A do Q̄t,h (s, a) ← Eq.(4) end

Theorem 4. Let γ ∈ (0, 1), α ∈ (0, 1/2), α < γ, ς ∈ (0, 1), with ς > 2α. If B(t, δ) satisfies B(t,δ) lim supδ→0 log ⩽ 1, then, PIPS coupled with the ηt log δ1 stopping time τδ described in (5) and run with parameterss γ, α and learning rate ηt = O(t−ς ), satisfies

2 3 4 5 6

// Optimistic policy evaluation 7

V̄P t,h (s) ← exp a∈A πt (a|s, h) Q̄t,h (s, a) end

The number of resampling is calibrated to ensure the δ-correctness of this stopping rule.

1 EM τδ ⩽ Γ−1 + o log 1δ , and M log δ τδ −1 PrM lim sup ⩽ Γ =1. M 1 δ→0 log δ

∀s∈S

// Policy improvement via online Mirror descent learner

foreach (s, h) ∈ S × [H] do Feed learner Ls,h with gain Q̄t,h (s, ·) exp Update πt+1 (· | s, h) from Ls,h end

exp return policy πt+1

9 10 11

Theorem 4 provides a suﬀicient condition on B(t, δ) to guarantee optimal sample complexity, but it does not ensure δ-correctness with this threshold. The lemma below provides the correctness guarantee when δ is small. Lemma 5. Let β be a threshold such that for any δ ∈ (0, 1), the event

Theorem 3. Consider an MDP M ∈ M with a unique optimal policy π ⋆ . Let γ ∈ (0, 1), α ∈ (0, 1/2), α < γ, ς ∈ (0, 1), with ς > 2α. Run with parameterss γ, α and learning rate ηt = O(t−ς ), letting νt denote the non-inflated posterior distribution, Algorithm 1 satisfies with probability one 1 ⋆ f) = ΓM . / Π⋆ ( M lim sup − log PrM f∼νt |Ht−1 π ∈ t t→∞ Combined with Theorem 2, this result ensures that when running Algorithm 1, the posterior probability of mis-identifying the optimal policy decays exponentially fast with an optimal rate. The next result we prove is in terms of sample complexity. We show that when equipped with a valid stopping procedure, Algorithm 1 attains the optimal sample complexity prescribed by Theorem 1. Sample Complexity. PIPS can be coupled with a posterior-sampling-based stopping rule similar to the stopping rule used in Kone et al. 2025 for bandits. This stopping rule is based on the number of i.i.d samples collected from νtηt before finding a challenger MDP, i.e., by clipping the number of iterations in line 3 of Algorithm 1 by a threshold B(t, δ) to be defined. ft,1 , M ft,2 , . . . are sampled i.i.d from Recalling that M the posterior distribution νtηt , the stopping time is for-

( Eδ ≜

∀t ⩾ 1,

) nt (x) KL(M̂t k M )x ⩽ β(t, δ)

x∈X

holds with probability at least 1 − δ, where Rh (s, a) ≜ N (µ(s, a, 1), 1), R̂ht (s, a) ≜ N (µ̂t (s, a, h), 1/nt (s, a, h)) and KL(M̂t k M )s,a,h ≜ KL(R̂ht (s, a) ⊗ p̂t (·|s, a, h) k Rh (s, a) ⊗ p(·|s, a, h)). x denotes a generic s, a, h and X = S × A × [H]. For B(t, δ) = O(exp(ηt β(t, δ)) log(t/δ)) the posterior stopping time defined in Eq.(5)satisfies lim supδ→0 PrM (τδ < ∞, π̂τδ 6= π ⋆ )/δ ⩽ 1. Kaufmann & Koolen 2021 prescribes correct thresholds β(t, δ) = log(1/δ) + O(log(t)). Thus B(t, δ) as B(t,δ) defined in Lemma 5 satisfies lim supδ→0 log ⩽ 1, ηt log δ1 so that Theorem 4 applies and PIPS is asymptically optimal and correct with the value of B(t, δ) prescribed by Lemma 5.

EXPERIMENTS

In this section, we evaluate our algorithm against a range of competitors. We examine two performance metrics : (i) the correct identification rate, defined as 1 if the recommended policy is optimal and 0 otherV ⋆ −V π̂

wise, and (ii) the performance factor 1 − 0 V ⋆ . The 0 average identification rate reflects the fraction of times

Optimal Posterior Sampling for Policy Identification in Tabular Markov Decision Processes

identification rate

1.0 0.8 PIPS UCBVI PSRL SSR

0.6 0.4 0.2 0.0

1.0

0.8

0.6

0.4

0.2

0.0 0

5000

10000 episode

15000

0.0 0

5000

10000 15000 20000 25000 30000 episode

2000

4000 6000 episode

8000

10000

Figure 1: Empirical with 95% Wilson’s confidence interval. From left to 1 identification rate averaged on 64 runs 1 1 right: MOCA instance, RiverSwim, CombinationLock. the recommended policy is optimal, while the performance factor measures the value of the recommended policy relative to the optimal value. To ensure fairness, all algorithms are run without a stopping rule. Each experiment is repeated 50 times, and we report the average value of both metrics. Baselines. Algorithm 1, is compared to several established baselines. We evaluate BPI-UCBVI (Menard et al. 2021). Its sampling rule corresponds to the greedy policy with respect to an optimistic estimate of the optimal Q-function. Another baseline is PSRL (Osband et al. 2013), a randomized exploration algorithm that maintains a posterior distribution over MDPs and acts greedily wrt a sampled MDP, and SSR (Xiong et al. 2022), a variant of PSRL with distributions over Q-functions. 5.1

Environments

We evaluate the algorithms on challenging environments that require eﬀicient exploration. MOCA The first environment, adapted from A. J. Wagenmaker, Simchowitz, et al. 2022 and shown in Fig. 2, contains (L + 1) states and (L + 1) actions. Among them, a⋆ is the optimal action at each stage but transitions to the next states with small probability. In contrast, the other actions a1 , . . . , aL are suboptimal, but taking al in sinit deterministically leads to state sl . This environment poses a subtle challenge: identifying the optimal action in states s1 , . . . , sL may require first taking suboptimal actions in sinit . As a result, it is particularly diﬀicult for regret-minimizing algorithms applied to BPI, since they tend to avoid actions with poor short-term performance. We also evaluate on the classic RiverSwim MDP (see e.g. Strehl & Littman 2008), and the CombinationLock MDP (see e.g. Zhang et al. 2022). RiverSwim. In the RiverSwim environment, the optimal policy always selects the RIGHT action to reach

the high-reward state at the end of the chain, despite the lower immediate rewards along the way. We use a chain of length L = 8, set the horizon to H = 10, and take sinit = s1 . The rewards are Gaussian with variance 1/1000. The environment is described in Appendix M.2, Fig.8.

Combination Lock. We consider a unichain, episodic combination-lock MDP with S = H states s1 , . . . , sH and an action set of size A ⩾ 2. Each nonterminal state sl (l = 1, . . . , H − 1) has a unique optimal action denoted a⋆l . If the agent selects the optimal action a⋆l in state sl , it transitions to sl+1 with probability 1 − ε and returns to the initial state s1 with probability ε. Any suboptimal action a 6= a⋆l at sl sends the agent back to s1 with probability 1. The terminal state sH is absorbing, and the agent receives a reward of 1 only upon reaching sH (i.e., after taking the optimal action at every step); all other transitions yield a reward of 0. In our experiments, we set H = 10, A = 3, and ε = 0.01.

5.2

Summary

Fig. 1 plots the identification rate vs the number of episodes for the three MDPs. The most significant improvement of PIPS over its competitors is achieved in the MOCA MDP. This is expected as in this MDP, optimistic algorithms essentially explore by following the policy with the highest value, which, as described in Fig. 2, only reach some states with probability exponentially small in S. Fig. 4 shows the performance factor results. Additional experiments, presented in Appendix M.2, Fig. 5,6, 7, compare the posterior contraction rates of PIPS and PSRL on the MOCA instance, highlighting the significant advantage of PIPS over its competitors. Further implementation details are provided in Appendix M.1.

Cyrille Kone, Kevin Jamieson

Figure 2: Custom MOCA environment where reaching the optimal policy requires exploratory steps involving sub-optimal actions a1 , . . . , aL (red dotted lines).

Figure 3: The H-state CombinationLock environment. Each nonterminal state offers A actions. The unique optimal action at state sl advances the agent to sl+1 with probability 1 − ε (solid black arrow) and returns it to the initial state sinit with probability ε (dashed arrow); any suboptimal action (red arrow) sends the agent back to sinit with probability 1. The terminal state sH is absorbing and yields a reward 1 (all other transitions give a reward 0). Transition probabilities are annotated on the corresponding arrows. izing the optimal lower bound for best policy identification. Our analysis establishes optimal asymptotic guarantees on posterior contraction and sample complexity when the algorithm is equipped with a δ-PACcalibrated threshold for the posterior stopping rule.

performance factor

1.0 0.8 0.6 0.4

PIPS UCBVI PSRL SSR

0.2 0.0 0

5000

10000 episode

15000

1.0

0.8 0.6 0.4 0.2 0.0 2000

4000 6000 episode

8000

10000

Figure 4: Performance factor 1 averaged on 64 runs with 95% Wilson’s confidence interval. MOCA instance (top), and CombinationLock (bottom).

DISCUSSION AND OPEN PROBLEMS

We have studied the problem of best policy identification in finite-horizon episodic tabular MDPs. In this work, we introduced a computationally eﬀicient algorithm that leverages posterior sampling combined with an online learning algorithm to explore the MDP and iteratively solve the optimization problem character-

Despite its computational eﬀiciency, our algorithm still requires sampling O(S 2 AH) parameters per round to generate a challenger MDP, which can be a bottleneck in tabular MDPs with large state spaces. Furthermore, while its theoretical guarantees are tight in the asymptotic regime, they do not provide explicit dependence on the state, action, or horizon parameters, which may not be negligible in the nonasymptotic setting, particularly when δ = Ω(1). Obtaining tight finite-sample guarantees in the moderate confidence regime remains an open challenge as current algorithms scale with diverse and often incomparable complexity terms, including substantial ”second-order” terms, such as those polynomial in log 1/δ (A. Wagenmaker & Jamieson 2022; A. J. Wagenmaker, Simchowitz, et al. 2022; Al-Marjani et al. 2023a; Narang et al. 2024). This work can be extended in several promising directions. For instance, one could consider extensions to the linear function approximation setting (Jin, Z. Yang, et al. 2020), or to the (ε, δ)-PAC framework, where the objective is to identify an ε-optimal policy for some ε > 0. We briefly discuss an extension of PIPS to the (ε, δ)-PAC setting in Appendix L, although the resulting guarantees are sub-optimal. Extending our main results to this framework remains an open question.

Optimal Posterior Sampling for Policy Identification in Tabular Markov Decision Processes

Acknowledgements KJ was funded in part by NSF Awards 2141511, 2023239. This work was done while CK was visiting KJ at the University of Washington.

References Fiechter, C.-N. (1994). “Eﬀicient reinforcement learning”. Proceedings of the Seventh Annual Conference on Computational Learning Theory. COLT ’94. Association for Computing Machinery. Puterman, M. L. (1994). Markov Decision Processes: Discrete Stochastic Dynamic Programming. 1st. John Wiley & Sons, Inc. Kearns, M. & S. Singh (1998). “Finite-Sample Convergence Rates for Q-Learning and Indirect Algorithms”. Advances in Neural Information Processing Systems. Vol. 11. MIT Press. Strehl, A. L. & M. L. Littman (2008). “An analysis of model-based Interval Estimation for Markov Decision Processes”. Journal of Computer and System Sciences 74. Lu, D. & W. V. Li (2009). “A note on multivariate Gaussian estimates”. Journal of Mathematical Analysis and Applications. Jaksch, T., R. Ortner & P. Auer (2010). “Near-optimal Regret Bounds for Reinforcement Learning”. Journal of Machine Learning Research 11. Azar, M. G., R. Munos & H. J. Kappen (2012). “On the sample complexity of reinforcement learning with a generative model”. Proceedings of the 29th International Coference on International Conference on Machine Learning. ICML’12. Omnipress. – (2013). “Minimax PAC bounds on the sample complexity of reinforcement learning with a generative model”. Mach. Learn. 91. Osband, I., D. Russo & B. V. Roy (2013). (More) Efficient Reinforcement Learning via Posterior Sampling. De Rooij, S., T. Van Erven, P. D. Grünwald & W. M. Koolen (2014). “Follow the leader if you can, hedge if you must”. J. Mach. Learn. Res. Dann, C. & E. Brunskill (2015). “Sample complexity of episodic fixed-horizon reinforcement learning”. Proceedings of the 29th International Conference on Neural Information Processing Systems - Volume 2. NIPS’15. MIT Press. Krichene, W., M. Balandat, C. Tomlin & A. Bayen (2015). “The Hedge Algorithm on a Continuum”. Proceedings of the 32nd International Conference on Machine Learning. Vol. 37. Proceedings of Machine Learning Research. PMLR. Garivier, A. & E. Kaufmann (2016). “Optimal Best Arm Identification with Fixed Confidence”. 29th An-

nual Conference on Learning Theory. Vol. 49. Proceedings of Machine Learning Research. PMLR. Russo, D. (2016). “Simple Bayesian Algorithms for Best Arm Identification”. 29th Annual Conference on Learning Theory. PMLR. Agrawal, S. & R. Jia (2017). “Optimistic posterior sampling for reinforcement learning: worst-case regret bounds”. Advances in Neural Information Processing Systems. Vol. 30. Curran Associates, Inc. Azar, M. G., I. Osband & R. Munos (2017). “Minimax Regret Bounds for Reinforcement Learning”. Proceedings of the 34th International Conference on Machine Learning. Vol. 70. Proceedings of Machine Learning Research. PMLR. Dann, C., T. Lattimore & E. Brunskill (2017). “Unifying PAC and Regret: Uniform PAC Bounds for Episodic Reinforcement Learning”. Advances in Neural Information Processing Systems. Curran Associates, Inc. Yang, Z.-H., W. Qian, Y. Chu & W. Zhang (2017). “On Rational Bounds for the Gamma Function”. Journal of Inequalities and Applications 2017. Jin, C., Z. Allen-Zhu, S. Bubeck & M. I. Jordan (2018). “Is Q-Learning Provably Eﬀicient?” Advances in Neural Information Processing Systems. Vol. 31. Curran Associates, Inc. Sidford, A., M. Wang, X. Wu, L. Yang & Y. Ye (2018). “Near-Optimal Time and Sample Complexities for Solving Markov Decision Processes with a Generative Model”. Advances in Neural Information Processing Systems. Vol. 31. Curran Associates, Inc. Sidford, A., M. Wang, X. Wu & Y. Ye (2018). “Variance reduced value iteration and faster algorithms for solving markov decision processes”. Proceedings of the Twenty-Ninth Annual ACM-SIAM Symposium on Discrete Algorithms. SODA ’18. Society for Industrial and Applied Mathematics. Russo, D. (2019). “Worst-Case Regret Bounds for Exploration via Randomized Value Functions”. Advances in Neural Information Processing Systems. Vol. 32. Curran Associates, Inc. Agarwal, A., S. Kakade & L. F. Yang (2020). “ModelBased Reinforcement Learning with a Generative Model is Minimax Optimal”. Proceedings of Thirty Third Conference on Learning Theory. Vol. 125. Proceedings of Machine Learning Research. PMLR. Jin, C., Z. Yang, Z. Wang & M. I. Jordan (2020). “Provably eﬀicient reinforcement learning with linear function approximation”. Proceedings of Thirty Third Conference on Learning Theory. Vol. 125. Proceedings of Machine Learning Research. PMLR. Al Marjani, A., A. Garivier & A. Proutiere (2021). “Navigating to the Best Policy in Markov Decision Processes”. Advances in Neural Information Processing Systems. Vol. 34. Curran Associates, Inc.

Cyrille Kone, Kevin Jamieson

Domingues, O. D., P. Ménard, E. Kaufmann & M. Valko (2021). “Episodic Reinforcement Learning in Finite MDPs: Minimax Lower Bounds Revisited”. Proceedings of the 32nd International Conference on Algorithmic Learning Theory. Vol. 132. Proceedings of Machine Learning Research. PMLR. Kaufmann, E. & W. M. Koolen (2021). “Mixture Martingales Revisited with Applications to Sequential Tests and Confidence Intervals”. Journal of Machine Learning Research 22. Kaufmann, E., P. Ménard, O. Darwiche Domingues, A. Jonsson, E. Leurent & M. Valko (2021a). “Adaptive Reward-Free Exploration”. Proceedings of the 32nd International Conference on Algorithmic Learning Theory. Proceedings of Machine Learning Research. PMLR. – (2021b). “Adaptive Reward-Free Exploration”. Proceedings of the 32nd International Conference on Algorithmic Learning Theory. Vol. 132. Proceedings of Machine Learning Research. PMLR. Marjani, A. A. & A. Proutiere (2021). “Adaptive Sampling for Best Policy Identification in Markov Decision Processes”. Proceedings of the 38th International Conference on Machine Learning. Vol. 139. Proceedings of Machine Learning Research. PMLR. Menard, P., O. D. Domingues, A. Jonsson, E. Kaufmann, E. Leurent & M. Valko (2021). “Fast active learning for pure exploration in reinforcement learning”. Proceedings of the 38th International Conference on Machine Learning. Vol. 139. Proceedings of Machine Learning Research. PMLR. Jourdan, M., R. Degenne, D. Baudry, R. de Heide & E. Kaufmann (2022). “Top Two Algorithms Revisited”. Advances in Neural Information Processing Systems. Vol. 35. Curran Associates, Inc. Tiapkin, D., D. Belomestny, D. Calandriello, E. Moulines, R. Munos, A. Naumov, M. Rowland, M. Valko & P. Ménard (2022). “Optimistic Posterior Sampling for Reinforcement Learning with Few Samples and Tight Guarantees”. Advances in Neural Information Processing Systems. Vol. 35. Curran Associates, Inc. Tirinzoni, A., A. Al Marjani & E. Kaufmann (2022). “Near Instance-Optimal PAC Reinforcement Learning for Deterministic MDPs”. Advances in Neural Information Processing Systems. Vol. 35. Curran Associates, Inc. Wagenmaker, A. & K. G. Jamieson (2022). “InstanceDependent Near-Optimal Policy Identification in Linear MDPs via Online Experiment Design”. Advances in Neural Information Processing Systems. Vol. 35. Curran Associates, Inc. Wagenmaker, A. J., Y. Chen, M. Simchowitz, S. Du & K. Jamieson (2022). “Reward-Free RL is No Harder Than Reward-Aware RL in Linear Markov Deci-

sion Processes”. Proceedings of the 39th International Conference on Machine Learning. Proceedings of Machine Learning Research. Wagenmaker, A. J., M. Simchowitz & K. Jamieson (2022). “Beyond No Regret: Instance-Dependent PAC Reinforcement Learning”. Proceedings of Thirty Fifth Conference on Learning Theory. Vol. 178. Proceedings of Machine Learning Research. PMLR. Xiong, Z., R. Shen, Q. Cui, M. Fazel & S. S. Du (2022). “Near-Optimal Randomized Exploration for Tabular Markov Decision Processes”. Advances in Neural Information Processing Systems. Vol. 35. Curran Associates, Inc. Zhang, X., Y. Song, M. Uehara, M. Wang, A. Agarwal & W. Sun (2022). “Eﬀicient Reinforcement Learning in Block MDPs: A Model-free Representation Learning approach”. Proceedings of the 39th International Conference on Machine Learning. Vol. 162. Proceedings of Machine Learning Research. PMLR. Al-Marjani, A., A. Tirinzoni & E. Kaufmann (2023a). “Active Coverage for PAC Reinforcement Learning”. Proceedings of Thirty Sixth Conference on Learning Theory. Proceedings of Machine Learning Research. PMLR. – (2023b). Towards Instance-Optimality in Online PAC Reinforcement Learning. Tirinzoni, A., A. Al-Marjani & E. Kaufmann (2023). “Optimistic PAC Reinforcement Learning: the Instance-Dependent View”. Proceedings of The 34th International Conference on Algorithmic Learning Theory. Vol. 201. Proceedings of Machine Learning Research. PMLR. Li, Z., K. Jamieson & L. Jain (2024). “Optimal Exploration is no harder than Thompson Sampling”. Proceedings of The 27th International Conference on Artificial Intelligence and Statistics. PMLR. Narang, A., A. Wagenmaker, L. J. Ratliff & K. Jamieson (2024). “Sample Complexity Reduction via Policy Difference Estimation in Tabular Reinforcement Learning”. Advances in Neural Information Processing Systems. Vol. 37. Curran Associates, Inc. Kone, C., M. Jourdan & E. Kaufmann (2025). “Pareto Set Identification With Posterior Sampling”. Proceedings of The 28th International Conference on Artificial Intelligence and Statistics. Vol. 258. Proceedings of Machine Learning Research. PMLR.

Checklist 1. For all models and algorithms presented, check if you include: (a) A clear description of the mathematical setting, assumptions, algorithm, and/or model.

Optimal Posterior Sampling for Policy Identification in Tabular Markov Decision Processes

[Yes] All claims are proven in the Appendix. (b) An analysis of the properties and complexity (time, space, sample size) of any algorithm. [Yes] The complexity is analyzed in a paragraph in the main, and all sample complexity claims are proven in the Appendix. (c) (Optional) Anonymized source code, with specification of all dependencies, including external libraries. [Not Applicable] 2. For any theoretical claim, check if you include: (a) Statements of the full set of assumptions of all theoretical results. [Yes] Statements made in the main are proven in the Appendix (b) Complete proofs of all theoretical results. [Yes] The results are proven in the Appendix (c) Clear explanations of any assumptions. [Yes] 3. For all figures and tables that present empirical results, check if you include: (a) The code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL). [Yes] Section M describes the experimental setup and other implementation details for reproducibility. (b) All the training details (e.g., data splits, hyperparameters, how they were chosen). [Yes] Described in Section M (c) A clear definition of the specific measure or statistics and error bars (e.g., with respect to the random seed after running experiments multiple times). [Yes] Described in Section 5 (d) A description of the computing infrastructure used. (e.g., type of GPUs, internal cluster, or cloud provider). [Yes] Described in Section M 4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets, check if you include: (a) Citations of the creator If your work uses existing assets. [Not Applicable] (b) The license information of the assets, if applicable. [Not Applicable] (c) New assets either in the supplemental material or as a URL, if applicable. [Not Applicable] (d) Information about consent from data providers/curators. [Not Applicable] (e) Discussion of sensible content if applicable, e.g., personally identifiable information or offensive content. [Not Applicable]

5. If you used crowdsourcing or conducted research with human subjects, check if you include: (a) The full text of instructions given to participants and screenshots. [Not Applicable] (b) Descriptions of potential participant risks, with links to Institutional Review Board (IRB) approvals if applicable. [Not Applicable] (c) The estimated hourly wage paid to participants and the total amount spent on participant compensation. [Not Applicable]

Contents 1 INTRODUCTION

1.1

Related Work . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

1.2

Main Contributions . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

2 PROBLEM FORMULATION

3 POLICY IDENTIFICATION VIA POSTERIOR SAMPLING

4 MAIN THEORETICAL RESULTS

5 EXPERIMENTS

5.1

Environments . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

5.2

Summary . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

6 DISCUSSION AND OPEN PROBLEMS

A EXTENDED NOTATION TABLE

B SAMPLE COMPLEXITY LOWER BOUND

C POSTERIOR STOPPING RULE

D OPTIMAL EXPLORATION VIA ONLINE LEARNING

D.1 Setup and Goal . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

D.2 Algorithm . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

D.3 Regret Upper Bound . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

E SUFFICIENT EXPLORATION AND MODEL ESTIMATION

E.1 Online Regret of UCBVI with Changing Rewards . . . . . . . . . . . . . . . . . . . . . . . . . . .

E.1.1 Algorithm . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

E.1.2 Regret Upper Bound . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

E.2 Suffficient Exploration of Reachable States . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

E.3 Guarantees on Model Estimation . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

F ANALYSING CONDITIONAL POSTERIOR SAMPLING

F.1 Setup . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

F.2 Posterior Density Ratio and Stochastic Hedge Loss . . . . . . . . . . . . . . . . . . . . . . . . . .

F.3 Hedge with Decreasing Learning Rate . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

F.4 Best Model in Hindsight . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

F.5 Technical and Concentration Lemmas . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

Optimal Posterior Sampling for Policy Identification in Tabular Markov Decision Processes

F.6 Proof of Theorem 15 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . G SADDLE-POINT CONVERGENCE G.1 Concentration Lemmas . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . H POSTERIOR CONVERGENCE

40 40 44 50

H.1 Asymptotic Limit of Likelihood Ratio . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

H.2 Proof of Theorem 2 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

SAMPLE COMPLEXITY GUARANTEES

J CONCENTRATION RESULTS

J.1

Self-noramlized Concentration . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

J.2

Gaussian and Dirichlet Concentration . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

K TECHNICAL RESULTS

L BEYOND EXACT POLICY IDENTIFICATION

M IMPLEMENTATION DETAILS AND ADDITIONAL EXPERIMENTS

M.1 Implementation Details . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

M.2 Additional Experiments . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

Cyrille Kone, Kevin Jamieson

EXTENDED NOTATION TABLE

We recall the notation used in the main and introduce some notation used in the appendix. In what follows, M is always fixed, and its optimal policy is also fixed. In general, ˆ· denotes empirical quantities, whereas e· denotes quantities and variables related to the posterior distribution. Table 1: Additional notations used in the main and in the appendix Symbol

Definition

M M̂t π⋆ 4 M Mε R P Aπ Aπ (q) Πdet (ηt )t πtexp ct nt (s, a, t) nt (s′ | s, a, h) µ̂t (s, a, h) Π⋆ (M ) Π⋆det (M ) Alt(M ) ν t (·| Alt(M̂t )) ηt t t t t H (sh , ah , rh , sh+1 ) h=1 p̂t (· | s, a, h) ≜ p̂t (· | s, a, h) Wh (s, a)

(S, A, H, P, µ, sinit ) (S, A, H, p̂t , µ̂t , sinit ) Unique (deterministic) optimal policy of M S-Probability simplex Set of parameters of MDPs : {(0, 1) × 4}SAH {(r, q) ∈ M : ∀(s, a, h, s′ ) , q(s′ | s, a, h) ⩾ ε} Expected reward kernel space: (0, 1)SAH Transition kernel space: 4SAH ⋆ (r, q) ∈ M : V0π (r, q) ⩾ V0π (r, q) {q : ∃r ∈ R such that (r, q) ∈ Aπ } Deterministic Markovian policies Learning rate or posterior inflation rate Behavioral or exploration policy at episode t Forced exploration or coverage policy at episode t Number of pulls of (s, a) up to the end of episode t Number of transitions (s, a, s′ ) at stage h Empirical reward at (s, a, h) based upon the first t episodes Optimal start-state policies of M Optimal deterministicpolicies of M Alternative models : M ′ ∈ M : M M ′ and Π⋆ (M ) ∩ Π⋆ (M ′ ) = ∅ Truncated and inflated posterior at the start episode t Trajectory collected during episode t Empirical transition probabilities at (s, a, h) during episode t Maximum visitation probability of (s, a, h): maxρ whρ (s, a)

Value gaps. We recall that a policy π is globally optimal if its support at state and stage is included in the argmax set of the optimal Q-function. For a policy π, we define the gap at (s, a, h) as ∆πh (s, a) = Vhπ (s) − Qπh (s, a). Note that, if π is globally optimal then the above gap is called the value gap and quantifies the suboptimality of action a in state s at stage h relative to the optimal action. We say that an MDP admits a strongly globally policy if, for every (s, h), the argmax set of the optimal Q-function is a singleton : | argmaxa∈A Q⋆h (s, a)| = 1 i.e., there is a unique optimal action at each state and stage. If in addition, the start-state optimal policy is unique, then it must coincide with the greedy policy induced by Q⋆ . In this case, we say the MDP admits a unique start-state optimal policy, which is also strongly optimal. Finally, writing a⋆ ≜ argmaxa Q⋆h (s, a), we define at (s, h), ∆πh (s, a⋆ ) = mina̸=a⋆ ∆h (s, a) and we define the minimal value gap as ∆πmin ≜ min ∆πh (s, a) , s,a,h

(A.1)

which is postive when π is the strongly globally optimal policy in M . When the MDP is not clear from the context, we write ∆πmin (M ).

Optimal Posterior Sampling for Policy Identification in Tabular Markov Decision Processes

Additional notation. Given a (Markov) policy ρ, we denote by wρ the visitation probabilities induced on the MDP, i.e., for any s, a, h, whρ (s, a) is the probability that the state-action (s, a) is visited at episode h when ρ is followed. The reward space is R and the transitions space is P. We have M = R ⊗ P. Since our results are targeting the first-order asymptotic bound, we let T0 denote an arbitrary constant that may appear multiple times in our calculations. O(f (t)) denotes an asymptotic upper bound when t → ∞, hiding constants, and g(t) = o(f (t)) if and only if g(t)/f (t) → 0 when t → ∞. Finally, we define the following ℓ1 -distance on M, X kM − M ′ k1 = µ(s, a, h) − µ′ (s, a, h) + kp(·|s, a, h) − p′ (·|s, a, h)k1 ,

∀ M, M ′ ∈ M .

s,a,h

SAMPLE COMPLEXITY LOWER BOUND

We now state a lower bound on the sample complexity of any δ-PAC algorithm for best policy identification, derived from the celebrated information-contraction principle. Let

Alt(M ) ≜ M ′ ∈ M : M M ′

and

Π⋆ (M ) ∩ Π⋆ (M ′ ) = ∅ ,

that is, the set of MDPs M ′ , so that M is absolutely continuous with respect to M ′ but have no optimal policy in common with M . At episode t, an adaptive algorithm selects an exploration policy ρt based on past observations. Executing ρt in the environment generates the trajectory {sth , ath , rht , sth+1 }H h=1 ,

st1 = sinit ,

ath | sth ∼ πtexp (· | sth , h),

with transitions sth+1 | sth , ath ∼ p(· | sth , ath , h) and rewards rht | sth , ath ∼ Rh (sth , ath ). Let Ht denote the σ-algebra generated by all observations and exploration policies up to the end of episode t, including external randomness (U1 , . . . , Ut ) used for tie-breaking and randomization. For two MDPs M = (S, A, H, p, µ, sinit ),

f = (S, A, H, p̃, µ̃, sinit ), M

we define the KL divergence between their kernels at (s, a, h) as f)s,a,h ≜ KL Rh (s, a) ⊗ p(· | s, a, h) R eh (s, a) ⊗ pe(· | s, a, h) . KL(M k M An algorithm is said to be δ-PAC if, for any M ∈ M, its stopping time τ and recommendation π̂ τ satisfy PrM τ < ∞, π̂τ ∈ / Π⋆ (M ) ⩽ δ. Using the data-processing inequality and change-of-distribution arguments (see Garivier & Kaufmann 2016), we obtain the following fundamental bound: Lemma 6. Consider an MDP M with a unique optimal policy. Then, for any δ-PAC algorithm with stopping time τ , it holds that X f ∈ Alt(M ). f)s,a,h ⩾ log 1 , ∀M EM [nτ +1 (s, a, h)] KL(M k M 2.4δ s,a,h

Using this result, we prove the sample complexity lower bound. We recall that ΩM denotes the set of visitation probabilities induced by Markovian policies on M ΩM ≜

{wh (s, a)}s,a,h : ∃ π ∈ Π such that wh (s, a) = PrπM (sh = s, ah = a), ∀(s, a, h) .

(B.1)

Cyrille Kone, Kevin Jamieson

Theorem 1. Any δ-PAC algorithm with stopping time τ satisfies, 1 EM [τ ] ⩾ Γ−1 , where M log 2.4δ " # X f ΓM ≜ sup inf wh (s, a) KL(M kM )(s,a,h) . f∈Alt(M ) w∈ΩM M

s,a,h

Proof. Fix M with a unique optimal policy and let τ be the stopping time of any δ-PAC algorithm. By Lemma 6, f ∈ Alt(M ), for every M X f)s,a,h ⩾ log 1 . EM [nτ +1 (s, a, h)] KL(M k M (B.2) 2.4δ s,a,h

We recall that ρt is predictable (i.e., measurable wrt Ht−1 ). Therefore   X EM [nτ +1 (s, a, h)] = EM  1(τ ⩾ t) · 1 sth = s, ath = a  t⩾1

  =

EM E

1(τ ⩾ t) · 1 sth = s, ath = a

 Ht−1 

t⩾1

which follows from the tower rule of expectations. Next observe that the event 1(t ⩽ τ ) is the complementary of the event 1(τ < t) which is the same as 1(τ ⩽ t − 1) so it is measurable wrt Ht−1 . Thus E 1(τ ⩾ t) 1 sth = s, ath = a Ht−1 = 1(τ ⩾ t) · E 1 sth = s, ath = a Ht−1 =

π exp

1(τ ⩾ t) · wht (s, a).

Combining the above displays, EM [nτ +1 (s, a, h)] = EM

" τ X

# π exp wht (s, a)

t=1

then dividing both sides by EM [τ ] yields hP i " # πtexp τ τ EM (s, a) t=1 wh 1 X πtexp = EM w (s, a) . wh (s, a) ≜ EM [τ ] EM [τ ] t=1 h Next, we remark that using Caratheodory’s theorem, A. J. Wagenmaker, Simchowitz, et al. 2022 justifies in the proof of their Lemma F.4 that {wh (s, a)}s,a,h as defined above belongs to ΩM . Rearranging gives EM [τ ]

X s,a,h

f)s,a,h ⩾ log 1 , ∀ M f ∈ Alt(M ), wh (s, a) KL(M kM 2.4δ

(B.3)

the result follows by taking on the left-hand side the inf over Alt(M ) followed by a sup over ΩM .

C POSTERIOR STOPPING RULE We discuss the proof of the stopping rule. We prove that under mild conditions, we can prove that a simple posterior stopping algorithm can be used to calibrate an eﬀicient stopping rule. The stopping rule for this setting relies on the posterior distribution. Letting the empirical MDP M̂t , define the π optimal policy in the empirical MDP as π̂t , and the value function for policy π as V̂t,0 .

Optimal Posterior Sampling for Policy Identification in Tabular Markov Decision Processes

ft,2 , M ft,2 , . . . , be i.i.d samples from the posteiror Letting B(t, δ) be a function to be defined, and letting M ⋆ distribution over the MDPs, we denote by Π (M ) the set of optimal policies of M . We introduce the following stopping time : n o ft,k ) , τ ≜ inf t ⩾ 1 : ∀k ⩽ B(t, δ), π̂t ∈ Π⋆ (M (C.1) and after stopping, we recommend the empirical optimal planning π̂ t . We prove the following lemma. Let us introduce an event Eδ that holds with probability at least (1 − δ) for any δ ∈ (0, 1). Let Altπ be the set of MDPs for which π is not an optimal policy. Lemma 7. We have ∞ h i X f t PrM(τ < ∞, π̂τ ∈ / Π⋆ (M )) ⩽ δ/2 + EM 1 Eδ/2 , π̂t ∈ / Π⋆ (M ) exp −PrM M ∈ Alt B(t, δ) π̂ f∼νt |Ht−1 t=1

Proof. We should expect to show that this stopping rule is δ-PAC for a specific choice of B and possibly an inflation of the posterior. To start, we assume that each MDP admits a unique optimal policy and the posterior is defined over such a space, then observe that n o fτ,k ) PrM (τ < ∞, π̂τ ∈ / Π⋆ (M )) ⩽ PrM τ < ∞, π̂τ ∈ / Π⋆ (M ), ∀ k ⩽ B(τ, δ), π̂τ ∈ Π⋆ (M ⩽

δ/2 +

∞ X

h ii h ft,k ) Ht−1 . EM 1 Eδ/2 , π̂t ∈ / Π⋆ (M ) EM 1 ∀ k ⩽ B(t, δ), π̂t ∈ Π⋆ (M

t=1

Next, we have to control

h i ft,k ) | Ht−1 , Lt,δ ≜ 1 Eδ/2 , π̂t ∈ / Π⋆M EM 1 ∀ k ⩽ B(t, δ), π̂t ∈ Π⋆ (M

for which a simple rewriting yields Lt,δ

= = ⩽ ⩽

B(t,δ) ⋆ f 1 Eδ/2 , π̂t ∈ / Π⋆ (M ) PrM f∼νt |Ht−1 π̂t ∈ Π (M ) B(t,δ) f) 1 Eδ/2 , π̂t ∈ / Π⋆ (M ) 1 − PrM / Π⋆ ( M f∼νt |Ht−1 π̂t ∈ o n f) B(t, δ) / Π⋆ ( M 1 Eδ/2 , π̂t ∈ / Π⋆ (M ) exp −PrM f∼νt |Ht−1 π̂t ∈ o n f . 1 Eδ/2 , π̂t ∈ / Π⋆ (M ) exp −PrM f∼νt |Ht−1 M ∈ Altπ̂t B(t, δ)

Thus, calibrating the stopping rule properly requirestuning B(t,δ) such that the sum in the RHS of Lemma 7 is f ∈ Altπ̂ involves evaluating high-dimensional integrals M smaller than δ/2. Preceisly computing Pr f M ∼νt |Ht−1

over the complex domain Altπ̂t .

OPTIMAL EXPLORATION VIA ONLINE LEARNING

In this section, we analyze a no-regret RL algorithm for online learning in the setting where rewards change over time and their range is not fixed in advance. We show that such an algorithm generates a sequence of exploration policies that accumulates information optimally. Our main result is a generic regret bound for online RL with changing rewards coupled with a mirror descent update, which is of independent interest beyond the present application. D.1

Setup and Goal

We analyze a generic version of Algorithm 2 where the policy improvement step uses a generic online learner on the action set A. The KL-based reward kernel fed to the online learner at episode t is 2 X θt (s, a, h) − µ̂t (s, a, h) p̂t−1 (s′ | s, a, h) exp , (D.1) rt (s, a, h) ≜ + p̂t (s′ | s, a, h) log 2 max (Φt (s′ | s, a, h), e−tα ) ′ s

Cyrille Kone, Kevin Jamieson

where the transition probability clipping ensures rtexp (s, a, h) ∈ (0, 12 + tα ). To ensure optimal data collection, the exploration algorithm must ultimately guarantee, with high probability, the sequence of exploration policies (πtexp )t⩾1 satisfies

sup

T X X

w∈ΩM t=1

wh (s, a) rtexp (s, a, h) −

T X X

π exp

wht (s, a) rtexp (s, a, h) ⩽ o(T ).

(D.2)

t=1 s,a,h

s,a,h

An online RL learner with changing rewards has been analyzed in Al-Marjani et al. 2023a, but under two restrictive assumptions: the horizon T is known in advance, and the reward range is fixed in (0, 1). In our setting, the reward range (0, 21 + tα ) grows with t, requiring a more careful analysis. D.2

Algorithm

We present a general online RL algorithm, whose pseudo-code is given in Algorithm 1. It maintains SH independent online learners (Ls,h )s,h on the simplex 4A , one per state-stage pair, combined with an optimistic policy evaluation step to compute the loss fed to each learner. The key features are: • Unknown dynamics. The transition kernel p is unknown and fixed; the agent uses the empirical estimate p̂t together with Bernstein-type bonuses to ensure optimism. • Changing rewards. The reward kernel rtexp is deterministic and revealed to the agent at the end of each episode, but changes with t. • Independent per-state policy learning. The exploration policy πtexp (· | s, h) is updated independently for each (s, h) via the corresponding learner Ls,h , initialized at the uniform distribution over A.

exp Optimistic policy evaluation. At episode t, the agent collects the trajectory {(sth , ath , sth+1 )}H h=1 using πt , exp exp after which the reward kernel rt is revealed. The optimistic Q-function for πt is computed by backward induction using the Bernstein-type bonus

v h exp i u p 3 u 2 Varp̂ V̄ πt t,h t,h+1 (s, a) · β (nt (s, a, h), 1/t ) t 2β p (nt (s, a, h), 1/t3 ) ght + , bonust (s, a, h) ≜ max (1, nt (s, a, h)) 3 max (1, nt (s, a, h))

(D.3)

where gt,h ≜ maxs V̄t,h+1 (s). The optimistic Q-function and its associated value function are defined recursively as π exp

t Q̄t,h (s, a) ≜ rtexp (s, a, h) + p̂t,h V̄t,h+1 (s, a) + bonust (s, a, h),

V̄t,h (s) ≜

πtexp (a | s, h) Q̄t,h (s, a),

(D.4)

a∈A

with V̄t,H+1 ≡ 0. In the policy improvement phase, each learner Ls,h is updated with Q̄t,h (s, ·) and the next exp is computed from the updated learner. policy πt+1 Choice of online learner. We use Bernstein-type bonuses (D.3) together with AdaHedge (De Rooij et al. 2014) as the per-state online learner. AdaHedge is parameter-free, adapts its learning rate on the fly, and guarantees sublinear regret even for unbounded losses, which is essential here since the range of ιt grows with t. Hedge with a scaled learning rate is also valid but yields less tight finite-time guarantees.

Optimal Posterior Sampling for Policy Identification in Tabular Markov Decision Processes

Algorithm 3: Policy Update Helper Function Function ComputeExplorationPolicy t, πtexp , rtexp , p̂t , nt 2 Set V̄t,H+1 (s) ≡ 0 for all s ∈ S 3 for episode stage h = H, H − 1, . . . , 1 do 4 foreach (s, a) ∈ S × A do s

Q̄t,h (s, a) ←

rtexp (s, a, h)

+ p̂t,h V̄t,h+1 (s, a) +

2 Varp̂t,h V̄t,h+1 (s, a) · β p (nt (s, a, h), 1/t3 ) + max (1, nt (s, a, h))

2 β p (nt (s, a, h), 1/t3 ) 3 max (1, nt (s, a, h)) // Optimistic P policy evaluation

V̄t,h (s) ←

6 7

exp a∈A πt (a | s, h) Q̄t,h (s, a)

∀s ∈ S

foreach (s, h) ∈ S × [H] do // Policy improvement

Feed learner Ls,h with gain Q̄t,h (s, ·) exp (· | s, h) from Ls,h Update πt+1

8 9 10

exp return policy πt+1

D.3

Regret Upper Bound

We present the analysis below. We define the events : n o p (T ) ≜ ∀ t ⩽ T , ∀ (s, a, h) , nt (s, a, h) KL(p̂t (·|s, a, h) k p(·|s, a, h)) ⩽ β p (nt (s, a, h), 1/T 2 ) , ED n o 1 cnt ED (T ) ≜ ∀ t ⩽ T , ∀ (s, a, h) , nt (s, a, h) ⩾ n̄t (s, a, h) − log(2SAHT 2 ) , 2 p cnt ED (T ) ≜ ED (T ) ∩ ED (T ) .

(D.5) (D.6) (D.7)

ft ≜ (θt , Φt ) is sampled from the inflated posterior, and M̂t = (µ̂t , p̂t ) denotes At episode t, the challenger MDP M the empirical MDP. The reward kernel rtexp is defined in (D.1). We prove the following result. Proposition 8. Let (Ls,h )s,h be some online learners on 4 such that for any sequence of losses, the learner guarantees some regret RegretL (T ). Then Algorithm 1 guarantees the following regret bound on the event ED , 1

Regret(T ) ⩽ O(T α+ 2 β p (T, 1/T 3 )) + HRegretL (T ) . π Proof. Below, we let Qπt,h and Vt,h denote the state–action and state value functions in the MDP Mtexp ≜ (rtexp , p). π Q̄πt,h , V̄t,h denote the state–action and state value functions in the ”optimistic MDP” M̄texp ≜ (rtexp + bonust , p). For any fixed comparator ρ,

T X t=1

π exp ρ − Vt,0t = Vt,0

T X

" π exp π exp π exp ρ ρ ρ − Vt,0t + V̄t,0t − V̄t,0t + V̄t,0 − V̄t,0 Vt,0

(D.8)

t=1

√ ρ ρ holds (Lemma 9), justifying that on ⩾ Vt,0 On the event ED (T ), for t ∈ ( T , T ), the optimism property V̄t,0 ED (T ) T X t=1

# ρ ρ Vt,0 − V̄t,0

⩽ O(T α+ 2 ).

Cyrille Kone, Kevin Jamieson

It remains to bound the remaining terms appearing in Eq. (D.8). Let us define " # " # T T X X πtexp πtexp πtexp ρ (A) = V̄t,0 − V̄t,0 (B) = V̄t,0 − Vt,0 t=1

t=1

Bounding (A). By the performance difference lemma, "H X πexp t,πtexp ρ ρ t V̄t,0 − V̄0 =E Q̄t,h (sh , ·), ρ(· | sh , h) − πtexp (· | sh , h)

# (D.9)

s1 = sinit .

h=1

The inner product in (D.9) is exactly the loss fed to learner Lsh ,h at round t against comparator ρ(· | sh , h) = ρh (· | sh ). By the regret guarantee of each learner, T X

π exp

t Q̄t,h (s, ·), ρ(· | s, h) − πtexp (· | s, h) ⩽ RegretL (T ),

∀(s, h) ∈ S × [H].

t=1

Combining with (D.9) and summing over h, (A) ⩽ H RegretL (T ).

(D.10)

Bounding (B). Term (B) controls the total size of the Bernstein bonuses. On ED (T ), by induction, X πexp π exp π exp V̄t,0t − Vt,0t ⩽ 2 wht (s, a) bonust (s, a, h), s,a,h

where, using Varp̂t,h (V̄t,h+1 )[s, a] ⩽ (ght )2 , "s bonust (s, a, h) ⩽ gt,h

# 2β p (nt (s, a, h), 1/t3 ) 2β p (nt (s, a, h), 1/t3 ) + . max (1, nt (s, a, h)) 3 max (1, nt (s, a, h))

(D.11)

We split the sum over (s, a, h) into two parts using the set Zt ≜ (s, a, h) : n̄t (s, a, h) ⩾ 4 log(2SAHT 2 ) ,

n̄t (s, a, h) ≜

t−1 h X

i π exp (1 − εk )whk (s, a) + εk whck (s, a) .

k=1

On ED (T ), for (s, a, h) ∈ Zt : nt (s, a, h) ⩾ 14 n̄t (s, a, h). Moreover, since εt = t−γ , for t ⩾ 21/γ , 1 − εt+1 ⩾ 12 . Using maxt⩽T,h [gt,h ] ⩽ T α and monotonicity of β p in its second argument, T X

T X (nt (s, a, h), 1/t3 ) ⩽ T α β p (T, 1/T 3 ) max (1, nt (s, a, h)) t=1

gt,h β π exp wht (s, a)

t=1 (s,a,h)∈Zt

X (s,a,h)∈Zt

π exp

wht (s, a) . max 1, 14 n̄t (s, a, h)

Applying the elliptic potential (Lemma 46) and Cauchy–Schwarz, T X

t=1 (s,a,h)∈Zt

πtexp 1 X 4 (1 − εt )wh ⩽4 log 1 max 1, 4 n̄t (s, a, h) s,a,h

and similarly, for any γ ∈ (0, 1)



⩽O

√

T +1 X

! π exp (1 − εt )wht (s, a)

t=1





πtexp

T X  X wh (s, a) r p ⩽ O   max (1, n (s, a, h)) t t=1 (s,a,h)∈Z t=1 (s,a,h)∈Z

T X

1 + 14

T +1 .

πtexp

 (1 − εt )wh (s, a)  P  exp πk t−1 max 1, k=1 (1 − εk )wh (s, a)

(D.12)

Optimal Posterior Sampling for Policy Identification in Tabular Markov Decision Processes

Combining these bounds,

1 Term (B) Z ⩽ O β p (T, 1/T 3 ) T α+ 2

(D.13)

For terms outside Zt , we note that for (s, a, h) ∈ / Zt , n̄t (s, a, h) < 4 log(2SAHT 2 ), so the total visits to such pairs is bounded. We have, T X

π exp

wht (s, a) bonust (s, a, h) ⩽

t=1 (s,a,h)∈Z / t

max bonust (s, a, h)

s,a,h

X T

t⩽T

π exp wht (s, a)1

t=1

max bonust (s, a, h) O

t⩽T,s,a,h

t−1 X

π exp (1 − εk )whk (s, a) < 4 log(SAHT 2 )

! ⩽

k=1

T X

π exp (1 − εt )wht (s, a)1

t=1

t−1 X

!! π exp (1 − εk )whk (s, a) < 4 log(SAHT 2 )

⩽

k=1

max bonust (s, a, h) · O(log T )

t⩽T,s,a,h

where the latter inequality holds for γ ∈ (0, 1), and we recall that εt ≜ t−γ . Using

max bonust (s, a, h) ⩽ T α

t⩽T,s,a,h

we have

T X

π exp

p 2β p (T, 1/T 3 ) 2β p (T, 1/T 3 ) + , 3

wht (s, a) bonust (s, a, h) ⩽ O(T α β p (T, 1/T 3 ) log T ).

(D.14)

t=1 (s,a,h)∈Z / t

√ Combining (D.10), (D.13), and (D.14), and noting that β p is logarithmic in T so all bonus terms are O( T ), 1 Regret(T ) ⩽ H RegretL (T ) + O β p (T, 1/T 3 ) T α+ 2 , on the event ED (T ). √ Lemma 9. On the event ED (T ), for any t ∈ ( T , T ) and s, a, h, and any policy ρ, we have Q̄ρt,h (s, a) ⩾ Qρt,h (s, a). √ Proof. We justify that for t ∈ ( T , T ), the bonus used in the optimistic Q-function is suﬀicient to ensure optimism on the event ED (T ). We proceed by induction; assuming Q̄ρt,h+1 (s, a) ⩾ Qρt,h+1 (s, a) holds for all s, a and some h fixed. Note that this ρ ρ directly implies that V̄t,h+1 (s) ⩾ Vt,h+1 (s) and we have s ρ 2 Varp̂t,h V t,h+1 (s, a) · β p (nt (s, a, h), 1/t3 ) 2β p (nt (s, a, h), 1/t3 )gt,h ρ exp ρ + Q̄t,h (s, a) = rt (s, a, h) + p̂t,h V̄t,h+1 (s, a) + max (1, nt (s, a, h)) 3 max (1, nt (s, a, h)) s (a) 2 Varp̂t,h V t,h+1 (s, a) · β p (nt (s, a, h), 1/T 2 ) 2β p (nt (s, a, h), 1/T 2 )gt,h ⩾ rtexp (s, a, h) + p̂t,h V̄t,h+1 (s, a) + + max (1, nt (s, a, h)) 3 max (1, nt (s, a, h)) (b)

⩾

⋆ rtexp (s, a, h) + ph Vh+1 (s, a)

(induction + monotonicity, see e.g. Menard et al. 2021) √ where the last inequality holds on the event ED (T ) and we recall that for t ∈ ( T , T ), we have t3 ⩾ T 2 , combining with the fact that δ 7→ β(x, δ) is decreasing. This concludes the induction step, and as the induction hypothesis trivially holds for H + 1, we have proved the claimed result.

Cyrille Kone, Kevin Jamieson

We now state the main result of this section, which is the regret of the exploration algorithm. Proposition 10. Running Algorithm 2 on an MDP M with the exploration function ComputeExplorationPolicy using AdaHedge as a subroutine generates a sequence of policies (πtexp )t⩾1 satisfying on the event ED (T ) XX X X πexp 1 wh (s, a)rtexp (s, a, h) − wht (s, a)rtexp (s, a, h) ⩽ O(β p (T, 1/T 3 )T α+ 2 ) . sup w∈ΩM

t⩽T s,a,h

Proof. This is a direct consequence of Proposition 8. Indeed, using AdaHedge as subroutine ensures that after T iterations, we have (cf De Rooij et al. 2014, ) p RegretL (T ) ⩽ σT T log S + 16σT (2 + log(S)/3), π exp

t where σT ≜ maxt⩽T,s,h,a Q̄t,h (s, a). We note that ,

σT ⩽ O(T α β p (T, 1/T 3 )), and the conclusion follows thanks to Proposition 8.

SUFFICIENT EXPLORATION AND MODEL ESTIMATION

In this section, we show that every reachable state–action pair is visited suﬀiciently often to guarantee accurate estimation of the unknown model parameters. Our approach achieves this by leveraging a regret minimization algorithm operating under changing reward functions. The use of regret minimization for exploration in MDPs has been widely studied. However, existing analyses typically rely on additional structural assumptions, such as the MDP being communicating or ergodic (A. Wagenmaker & Jamieson 2022; Al-Marjani et al. 2023a), the availability of lower bounds on visitation probabilities (A. J. Wagenmaker, Simchowitz, et al. 2022), or the use of periodic restarts and re-initialization of the learning algorithm (A. J. Wagenmaker, Simchowitz, et al. 2022). In contrast, our analysis does not require any of these assumptions. We define the following high-probability events: n o EEp (T ) ≜ ∀t ⩽ T, ∀(s, a, h) : nt (s, a, h) KL p̂t (· | s, a, h) k p(· | s, a, h) ⩽ β p (nt (s, a, h), 1/T 2 ) , n o EEµ (T ) ≜ ∀t ⩽ T, ∀(s, a, h) : nt (s, a, h) KL R̂ht (s, a) k Rh (s, a) ⩽ β µ (nt (s, a, h), 1/T 2 ) , n o EEcnt (T ) ≜ ∀t ⩽ T, ∀(s, a, h) : nt (s, a, h) ⩾ 12 n̄t (s, a, h) − log(2SAHT 2 ) . Here, n̄t (s, a, h) ≜

t−1 X

π exp (1 − εk ) whk (s, a) + εk whck (s, a) .

(E.1) (E.2) (E.3)

(E.4)

k=1

We recall that Zt ∼ Bernoulli(t−γ ) and π̃texp (·|s, h)

←

(

ct (·|s, h), πtexp (·|s, h),

if Zt = 1, if Zt = 0,

∀(s, h) ∈ S × [H].

We introduce the event ( EEZ (T ) ≜

∀(s, a, h), ∀t ⩽ T,

t X k=1

εk · 1 s̃k = s, ãk = a, h̃k = h ⩾

) t 1 X εk − log(2SAHT 2 ) . 2SAH

(E.5)

k=1

We further define EE (T ) ≜ EEp (T ) ∩ EEµ (T ) ∩ EEcnt (T ) ∩ EEZ (T )

(E.6)

Optimal Posterior Sampling for Policy Identification in Tabular Markov Decision Processes

E.1

Online Regret of UCBVI with Changing Rewards

We recall that a state–action pair (s, a) is said to be reachable at stage h if maxρ whρ (s, a) > 0, i.e., if there exists a policy that visits (s, a) at stage h with positive probability. We now show that, on event EE (T ), every reachable state–action pair is visited suﬀiciently often to ensure the consistency of the empirical estimates used by the algorithm. To enforce exploration, we introduce a forced exploration mechanism based on a single step of UCBVI with a known reward function defined as the indicator of a randomly selected state–action–stage triplet. Formally, at episode t, a triplet (s̃t , ãt , h̃t ) is sampled uniformly at random, and we define the reward kernel rtf as rtf (s, a, h) ≜ 1 s = s̃t , a = ãt , h = h̃t .

(E.7)

Let Mtf ≜ (S, A, H, p, rtf , sinit ) denote the corresponding MDP. For any policy π, the value at the initial state satisfies π Vt,0 = wh̃πt (s̃t , ãt ),

that is, the visitation probability of (s̃t , ãt ) at stage h̃t under policy π. Consequently, an optimal policy for Mtf maximizes the visitation probability of the sampled triplet (s̃t , ãt , h̃t ). Finally, we note that maximizing whπ (s, a) is equivalent to maximizing whπ (s), since the action chosen at stage h does not affect the probability of reaching state s at that stage. Therefore, Wh (s, a) ≜ max whπ (s, a) = max whπ (s). π

E.1.1

Algorithm

The forced exploration policy is computed via a single optimistic planning step using the empirical transition kernel p̂t . The goal is to construct an optimistic estimate of the optimal Q-function in the MDP Mtf ≜ (rtf , p). In this section, let Q̄t and V̄t denote the resulting optimistic state–action and value functions. For notational simplicity, we omit the explicit dependence on the reward kernel rtf . We define the exploration bonus s bonust (s, a, h) ≜

2 Varp̂t,h V̄t,h+1 (s, a) β p nt (s, a, h), 1/t3 2 β p nt (s, a, h), 1/t3 + . max (1, nt (s, a, h)) 3 max (1, nt (s, a, h))

(E.8)

We then construct the optimistic Q-function as Q̄t,h (s, a) ≜ rtf (s, a, h) + p̂t,h V̄t,h+1 (s, a) + bonust (s, a, h) ∧ 1 .

(E.9)

The clipping by 1 ensures consistency with the reward range. The exploration policy ct is defined as the greedy policy with respect to Q̄t , namely ct (s, h) ≜ argmax Q̄t,h (s, a),

∀(s, h) ∈ S × [H].

a∈A

ρ ρ In the sequel, we denote by Vt,· the value function in Mtf , and, by V̄t,· the value function in the ”optimistic MDP” M̄tf ≜ (rtf + bonust , p̂t ). The associated Q-value functions are defined similarly.

Cyrille Kone, Kevin Jamieson

Algorithm 4: Forced Exploration Helper Function Function ComputeForcedExplorationPolicy(t, p̂t , nt ) 2 Sample (s̃t , ãt , h̃t ) uniformly at random from S × A × [H]

3 4 5 6 7

Set reward function rtf (s, a, h) ← 1 s = s̃t , a = ãt , h = h̃t

Initialize V̄t,H+1 (s) ≡ 0, ∀s ∈ S for h = H, H − 1, . . . , 1 do for (s, a) ∈ S × A do Q̄t,h (s, a) ← rtf (s, a, h) + p̂t,h V̄t,h+1 (s, a) + bonust (s, a, h) V̄t,h (s) ← maxa∈A Q̄t,h (s, a),

∀s ∈ S // Greedy policy with respect to optimistic Q

for (s, h) ∈ S × [H] do ct (s, h) ← argmaxa∈A Q̄t,h (s, a)

return policy ct

E.1.2

Regret Upper Bound

We define the well-explored set for a fixed T n o Zt ≜ (s, a, h) : n̄t (s, a, h) ⩾ 4 log(2SAHT 2 ) .

(E.10)

On the event EEcnt (T ), we have for all (s, a, h) ∈ Zt : nt (s, a, h) ⩾ 14 n̄t (s, a, h).

(E.11)

We also introduce for any sequence (κt )t⩾1 , the sum operator κ1:T ≜

t=1 κt .

We prove the following result. Lemma 11. Let (St )t⩾1 be a sequence of binary √ random variables. On the event EE (T ), the policies (ct )t⩾1 produced by Algorithm 4 satisfy for any K ∈ ( T , T ), v   u K K u X X √ p ct ⋆ St εt Vt,0 − Vt,0 κt whct (s, a) , ⩽ T 1−γ + O(log T ) + O β p (T, 1/T 2 )t1 + t=1

t=1

where κt = St εt . Proof. We decompose K X

ct ⋆ κt Vt,0 − Vt,0

t=1

√

⩽

T X

κt +

t=1

| {z } (I)

K X √ t= T +1

ct κt V̄t,0 − Vt,0 . {z

(E.12)

}

(II)

We focus on controlling (II). √ For t ⩾ T , we have 1/t3 ⩽ 1/T 2 . Since δ 7→ β p (·, δ) is non-increasing, the exploration bonus used in Algorithm 4 dominates the one associated with confidence level 1/T 2 . Therefore, when EEp (T ) holds, from classical UCBVI analysis technique (see e.g. Azar, Osband, et al. 2017) √ Q̄t,h (s, a) ⩾ Q⋆t,h (s, a), ∀(s, a, h), ∀t ∈ ( T , T ). By backward induction, this implies ⋆ V̄t,h (s) ⩾ Vt,h (s),

∀(s, h).

Optimal Posterior Sampling for Policy Identification in Tabular Markov Decision Processes

Since ct is greedy with respect to Q̄t , standard decomposition with value difference lemma yields X ct 2 whct (s, a) bonust (s, a, h) ∧ 1 . ⩽ V̄t,0 − Vt,0 s,a,h

We then obtain (II) ⩽

K X X

2 κt whct (s, a) bonust (s, a, h) ∧ 1 .

t=1 s,a,h

For (s, a, h) ∈ / Zt , we have n̄t (s, a, h) < 4 log(SAHT 2 ), hence K X

2κt whct (s, a) bonust (s, a, h) ∧ 1

⩽

t=1 (s,a,h)∈Z / t

K X X

2εt whct (s, a)1

t=1 s,a,h

t−1 X

! εk whck (s, a) < 4 log(SAHT 2 )

k=1

⩽ 2 + 8 log(SAHT ). 2

(E.13)

For the well-explored triplets, let us define (III) ≜

K X

2κt whct (s, a) bonust (s, a, h) ∧ 1 .

t=1 (s,a,h)∈Zt

Using Eq. (E.11) K X

! K X (nt (s, a, h), 1/T 2 ) c ⩽ O β p (T, 1/T 2 ) log 1 + 14 κt wht (s, a) max (1, nt (s, a, h)) t=1

β κt whct (s, a)

t=1 (s,a,h)∈Zt

(E.14)

which follows by elliptic potential (Lemma 46), the definition of Zt , and monotonicity of β p . Similarly, since V̄t,h+1 (·) ∈ (0, 1), we have Varp̂t,h V̄t,h+1 (s, a) ⩽ 1. Hence, K X

t=1 (s,a,h)∈Zt K X

2 Varp̂t,h [V̄t,h+1 ](s, a) β p (nt (s, a, h), 1/T 2 ) ⩽ max (1, nt (s, a, h)) s 2 β p (T, 1/T 2 ) ct ⩽ κt wh (s, a) max 1, 14 n̄t (s, a, h)

κt whct (s, a)

t=1 (s,a,h)∈Zt

K X X p r 4 2 β p (T, 1/T 2 ) t=1 s,a,h

κt whct (s, a) , Pt−1 ck 1 max 1, 4 k=1 κk wh (s, a)

where the last inequality uses

n̄t (s, a, h) ⩾

t−1 X

κk whct (s, a),

k=1

which follows from Eq. (E.4), where κk = εk Sk and Sk is binary. Applying Lemma 16 of Kaufmann, Ménard, et al. 2021b, we obtain K X

t=1 (s,a,h)∈Zt

s κt whct (s, a)

2 Varp̂t,h [V̄t,h+1 ](s, a) β p (nt (s, a, h), 1/T 2 ) max (1, nt (s, a, h))

v   u K u X p ⩽ O β p (T, 1/T 2 )t1 + κt wct (s, a) . h

t=1

Cyrille Kone, Kevin Jamieson

Combining this bound with Eq. (E.14) yields, on the event EE (T ), v   u K u X p κt wct (s, a) (III) ⩽ O β p (T, 1/T 2 )t1 + h

(E.15)

t=1

Finally, combining Eq. (E.12), Eq. (E.13), and Eq. (E.15) concludes the proof. E.2

Suffficient Exploration of Reachable States

In this section, we prove suﬀicient exploration of reachable state–actions. We recall that Wh (s, a) ≜ sup whρ (s, a), ρ

is the maximum visitation probability of triplet (s, a, h). We prove the following lemma. Lemma 12. Let γ ∈ (0, 1). There exists an explicit finite constant k0 < ∞, such that on the event EE (T ), for all T ⩾ k0 , all t ∈ (T 2/3 , T ), and all (s, a, h), nt (s, a, h) ⩾

Wh (s, a) −γ t . 4SAH

The explicit expression of k0 is given in the proof. Proof. Let us fix a triplet (s, a, h) and t ∈ (T 2/3 , T ). We note that the claim is trivially satisfied for unreachable triplets. Hence, in the sequel, we assume Wh (s, a, h) > 0. We have E 1 sth = s, ath = a Ht−1

(1 − εt ) E 1(sh = s, ah = a) Ht−1 , Zt = 0 + εt E 1 sth = s, ath = a Ht−1 , Zt = 1 ,

⩾

εt whct (s, a) .

Next, when EEcnt (T ) holds, nt (s, a, h)

⩾

t−1 1X E 1 skh = s, akh = a Hk−1 − log(6SAHT 2 ) , 2 k=1

1X εk whck (s, a) − log(6SAHT 2 ), 2 t−1

⩾

k=1

and the right-hand side is related to the dynamic regret of the UCBVI forced exploration routine as t−1 X

εk whci (s, a)

k=1

⩾

t−1 X

ck εk 1 s̃k = s, ãk = a, h̃k = h Vt,0 .

k=1

Thanks to Lemma 11, on the event EE (T ), with Sk ≜ 1 s̃k = s, ãk = a, h̃k = h , and κk ≜ εk Sk t−1 X k=1

ck κk Vk,0 ⩾

t−1 X k=1

p √ ⋆ t1−γ · β p (T, 1/T 2 ) . κk Vk,0 − O( T 1−γ ) − O(log T ) − O

(E.16)

Optimal Posterior Sampling for Policy Identification in Tabular Markov Decision Processes

Rearranging the above displays yields nt (s, a, h) ⩾ t−1 p √ 1X k k k εk 1 s̃ = s, ã = a, h̃ = h max whρ (s, a) − O( T 1−γ ) − O(log T ) − O t1−γ · β p (T, 1/T 2 ) . ρ 2 k=1

We remark that on the event EEZ (T ), t−1 X k=1

1 X εk · 1 s̃ = s, ã = a, h̃ = h ⩾ εk − log(2SAHT 2 ) 2SAH k

t−1

k=1

we have for any t ∈ (T 2/3 , T ), and any triplet (s, a, h) q q 3(1−γ) Wh (s, a) nt (s, a, h) ⩾ ε1...t−1 − O t 2 − O(log t) − O t1−γ · β p (t3/2 , 1/t3 ) . 2SAH

(E.17)

Now, let us define Wh (s, a) 1−γ T0 (s, a, h) ≜ sup t ⩾ 1 : nt (s, a, h) ⩽ t . 4SAH

(E.18)

1−γ

From Eq. (E.17), and, as ε1...t ∼ t1−γ , we know that on EE (T ), k0 (s, a, h) is upper bounded by a deterministic constant k0 (s, a, h) defined as q q 3(1−γ) ε1...t t1−γ 2/3 t 2 + O(log t) + O t1−γ · β p (t3/2 , 1/t3 ) . k0 (s, a, h) ≜ sup t ⩾ 1 : Wh (s, a) − Wh (s, a) ⩽ O 2SAH 4SAH Indeed, the expressions in O(·) can be recovered from the proof to make the above well-defined, moreover, it is finite for any γ ∈ (0, 1) and any reachable (s, a, h). Next we define k0 ≜

max (s,a,h):Wh (s,a)>0

k0 (s, a, h),

which concludes the proof. E.3

Guarantees on Model Estimation

We prove that after a finite number of episodes the empirical optimal policy coincides with the true optimal policy on most episodes, when EE (T ) holds. The key tool is the following value difference lemma. Lemma 13 (Value Difference). For any policy π ∈ Π, "H # X π π π V0 − V̂t,0 = EM bt (sh , ah , h) , h=1

π where bt (s, a, h) ≜ µ(s, a, h) − µ̂t (s, a, h) + ph − p̂t,h V̂t,h+1 (s, a).

Proof. By Bellman’s equations, h i π π π Vhπ (s) − V̂t,h (s) = Ea∼πh (·|s) µh (s, a) − µ̂t (s, a, h) + ph Vh+1 (s, a) − p̂t,h V̂t,h+1 (s, a) h i π π = Ea∼πh (·|s) bt (s, a, h) + ph Vh+1 (s, a) . − V̂t,h+1 Applying this recursion from h down to H and using s1 = sinit gives the claimed result.

Cyrille Kone, Kevin Jamieson

Combining Lemma 13 with the suﬀicient exploration guarantee of Section E.2 and standard concentration bounds shows that the empirical value of any policy converges quickly to its true value. In particular, once the right-hand side of Lemma 13 falls below ∆min (Π) ≜ V0⋆ − max ⋆ V0π , (E.19) π∈Π\{π }

⋆

the empirical optimal policy π̂t must equal π . This is formalized in the next lemma. Lemma 14. There exists an explicit finite constant k0 < ∞ such that on the event EE (T ), for all T > k0 and all t ∈ (T 2/3 , T ), π̂t = π ⋆ . The constant k0 is given explicitly in the proof. ⋆

π π for all π 6= π ⋆ . > V̂t,0 Proof. Since π ⋆ is the unique optimal start-state policy for M , it suﬀices to show that V̂t,0 π By Lemma 13, for any π ∈ Π, V0π − V̂t,0 = EπM

H h=1 bt (sh , ah , h) .

π Since V̂t,h+1 (s) ⩽ H for all (s, h), Pinsker’s inequality gives, on EE (T ),

s |bt (s, a, h)| ⩽

2β µ (t, 1/t2 ) + nt (s, a, h)

2H 2 β p (t, 1/t2 ) . nt (s, a, h)

By Lemma 12, on EE (T ), there exists another k0′ < ∞ such that for T ⩾ k0′ , and t ∈ (T 2/3 , T ), and any reachable (s, a, h), Wh (s, a) 1−γ t nt (s, a, h) ⩾ 4SAH . Therefore, when EE (T ) holds, for T ⩾ k0′ , and t ∈ (T 2/3 , T ), s s 8SAH β µ (t, 1/t2 ) 8SAH 3 β p (t, 1/t2 ) |bt (s, a, h)| ⩽ + . 1−γ Wh (s, a) t Wh (s, a) t1−γ

(E.20)

Since whπ (s, a) = 0 for unreachable triplets (s, a, h), π V0π − V̂t,0

whπ (s, a) bt (s, a, h)

s,a,h: Wh (s,a)>0

⩽

whπ (s, a)

s,a,h: Wh (s,a)>0

8SAH β µ (t, 1/t2 ) + Wh (s, a) t1−γ

# 8SAH 3 β p (t, 1/t2 ) . Wh (s, a) t1−γ

We define ( 2/3 T0 ≜

max

s,a,h: Wh (s,a)>0

sup t ⩾ 1 :

8SAH β µ (t, 1/t2 ) + Wh (s, a) t1−γ

8SAH 3 β p (t, 1/t2 ) ∆min (Π) ⩾ Wh (s, a) t1−γ 2

) ,

(E.21)

which is finite since β µ , β p are logarithmic in t, γ ∈ (0, 1), and ∆min (Π) > 0. We further let k0 = max (k0′ , T0 ). Thus, when EE holds, for T ⩾ k0 , t ∈ (T 2/3 , T ) and any π 6= π ⋆ , ⋆

⋆

π π π > 0, V̂t,0 − V̂0t,π ⩾ ∆min (Π) + V̂t,0 − V0π + V0π − V̂t,0 | {z } | {z } > −∆min (Π)/2

so that π̂t = π ⋆ .

> −∆min (Π)/2

Optimal Posterior Sampling for Policy Identification in Tabular Markov Decision Processes

F ANALYSING CONDITIONAL POSTERIOR SAMPLING In this section, we analyze the guarantees on the inf player (distribution over Alt(M̂T )) by interpreting it as an instance of continuous exponential weights (CEW). The goal of these results is to demonstrate that conditional posterior sampling, akin to continuous exponential weights, is a no-regret learner of the best challenger MDP. We recall the log-likelihood ratio between a model (θ, Φ) and the empirical MDP M̂T after (T − 1) episodes " # 2 X X θ(s, a, h) − µ̂T (s, a, h) p̂t (s′ | s, a, h) ′ ΥT (θ, Φ) ≜ nT (s, a, h) + p̂T (s | s, a, h) log . (F.1) 2 Φ(s′ | s, a, h) ′ s

s,a,h

The conclusion of this section is to prove the result below. Theorem 15. Let ηt ∝ t−ς with ς > 2α. There exists a constant k0 < ∞ and a deterministic function f (T ) = o(T ) such for any T ⩾ k0 , the inequality T −1 X

E(θ,Φ)∼ν ηt (·|Alt(M̂t )) [ℓα t (θ, Φ)] ⩽ t

t=T 2/3

T −1 X X

inf

(θ,Φ)∈Alt(M̂T ) t=1

ℓα t (θ, Φ) ≜

1 (µ̂T (sth , ath , h) − θ(sth , ath , h))2 + log + f (T ) , where 2 Φ(sth+1 , | sth , ath , h)

X (µ̂t sth , ath , h) − θ(sth , ath , h) 2 h

+ log

1 , max Φ(sth+1 | sth , ath , h), e−tα

holds with probability at least 1 − O(1/T 2 ). Moreover, f can be made explicit from the proof. F.1

Setup

Before introducing the proof of this result, let us introduce the function over reward parameters and the function over the transition kernels defined for some η > 0 by    X η · n (s, a, h)  t η ft (θ) ≜ exp − (µ̂t (s, a, h) − θ(s, a, h))2 ,   2 s,a,h   (F.2) X X  η gt (Φ) ≜ exp η · nt (s′ |s, a, h) log Φ(s′ |s, a, h) .   ′ s,a,h s

Since ftη , gtη are random variables, all results claimed in the next lemma hold with probability 1. Lemma 16. Let η > 0 and let M ≜ (θ, Φ) ∈ M be any MDP with positive transitions. After observing the H trajectory (sth , ath , rht , sth+1 ) h=1 at episode t, the following identities hold with probability 1: ( ) H η X (Φ) gt+1 1 = exp −η , log gtη (Φ) Φ(sth+1 | sth , ath , h)

(F.3)

( ) H η 2 η (θ) ft+1 ηX t t t t µ̂t (sh , ah , h) − θ(sh , ah , h) + σt+1 (θ) , = exp − ftη (θ) 2 2

(F.4)

h=1

where σt+1 (θ) ≜

H X h=1

h 2 2 i nt+1 (sth , ath , h) µ̂t (sth , ath , h) − θ(sth , ath , h) − µ̂t+1 (sth , ath , h) − θ(sth , ath , h) .

(F.5)

Cyrille Kone, Kevin Jamieson

exp −η

X h

1 log t Φ(sh+1 | sth , ath , h)

) ,

where (a) follows since the observed transition (sth , ath , sth+1 ) contributes exactly +1 to the count. The proof of the second statement of the lemma follows by noting that ( η X η · nt+1 (st , at , h) ft+1 (θ) h h = exp − (µ̂t+1 (sth , ath , h) − θ(sth , ath , h))2 + η ft (θ) 2 h ) X η · nt (st , at , h) t t t t 2 h h (µ̂t (sh , ah , h) − θ(sh , ah , h)) , 2 h ( i X η · nt+1 (st , at , h) h (b) h h = exp (µ̂t (sth , ath , h) − θ(sth , ath , h))2 − (µ̂t+1 (sth , ath , h) − θ(sth , ath , h))2 − 2 h ) Xη t 2 t t t (µ̂t (sh , ah , h) − θ(sh , ah , h)) , 2 h

where (b) follows since nt+1 (sth , ath , h) = nt (sth , ath , h) + 1. Indeed, since (sth , ath ) is visited at stage h during episode t, the count increases by exactly 1. With the definition of σt+1 , we have proven the claimed statement. F.2

Posterior Density Ratio and Stochastic Hedge Loss

We introduce further notation related to the posterior density. As the transition and rewards are independent, it is a simple exercise to show that the posterior density is simply the product of gtηt and ftηt . For η > 0, we introduce λη defined on M as 1 λt (θ, Φ) ≜ − log {ftη (θ)gtη (Φ)} , η

(F.6)

f ≜ (θ, Φ), and we define a density w.r.t. canonical measure of M proportional to for an MDP M ftη (θ)gtη (Φ) · 1 (θ, Φ) ∈ Alt(M̂t ) whose normalizing constant is Z Z Ληt ≜ ftη (θ)gtη (Φ) · 1 (θ, Φ) ∈ Alt(M̂t ) dθdΦ = e−ηλt (θ,Φ) · 1 (θ, Φ) ∈ Alt(M̂t ) dθdΦ ,

(F.7)

where we recall that Alt(M̂t ) ⊂ M and dθdΦ denotes the standard (product) measure on this space. We denote by νtη (·) the distribution over M with product density ftη · gtη and by νtη (· | X ) its truncattion to the set X ⊆ M. f ≜ (θ, Φ) from νtη (· | Alt(M̂t )) is proportional to Given Ht−1 , the probability of sampling an MDP M e−ηλt (θ,Φ) · 1 (θ, Φ) ∈ Alt(M̂t ) .

Optimal Posterior Sampling for Policy Identification in Tabular Markov Decision Processes

Lemma 17. Let η > 0. If π̂t = π̂t+1 then, we have log

h i 1 h i Ληt+1 α 1 log E(θ,Φ)∼ν η (·|Alt(M̂t )) e−2η·ℓt (θ,Φ) + log E(θ,Φ)∼ν η (·|Alt(M̂t )) eη·σt+1 (θ) . η ⩽ t t Λt 2 2

Proof. First, we note that given and MDP M , the set of MDPs M ′ such that M is not absolutely continuous w.r.t. M ′ has empty interior. Therefore Alt(M̂t ) = Alt(M̂t+1 ) (up to a null set) whenever π̂t = π̂t+1 . We have

Λη log t+1 = log Ληt

η η (Φ) · 1 (θ, Φ) ∈ Alt(M̂t ) dθ dΦ (θ) gt+1 ft+1 R η ft (θ) gtη (Φ) · 1 (θ, Φ) ∈ Alt(M̂t ) dθ dΦ

η η (Φ) (θ) gt+1 ft+1 = log E(θ,Φ)∼ν η (·|Alt(M̂t )) · η t ftη (θ) gt (Φ) " ( )# ηX 1 η 2 = log E(θ,Φ)∼ν η (·|Alt(M̂t )) exp − µ̂t (sth , ath , h) − θ(sth , ath , h) + 2 log + σt+1 (θ) t 2 Φ(sth+1 | sth , ath , h) 2 h " ( )# X (a) 1 2 1 t t t t µ̂t (sh , ah , h) − θ(sh , ah , h) + 2 log ⩽ log E(θ,Φ)∼ν η (·|Alt(M̂t )) exp −η t 2 Φ(sth+1 | sht , ath , h) h h i 1 + log E(θ,Φ)∼ν η (·|Alt(M̂t )) eη·σt+1 (θ) , t 2

where (a) follows from the Cauchy–Schwarz inequality log E[XY ] ⩽ 21 log E[X 2 ] + 12 log E[Y 2 ]. Intuitively, everything proceeds as if we were running continual exponential weights with learning rate η and stochastic loss # " 2 H X µ̂t (sth , ath , h) − θ(sht , ath , h) 1 + log . (F.8) lt (θ, Φ) ≜ 2 Φ(sth+1 | sth , ath , h) h=1

However, to control the magnitude of the stochastic loss, which might be unbounded due to low transition probabilities, we work with the clipped loss ℓα t (θ, Φ) ≜

H X h=1

µ̂t (sth , ath , h) − θ(sth , ath , h) 2

# 1 , + log max Φ(sth+1 | sth , ath , h), e−tα

(F.9)

which satisfies ℓα t (θ, Φ) ⩽ lt (θ, Φ) by definition. Combining with the calculations above, we have with probability 1, log

i h i 1 h Ληt+1 α 1 log E(θ,Φ)∼ν η (·|Alt(M̂t )) e−2η ℓt (θ,Φ) + log E(θ,Φ)∼ν η (·|Alt(M̂t )) eη σt+1 (θ) . η ⩽ t t Λt 2 2

We further define φηt+1 ≜

i h 1 log E(θ,Φ)∼ν η (·|Alt(M̂t )) eη σt+1 (θ) , t 2

which will be later upper bounded in Lemma 21.

(F.10)

Cyrille Kone, Kevin Jamieson

F.3

Hedge with Decreasing Learning Rate

We use a decreasing learning rate schedule (ηt )t⩾1 rather than a constant one, since a constant rate would require the doubling trick and discarding of past data. Throughout, we assume (ηt )t⩾1 is non-increasing and satisfies PT −ς t=1 ηt = o(T ). For generality, we assume ηt = t . Lemma 18. Let T0 ⩽ T and assume the event EFT0 (T ) ≜ {∀ T0 ⩽ t ⩽ T, π̂t = π̂T0 )} holds. Define α ψt+1 ≜

where Ct,α ≜ H 12 + t T −1 X t=T0

2 ηt Ct,α t , + ηt−1 φηt+1 2

ΨTα (T0 ) ≜

T −1 X

α ψt+1 ,

t=T0

t and φηt+1 is defined in (F.10). Then

1 1 η log ΛηTT + ΨTα (T0 ) + log ΛTT00 + E(θ,Φ)∼ν ηt (·|Alt(M̂t )) ℓα t (θ, Φ) ⩽ − t ηT ηT0 T −1 X t=T0

1 1 ηt+1 ηt log Λt+1 − log Λt+1 . ηt ηt+1

Proof. Thanks to Lemma 17 we have log

h i Ληt+1 1 η −2η ℓα t (θ,Φ) + φ η ⩽ log E e t+1 (θ,Φ)∼νt (·|Alt(M̂t )) Ληt 2

1 1 α α where φηt+1 is defined in Eq. (F.10). Next, we note that ℓα t ∈ (0, ( 2 + t ) · H), and we define Ct,α ≜ H ( 2 + t ).

By Hoeffding’s lemma we have for any η, h

E(θ,Φ)∼ν η (·|Alt(M̂t )) e

−2η ℓα t (θ,Φ)

Therefore,

( ⩽ exp

−2ηE(θ,Φ)∼ν η (·|Alt(M̂t )) [ℓα t (θ, Φ)] + t

2 η 2 Ct,α 2

) .

α 1 Ληt 1 −1 ηt 2 log t+1 ηt ⩽ −E(θ,Φ)∼νtηt (·|Alt(M̂t )) ℓt (θ, Φ) + ηt Ct,α + ηt φt+1 . ηt Λt 2 | {z }

(F.11)

(F.12)

α ψt+1

Assuming EFT0 (T ) and summing (F.12) from t = T0 to T − 1 gives T −1 X

T −1 −1 X α TX Λ ηt 1 α ηt E ℓ ψt+1 log t+1 ⩽ − (θ, Φ) + . (θ,Φ)∼νt (·|Alt(M̂t )) t ηt Ληt t t=T0 t=T0 t=T0 | {z }

(F.13)

ΨTα (T0 )

The left-hand side is not directly telescoping. We decompose it as T −1 X t=T0

T −1 X Ληt 1 1 1 ηt ηt log t+1 = log Λ − log Λ t t+1 ηt Ληt t ηt ηt t=T0

−1 T X

t=T0

1 ηt+1

t+1 log Λt+1 −

TX −1 1 1 1 ηt+1 t log Ληt t + log Ληt+1 − log Λt+1 ηt ηt ηt+1

1 1 η log ΛηTT − log ΛTT00 + ηT ηT0

T −1 X t=T0

t=T0

1 1 ηt+1 t − , log Ληt+1 log Λt+1 ηt ηt+1

where the last step uses that the first sum telescopes. Substituting back into Eq. (F.13) and rearranging gives the claimed inequality.

Optimal Posterior Sampling for Policy Identification in Tabular Markov Decision Processes

F.4

Best Model in Hindsight

The idea here is to show a lower bound on this quantity when there is enough mass around the best model. While such bounds are standard to derive for Hedge on convex sets (with convex loss), in our case, although the loss is convex, the set Alt(M̂t ) is non-convex and a priori not a simple union of convex sets. 1 Lemma 19. Let γ ≜ 6T SAH(H+1) . We have

T −1 X X (µ̂T (sth , ath , h) − θ(sth , ath , h))2 1 1 ηT inf log ΛT ⩾ − + log ηT 2 Φ(sth+1 , | sht , ath , h) (θ,Φ)∈Alt(M̂T ) t=1 h

√

TH log vol(Alt(M̂T ) ∩ M1/S· T ) 2(SAH + S AH) log γ −γ log(eS 2 T ) + − O(1) . η 2 η 2

Proof. The proof of this result relies on a lower bound on the inflated posterior over Alt(M̂T ). Simplifying notation, let η ≜ ηT and we have Z ΛηT =

e−ηλT (θ,Φ) 1 θ, Φ) ∈ Alt(M̂T ) dθd Φ ,

(F.14)

where ληT is a convex function as defined earlier. We prove that the posterior puts the majority of its mass around the mode. First, we need to prove that there is enough volume around the best model (θT⋆ , Φ⋆T ) ∈

argmin

λT (θ, Φ),

(F.15)

(θ,Φ)∈cl(Alt(M̂T ))

f⋆ = (θ⋆ , Φ⋆ ). so that the integral in Eq. (F.14) will concentrate around the mode. We define M T T T First, note that this argmin is well defined as cl(Alt(M̂T )) is a compact set. Because the instance in (F.15) f⋆ which is still contained in the might be close to the boundary of Alt(M̂T ), we build another instance around M T alternative set Alt(M̂T ), and with a non-negligible enclosing volume around it. f⋆ belongs to the closure of the alternative set, there exists a deterministic policy π̃ ∈ Πdet \{π̂t } such Because M T f⋆ . that π̃ is globally optimal in M T f′ ≜ (θ′ , Φ′ ) such that π̃ is the unique Next, we invoke Lemma 43 which shows that there exists some MDP M global optimal and start-state policy, and its transition and reward kernels satisfy  f′ )  ∆(M ′ Φ (· | s, a, h)  ⋆ |θT (s, a, h) − θ′ (s, a, h)|

⩾ T1 ≜ (1 − γs,a,h )Φ⋆T (· | s, a, h) + γs,a,h Φ⋆T (· | s, πh (s), h) , ⩽ ds,a,h ≜ 2H−h+1 ε

(F.16)

i h (1+2H−h )ε where γs,a,h ∈ 0, 1−(2 H−h+1 +3)ε . f′ ∈ Alt(M̂T ). Next, we relate Λη to the (near optimal) λη (θ′ , Φ′ ). Thanks to Lemma 44, for In particular, M T T 1 fT ≜ (θT , ΦT ) the instance built by this lemma. We then have γ ≜ 6T SAH(H+1) , and calling M f′ + γ Alt(M̂T ) ⊂ Alt(M̂T ). (1 − γ)M

(F.17)

Cyrille Kone, Kevin Jamieson

Therefore, Z log ΛηT = log

⩾ log Z

e−ηλT (θ,Φ) 1 (θ, Φ) ∈ Alt(M̂T ) dθdΦ f′ +γ Alt(M̂T ) (1−γ)M

e−ηλT (θ,Φ) dθdΦ

2 γ S AH+SAH exp −ηλT ((1 − γ) · θ′ , Φ′ + γ · θ, Φ) dθdΦ Alt(M̂T ) Z 2 ′ ′ ⩾ log γ S AH+SAH e−η(1−γ)·λT (θ ,Φ ) + log e−ηγ·λT (θ,Φ) dθdΦ Alt(M̂T ) Z ′ ′ 2 = −η(1 − γ)λT (θ , Φ ) + (S AH + SAH) log γ + log e−ηγλT (θ,Φ) dθdΦ

⩾ log

(F.18)

Alt(M̂T )

which follows by convexity of λT . Combining the results above, in particular in Eq. (F.16), ′

′

λT (θ , Φ )

⩽

T −1 X X t=1

⩽

(µ̂T (sth , ath , h) − θ′ (sth , ath , h))2 1 + log 2 (1 − γh )Φ⋆T (sth+1 | sth , ath , h)

′

1 (µ̂T (sth , ath , h) − θT⋆ (sth , ath , h))2 + log 2 (1 − γh )Φ⋆T (sth+1 | sth , ath , h) t=1 h +T H 4H ε2 + 2H ε 1 λT (θT⋆ , Φ⋆T ) + T H 4H ε2 + 2H ε + log (monotonicity of h 7→ γh ) . 1 − γ1 T −1 X X

Therefore

λT (θ , Φ ) ⩽ λT (θT⋆ , Φ⋆T ) + T H

1 4 ε + 2 ε + log 1 − γ1 H 2

(F.19)

It remains to control the residual integral Z

e−ηγλT (θ,Φ) dθdΦ.

log Alt(M )

We note that if this is too small the bound above will be meaningless. Also, observe that the integraand is √ 1/S· T to be the unbounded, (due to the transition probabilities) hence to control this quantity we introduce M √ set ofs MDPs for which each transition is at least 1/S · T . We have Z Z −ηγλT (θ,Φ) e−ηγλT (θ,Φ) dθdΦ e dθdΦ ⩾ √ Alt(M ) Alt(M )∩M1/ T   Z  X  1 1 ⩾ exp −ηγ nt (s, a, h) + log(S 2 T ) dθdΦ √   2 2 Alt(M̂T )∩M1/S· T s,a,h

−ηγ T2H log(eS 2 T )

√

vol(Alt(M̂T ) ∩ M1/S· T )

Combining the above displays yields, Z √ TH log(eS 2 T ) + log vol(Alt(M̂T ) ∩ M1/S· T ). log e−ηγλT (θ,Φ) dθdΦ ⩾ −ηγ 2 Alt(M )

(F.20)

Plugging-in this back into Eq. (F.18) yields √

1 2(SAH + S 2 AH) log γ log(eS 2 T ) log vol(Alt(M̂T ) ∩ M1/S· T ) log ΛT ⩾ −(1 − γ)λT (θ′ , Φ′ ) + −γ + η η 2 η

Optimal Posterior Sampling for Policy Identification in Tabular Markov Decision Processes

and combining with Eq. (F.19) yields 1 log ΛηT ⩾ η 2(SAH + S − λT (θT⋆ , Φ⋆T ) +

T H 4H ε2 + 2H ε + log

√

AH) log γ

TH log vol(Alt(M̂T ) ∩ M1/S· T ) −γ log(eS 2 T ) + − 2 η

(F.21)

1 . 1 − γ1

i h H−1 ) Plugging ε = 1/T in the above, using the the inequality log(1 + x) ⩽ x and recalling that γ1 ∈ 0, T(1+2 H −(2 +3) yields H 4 γ1 H + 2H + T = O(1). T 1 − γ1 Rewriting the terms above, we have √

2(SAH + S 2 AH) log γ TH log vol(Alt(M̂T ) ∩ M1/S· T ) 1 log ΛηT ⩾ −λT (θT⋆ , Φ⋆T ) − −γ log(eS 2 T ) + − O(1) η η 2 η which achieves to prove the claimed result. F.5

Technical and Concentration Lemmas

We show a series of lemmas to complete the analysis of the conditional posterior resampling. Lemma 20. We have " # TX T −1 −1 X 1 1 log vol(Alt(M̂t+1 )) log vol(Alt(M̂t+1 )) ηt+1 ηt log Λt+1 − log Λt+1 ⩽ − . ηt ηt+1 ηt ηt+1 t=T0

t=T0

Proof. Let ρt (θ, Φ) ≜

1 vol(Alt(M̂t ))

1 (θ, Φ) ∈ Alt(M̂t )

be the uniform measure on Alt(M̂t ). By definition of Ληt , Z 1 log vol(Alt(M̂t )) 1 log Ληt = log ρt (θ, Φ)e−η λt (θ,Φ) dθ dΦ + . η η η R Introducing ft (η) ≜ η1 log ρt (θ, Φ)e−η λt (θ,Φ) dθ dΦ, we decompose 1 1 ηt+1 t log Ληt+1 − log Λt+1 = ηt ηt+1 t=T0 # " T −1 T −1 X X log vol(Alt(M̂t+1 )) log vol(Alt(M̂t+1 )) . [ft+1 (ηt ) − ft+1 (ηt+1 )] + − ηt ηt+1 T −1 X

t=T0

By Lemma 38, ft+1 is non-decreasing. Since (ηt )t is non-increasing, ηt ⩾ ηt+1 for all t, so ft+1 (ηt )−ft+1 (ηt+1 ) ⩾ 0, which concludes the proof. Lemma 21. Define Ct,α ≜ H 12 + tα and assume ηt ∝ t−ς , with ς > 2α. There exists deterministic fΨ such that the event ( # ) " T −1 2 X ηt Ct,α −1 ηt Ψ α + ηt φt+1 ⩽ fΨ (T ) EF (T ) ≜ ΨT (T0 ) ≜ 2 t=T0

holds with probability 1 − O(1/T ). Moreover, fΦ is given in the proof and satisfies fΨ = o(T ) 2

Cyrille Kone, Kevin Jamieson

Proof. We bound each term in ΨTα (T0 ) separately. For the first term, we observe that ! T −1 T −1 X X 2 2α −ς ηt Ct,α ⩽ O t t = o(T ), t=1

t=T0

when ς > 2α. t To control the second term in Ψ , recalling the definition of φηt+1 from (F.10) and σt+1 (θ) from (F.5), we have X σt+1 (θ) ⩽ 2 nt+1 (sth , ath , h) µ̂t (sth , ath , h) − θ(sth , ath , h) µ̂t (sth , ath , h) − µ̂t+1 (sth , ath , h)

2 nt+1 (sth , ath , h) µ̂t (sth , ath , h) µ̂t (sth , ath , h) − µ̂t+1 (sth , ath , h)

−

X |

}

Jt+1

2 nt+1 (sth , ath , h) θ(sth , ath , h) µ̂t (sth , ath , h) − µ̂t+1 (sth , ath , h) .

}

κt+1 (θ)

The term Jt+1 is independent of (θ, Φ) and is controlled bypAzuma–Hoeffding (see Lemma 22 below). For κt+1 (θ), applying Hoeffding’s lemma with the bound |κt+1 (θ)| ⩽ ϕt+1 , where !2 X t t t t t t , ϕt+1 ≜ 2 nt+1 (sh , ah , h) µ̂t (sh , ah , h) − µ̂t+1 (sh , ah , h) h

gives

h i 1 ηt−1 log E(θ,Φ)∼ν ηt (·|Alt(M̂t )) e−ηt κt+1 (θ) ⩽ −E(θ,Φ)∼ν ηt (·|Alt(M̂t )) [κt+1 (θ)] + ηt ϕt+1 . t t 2 Summing over t and applying concentration to both Jt+1 and ϕt+1 (via Lemmas 22, 23, and 24) shows that PT −1 −1 ηt 2 t=K ηt φt+1 = o(T ) on an event with probability larger than 1 − O(1/T ). The conclusion then follows. Lemma 22. There exists a deterministic function fJ such that the event ) ( H T X X J t t t t t t t t 2 nt+1 (sh , ah , h) µ̂t (sh , ah , h) µ̂t (sh , ah , h) − µ̂t+1 (sh , ah , h) ⩽ fJ (T ) , EF (T ) JT ≜ t=1 h=1

holds with probability at least 1 − O(1/T 2 ). Moreover fJ is given in the proof and satisfies fJ = o(T ). P Proof. Define St (s, a, h) ≜ k∈[t−1] 1 skh = s, akh = a Rhk , where Rhk | skh , akh ∼ µ(skh , akh , h)+ζhk and ζhk ∼ N (0, 1) is independent noise. A direct computation gives nt+1 (sth , ath , h) µ̂t (sht , ath , h) − µ̂t+1 (sth , ath , h) = St (sth , ath , h) − St+1 (sth , ath , h) + µ̂t (sth , ath , h) = µ̂t (sth , ath , h) − µ(sth , ath , h) + ζht , so that JT = 2 |

T X H X t=1 h=1

µ̂t (sth , ath , h) µ̂t (sth , ath , h) − µ(sth , ath , h) {z

}

(I)

+2 |

T X H X t=1 h=1

µ̂t (sth , ath , h) ζht . {z

(F.22)

}

(II)

Term (II). Since µ̂t is Ht−1 -measurable, (sth , ath , h) is independent of ζht which is itself independent of Ht−1 , and centered. Therefore, the partial sums form a martingale with subgaussian increments. By Azuma–Hoeffding’s inequality, the event v   u T X H   u X (II) ⩽ 2t2 log(T 2 ) µ̂t (sth , ath , h)2   t=1 h=1

Optimal Posterior Sampling for Policy Identification in Tabular Markov Decision Processes

holds with probability at least 1 − 1/T 2 . Term (I). By Cauchy–Schwarz, v v u T H u T t t 2 uX uX X 2 µ̂t (sh , ah , h) t (I) ⩽ 2 ·t max (1, nt (sth , ath , h)) µ̂t (sth , ath , h) − µ(sth , ath , h) t t max (1, n (s , a , h)) t h h t=1 t=1 h=1 v u H uX T X p µ̂t (sth , ath , h)2 t ⩽2 · T β µ (T, 1/T 2 ), t , at , h)) max (1, n (s t h h t=1 h=1

where the second inequality uses the self-normalized concentration event     X EFµ (T ) ≜ ∀t ⩽ T, nt (s, a, h) KL(R̂ht (s, a) k Rh (s, a)) ⩽ β µ (T, 1/T 2 ) ,   s,a,h

which holds with probability 1 − 1/T 2 . Since rewards are bounded in (0, 1), the elliptic potential yields T X

µ̂t (sth , ath , h)2 ⩽ 4 log(T + 1), max (1, nt (sth , ath , h)) t=1

giving (I) ⩽ 4H

p 2T β µ (T, 1/T 2 ) log(T + 1).

PT P Combining the bounds on (I) √ and (II), and using that t=1 h µ̂t (sth , ath , h)2 ⩽ HT together with β µ (T, 1/T 2 ) = O(log T ), we obtain JT ⩽ O T log T = o(T ) on an event holding with probability at least 1 − 1/T 2 . Lemma 23. There exists a deterministic sequence ϕ(T ) = o(T ) such that the event ) ( T H i2 XX h ϕ t t t t t t ⩽ ϕ(T ) EF (T ) ≜ ηt 2 nt+1 (sh , ah , h) µ̂t (sh , ah , h) − µ̂t+1 (sh , ah , h) t=1 h=1

holds with probability at least 1 − O(1/T 2 ), where ϕ(T ) is given explicitly in the proof. Proof. Using the identity established in the proof of Lemma 22, nt+1 (sth , ath , h) µ̂t (sth , ath , h) − µ̂t+1 (sth , ath , h) = µ̂t (sth , ath , h) − µ(sth , ath , h) + ζht , and the inequality (x + y)2 ⩽ 2x2 + 2y 2 , we obtain h

2 nt+1 (sht , ath , h) µ̂t (sth , ath , h) − µ̂t+1 (sth , ath , h)

⩽ 8 µ̂t (sth , ath , h) − µ(sth , ath , h)

2 + 8 ζht .

Using the self-normalized concentration event     X EFµ (T ) ≜ ∀t ⩽ T, nt (s, a, h) KL(R̂ht (s, a) k Rh (s, a)) ⩽ β µ (T, 1/T 2 ) ,   s,a,h

it follows T X

2 µ̂t (sth , ath , h) − µ(sth , ath , h) =

t=1

2 T X max (1, nt (sth , ath , h)) µ̂t (sth , ath , h) − µ(sth , ath , h) t=1

⩽

T X

max{1, nt (sth , ath , h)}

2β µ (T, 1/T 2 ) , max (1, nt (sth , ath , h)) t=1

(F.23)

Cyrille Kone, Kevin Jamieson

which thus holds with probability at least 1 − 1/T 2 . For the χ2 terms, by a union bound over SAHT events, the following o n p EFζ (T ) ≜ ∀ t ⩽ T, ∀ (s, a, h), |ζht | ⩽ 2 log(2SAHT 3 ) holds with probability at least 1 − 1/T 2 . On EFµ (T )∩EFζ (T ), which holds with probability at least 1−2/T 2 , substituting both bounds into (F.23), summing PT over t and h, and applying the elliptic potential bound t=1 max 1,n 1(st ,at ,h) ⩽ 4 log(T + 1) gives ( t h h ) T X H X

h i2 ηt 2 nt+1 (sth , ath , h) µ̂t − µ̂t+1 ((sth , ath , h))

t=1 h=1

⩽ 16H log(2SAHT 3 )

T X

ηt + 64H β µ (T, 1/T 2 ) log(T + 1),

t=1

which is o(T ) since

t=1 ηt = o(T ) and β

(T, 1/T ) = O(log T ).

Lemma 24. There exists a deterministic sequence d(T ) = o(T ) such that the event ( T ) X D EF (T ) ≜ −E(θ,Φ)∼ν ηt (·|Alt(M̂t )) κt+1 (θ) ⩽ d(T ) t

t=K

holds with probability at least 1 − 1/T 2 , where d(T ) is given explicitly in the proof. Proof. Recalling the definition of κt+1 (θ) from the proof of Lemma 21 and introducing ut (s, a, h) ≜ E(θ,Φ)∼ν ηt (·|Alt(M̂t )) θ(s, a, h) , t

we have DT (T0 ) =

H T −1 X X

2 nt+1 (sth , ath , h) ut (sth , ath , h) µ̂t+1 (sth , ath , h) − µ̂t (sth , ath , h)

t=T0 h=1

H T −1 X X

2 µ(sth , ath , h) − µ̂t (sth , ath , h) − ζht ut (sth , ath , h),

t=T0 h=1

where the second equality uses the identity from Lemma 22. We bound the two resulting terms separately. Since ut (s, a, h) ⩽ 1, Cauchy–Schwarz and the elliptic potential (Lemma 46) give, on     X EFµ (T ) ≜ ∀t ⩽ T, nt (s, a, h) KL(R̂ht (s, a) k Rh (s, a)) ⩽ β µ (T, 1/T 2 ) ,   s,a,h

the following T −1 X H X

v uX u t t t t µ − µ̂t (sh , ah , h) · ut (sh , ah , h) ⩽ t

sX 2 ut (sth , ath , h)2 · max (1, nt (sth , ath , h)) µ − µ̂t (sth , ath , h) max (1, nt (sth , ath , h)) t,h t,h p p ⩽ 4H log(T + 1) · 2T β µ (T, 1/T 2 ).

t=T0 h=1

The sequence (−ζht ut (sth , ath , h))t,h is a martingale difference sequence. By Azuma–Hoeffding’s inequality, the event   s −1 X H   TX X ut (sth , ath , h)2 −ζht+1 ut (sth , ath , h) ⩽ 2 log(T 2 )   t=T0 h=1

t,h

Optimal Posterior Sampling for Policy Identification in Tabular Markov Decision Processes

holds with probability at least 1 − 1/T 2 . Combining both bounds and using ut ⩽ 1 gives, on an event of probability at least 1 − 2/T 2 , p p DT (T0 ) ⩽ 2 8HT β µ (T, 1/T 2 ) log(T + 1) + 2 2HT log(T 2 ) = o(T ), which defines the event EFD (T ). F.6

Proof of Theorem 15

We can now prove Theorem 15 with the previous intermediate results. Let us define

EF (T ) = EE (T ) ∩ EFµ (T ) ∩ EFϕ (T ) ∩ EFD (T ) ∩ EFJ (T ) ∩ EFΨ (T ),

(F.24)

where EE (T ) is defined in Appendix E, and the other events are defined above. By Lemma 14, there exists k0 < ∞ such that for T ⩾ k0 , when EE (T ) holds, π̂t = π ⋆ , Setting T0 ≜ T 2/3 in Lemma 18, EE (T ) =⇒ EFT

2/3

∀t ∈ (T 2 /3, T ).

(T ), so that

T −1 X t=T

1 1 η 2/3 log ΛηTT + ΨTα (T 2/3 ) + log ΛTT2/3 + E(θ,Φ)∼ν ηt (·|Alt(M̂t )) ℓα t (θ, Φ) ⩽ − t η η 2/3 T T 2/3 T −1 X t=T 2/3

1 1 ηt+1 ηt log Λt+1 − log Λt+1 . ηt ηt+1

Thanks to Lemma 21, Lemma 20, and Lemma 19, and noting that EF (T ) holds with probability at least 1 − O(1/T 2 ) concludes the proof.

SADDLE-POINT CONVERGENCE

In this section, we prove the saddle-point convergence result that characterizes the optimality of the sampling rule in terms of posterior contraction. We work on the intersection of the following concentration events:     X p EG (T ) ≜ ∀ t ⩽ T : nt (s, a, h) KL(p̂t (· | s, a, h) k p(· | s, a, h)) ⩽ β p (T, 1/T 2 ) ,   s,a,h     X µ EG (T ) ≜ ∀ t ⩽ T : nt (s, a, h) KL R̂ht (s, a) k Rh (s, a) ⩽ β µ (T, 1/T 2 ) ,   s,a,h     X µ+p EG (T ) ≜ ∀ t ⩽ T : nt (s, a, h) KL R̂ht (s, a) ⊗ p̂th (· | s, a) k Rh (s, a) ⊗ p(· | s, a, h) ⩽ β µ+p (T, 1/T 2 ) .   s,a,h

(G.1) Explicit thresholds β p , β µ , β µ+p ensuring high-probability coverage follow from self-normalized concentration (see e.g. Kaufmann & Koolen 2021). Each event holds with probability at least 1 − 1/T 2 , so their intersection holds with probability at least 1 − 3/T 2 . We recall the log-likelihood ratio between a model (θ, Φ) and the empirical MDP M̂t : " # 2 ′ X X θ(s, a, h) − µ̂t (s, a, h) p̂ (s | s, a, h) t Υt (θ, Φ) ≜ nt (s, a, h) . + p̂t (s′ | s, a, h) log 2 Φ(s′ | s, a, h) ′ s,a,h

(G.2)

Cyrille Kone, Kevin Jamieson

The GLR (Generalized Likelihood Ratio) is defined as GLR(T ) ≜

inf

(G.3)

ΥT (θ, Φ)

(θ,Φ)∈Alt(M̂T )

The main result of this section is the following. Proposition 25. Let γ ∈ (0, 1), α ∈ (0, 12 ), and ς > 2α and ηt = O(t−ς ). There exists k0 < ∞ such that for any T > k0 , γ+1 p GLR(T ) ⩾ T · ΓM − O(log T ) − O(T 1+α−γ ) − o(T ) − O T α+ 2 log T , holds with probability at least 1 − O(1/T 2 ). Proof. We note that ≜

GLR(T )

inf

ΥT (θ, Φ) ,

(θ,Φ)∈Alt(M̂T )

inf (θ,Φ)∈Alt(M̂T )

s,a,h

(θh (s, a) − µ̂T (s, a, h))2 X 1 nT (s, a, h) + p̂T (s′ | s, a, h) log ′ | s, a, h) 2 Φ(s ′

′

nT (s, a)p̂T (s | s, a, h) log p̂T (s | s, a, h) .

s,a,h s′

We recall that in the equation above, for unobserved transitions, 0 · log 0 = 0. Next, since nT (s, a, h)p̂T (s′ | s, a, h) = nT (s′ | s, a, h) (and this is always true, even at initialization), we have XX

nT (s, a, h)p̂T (s′ | s, a, h) log

s,a,h s′

1 Φ(s′ | s, a, h)

X X

t∈[T −1] h

log

1 . Φ(sth+1 | sth , ath , h)

Similarly we have X s,a,h

θ(s, a, h) − µ̂T (s, a, h) nT (s, a, h) 2

2 =

X X (θ(st , at , h) − µ̂T (st , at , h))2 h h h h . 2

t∈[T −1] h

Combining these yields " # X X (θ(st , at , h) − µ̂T (st , at , h))2 X X 1 h h h h ΥT (θ, Φ) = + log + log p̂T (sth+1 | sth , ath , h) . t t t 2 Φ(sh+1 | sh , ah , h) t∈[T −1]

t∈[T −1] h

Moreover, the above is well-defined since ∀t < T, pT (sth+1 | sth , ath , h) > 0. Next, we invoke Theorem 15 using the regret properties of the conditional posterior sampling algorithm, we have with probability larger than 1 − O(1/T 2 ), # " X X θ(sth , ath , h) − µ̂T (sth , ath , h) 2 1 inf + log ⩾ 2 Φ(sth+1 | sth , ath , h) (θ,Φ)∈Alt(M̂T ) t∈[T −1] h (G.4) T −1 X α E(θ,Φ)∼ν ηt (·|Alt(M̂t )) ℓt (θ, Φ) − o(T ) , t

t=T 2/3

which yields GLR(T ) ⩾

T −1 X

E(θ,Φ)∼ν ηt (·|Alt(M̂t )) [ℓα t (θ, Φ)] + t

t=1

X X t∈[T −1] h

log p̂T (sth+1 | sth , ath , h) − o(T ),

(G.5)

Optimal Posterior Sampling for Policy Identification in Tabular Markov Decision Processes 1 1 α 2/3 where, since ℓα terms of the sum contribute for O(T t ∈ (0, H( 2 + t )) and α < 2 , the first T

2+2α 3

) = o(T ).

By Lemma 31, with probability at least 1 − O(1/T 2 ), we have X

ℓα t (θt , Φt ) ⩽

t∈[T −1]

v u T X α u 2 , E(θ,Φ)∼ν ηt (·|Alt(M̂t )) ℓt (θ, Φ) + t2 log(T 2 ) Ct,α

(G.6)

t=1

ft ≜ (θt , Φt ) is the challenger MDP sampled at time t in the pseudocode with Ct,α = H( 12 + tα ). We recall that M of Algorithm 1. Substituting (G.6) into (G.5) and recalling the definition of ℓα t from (F.9), we obtain H X X

GLR(T ) ⩾

µ̂t (sth , ath , h) − θt (sth , ath , h) 2

t∈[T −1] h=1 H X X

µ̂t (sth , ath , h) − θt (sth , ath , h) 2

t∈[T −1] h=1 H X X

1 + log t max Φt (sh+1 | sth , ath , h), e−tα

log p̂T (sth+1 | sth , ath , h) − o(T )

t∈[T −1] h=1 H X X

log

t∈[T −1] h=1

# max p̂t (sth+1 | sth , ath , h), ε + log max Φt (sth+1 | sth , ath , h), e−tα

p̂T (sth+1 | sth , ath , h) − o(T ), max p̂t (sth+1 | sth , ath , h), ε

(G.7)

where the equality rearranges the logarithms by adding and subtracting max p̂t (sth+1 | sth , ath , h), ε . Thanks to Lemma 29, with probability larger than 1 − O(1/T 2 ). # 2 X µ̂t (s, a, h) − θt (s, a, h) max (p̂t (s′ | s, a, h), ε) ′ p(s | s, a, h) log + 2 max (Φt (s′ | s, a, h), e−tα ) t=1 s,a,h s′ " # 2 H T −1 X X max p̂t (sth+1 | sth , ath , h), ε µ̂t (sth , ath , h) − θt (sth , ath , h) ⩽ − + log 2 max Φt (sth+1 | sth , ath , h), e−tα t=1 h=1 v u T −1 u X t2 H 2 (tα + 12 )2 log(T 2 ) T −1 X X

π̃ exp wht (s, a)

(G.8)

t=1

Combining (G.7) with (G.8) gives

GLR(T ) ⩾

T −1 X X t=1 s,a,h

" π̃ exp wht (s, a)

µ̂t (s, a, h) − θt (s, a, h) 2

max (p̂t (s′ | s, a, h), ε) + p(s′ | s, a, h) log max (Φt (s′ | s, a, h), e−tα ) ′

T −1 X H X

p̂T (sth+1 | sth , ath , h) − o(T ). + log max p̂t (sth+1 | sth , ath , h), ε t=1 h=1 (G.9)

Next, we relate the true transition p(s′ | s, a, h) to the empirical p̂t (s′ | s, a, h) in the logarithm above.

Cyrille Kone, Kevin Jamieson

We then invoke Lemma 26, to show that with probability at least 1 − O(1/T 2 ), T −1 X X

π̃ exp

wht (s, a)

p(s′ | s, a, h) log

s′

t=1 s,a,h T −1 X X

π̃ exp wht (s, a)

p max (p̂t (s′ | s, a, h), ε) α+ γ+1 2 p̂t (s | s, a, h) log − O T log T − o(T ) max (Φt (s′ | s, a, h), e−tα ) ′

t=1 s,a,h

max (p̂t (s′ | s, a, h), ε) ⩾ max (Φt (s′ | s, a, h), e−tα ) (G.10)

′

Substituting (G.10) into (G.9) and noting that p̂t (s′ | s, a, h) log max (p̂t (s′ | s, a, h), ε) ⩾ p̂t (s′ | s, a, h) log p̂t (s′ | s, a, h) gives GLR(T ) ⩾

T X X

" t whρ̃ (s, a)

t=1 s,a,h

T X H X

log

t=1 h=1

µ̂t (s, a, h) − θt (s, a, h) 2

max (p̂t (s′ | s, a, h), ε) + p̂t (s | s, a, h) log max (Φt (s′ | s, a, h), e−tα ) ′

′

p̂T (sth+1 | sth , ath , h) γ+1 p − O T α+ 2 log T − o(T ). t t t max p̂t (sh+1 | sh , ah , h), ε

(G.11)

Recalling the definition of the reward kernel rtexp (s, a, h) ≜

µ̂t (s, a, h) − θt (s, a, h) 2

2 +

p̂t (s, a, s′ , h) log

s′

p̂t (s′ | s, a, h) , max (Φt (s′ | s, a, h, e−tα )

(G.12)

and substituting into (G.11), we obtain, with probability 1 − O(1/T 2 ) GLR(T ) ⩾

T X X

π̃ exp

wht (s, a) rtexp (s, a, h) +

t=1 s,a,h

T X H X

log

t=1 h=1

p p̂T (sth+1 | sth , ath , h) α+ γ+1 2 log T − o(T ). − O T max p̂t (sth+1 | sth , ath , h), ε (G.13)

Decomposing the mixed policy π̃texp = (1 − εt )πtexp + εt ct , T −1 X X

π̃ exp wht (s, a) rtexp (s, a, h) =

t=1 s,a,h

T −1 X X

π exp wht (s, a) rtexp (s, a, h) +

t=1 s,a,h

T −1 X X

π exp εt whct (s, a) − wht (s, a) rtexp (s, a, h).

t=1 s,a,h

(G.14) By Lemma 27, the correction term satisfies T −1 X X

π exp εt whct (s, a) − wht (s, a) rtexp (s, a, h) ⩽

t=1 s,a,h

2 T 1+α−γ . 1+α−γ

(G.15)

Substituting (G.14) and (G.15) into (G.13), GLR(T ) ⩾

T X X

π exp

wht (s, a) rtexp (s, a, h)+

t=1 s,a,h

T −1 X H X t=1 h=1

log

p̂T (sth+1 | sth , ath , h) γ+1 p log T −o(T ). −O T 1+α−γ −O T α+ 2 t t t p̂t (sh+1 | sh , ah , h) (G.16)

By Lemma 30, with probability larger than 1 − O(1/T 2 ), T −1 X H X t=1 h=1

log

p p̂T (sth+1 | sth , ath , h) α+ γ+1 2 ⩽ O T log T /ε max p̂t (sth+1 | sth , ath , h), ε

(G.17)

Optimal Posterior Sampling for Policy Identification in Tabular Markov Decision Processes

Combining (G.16) and (G.17) gives GLR(T ) ⩾

T −1 X X

γ+1 p π exp wht (s, a) rtexp (s, a, h) − O T 1+α−γ − O T α+ 2 log T − o(T ).

(G.18)

t=1 s,a,h

Using the no-regret property of the online RL learner (Proposition 10), we have for α ∈ (0, 21 ) GLR(T )

⩾

sup w∈ΩM

γ+1 p log T − o(T ). wh (s, a)rtexp (s, a, h) − O(T 1+α−γ ) − O T α+ 2

t∈[T −1] s,a,h

Next, we invoke Lemma 28 to relate the empirical quantities in rtexp to their actual values. Indeed Lemma 28 shows that on an event that holds with probability 1 − O(1/T 2 ) T −1 X X

wh (s, a) rtexp (s, a, h) ⩾

t=1 s,a,h T −1 X X

" wh (s, a)

t=1 s,a,h

θt (s, a, h) − µ(s, a, h) 2

# p(s′ | s, a, h) ′ + p(s | s, a, h) log + o(T ) max (Φt (s′ | s, a, h), e−tα ) ′ X

(G.19)

Therefore, GLR(T )

⩾

T −1 X X

θt (s, a, h) − µ(s, a, h) sup wh (s, a) 2 w∈ΩM t=1 s,a,h γ+1 p −O(T 1+α−γ ) − o(T ) − O T α+ 2 log T

p(s′ | s, a, h) p(s | s, a, h) log + max (Φt (s′ | s, a, h), e−tα ) ′ ′

The final step to conclude the proof invokes Lemma 32 to prove that there exists k0 < ∞ such that for T ⩾ k0 , we have for any w ∈ ΩM " # 2 X X θt (s, a, h) − µ(s, a, h) p(s′ | s, a, h) ′ wh (s, a) + ⩾ p(s | s, a, h) log 2 max (Φt (s′ | s, a, h), e−tα ) s′ s,a,h " # 2 X X θ(s, a, h) − µ(s, a, h) p(s′ | s, a, h) ′ inf wh (s, a) + p(s | s, a, h) log − O(1/t). 2 Φ(s′ | s, a, h) (θ,Φ)∈Alt(M ) ′ s

s,a,h

We then have GLR(T )

⩾

T −1 X

(θ(s, a, h) − µ(s, a, h))2 X p(s′ | s, a, h) sup inf wh (s, a) + p(s′ | s, a, h) log 2 Φh (s, a, s′ ) w∈ΩM t=1 (θ,Φ)∈Alt(M ) s′ s,a,h γ+1 p −O(log T ) − O(T 1+α−γ ) − o(T ) − O T α+ 2 log T γ+1 p log T , T · ΓM − O(log T ) − O(T 1+α−γ ) − o(T ) − O T α+ 2

which concludes the proof.

G.1

Concentration Lemmas

We now prove the concentrations lemmas used in this section.

Cyrille Kone, Kevin Jamieson

Lemma 26. There exists k0 < ∞ such that for T ⩾ k0 , with probability 1 − O(1/T 2 ), the following holds: T X X

π̃ exp

wht (s, a)

p̂t (s′ | s, a, h) ⩾ max (Φt (s′ | s, a, h), e−tα )

p(s′ | s, a, h) log

s′

t=1 s,a,h T X X

π̃ exp

wht (s, a)

p̂t (s′ | s, a, h) log

s′

t=1 s,a,h

p p̂t (s′ | s, a, h) α+ γ+1 2 − O T log T − o(T ). α max (Φt (s′ | s, a, h), e−t ) ′

(s,a,s ,h) Proof. For fixed (t, s, a, s′ , h), letting Lt (s, a, s′ ) ≜ log max(Φp̂tt(s ′ |s,a,h),e−tα ) ,

p 2+2α t2α+γ−1 β p (t2 , 1/t3 ) . The first T 2/3 terms contributes for O(T 3 ) = π̃ exp

o(T ) (as α ∈ (0, 12 )). Summing over t, s, a, h and weighting by wht , the error term is at most γ+1 p O T α+ 2 β p (T 2 , 1/T 3 ) . The conclusion follows as β p (T 2 , 1/T 3 ) = O(log T ). Lemma 27. We have T X X

π exp εt whct (s, a) − wht (s, a)

rtexp (s, a, h) ⩽ O(T 1−γ+α ).

t=1 s,a,h

Proof. Since |rtexp (s, a, h)| ⩽ 12 + tα for all (s, a, h) and εt = t−γ , the triangle inequality gives T X X

T X X π exp π exp εt whct (s, a) − wht (s, a) rtexp (s, a, h) ⩽ εt whct (s, a) − wht (s, a) rtexp (s, a, h)

t=1 s,a,h

t=1

⩽2

s,a,h

1 t−γ tα + 2 t=1

T X

⩽ O(T 1−γ+α ), where the second inequality uses that integrals.

πtexp ct (s, a)| ⩽ 2, and the last step bounds the sums by s,a,h |wh (s, a) − wh

Lemma 28. There exists k0 < ∞ such that for T ⩾ k0 , for any valid state-action allocation w ∈ ΩM , the following , T X X

wh (s, a) rtexp (s, a, h) ⩾

t=1 s,a,h T X X t=1 s,a,h

" wh (s, a)

θt (s, a, h) − µ(s, a, h) 2

# 1+γ p p(s′ | s, a, h) α+ 2 + O T log T + p(s | s, a, h) log max (Φt (s′ | s, a, h), e−tα ) ′ X

′

(G.20) holds with probability larger than 1 − O(1/T 2 ).

Optimal Posterior Sampling for Policy Identification in Tabular Markov Decision Processes

Proof. Since wh (s, a) = 0 for any unreachable (s, a, h), we may restrict attention to reachable triplets throughout. For each (s, a, s′ , h), define Ct (s, a, s′ , h) ≜ p̂t (s′ | s, a, h) log

Applying a similar resaoning to the rewards (bounded in (0, 1)) yield θt (s, a, h) − µ̂t (s, a, h)

⩾ θt (s, a, h) − µ(s, a, h)

+ µ̂t (s, a, h) − µ(s, a, h)

− 2 µ̂t (s, a, h) − µ(s, a, h) ,

which, again, due to Lemma 12 yields µ̂t (s, a, h) − µ(s, a, h) ⩽ O

p tγ−1 β p (t2 , 1/t3 ) .

Summing over t, s, a, s′ , h and weighting by wh (s, a), T X X

wh (s, a) rtexp (s, a, h) ⩾

t=1 s,a,h T X X t=1 s,a,h

" wh (s, a)

θt (s, a, h) − µ(s, a, h) 2

Further noting that

tα O

# p X p(s′ | s, a, h) + − tα O p(s | s, a, h) log tγ−1 β p (t2 , 1/t3 ) . α ′ −t max (Φt (s | s, a, h), e ) ′ X

′

t⩽T

p 1+γ p tγ−1 β p (t2 , 1/t3 ) ⩽ O T α+ 2 log T

t⩽T

concludes the proof. Lemma 29. Fix ε ∈ (0, 1). The event # ( T " 2 X X X π̃exp µ̂t (s, a, h) − θt (s, a, h) max (p̂t (s′ | s, a, h), ε) ′ t p(s | s, a, h) log + Ew (T ) ≜ wh (s, a) 2 max (Φt (s′ | s, a, h), e−tα ) t=1 s,a,h s′ " # 2 T X H X max p̂t (sth+1 | sth , ath , h), ε µ̂t (sht , ath , h) − θt (sth , ath , h) ⩽ − (G.21) + log 2 max Φt (sth+1 | sth , ath , h), e−tα t=1 h=1 v ) u T u X t2 (G.22) H 2 (tα + 1 )2 log(T 2 ) 2

t=1

holds with probability at least 1 − 1/T 2 .

Cyrille Kone, Kevin Jamieson

Proof. Let us define, for each t and (s, a, s′ , h), µ̂t (s, a, h) − θt (s, a, h) Ψt (s, a, s , h) ≜ 2

′

+ log

max (p̂t (s′ | s, a, h), ε) , max (Φt (s′ | s, a, h), e−tα )

and note that |Ψt (s, a, s′ , h)| ⩽ tα + 12 . Since µ̂t is Ht−1 -measurable and (θt , Φt ) depends only on Ht−1 plus et−1 ≜ Ht−1 ∪ {θt , Φt }. Conditioned on additional independent randomization, we augment the filtration to H e Ht−1 , the only source of randomness is the trajectory at episode t, generated under π̃texp , so # "H H X X X π̃ exp t t t e Ψt (sh , ah , sh+1 , h) Ht−1 = E w t (s, a) p(s′ | s, a, h) Ψt (s, a, s′ , h). h

h=1 s,a,s′

h=1

Therefore the process

  t H H X X X X exp π̃  mt ≜ Ψk (skh , akh , skh+1 , h) − whk (s, a) p(s′ | s, a, h) Ψk (s, a, s′ , h) k=1

h=1

h=1 s,a,s′

is a martingale with, increments bounded in absolute value by H(tα + 21 ). By Azuma–Hoeffding’s maximal inequality, v   u T u X 1 H 2 (tα + 12 )2 log(T 2 ) ⩽ 2 , PrmT −1 > t2 T t=1 which gives the claimed event Ew (T ). Lemma 30. Let ε > 0. With probability at least 1 − O(1/T 2 ), 1+γ p p̂T (sth+1 | sth , ath , h) α+ 2 ⩽O T log T /ε . log max p̂t (sth+1 | sth , ath , h), ε t=1 h=1

T −1 X H X

Proof. Since p̂T (sth+1 | sth , ath , h) ⩾ 1/T > 0 (the transition is observed at episode t ⩽ T ) and max (p̂t , ε) ⩾ ε > 0, every logarithm is finite. By the triangle inequality, H T −1 X X p̂T (sth+1 | sth , ath , h) p̂T (sth+1 | sth , ath , h) . log ⩽ log t t , at , h), ε t t , at , h), ε max p̂ (s | s max p̂ (s | s t t h+1 h h h+1 h h t=1 h=1 t=1 h=1

H T −1 X X

For T0 = o(T ), the contribution from t < T0 is o(T ) since max (p̂t , ε) ⩾ ε, giving | log(p̂T /ε)| ⩽ log(1/ε) per term, times T0 . p Thanks to Lemma 12, there exists k0 < ∞, such that for T ⩾ k0 for all t ∈ (T 2/3 , T ) and all (s, a, h), if EG (T ) holds, then p kp̂t (· | s, a, h) − p(· | s, a, h)k1 ⩽ O tγ−1 β p (t2 , 1/t3 ) .

We assume t ⩾ T 2/3 . We split into two cases. In the first case, assume p̂t (sth+1 | sth , ath , h) ⩾ ε. Then max (p̂t , ε) = p̂t , and since γ < 21 both p̂T and p̂t lie in [(1 − tγ−1/2 )ph , (1 + tγ−1/2 )ph ], so p p 1+O tγ−1 β p (t2 , 1/t3 ) p̂T p ⩽O ⩽ log log tγ−1 β p (t2 , 1/t3 ) p̂t 1−O tγ−1 β p (t2 , 1/t3 ) In the second case, assuming p̂t (sth+1 | sth , ath , h) < ε, then max (p̂t , ε) = ε. Since p̂t < ε, the concentration bound gives p p p(sth+1 | sth , ath , h) ⩽ p̂t (sth+1 | sth , ath , h) + O tγ−1 β p (t2 , 1/t3 ) < ε + O tγ−1 β p (t2 , 1/t3 ) .

Optimal Posterior Sampling for Policy Identification in Tabular Markov Decision Processes

Applying the concentration bound again to p̂T , yields p̂T (sth+1 | sth , ath , h) ⩽ p(sth+1 | sth , ath , h) + O

p p T γ−1 β p (T 2 , 1/T 3 ) ⩽ ε + O tγ−1 β p (t2 , 1/t3 ) .

Therefore, using log(1 + x) ⩽ x, log

p̂T (sth+1 | sth , ath , h) ⩽ pε ε+O tγ−1 β p (t2 , 1/t3 )

log O

ε tγ−1 β p (t2 , 1/t3 )

 = log1 +

tγ−1 β p (t2 , 1/t3 ) ε

 ⩽

tγ−1 β p (t2 , 1/t3 )/ε ,

where the last step uses that ε is a fixed positive constant. In both cases the contribution per (t, h) is O T −1 X H X t=T0 h=1

log

p tγ−1 β p (t2 , 1/t3 ) . Summing gives

1+γ p p̂T ⩽ O T α+ 2 log T /ε max (p̂t , ε)

Lemma 31. The event v   T   X h i u u X conv α t2 log(T 2 ) 2 η C EG (T ) ≜ ℓα (θ , Φ ) − E ℓ (θ, Φ) ⩽ t t t t,α t (θ,Φ)∼νt (·|Alt(M̂t )) t   t∈[T −1]

(G.23)

t=1

1 2 α holds with probability at least 1 − 1/T , where C ≜ H + t . t,α 2 p 1 O T 1+2α log T , which is o(T ) whenever α < 2 .

In particular, the right-hand side is

α Proof. Recalling the definition of ℓα t from (F.9), its increments satisfy ℓt (θ, Φ) ∈ (0, Ct,α ) for all (θ, Φ). Condiα tioned on Ht−1 , the only source of randomness in ℓt (θt , Φt ) is the sample (θt , Φt ) ∼ νtηt (· | Alt(M̂t )), so

α E[ℓα t (θt , Φt ) | Ht−1 ] = E(θ,Φ)∼ν ηt (·|Alt(M̂t )) ℓt (θ, Φ) . t

Therefore the process mt ≜

t h X

α i ηk ℓα (θ , Φ ) − E ℓ (θ, Φ) k k k (θ,Φ)∼ν (·|Alt(M̂k )) k k

k=1

is a martingale with increments bounded in (−Ct,α , Ct,α ). By Azuma–Hoeffding’s maximal inequality, v  u T u X 2 ⩽ 1 , PrmT −1 > t2 log(T 2 ) Ct,α T2 t=1 

which gives the claimed event. The result below proves that the clipping of the transition probabilities on the sampled instance has a negligible impact on the loss of the regret minimizer.

Cyrille Kone, Kevin Jamieson

Lemma 32. There exists a constant k0 < ∞ such that for all t ⩾ k0 " # 2 X X θt (s, a, h) − µ(s, a, h) p(s′ | s, a, h) ′ wh (s, a) + p(s | s, a, h) log ⩾ 2 max (Φt (s′ | s, a, h), e−tα ) s′ s,a,h " # X (θ(s, a, h) − µ(s, a, h))2 X p(s′ | s, a, h) ′ inf wh (s, a) + p(s | s, a, h) log − O(1/t), 2 Φh (s′ | s, a, h) f≜(θ,Φ)∈Alt(M ) M ′ s

s,a,h

holds with probability 1. Proof. To ease notation, we introduce " # 2 ′ X X (θ (s, a, h) − µ(s, a, h)) p(s | s, a, h) t dα wh (s, a) + p(s′ | s, a, h) log . w ((θt , Φt ), (µ, p)) ≜ 2 max (Φt (s′ | s, a, h), e−tα ) ′ s

s,a,h

ft ≜ (θt , Φt ) ∈ Alt(M ) there exists a deterministic policy π̃ 6= π ⋆ such that π̃ is a global optimal policy Since M ft . in M We invoke Lemma 43, which provides an MDP M ′ ≜ (p̃, µ̃) for which π̃ is the only global optimal policy and there is a unique optimal action per stage and state. Assuming 1t ⩽ 13 1+21H−1 i.e., t ⩾ 3(1 + 2H−1 ) then we have by the construction in Lemma 43 with ε = 1/t, p̃h (· | s, a) ≜ (1 − γs,a,h )Φt (· | s, a, h) + γs,a,h Φt (· | s, πh (s), h) |µ̃(s, a, h) − θt (s, a, h)| ⩽ ds,a,h ≜ 2H−h+1 /t h i (1+2H−h )ε where γs,a,h ∈ 0, 1−(2 H−h+1 +3)ε . Therefore, we have

dα w ((µ̃, p̃), (µ, p))

⩽ ⩽

p(s′ | s, a, h) (µ̃(s, a, h) − µ(s, a, h))2 X + p(s′ | s, a, h) log wh (s, a) 2 max ((1 − γh )Φt (s′ | s, a, h), e−tα ) s′ s,a,h H−h+1 X 4 2H−h+1 1 α dw ((θt , Φt ), (µ, p)) + wh (s, a) + log 2t2 t 1 − γh X

s,a,h

⩽ ⩽

2H γ1 4H + + 2 2t t 1 − γ1 dα w ((θt , Φt ), (µ, p)) + O(1/t).

dα w ((θt , Φt ), (µ, p)) +

(by monotonicity of h 7→ γh and log(1 + x) ⩽ x)

Now we invoke Lemma 45 to build an instance Mγ′ ≜ (µ̃, tildepγ ) close to (µ̃, p̃), but for which π̃ is still the unique optimal policy and such that each transition probability is larger than some γ. We invoke Lemma 45 with γ = 1/t. This instance is obtained by a mixture of (µ̃, p̃) with (µ̃, u) with weights 1 − γ, γ, where u is the uniform transition kernel. We then have thanks to log(1 + x) ⩽ x, " # X (µ̃(s, a, h) − µ(s, a, h))2 X p(s′ | s, a, h) α γ ′ dw ((µ̃, p̃ ), (µ, p)) ⩽ wh (s, a) + p(s | s, a, h) log 2 max ((1 − γ)p̃(s′ | s, a, h), e−tα ) ′ s

s,a,h

⩽ = =

γ 1−γ 1/t dα w ((µ̃, p̃), (µ, p)) + 1 − 1/t dα ((µ̃, p̃), (µ, p)) + O(1/t) . w dα w ((µ̃, p̃), (µ, p)) +

Thus, combining the two displays above, we have α γ dα w ((θt , Φt ), (µ, p)) ⩾ dw ((µ̃, p̃ ), (µ, p)) − O(1/t).

Optimal Posterior Sampling for Policy Identification in Tabular Markov Decision Processes

However, the transition probabilities in p̃γ are larger than γ/S since p̃γ is the mixture with the uniform transition α kernel with weight γ. Thus, assuming t is such that 1/(tS) > e−t , we have for all (s, a, h, s′ ) p̃γ (s′ | s, a, h) ⩾ Therefore, γ dα w ((µ̃, p̃ ), (µ, p))

X s,a,h

α 1 ⩾ e−t tS

(θh (s, a) − µ(s, a, h))2 X p(s′ | s, a, h) + p(s′ | s, a, h) log γ ′ wh (s, a) 2 p̃ (s | s, a, h) ′

# .

Thus, since Lemma 45 ensures that π̃ 6= π ⋆ is the unique optimal policy for each stage and state, we have (µ̃, p̃γ ) ∈ Alt(M ). Putting the above results together, we have " # 2 ′ X X (θ(s, a, h) − µ(s, a, h)) p(s | s, a, h) dα wh (s, a) + p(s′ | s, a, h) log γ ′ − O(1/t) w ((θt , Φt ), (µ, p)) ⩾ 2 p̃ (s | s, a, h) ′ s s,a,h " # ′ X (θ(s, a, h) − µ(s, a, h))2 X p(s | s, a, h) ⩾ inf wh (s, a) + p(s′ | s, a, h) log − O(1/t). 2 Φ(s′ | s, a, h) (θ,Φ)∈Alt(M ) ′ s

s,a,h

H POSTERIOR CONVERGENCE In this section, we prove the results related to the posterior contraction rate. First, we prove the lower bound on the posterior anti-concentration rate. More precisely, we prove the result below. Theorem 2. Fix an MDP M ∈ M with optimal policy π ⋆ and let νt denote the posterior distribution defined in Eq. (1). For any adaptive algorithm for best policy identification, with probability one, ⋆ f) ⩽ ΓM . / Π⋆ ( M lim sup − 1t log PrM f∼νt | Ht−1 π ∈ t→∞

Let (πtexp )t⩾1 be the sequence of exploration policies used by an adaptive algorithm, i.e., at episode t, the learner selects a policy πtexp depending on past observations summarized in Ht−1 . We that recall Υt (θ, Φ) ≜

X s,a,h

" nt (s, a, h)

µ̂t (s, a, h) − θ(s, a, h) 2

p̂t (s′ | s, a, h) + p̂t (s | s, a, h) log Φ(s′ | s, a, h) ′

′

(H.1)

f ≜ (θ, Φ) ∈ M. is the log-likelihood ratio between the empirical MDP M̂t ≜ (µ̂t , p̂t ) and the MDP M H.1

Asymptotic Limit of Likelihood Ratio

Before proving the result above, we show the following intermediate proposition. Proposition 33. For any adaptive algorithm for best policy identification, with probability one, 1 lim sup inf Υt (θ, Φ) ⩽ ΓM . t (θ,Φ)∈Alt(M ) t→∞ Proof. First, let us justify that the empirical frequency vector ( nt (s,a,h) )s,a,h is an almost valid state–action t−1 allocation for M . For this, let us define the t-indexed stochastic process mt (s, a, h)

≜

t h X

i π exp 1 skh = s, akh = a − whk (s, a) ,

k=1

nt+1 (s, a, h) −

X k⩽t

π exp

whk (s, a) ,

Cyrille Kone, Kevin Jamieson exp

where wπt is the state–action visitation measure induced by πtexp on M . Note that (mt (s, a, h))t is a martingale and its increments are bounded in (−1, 1). Thanks to Azuma-Hoeffding, the event

s,a,h EH (T ) ≜

  

∀ t ⩽ T, nt+1 (s, a, h) −

exp πk

(s, a) ⩽

  2T log(2T 2 SAH)

k⩽t

(H.2)



holds with probability at least 1 − 1/(SAHT 2 ) so that EH (T ) ≜

s,a,h EH (T ),

(H.3)

s,a,h

holds with probability at least 1 − 1/T 2 . Therefore, thanks to Borell-Cantelli’s lemma, there exits Te such that with probability 1, for T ⩾ Te, EH (T ) holds which implies that for any s, a, h and T ⩾ Te, 1 nT +1 (s,a,h) − T T

r π exp wht (s, a)

⩽

t⩽T

2 log(2T 2 SAH) , T

(H.4)

and further remark that, as the space of valid-action allocation is convex (see e.g.,A. J. Wagenmaker, Chen, et al. P π exp 2022), T1 t∈[T ] wht is a valid state-action allocation for M . Since transition probabilities may be arbitrarily small, the KL divergence in Υt is potentially unbounded. To handle this, we restrict to alternatives with bounded transitions. Concretely, observe that the model obtained by setting all transitions to uniform and all rewards to zero on states visited by the optimal policy of M lies in Alt(M ), as M is trivially absolutely continuous with respect to it. Therefore, for any ε > log S, the restricted alternative set Alt(M, ε) ≜ (θ, Φ) ∈ Alt(M ) : ∀ (s, a, h, s′ ), Φ(s′ | s, a, h) ⩾ e−ε is non-empty. Since Alt(M, ε) ⊂ Alt(M ), inf (θ,Φ)∈Alt(M )

Υt (θ, Φ) ⩽ X

inf (θ,Φ)∈Alt(M,T 1/4 )

" nT (s, a, h)

s,a,h

µ̂T (s, a, h) − θ(s, a, h) 2

# p̂T (s′ | s, a, h) + . p̂T (s | s, a, h) log Φ(s′ | s, a, h) ′ X

′

(H.5)

By definition of Alt(M, T 1/4 ), every (θ, Φ) in this set satisfies ∀ (s, a, h, s′ ) :

log

1 Φ(s′ | s, a, h)

⩽ T 1/4 ,

(H.6)

which allows concentration arguments to be applied in the next step. µ+p µ p We define the event EH (T ) ≜ EH (T ) ∩ EH (T ), where

p EH (T ) ≜

µ EH (T ) ≜

     

∀t ⩽ T :

nt (s, a, h) KL(p̂t (· | s, a, h) k p(· | s, a, h)) ⩽ β p (T, 1/T 2 )

s,a,h

∀t ⩽ T :

X s,a,h

nt (s, a, h) KL R̂ht (s, a) k Rh (s, a) ⩽ β µ (T, 1/T 2 )

  

(H.7)

(H.8)

Optimal Posterior Sampling for Policy Identification in Tabular Markov Decision Processes

On E µ (T ) we have, µ̂T (s, a, h) − θ(s, a, h) 2

(µ(s, a, h) − θ(s, a, h))2 (µ(s, a, h) − µ̂T (s, a, h))2 + + |µ(s, a, h) − µ̂T (s, a, h)| 2 2 s 2 µ(s, a, h) − θ(s, a, h) β µ (T, 1/T 2 ) 2β µ (T, 1/T 2 ) + + . ⩽ 2 nT (s, a, h) nT (s, a, h) (H.9) ⩽

Next, for any q ∈ 4, we let ST+ ≜ {s′ : p̂T (s′ | s, a, h) > 0}. For s′ ∈ / ST+ , the convention 0 log 0 = 0 means the left-hand side contributes 0. By convexity of x 7→ x log x, p(s′ | s, a, h) p̂T (s′ | s, a, h) − p(s′ | s, a, h) log ⩽ p̂T (s′ | s, a, h) log ′ q(s ) q(s′ ) p̂T (s′ | s, a, h) p̂T (s′ | s, a, h) − p(s′ | s, a, h) , s′ ∈ ST+ . 1 + log ′ q(s ) Summing over s′ ∈ ST+ yields X

p̂T (s′ | s, a, h) log

s′

p(s′ | s, a, h) log

s′

p̂T (s′ | s, a, h) ⩽ q(s′ )

X p(s′ | s, a, h) p̂T (s′ | s, a, h) + 1 + log p̂T (s′ | s, a, h) − p(s′ | s, a, h) . ′ ′ q(s ) q(s ) + ′

(H.10)

s ∈ST

We note that for s′ ∈ ST+ , either the transition s′ has been observed at least once from (s, a, h), or (s, a, h) has never been visited and p̂T (s′ | s, a, h) = 1/S, by initialization. In all cases p̂T (s′ | s, a, h) ⩾ min(1/T, 1/S). Therefore, p̂T (s′ | s, a, h) 1 + log ⩽ 1 + log(ST ). q(s′ ) Substituting into (H.10) and using the triangle inequality, X

p̂T (s′ | s, a, h) log

s′

p Applying Pinsker’s inequality on EH (T ),

X s

p̂T (s′ | s, a, h) X p(s′ | s, a, h) p̂T (s | s, a, h) log ⩽ p(s′ | s, a, h) log + 1 + log(ST ) ′ ′ q(s ) q(s ) ′ ′ ′

2β p (T, 1/T 2 ) . (H.12) nT (s, a, h)

Substituting (H.9) and (H.12) into (H.5), applying Cauchy–Schwarz to the error terms, and using P n (s, a, h) = T H, T s,a,h inf (θ,Φ)∈Alt(M )

Υt (θ, Φ) ⩽ X

p(s′ | s, a, h) inf p(s′ | s, a, h) log nT (s, a, h) + Φ(s′ | s, a, h) (θ,Φ)∈Alt(M,T 1/4 ) s′ s,a,h p p SAHβ µ (T, 1/T 2 ) + 2SAH 2 T β µ (T, 1/T 2 ) + 1 + log(ST ) 2SAH 2 T β p (T, 1/T 2 ). µ(s, a, h) − θ(s, a, h) 2

Since β µ and β p are at most logarithmic in T , the three error terms in (H.13) are all o(T ).

# +

(H.13)

Cyrille Kone, Kevin Jamieson

We then invoke Eq. (H.2), which proves that on the event EH (T ) we have " # X (µ(s, a, h) − θ(s, a, h))2 X p(s′ | s, a, h) ′ nT (s, a, h) + ⩽ inf p(s | s, a, h) log 2 Φ(s′ | s, a, h) (θ,Φ)∈Alt(M,T 1/4 ) s′ s,a,h # " X X πexp p(s′ | s, a, h) (µ(s, a, h) − θ(s, a, h))2 X ′ t inf wh (s, a) + p(s | s, a, h) log 2 Φ(s′ | s, a, h) (θ,Φ)∈Alt(M,T 1/4 ) ′ s

s,a,h t∈[T ]

+ SAH

1 2T log(2T 2 SAH) + T 1/4 , 2

Combining with Eq. (H.13) yields inf (θ,Φ)∈Alt(M )

Υt (θ, Φ) ⩽ X

inf (θ,Φ)∈Alt(M,T 1/4 )

π exp wht (s, a)

s,a,h t∈[T −1]

# (µh (s, a) − θh (s, a))2 X p(s′ | s, a, h) ′ + p(s | s, a, h) log + 2 Φ(s′ | s, a, h) ′

(1 + log(ST )) 2SAH 2 T β p (T, 1/T 2 ) h i p p 1 + SAHβ µ (T, 1/T 2 ) + 2SAH 2 T β µ (T, 1/T 2 ) + SAH 2T log(2T 2 SAH) + T 1/4 . 2

Next, noting that β µ , β p ’s dependency on T is at most logarithmic (see e.g. Tirinzoni, Al-Marjani, et al. 2023), the second order term in the equation above is o(T ), which we rewrite as: on the event EH (T ),

inf (θ,Φ)∈Alt(M )

Υt (θ, Φ) ⩽ X

inf (θ,Φ)∈Alt(M,T 1/4 )

" π exp wht (s, a)

s,a,h t∈[T −1]

# 2 X µ(s, a, h) − θ(s, a, h) p(s′ | s, a, h) ′ + + o(T ) p(s | s, a, h) log 2 Φ(s′ | s, a, h) s′ " 2 X µ(s, a, h) − θ(s, a, h) πtexp wh (s, a) + 2

1 T −1 s,a,h t∈[T −1] # ′ X p(s | s, a, h) p(s′ | s, a, h) log + o(T ) Φh (s, a, s′ ) ′

= (T − 1) ·

inf

(θ,Φ)∈Alt(M,T 1/4 )

⩽ (T − 1) · sup

inf

w∈ΩM (θ,Φ)∈Alt(M,T 1/4 )

s,a,h

# (µ(s, a, h) − θ(s, a, h))2 X p(s′ | s, a, h) ′ wh (s, a) p(s | s, a, h) log + + o(T ) 2 Φ(s′ | s, a, h) ′ s

1 which follows as, from the convexity of ΩM we have T −1

PT −1 t=1

π exp

w· t

∈ ΩM . Finally, remark that the sequence

(uT )T ⩾1 defined by uT ≜ sup

inf

w∈ΩM (θ,Φ)∈Alt(M,T 1/4 )

X s,a,h

" wh (s, a)

µ(s, a, h) − θ(s, a, h) 2

p(s′ | s, a, h) p(s | s, a, h) log + Φ(s′ | s, a, h) ′

′

µ+p is non-increasing and inf {uT : T ⩾ 1} = ΓM . Thus, noting Pr(EH (T ) ∩ EH (T )) ⩾ 1 − O(1/T 2 ), thanks to Borel-Cantelli’s lemma, we have with probability one, 1 lim sup inf ΥT (θ, Φ) ⩽ ΓM . T →∞ T (θ,Φ)∈Alt(M )

Optimal Posterior Sampling for Policy Identification in Tabular Markov Decision Processes

H.2

Proof of Theorem 2

Proof. The result follows from Proposition 33, Lemma 19, and an anti-concentration bound on the posterior. We proceed in two steps: first expressing the posterior probability of Alt(M ) in terms of the GLR, then lowerbounding it. On can verify that the posterior probability of Alt(M ) is proportional to Z f Pr M ∈ Alt(M ) ∝ e−Υt (θ,Φ) dθ dΦ. f∼νt |Ht−1 M

(H.14)

Alt(M )

Defining

e−Υt (θ,Φ) dθ dΦ,

Λt ≜

(H.15)

Alt(M )

we obtain

f∼νt |Ht−1 M

f ∈ Alt(M ) = Z M

Λt e−Υt (θ,Φ) dθ dΦ

(H.16)

Since Υt is non-negative (it is a sum of KL divergences and squared differences) and M is bounded, Z e−Υt (θ,Φ) dθ dΦ ⩽ vol(M). M

Substituting into (H.16),

f∼νt |Ht−1 M

f ∈ Alt(M ) ⩾ M

Λt . vol(M)

(H.17)

By Lemma 19 (with η = 1), with probability one, "

# 2 p̂t (skh+1 | skh , akh , h) µ̂t (skh , akh , h) − θ(skh , akh , h) log Λt ⩾ − inf + log − o(t) 2 (θ,Φ)∈Alt(M ) Φ(skh+1 | skh , akh , h) k=1 h=1 # " 2 X X µ̂t (s, a, h) − θ(s, a, h) p̂t (s′ | s, a, h) ′ + p̂t (s | s, a, h) log − o(t). =− inf nt (s, a, h) 2 Φ(s′ | s, a, h) (θ,Φ)∈Alt(M ) ′ t−1 X H X

s,a,h

(H.18)

Combining (H.17) and (H.18), − log

f∼νt |Ht M

f ∈ Alt(M ) ⩽ M

inf (θ,Φ)∈Alt(M )

Υt (θ, Φ) + log vol(M) − o(t).

Since log vol(M) = O(1), dividing by t and taking the lim sup, Proposition 33 gives, with probability one, 1 f ∈ Alt(M ) ⩽ lim sup 1 lim sup − log Pr M inf Υt (θ, Φ) ⩽ ΓM , t f∼νt |Ht−1 (θ,Φ)∈Alt(M ) t→∞ t→∞ t M which completes the proof of Theorem 2. We now prove the main result of this section, showing that PIPS reaches the optimal posterior contraction rate with probability one.

Cyrille Kone, Kevin Jamieson

Thanks to Lemma 36 we have for − GLR(t)) f PrM , f∼νt |Ht−1 (M ∈ Alt(M )) ⩽ h(t)e

where h is polynomial in t, where due to Lemma 12, we will have Alt(M̂t ) = Alt(M ) (up to a zero measure set), for t large enough. Therefore, with probability 1 − O(1/t2 ), 1+α−γ f log PrM ) + o(t) + O(tα+ 2 f∼νt |Ht−1 (M ∈ Alt(M )) ⩽ −tΓM + log h(t) + O(log t) + O(t

γ+1

log t),

then, thanks to Borel-Cantelli’s lemma, with probability one, 1 f lim inf − log PrM f∼νt |Ht−1 (M ∈ Alt(M )) ⩾ ΓM , t→∞ t which combined with Theorem 2 completes the proof.

SAMPLE COMPLEXITY GUARANTEES

In this section, we establish sample complexity guarantees regardless of the algorithm’s correctness for the specified threshold, showing that when B(t, δ) is correctly calibrated, the posterior sampling rule combined with the sampling rule of PIPS attains the optimal expected sample complexity. B(t,δ) Theorem 4. Let γ ∈ (0, 1), α ∈ (0, 1/2), α < γ, ς ∈ (0, 1), with ς > 2α. If B(t, δ) satisfies lim supδ→0 log ⩽ ηt log δ1 1, then, PIPS coupled with the stopping time τδ described in (5) and run with parameterss γ, α and learning rate ηt = O(t−ς ), satisfies

1 + o log 1δ , and EM τδ ⩽ Γ−1 M log δ τδ −1 PrM lim sup ⩽ Γ =1. M 1 δ→0 log δ Proof. Let (E(t))t⩾1 be some events that will be explicit below. We have EM [τδ ]

⩽ ⩽ ⩽

1+ 1+ 1+

∞ X t=1 ∞ X t=1 ∞ X

PrM (τδ > t) , ft,k )) , PrM (∃ k ⩽ B(t, δ) : π̂t ∈ / Π⋆ ( M ∞ h i X ft,k )|Ht−1 ) + EM 1(E(t)) PrM (∃ k ⩽ B(t, δ) : π̂t ∈ / Π⋆ ( M PrM (E(t)c ) ,

t=1

⩽

∞ X t=1

PrM (E(t)c ) +

∞ X

f) EM B(t, δ)1(E(t)) PrM / Π⋆ ( M f∼ν ηt |Ht−1 π̂t ∈ t

t=1

i ,

Optimal Posterior Sampling for Policy Identification in Tabular Markov Decision Processes

ft,1 , M ft,2 , . . . are i.i.d. samples from distribution with conditional which follows since conditionally on Ht−1 , M ηt density νt (·). Thus taking k0 such that π̂t = π ⋆ for any t ⩾ k0 , we have EM [τδ ]

⩽

1 + k0 +

∞ X

PrM (E(t)c ) +

t=1

∞ X

h i f η E B(t, δ)PrM M ∈ Alt(M ) . t f∼ν |Ht−1 t

t=1

Next, thanks to Lemma 36 we have −ηt GLR(t) f PrM , f∼ν ηt |Ht−1 (M ∈ Alt(M )) ⩽ f (t)e t

where f is at most polynomial in t. By Proposition 25, there exists an event E(t), which holds with probability 1 − O(1/t2 ) such that we have GLR(t) ⩾ ΓM · t − h(t) where h is a deterministic sublinear function of t satisfying h(t) = o(t). Combining the above results, we have f 1(E(t)) PrM f∼ν ηt |Ht−1 (M ∈ Alt(M )) ⩽ f (t) exp (−ΓM tηt + ηt h(t)) . t

Thus we define T (δ) ≜ sup t ⩾ 1 : − log(B(t, δ)) − log f (t) + ηt t · ΓM − ηt h(t) ⩽ log(t log(t)1+σ ) , so we have EM [τδ ]

⩽

T0 + T (δ) +

1 , 1+σ t log(t) t>2

(I.1)

where T0 < ∞ is a constant. We note that

1 <∞ 1+σ t log(t) t>2

for σ > 0. Therefore, it remains to control the quantity T (δ) and for this remark that T (δ) = sup t ⩾ 1 tΓM ηt − log(B(t, δ)) − log f (t) − ηt h(t) ⩽ log(t log(t)1+σ ) , 1 1 1 −1 1+σ = sup t ⩾ 1 t ⩽ ΓM log(B(t, δ)) + log f (t) + h(t) − log(t log(t) ) , ηt ηt ηt Now we remark in the above the leading term in δ dependency is the term in η1t log(B(t, δ)). Bounding T (δ) t. Thus,

First note that f is at most polynomial and h(t) is sublinear in t, as well as η1t is sublinear in q(t) ≜

1 1 1 log B(t, δ) + log f (t) + h(t) − log t log(t)1+σ ηt ηt ηt

is sublinear in t. Hence, there exists ε ∈ (0, 1) such that q(t) = ot→∞ (tε ). Next, observe that q(log(1/δ)1/ε ) = o(log(1/δ)). Thus, we have for tδ = log(1/δ)1/ε 1 ηtδ log(B(tδ , δ)) + q(tδ ) ⩽ 1. lim sup log 1δ δ→0

Cyrille Kone, Kevin Jamieson

Let us introduce b(tδ ) ≜

1 log(B(tδ , δ)) + q(tδ ) . ηtδ

Let δmin ∈ (0, 1) be defined as ≜

δmin

n o inf δ ∈ (0, 1) | b(tδ ) > log(1/δ)1/ε ΓM inf δ ∈ (0, 1) | Γ−1 M b(tδ ) > tδ ,

b(tδ ) 1/ε . For all t ⩾ Tmax , there which is well defined as lim supδ→0 log 1 ⩽ 1 and ε ∈ (0, 1). Let Tmax = log(1/δmin ) δ

exists (0, 1) 3 δ ′ ⩽ δmin such that btδ′ c = t and Γ−1 M b(tδ ′ ) < tδ ′ . Therefore, for all δ ⩽ δmin T (δ) ⩽ log(1/δ)1/ε .

(I.2)

Moreover, by definition T (δ) ⩽ Γ−1 M b(T (δ)) and b is increasing for small δ, it follows that 1/ε T (δ) ⩽ Γ−1 ). M · b(log(1/δ)

(I.3)

Combining (I.1), (I.2) and (I.3), it follows that for all δ ⩽ δmin , EM [τδ ] Finally, since lim supδ→0 b(log(1/δ) log 1

1/ε

)

⩽

1/ε Γ−1 ) + O(1). M · b(log(1/δ)

(I.4)

⩽ 1, we conclude that lim sup δ→0

EM [τδ ] ⩽ Γ−1 M . log(1/δ)

Almost-sure upper bound The almost-sure bound on the sample complexity is derived from an application of Borel-Cantelli’s lemma. Indeed, introducing for any T the event ( T ) T X X p E(T ) ≜ 1(τδ > t) − Pr(τδ > t | Ht−1 ) ⩽ 2T log T 2 , t=1

t=1

we have when E(T ) holds min(τδ , T )

⩽

T X

1(τδ > t)

t=1

⩽

T X t=1

⩽

T (δ) +

Pr(τδ > t | Ht−1 ) + p

2T log T 2

2T log T 2 + T0 ,

where the last inequality follows from the derivations above on the event E(T ) and holds with probability larger than 1 − O(1/T 2 ). Next, assume E(T ) holds. Then, if n o p T ⩾ T ′ (δ) ≜ sup t ⩾ 1 : t < T0 + T (δ) + 2t log t2 , and E(T ) holds then the stopping rule would trigger before episode T . Thanks to Azuma-Hoeffding E(T ) holds with probability 1 − 1/T 2 for any T ⩾ 1. Thus E(T )c holds with probability at most 1/T 2 , similarly for P at least 1 c E(T ) and, as T ⩾1 T 2 < ∞, we invoke Borel-Cantelli’s lemma to justify that Pr(lim supT →∞ (E(T )∩E(T ))c ) = 0 so with probability 1, there exists Te (possibly random) such that for T ⩾ Te, E(T ) ∩ E(T ) holds. ′ From the calculations above, we remark that the leading term in T ′ (δ) is at most of order log(1/δ). Let δmin be ′ such that ∀ δ ⩽ δmin , dlog(1/δ)3/2 e > T ′ (δ),

Optimal Posterior Sampling for Policy Identification in Tabular Markov Decision Processes ′ which is well-defined. Let δ ⩽ δmin such that additionally dlog(1/δ)3/2 e > Te. We then have q τδ ⩽ T0 + T (δ) + 2dlog(1/δ)3/2 e logdlog(1/δ)3/2 e2 q ⩽ T0 + T (δ) + log(1/δ)3/4 8 log(1 + log(1/δ)3/2 ) .

Finally recalling that lim sup δ→0

T (δ) ⩽ Γ−1 M . log 1δ

We conclude by noting that τδ lim sup 1 δ→0 log δ

⩽

p T (δ) T0 + log(1/δ)3/4 8 log(1 + log(1/δ)3/2 ) lim sup lim 1 + δ→0 log 1δ δ→0 log δ

⩽

Γ−1 M .

CONCENTRATION RESULTS

J.1

Self-noramlized Concentration

We prove some concentration results used in other sections. We define the following concentration events of the self-normalized quantities related to the GLR     X E p (δ) ≜ ∀t ⩾ 1, nt (s, a, h) KL(p̂t (· | s, a, h) k p(· | s, a, h)) ⩽ β p (t, δ)   s,a,h     X E µ (δ) ≜ ∀t ⩾ 1, nt (s, a, h) KL(R̂ht (s, a) k Rh (s, a)) ⩽ β µ (t, δ)   s,a,h     X E µ+p (δ) ≜ ∀t ⩾ 1, nt (s, a, h) KL(R̂ht (s, a) ⊗ p̂t (· | s, a, h) k Rh (s, a) ⊗ p(· | s, a, h)) ⩽ β µ+p (t, δ)   s,a,h

It is possible to choose β p , β µ , β µ+p with logarithmic dependency to ensure that each event hold with probability 1 − δ (cf e.g., Kaufmann & Koolen 2021). Since our focus is on characterizing the sample complexity in the asymptotic regime, we do not make these terms explicit. The interested reader may refer to Kaufmann & Koolen 2021 for explicit expressions of β p , β µ , and β µ+p . Moreover, the event o n E cnt (T ) ≜ ∀t ⩾ 1, ∀(s, a, h) : nt (s, a, h) ⩾ 12 n̄t (s, a, h) − log(2SAH/δ) . holds with probability at least 1 − δ, thanks to Lemma 34. Lemma 34 (Dann, Lattimore, et al. 2017). Let X1 , X2 , . . . be a sequence of Bernoulli random variables adapted to a filtration {Ht }t⩾1 and such that for any i, Pr(Xi = 1 | Hi−1 ) = Pi . We have X 1X Xt ⩽ Pt − W ⩽ exp(−W ) . Pr ∃n ⩾ 1 : 2 t⩽n

J.2

t⩽n

Gaussian and Dirichlet Concentration

We leverage Lemma 35 to prove concentration of distribution over MDPs in M. Lemma 35 (Lu & W. V. Li 2009). Let X ∼ N (θ, Σ) be a multivariate normal vector in dimension d. For any convex set C ⊂ Rd 2 1 1 PrX∼N (θ,Σ) (X ∈ C) ⩽ e− inf λ∈C 2 ∥λ−θ∥Σ−1 . 2

Cyrille Kone, Kevin Jamieson

We note that from the proof of Lemma 35 in Lu & W. V. Li 2009, the result extends to multivariate normals truncated to to a convex set. We prove the following crucial concentration. Lemma 36. Let P ≜ (P (s, a, h))s,a,h be a distribution whose marginals (P (s, a, h))s,a,h are independent Dir((α(s′ | s, a, h) + 1)s′ ∈S ) and R ≜ (R(s, a, h))s,a,h where each R(s, a, h) is an independent N (µ(s, a, h),σ 2 (s, a, h), possibly truncated to a convex set C. Let X ⊂ M such that for any q ∈ P, Xq ≜ θ ∈ (0, 1)SAH : (θ, q) ∈ X is convex. Introducing (κh )h satisfying κ(s′ | s, a, h)α(s, a, h) = α(s′ | s, a, h) , ∀(s, a, h, s′ ), we have Pr((R, Q) ∈ X )   ( " X (µ(s, a, h) − θ(s, a, h))2 1 Y ⩽ max 1, Sα(s, a, h)S−1 e−Ent(κ(·|s,a,h))  exp − inf + 2 2σ 2 (s, a, h) (θ,q)∈X s,ah s,a,h #) X α(s, a, h) KL(κ(· | s, a, h)kq(· | s, a, h)) . s,a,h

Proof. We have X ⊂ M ≜ (0, 1)SAH × 4SAH and by assumption, for any q ∈ 4SAH , the set Xq = {r ∈ R : (r, q) ∈ X } is convex. Next, P has the same distribution as a multivariate normal in dimension SAH with independent marginals. Noting that Pr((R, P ) ∈ X ) = E [Pr((R, P ) ∈ X | P )], and, since R, P are independent, leveraging Lemma 35 we have Pr((R, P ) ∈ X ) = E Pr((R, P ) ∈ X P )     X (µ(s, a, h) − θ(s, a, h))2  1   ⩽ E exp − inf (J.1)  θ s.t (θ,P )∈X  2 2σ 2 (s, a, h) s,a,h

which invokes Lemma 35 since given P ∈ P, XP as defined above is convex and R is a multivariate normal vector. We recall that for α ∈ Rd , the density of Dir(α) is proportional to e

−

1 i (αi −1) log xi

Next, we have for any (s, a, h), letting q(· | s, a, h) ∈ 4, X s′

α(s′ | s, a, h) log

1 q(s′ | s, a, h)

= α(s, a, h)

κ(s′ | s, a, h) log

1 q(s′ | s, a, h)

κ(s′ | s, a, h) log

κ(s′ | s, a, h) − q(s′ | s, a, h)

s′

= α(s, a, h)

X s′

α(s, a, h)

κ(s′ | s, a, h) log κ(s′ | s, a, h)

s′

= α(s, a, h)

κ(s′ | s, a, h) log

s′

Thus, the density of Dir(1 + α(· | s, a, h)) is also proportional to ′

′

e−α(s,a,h) KL(κ(s ·|s,a,h) ∥ q(s ·|s,a,h))) .

Optimal Posterior Sampling for Policy Identification in Tabular Markov Decision Processes

Combining with Eq.(J.1) we have

 

 X (µ(s, a, h) − θ(s, a, h))2 

1 exp − inf ·  θ s.t (θ,q)∈X  2N P 2σ 2 (s, a, h) s,a,h    X  exp − α(s, a, h) KL(κ(· | s, a, h) k q(· | s, a, h)) dq ,  

Pr((R, P ) ∈ X ) ⩽

s,a,h

where

 

Z N≜

exp P

Next, Pr((R, P ) ∈ X )

⩽



−

α(s, a, h) KL(κ(· | s, a, h) k q(· | s, a, h))

s,a,h

  

(J.2)

dq .

   X (µ(s, a, h) − θ(s, a, h))2 vol(P) · exp − inf  inf +  q∈P θ s.t. (θ,q)∈X 2N 2σ 2 (s, a, h) s,a,h   X αh (s, a) KL(κ(· | s, a, h) k q(· | s, a, h)) ,  s,a,h

then applying Lemma 40 yields   X (µ(s, a, h) − θ(s, a, h))2 X inf inf  + α(s, a, h) KL(κ(· | s, a, h) k q(· | s, a, h)) = q∈P θ s.t (θ,q)∈X 2σ 2 (s, a, h) s,a,h s,a,h   X X (µ(s, a, h) − θ(s, a, h))2 inf  + α(s, a, h) KL(κ(· | s, a, h) k q(· | s, a, h)) , 2σ 2 (s, a, h) (θ,q)∈X s,a,h

thus we have Pr((R, Q) ∈ X ) ⩽

s,a,h

 



vol(P) exp − inf   (θ,q)∈X 2N

We remark that N rewrites as

X (µ(s, a, h) − θ(s, a, h))2 2σ 2 (s, a, h)

s,a,h

  α(s, a, h) KL(κ(· | s, a, h) k q(· | s, a, h)) , 

# e

−α(s,a,h) KL(κ(·|s,a,h)∥q)

△

then we invoke Lemma 37 which yields

( e

−α(s,a,h) KL(κ(·|s,a,h)∥q)

△

dq ⩾ min

eEnt(κ(·|s,a,h)) α(s, a, h) S!

−(S−1)

1 , (S − 1)!

) ,

where Ent(p) denotes the entropy. 1 Combining the above displays and noting that vol(4) = (S−1)! and vol(P) = (vol(4))SAH yields

Pr((R, Q) ∈ X )   ( " Y X (µ(s, a, h) − θ(s, a, h))2 1 S−1 −Ent(κ(·|s,a,h))  exp − inf max 1, Sα(s, a, h) e + ⩽ 2 2σ 2 (s, a, h) (θ,q)∈X s,ah s,a,h #) X α(s, a, h) KL(κ(· | s, a, h)kq(· | s, a, h)) . s,a,h

Cyrille Kone, Kevin Jamieson

The following lemma bounds the normalizing constant of the posterior Dirichlet distribution when its density is expressed in terms of KL divergence. P Lemma 37. Let q ∈ 4 (the probability simplex of Rd ) and denote by Ent(q) ≜ i −qi log qi the entropy of q. We have Ent(q) −(d−1) Z e n 1 e−n KL(q ∥ x) dx ⩾ min , . d! (d − 1)! △ Proof. Below dx denotes a Hausdorff-measure differential element and, (only) below, Γ is the gamma function. R 1 1 = (d−1)! . Next, we assume n ⩾ 1. Let ζ ∈ (0, 1) and For n = 0, the bound follows as △ dx = Vol(4) = Γ(d) observe that Z Z e−n KL(q ∥ x) dx ⩾ e−n KL(q ∥ x) ddx △ (1−ζ)q+ζ△ Z d−1 = ζ e−n KL(q ∥ (1−ζ)q+ζx) dx △ Z d−1 ⩾ ζ e−n(1−ζ) KL(q ∥ q)−nζ KL(q ∥ x) dx △

which follows by convexity of x 7→ KL(q k x). Combining with KL(q k q) = 0 and letting ζ = 1/n it follows that Z e

−n KL(q ∥ x)

△

⩾

−(d−1)

e− KL(q∥x) dx. △

Next, we bound the right-hand integral. We observe that KL(qkx) =

qi log

where Ent(q) ≜

X qi = −Ent(q) − qi log xi , xi i

i −qi log qi . Combining above displays yield

e− KL(q∥x) dx = eEnt(q)

Z Y

△

Next, observe that

R Q △

△

xqi i dx

qi i xi dx is the normalizing constant of Dir(1 + q) thus

Z Y △

Q xqi i dx

Γ(1 + qi ) iP

Γ( Q

i 1 + qi )

i Γ(1 + qi )

Γ(d + 1)

Q =

i Γ(1 + qi )

where Γ is the gamma function. As proven in Z.-H. Yang et al. 2017, Γ(1+x) ⩾ xx for any 0 ⩽ x ⩽ 1. Combining the above displays, we have Z

e− KL(q ∥ x) dx △

Therefore

Z △

⩾

eEnt(q) Y Γ(1 + qi ) d! qiqi i

⩾

eEnt(q) Y qiqi eEnt(q) . qi = d! qi d! i

e−n KL(q ∥ x) dx ⩾

eEnt(q) n−(d−1) . d!

Optimal Posterior Sampling for Policy Identification in Tabular Markov Decision Processes

TECHNICAL RESULTS

We state different technical results used in the other sections. The following result is known and appears in Krichene et al. 2015; we include a proof here for completeness. Lemma 38. Let ρ be a density wrt canonical measure of X and g : X → R Z 1 η 7→ log ρ(x) exp(−ηg(x))dx η is non-decreasing. R Proof. Let Λη ≜ ρ(x) exp(−ηg(x))dx and f (η) = η1 log Λη and let the density wη (x) = ρ(x) exp(−ηg(x)) and Λη X ∼ wη . By derivation under the integral sign R 1 1 −g(x)ρ(x) exp (−ηg(x)) dx ′ R f (η) = − 2 log Λη + η η ρ(x) exp (−ηg(x)) dx 1 1 = − 2 log Λη + 2 EX [log exp(−ηg(X))] η η 1 exp(−ηg(X)) EX log = η2 Λη 1 wη (X) = EX log η2 ρ(X) 1 KL(νη , νρ ) ⩾ 0 , = η2 where νη is the distribution with density wη and νρ the distribution with distribution ρ. Lemma 39. For any mesaurable set A ⊂ M there exists an aboluste constant k0 < ∞ such that for t ⩾ k0 , √

vol(A ∩ M1/S· t ) ⩾

1 vol(A) . 2

√

Proof. The sequence at ≜ vol(A ∩ M1/S· t ) is non-decreasing and upper bounded by vol(A), which is also its supremum, thus it converges to this value, and the conclusion follows from this observation. Lemma 40. Let ∆ be a subset of X × Y and f : S → R be such that for all (x, y) ∈ X × Y, f (x, y) = g(x) + h(y) with g : X → R and h : Y → R. Then " # g(x) +

inf

x∈X

inf

y∈Y s.t (x,y)∈∆

h(y) =

inf

f (x, y)

(x,y)∈∆

Proof. inf

f (x, y)

(x,y)∈∆

inf

x∈X

inf

x∈X

inf

f (x, y)

inf

[g(x) + h(y)]

y∈Y s.t (x,y)∈∆

" =

inf

x∈X

g(x) +

# inf

y∈Y s.t (x,y)∈∆

h(y) .

f ∈ M satisfy kM − M fk1 ⩽ γ. Then for any (s, a, h) ∈ S × A × [H] and any π ∈ Πdet , Lemma 41. Let M, M Q̃πh (s, a) − Qπh (s, a) ⩽ (H − h + 1) γ.

Cyrille Kone, Kevin Jamieson

Proof. Conditioning on sh = s, ah = a and following π thereafter, the value difference lemma (Lemma 13) f gives applied between M and M # " H X π π π δk (sk , ak ) sh = s, ah = a , Qh (s, a) − Q̃h (s, a) = EM k=h

where

⊤ π (·). δk (s, a) ≜ µ(s, a, k) − µ̃(s, a, k) + pk (· | s, a) − p̃k (· | s, a) Ṽk+1

π Since rewards lie in (0, 1), kṼk+1 k∞ ⩽ H − k, so

H X

EπM [|δk (sk , ak )| | sh = s, ah = a]

k=h

⩽

H XX

(H − k + 1)γs′ ,a′ ,k

s′ ,a′ k=h

fk1 ⩽ (H − h + 1) kM − M ⩽ (H − h + 1) γ.

f ∈ Alt(M ) does not have enough enclosing volume The purpose of the following lemma is to show that if M around it, then it is a diﬀicult instance, i.e., the sup-optimality gaps are small. Lemma 42. Let M ∈ M and let π ∈ Πdet be its unique globally optimal policy. Denote by ∆π (M ) the smallest value gap in M for policy π, and let Aπ := M ′ ∈ M : π is globally optimal in M ′ . δ If the ℓ1 -ball of radius 2(H+1) around M is not entirely contained in Aπ , i.e.,

M ′ ∈ M : kM − M ′ k1 ⩽

δ 2(H + 1)

6⊂ Aπ ,

then the smallest gap of M must satisfy ∆π (M ) ⩽ δ . Proof. Fix δ, δ ′ ⩾ 0. Assume that ∆π (M ) > δ ′ , and let M ′ ∈ M satisfy kM − M ′ k1 ⩽ δ. e h and Veh the state–action and state value functions in M ′ , and by Qh and Vh their counterparts in Denote by Q M. Then, for any action a 6= πh (s), we have

⩾

e πh (s, a) Qπh (s, πh (s)) − Qπh (s, a) + Vehπ (s) − Vhπ (s) + Qπh (s, a) − Q e πh (s, a) δ ′ + Vehπ (s) − Vhπ (s) + Qπh (s, a) − Q

⩾

π π δ ′ − δ − |Vehπ (s) − Vhπ (s)| − |Veh+1 (s) − Vh+1 (s)|

⩾

δ ′ − (2H − 2h + 2)δ,

e π (s, πh (s)) − Q e π (s, a) = Q h h

where the last inequality follows by value difference lemma as kM − M ′ k1 ⩽ δ. Thus, it suﬀices to take δ ′ ≜ 2(H + 1)δ

Optimal Posterior Sampling for Policy Identification in Tabular Markov Decision Processes

to guarantee that π is the unique globally optimal policy in the MDP M ′ . Indeed, for any other policy ρ, by induction, we have at step h, Vhρ (s)

ρ µ(s, ρh (s), h) + ph Vh+1 (s, ρh (s))

π µ(s, ρh (s), h) + ph Vh+1 (s, ρh (s))

π µ(s, πh (s), h) + ph Vh+1 (s, πh (s)) = Vhπ (s),

which concludes the proof.

The goal of the next lemma is to show that from an instance M , one can pick a policy π ∈ Π⋆det (M ) and build an easier instance M ′ (with larger value gaps) close to M and such that π is still optimal in M ′ . This will be f) (so a diﬀicult instance with nearly zero gap) used to show that if an instance M lies on the boundary of Alt(M ′ f f), showing that of another instance M , one can build an easier instance M close to M , which is still in Alt(M there is enough volume around any point in the interior of the alternative set of an MDP. Lemma 43. Let M ∈ M and let π ∈ Πdet be a globally optimal policy for M . For any ε ⩽ 3(1+21H−1 ) , there exists an MDP M̃ ≜ (µ̃, p̃) and coeﬀicients (γs,a,h )s,a,h ∈ [0, 1] such that for all (s, a, h), p̃h (· | s, a) ≜ (1 − γs,a,h )ph (· | s, a) + γs,a,h ph (· | s, πh (s)), with

|µ(s, a, h) − µ̃(s, a, h)| ⩽ ds,a,h ≜ 2H−h+1 ε, (K.1)

(1 + 2H−h )ε γs,a,h ∈ 0, . 1 − (2H−h+1 + 3)ε

Moreover, π is strongly globally optimal in M̃ : for all (s, h), Q̃πh (s, πh (s)) − max Q̃πh (s, a) ⩾ ε. a̸=πh (s)

Proof. We construct M̃ bottom-up from stage H to stage 1. Throughout, Q, V (resp. Q̃, Ṽ ) denote value functions in M (resp. M̃ ), and we maintain the induction hypothesis 0 ⩽ Ṽhπ (s) − Vhπ (s) ⩽ Ωh (ε) ≜ (2H−h+1 − 1)ε,

∀ s.

(K.2)

Base case h = H. If µ(s, πH (s), H) − µ(s, a, H) ⩽ ε, set µ̃(s, a, H) ≜ max(0, µ(s, a, H) − ε),

µ̃(s, πH (s), H) ≜ min(1, µ(s, πH (s), H) + ε),

(K.3)

and γs,a,H = 0. This ensures Q̃πH (s, πH (s)) − Q̃πH (s, a) ⩾ ε and 0 ⩽ ṼHπ (s) − VHπ (s) ⩽ ε = ΩH (ε), so (K.2) holds at h = H. Induction step. Assume (K.2) holds at stage h + 1. At stage h, for each (s, a) we consider three cases. Case 1. If π π µ(s, πh (s), h) − µ(s, a, h) + ph Ṽh+1 (s, πh (s)) − ph Ṽh+1 (s, a) ⩾ ε,

(K.4)

set p̃h (· | s, a) ≜ ph (· | s, a) and µ̃(s, a, h) ≜ µ(s, a, h), i.e., γs,a,h = 0. The gap is already at least ε without any perturbation. Case 2. If (K.4) does not hold but µ̃(s, πh (s), h) − µ̃(s, a, h) ⩾ µ(s, πh (s), h) − µ(s, a, h) + Ωh+1 (ε) + ε,

(K.5)

Cyrille Kone, Kevin Jamieson

set γs,a,h = 0. Shifting rewards alone is suﬀicient to guarantee the gap. Case 3. Otherwise, neither (K.4) nor (K.5) holds. Since (K.5) does not hold, we have µ(s, πh (s), h) + Ωh+1 (ε) + ε ⩾ 1

and

µ(s, a, h) − Ωh+1 (ε) − ε ⩽ 0,

(K.6)

Indeed, if any of the statements above above didn’t hold, we would have Eq.(K.5). This further implies µh (s, πh (s)) − µh (s, a) ⩾ 1 − 2 (Ωh+1 (ε) + ε) ,

(K.7)

and since (K.4) does not hold either, we have π π ph Ṽh+1 (s, πh (s)) − ph Ṽh+1 (s, a) ⩽ ε + 2(Ωh+1 (ε) + ε) − 1.

(K.8)

By the mixture definition of p̃h and the induction hypothesis, Q̃πh (s, πh (s)) − Q̃πh (s, a) ⩾ Qπh (s, πh (s)) − Qπh (s, a) − Ωh+1 (ε) π π + γ ph Ṽh+1 (s, a) − ph Ṽh+1 (s, πh (s))

⩾ −Ωh+1 (ε) + γ 1 − ε − 2(Ωh+1 (ε) + ε) , where the first inequality uses π ∈ Π⋆det (M ) (so Qπh (s, πh (s)) − Qπh (s, a) ⩾ 0) and (K.2), and the second uses (K.8). Setting ε + Ωh+1 (ε) (1 + 2H−h )ε γs,a,h ≜ = (K.9) 1 − ε − 2(Ωh+1 (ε) + ε) 1 − (2H−h+1 + 3)ε gives Q̃πh (s, πh (s)) − Q̃πh (s, a) ⩾ ε. The condition γs,a,h ∈ [0, 1] requires 1 − (2H−h+1 + 3)ε ⩾ (1 + 2H−h )ε, which 1 1 holds whenever ε ⩽ 3+2H−h +2 H−h+1 . This is satisfied for all h ∈ [H] if ε ⩽ 3(1+2H−1 ) . Verifying the induction hypothesis at stage h. Since p̃h (· | s, πh (s)) = ph (· | s, πh (s)) and µ̃(s, πh (s), h) ⩽ µ(s, πh (s), h) + ε, π Ṽhπ (s) = µ̃(s, πh (s), h) + ph Ṽh+1 (s, πh (s)) π ⩽ µ(s, πh (s), ) + ε + Ωh+1 (ε) + ph Vh+1 (s, πh (s))

= Vhπ (s) + 2Ωh+1 (ε) + ε, | {z } = Ωh (ε)

and, it is simple to verify that Ṽhπ (s) ⩾ Vhπ (s) since both rewards and transitions are translated in favour of π. This confirms (K.2) at stage h, with the closed form Ωh (ε) = (2H−h+1 − 1)ε. From the update rule (K.3) and the definition of Ωh+1 , |µ̃(s, a, h) − µ(s, a, h)| ⩽ Ωh+1 (ε) + ε = 2H−h+1 ε = ds,a,h , which holds for all (s, a, h), completing the proof. Lemma 44. Let an MDP M ∈ M with a unique stage-optimal policy π ∈ Πdet . If π is a global optimal policy ∆ in M and ∆π (M ) > 0, then for any γ ∈ 0, 6SAH(H+1) , π is a global optimal policy for any MDP in the set (1 − γ)M + γM. Proof. Let π ∈ Πdet be a global optimal policy in M . From Lemma 42 we have for ε ⩽ ∆ and for any ε M ′ ∈ B1 (M, 2(H+1) ), π is an optimal policy in M ′ . Now it suﬀices to take γ such that (1−γ)M +γM ⊂ B1 (M, γ ′ ). For this, observe that for any M ′ ∈ M we have kM − ((1 − γ)M + γM ′ )k1

γ kM − M ′ k1

⩽

2γSAH.

∆ and using Lemma 42 Thus we have (1 − γ)M + γM ⊂ B1 (M, 3γSAH) and taking γ such that 3γSAH ⩽ 2(H+1) concludes the proof.

Optimal Posterior Sampling for Policy Identification in Tabular Markov Decision Processes

The goal of the following lemma is to, given an MDP M , generate M ′ ∈ M such that the transition probabilities in M ′ are larger than some threshold while M and M ′ still have the same optimal policy. Lemma 45. Let an MDP M ≜ (µ, p) with optimal policy π such that its minimum policy gap ∆π (M ) > 0. Let u denote the uniform transition kernel and define for any γ ∈ (0, 1) the MDP Mγ = (µ, (1 − γ)p + γu). For γ ∈ (0, ∆M (π)/(3H), π is state-wise striclty optimal in Mγ . f = Mγ , ∆M (π) = ∆, and define Q, e Ve the Q-function Proof. Let γ be fixed and simplify notation by letting M f and value function in M . We have for any a = 6 πh (s) e πh (s, πh (s)) − Q e πh (s, a) Q

e πh (s, πh (s)) − Qπh (s, πh (s)) + Qπh (s, a) − Q e πh (s, a) ∆+Q π π ∆ + Vehπ (s) − Vhπ (s) + ph Vh+1 (s, a) − p̃h Veh+1 (s, a)

π π π ∆ + Vehπ (s) − Vhπ (s) + ph Vh+1 (s, a) − ph Veh+1 (s, a) + γ(ph − u)Veh+1 (s, a)

⩾

∆ − 3Hγ

⩾

so it suﬀices for γ to be smaller than ∆/(3H) to guarantee the claimed property.

The next lemma is proven in Menard et al. 2021. We restate it for completeness. Pt Lemma 46. For a sequence (ut )t⩾1 ∈ (0, 1)N and Ut ≜ l=1 ul we have T X ut+1 t=1

Ut ∨ 1

⩽ 4 log(UT +1 + 1).

BEYOND EXACT POLICY IDENTIFICATION

In this section, we discuss the extension of our algorithm to the setting with (ε, δ)-PAC policy identification with ε > 0. In this case, we let Π⋆ε (M ) be the set of (Markov, deterministic and stochastic) policies π such that ′

max V0π ⩽ V0π + ε . ′ π

The goal of the agent is then to identify a policy π ⋆ ∈ Π⋆ε (M ) with high probability, i.e., given an error rate δ, her final recommendation π̂ τ at stopping time τ should satisfy PrM (τ < ∞, π̂τ ∈ / Π⋆ε (M )) ⩽ δ. Al-Marjani et al. 2023b derived a lower bound for this setting, which we state below. To introduce n this lower bound, we introduce the following notation. For a policy π and given ε ⩾ 0 let Altπε (M ) ≜ M ′ : M o M ′ and π ∈ / Π⋆ε (M ′ ) . In words, it is the set of M ′ , for which M M ′ and where the π is not a near-optimal policy. Theorem 47. Fix an M ∈ M. For any adaptive stopping time τ for (ε, δ)-PAC policy identification lim inf δ→0

EM [τ ] ⩾ Γε (M ) , log 1δ

where

 Γ−1 ε (M ) ≜

max

inf

π ′ π∈Π⋆ ε (M ) w∈ΩM M ∈Altε (M )



 wh (s, a) KL(M, M ′ )s,a,h  .

s,a,h

Intuitively, the lower bound states that the optimal sample complexity should scale with the easiest to identify policy among the ε-optimal set of policies. From previous works in bandits, asymptotically matching this bound often required ”sticking procedures” to stick to an answer to verify. Indeed, for ε = 0 and |Π⋆ε (M )| = 1, sticking

Cyrille Kone, Kevin Jamieson

at each episode, the algorithm’s sampling rule is meant to collect evidence for assessing the true optimality of b t,ε | > 1. The challenge is then to pick its current best policy guess π̂t . For ε > 0, it is likely that at episode t, |Π which answer in this set the sampling rule should aim to verify. Following the approach developed in Appendix H it’s simple to prove that Theorem 48. Fix an MDP M ∈ M and let νt denote the posterior distribution introduced in Eq.(1). For any adaptive algorithm for best policy identification, with probability one, 1 ⋆ f lim sup − log min Pr π ∈ / Π ( M ) ⩽ Γε (M ). f ε M ∼νt |Ht−1 t π∈Π⋆ t→∞ ε (M ) c and To state the modification of the Algorithm 1, let a generic sticking procedure take as input the empirical M b output a policy in Πt,ε . We require the procedure to be deterministic and stable, that is if the sticking recommends b t′ ,ε for t′ ⩾ t, it should keep recommending it. In an answer at time t, then as long as this answer belongs to Π case of ties, the sticking break rules with consistency, so that this defines a strict total order on Π⋆ (M ). b t,ε → Π⋆ε (M f) so that the Because of the smooth forced exploration, we should have that on a good event Π sticking rule is simply a rule to select an answer to collect samples to verify it. We restate below the slight modification of Algorithm 1 to account for the ε > 0 case. Remark that compared to the initial algorithm, only Line 2 has changed, to pick an answer to verify among the set of empirically admissible answers. Intuitively, the sticking procedure transforms the problem into verifying a single answer, as opposed to verifying multiple correct answers. As the algorithm sticks to an answer π, the conditional posterior sampling is now a no-regret learner on Altπε (M ). Following the same lines as in previous sections, one could derive the no-regret property of the conditional posterior sampling on Altπε (M ). Combining these observations, the following will hold with probability one A posterior-based stopping time for the multiple correct answers setting. We propose the following stopping rule for the (ε, δ)-PAC policy identifcation with ε > 0. n o b t,ε such that ∀ k ⩽ B(t, δ), π ∈ Π⋆ε (M ft,k ) . τ ≜ inf t ⩾ 1 : ∃π ∈ Π (L.1) In words, the stopping rule triggers when there exists an answer for which finding a challenger by conditional resampling takes more than B(t, δ) samples. Observe that for this stopping rule, we have Pr(τ > t | Ht−1 )

⩽ ⩽ ⩽

b t,ε , ∃ k ⩽ B(t, δ) : π ∈ ft,k ) | Ht−1 ) Pr(∀ π ∈ Π / Π⋆ε (M ⋆ f / Π (Mt,k ) | Ht−1 ) min Pr(∃ k ⩽ B(t, δ) : π ∈ ε

b t,ε π∈Π

ft ∈ Altπε (M̂t ) | Ht−1 ) , B(t, δ) min Pr(M b t,ε π∈Π

ft , M ft,1 , . . . , M ft,k are i.i.d . where M By the concentration of the posterior distribution, we would have with probability one,        X e B(t, δ) min exp − f)s,a,h    Pr(τ > t | Ht−1 ) ⪅ O inf nt (s, a, h) KL(M̂t kM  M  b t,ε f∈Altπ (M̂t ) π∈Π ε s,a,h        X e B(t, δ) exp − max f)s,a,h    ⪅ O inf nt (s, a, h) KL(M̂t kM  π∈Π  b t,ε M f∈Altπ (M̂t ) ε s,a,h

where Õ hides constant and terms that are at most polynomial in t. As the empirical estimates converge, the term in the exponent will be equivalent (up to logarithmic terms in t)   X f)s,a,h   max inf nt (s, a, h) KL(M kM π∈Π⋆ f∈Altπ (M ) ε (M ) M ε

s,a,h

Optimal Posterior Sampling for Policy Identification in Tabular Markov Decision Processes

which is larger than evaluating it at any response computed by the sticking procedure. Combining the convergence of the algorithm with ”sticking” with the stopping time τ , one can prove results similar to the previous section in terms of sample complexity. It is also possible to ensure that if the sampling rule guarantees saddle-point convergence, then plugging in this stopping rule will ensure optimality. We leave as an open question the design of a converging, computationally eﬀicient sampling rule for the setting with ε > 0. Given m posterior samples, checking if the stopping triggers require solving the feasibility problem k,π k,⋆ max min (Vet,0 + ε − Vet,0 )>0

b t,ε k⩽m π∈Π

b t,ε , which is a major caveat. For the computation of the and solving this might ultimately require enumerating Π stopping rule. Algorithm 5: ε-PIPS: Policy Identification via Posterior Sampling in Tabular MDPs Require: Initial exploration π1exp = {π1exp (·|s, h)}s,h ; initial forced exploration c1 ; mixing parameter γ ∈ (0, 1); clipping parameter α ∈ (0, 12 ); posterior inflation (ηt )t⩾1 ; initial MDP M̂0 ∈ M (arbitrary); ε ⩾ 0 1

for episode t = 1, 2, . . . do // Stick to an answer among current empirical ε-best guesses

Compute π̂t ← sticking(M̂t , ε)

// Computing a challenger MDP

for k = 1, 2, . . . do ft,k ∼ νtηt Sample candidate challenger M e k,⋆ ft,k by backward induction Compute the optimal Q-function Q of M t k,⋆ e if π̂t is ε-suboptimal for Qt then ft ← M ft,k and break Set challenger M end end

3 4 5 6 7 8 9

Compute ct as in Algorithm 4 Draw Zt ∼ Bernoulli(t−γ ) and define ∀(s, h) ∈ S × [H]

10 11

( π̃texp (·|s, h)

←

ct (·|s, h), πtexp (·|s, h),

if Zt = 1, if Zt = 0,

// Rollout and data collection

Set st1 ← sinit for h = 1 to H do Sample action ath ∼ π̃texp (·|sth , h) Observe reward rht and transition to sth+1 end

12 13 14 15 16

// KL-based bonuses for the online learner

ft as in Eq.(3) Construct bonus kernel rtexp (s, a, h) from M̂t and M

// Behaviour policy improvement exp ← ComputeExplorationPolicy t, πtexp , rtexp , p̂t , nt Update main exploration policy πt+1 19 end 18

IMPLEMENTATION DETAILS AND ADDITIONAL EXPERIMENTS

This section contains implementation details and additional experimental results comparing PIPS to PSRL.

Cyrille Kone, Kevin Jamieson

M.1

Implementation Details

In this section, we provide details on the implementation and experimental setup for reproducibility.

Experimental setup. For each environment shown in Figure 2 and Figure 8, rewards are modeled as Gaussian random variables with means specified in the figures. In the Riverswim environment, the reward variance is set to 1/20, while in the MOCA instance it is set to 1/4. Although both environments are stationary, we did not include this information in our algorithm, which allows us to estimate the model and maintain distributions over non-stationary MDPs.

Online learner. For the policy improvement step of the online learner, we use AdaHedge, implementing the pseudo-code described in De Rooij et al. 2014, which automatically adjusts its learning rate at each round. We recall that the initialization of Algorithm 1 with a generic online learner is described and analyzed in Appendix D, with its pseudo-code in Algorithm 1. In practice, we observe little difference in performance compared to the natural Hedge update. As argued in Appendix D, both algorithms ensure sublinear regret and thus preserve the theoretical guarantees of Algorithm 1. Posterior inflation and thresholds calibration. The posterior is inflated using ηt = 1/(t + 1)4 , which provides a good trade-off between empirical performance and runtime. Running without inflation (i.e., ηt = 1) tends to yield slower experiments without noticeable benefits. However, for simplicity, we recommend using ηt ← 1 with B(t, δ) ≈ SAH δ log(tSAH/δ) which experimentally ensure δ-correctness. Other algorithms implementation. For each algorithm, we refer to the pseudo-code in the original paper for the implementation. For the UCBVI algorithm (Azar, Osband, et al. 2017) and SSR (Xiong et al. 2022) we implement the following idealized Bernstein confidence bonuses s bonust (s, a, h) = min

Varp̂t [V̄t,h+1 ](s, a) log(t + 1) log(t + 1) + (H − h + 1) ,H − h + 1 max(1, nt (s, a, h)) max (1, nt (s, a, h))

! .

For PSRL, the prior distribution is the same as for Algorithm 1.

M.2

Additional Experiments

We present further experimental results on the MOCA instance. Specifically, we study the number of posterior samples required before identifying an alternative model. This experiment evaluates the behavior of the posterior stopping time. Since both PIPS and PSRL rely on posterior sampling, we compare their performance in this setting. Each algorithm is run on the MOCA instance for T = 10,000 episodes. At every episode, we record the number of posterior resamples needed to obtain an alternative model. For PSRL, this quantity is immaterial, but for PIPS it directly influences the sampling rule. We report the average (over 50 independent runs). The results in Figure 5 show that for PIPS, the number of resamples grows exponentially with t, as further confirmed by the log-plot in Figure 6. This behavior is consistent with theory: conditionally on the history, the number of resamples follows a geometric law with parameter scaling with e−tΓ , so its expectation grows exponentially. This indicates that evidence accumulates rapidly, enabling the algorithm to certify the optimal policy. In contrast, PSRL exhibits only polynomial growth in the number of resamples (also evident from the log-log plot in Figure 7). This suggests that its posterior concentrates at a polynomial rate.

Optimal Posterior Sampling for Policy Identification in Tabular Markov Decision Processes

300 200 100

101

100

0 0

2000

4000 6000 episode

8000

10000

PSRL

nbr. of samples

101

PIPS nbr. of samples

PIPS PSRL

400

100 0

2000

4000 6000 episode

8000

10000

103

104 episode

Figure 5: Number of samples Figure 7: Log-log 1 of samples for Figure 6: Log number 1 1 plot of number finding an alternative model aver- for finding an alternative model av- of samples for finding an alternative aged on 100 runs. eraged on 100 runs. model averaged on 100 runs.

Figure 8: The L-state RiverSwim environment. Each state allows two actions: left (red arrows) and right (black arrows). Rewards are placed at the two ends of the chain. Transition probabilities are annotated above the corresponding arrows.

Record · ID 155281 · SHA-256 507be443b5e6da2c
Retrieved via Conceptio — every document is proof-bundled with source, license, and retrieval metadata.