Posterior Sampling Model-based Policy Optimization under Approximate Inference
Authors: Chaoqi Wang
Yuxin Chen
Kevin Patrick Murphy
Venue: Submitted to ICLR 2023 (2023)
Keywords: Reinforcement learning
Posterior
Model-based reinforcement learning

Abstract: Model-based reinforcement learning algorithms (MBRL) hold tremendous promise for improving the sample efficiency in online RL. However, many existing popular MBRL algorithms cannot deal with exploration and exploitation properly. Posterior sampling reinforcement learning (PSRL) serves as a promising approach for automatically trading off the exploration and exploitation, but the theoretical guarantees only hold under exact inference. In this paper, we show that adopting the same methodology as in exact PSRL can be fairly suboptimal under approximate inference. Motivated by the analysis, we propose an improved factorization for the posterior distribution of polices by removing the conditional independence between the policy and data given the model. By adopting such a posterior factorization, we further propose a general algorithmic framework for PSRL under approximate inference and a practical instantiation of it. Empirically, our algorithm can surpass the baseline methods by a significant margin on both dense rewards and sparse rewards tasks from DM control suite, OpenAI Gym and Metaworld benchmarks.

Under review as a conference paper at ICLR 2023
POSTERIOR SAMPLING MODEL-BASED POLICY OPTI- MIZATION UNDER APPROXIMATE INFERENCE
Anonymous authors Paper under double-blind review
ABSTRACT
Model-based reinforcement learning algorithms (MBRL) hold tremendous promise for improving the sample efficiency in online RL. However, many existing popular MBRL algorithms cannot deal with exploration and exploitation properly. Pos- terior sampling reinforcement learning (PSRL) serves as a promising approach for automatically trading off the exploration and exploitation, but the theoretical guarantees only hold under exact inference. In this paper, we show that adopting the same methodology as in exact PSRL can be suboptimal under approximate inference. Motivated by the analysis, we propose improved factorization for the posterior distribution of polices by removing the conditional independence between the policy and data given the model. By adopting such a posterior factorization, we further propose a general algorithmic framework for PSRL under approximate inference and a practical instantiation of it. Empirically, our algorithm baseline methods by a significant margin on both dense rewards and sparse rewards tasks from the Deepmind control suite, OpenAl Gym and Metaworld benchmarks.
surpass
INTRODUCTION
Model-based reinforcement learning has demonstrated great success in improving the sample effi- ciency of RL. However, many existing popular model-based algorithms (Kurutach et al., 2018; Chua al., 2018; Janner et al., 2019) cannot deal with exploration and exploitation properly, and hence may lead to poor performance when exploration is crucial. To trade off exploration and exploitation, most of the existing algorithms can be categorized into 1) optimism- baased (Jaksch al., 2010; Pacchiano (Strens, 2000; Osband et al., 2013; 2018;
et al., 2021; Curi et al., 2020); 2) posterior-sane posterior-sampling Fan & Ming. 2021); and 3) information- directed sampling (Russo & Van Roy, 2014) approaches.
As shown by Osband & Van Roy (2017), posterior sampling reinforcement learning (PSRL) can match the statistical efficiency (or regret bound) of optimism- based algorithms, but enjoys better computational efficiency. Information-directed sampling methods can be more statistically efficient when faced with complex information structure (Russo & Van Roy, 2014), but require estimators for the mutual information, which is difficult for high-dimensional random variables. Hence, in this paper, we focus on posterior sampling. For simplicity, we restrict attention to the episodic RL setting.
Under the PSRL framework, maintains a posterior p(M)Ds) of the Markov decision process (MDP) M given the observations De collected in real environment E. At the beginning of each episode, MDP is sampled from the posterior, and then compute the optimal policy x (MM) for the sampled model M. Equivalently, we can also view this policy as a sample from a "degenerate" posterior! over policies of the form p(r[De) = f MM )dM, where 8( (M) 8(mT
T (M)) is a Dirac delta distribution. This policy is then executed in the real environment to collect new data. Theoretically, such a simple strategy has been shown to achieve a Bayesian regret of Ö(VK) for K episodes (Osband et al., 2013; Osband & Van Roy, 2014). However, the theoretical guarantees only hold under exact inference, i.e., when we have access to the true posterior over models p(M (De), and when we can compute the optimal policy, which is very unlikely in practice.
common heuristic approximation PSRL (see e.g., Fan Ming (2021)) to replace the posterior over models with some approximation, such as Bayesian linear regression on top of the representations
For simplicity, we assume that there is only one optimal policy for each MDP.

---

Under review as a conference paper at ICLR 2023
learned by neural network. In this way, the resulting policy is then sampled from q ( D2) where we replace p(M[De) with the approximate posterior qf (M(Ds)
J ((M
At first glance, such a heuristic choice is natural it shares the same form as the true posterior p(z]De). However, in this paper, we prove that we can get lower regret if we replace the degenerate ((([M) with non-degenerate distribution of the form ( Dz). that depends the model as well as empirical data Dz; this is necessary to compensate for the fact that the posterior over models (MM))) may be suboptimal. By tuning the relative strength of the dependence of on DE and M, we can find sweet spot between maximizing the data efficiency and minimizing the effect of approximate inference error. Furthermore, such decomposition is guaranteed be no worse than the standard approach of using q" (P)
Building upon the above results, we come up with generic framework for PSRL under approximate inference. To implement the method in practice, we combine deep ensembles (Lakshminarayanan al., 2017) and Model-based based Policy Optimization (MBPO) (Janner et al., 2019). We also propose two different sampling strategies for policy selection that exploit posterior approximation. Empirically, our algorithm significantly outperforms the baselines on both dense reward and sparse reward tasks (Brockman et al., 2016; Tunyasuvunakool et al., 2020; Yu et al., 2020), Additionally, we also conduct various ablation studies to provide a better understanding of our algorithm.
In summary, contributions are
1. We conduct a rigorous study on how approximate inference affects the Bayesian regret in PSRL., showing that adopting the same methodology as in exact PSRL may be suboptimal when the true posterior unavailable (Section 2).
2. Motivated by analysis, develop generic framework for PSRL under approximate
inference as well as a practical version of it based on deep ensembles and (optimistic) sampling approaches for the policies (Section 3).
3. We present empirical results on DM control suite, OpenAI Gym and Metaworld benchmarks to demonstrate the efficacy of the proposed approach (Section 4).
2 PROBLEM STATEMENT AND THEORETICAL RESULTS
We start by introducing some notation, and summarizing prior work, before presenting our new theoretical results, which forms basis for our algorithm.
Notation. We consider the finite-horizon episodic Markov Decision Process (MDP) problem, of
which we denote an instance as M := AS, TM, PM, H, p}. For each instance M, and A denote the set of states and actions, respectively. TM S x A + [0, Rmaa]] is the reward function, PM is the transition distribution, / is the length of the episode, and is the distribution the initial state. We further define the value function of policy и under MDP M at timestep i
(s) E.M
VaM
Σ rM a:)l
where 8:41 (sls, a:) and a; x())s)). We define -* as the optimal policy for an MDP M if V.M (3) = max, Fu (s) for all s E S and i E [1, H.. We define the cumulative reward obtained by policy x over H steps sampled from model M as follows:
h
rr a.)
RM(T)= = EMM.
Σ
where 8.41 ~ PPM(s)s az) and a, ~ (a)s:), (2)
where the initial state is sampled from 81 ~ p(s).
Regret. For given MDP M, the regret is defined as the difference between value function of the optimal policy in hindsight and that the actual policy executed by the algorithm m
K
Regret(T. (T, a .M) = Σ
p(si) (VM), - V-M1 (s1) ds1,
(3)
2

---

Under review as a conference paper at ICLR 2023
where x² is the optimal policy for M,, and x² is the policy employed by the algorithm for ky episode. Correspondingly, the Bayesian regret is defined as the expectation the above regret, i.e.,
K
BayesianRegret(T,dd P = E [Regret(T, [f)
Σ
A
(4)
Here the expectation is taken over the prior distribution of dynamics models p(M) and the randomness in the algorithm of and environment.
PSRL. Posterior Sampling Reinforcement Learning or PSRL (Strens, 2000) serves as generic algorithmic framework for automatically trading off exploration and exploitation online RL. The core of PSRL is computing the posterior over MDPs (dynamics and reward models). Conditioned on the data Dg collected from the environment, we denote the posterior of the model as p (M)) Hence, the posterior distribution of policies is
p(a(De) = f (0000-p(TAMAA ((TIM))(M(DD)dMM where p (
=
(5)
where M(] is Dirac delta distribution, defined as ((T(M)
1. where (M)
arg max- RM (T) is the optimal policy for solving the MDP M. At the beginning of every episode, an MDP (or equivalently a policy) is sampled from the posterior distribution, and is then used for collecting new data. This simple algorithmic framework not only attains Bayesian regret of Ö(VK) (Osband et al., 2013), where K is the total number of episodes, but also enjoys better empirical performance than optimism-based methods bandits (Chapelle & Li, 201 I; Osband Van Roy, 2017). Nevertheless, the theoretical results only hold under exact inference. In the rest of this section, we discuss how performance is affected when we use approximate inference.
Bayesian regret under approximate inference. Let denote the approximate and true posterior distribution of polices at kes, episode by
q ( (D)) and (T( = f (0000-8((MMM////MM ((T
(6)
where DIS denotes all the data collected from the environment till the k-th episode. Next we characterize how approximate posterior inference affects the Bayesian regret.
Theorem For K episodes, the Bayesian regret of posterior sampling reinforcement learning algorithm of with any approximate posterior distribution qu episode k upper bounded by
K
V CCK(HRmax PH (T*) + 2H PAMAX
E
dRI. (qk (T()T)
where H(T*) is the entropy of the prior distribution of optimal polices, Le.. p(T) f [((m)M)((MMMMM C is some problem-dependent constant and dyz (-1.) is the KL-divergence.
Theorem (proved in the appendix) provides an upper bound of the Bayesian regret under approximate posterior inference. The first term is Ö(VK). The second term will be zero under exact posterior that should choose the
inference. However, when performing approximate inference, approximate posterior distribution g( (Dx) as "close" to the true distribution p( (De) as possible. A natural and common choice of q( 'De) would be the following one, which takes the same form as the true posterior
(m) = ((MM)
(8)
However, the following proposition shows that such a choice suboptimal.
Proposition 1 Under approximate inference (i.e., q(MiDe) * p(M(De)) the optimal M( may not be Dirac delta distribution, i.e., there exists other ((T(DE) such that
dxa (g(m (D)) dRa (4 (mT (D)) p(a(D=))
(9)
3

---

Under review as a conference paper at ICLR 2023
As illustration, provide the following example, which also serves constructive proof of Proposition
EXAMPLE 1. SUBOPTIMALITY OF (T([)M)
Consider toy setting, where the support set of MDPs is (M1 and the support set of policies Tı, ) Suppose that the true posterior distribution of MDPs is pl Mi (De) 1/3, p(M2) (gg) = 2/3, and the optimal policy per MDP is ((TT = and 1. This we get the following exact distribution
= p(malPe)
policies: p(T Dz) is
p( (D2)
5(m) M,]=1, (M))
((M,)P))
(10)
6(mz Mi)=0, M(malM2)=1
p(Ma|Pa)
((s))M))
p(M(De)
Now suppose we use the approximate posterior distribution over models, g(M, (Dg) 3) =
and (MM 1. We can optimize ((T[MM by minimizing dkı. (g(m (e) p( De)) One solution could be
(P)))))) (1) (6al a(Ma)Pe)=1] P2))
(1(]l/M))- ((x)
(m DE) =
(11)
((M()))
We see that the optimal ((..MM requires modeling uncertainty in the policy conditional on the model. By contrast, we adopt q" (л] (Dg) as our approximation, we will have
dxL
= log 3 = max dxL gEat
(12)
The above example tells that the (m² can perform arbitrarily poorly in terms of the KI. divergence. Given this observation, a natural follow-up question would be: is there a better choice other than od (r|De)? We provide answer in the next section.
3 METHOD
Motivated by the results the previous section, we first introduce a posterior decomposition which guaranteed to be better than q (x(De). Then introduce practical version of the algorithm built upon deep ensembles (Lakshminarayanan et al., 2017). Finally we propose two sampling approaches for encouraging efficient exploration.
3.1 POSTERIOR DECOMPOSITION
In section 2, we showed that (x)(D) is not favorable choice. This is because it assumes that, once M is given, the policy is determined. This motivates us consider the following more flexible posterior decomposition of (<((De),
q(m D2, - q(MMM DM(MM)))
(13)
Intuitively, such posterior dccomposition no longer assumes that the model can capture all the relevant properties of the data. We illustrate these two posterior approximations in Figure 1, The extra parameter E [0, 1] allows us balance the importance of fictitious data (Dx) from M and the real data (Dx) from the environment. In particular, we define
q( D,,MDDsAAA À 0) ( 8(xlMM
(14)
- ((T(De)
(15)
Thus when small, we trust our model more, whereas when À large we trust less. In the extreme case where À O, this framework reduces the degenerate posterior g (T D2). By adjusting X. we can find sweet spot in minimizing the effect of approximate inference error and maximizing the data efficiency. More formally, the following proposition illustrates the advantage of equation 13.
Proposition 2 By adopting the posterior decomposition of equation 13, we have min dya (q( De, A)) p(= < dxa. (()
(16)
4

---

Under review as a conference paper at ICLR 2023
Algorithm PSRL with approximate inference using Ensemble Sampling (PS-MBPO) Opti- mistic Ensemble Sampling (OPS-MBPO).
Require: Initialize an ensemble of dynamics models è (0., ) i.i.d. ~ g
Require:
Initialize an ensemble of policy networks + [f J. i.i.d. ~ q(Φ).
Require: Initialize empty datasets Dg and (DM J Jn.m=1 Real data vs. synthetic data ratio X.
1: for K episodes do
/ Dynamics training. (Line 2) 4/
2: Train the ensemble models on De under the objective in equation 17.
1 Policy sampling. (Line 3) */
3: Sample a policy from uniformly at random (equation 18) or based on the optimistic
distribution (equation 19).
4: Sample state 81 from the initial state distribution p(8)
5:
for / Data collection. (Lines 6-11) %/
steps do
6:
8/. rollout (oollu dynamics &, policy n, initial state -4-1, num. steps 1)
7:
Add to De
8:
Sample state 8 DE for each model policy do
9:
10:
D moollou(dynammie ,n, policy Pa,m. initial state s, num. steps R) Created mixed dataset D XD + (1 XyD
11:
b
/ Policy optimization (Line 12) */
12:
On.m = update-policy((o...... D, num. gradient steps G)
end for end for
14:
Update the optimistic policy distribution (equation 19).
16:
end for
3.5 le5
Lumulative Regret
A=
M
M
3.0
A=0.1
A=02
2.5
A-3.3
A=4.4
DE
DE
2.0
A
1.5
y-x
(a)
(b)
1.0
Steps
1e5
Figure 1: Graphical models for (a) the standard and (b) our posterior over policies π. Differences are shown red.
Figure 2: A comparison of cu- mulative regret for different X.
The above proposition informs us that we can minimize the KI. divergence, and hence the upper bound of the Bayesian regret, by carefully choosing the value of X. As empirical evidence, in Figure we plot the cumulative regret with different À in the cartpole-swingup environment with sparse reward. We that value of about À 0.4 X 0.5 is the best, and X 0 the worst in this case.
3.2 PROPOSED ALGORITHM
Building upon the above observations, we introduce a simple and general algorithmic framework for PSRL under approximate inference, which only differs from the standard PSRI. framework in the decomposition of policy posterior (see Algorithm in Appendix C). implement the method in practice, use deep ensembles (Lakshminarayanan et al., 2017) approximate the posterior distributions q {M) and q(T [,M Dz) with and & (see Algorithm 1). This is similar to ME- TRPO (Kurutach et al., 2018), 8), PETS (Chua et al., 2018) and MBPO (Janner et al., 2019), except also model the uncertainty over policies (i.e., ( De))
)) in addition to dynamics (i.e., (+))
By mixed dataset XD: (1 X)D M mean that for each data point in the training batch, it is with probability of À heing sampled from the real data D3 and probability of 1 from the fictitious data n
5

---

Under review as a conference paper at ICLR 2023
Dense reward tasks
Spurse reward tasks
Figure 3: We consider tasks from three benchmarks: OpenAl Gym, DM Control and Metaworld. These seven tasks cover both dense reward and sparse reward tasks.
In more detail, each member in the deep ensemble is conditional Gaussian distribution over outputs, characterized by its mean μ and variance a². For multi-dimensional predictions, we treat each dimension independently, and only predict the marginal mean and marginal variance for simplicity. Each ensemble member is trained independently by minimizing the negative log-likelihood,
po (v) = log (a(rr) (y ==)()²)²² + constant. (x)
(17)
2
We maintain N different dynamics models. For each such model. also compute M different policies; we use the soft actor-critic (SAC) (Haarnoja et al., 2018) method (see the appendix for details). The policy network nmm is updated based on synthetic data D generated by dynamics model and policy model m, as well as environmental data De generated by interacting the real world dynamics with sampled policy. See Algorithm 1 for the pseudo-code.
3.3 SAMPLING POLICIES
Ensemble Sampling, Given the posterior distributions, it remains to specify the sampling approach for policies. The simplest sampling strategy is uniform sampling the beginning of each episode.
(18)
In the case bandits (where N = 1, since there is no transition model), such a simple strategy has been shown to achieve a regret of Ö(VT + AAAM/) (Lu & Van Roy, 2017: Qin et al., 2022) for T steps in Gaussian linear bandits, where M is the size of the ensemble and A is the number of arms. This regret bound analysis tells us that adding more ensemble members will reduce the regret, although it's not clear how this theoretical result extends to the RL setting.
Optimistic Ensemble Sampling. Unfortunately, ensemble sampling may overly explore some unpromising regions, it treats each member in the ensemble model equally. This may lead to unnecessary or wasteful explorations. To reduce this, we propose an optimistic version of ensemble sampling which we call OPS-MBPO. Specifically, we keep track of the performance of each ensemble member in terms of the accumulated episodic return (alternatively, one could also use the value function for each policy (Agarwal & Zhang. 2022).) We then use this performance determine the probability of choosing each member, thus gradually discarding unpromising ensemble members. More precisely, at the beginning hit episode, we sample the policy from the following Boltzmann distribution, instead of uniformly at random,
Et Re RE(TT;,)) 1)/=,
exp
Pk(T (TT = mT)
(19)
N.M Ej=1
Et Re(T;;1)/71
exp
where 7 is temperature term for controlling the level of optimism, and Rg( 1) is the empirical cumulative reward of πं at the lu episode. When - oc, we recover uniform ensemble sampling, which we call PS-MBPO.
EXPERIMENTS
Our empirical evaluation aims to 1) verify the effectiveness of proposed methods: 2) offer deeper understanding about the mechanisms that are key to the improved performance; and 3) provide additional ablation studies non-key components. We start by introducing the experimental setup.

---

Under review as a conference paper at ICLR 2023
1e4 Halfcheetah
1e3
Ant
le3
Hopper
1ee
Walker2D
1.5
Average Reward
Average Reward
6
Average Reward
w
Average Reward
1.0
0.5
OPS-8D2
P5-HBPO
SLRD
NBPO
SAC
0.00
2
3
%
8.0
0.5
%
Steps
le5
Steps
Steps
le5
Steps
Figure 4: Comparisons on four tasks with dense rewards. The shaded region denotes the one-standard emror. The dashed green curve corresponds to the asymptotic performance of SAC at 3M steps. PS-MBPO improves over MBPO across all of the four tasks, and the improvement is more significant on Ant and Walker2D. OPS-MBPO achieves similar sample efficiency with PS-MBPO.
1e2
Ball-in-Cup
1e2 Cartpole-swingup
1e2 Window-open-v2
Average Reward
Average Reward
7.5
Average Reward
7.5
-10
5.0
OPS-MBPO
2.5
2.5
PS-MBPO
-20
MBPO
0.0
0.0
0.5 Steps
1.0
1.5
2 Steps
3
4
0.0
0.5
1.0
1.5
1e5
1e5
Steps
1e5
Figure 5: Comparisons three tasks with sparse rewards. PS-MBPO improves over MBPO, and OPS-MBPO further improves over PS-MBPO in terms of sample efficiency.
4.1 EXPERIMENTAL SETUP
We consider seven tasks from OpenAl Gym (Brockman et al., 2016), Deepmind Control Suite (Tun- yasuvunakool et al., 2020) and Metaworld (Yu et al., 2020). This includes four dense reward tasks (ant, hal f-cheetah, walk wer2d and hopper) where the agent receives an immediate reward at each step, and three sparse reward tasks (bal 1b - -----nuup ca cartpole e-sw ngup and window- oopoowo where the agent receives a reward only if it finishes the corresponding task. For sparse reward tasks, efficiently exploring the environment is more crucial than in dense reward tasks. See Figure: 3 for visualizations of the tasks.
For baseline methods, we consider the model-basee approaches SLBO (Luo et al., 2019), PETS (Chua al., 2018) and MBPO (Janner al., 2019), and model-free method SAC (Haarnoja et al., 2018). We compare each method in terms of the averaged episode reward, where each episode ends when the time step reaches 1, 000 or the agent reaches the terminal state. To draw more robust conclusions, we repeat each experiment with 10 different random seeds, and report the mean and the standard error. More details can be found in Appendix A and
4.2 COMPARISON WITH EXISTING METHODS
We report the results on the dense reward tasks Figure We see that our (O)PS-MBPO outperforms the baselines in all four tasks, including the MBPO method. (We verified that our implementation of MBPO comparable, or superior, to the performance of the original implementation; see appendix for details.) For example, hopper, method only requires roughly 40K iterations reach an average reward around 3,500, whereas MBPO needs around 150K steps.
We report the results on the sparse reward tasks in Figure 5. We first observe that both PS-MBPO and OPP-MBPO outperform MBPO across three tasks, and the improvement is more significant on Cartpole-swing OPS-MBPO further improves PS- PS-MBPO on these three tasks by a significant margin. This confirms the advantage of adopting an optimistic sampling strategy in sparse reward tasks.
and Window-open-
In contrast to Figure 4, also observe that
7

---

Under review as a conference paper at ICLR 2023
2
Figure 8: Visualization of the visited state space of PS-MBPO (top left) and MBPO (top right) at the initial stage of training. We also present two representative trajectories of PS-MBPO (bottom left) and MBPO (bottom right), which are four frames taken from the videos.
4.3 ABLATION STUDIES
In this section, we conduct a variety of ablation studies to better understand our proposed method. Additional experiments be found in the appendix, including ablation study with forced explo- ration (Phan et al., 2019), random function prior network (Osband et al., 2018) etc.
1e3
Walker2D
Does the gain come from posterior sampling ensemble? To assess the importance of posterior sampling, we compare perfor- mance where we sample policy from the posterior at each episode to the standard approach where we use the average policy, com- puted by averaging the distribution over actions across all ensemble members. The results are presented in Figure with namies networks and M (1,3, 5) policy networks. We see that when we disable the sampling procedure, the performance drops significantly, Another interesting observation is that, when with pos- terior sampling, the performance improves as we add more ensemble members (i.e., increase M). By contrast, when posterior sampling is disabled, increasing M seems doesn' 't improve the performance. This confirms that posterior sampling is the main factor behind the the improved performance, not just the fact that we have larger ensemble for both dynamics and policies.
Average Reward
M=1
3
5 dy-
2
%
2
Steps
1e5
Figure 6: Ablation study on the performance of with (solid curves) and without (dashed curves) the sampling step.
Effect of N and M. Since we are using the deep ensemble approx- imation, it is natural to wonder if we could get better performance by using a lager size of ensemble for dynamics models and policies.
Walker2d
5000
In Figure 7, we plot the average reward of the last 10 evaluations with 10 different random seeds, using N dynamics models and M policy networks for each dynamics model, by varying N and M {1, 2, 5). We see that increasing both N and M can improve performance, and that both forms of uncertainty (uncertainty on polices and dynamics) seem to matter.
M (T))
3
4500
4000
2
3 4
5
3500
N (M)
Figure 7: Average reward for varying number dynamics model (N) and policies (M).
Visualization of the state space. To gain insight into the explo- ration behavior. we project the high-dimen nigh-dimensional states of each tra- jectories collected by PS-MBPO and MBPO into a two-dimensional space using Umap (McInnes et al., 2018). The visualization of these embeddings can be found in Figure 8. We see that, at the initial phase, PS-MBPO (top left three figures) more explorative than MBPO (top right three figures), and that this also leads to better final performance (see Figure 5). In addition, we also plot two representative trajectories of PS-MBPO and MBPO for a qualitative comparison in Figure 8 at the same training iterations. Visually, move the robot arm to cover a more diversified regions, whereas MBPO fails to do so and the two trajectories look very similar each other (this is also reflected in the Umap visualization).
observe that PS-MBPO can
8

---

Under review as a conference paper at ICLR 2023
5 RELATED WORKS
Model-based reinforcement learning (MBRL). Research in MBRL mainly concerns how to learn the dynamics model given the data, and how approach for model learning is using the L2 loss for one-step transitions (Kurutach et al., 2018; Luo al., 2019; Chua et al., 2018: Janner et al., 2019; Rajeswaran et al., 2020), which is also equivalent to maximum likelihood estimation under a Gaussian assumption. In addition to the one-step training, Hafner et al. (2019); Asadi al. (2019); Lutter et al. (2021) show that multi-step training can further improves the prediction accuracy for long horizons, but with extra computational costs, which scale quadratically with the number of prediction steps. Another line of work focusses on the objective mismatch problem in MBRL (Ziebart, 2010; Farahmand et al., 2017; Luo et al., 2019; Lambert et al., 2020; Eysenbach et al., 2021), which modify the model training objective to provide performance guarantee for the induced policy in the unknown real environment.
the learned model. The most commonly adopted
Given the learned model, there are several ways to utilize it. Model predictive control (MPC) (Ca- macho & Alba, 2013) is derivative-free optimization method that has been adopted in many prior works (Nagabandi et al., 2018; Chua et al., 2018; Hafner et al., 2019; Fan & Ming. 2021; Lutter al., 2021). However, MPC is sensitive to the planning horizon and struggles with high-dimensional problems. As a mitigation, Kurutach et al. (2018); Luo et al. (2019); Janner et al. (2019) instead train policy on top of the model for amortizing the planning cost. Similarly, there are also works that utilize the model to facilitate learning of value functions (Feinberg et al., 2018: Buckman et al., 2018). However, few of these works have investigated the exploration and exploitation tradeoff.
Exploration and exploitation. Handling the exploration and exploitation tradeoff is the central problem in online learning. Typical methods can be categorized into the following three classes: 1) optimism-based (Jaksch et al., 2010; Pacchiano et al., 2021; Curi et al., 2020); 2) posterior-samplann based (Strens, 2000; Osband et al., 2013; Osband & Van Roy, 2014; Osband et al., 2018; Fan & Ming, 2021); and 3) information-directed sampling (Russo & Van Roy. 2014; Nikolov et al., 2019) approaches. Among them, optimism-baae based methods needs one construct the confidence set that contains the target model/policy with high probability, which suffers from scalability issues (Os- band & Van Roy, 2017); in addition this approach empirically performs worse than Thompson sampling (Chapelle & Li, 201 I). information-directt sampling can be better than optimism- based methods and Thompson sampling, as directly minimizes the "regret per information bit" (Russo Van Roy, 2014). However, it relies on estimating the mutual information between random variables, which is especially difficult for nigh-dimensionas continuous random variables (McAllester & Stratos, 2020). Therefore. focus posterior sampling. However, different from prior works, study the effect of approximate inference in an RI. setting.
SUMMARY AND FUTURE WORK
In this paper, we presented PS-MBPO and OPS-MBPO as two algorithms for efficient model-based reinforcement lcarning in complex environments. We demonstrate that both PS-MBPO and OPS- MBPO can greatly improve the sample efficiency in online reinforcement learning, and surpass various baseline methods by a large margin, especially on sparse reward tasks. We hope that, beyond our specific approach, our analysis can inspire future works to propose improved factorizations of the posterior over policies and models.
In the future, we would like to explore automatically adapting the value of in an online fashion. (See Section B.2 for some preliminary results.) In addition, we would like to extend the result of Phan et al. (2019) and prove a sublinear regret for PSRL. under approximate inference. It would also be interesting to explore epistemic neural networks (Osband et al., 2021) and transformers (Vaswani et al., 2017; Müller et al., 2021) as alternatives to deep ensembles for approximate posterior inference. Lastly, making information directed sampling (Russo & Van Roy, 2014) practical for reinforcement learning problems is another promising direction to further improve the sample efficiency.
Reproducibility Statement. We will open source the code for reproducing the results in our paper. For details about our experiments and algorithms, we encourage the reader to checkout Appdendix A for extended backgrounds and Appendix C for the hyperparameters used in our experiments.

---

Under review as a conference paper at ICLR 2023
REFERENCES
Alekh Agarwal and Tong Zhang. Model-based rl with optimistic posterior sampling: Structural conditions and sample complexity. arXiv preprint arXiv:2206. 200,,,22060,663, 2022.
Shipra Agrawal and Randy Jia. Optimistic posterior sampling reinforcement learning: worst-case regret bounds. Advances in Neural Information Processing Systems, 30, 2017.
Kavosh Asadi, Dipendra Misra, Seungchan Kim, and Michel L Littmaan. Combating the compounding- error problem with multi-step model. arXiv preprint 1rXiv 2019.
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary. Dougal Maclau- rin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018.
Greg Brockman, Vicki Cheung. Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojeiech Zaremba. Openai gym, 2016.
Jacob Buckman, Danijar Hafner, George Tucker, Eugene Brevdo, and Honglak Lee. Sample-efficient reinforcement learning with stochastic ensemble value expansion. Advances in neural information processing systems, 31, 2018.
Eduardo F Camacho and Carlos Bordons Alba. Model predictive control. Springer science & business media, 2013.
Olivier Chapelle and Lihong Li. An empirical evaluation of thompson sampling. Advances in neural information processing systems, 24, 2011.
Kurtland Chua, Roberto Calandra, Rowan McAllister, and Sergey Levine. Deep reinforcement learning in handful of trials using probabilistic dynamics models. Advances in neural information processing systems, 31, 2018.
Sebastian Curi, Felix Berkenkamp, and Andreas Krause. Efficient model-based reinforcement learning through optimistic policy search and planning. Advances in Neural Information Processing Systems, 33:14156-14177 2020.
Benjamin Eysenbach, Alexander Khazatsky, Sergey Levine, and Ruslan Salakhutdinov. Mismatched no more: Joint model-policy optimization for model-based rl. arXiv preprint arXiv:2110.02758. 2021.
Ying Fan and Yifei Ming. Model-based reinforcement learning for continuous control with posterior PMLR, 2021.
sampling. In International Conference Machine Learning, pp.
Amir-massoud Farahmand, Andre Barreto, and Daniel Nikovski. Value-aware loss function for model-base based reinforcement learning. In Artificial Intelligence and Statistics, pp. 1486-1494. PMLR, 2017.
Vladimir Feinberg, Alvin Wan, Ion Stoica, Michael Jordan, Joseph E Gonzalez, and Sergey Levine. Model-based value estimation for efficient model-free reinforcement learning. arXiv preprint 1rX1v:1803.00101, 2018
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with stochastic actor. In International conference on machine learning, pp. 1861-1870. PMLR, 2018.
Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In International conference on machine learning, pp. 2555-2565. PMLR, 2019.
Matt Hoffman, Bobak Shahriari, John Aslanides, Gabriel Barth-Maron, Feryal Behbahani, Tamara Norman, Abbas Abdolmaleki, Albin Cassirer, Fan Yang, Kate Baumli, Sarah Henderson, Alex Novikov, Sergio Gómez Colmenarejo, Serkan Cabi, Caglar Gulcehre, Tom Le Paine, Andrew Cowie, Ziyu Wang, Bilal Piot, and Nando de Freitas. Acme: research framework for distributed 2020.
reinforcement learning. arXiv preprint XrXiv:2006.000
10

---

Under review as a conference paper at ICLR 2023
Thomas Jaksch, Ronald Ortner, and Peter Auer. Near-optimal regret bounds for reinforcement learning. Journal of Machine Learning Research, 11(51):1563-1600 2010.
Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine. When to trust your model: Model-based policy optimization. Advances in Neural Information Processing Systems, 32, 2019.
Thanard Kurutach, Ignasi Clavera, Yan Duan, Aviv Tamar, and Pieter Abbeel. Model-ensemble trust-region policy optimization. In International Conference on Learning Representations, 2018.
Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles. Advances in neural information processing systems, 30, 2017.
Nathan Lambert, Brandon Amos, Omry Yadan, and Roberto Calandra. Objective mismatch in model-based reinforcement learning. arXiv preprint arXiv:2002.04523, 2020.
Xiuyuan Lu and Benjamin Van Roy. Ensemble sampling. Advances in neural information processing systems, 30, 2017.
Yuping Luo, Huazhe Xu, Yuanzhi Li, Yuandong Tian, Trevor Darrell, and Tengyu Ma. Algorith- mic framework for model-based deep reinforcement learning with theoretical guarantees. In International Conference on Learning Representations, 2019.
Michael Lutter, Leonard Hasenclever, Arunkumar Byravan, Gabriel Dulac-Arold, Piotr Trochim, Nicolas Heess, Josh Merel, and Yuval Tassa. Learning dynamics models for model predictive agents. arXiv preprint arXiv:2109. 14311, 2021.
David McAllester and Karl Stratos. Formal limitations on the measurement of mutual information. In International Conference on Artificial Intelligence and Statistics, pp. 875-884. PMLR, 2020.
Leland McInnes, John Healy, and James Melville. Umap: Uniform manifold approximation and projection for dimension reduction. arXiv preprint 1802.03326 2018.
Samuel Müller, Noah Hollmann, Sebastian Pineda Arango, Josif Grabocka, and Frank Hutter. Transformers can do bayesian inference. In International Conference on Learning Representations, 2021.
Anusha Nagabandi, Gregory Kahn, Ronald S Fearing, and Sergey Levine. Neural network dynam- ics for model-based deep reinforcement learning with model-free fine-tuning. In 2018 IEEE International Conference Robotics and Automation (ICRA), pp. 7559-7566. IEEE, 2018.
Nikolay Nikolov, Johannes Kirschner, Felix Berkenkamp, and Andreas Krause. Information- exploration for deep reinforcement learning. In International Conference on Learning Representa- tions, 2019.
Ian Oshand and Benjamin Van Roy. Model-based reinforcement learning and the eluder dimension. Advances in Neural Information Processing Systems, 27, 2014.
Ian Osband and Benjamin Van Roy. Why is posterior sampling better than optimism for reinforcement learning? In International conference machine learning, pp. 2701-2710. PMLR, 2017.
Ian Oshand, Daniel Russo, and Benjamin Van Roy. (more) efficient reinforcement learning via
posterior sampling. Advances in Neural Information Processing Systems, 26, 2013.
Ian Osband, John Aslanides, and Albin Cassirer. Randomized prior functions for deep reinforcement learning. Advances in Neural Information Processing Systems, 31, 2018.
Ian Osband, Zheng Wen, Mohammad Asghari, Morteza Ibrahimi, Xiyuan Lu, and Benjamin Van Roy. Epistemic neural networks. arXiv preprint arXiv:2107.08924 2021.
Aldo Pacchiano, Philip Ball, Jack Holder, Krzysztof Choromanski, and Stephen Roberts. Towards tractable optimism in model-based reinforcement learning. In Uncertainty in Artificial Intelligence, pp. 4413-1423, PMLR, 2021.
11

---

Under review as a conference paper at ICLR 2023
My Phan, Yasin Abbasi Yadkori, and Justin Domke. Advances in Neural Information Processing Systems, 32, 2019.
Thompson sampling and approximate inference.
Luis Pineda, Brandon Amos, Amy Zhang, Nathan O. Lambert, and Roberto Calandra. Mbrl-lib: A modular library for model-based reinforcement learning. Arxiv, 2021.
Chao Qin, Zheng Wen, Xiuyuan Lu, and Benjamin Van Roy. An analysis of ensemble sampling. arXiv preprint arXiv:2203.0/303, 2022.
Aravind Rajeswaran, Igor Mordatch, and Vikash Kumar. A game theoretic framework for model based reinforcement learning. In International conference on machine learning, pp. 7953-7963. PMLR, 2020.
Daniel Russo and Benjamin Van Roy. Learning optimize via information-direectd sampling. Advances in Neural Information Processing Systems, 27, 2014.
Daniel Russo and Benjamin Van Roy. An information-theoreti analysis of thompson sampling. The Journal of Machine Learning Research, 17(1):2442-247, 2016.
Malcolm Strens. A bayesian framework for reinforcement learning. In ICML, volume 2000, pp. 943-950, 2000.
Danii Tiapkin, Denis Belomestny, Daniele Calandriello, Eric Moulines, Remi Munos, Alexey Naumov, Mark Rowland, Michal Valko, and Pierre MENARD. Optimistic posterior sampling for reinforcement learning with few samples and tight guarantees. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho Advances in Neural Information Processing Systems, 2022a.
Daniil Tiapkin, Denis Belomestny, Eric Moulines, Alexey Naumov, Sergey Samsonov, Yunhao Tang, Michal Valko, and Pierre Menard. From Dirichlet to rubin: Optimistic exploration in RL without bonuses. In Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, 2022b.
Saran Tunyasuvunakool, Alistair Muldal, Yotam Doron, Siqi Liu, Steven Bohez, Josh Merel, Tom Erez, imothy Lillicrap, Nicolas Heess, and Yuval Tassa. dm_control: Software and tasks for continuous control. Software Impacts, 6. 2020.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
Tianhe Yu, Deirdre Quillen, Zhanpeng Hc, Ryan Julian, Karol Hausman, Chelsea Finn, and Sergey Levine. Meta-world A benchmark and evaluation for multi-task and meta reinforcement learning. In Conference on robot learning, pp. 1094 1100. PMLR, 2020.
Brian D Ziebart. Modeling purposeful adaptive behavior with the principle of maximum causal entropy. Carnegie Mellon University, 2010.
12

---

Under review as a conference paper at ICLR 2023
CONTENTS
Introduction
Problem statement and theoretical results
3
Method
3.1
Posterior decomposition
3.2
Proposed algorithm
3.3
Sampling policies
Experiments
4.1
Experimental Setup
4.2
Comparison with existing methods
4.3
Ablation Studies
5 Related Works
6
Summary and future work
Extended Backgrounds
15
A.1 Dynamics model
15
A.2 Soft actor-critic (SAC)
15
A.3 Revisiting MBPO, PETS and ME-TRPO
16
A.4 Posterior sampling reinforcement learning
16
Additional Results for (O)PS-MBPO
17
B.1
Visualization of the optimistic weights
17
B.2 How does the temperature term 7 affect the performance?
17
B.3 How does the schedule of À affect the performance?
17
Additional Details about the Algorithm and Experiments
18
C.I Algorithm Details
19
Implementation Details
19
C,3 Task Details
19
Forced Exploration
20
Random Function Prior
21
F Additional Umap Visualizations
23
G Proofs
24
G.1 Proof of Theorem
25
13

---

Under review as a conference paper at ICLR 2023
G.2 Supporting Lemmas
26
14

---

Under review as a conference paper at ICLR 2023
A EXTENDED BACKGROUNDS
A.1 DYNAMICS MODEL
We use deep ensemble for fitting the environment dynamics. For each network in the ensemble fo. it takes whitened state and action pair as input, and predicts the residual of the next state as well as the reward, i.e.,
Pa =N
Ao.
diag ( ]
0, (e. 0
)
fo
o
(20)
Fa
where M.s. pa are the empirical mean of the states and actions, d, and o are the empirical standard deviation of them. Then, the predictions of next state and reward will be
88+1|
( + Asi] N/
diag(o Be.). 0,
((((st)
o
(21)
Below is our implementation of each individual neural network in JAX (Bradbury et al., 2018).
class GausslanMLP (a.. Modulc):
"M"MLP with Gaussian distribution outputs, "**
self,
output_size int, hiaden_sizes, Sequence intl,
activation-jax.nt swish, min_logvar: float = -10.0,
max_logvar: float 2.0,
name: Optional [str] None,
() ().er()..nnit (name-name)
sclf. celffoouput_sizze = output_size
w_init hk. wag..ieerr..atiincceeeaaanggggi (1.0, ffaniin'
truncatedLnormal
self.mlp hk. .et. hidden_sizes, Li act activation-activati
activate_ saneeefiiaalaee (inal-True)
self logvar inp. inp.ones (cutput size) self.max self.max_logvar logvar iny. (output_ size)
min Logvar
max Logvar
logvar self delf.mean.d and logvar = hk. hk.Linear(
self.output_size 2, w_init-w_init,
mane="mean_and
def call_
x):
h self.mlppxx mean, legvar = hpp.splitt (self.mean_ant aeeeaaaananngeeeeeee (h), 2, axis=-1)
logvar logvaar
jax.nn, softplus
logvar)
logvar self.min logvar jax.nn. nn. sof.pplu (logvar self min logvar) return mean, Logvar
A.2 SOFT ACTOR-CRITIC (SAC)
We use SAC for learning the policies. In highlevel, SAC is maximum entropy RL algorithm, which typically optimizes the following objective,
( Σ E E (()((())))))
(22)
As result, maximum entropy RL algorithm will favor those policies that not only optimize for the reward, but also has large entropy. This can in turn improve the robustness of the optimized policy.
15

---

Under review as a conference paper at ICLR 2023
As for SAC, it searches the policy by iteratively solving the policy evaluation and policy improvement steps.
Policy Evaluation:
(s, a) + r(s, a) >pp a.a) a,a) (s){ V = Ba (s, a) log (a(a)s)];
(23)
(24)
Policy Improvement: 7++1 < arg min dg1. (T(( exp(O"
(25)
In the practical implementation of SAC, it uses a separate function approximator for the state value stabilize the training. Specifically, there are three components in SAC, a parameterized state value function V a soft Q-function Qe a), and policy To Ta(a))) The objectives for each component are
1 (Vo(s) - (1on) log (a)s)))
Jv(bb = ReND
(26)
2
1 (Qo(s.a) ((..
(0) (
(27)
2
(+) = EgwD [dkL (T {) (-())
exp(Qo(s) )))]
(28)
where (-) is normalizing constant, and (s a) defined as (().a) a) = r(s a) + (p(-laa)))))) [V*('')
(29)
Additionally, vp is the exponentially moving average of the weights the value network, and J= (Ø) be optimized with reparameterization trick under Gaussian case, which can further reduces the variance of the gradient estimator and hence stabilizes the training. We adopt the SAC implementation from Acme (Hoffman al., 2020).
A.3 REVISITING MBPO, PETS AND ME-TRPO
Popular model-based reinforcement learning algorithms such ME-TRPO (Kurutach et al., 2018), PETs (Chua et al., 2018) 8) and MBPO (Janner et al., 2019) typically repeat the following three steps: 1) train dynamics model (or an ensemble of models) {(M[Ds) 2) train/extract policy x* from the learned model; 3) collect data from the environment with the policy, Consequently, their policy is (approximately) equivalent to the one obtained by solving
/ RM(
n arg max [M[R = arg max TEll
(30)
where the posterior of the model M is approximated by ensemble of neural networks, IT is the search space of policies, and the cumulative reward (M( for an episode of length // of a policy π under dynamics model M is defined
Σ ( a.)
(T) = E
where
8.+1 (MM a.) and a, Talss))
(31)
However, the above strategy only accounts for exploitation, so will lead to low data efficiency.
A.4 POSTERIOR SAMPLING REINFORCEMENT LEARNING
The idea of PSRL is introduced by Strens (2000). The first regret bound Ö(HSVAT) for PSRL is proved by Osband et al. (2013) for a tabular case with S, A, T, H denotes the number of state, number of actions, number of time steps, and the length of each episode, respectively. In Osband Van Roy (2017), the bound was improved to Õ(H SAT). For the continuous case, Osband & Van Roy (2014) provides the first regret bound D(/dxde] based on the eluder dimension dg and Kolmogorov dimension dx. More recently, Fan & Ming (2021) study the regret bound for PSRL under Gaussian process assumption, and obtain regret bound of (HH32 In addition to the bound on Bayesian regret, there also a line of works studying the worst-case or frequentist regret bound for PSRL (Agrawal & Jia, 2017: Tiapkin et al., 2022b;a), and achieve regret bound of order Ö(VT). Nevertheless, all these regret bounds are derived under exact Thompson sampling.
16

---

Under review as a conference paper at ICLR 2023
Optimistic Weights Curve
Optimistic Weights Curve
Reward Curve
Figure 9: Visualization of the optimistic weights of the first 40K iterations (left) and during the entire training process (middle), and the reward curve (right) Hopper. Each chart in the left and middle figures corresponds to the weights each single policy.
Optimistic Weights Curve
Optimistic Weights Curve
Reward Curve
Figure 10: Visualization of the optimistic weights of the first 100K iterations (left) and during the entire training process (middle), and the reward curve (right) on Cartpole-Swa --Seingup. Each chart in the left and middle figures corresponds to the weights of each single policy.
B ADDITIONAL RESULTS FOR (O)PS-MBPO
B. VISUALIZATION THE OPTIMISTIC WEIGHTS
For OPS-MBPO, we will maintain the weights for each policy throughout the entire process of online learning. To investigate how those weights evolve, we plot the weights of each policy in Figure and Figure 10, which covers one dense reward task and one sparse reward task. The leftmost figure corresponds to the weights in the initial phase, which will change more rapidly than the later phases, We observe that the weights of some polices will first go up and than go down, and finally it will converge to a single policy. More interestingly, the reward curve in the rightmost figure is also consistent with the pattern in the optimistic weights curve.
B.2 How DOES THE TEMPERATURE TERM T AFFECT THE PERFORMANCE?
further study how does the temperature term will affect the performance on both dense reward and sparse reward tasks. The results can be found in Figure 11 and Figure 12. We observe that the temperature term will affect the convergence speed the reward on most of the tasks. For some of the tasks, such as Ant, Hopper and Walker2e cl, it will also affect the converged reward slightly. In general, we recommend the temperature term to be around five times of the best averaged episodic reward that can be achieved.
B.3 How DOES THE SCHEDULE OF À AFFECT THE PERFORMANCE?
Since plays a role in balancing the effect of approximate inference error and data efficiency, we are interested in how different schedules of À will affect the reward curve. We consider two schemes for adjusting X, i.e., 1) constant schedule; and 2) linear schedule. For constant schedule, we fix the value of À throughout training, whereas for linear schedule, decrease the value of X from 1 to 0 linearly. The rightmost figure in Figure 13 visualizes the difference between these two schedules. To make them comparable, ensure that the areas under both curves of the same size, that the
17

---

Under review as a conference paper at ICLR 2023
125 1e2
Halfcheetah
1e2
Ant
1e2
Happer
1e2
Walker2d
Average Reward
100
60
Average Reward
30
Avera flewaad
75
<- 300
40
300
300
300
1000
1003
1000
1000
50
3050
15000
20
20
5003
15020
5000
15000
20
5000
115000
25
50000
50000
50000
500000
152000
150200
150000
150000
%
0
8.0
0
Steps le5
Steps
le5
Steps
1e5
Steps 1e5
Figure 11: Ablation study on how the choice of the temperature will affect the performance on dense
reward tasks. All the experimental setups are the same as those experiments in our main paper, except for the temperature term.
10 1e2
Ball-in-Cup
1e2 Cartpole-swingup
1e2Window-open-22
Average Reward
Average Reward
6
Average Reward
-10
6
t= 10
4
T=
T=300
-15
3000
4
T- 100
1-100
T-5000
2
t= 1000
300
2
t= 1000
300
-20
50000
15000
5000
150000
8.0
%
-2000
1.0
1.5
2 Steps
0.5
1.0
1.5
Steps
1e5
1e5
Steps
1e5
Figure 12: Ablation study on how the choice of the temperature will affect the performance on sparse reward tasks. All the experimental setups are the same as those experiments in our main paper, except for the temperature term.
total amount real-world data the same. The comparison on three tasks are shown in the left three figures of Figure 13. We observe that the linear schedule has very little effect on the dense reward tasks, though slightly improves the the final reward in Hopper. For Cartpole-Sn swiiuususiguup.p.....oo nguu, the performance of linear schedule improves faster than constant schedule, but achieves similar rewards in the end. Nevertheless, we believe that there might be more sophisticated schedules for that can achieve better performance than the constant schedule, e.g., adapting the value of X based on the model's validation loss. For simplicity, we recommend to use a constant schedule in practice. For sparse reward tasks, the search range of X can be (0, 00,0.1, 0.3, 0.5, 0.7). and (0, 0.05, 0.15, 0.3) for dense reward tasks.
In addition to the grid search, believe it's also possible to adapt the value of À online. We cast the the problem of choosing the optimal value of x as bandit problem. The high-level idea of the algorithm is: 1) Initialize a of possible values for X, and treat each value of À as an arm in bandit. However, we haven't test this algorithm yet, and it would be interesting as future extension.
2) Apply any no-regret learning algorithms for solving it, e.g.,
C ADDITIONAL DETAILS ABOUT THE ALGORITHM AND EXPERIMENTS
Algorithm 2 PS-MBPO (abstract formulation)
Require:
Prior distributions q(M), q(T) and tuning hyperparameter
Require: Initialize an empty dataset De for storing data collected from the environment.
1: for K episodes do
2: Fit the posterior of the policy (T X) on data De using equation 13.
3: Sample a policy x² ~ (D, X) from the posterior distribution.
4: for H steps do
5: Run the policy x² in the environment and add the collected data to De.
6: end for
7: end for
18

---

Under review as a conference paper at ICLR 2023
le2
Ant
1e2
Hopper
1e2 Cartpole-swingup Constanc
À schedule
Average Reward
60
Average Reward
30
Averago Reward
6
Constant
Linear
Uieer
40
20
20
Constant Linear
10
Consrant Linear
0
10
20
0
10
20 Steps
30
Steps
Steps
1e4
1ed
Steps
Figure 13: Ablation study how the schedule of X affect the performance with PS-MBPO. All the experimental setups are the same as the experiments in main paper.
Uncertainty Propagation
Blocked
Model
Policy
(a)
(b)
(0)
Figure 14: An illustration of the differences between (a) MBPO, (b) PS-MBPO with = 0 and (c) PS-MBPO with X E (0. 1). MBPO adopt point estimation to the policy, which is obtained by MAP inference. Thus, the uncertainty propagation from the dynamics model the policy is blocked. For (b), it implicit assumes the dynamics model captures all the relevant properties of the data. For (c), we add a short cut from data directly to the policy, which controlled by the value of Hence, the policy further utilize the information in the data that not captured in the dynamics model.
C.I ALGORITHM DETAILS
In Algorithm 1, we approximate the posterior of MDPs and policies, i.e., g(M)Dg) and (x|M, De, using deep ensemble, which can be regarded as finite particle approximation to the posterior. Specifically, ((M[))) is approximated by {Ma I and g(T (M De, A) is approximated by {T&..... JM= for all n E [N], where both MM. and π T.... are implemented using
multi-layer perceptron (MLP) with parameters O, trained on De and Hn... trained on the mixed dataset XD: + (1 - X)D respectively. By mixed dataset XD + (1 A)) we mean that for each data point in the training batch, it is with probability of being sampled from the real data De
and probability of 1 - from the fictitious data D
C.2 IMPLEMENTATION DETAILS
In this section, we provide the additional details about our algorithm and experiments. We pro- vide a detailed description of our approach in Algorithm 1 and visual illustration about its difference with MBPO Figure 14. In terms of the hyperparameters, our choice of them are mostly the same the ones adopted in MBPO (Janner et al., 2019) and Pineda et al. (2021) for Ant, Halfcheetah, Hopper, Walker2D and Cartpole- swingup, and Eysenbach et al. (2021) for Window-open: Specifically, the hyperparameters of MBPO are directly adopted from https: /tt::///iiii /ggttub..../ com/ facebook weesearcc/mbrr--1
which are sufficiently optimized by the authors for MBPO.
for dense reward tasks. For sparse reward tasks, the hyperpa- rameters are adopted from Eysenbach et al. (2021). The hyperparameters for our method on each task reported in Table 1. We will also release our code for reproducing all the experiments. Next, we introduce the details about each tasks.
C.3 TASK DETAILS
Ant, Halfcheetah, Hopper and Walker2D, These tasks are taken from the official Github repository of MBPO (Janner al., 2019), https://gi httb.//gthhh.o..oo...... http//jnnnnnmmmmmmmmm rmpp..
19

---

Under review as a conference paper at ICLR 2023
1e4 Halfcheetah
1e3
Ant
Hopper
1e3
Walker2D
1.5
Average Reward
Average Reward
6
Average Reward
w
Average Reward
1.0
0.5
0P5P
5000
1SH0T0
0.00
%
8.0
0.5
%
Steps
le5
Steps
Steps
le5
Steps
Figure 15: Comparisons on four tasks with dense rewards including Halfcheetah, Ant, Hopper and Walker2D. MBPO" is the curve from the original paper by Janner et al. (2019). To be noted, in the original implementation of MBPO, they use networks for the ensemble of dynamics model, whereas our implementation only uses networks. But still, our implementation mostly reproduces their results and sometimes is even better.
Ball-in-Cup and Cartpole-swingup. These tasks are taken from the deepmind control suite (Tun- yasuvunakool et al., 2020). More details can be found in the Github repository at https: //github. com/ deepmind/ cim_contrll...
Window-open-v2. This task is based on the original Window open-wopen-eee v2 Metaworld bench- makr (Yu et al., 2020). The sparse reward is 0 only if the window is within 3 units of the open position, and 10 for other positions.
Table 1: Hyperparameters for each task. + y over episodes - b denotes segment lin- car function, f(k) = Imin(max(x: + (k ab/(b Cheetah, Walker and Window as abbreviations for Ball fall-inn-- iin--upp Cartpole SactppoeeSwiinng Swingup, Halfcheetah, WeDker2D and Window-open--oee so as to make the table fit the page.
(b))/(bbo0) - a) (y z), v)) We use Ball, Cart,
Ant
Bal11
Cart
Cheetah
Eopper
Walker
Window
Replay butfer capacity
10c
10a
105
10c
10a
106
10a
Episode length Number of episodes
1000
1000
1000
1000
1000
1000
250
300
150
400
300
125
300
600
Batch size for model
256
256
256
256
256
256
256
Model update frequency
250
250
250
250
250
250
250
Model Hidden dim.
200
200
200
200
200
200
200
Model #Hidden layers
Learning rate for model
10-3
10-3
10-3
10-3
10-3
10-3
10-3
Weight decay for model
10-4
5 10-5
10-5
10-4
10-1
10-4
10-1
#Model ensemble (N)
5
5
5
5
5
Validation ratio
20%
20%
20%
20%
20%
20%
20%
Rollout batch size
10
103
10
10a
10
10
5 x 10
- 25
+
Model horizon
over episodes
1
over episodes
20
100
20 100
Batch size for policy
256
256
256
256
256
256
Learning rate for policy
3e-4
3e-4
3e-4
3e-4
3e-4
3e-4
3e-4
#Policy per model (M)
5
5
5
5
5
Discount (1)
0.99
0.99
0.99
0.99
0.99
0.99
0.99
Target entropy
-0.05
0.05
-3
-3
Policy update frequency
Optimizer
Adam
Adam
Adam
Adam
Adam
Adam
Adam
0.05
0.5
0.5
0.05
0.05
0.05
0.5
7 OPS
50000
1000
5000
15000
15000
D FORCED EXPLORATION
Forced exploration is proposed in Phan et al. (2019) to improve approximate Thompson sampling for bandit problems. Without properly dealing with the approximate inference error, there will be an extra term in the regret that is linear in 7, regardless how small the error is. In their paper, they use
20

---

Under review as a conference paper at ICLR 2023
1e2
Walker2d
1e2
Hopper
1e2Window-open-2
60
Average Reward
Average Reward
30
Average Reward
-10
20
20
T=1
T=5
10
-20
T=
%
8.0
0.5
1.0
0.0
0.5
1.0
1.5
Steps
1e5
Steps
le5
Steps
1e5
Figure 16: Experiments of forced exploration on Walker2d, Hopper and Window-open--2.
The shaded region denotes the one-standard error. is the one without forced exploration.
60 le2
Walker2d
1e2
Hopper
1e2Window-open-22
Average Reward
Average Reward
30
Average Reward
40
-10
B=1
20
0.3
20
0.1
10
-20
%
8.0
0.5
1.0
0.0
0.5
1.0
1.5
Steps
1e5
Steps
1e5
Steps
1e5
Figure 17: Experiments with random function prior networks (RFPs) on Wa l ker2d, Hopper and
Wi diinddw-ppe ondow-open-v2. The shaded region denotes the one-standard error. 3 = 0 corresponds to the one without using random function prior networks.
the a-divergence for measuring the approximate inference error, defined as
Da 1- J-J (=))((()(()))>))))))))))))))))))
(32)
a(1 aa)
The a-divergence can capture many divergences, including forward KI. divergence (a + 1), back- ward KL divergence (a 0), Hellinger distance (a = 0.5) and x² divergence (a 2). Different inference methods will give error guarantee measured by a-divergence with different a.
We are interested in the error guarantee under the reverse KL-divergence, i.e., a sampling (Lu & Van Roy, 2017) provides error guarantees under the reverse Kl-divergence. In Phan al. (2019), they prove that forced exploration can make the posterior concentrate and hence restore the sub-linear regret bound, if the inference error is bounded by a-divergence with a < 0. The reverse Kl.-divereence falls in this case. Specifically, the forced exploration is simple method, that maintains probability of random exploration. This probability decays as t, the online steps, grows.
the ensemble
Though the above results only hold for bandit setting and it's unclear for reinforcement learning, we are interested in testing its empirical performance in RL. In our experiments, consider the following exploration rate
P& (random explore=True) = Bern(7/k),
(33)
where k is the index for the episode, and is the hyperparameter for controlling the frequency of foreed exploration. As & increases, the random exploration probability will decrease. In our experiments, consider E (1, 10}. All the other settings the same main paper. The results are presented in Figure 16. We observe that forced exploration is mostly not helpful in our experiments, except for the Hopper task. Moreover, increasing usually make the performance even worse. On the other hand, this may not be so surprising as the forced exploration is designed for approximate Thompson sampling in the bandit setting, and the result may not necessary generalize to the RL setting. We leave the theoretical analysis as a future work.
experiments in the
E RANDOM FUNCTION PRIOR
The random function prior (RFP) is proposed in Osband et al. (2018) for improving the uncertainty estimation. The prior network chosen for modelling the uncertainty that does not come from the
21

---

Under review as a conference paper at ICLR 2023
observed data. The RFPs also be viewed as regularization in the output space. In contrast to weight space regularization, RFP makes it easier to incorporate different property (e.g.. periodicity) of the function to be learned prior information. More importantly, when using deep ensemble, incorporating the RFP is fairly simple. It modifies the original training objective {fo, D) by adding an extra regularization term,
CRFP (fo. = l(fo D).
(34)
where is scaling term for adjusting the effect of the prior, fe, is the prior network which is held fixed during training. Hence, we also conduct experiments with RFPs our experiments to investigate how does the RFPs will affect the learning of dynamics models.
We vary the value of in 10.1, 0.3, 1}. The results are reported in Figure 17. Firstly, by properly choosing the value of B, RFPs slightly improve the performance performance a lot on Walker2) Howeee,,,,,,, for Window-opss hwwpw-ppen-y. -yp. RFPs will hurt the performance lot. We conjecture that this might because our choice of the prior function on Window findow--pen--x op is not suitable for the task, i.e., the reward is sparse in Window-open- but the RFPs don't t induce sparsity on the predictions.
Hopper. r, and don't affect the
Nevertheless, one interesting observation is that both forced exploration and RFPs seem to help on Hopper, r, and their overall pattern on three tasks a bit consistent. Therefore, it would be interesting to figure out if there is deep connection between the forced exploration and RFPs.
22

---

Under review as a conference paper at ICLR 2023
F ADDITIONAL UMAP VISUALIZATIONS
We provide the Umap visualization of the state embeddings of PS-MBPO and MBPO during training in Figure 18 and Figure 19. We observe that PS-MBPO will mostly cover more broad range of the embedding space, and its pattern also evolves more rapidly than MBPO.
PS-MBPO
MBPO
PS-MBPO
MBPO
PS-MBPO
MBPO
Figure 18: Visualization of the Umap embeddings of PS-MBPO and MBPO from consecutive training periods Hopper, Ant and Ha 1 fcheetah (from top bottom).
23

---

Under review as a conference paper at ICLR 2023
PS-MBPO
MBPO
PS-MBPO
MBPO
PS-MBPO
MBPO
Figure 19: Visualization of the Umap embeddings of PS-MBPO and MBPO from consecutive training periods on Walker2d, hall-in-Cupp and Cartpole-swingee (from top to bottom).
G PROOFS
In this section, we present the proof of Theorem 1. The proof of this theorem is inspired by the techniques in Russo & Van Roy (2016), with some additional modifications to extend the results from bandit setting to the reinforcement learning setting.
24

---

Under review as a conference paper at ICLR 2023
G.1 PROOF OF THEOREM
Theorem For K episodes, the Bayesian regret of posterior sampling reinforcement learning algorithm of with any approximate posterior distribution qu at episode k upper bounded by
K
(C max PH (x*) + 2H R
E
dAL (qk (T) Ph (T)
TELAX
where H(TT)) is the entropy of the prior distribution of optimal polices, i.e, p(α)) I M(MMMM C is some problem-dependnn constant and dga. (-1.) is the KL-divergence.
Proof: Recall the definition of Bayesian regret.
K A
BayesianRegrea(T,e a .(MM)) = E [Regret(TT M)] E
(35)
Let's denote history the beginning of episode k as Hg. Then, can rewrite the Bayesian regret as BayesianRegret(,, ot
K
E
[
(36)
By doing so, bound each term E/A,| H&] separately. For convenience, define E₁ =
E[A,|
Then, by Lemma 1, we can further decompose it into,
E|
Gx + Dr:
(37)
where
GN = / Vau(T)p((T) (Ex [V+((= = -E (si)) dv
(38)
and
D& = / -VPalT)-Vaaa (VPn(T)BB (s:)( = + v (a (T))E VM (81)) dr.
(39)
Then, it remains to bound Σ E B(Ga] and Σk E[Da). By Lemma bound the sum of expectation of D3 by
B[Dx] 2
Raxx
E |kL (4kl pk]]
(40)
By Lemma 3, the upper bound for the sum the expectation of G3 is
Σ LEGG] S VCK (22
(41)
Hence, the term D3 captures the regret incurred by the approximate inference error, and Gg captures the standard regret for Thompson sampling, which of order Ō(VK) By combining them together, we finally arrive at the upper bound of the Bayesian regret
BayesianRegret(T:, o .(M)) SVCK (C Rmax
(H())*)
+ 2H R
E
dxL. (gx(T)) Pa (T))
(42)
25

---

Under review as a conference paper at ICLR 2023
G.2 SUPPORTING LEMMAS
Lemma For each time k = .,., K, we have
E [ = E (81) - Va (8)))(K) = En (v1) (81) = GR + Da. (43)
where
Ga = / ((((F))a(TT) (Ex [811(= (B = - -Ba [Vr.((1)]] dr
(44)
and
Dje == (VPa(T) - VT((T) (VPL(T)E& [V., = -a] + v 2x (T))E [V2,1 (s1)) dx.
(45)
Proof: Conditioning on the history Ha, we can write the conditional Bayesian regret as VM x (81)
Ex
(81)
VM (81))
(46)
f Pk (T) [V (51)) = π dv - (T) [V (31)| π| dm
(47)
= (T) [V (x = π| dr - /a(TT)E [V (s1)]
dr
(48)
G Dr.
(49)
where the second equality holds because the value function is independent of the instantiation of the policy T when given the history Hg.
Lemma For any ...,, K, we have
K
E DDd
S 2HRmaas
v E dxx (gx1ps)]
(50)
Proof:
Recall DR.
Dr = / (VPa(T)- - (VPL(T)E& (()(())) = =x] (xx (V%((a))) dx (51)
By using the Cauchy-Schwarr inequality, we have
D& - (Vpa(T)) V(((T)) dr
(W/ Pe
= =]]² dm + /
(T) [81]] dm
(52)
By the definition of Hellinger distance du (-1-) between two random variables, we have
D& du (gul (gelpe) (V/ p& (T) [VM = m] dr
( (81) dr
(53)
Since [dH (-1-)12 < dkL (-1-) and ViM is a bounded random variable with H Rmax as its upper bound, we have
Dr 2dh (9xl (p) H Rmax < 2 dKL (gx PH
(54)
Hence, we have
B[Dx] 2H R
12 [KKL (gk) Px]]
(55)
26

---

Under review as a conference paper at ICLR 2023
Lemma 3 For each lo 1,... K, we have
K
BIGa S VCK
(56)
Proof: Recall the definition of Gk.
Gk I Vul(T))p((T) (Ex (81)) = = -E E& (x1)) dr.
(57)
Since V.1 (here, drop the dependency 81 for clearesss) is bounded random variable, and more specifically, it's ((ax))))) (2)-sub-Guussia Hence, by Lemma 4, the following holds,
H. Rmax
Bu [ =x] E& [V=1]]
V (Pn(V = T) (v (58)
This gives us that
Gk S f (u((T)p((T)
Bmax
22kk. (VAM +)] pa (V+)) dr.
(59)
2
Then, we can further rewrite the KL-divergence using the conditional mutual information Bg (+: (i.e.. conditioning on the history Hş),
/f qu: (T) (T' (KK) (D&(VMM = px))) drdw' / (T))4( V.M dT
(60)
When conditioning on the history Hę, the optimal policy nt and M is independent of the x hence we have
VM (ddd = qx((T))))) (T(T (T'; V.M 1 = x)dn = I<(T=; V. I**) / Q (T)T& VVM
(61)
By the fact that x* is jointly independent of nt,1 and x when conditioning on the history H3. hence
we have
A VM m²) (TT"; V + Ia(==; : By the chain rule of mutual information, finally get
(62)
(m<; V.M 7
Irr) Te((r (+"))====>>)))))""" Tu(mm (V= (VM,i.=*))
(63)
Now, let's define the following function Is and C,
In (n, (T, T') == v qu (T)P (Ex [VM x1 -E, [Val])
(64)
( gT. x)
C = max
(65)
KEZ, [5 T' (T) doTdd"
Thus, we further have
H Rmaxx /f (T drddd'.
2
Te(m" (V- x x))
2
(66)
On the other hand, we can rewrite G5 as
G+ = f 9h (T,, ())..
(67)
By rearranging the terms, we get
G
((H 12/2) (ʃ ga (), If (,,(","')|" m'Il drdd'
C ((Hxx
(68)
1 ( V. ()
27

---

Under review as a conference paper at ICLR 2023
Hence,
G. S
(69)
Hence, we have
K
K
Σ G
S
E
V (CH (iax) I
(TT: (TTh
(70)
R=11
K
= =VC((H /2)2)
E
1 (TT (o V [)]]
(71)
k=1
K
mIH 1/22) )2)
KE
( (Tt v2(
(72)
K ( (Tr v2(,))
VCK ((H )2/2)
(73)
S VVC ((H 12/2) (TT),
(74)
where the last inequality holds because of the chain rule of mutual information
K
I (T= (x*h, Vila
E
I(
VM x
(75)
E [(TT's (TT V-M
Vo T², V 2 x K (xK,.))
(76)
H(T²) (= VA Vin VM, 12,
K
7 1))
(77)
H(r*).
(78)
Remark 1 When the number of policies [II] is finite, and the value function V.M. is linear with its parameter lives in Rd, then C can be upper bounded by d, i.e., C d.
Lemma 4 (Russo & Van Roy (2016)) Suppose that there is Hy-measurable random variable such that for each π E II, V. is a suuu-Gaussiin random variable when conditioned on Hy. then for every n. n² E I1, the following holds with probability 1,
VV |T = T'] - E& [V S my 2dke (Pa(VaM|T** = x²) (V2,))
(79)
28


Reviews:
[{"review_id": "yk8O-Fibsak", "reviewer": "Reviewer_E1oc", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "The paper proposes an approximate version of posterior sampling reinforcement learning (PSRL). Empirical results show that the proposed algorithm outperforms the chosen baseline methods in some benchmark tasks.", "summary_of_the_review": "As mentioned above, both analysis and experiments need to be improved.", "strength_and_weaknesses": "The paper attempts to provide regret guarantees for approximate versions of PSRL, and proposes one simple version motivated by the analysis. To understand the proposed algorithm better, the authors conducted various ablation studies.\n\nThere are several weaknesses in the analysis. First I am not sure whether the regret bound in Theorem 1 is correct. For example, why the inequality (74) holds? I guess the chain rule is used here, but I do not see how to use it. Now suppose the regret bound is correct. It is not discussed in the paper how the first term on equation 7 depends on S, A, H. For example, the entropy of the optimal policy could be pretty large; also how the constant C depends the problem parameters? In addition, it seems that the paper does not try to bound the second term on equation 7 for the proposed algorithm. I think the bound on the second term is really needed; otherwise, I could not see how tight the regret bound in terms of the problem parameters. \n\nIn terms of experiments, I am not sure why other approximate versions of PSRL are not included as baseline methods, for example, bootstrap DQN [1], RLSVI [2], hypermodels [3], hyperDQN [4]. Although most of them are value-based methods, I think they still need to be included in the experiments for comparison.\n\n[1] Ian Osband, Charles Blundell, Alexander Pritzel, Benjamin Van Roy, Deep Exploration via Bootstrapped DQN, NeurIPS 2016.\\\n[2] Ian Osband, Benjamin Van Roy, Daniel J. Russo, Zheng Wen, Deep Exploration via Randomized Value Functions, JMLR 2019.\\\n[3] Vikranth Dwaracherla, Xiuyuan Lu, Morteza Ibrahimi, Ian Osband, Zheng Wen, Benjamin Van Roy, Hypermodels for Exploration, ICLR 2020.\\\n[4] HyperDQN: A Randomized Exploration Method for Deep Reinforcement Learning, ICLR 2021.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "UwvongPbKsM", "reviewer": "Reviewer_eWBV", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "The present paper is concerned about posterior sampling reinforcement learning (PSRL), which is an extension of Thompson sampling to RL; an MDP is sampled from the posterior distribution of MDP given data, the optimal policy for the sampled MDP is computed and is used to collect more data. The performance of this algorithm is analyzed in terms of Bayesian regret.\n\nThe authors are concerned about the case where only an approximate inference is feasible, where an approximate posterior $q$ is substituted for the true posterior. The authors first show that in this case there could be a better alternative than utilizing the optimal policy for the MDP sampled from the approximate posterior in Proposition 1. Motivated by this observation, the authors propose to derive the policy not only from the sampled model but also from the raw data with hyperparameter $\\lambda\\in[0,1]$ balancing between them. In the proposed algorithm, $\\lambda$ is used to construct a mixed dataset, where the original data and the synthetic data are mixed in the ratio of $\\lambda:1-\\lambda$. The authors provide two approaches to sample policies from the posterior.\n\nThe authors empirically evaluate the performance of the proposed method. The result shows that the proposed algorithm achieves better performance than the baseline algorithms, especially for the tasks with sparse rewards. The authors conduct ablation studies to understand the mechanism of the proposed method.", "summary_of_the_review": "I would like to suggest to accept this paper, given that the proposed method is well motivated and is well investigated empirically. But I would encourage the authors to include more detailed discussion on how to decide $\\lambda$ before deployment.", "strength_and_weaknesses": "## Strengths\n- The proposed method is well-motivated from a theoretical perspective. Although the theoretical finding is simple, it is easy to understand and is actually useful to develop the proposed method.\n- The proposed method achieves solid improvement over the baseline methods.\n- The paper is well-organized and is easy to read.\n- The additional results in Section 4.3 are insightful to understand the behavior of the proposed method.\n\n## Weaknesses\n- Optimal value or scheduling of $\\lambda$ is not clear, as the authors discuss in the future work section. In real applications where it is costly to deploy an agent to the environment, manual tuning of $\\lambda$ should be avoided. Although I don't ask the authors to completely address the issues, I would like to ask the authors to discuss how to roughly decide $\\lambda$ without deploying the agent to the environment. Will the quality of the posterior be relevant?", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "QTzgnPEXaKb", "reviewer": "Reviewer_PkUu", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "The authors study model-based reinforcement learning in episodic Markov decision process (MDP) and in particular the exploration-exploitation dilemma.\n\nPrecisely they consider Posterior sampling reinforcement learning (PSRL) under approximate inference. They show that PSRL can be sub-optimal under approximate inference.  To solve this issue they propose an improved factorization of the posterior over the policy where the sampled policy is not anymore an optimal policy of a sampled MDP  from the posterior but a sample from a distribution that depends on the  sampled MDP and the observations collected in the real environment. \n\nThey provide a general algorithmic framework under approximate inference and a practical instantiation of it. They also show that empirically the proposed algorithms improve over various baselines on both dense and sparse rewards task from the Deepmind control suite, OpenAI Gym and Metaworld benchmarks.", "summary_of_the_review": "See above", "strength_and_weaknesses": "Strengths:\n- PSRL sub-optimal under approximate inference.\n- New posterior factorization for PSRL under approximate inference.\n\nWeaknesses:\n- Description of PS-MBPO and OPS-MBPO.\n- Intuition behind the posterior factorization not completely clear.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "MmergoCbdSl", "reviewer": "Reviewer_m7LJ", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "This paper considers the problem of approximate posterior sampling in model-based policy optimization. The main contribution of this paper is a novel algorithm called PS-MBPO that utilizes the structure within the approximated posterior. By decomposing the posterior into a posterior on the model $\\mathcal{M}$ and a delta function indicating the procedure of choosing $\\pi$ according to $\\mathcal{M}$, the authors proposes to use an alternative function that interpolates between the delta function and the posterior of $\\pi$. The proposed algorithm, PS-MBPO, together with its optimistic variant, OPS-MBPO, shows improvements on tasks with both dense rewards and sparse rewards. Moreover, OPS-MBPO outperforms PS-MBPO more significantly in sparse reward tasks.\n", "summary_of_the_review": "This paper provides strong empirical validations of the proposed algorithm, but the theoretical part is not written very clearly. Still a solid contribution.", "strength_and_weaknesses": "The intuition is fairly reasonable, introducing a parameter that extends the hypothesis class reduces the loss, so I trust the effectiveness of the design and the superiority of the empirical performance.\n\nHowever, I think the procedure of the algorithm has not been explained properly, and the notations are disorganized.\n\nDespite the strong empirical justification, I still find the following parts somewhat hindering the understanding of this work. Especially, some notations throughout the methodology and Algorithm 1 is misleading.\n\n- Throughout the paper, the authors are adopting the notation “$q$” to represent both the dynamics models $q(M\\mid D_M)$ and the policy networks $q(\\mathcal{M}\\mid \\mathcal{D}_{\\mathcal{E}})$, and sometimes $q(\\pi \\mid \\mathcal{M})$. As here $q$ here is not the exact probability and are approximated via different approaches, it seems to me that it is better to use different notations for different approximation networks.\n- In (14) and (15), it is clear that $q(\\pi \\mid \\mathcal{M})$ is usually the optimal policy of $\\mathcal{M}$ and is estimated using network $q(\\theta)$, but what does $q(\\pi \\mid \\mathcal{D}_{\\mathcal{E}})$ represent? What is the corresponding procedure in Algorithm 1? (In my understanding $q(\\phi)$ is $q(\\pi \\mid \\mathcal{M})$? So what is $q(\\pi \\mid D)$?)\n- In line 11 of Algorithm 1, how to create mixed dataset $\\lambda \\mathcal{D}_1 + (1 - \\lambda) \\mathcal{D}_2$? With the presence of $\\lambda$, a simple merge does not seem to do the job. Does it mean with probability $\\lambda$ sampling from the first and with probability $1 - \\lambda$ sampling from the second?\n- Notations in (4), when the regret is defined on $M^*$, one cannot take expectation over $p(\\mathcal{M})$, I think the notations should be unified (either using $M^*$ or $M$ for both).\n\nI would be willing to raise my score if the algorithm can be clarified.\n\nMinor questions:\n\n- How does the optimal $\\lambda$ change across tasks. i.e. how hard is it to tune $\\lambda$?", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject