Model-based RL with Optimistic Posterior Sampling: Structural Conditions and Sample Complexity
Authors: Alekh Agarwal
Tong Zhang
Venue: NeurIPS 2022 Accept (2022)
Keywords: Reinforcement Learning
Model-based RL
Sample Complexity

Abstract: We propose a general framework to design posterior sampling methods for model-based RL. We show that the proposed algorithms can be analyzed by reducing regret to Hellinger distance in conditional probability estimation. We further show that optimistic posterior sampling can control this Hellinger distance, when we measure model error via data likelihood. This technique allows us to design and analyze unified posterior sampling algorithms with state-of-the-art sample complexity guarantees for many model-based RL settings. We illustrate our general result in many special cases, demonstrating the versatility of our framework.

Model-based RL with Optimistic Posterior Sampling: Structural Conditions and Sample Complexity
Alekh Agarwal Google Research
Tong Zhang Google Research and HKUST
Abstract
We propose general framework to design posterior sampling methods for model- based RL. We show that the proposed algorithms can be analyzed by reducing re- gret to Hellinger distance in conditional probability estimation. We further show that optimistic posterior sampling can control this Hellinger distance, when we measure model error via data likelihood. This technique allows to design and analyze unified posterior sampling algorithms with state-of-the-art sample com- plexity guarantees for many model-based RL settings. We illustrate our general result in many special cases, demonstrating the versatility of our framework.
Introduction
In model-based RL, learning agent interacts with its environment to build an increasingly more accurate estimate of the underlying dynamics and rewards, and uses such estimate to progressively improve its policy. This paradigm is attractive as it is amenable to sample-efficiency in both the- ory [Kearns and Singh, 2002, Brafman and Tennenholtz, 2002, Auer al., 2008, Azar al., 2017, Sun et al., 2019, Foster et al., 2021] ] and practice [Chua et al., 2018, 8, Nagabandi et al., 2020, Schrit- twieser et al., 2020J. The learned models also offer the possibility of use beyond individual tasks (Ha and Schmidhuber, 2018], effectively providing learned simulators. Given the importance of this paradigm, it is vital to understand how and when can sample-efficient model-based RL be achieved.
Two key questions any moddll-based RL agent has to address arc: i) how to collect data from the environment, given the current learning state, and ii) how to define the quality of model, given the currently acquired dataset. In simple theoretical settings such as tabular MDPs, the first question typically addressed through optimism, typically via an uncertainty bonus, while likelihood of the data under model is a typical answer to the second. While the empirical literature differs on its answer to the first question in problems with complex state spaces, it still largely adopts likelihood, VAE-style approximation [e.g. Chua et al., 2018, Ha and Schmidhuber, 2018, Sekar et al., 2020], as the measure model quality. On the other hand, the theoretical literature for rich observation RL is much more varied in the loss used for model fitting, ranging from a direct squared errrr in parameters [Yang and Wang. 2020, Kakade et al., 2020] to more complicated divergence measures [Sun et al., 2019, Du et al., 2021, Foster et al., 2021]. Correspondingly, the literature also has a variety of structural conditions which enable model-based RL.
In this paper, we seek to unify the theory of model-based RL under a general theoretical and algo- rithmic framework. We address the aforementioned two questions by always using data likelihood as model quality estimate, irrespective of the observation complexity, and use optimistic poste- rior sampling approach for data collection. For this approach, we define a Bellman error decoupling framework which governs the sample complexity of finding near-optimal policy. Our main result establishes that when the dccoupling cocfficient in a model-based RL setting is small, our Model- based Optimistic Posterior Sampling algorithm (MOPS) is sample-efficient.
36th Conference on Noural Information Processing Systems (NeurIPS 2022).

---

A key conceptual simplification in MOPS is that the model quality is always measured using like- lihood, unlike other general frameworks [Du et al., 2021, Foster et al., 2021] which need to mod- ify their loss functions for different settings. Practically, posterior sampling is relatively amenable to tractable implementation via ensemble approximations [see e.g. Osband et al., 2016a, Lu and Van Roy, 2017. Chua et al., 2018. Nagabandi et al., 2020] sampling methods such stochastic gradient Langevin dynamics [Welling and Tch, 2011]. This is in contrast with the version-space based optimization methods used in most prior works. We further develop broad structural condi- tions on the underlying MDP and model class used, under which the decoupling coefficient admits good bounds. This includes many prominent examples, such
Finite action problems with a small witness rank [Sun et al., 2019]
Q-type witness rank problems (new) Kernelized Non-linear Regulators [Kakade et al., 2020]
Linear MDPs with infinite actions
Small witness rank and linearly embedded backups (new)
Linear mixture MDPs [Modi ct al., 2020, Ay- oub ct al., 2020]
Remarkably, MOPS simultaneously addresses all the scenarios listed here and more, only requiring one change in the algorithm regarding the data collection policy induced as function of the poste- rior. The analysis of model estimation itself, which is shared across all these problems, follows from a common online learning analysis of the convergence of our posterior. Taken together, our results provide a coherent and unified approach to model-based RL, and we believe that the conceptual simplicity our approach will spur future development in obtaining more refined guarantees, more efficient algorithms, and including new examples under the umbrella of statistical tractability.
2 Related work
Model-based RL. There is rich literature model-based RL for obtaining strong sample com- plexity results, from the seminal early works [Kearns and Singh, 2002, Brafman and Tennenholtz, 2002, Auer et al., 2008] the more recent minimax optimal guarantees [Azar et al., 2017, Zanette and Brunskill, 2019]. Beyond tabular problems, techniques have also been developed for rich ob- servation spaces with function approximation in linear [Yang and Wang, 2020, Ayoub et al., 2020] and non-linear [Sun et al., 2019, Agarwal et al., 2020, Uehara et al., 2021, Du et al., 2021] settings. Of these, our work builds most directly on that of Sun et al. [2019] in terms of the structural prop- erties used. However, the algorithmic techniques are notably different. Sun et al. [2019] repeatedly solve an optimization problem to find an optimistic model consistent with the prior data, while we use optimistic posterior sampling, which scales to large action spaces unlike their use of uniform randomization over the actions. We also measure the consistency with prior data in terms of the likelihood of the observations under a model, as opposed to their integral probably metric losses. Du et al. [2021] effectively reuse the algorithm of Sun et al. [2019] in the model-based setting, so the same comparison applies. We note that recent model-based feature learning works [Agarwal et al., 2020, Uehara et al., 2021] do measure model fit using log- likeliod,oo,, and some of their tech- nical analysis shares similarities with our proofs, though there are notable differences in the MLE versus posterior sampling approaches. Finally, we note that a parallel line of work has developed model-free approaches for function approximation settings [Jiang et al., 2017, Du et al., 2021, Jin et al., 2021] 1], and while our structural complexity measure is always smaller than the Bellman rank of Jiang et al. [2017], our model-based realizability is often stronger assumption than realizability of just the Q' function. For instance, Jin et al. [2020], Du et al. [2019] and Misra et al. [2019] do not model the entire transition dynamics in linear and block MDP models.
Posterior sampling in RL. Posterior sampling methods for RL, motivated by Thompson sampling (TS) [Thompson, 1933], have been extensively developed and analyzed in terms of their expected regret under a Bayesian prior by many authors [see c.g. Osband et al., 2013, Russo et al., 2017, Osband et al., 2016b] and often popular they offer simple implementation heuristic through approximation by bootstrapped ensembles [Osband et al., 201 6a]. Worst-case analysis of TS in RL settings has also been done for both tabular [Russo, 2019, Agrawal and Jia, 2017] and linear [Zanette et al., 2020] settings. Our work is most closely related to the recent Feel-Good Thompson Sampling strategy proposed and analyzed in [Zhang, 2021], and its extensions [Zhang et al., 2021, Agarwal and Zhang, 2022]. Note that our moodel-aaeed setting does not require the two timestep strategy to solve a minimax problem, as was required in the model-free work of Agarwal and Zhang [2022].
2

---

Model-based control. Model-based techniques are widely used in control, and many recent works analyze the Linear Quadratic Regulator [see e.g. Dean et al., 2020, Mania et al., 2019, Agarwal et al., 2019, Simchowitz and Foster, 2020], as well as some non-linear generalizations [Kakade et al., 2020, Mania et al., 2020, Mhammedi et al., 2020]. Our framework does capture many of these settings we demonstrate in Section 6.
Relationship with Foster et al. [2021]. This recent work studies a broad class of decision mak- ing problems including bandits and RL. For RL, they consider model-based framing, and provide upper and lower bounds on the sample complexity in terms of a new parameter called the decision estimation coefficient (DEC). Structurally, DEC is closely related to the Hellinger decoupling con- cept introduced in this paper (see Definition 2). However, while Hellinger decoupling is only used our analysis (and the distance is measured to the true model), the DEC analysis of Foster et al. [2021] measures distance to plug-in estimator for the true model, and needs complicated algorithms to explicitly bound the DEC. In particular, it is not known if posterior sampling is admissible in their framework, which requires more carcful control of a minimax objective. The Hellinger dccoupling coefficient, in contrast, admits conceptually simpler optimistic posterior sampling techniques.
3 Setting and Preliminaries
We study RL in an episodic, finite horizon Contextual Markov Decision Process (MDP) that is parameterized as (X, R,, P.), where X' is state space, A is action space, D is the dis- tribution over the initial context, R, is the expected reward function and P. denotes the transition dynamics. An agent observes context x. ~ 'D for some fixed distribution D. At each time step h E (1,.... H), the agent observes the state xh, chooses an action a², observes r with Elr h D1, ax] = R (z), an) and transitions to hh+1 ~ ( I², a²). We assume that z² for any h 1 always includes the context z! to allow arbitrary dependence of the dynamics and rewards on x Following prior works [e.g. Jiang et al., 2017, Sun et al., 2019], we assume that r E [0, 1] and E rd E [0, to capture reward settings [Jiang and Agarwal, 2018]. We make no assumption on the cardinality the state and/or action spaces, allowing both to be potentially in- finite. We use to denote the agent's 's decision policy, which maps from X' + A(A), where A(-) represents probability distributions over a set. The goal of learning is to discover an optimal policy T.. which is always deterministic and conveniently defined in terms of the Q, function [see e.g. Puterman, 2014, Bertsekas and Tsitsiklis, 1996]
(x²h = argmax Qh ,), (x a²) = E max +1 a') x² aa", acAA
(1)
a'EAA
where we define Q4t a) 0 for all I, a. We also define ((*(a* max, (zt,aa)
In the model-bascd RL setting of this paper, the learner has access to a model class M consisting of tuples (Pxs, ). denoting the transition dynamics and expected reward functions according to the model M E M.. For any model M, we use T&e and Vh to denote the optimal policy and value function, respectively, at level h in the model M. We assume that like ViA, V also satisfies the normalization assumption V (x) E [0, 1] for all M E M. We use E to denote expectations evaluated in the model M and E denote expectations in the true model.
We make two common assumptions on the class M.
Assumption (Realizabilitty M, E M such that P = Ph, and R = -Rh,, for all h E [M]. The realizability assumption ensures that the model estimation is well-specified. We also assume access to a planning oracle for any fixed model M E M.
Assumption 2 (Planning oracle) There is a planning oracle which given a model M, returns its optimal policy πM = [ Tas = arg max, MM,x |E R& (xxh, x (2>)) where MM is the expectation over trajectories obtained by following the policy in the model M.
Given M E M, we define model-based Bellman сттоr as:
EM a²) Qhr (w0 a²)) (Ph 1
Vi++ aa)),
(2)
We intentionally call context and not an initial state of the MDP as we will soon make certain structural assumptions which depend on the context, but take expectation over the states. We use PM and H&s to denote the sets { M and & Him respectively.

---

Algorithm 1 Model- based Optimistic Posterior Sampling (MOPS) for model-based RL
Require: Model class M, prior Po E A(M), policy generator Wgen. learning rates 17. r/ and opti-
mism coefficient
1: Set 5n
2: for 1= do
3: Observe oi ~ and draw hy ~ {I,............... .,. H) uniformly at random.
4: Let (MM - (R (x a²) r2))2 + r/' In Ph (a h+1 x a Likelihood function
5:
Define Pi(M = O(MSS-1) x Po(M) exp S (yVas (a:)) + - (M)) as the posterior D Optimistic posterior sampling update
6: Let गΙ ggen (gen (he.Pt)
b policy generation
7: Play iteration t using for h = 1,.... hr and observe ((* at rf, )/=
8: Update S. St (1t af r1, x 1.,t1) for h4-
9: end for
return (TI TT).
where for function f X x A x [0, x X, define (PM)) BM a, This quantity plays a central role in our analysis due to the simulation lemma for model-based RL.
(,,,,, x') a].
Lemma (Lemma 10 of [Sun et al., 2019I.. For any distribution p E A(M) and xt, we have Emax [V*(x!)-VV V (xx)] EM-P E at EB(M, a²) (VM))
where AVM(x)) Vas( V*'('')
The Bellman error can in turn be related to Hellinger error in conditional probability estimation via the decoupling coefficient in Definition 2, which captures the structural properties of the MDP.
We use typical measures distance between probability distributions to capture the er- in dynamcis, and for any two distributions P and Q over samples note TV(P, Q) = 1 1/2 (P, Ez-r( (dQ(=)/dP(z) - 1) We use (SS) to denote the space of all probabil- ity distributions over set S (under suitable o-algebra) and [H] [I....,
Z, we de- - 1/. KL(PIIQ) = Eg<P In d22))/dQ(2) and
Effective dimensionality. We often consider ciiintt-dimmesiinne maps X(21, x2) over a pair of inputs E Z1 and 22 E Zz. We define the effective dimensionality of such maps as follows.
Definition 1 (Effective dimension) Given any measure over 21 x 22, and feature map x, define:
E(p, 21) =
E x
x(21,22) 22) @ (2), K(A)= = sup trace ((E(()())))))))))))) 21) + X1) -E(p 21)). p.=.
For any O, define the effective dimension of x (x. inf {K()) XK(A) < }
If dim(x) d, dest (x. 0) < d, and dea (x, e) more generally be bounded in terms of spectral decay assumptions (see e.g. Proposition in Agarwal and Zhang [2022]).
Model-based Optimistic Posterior Sampling
We now describe algorithm MOPS for model-based RL in Algorithm 1. The algorithm defines an optimistic posterior over the model class and acts according to a policy generated from this posterior. Specifically, the algorithm requires prior po over the model class and uses optimistic model-error measure to induce the posterior distribution. We now highlight some of the salient aspects of our algorithm design.
Likelihood- dynamics prediction. At cach round, MOPS computes a likelihood over the space of models as defined in Line 4. The likelihood of model M includes two terms. The first term measures the squared error of the expected reward function R, in predicting the previously ob- served rewards. The second term measures the loss of the dynamics in predicting the observed states rh+1 given z and at each previous round s. For the dynamics, we use negative log-likelihood as the loss, and the reward and dynamics terms weighted by respective learning rates 1) and 1/. Prior works of Sun et al. [2019] and Du et al. [2021] use an integral probability metric divergence, and require more complicated Scheflé tournament algorithmically handle model fitting under total
4

---

variation unlike our approach. The more recent work of Foster et al. [2021] needs to incorporate the Hellinger distance their algorithm. In contrast, we directly learn a good model in the Hellinger distance (and hence total variation) as our analysis shows, by likelihood driven sampling.
Optimistic posterior updates. Prior works in tabular [Agrawal and Jia, 2017], linear [Zanette et al.. 2020] and model-free [Zhang et al., 2021, Agarwal and Zhang, 2022] RL make optimistic modifications to the vanilla posterior to obtain worst-case guarantees. and we perform a similar modification in our algorithm. Concretely, for every model M, we add the predicted optimal value in the initial context all previous rounds s to the likelihood term, weighted by a parameter y in Line Subsequently, we define the posterior using this optimistic likelihood. As learning progresses, the posterior concentrates on models which predict the history well, and whose optimal value function predicts large average value on the context distribution. Consequently any model sampled from the posterior has an optimal policy that either attains a high value in the true MDP M or visits some parts of the state space not previously explored, the model predicts the history reasonably well. Incorporating this fresh data into our likelihood further sharpens our posterior, and leads to the typical exploit- lr-le behavior that optimistic algorithms manifest in RL.
Given a sampling distribution p E A(M) (which is taken as the optimistic posterior in MOPS), and a time step h, we assume access to policy generator Tgen that takes parameters and p, and returns a policy Tgen (gen(h. X' + A (line 6). MOPS executes this policy up to a random time ht (line 7), which we denote as (a**t, "a"")""""""" ~ Tgen (h), The MDP then returns the tuple (a"s,a"t, مhر ..het)), which is used in our algorithm to update the posterior. The choice of the policy generator plays a crucial role in our sample complexity guarantees, and we shortly present a decoupling condition on the generator which is a vital component of our analysis.
Policy Generator.
For the examples considered in the paper (see Section 6). policy generators that lead to good regret bounds are given as follows.
Q-type problems: (h, follows a sample from the posterior, Tgen (h,p) TM, M ~ p- V-tyee problems with finite actions: Tgen (h. p) generates a trajectory up to x² using a single sample from posterior TM with M ~ p and then samples a' ~ Un)
V-type problems with infinite actions: Tgen p) draws two independent samples M, M' p from the posterior. It generates a trajectory up to x using TM. and samples a² using TMy (-1z²").
It worth mentioning that for Q-type problems, Tgen (h. p) does not depend on h. Hence we can replace random choice of h, by executing length H trajectories in Algorithm I and using all the samples in the loss. This modification of MOPS has better regret bound in terms of H dependency.
Main Result
We now present the main structural condition that we introduce in this paper, which is used to char- acterize the quality of the policy generator used in Algorithm 1. We will present several examples of concrete models which can be captured by this definition in Section 6. The assumption is inspired by prior decoupling conditions [Zhang, 2021, Agarwal and Zhang, 2022] used in the analysis of contextual bandits and some forms of model-free RL.
Definition 2 (Hellinger Decoupling of Bellman Error) Let distribution p E A(M) and a policy o (x*,, aa|x') be given. For any é >> O, E (0, 1] and h E [H), we define the Hellinger decoupling coefficient de (e,,,, a.) of an MDP M as the smallest number ch 2 0 so that for all xt:
EMap B(sh, ( (B(M, th ,ah) (c BB
e (* ,a²))
(aa) )* +
E M (a (alutaaarrtt (-
where (M, Ih, aa DH d²), P.( a²))² + (Z a²) - R. aaa))²
Intuitively, the distribution p in Definition 2 plays the role of our estimate for M*, and we seek a low regret for the optimal policies of models M ~ p, which is closely related to the madel-based Bellman error under samples x², a are drawn from n AM (the LHS of Lemma 1). The decoupling incquality relates the Bellman error to the estimation error of in terms of mcan-squared error of
TThis can be extended to more general experimental design strategy as we show in Appendix E.
5

---

the rewards and Hellinger distance to the true dynamics P.. However, it is crucial to measure this error under distribution of the data which is used under model-fitting. The policy æ = ggen (h, plays this role of the data distribution, and is typically chosen in manner closely related to p our examples. The decoupling inequality bounds the regret of p in terms of estimation error under x², a" π,, for all p, and allows to find good distribution p via online learning.
For stating our main result, we define a standard quantity for posterior sampling, measuring how well the prior distribution Po used in MOPS covers the optimal model M..
Definition 3 (Prior around true model) Given a > 0 and po on M, define
w(a, = inf [are In Inpo(Me))). <>0
where M(e) {MM E M sup sup 8th (M, a²) S }). with (M, aaa)
of hazhiah
KL((P.-|x7, (-")),,,,="))111111 al))|| aa))) + (RM("", a²) (art""a))
Definition implicitly uses model realizability to ensure that MMe) is non-emply for any é > O. However, note that our bound based on w(a, Po) can still be applied even if the model is mis- specified, whence the optimization over e naturally gets limited above the approximation error. Un- derstanding the dependence of such approximation error in the final bounds and instantiating it with various models of misppecification/orrapiin is an interesting direction for future work.
Before stating our main result, we state useful property of w(a,po) which illustrates its behavior: also says that for concrete problems, we can replace the KL ball by the better behaved Hellinger ball and only pay an extra logarithmic penalty.
Lemma 2. If M is finite, po is uniform over M, and M E M, then w(a, po) S. In |M for all > O. More generally, suppose p < Po with [|dP /dPollo S. B for any reference probability measure Po. and that .M admits an -oveerin under the metric e (given in Definition 2) for all h [H), of a size (c) at radius c. Then for any S 2/3 such that B 2 (og((B" /c). there exists a prior pó such that p6) & ar + log (</((6 log(B/v)) where v = </6 (6 log (6B² /e)).
The prior ph adds a small perturbutation to the dynamics in M in order to ensure the boundedness of the KL divergence P* Detailed choice of pó and definition of ((B) are based on the result comparing KL divergence and Hellinger distance in Sason and Verdú [2016, Theorem 9], and we provide a proof of Lemma 2 in Appendix C.
We now state the main result of the paper.
Theorem (Sample complexity under decoupling) Under Assumptions and 2, suppose that there exists such that for all p. dc (c. dc. Agen p). a) I dc (c, a). Define
de(e, a) - - H1
(1 (1-a)/a
de (d
If we take 17 = 1/6 and 7 I 0.5, then the following bound holds for MOPS: T
w(3HI.. + 2-T + HT [c + (1 - -2)0((20H70) o- de(e, a)) T-
Σ E Vx(xt) - E
vxr(())) Map.
To simplify the result, we consider finite model classes with po as the uniform prior on M. For a < 0.5, by taking 7 = min(0.5, [M/TT (InlM//T ddce,) a/H), we obtain
T
(=( = o min (H dc(e, a) In MM In MM| ) +H). (3)
T Σ v2(att) -
Mopp
T
We note that in Theorem 1, the decoupling coefficient fully characterizes the structural properties of the MDP. Once dc(c, a) is estimated, Theorem 1 can be immediately applied. We will instantiate this general result with concrete examples in Section 6. Definition 2 appears related to the decision estimation coefficient (DEC) of Foster al. [2021]. As expalined in Section 2, our definition is only needed in the analysis, and more suitable to posterior sampling based algorithmic design. The definition also related to the Bilinear classes model of Du et al. [2021], J, since the bilinear structures can be turned into decoupling results as we will see in our examples. Compared to these earlier results, our definition is more amenable algorithmically.
6

---

Proof of Theorem
We now give a proof sketch for Theorem As in prior works, we start from bounding the regret of any policy MM in terms of Bellman error term and an optimism gap via Lemma 1. We note that in the definition of Bellman error Lemma I, model M being evaluated the same model that also generates the data, and this coupling cannot be handled directly using online learning. This is where the decoupling argument is used, which shows that the coupled Bellman error can be bounded by decoupled loss. In the decoupled loss, the data is generated according to Tgen (4,, while the model being evaluated is drawn from Pt independently of data generation. This intuition is captured in the following proposition, proved in Appendix A.
Proposition 1 (Decoupling the regret) Under conditions of Theorem I, the regret of Algorithm I at any round t can be bounded, for any ş > O and > 0, as
E
(xtt)
E M-pi
(ar(a;)
S
E Mapp
|H (M he a (vM
H [ (1- (11-a)(())) a/0 (1-011))60 -000- dele, a) a))/(1-0))
The proposition involves error terms involving the observed samples (x a, r ot sf which algorithm controls via the posterior updates. Specifically, we expect the regret to be small whenever the posterior has a small average crror of models M ~ Pt, relative to M.. This indeed happens evidenced by our next result, which we prove in Appendix B.
Proposition 2 (Convergence of online learning) With n/ 1/6 and y O.5. MOPS ensures:
T
E
Mh (M, x/, "o ait)) [x s (3IIT.)) + 2yT.
Mrpt
Armed with Proposition 1 and Proposition 2, we are ready to prove the main theorem follows.
Proof of Theorem 1. Combining Propositions and 2 with μ = =0.3n/7 gives the desired result.
6 MDP Structural Assumptions and Decoupling Coefficients Estimates
Since Definition is fairly abstract, we now instantiate conerete models where the dccoupling co- efficient can be bounded in terms of standard problem complexity measures. We give examples of V-type and Q-type decouplings, a distinction highlighted in many recent works [e.g. Jin et al., 2021, Du et al., 2021]. The V-type setting captures more non-linear scenarios at the expense of slightly higher algorithmic complexity, while Q-type is more elegant for (nearly) linear settings.
6.1 V type decoupling and witness rank
Sun et al. [2019] introduced the notion of witness rank to capture the tractability of model-based RL with general function approximation, building on the earlier Bellman rank work of Jiang et al. [2017] for model-free scenarios. For finite action problems, they give an algorithm whose sample complexity is controlled in terms of the witness rank, independent of the number of states, and show that the witness rank is always smaller than Bellman rank for model-free RL. The measure is based on a quantity called witesssed model misfit that captures the difference between two probability models in terms of the differences in expectations they induce over test functions chosen from some class. We next state a quantity closely related to witness rank.
Assumption (Generalized witness factorization) Let F = [f(a, a, o, x'),, =r g(x, a, x') g E G}. with g(I, a, x²) E [0, 1). be given. Then there exist maps 16th (M,x*) and ud uM(M x²), and a constant K E (0, 1], such that for any context x. level h and models M, M' E M,, we have
(M, M', ,,,,')) S ((B"(M, x²), 3, (MMM,xx" x'), ,aa) - f) ,a)))))))),,
<. sup E IEF2 zh (x atammar (un)
E
(Bellman domination)
7

---

E where (M,M, M', ",,, h,xxt) 2 oe at spr last [M ,x", a We assume that I|ud (M, x x2)122 s
B1 for all M and x!.
Sun et al. [2019] define a similar factorization, but allow arbitrary dependence of f on the reward learn the full distribution of rewards, addition to the dynamics. We focus on only additive reward term, as we only need to estimate the reward in expectation, for which this structure of test functions is sufficient. The additional dependence on the context x! allows us to capture contextual RL setups [Hallak et al., 2015]. This assumption captures a wide range of structures including tabular, factored, linear and low-rank MDPs (see Sun et [2019] for further examples).
The bilinear structure of the factorization enables us to decouple the Bellman error. We begin with the case of finite action sets studied in Sun et al. [2019]. Let po" q be a non-stationary policy which follows т ~ p for h - I steps, and chooses a" ~ x (-1x") for n ~ q.
Proposition 3 Under Assumption 3. suppose further that |A] = K. Let us define 21 = x² 22 = M and x v/h (M, ) in Definition 1. Then for any c <> O, we have
dc(e, P. ggen (h, (), 0.5) S
4K R2
detf
K
where Tgen (h. Tgen
-p o Uni
B1
The proofs of Proposition 3 and all other results in this section are in Appendix E.
Sample complexity under low witness rank and finite actions. For ease of discussion, let dim(1/ (/5h) S for all h E [H] and [M] < co. Plugging Proposition 3 into Theorem 1 gives:
Corollary Under conditions of Theorem I, suppose further that Assumption 3 holds with dim(tb))
S d for all h E [H]. Let the model class M an action space A have finite cardi-
nality with |A| = K. Then MOPS satisfies 7 +I= V= SO S (V |Haekl IM]
(V
R²T
With standard online-to-batch conversion argument [Cesa-Bianchi et al., 2004], this implies a sample complexity bound to find an e-suboptimal policy of O [I la [MI) when the contexts are i..dd from a distribution. This bound improves upon those of Sun al. [2019],and Du al. [2021], who require Ô HH K²C² their algorithm, while we only use the discriminators implicitly in our analysis.
x²²²"
In TIMIF samples, where F is discriminator class explicitly used in
Factored MDPs Sun et al. [2019] show exponential separation between sample complexity of model-based and model-free learning in factored MDPs [Boutilier et al., 1995] by controlling the eπrrr of each factor independently. The gap is demonstrated by choosing discriminator class that measures the error separately on each factor in their algorithm. A similar adaptation of our ap- proach to measure the likelihood of each factor separately in the setting of Proposition allows our technique to handle factored MDPs.
Next, we further generalize this decoupling result to large action spaces by making linear embed- ding assumption that can simultaneously capture all finite action problems, as well as some more general settings [Zhang, 2021, Zhang et al., 2021, Agarwal and Zhang, 2022].
Assumption 4 (Lincar embeddability of backup errors) Let F = {f(x,a,r,z)) = r + g(x a, a.r') : E g). with g(x, a, x') E [o, 1), be given. There exist (unknown)) functions &h (h a and wh (M, f, xh) such that f E F, M M, E [H] and xh,aa".
(M,
(Ph d²) (P (x"", a²) = (w.(M, (M, x²),, d² da(" aa")).
We assume that (M, x2) S Bg for all M, and [H].
Since the weights " can depend on both the f and xA for finite action problems it suffices to For linear MDPs, the assumption holds with 6² being the MDP features and w being independent of h. A similar assumption on Bellman errors has recently been used in the analysis of model-free strategies for non-linear RL scaling to large action spaces [Zhang, 2021, Zhang et al., 2021, Agarwal and Zhang, 2022]. Note that Assumption 4 posits a pointwise factorization for each x² a and for individual test functions f E F, and in general is not directly comparable to Assumption 3, which assumes a factorization of the average error in backups the worst-case over all test functions. Assumption 3
choose 6x(z, a) = c. and w"(M, (M,,, = (x*", aa) (Ph f)" a)) 1-

---

allows us to decouple the distribution of ε from the model being evaluated, while Assumption 4 needed to decouple the choice of actions a". Obtaining a decoupling as per Definition requires using the two assumptions together.
We now state a general bound on the decoupling coefficient under Assumptions 3 and 4. Proposition 4 Suppose Assumptions 3 and 4 hold. For 6 in Definition 1, we define 21 = x² and 22 a², with same choices for w as Proposition 3. Then for any e. > 0, we have
dc (c P, (h, p), 0.25) & 16 16x== (v.
.2),
where Tgen (h, = p po
for any 61,62 > 0 satisfying Bres/x + 2B2c2 duar( 61)//=
Compared with Proposition 3, we see that the exponent a changes to 0.25 in Proposition 4. This hap- pens because we now change the action choice at step h to be from TM,, where M ~ p independent of M. To carry out decoupling for this choice, we need to use both the factorizations in Assump- tions 3 and 4, which introduces an additional Cauchy- -Shuchy step. This change is necessary as no obvious exploration strategy like uniform exploration of Proposition 3 is available here.
Sample complexity for low witness rank and (unknown) linear embedding, Similar to Corol- lary 1, we can obtain a concrete result for this setting by combining Proposition 4 and Theorem
Corollary 2 Under conditions of Theorem 1, suppose further that Assumptions and 4 hold, and [M < Co. For any 61 > 0, let der(v .<1) S dest (V,, E1) and den .2) S
( 62). for all h [H]. Then MOPS satisfies 7 DEE (*!x) V* (x!)
1/4 )
Bia
(2Barzdat(Vin) x
H
T
When the maps wha tt are both finite dimensional with dim(( S d, and dim(v² S d2 for all h, MOPS enjoys a sample complexity of O [8d do H1' in [M] in this setting. We are not aware of x4.² any prior methods for this setting. We note that similar models have been studied in a model-free framework in Agarwal and Zhang [2022], and that result obtains the same suboptimal e e rate as shown here. The loss of rates in T arises due to the worse exponent of 0.25 in the decoupling bound of Proposition 4, which is due to the extra Cauchy-Schwarz slep as mentioned earlier. For more general infinite dimensional cases, it is straightfowward to develop results based on an exponential or polynomial spectral decay analogous to the prior model-free work.
6.2 Q-type decoupling and linear models
We now give examples of two other structural assumptions, where the decoupling holds pointwise for all r and not just in expectation. As this is somewhat analogous to similar phenomena in Q-type Bellman rank [Jin et al., 2021], we call such results Q-type decouplings. We begin with the first assumption which applies to linear MDPs as well as certain models in continuous control, including Linear Quadratic Regulator (LQR) and the KNR model.
Assumption 5 (Q-type witness factorization) Let F be function class such that f(x, o, x²) T + g(I a, x²) for E G. with g(z, xx) E [0, 1], for all x, x'. Then there exist maps (x" aa)) and ud (M, f) and a constant K. > 0, such that for any h,
seh, ,,,,, E F and M e M:
((PM)) (aa))))))aa))))))))))))))))))( (P- faa",, a²)) 2 E(M,
f. th d'),,
sup E(M, f,xh, a²) 2 KEB((M, aa"), FEF
By for all M, f, h. Assumption 5 is clearly satistied by a linear MDP when d² are the linear MDP features and FF is any arbitrary function class. We show in Appendix D.1 that this assumption also includes the Kernelized Non-lincar Regulator (KNR), introduced in Kakade et al. [2020] as a model for contin- uous control that generalizes LQRs to handle some non-linearity. In KNR, the dynamics follow gh+11 W ">*" a²) + c, with e ~ (O, o² I), and the features s0(**, a²) are known and lie in an RKHS. In this case, for an appropriate class J, we again get Assumption with features &² = from the KNR dynamics.
with E(M, a²) Kab
(())(((>) (xr a²)) n (M, f)>l- We assume
Under this assumption, we gel an immediate decoupling result, with the proof in Appendix F.

---

Proposition 5 Under Assumption 5, let define a²)) ,x","") and = (a , a²)) in Definition 1. Then for any > O, we have
dc(c, P. geen (), - delt v
Pi 1 B1
where (h, = TM with M ~
x²
Here that the decoupling coefficient scales with the effective feature dimension, which now simultaneously captures the exploration complexity over both states and actions, consistent with existing results for Q-type settings such as linear MDPs.
Sample complexity for the KNR model. We now instantiate a concrete corollary of Theorem for the KNR model, under the assumption that x² Rdx, f0(x", a²) E R4 and || a²))l2 B for all w", al and h E [H]. A similar result also holds for all problems where Assumption 5 holds, but we state concrete result for KNRs to illustrate the handling of infinite model classes.
Corollary 3 Under conditions of Theorem 1, suppose further that we apply MOPS to the KNR model with the model class MKNR {W E Raxxdd [Wll I R}. The MOPS satisfies:
V-xx))
< O HV
Bdxx (og(R/a,BHT// To
Structurally, the result is a similar to Corollary 1, except that there is no action set dependence any more, since the feature dimension captures both state and action space complexities in the Q- type setting as remarked before. It is proved by using the bound on w Po) in Lemma 11 in Appendix D.I and the dccoupling coefficient bound in Proposition 7 in Appendix F with deat de-
We also do not make finite model space assumption in this result, as mentioned earlier. To apply Corollary 3 to this setting, we bound (T) in Lemma 11 in Appendix D.1. We notice that Corollary has slightly inferior d& dimension dependence compared to the de (dx+ +d3) scaling in Kakade al. [2020]. It is possible bridge this gap by direct analysis of the algorithm in this casc, with similar arguments as the KNR paper, but our decoupling argument loses an extra dimension factor. Note that unclear how to cast the broader setting of Assumption 5 in the frameworks of Bilinear classes or DEC.
Linear Mixture MDPs. A slightly different Q-type factorization assumption which includes linear mixture MDPs [Modi et al., 2020, Ayoub et al., 2020], also amenable to decoupling, is discussed in
Appendix D.2. For this model, we get an error bound o H
(3 Tea
Given our assump-
tion on value functions normalization by 1, this suggests suboptimal scaling in H factors [Zhou et al., 2021], because our algorithm uses samples from a randomly chosen time step only and our analysis does not currently leverage the Bellman property of variance, which is crucial to sharper analysis. While addressing the former under the Q-type assumptions is easy, improving the latter for general RL settings is exciting research direction.
7 Conclusion
This paper proposes general algorithmic and statistical framework for model-based RL bassd on optimistic posterior sampling. The development yiclds state-of-the-att sample complexity results under several structural assumptions. Our techniques are also amenable to practical adaptations, as opposed to prior attempts relying complicated constrained optimization objectives that may be difficult to solve. Empirical evaluation of the proposed algorithms would be interesting for future work. As another future direction, our analysis (and that of others in the rich observation setting) does not leverage the Bellman property of variance, which is essential for sharp horizon dependence in tabular [Azar et al., 2017] and some linear settings [Zhou al., 2021]. Extending these ideas general non-linear function approximation is an important direction for future work. More generally, understanding the guarantees can be more adaptive on per-instance basis, instead of worst-case, is critical for making this theory more practical.
Acknowledgements
The authors thank Wen Sun for giving feedback on early draft of this work.
10

---

References
Alekh Agarwal and Tong Zhang. Non-linear reinforcement learning in large action spaces: Struc- tural conditions and sample-effciency of posterior sampling. arXiv preprint arXiv:2203.0824. 2022.
Alekh Agarwal, Sham Kakade, Akshay Krishnamurthy, and Wen Sun. Flambe: Structural com- plexity and representation learning of low rank mdps. Advances in neural information processing systems, 2020.
Naman Agarwal, Brian Bullins, Elad Hazan, Sham Kakade, and Karan Singh. Online control with adversarial disturbances. In International Conference on Machine Learning. pages 111-119. PMLR, 2019.
Shipra Agrawal and Randy Jia. Posterior sampling for reinforcement learning: worst-case regret bounds. In Advances in Neural Information Processing Systems, pages 1184-1194, 2017.
Peter Auer, Thomas Jaksch, and Ronald Ortner. Near-optimal regret bounds for reinforcement learning. Advances in neural information processing systems, 21, 2008
Alex Ayoub, Zeyu Jia, Csaba Szepesvari, Mengdi Wang, and Lin Yang. Model-based reinforcement learning with value-targeted regression. In International Conference on Machine Learning, pages 463-474. PMLR, 2020.
Mohammad Gheshlaghi Azar, Ian Oshand, and Rémi Munos. Minimax regret bounds for reinforce- ment learning. In International Conference on Machine Learning, pages 263-272. PMLR, 2017.
Dimitri P. Bertsekas and John N. Tsitsiklis. Neuro-Dynamic Programming. Athena Scientific, 1st edition, 1996. ISBN 1886529108.
Craig Boutilier, Richard Dearden, Moisés Goldszmidt, al. Exploiting structure in policy construc- tion. In IJCAI, volume 14, pages 1104-1113, 1995.
Ronen Brafman and Moshe Tennenholtz. R-max-a general polynomial time algorithm for near- optimal reinforcement learning. Journal of Machine Learning Research, 3(0ct)):13-231,, 2002.
Nicolo Desa-Bianchi, Alex Conconi, and Claudio Gentile. On the generalization ability of on-line learning algorithms. IEEE Transactions Information Theory, 5009))2050-2057 2004.
Kurtland Chua, Roberto Calandra, Rowan McAllister, and Sergey Levine. Deep reinforcement leam- ing in a handful of trials using probabilistic dynamics models. Advances in neural information processing systems, 31, 2018.
Sarah Dean, Horia Mania, Nikolai Matni, Benjamin Recht, and Stephen Tu. On the sample complex- ity of the linear quadratic regulator. Foundations of Computational Mathematics, 20(4):6336679, 2020.
Simon Du, Akshay Krishnamurthy, Nan Jiang, Alekh Agarwal, Miroslav Dudik, and John Langford. Provably efficient rl with rich observations via latent state decoding. In International Conference on Machine Learning, pages 1665-1674. PMLR, 2019.
Simon S Du, Sham Kakade, Jason D Lee, Shachar Lovett, Gaurav Mahajan, Wen Sun, and Ruosong Wang. Bilinear classes: A structural framework for provable generalization in rl. Inter- national Conference on Machine Learning, 2021.
Dylan J Foster, Sham M Kakade, Jian Qian, and Alexander Rakhlin. The statistical complexity of interactive decision making, arXiv preprint arXiv:2112.13487 2021.
David Ha and Jürgen Schmidhuber. World models. arXiv preprint arXiv:1803.10122, 2018.
Assaf Hallak, Dotan Di Castro, and Shie Mannor. Contextual markov decision processes. arXiv
preprint arXiv: (50.::2202002299 2015.
Nan Jiang and Alekh Agarwal. Open problem: The dependence of sample complexity lower bounds on planning horizon. In Conference On Learning Theory, pages 3395-3398. PMLR, 2018.
11

---

Nan Jiang, Akshay Krishnamurthy, Alekh Agarwal, John Langford, and Robert E Schapire. Con- textual decision processes with low bellman rank are pac-learnable. In International Conference on Machine Learning, pages 1704-1713. PMLR, 2017.
Chi Jin, Zhuoran Yang, Zhaoran Wang, and Michael I Jordan. Provably efficient reinforcement learning with linear function approximation. Conference on Learning Theory, pages 2137- 2143. PMLR. 2020.
Chi Jin, Qinghua Liu, and Sobhan Miryoosefi. Bellman eluder dimension: New rich classes of rl problems, and sample-efficiet algorithms. Advances neural information processing systems, 2021.
Sham M. Kakade, Akshay Krishnamurthy, Kendall Lowrey, Motoya Ohnishi, and Wen Sun. Infor- mation theoretic regret bounds for online nonlinear control. In NeurIPS, 2020.
Michael Kearns and Satinder Singh. Near-optimal reinforcement learning in polynomial time. Ma- chine learning, 49(2):209-232 2002.
Jack Kiefer and Jacob Wolfowitz. The equivalence of two extremum problems. Canadian Journal of Mathematics, 12:363- 366, 1960.
processsnn Xiuyuan Lu and Benjamin Van Roy. Ensemble sampling, Advances in neural information processing systems, 30, 2017.
Horia Mania, Stephen Tu, and Benjamin Recht. Certainty equivalence is efficient for linear quadratic control. arXiv preprint arXiv:1902 07826, 2019.
Horia Mania, Michael Jordan, and Benjamin Recht. Active learning for nonlincar system identifi- cation with guarantces. arXiv preprint 1iXi:20006102777 2020.
Zakaria Mhammedi, Dylan J Foster, Max Simchowitz, Dipendra Misra, Wen Sun, Akshay Krish- namurthy, Alexander Rakhlin, and John Langford. Learning the lincar quadratic regulator from nonlinear observations. arXiv preprint arXiv:2010.03799, 2020.
Dipendra Misra, Mikacl Henaff, Akshay Krishnamurthy, and John Langford.
Kinematic state abstraction and provably efficient rich-observation reinforcement learning. arXiv preprint 1rXiv:/91.0058/ 2019.
Aditya Modi, Nan Jiang, Ambuj Tewari, and Satinder Singh. Sample complexity of reinforcement learning using linearly combined model ensembles. In International Conference on Artificial Intelligence and Statistics, pages 2010-2020. PMLR, 2020.
Anusha Nagabandi, Kurt Konolige, Sergey Levine, and Vikash Kumar. Deep dynamics models for learning dexterous manipulation. In Conference on Robot Learning, pages 1101-1112. PMLR, 2020.
lan Osband, Daniel Russo, and Benjamin Van Roy. (more) efficient reinforcement learning via posterior sampling. arXiv preprint arXiv:1306.099 0940, 2013.
lan Osband, Charles Blundell, Alexander Pritzel, and Benjamin Van Roy. Deep exploration via bootstrapped dqn. Advances in neural information processing systems, 29, 2016a.
Ian Osband, Benjamin Van Roy, and Zheng Wen. Generalization and exploration via randomized value functions. In International Conference on Machine Learning, pages 2377-2386. PMLR, 2016b.
M.L. Puterman. Markov Decision Processes: Discrete Stochastic Dynamic Programming. Wiley Series in Probability and Statistics. Wiley, 2014. ISBN 97811186222 URL https://books.. google. wtt.lee.comooooooeee cctp///botteiicceeBAAAAAAA..
Daniel Russo. Worst-case regret bounds for exploration via randomized value functions. arXiv
preprint arXiv:1906.02877, 2019.
Daniel Russo, Benjamin Van Roy, Abbas Kazerouni, Ian Osband, and Zheng Wen. A tutorial on thompson sampling. arXiv preprint arXiv:1707.00 02038, 2017.
12

---

Igal Sason and Sergio Verdú. f-divergence inequalities. IEEE Transactions on Information Theory, 62(11):5973,6006 2016.
Julian Schritwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, et al. Mastering atari, go, chess and shogi by planning with a learned model. Nature, 588(7833):604--699, 2020.
Ramanan Sekar, Oleh Rybkin, Kostas Daniilidis, Pieter Abbeel, Danijar Hafner, and Deepak Pathak. Planning to explore via self-supervised world models. ICML, 2020.
Max Simchowitz and Dylan Foster. Naive exploration is optimal for online lqr. In International Conference on Machine Learning, pages 8937-8948. PMLR, 2020.
Wen Sun, Nan Jiang, Akshay Krishnamurthy. Alekh Agarwal, and John Langford. Model-based in contextual decision processes: Pac bounds and exponential improvements over model-free approaches. In Conference learning theory, pages 2898-2933. PMLR, 2019.
William R Thompson. On the likelihood that one unknown probability exceeds another in view of the evidence of samples. Biometrika, 25(3-4):285 294, 1933.
Nicole Komczak-Jaegermann Banach-Mazur distances and finite fimite-dimensioonal operator ideals, vol- ume 38. Longman Sc Tech, 1989.
Masatoshi Uchara, Xuezhou Zhang, and Wen Sun. Representation learning for online and offline rl in low-rank mdps. arXiv preprint arXiv:2 10.00422 2021.
Sara van de Geer. Empirical Processes in M-estimation, volume Cambridge university press, 2000.
Roman Vershynin. High-dimensional probability: An introduction with applications in data science, volume 47. Cambridge university press, 2018.
Max Welling and Yee W Teh. Bayesian learning via stochastic gradient langevin dynamics. In Proceedings of the 28th international conference machine learning (ICML-I1), pages 681- 688. Citeseer, 2011.
Lin Yang and Mengdi Wang. Reinforcement learning in feature space: Matrix bandit, kernels, and regret bound. In International Conference on Machine Learning, pages 10746-10756. PMLR, 2020.
Andrea Zanette and Emma Brunskill. Tighter problem-dependent regret bounds in reinforcement learning without domain knowledge using value function bounds. In International Conference on Machine Learning, pages 7304-7312. PMLR, 2019.
Andrea Zanette, David Brandfonbrener, Emma Brunskill, Matteo Pirotta, and Alessandro Lazaric. Frequentist regret bounds for randomized least-squares value iteration. In International Confer- ence on Artificial Intelligence and Statistics, pages 1954-1964. PMLR, 2020.
Tong Zhang. From c-entropy to kl-entropy: Analysis of minimum information complexity density estimation. The Annals of Statistics, 34(5):2180-2210, 2006.
Tong Zhang. Feel-good thompson sampling for contextual bandits and reinforcement learning. ArXiv, abs/2110.0081 2021.
Weitong Zhang, Jiafan He, Dongruo Zhou, Amy Zhang, and Quanquan Gu. Provably efficient representation learning in low-rank markov decision processes. arXiv preprint arXiv:2106. 11355, 2021.
Dongruo Zhou, Quanquan Gu, and Csaba Szepesvari. Nearly minimax optimal reinforcement learn- ing for linear mixture markov decision processes. In Conference on Learning Theory, pages 4532-4576. PMLR, 2021.
13

---

Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's
contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes] See Section 7.
(c) Did you discuss any potential negative societal impacts of your work? [N/A] This work is theoretical nature on a basic question, relatively far removed from direct negative impacts.
(d) Have you read the ethics review guidelines and ensured that your paper conforms them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [Yes] See Sections 5 and 6.
(b) Did you include complete proofs of all theoretical results? [Yes] See Section 5 and Appendices A-F.
3. If you ran experiments..
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material or as URL)? [N/A]
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [N/A]
(c) Did you report error bars (e.g., with respect to the random seed after running experi- ments multiple times)? [N/A]
(d) Did you include the total amount of compute and the type of resources used (c.g., type of GPUs, internal cluster, or cloud provider)? [N/A]
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? [N/A]
(b) Did you mention the license of the assets? [N/A]
(c) Did you include any new assets either in the supplcmental material or as a URL? [N/A]
(d) Did you discuss whether and how consent was obtained from people whose data you're re using/curating? [N/A]
(e) Did you discuss whether the data you using/curating contains personally identifi- able information or offensive content? [N/A]
5. If you used crowdsourcing or conducted research with human subjccts...
(a) Did you include the full of instructions given to participants and screenshots, if applicable? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]
14


Reviews:
[{"review_id": "v3u5egM2a6i", "reviewer": "Reviewer_93tk", "summary": "The work proposed a unified framework for model-based RL by designing likelihood based posterior sampling methods. The theoretical framework unified many previously studied example scenarios and give state-of-art sample complexity in some cases. The authors also claim by combining variance reduction techniques, the sample complexity can be further improved.", "questions": "1. Extra T in Eq. (3)?", "limitations": "Yes.", "rating": 7, "confidence": 2, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The paper discusses a meaningful unified approach for model-based RL. The algorithm and results are well presented. The theoretical analysis looks solid.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "nYaiQLFPYMg", "reviewer": "Reviewer_ysuF", "summary": "The paper develops and formulates a statistical framework for model-based reinforcement learning with an optimistic posterior sampling mechanism. The paper develops a theoretical framework to solve the exploration-exploitation balance in an MBRL framework by incorporating an exploit-or-learn behavior in the learning framework- typical of optimistic posterior sampling algorithms. The paper draws an interesting connection between the regret of the optimistic PSRL algorithm with Hellinger distance and thereby derives a unified posterior sampling algorithm with state-of-the-art sample complexity under several structural assumptions.", "questions": "1. How does the bound in Definition 3 holds even under misspecified prior, can you please discuss that clearly?\n\n2. How $\\eta, \\eta'$ in 4 in Algorithm 1 is obtained, and how much impact will it have on the analysis?\n\n3. In [1], they show that the Bellman error is a poor proxy for the accuracy of the value function and the magnitude of the Bellman error is weakly related to the distance to the true value function. However, here the analysis has a very strong dependence on the Bellman error. Although it's a general issue and not specific to this research, please can you share your thoughts on the same.\n", "limitations": "The authors have mentioned the limitations of the current theoretical analysis including the need to perform empirical evidence, extension to non-linear function approximations in the Conclusion.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The primary strength lies in the statistical formulation of an optimistic model-based posterior sampling reinforcement algorithm and the connection of the regret with the Hellinger distance between the conditional probabilities to derive the state-of-the-art sample complexity for MBRL under structural assumptions. As done in optimistic algorithms, the algorithm assigns an optimistic bonus as done by $p_t(M) \\propto p_0(M)\\exp(\\sum_{s=1}^{t-1}(\\gamma V_M(x_s^1) + L_s^{h_s}(M))$ which means that the posterior would concentrate on model which predicts the history well or models under which the policy will lead to high value function. The novelty lies in the decoupling framework which builds the connection between the Bellman error and the Hellinger distance $D_h (P_M(\\cdot | x^h, a^h), P_\\star(\\cdot | x^h, a^h))$, with a dependence on the true dynamics. As it can be observed that it has a dependence on the true dynamics, hence it is critical to understand the confidence set creation to measure how\nwell the prior distribution covers the optimal model $M^*$ which is shown in Definition 3. Although, it needs a little more explicit explanation on how $M(\\epsilon)$ is non-empty for $\\epsilon > 0$ and also the bound still holds for misspecified prior. Overall the paper proposes a very interesting theoretical framework leveraging a Hellinger decoupling framework to derive the state-of-the-art sample complexity for optimistic PSRL-type algorithms and is also simplistic and unique in its formulation. The paper is well-written and the mathematical analysis is clear and precise. The main strength lies in the decoupling framework which is novel and very interesting and helps in controlling the Bellman error optimally. However, as pointed out by the authors it will be interesting to see the empirical evidence of the algorithm especially given that the theory has a strong dependence on the Bellman error and recent literature has shown that the Bellman error is a poor proxy for the accuracy of the value function [1]. \n\nReferences : \n[1]. Why Should I Trust You, Bellman? The Bellman Error is a Poor Replacement for Value Error by Scott Fujimoto, David Meger, Doina Precup, Ofir Nachum, Shixiang Shane Gu. https://arxiv.org/abs/2201.12417\n\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "UIVvFRbH_sM", "reviewer": "Reviewer_XV4J", "summary": "The paper proposes a general theoretical and algorithmic framework for model-based RL in which the authors use data likelihood to estimate a model quality and use an optimistic posterior sampling for data collection.  The combination of Heliginer distance-based likelihood and posterior sampling brings some advantages for their algorithmic framework as well as for regret analysis. They show that in many special cases, the sample complexity of their proposed algorihtms improves previous related works. \n", "questions": "It would be nice if the authors address my some questions:\n- What is the connection between the likelihood estimate at line 4 of the algorithm and the Hellinger distance defined in Definition 2 as the authors mentioned?\n- Does the coefficient $dc(\\epsilon, \\alpha)$ depend on the size of the action space?\n- Can the authors explain intuitively why their regret bound improves that of Sun et al. [2019] in V-type setting?\n- Are your results able to compare with those using the UCB approach in model-based RL? ", "limitations": "Yes", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Weaknesses:\n- No experiment is provided to demonstrate the efficiency of the proposed algorithm.\n- The framework is designed for only the posterior sampling approach.\n\nStrengths:\n- The paper is well written. \n- The technique of combining the likelihood and optimistic posterior sampling is novel. The technique allows avoiding the minimax problem which was required in the model-free work of Agarwal and Zhang [2022], and avoiding to incorporate the Hellinger distance in the algorithm which is complicated to estimate as in the work of Foster et al. [2021]. \n- In some special settings, the sample complexity of their proposed general algorithm improves that of some previous works.  ", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "So2g7rTJRc", "reviewer": "Reviewer_sfzn", "summary": "The authors propose a general algorithm for model-based RL based on optimistic posterior sampling. Theoretical analysis is provided to show that it achieves state-of-the-art sample complexity under some structural assumptions. ", "questions": "It's related with the weakness discussed above. Can the authors provide further discussions on how we can reasonably quantify (or even estimate) these quantities in order to better understand the sample complexity in some general settings.", "limitations": "Yes", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strength: The proposed algorithm is an interesting effort to solve the problems in model-based RL that collects data from the environment given the current learning state and defines the quality of a model given the current acquired dataset. The idea is to use the data likelihood as the estimate for model quality, and use an optimistic posterior sampling method to collect data. A so-called Bellman error decoupling method is used for the sample complexity of near-optimal policy, and it is sample efficient under the condition that decoupling coefficient of a model-based RL is small.\n\nWeakness: The assumptions in the sample complexity is somewhat stringent for the Hellinger decoupling. It's not clear under the general settings how to measure the Hellinger distance and Bellman Error to quantify such condition. It help better understand such condition if the authors can provide further discussions on how we can reasonably quantify (or even estimate) these quantities in order to better understand the sample complexity.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ["I don’t know"]}, "extra_text": {}}]

Author rebuttal:

Decision: Accept