Original: P OSTERIOR S AMPLING M ODEL - BASED P OLICY O PTIMIZATION UNDER A PPROXIMATE I NFERENCE
Revised: P OSTERIOR S AMPLING M ODEL - BASED P OLICY O PTIMIZATION UNDER A PPROXIMATE I NFERENCE

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Model -based reinforcement learning algorithms (MBRL) hold tremendous promise for improving the sample efficiency in online RL.
Revised: Model-based reinforcement learning algorithms (MBRL) hold tremendous promise for improving the sample efficiency in online RL.

Original: However, many existing popular
Revised: However, many existing popular

Original: MBRL algorithms cannot deal with exploration and exploitation properly
Revised: MBRL algorithms cannot deal with exploration and exploitation properly.

Original: 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
Revised: 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.

Original: In this paper, we show that adopting the same methodology as in exact PSRL can be suboptimal under approximate inference.
Revised: In this paper, we show that adopting the same methodology as in exact PSRL can be suboptimal under approximate inference.

Original: 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.
Revised: 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.

Original: By adopting such a posterior factorization, we further propose a general algorithmic framework for PSRL under approximate inference and a practical instantiation of it.
Revised: By adopting such a posterior factorization, we further propose a general algorithmic framework for PSRL under approximate inference and a practical instantiation of it.

Original: Empirically, our algorithm can surpass baseline methods by a significant margin on both dense rewards and sparse rewards tasks from the Deepmind control suite, OpenAI Gym and Metaworld benchmarks
Revised: Empirically, our algorithm can surpass baseline methods by a significant margin on both dense rewards and sparse rewards tasks from the Deepmind control suite, OpenAI Gym and Metaworld benchmarks.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Model-based reinforcement learning has demonstrated great success in improving the sample efficiency of RL.
Revised: Model-based reinforcement learning has demonstrated great success in improving the sample efficiency of RL.

Original: However, many existing popular model-based algorithms (Kurutach et al., 2018; Chua et al., 2018; Janner et al., 2019) cannot deal with exploration and exploitation properly, and hence may lead to poor performance when exploration is crucial.
Revised: However, many existing popular model-based algorithms (Kurutach et al., 2018; Chua et al., 2018; Janner et al., 2019) cannot deal with exploration and exploitation properly, and hence may lead to poor performance when exploration is crucial.

Original: To trade off exploration and exploitation, most of the existing algorithms can be categorized into 1) optimism-based (Jaksch et al., 2010; Pacchiano et al., 2021; Curi et al., 2020); 2) posterior-sampling-based (Strens, 2000; Osband et al., 2013; 2018;
Revised: To trade off exploration and exploitation, most of the existing algorithms can be categorized into 1) optimism-based (Jaksch et al., 2010; Pacchiano et al., 2021; Curi et al., 2020); 2) posterior-sampling-based (Strens, 2000; Osband et al., 2013; 2018;

Original: Fan & Ming, 2021); and 3) information-directed sampling (Russo & Van Roy, 2014) approaches.
Revised: Fan & Ming, 2021); and 3) information-directed sampling (Russo & Van Roy, 2014) approaches.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: Hence, in this paper, we focus on posterior sampling.
Revised: Hence, in this paper, we focus on posterior sampling.

Original: For simplicity, we restrict attention to the episodic RL setting.
Revised: For simplicity, we restrict attention to the episodic RL setting.

Original: Under the PSRL framework, one maintains a posterior p ( M|D E ) of the Markov decision process (MDP) M given the observations D E collected in real environment E .
Revised: Under the PSRL framework, one maintains a posterior p ( M|D E ) of the Markov decision process (MDP) M given the observations D E collected in real environment E .

Original: At the beginning of each episode, an MDP is sampled from the posterior, and then we compute the optimal policy π ( M ) for the sampled model M .
Revised: At the beginning of each episode, an MDP is sampled from the posterior, and then we compute the optimal policy π ( M ) for the sampled model M .

Original: Equivalently, we can also view this policy as a sample from a “degenerate” posterior  over policies of the form p ( π |D E ) =
Revised: Equivalently, we can also view this policy as a sample from a “degenerate” posterior 1 over policies of the form p ( π |D E ) =

Original: (cid:82) δ ( π |M ) p ( M|D E )
Revised: (cid:82) δ ( π |M ) p ( M|D E )

Original: d M , where δ ( π |M ) =
Revised: d M , where δ ( π |M ) =

Original: δ ( π − π ( M )) is a Dirac delta distribution.
Revised: δ ( π − π ( M )) is a Dirac delta distribution.

Original: This policy is then executed in the real environment to collect new data.
Revised: This policy is then executed in the real environment to collect new data.

Original: Theoretically, such a simple strategy has been shown to achieve a Bayesian regret of ˜
Revised: Theoretically, such a simple strategy has been shown to achieve a Bayesian regret of ˜

Original: O ( √ K ) for K episodes (Osband et al., 2013; Osband & Van Roy, 2014).
Revised: O ( √ K ) for K episodes (Osband et al., 2013; Osband & Van Roy, 2014).

Original: However, the theoretical guarantees only hold under exact inference, i.e., when we have access to the true posterior over models p ( M|D E ) , and when we can compute the optimal policy, which is very unlikely in practice.
Revised: However, the theoretical guarantees only hold under exact inference, i.e., when we have access to the true posterior over models p ( M|D E ) , and when we can compute the optimal policy, which is very unlikely in practice.

Original: A common heuristic approximation to PSRL (see e.g., Fan & Ming (2021)) is to replace the posterior over models with some approximation, such as Bayesian linear regression on top of the representations learned by a neural network.
Revised: A common heuristic approximation to PSRL (see e.g., Fan & Ming (2021)) is to replace the posterior over models with some approximation, such as Bayesian linear regression on top of the representations learned by a neural network.

Original: In this way, the resulting policy is then sampled from q δ ( π |D E ) =
Revised: In this way, the resulting policy is then sampled from q δ ( π |D E ) =

Original: (cid:82) δ ( π |M ))
Revised: (cid:82) δ ( π |M ))

Original: q ( M|D E )
Revised: q ( M|D E )

Original: d M , where we replace p ( M|D E ) with the approximate posterior q ( M|D E ) .
Revised: d M , where we replace p ( M|D E ) with the approximate posterior q ( M|D E ) .

Original: At first glance, such a heuristic choice is natural as it shares the same form as the true posterior p ( π |D E ) .
Revised: At first glance, such a heuristic choice is natural as it shares the same form as the true posterior p ( π |D E ) .

Original: However, in this paper, we prove that we can get lower regret if we replace the degenerate δ ( π |M ) with a non-degenerate distribution of the form q ( π |M , D E ) , that depends on the model as well as empirical data D E ; this is necessary to compensate for the fact that the posterior over models q ( M|D E ) may be suboptimal.
Revised: However, in this paper, we prove that we can get lower regret if we replace the degenerate δ ( π |M ) with a non-degenerate distribution of the form q ( π |M , D E ) , that depends on the model as well as empirical data D E ; this is necessary to compensate for the fact that the posterior over models q ( M|D E ) may be suboptimal.

Original: By tuning the relative strength of the dependence of π on D E and
Revised: By tuning the relative strength of the dependence of π on D E and

Original: M , we can find a sweet spot between maximizing the data efficiency and minimizing the effect of approximate inference error.
Revised: M , we can find a sweet spot between maximizing the data efficiency and minimizing the effect of approximate inference error.

Original: Furthermore, such a decomposition is guaranteed to be no worse than the standard approach of using q δ ( π |D E ) .
Revised: Furthermore, such a decomposition is guaranteed to be no worse than the standard approach of using q δ ( π |D E ) .

Original: Building upon the above results, we come up with a generic framework for PSRL under approximate inference.
Revised: Building upon the above results, we come up with a generic framework for PSRL under approximate inference.

Original: To implement the method in practice, we combine deep ensembles (Lakshminarayanan et al., 2017) and MBPO (Janner et al., 2019).
Revised: To implement the method in practice, we combine deep ensembles (Lakshminarayanan et al., 2017) and Model-based Policy Optimization (MBPO) (Janner et al., 2019).

Original: We also propose two different sampling strategies for policy selection that exploit our posterior approximation.
Revised: We also propose two different sampling strategies for policy selection that exploit our posterior approximation.

Original: 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).
Revised: 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).

Original: Additionally, we also conduct various ablation studies to provide a better understanding of our algorithm.
Revised: Additionally, we also conduct various ablation studies to provide a better understanding of our algorithm.

Original: In summary, our contributions are
Revised: In summary, our contributions are

Original: 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 is unavailable (Section 2).
Revised: 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 is unavailable (Section 2).

Original: Motivated by our analysis, we develop a 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).
Revised: Motivated by our analysis, we develop a 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).

Original: We present empirical results on DM control suite, OpenAI Gym and Metaworld benchmarks to demonstrate the efficacy of the proposed approach (Section 4).
Revised: We present empirical results on DM control suite, OpenAI Gym and Metaworld benchmarks to demonstrate the efficacy of the proposed approach (Section 4).

Original:  P ROBLEM STATEMENT AND THEORETICAL RESULTS
Revised:  P ROBLEM STATEMENT AND THEORETICAL RESULTS

Original: We start by introducing some notation, and summarizing prior work, before presenting our new theoretical results, which forms a basis for our algorithm.
Revised: We start by introducing some notation, and summarizing prior work, before presenting our new theoretical results, which forms a basis for our algorithm.

Original: Notation.
Revised: Notation.

Original: We consider the finite-horizon episodic Markov Decision Process (MDP) problem, of which we denote an instance as M := {S , A , r M , p M , H, ρ } .
Revised: We consider the finite-horizon episodic Markov Decision Process (MDP) problem, of which we denote an instance as M := {S , A , r M , p M , H, ρ } .

Original: For each instance M , S and A denote the set of states and actions, respectively.
Revised: For each instance M , S and A denote the set of states and actions, respectively.

Original: r M : S ×
Revised: r M : S ×

Original: A → [0 , R max ] is the reward function, p M is the transition distribution, H is the length of the episode, and ρ is the distribution of the initial state.
Revised: A → [0 , R max ] is the reward function, p M is the transition distribution, H is the length of the episode, and ρ is the distribution of the initial state.

Original: We further define the value function of a policy π under MDP M at timestep i as
Revised: We further define the value function of a policy π under MDP M at timestep i as

Original: [Equation]
Revised: [Equation]

Original:  We define π ⋆ as the optimal policy for an MDP M if V M π ⋆ ,i ( s ) = max π V M π,i ( s ) for all s ∈ S and i ∈ [1 , H ] .
Revised: where s t +1 ∼ p M ( s | s t , a t ) and a t ∼ π ( a | s t ) . We define π ⋆ as the optimal policy for an MDP M if V M π ⋆ ,i ( s ) = max π V M π,i ( s ) for all s ∈ S and i ∈ [1 , H ] .

Original: We define the cumulative reward obtained by policy π over H steps sampled from model M as follows:
Revised: We define the cumulative reward obtained by policy π over H steps sampled from model M as follows:

Original: [Equation]
Revised: [Equation]

Original: where s t +1 ∼ p M ( s | s t , a t ) and a t ∼ π ( a | s t ) . [Equation] 2) where the initial state is sampled from s 1 ∼ ρ ( s ) .
Revised: where s t +1 ∼ p M ( s | s t , a t ) and a t ∼ π ( a | s t ) , (2) where the initial state is sampled from s 1 ∼ ρ ( s ) .

Original: Regret.
Revised: Regret.

Original: For a given MDP M , the regret is defined as the difference between value function of the optimal policy in hindsight and that of the actual policy executed by the algorithm A ,
Revised: For a given MDP M , the regret is defined as the difference between value function of the optimal policy in hindsight and that of the actual policy executed by the algorithm A ,

Original: Regret ( T, A , M [Equation]
Revised: Regret ( T, A , M ) :=

Original: [Equation]
Revised: [Equation]

Original: where π ⋆ is the optimal policy for M , and π k is the policy employed by the algorithm for k th episode.
Revised: where π ⋆ is the optimal policy for M , and π k is the policy employed by the algorithm for k th episode.

Original: Correspondingly, the Bayesian regret is defined as the expectation of the above regret, i.e.,
Revised: Correspondingly, the Bayesian regret is defined as the expectation of the above regret, i.e.,

Original: BayesianRegret [Equation]
Revised: BayesianRegret ( T, A , p ( M )) :

Original: (cid:34)
Revised: = E [ Regret ( T

Original: [Equation]
Revised: [Equation]

Original: k
Revised: A

Original: [Equation]
Revised: [Equation]

Original: ∆
Revised: M )] = E (cid:34) K (cid:88) k

Original: [Equation] (cid:35) .
Revised: = ∆ k (cid:35) .

Original: Here the expectation is taken over the prior distribution of dynamics models p ( M ) and the randomness in the algorithm A and environment.
Revised: Here the expectation is taken over the prior distribution of dynamics models p ( M ) and the randomness in the algorithm A and environment.

Original: PSRL.
Revised: PSRL.

Original: Posterior Sampling Reinforcement Learning or PSRL (Strens, 2000) serves as a generic algorithmic framework for automatically trading off exploration and exploitation in online RL.
Revised: Posterior Sampling Reinforcement Learning or PSRL (Strens, 2000) serves as a generic algorithmic framework for automatically trading off exploration and exploitation in online RL.

Original: The core of PSRL is computing the posterior over MDPs (dynamics and reward models).
Revised: The core of PSRL is computing the posterior over MDPs (dynamics and reward models).

Original: Conditioned on the data D E collected from the environment, we denote the posterior of the model as p ( M|D E ) .
Revised: Conditioned on the data D E collected from the environment, we denote the posterior of the model as p ( M|D E ) .

Original: Hence, the posterior distribution of policies is
Revised: Hence, the posterior distribution of policies is

Original: [Equation]
Revised: [Equation]

Original: where p ( π |M ) =
Revised: where p ( π |M ) =

Original: δ ( π |M ) (5) where δ ( π |M ) is a  delta distribution, defined as δ(arg max π R M ( π  ) |M )
Revised: δ ( π |M ) (5) where δ ( π |M ) is a Dirac delta distribution, defined as δ ( π ( M ) |M )

Original: = 1  .
Revised: = 1 , where π ( M ) = arg max π R M ( π ) is the optimal policy for solving the MDP M .

Original: 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.
Revised: 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.

Original: This simple algorithmic framework not only attains a  regret of ˜ O ( √ K )
Revised: This simple algorithmic framework not only attains a Bayesian regret of ˜ O ( √ K )

Original: (Osband et al., 2013), where K is the total number of episodes, but also enjoys better empirical performance than optimism-based methods in bandits (Chapelle & Li, 2011).
Revised: (Osband et al., 2013), where K is the total number of episodes, but also enjoys better empirical performance than optimism-based methods in bandits (Chapelle & Li, 2011; Osband & Van Roy, 2017).

Original: Nevertheless, the theoretical results only hold under exact inference.
Revised: Nevertheless, the theoretical results only hold under exact inference.

Original: In the rest of this section, we discuss how performance is affected when we use approximate inference.
Revised: In the rest of this section, we discuss how performance is affected when we use approximate inference.

Original: Bayesian regret under approximate inference.
Revised: Bayesian regret under approximate inference.

Original: Let us denote the approximate and true posterior distribution of polices at k th episode by
Revised: Let us denote the approximate and true posterior distribution of polices at k th episode by

Original: [Equation]
Revised: [Equation]

Original: where D k E denotes all the data collected from the environment E till the k -th episode .
Revised: where D k E denotes all the data collected from the environment E till the k -th episode .

Original: Next we characterize how approximate posterior inference affects the Bayesian regret.
Revised: Next we characterize how approximate posterior inference affects the Bayesian regret.

Original: Theorem 1 For K episodes, the Bayesian regret of posterior sampling reinforcement learning algorithm A with any approximate posterior distribution q k at episode k is upper bounded by
Revised: Theorem 1 For K episodes, the Bayesian regret of posterior sampling reinforcement learning algorithm A with any approximate posterior distribution q k at episode k is upper bounded by

Original: [Equation]
Revised: [Equation]

Original: where H ( π ⋆ ) is the entropy of the prior distribution of  polices,
Revised: where H ( π ⋆ ) is the entropy of the prior distribution of optimal polices,

Original: i.e., p ( π ) = (cid:82) δ ( π |M ) p ( M ) d M ,and C is some problem-dependent constant.
Revised: i.e., p ( π ) = (cid:82) δ ( π |M ) p ( M ) d M , C is some problem-dependent constant and d KL ( ·| · ) is the KL-divergence.

Original: Theorem 1 (proved in the appendix) provides an upper bound of the Bayesian regret under approximate posterior inference.
Revised: Theorem 1 (proved in the appendix) provides an upper bound of the Bayesian regret under approximate posterior inference.

Original: The first term is ˜ O ( √ K ) .
Revised: The first term is ˜ O ( √ K ) .

Original: The second term will be zero under exact posterior inference.
Revised: The second term will be zero under exact posterior inference.

Original: However, when performing approximate inference, we see that we should choose the approximate posterior distribution q ( π |D E ) as “close” to the true distribution p ( π |D E ) as possible.
Revised: However, when performing approximate inference, we see that we should choose the approximate posterior distribution q ( π |D E ) as “close” to the true distribution p ( π |D E ) as possible.

Original: A natural and common choice of q ( π |D E ) would be the following one, which takes the same form as the true posterior,
Revised: A natural and common choice of q ( π |D E ) would be the following one, which takes the same form as the true posterior,

Original: [Equation]
Revised: [Equation]

Original: However, the following proposition shows that such a choice is suboptimal.
Revised: However, the following proposition shows that such a choice is suboptimal.

Original: Proposition 1 Under approximate inference (i.e., q ( M|D E ) ̸
Revised: Proposition 1 Under approximate inference (i.e., q ( M|D E ) ̸

Original: = p ( M|D E ) ), the optimal q ( π |M ) may not be a Dirac delta distribution, i.e., there exists other q ( π |D E ) such that
Revised: = p ( M|D E ) ), the optimal q ( π |M ) may not be a Dirac delta distribution, i.e., there exists other q ( π |D E ) such that

Original: [Equation]
Revised: [Equation]

Original: As an illustration, we provide the following example, which also serves as a constructive proof of
Revised: As an illustration, we provide the following example, which also serves as a constructive proof of

Original: Proposition 1.
Revised: Proposition 1.

Original: EXAMPLE 1. S UBOPTIMALITY OF q δ ( π |M ) .
Revised: EXAMPLE 1. S UBOPTIMALITY OF q δ ( π |M ) .

Original: Consider a toy setting, where the support set of MDPs is {M 1 , M 2 } , and the support set of policies is { π 1 , π 2 } .
Revised: Consider a toy setting, where the support set of MDPs is {M 1 , M 2 } , and the support set of policies is { π 1 , π 2 } .

Original: Suppose that the true posterior distribution of MDPs is p ( M 1 |D E )
Revised: Suppose that the true posterior distribution of MDPs is p ( M 1 |D E )

Original: = 1 / 3 , p ( M 2 |D E )
Revised: = 1 / 3 , p ( M 2 |D E )

Original: = 2 / 3 , and the optimal policy per MDP is δ ( π 1 |M 1 ) = 

Revised: = 2 / 3 , and the optimal policy per MDP is δ ( π 1 |M 1 ) =

Original: and δ ( π 2 |M 2 ) = 1 .
Revised: and δ ( π 2 |M 2 ) = 1 .

Original: This we get the following exact distribution over policies: p ( π |D E ) is
Revised: This we get the following exact distribution over policies: p ( π |D E ) is

Original: [Equation]
Revised: [Equation]

Original: Now suppose we use the approximate posterior distribution over models, q ( M 1 |D E )
Revised: Now suppose we use the approximate posterior distribution over models, q ( M 1 |D E )

Original: = 0 and q ( M 2 |D E )
Revised: = 0 and q ( M 2 |D E )

Original: = 1 .
Revised: = 1 .

Original: We can optimize q ( π |M ) by minimizing d KL ( q ( π |D E ) | p ( π |D E )) .
Revised: We can optimize q ( π |M ) by minimizing d KL ( q ( π |D E ) | p ( π |D E )) .

Original: One solution could be
Revised: One solution could be

Original: [Equation]
Revised: [Equation]

Original: (cid:125)
Revised: (cid:125)

Original: q
Revised: q

Original: [Equation]
Revised: [Equation]

Original: |
Revised: |

Original: [Equation]
Revised: [Equation]

Original: (cid:104)
Revised: (cid:104)

Original: q
Revised: q

Original: [Equation]
Revised: [Equation]

Original: M 1 |D E )
Revised: M 1 |D E )

Original: [Equation]
Revised: [Equation]

Original: 0 q
Revised: 0 q

Original: [Equation]
Revised: [Equation]

Original: M 2 |D E )
Revised: M 2 |D E )

Original: [Equation]
Revised: [Equation]

Original: 1 (cid:105) (cid:124) (cid:123)(cid:122) (cid:125) q
Revised: 1 (cid:105) (cid:124) (cid:123)(cid:122) (cid:125) q

Original: ( M|D E ) =
Revised: ( M|D E ) =

Original: (cid:104) q
Revised: (cid:104) q

Original: ( π 1 |D E )
Revised: ( π 1 |D E )

Original: [Equation]
Revised: [Equation]

Original: 23 q ( π 2 |D E )
Revised: 23 q ( π 2 |D E )

Original: [Equation]
Revised: [Equation]

Original: 13 (cid:105)
Revised: 13 (cid:105)

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: We see that the optimal q ( π |M ) requires modeling uncertainty in the policy even conditional on the model.
Revised: We see that the optimal q ( π |M ) requires modeling uncertainty in the policy even conditional on the model.

Original: By contrast, if we adopt q δ ( π |D E ) as our approximation, we will have
Revised: By contrast, if we adopt q δ ( π |D E ) as our approximation, we will have

Original: [Equation]
Revised: [Equation]

Original: The above example tells us that the q δ ( π |D E ) can perform arbitrarily poorly in terms of the KL divergence.
Revised: The above example tells us that the q δ ( π |D E ) can perform arbitrarily poorly in terms of the KL divergence.

Original: Given this observation, a natural follow-up question would be: is there a better choice other than q δ ( π |D E ) ?
Revised: Given this observation, a natural follow-up question would be: is there a better choice other than q δ ( π |D E ) ?

Original: We provide an answer in the next section.
Revised: We provide an answer in the next section.

Original:  M ETHOD
Revised:  M ETHOD

Original: Motivated by the results in the previous section, we first introduce a posterior decomposition which is guaranteed to be better than q δ ( π |D E ) .
Revised: Motivated by the results in the previous section, we first introduce a posterior decomposition which is guaranteed to be better than q δ ( π |D E ) .

Original: Then we introduce a practical version of the algorithm built upon deep ensembles (Lakshminarayanan et al., 2017).
Revised: Then we introduce a practical version of the algorithm built upon deep ensembles (Lakshminarayanan et al., 2017).

Original: Finally we propose two sampling approaches for encouraging efficient exploration.
Revised: Finally we propose two sampling approaches for encouraging efficient exploration.

Original: 3. P OSTERIOR DECOMPOSITION
Revised: 3. P OSTERIOR DECOMPOSITION

Original: In section 2, we showed that q δ ( π |D E ) is not a favorable choice.
Revised: In section 2, we showed that q δ ( π |D E ) is not a favorable choice.

Original: This is because it assumes that, once
Revised: This is because it assumes that, once

Original: M is given, the policy π is determined.
Revised: M is given, the policy π is determined.

Original: This motivates us to consider the following more flexible posterior decomposition of q ( π |D E ) ,
Revised: This motivates us to consider the following more flexible posterior decomposition of q ( π |D E ) ,

Original: [Equation]
Revised: [Equation]

Original: Intuitively, such a posterior decomposition no longer assumes that the model can capture all the relevant properties of the data.
Revised: Intuitively, such a posterior decomposition no longer assumes that the model can capture all the relevant properties of the data.

Original: We illustrate these two posterior approximations in Figure 1.
Revised: We illustrate these two posterior approximations in Figure 1.

Original: The extra parameter λ ∈ [0 , 1] allows us to balance the importance of fictitious data ( D M ) from M and the real data ( D E ) from the environment.
Revised: The extra parameter λ ∈ [0 , 1] allows us to balance the importance of fictitious data ( D M ) from M and the real data ( D E ) from the environment.

Original: In particular, we define
Revised: In particular, we define

Original: [Equation]
Revised: [Equation]

Original: Thus when λ is small, we trust our model more, whereas when λ is large we trust it less.
Revised: Thus when λ is small, we trust our model more, whereas when λ is large we trust it less.

Original: In the extreme case where λ = 0 , this framework reduces to the degenerate posterior q δ ( π |D E ) .
Revised: In the extreme case where λ = 0 , this framework reduces to the degenerate posterior q δ ( π |D E ) .

Original: By adjusting λ , we can find a sweet spot in minimizing the effect of approximate inference error and maximizing the data efficiency.
Revised: By adjusting λ , we can find a sweet spot in minimizing the effect of approximate inference error and maximizing the data efficiency.

Original: More formally, the following proposition illustrates the advantage of equation 13.
Revised: More formally, the following proposition illustrates the advantage of equation 13.

Original: Proposition 2 By adopting the posterior decomposition of equation 13, we have
Revised: Proposition 2 By adopting the posterior decomposition of equation 13, we have

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: Algorithm

Original: Require: Initialize an ensemble of dynamics models Θ
Revised: Require: Initialize an ensemble of dynamics models Θ

Original: = { ˆ θ n } Nn =1 i.i.d.
Revised: = { ˆ θ n } Nn =1 i.i.d.

Original: ∼ q ( θ ) .
Revised: ∼ q ( θ ) .

Original: Require: Initialize an ensemble of policy networks Φ = { ˆ ϕ n,m } N,Mn,m =1 i.i.d.
Revised: Require: Initialize an ensemble of policy networks Φ = { ˆ ϕ n,m } N,Mn,m =1 i.i.d.

Original: ∼ q ( ϕ ) .
Revised: ∼ q ( ϕ ) .

Original: Require:
Revised: Require:

Original: Initialize empty datasets D E and {D n,m M } N,Mn,m =1 .
Revised: Initialize empty datasets D E and {D n,m M } N,Mn,m =1 .

Original: Real data vs. synthetic data ratio λ .
Revised: Real data vs. synthetic data ratio λ .

Original: 1: for K episodes do ▷ /* Dynamics training.
Revised: 1: for K episodes do ▷ /* Dynamics training.

Original: (Line 2)
Revised: (Line 2)

Original: Train the ensemble models Θ on D E under the objective in equation 17.
Revised: Train the ensemble models Θ on D E under the objective in equation 17.

Original: ▷ /* Policy sampling.
Revised: ▷ /* Policy sampling.

Original: (Line 3)
Revised: (Line 3)

Original: Sample a policy π from Φ uniformly at random (equation 18) or based on the optimistic distribution (equation 19).
Revised: Sample a policy π from Φ uniformly at random (equation 18) or based on the optimistic distribution (equation 19).

Original: Sample state s 1 from the initial state distribution ρ ( s ) 5: for h = 2 : H steps do ▷ /* Data collection.
Revised: Sample state s 1 from the initial state distribution ρ ( s ) 5: for h = 2 : H steps do ▷ /* Data collection.

Original: (Lines 6-11) */ 6: s h = rollout(world dynamics E , policy π , initial state s h − 1 , num.
Revised: (Lines 6-11) */ 6: s h = rollout(world dynamics E , policy π , initial state s h − 1 , num.

Original: steps 1 ) 7:
Revised: steps 1 ) 7:

Original: Add s h to D E 8:
Revised: Add s h to D E 8:

Original: Sample state s ∼ D E 9: for each model n , policy m do 10:
Revised: Sample state s ∼ D E 9: for each model n , policy m do 10:

Original: D n,m M = rollout(dynamics ˆ θ n , policy ˆ ϕ n,m , initial state s , num. steps R ) 11:
Revised: D n,m M = rollout(dynamics ˆ θ n , policy ˆ ϕ n,m , initial state s , num. steps R ) 11:

Original: Created mixed dataset D = λ D E + (1 − λ ) D n,m M ▷ /* Policy optimization (Line 12) */ 12: ˆ ϕ n,m = update-policy( ˆ ϕ n,m , D , num.
Revised: Created mixed dataset D = λ D E + (1 − λ ) D n,m M▷ / * Policy optimization (Line 12) */ 12: ˆ ϕ n,m = update-policy( ˆ ϕ n,m , D , num.

Original: gradient steps G ) 13: end for 14: end for 15:
Revised: gradient steps G ) 13: end for 14: end for 15:

Original: Update the optimistic policy distribution (equation 19).
Revised: Update the optimistic policy distribution (equation 19).

Original: 16: end for 
Revised: 16: end for π

Revised: M

Revised: D

Original: [Equation]
Revised: [Equation]

Original: standard and
Revised: (a) π M D

Original: ) our posterior over policies π 
Revised: [Equation] λ (b) the standard and (b) our posterior over policies π .

Original: Differences are shown in red.
Revised: Differences are shown in red.

Original: The above proposition informs us that we can minimize the KL divergence, and hence the upper bound of the Bayesian regret, by carefully choosing the value of λ .
Revised: The above proposition informs us that we can minimize the KL divergence, and hence the upper bound of the Bayesian regret, by carefully choosing the value of λ .

Original: As empirical evidence, in Figurewe plot the cumulative regret with different λ in the cartpole-swingup environment with sparse reward.
Revised: As empirical evidence, in Figurewe plot the cumulative regret with different λ in the cartpole-swingup environment with sparse reward.

Original: We see that a value of about λ = 0 .
Revised: We see that a value of about λ = 0 .

Original: 4 or λ = 0 .
Revised: 4 or λ = 0 .

Original: 5 is the best, and λ = 0 is the worst in this case.
Revised: 5 is the best, and λ = 0 is the worst in this case.

Original: 3. P ROPOSED ALGORITHM
Revised: 3. P ROPOSED ALGORITHM

Original: Building upon the above observations, we introduce a simple and general algorithmic framework for PSRL under approximate inference, which only differs from the standard PSRL framework in the decomposition of policy posterior (see Algorithm 2 in Appendix C).
Revised: Building upon the above observations, we introduce a simple and general algorithmic framework for PSRL under approximate inference, which only differs from the standard PSRL framework in the decomposition of policy posterior (see Algorithm 2 in Appendix C).

Original: To implement the method in practice, we use deep ensembles (Lakshminarayanan et al., 2017) to approximate the posterior distributions q ( π |M , D E ) and q ( M|D E ) .
Revised: To implement the method in practice, we use deep ensembles (Lakshminarayanan et al., 2017) to approximate the posterior distributions q ( M|D E ) and q ( π |M , D E ) with Θ and Φ (see Algorithm 1).

Original: This is similar to ME-TRPO (Kurutach et al., 2018), PETS (Chua et al., 2018) and MBPO (Janner et al., 2019), except we also model the uncertainty over policies (i.e., q ( π |M , D E ) ) in addition to dynamics (i.e., q ( M|D E ) ). 
Revised: This is similar to METRPO (Kurutach et al., 2018), PETS (Chua et al., 2018) and MBPO (Janner et al., 2019), except we also model the uncertainty over policies (i.e., q ( π |M , D E ) ) in addition to dynamics (i.e., q ( M|D E ) ). [Figure]

Original: In more detail, each member in the deep ensemble is a conditional Gaussian distribution over outputs, characterized by its mean µ and variance σ 2 .
Revised: In more detail, each member in the deep ensemble is a conditional Gaussian distribution over outputs, characterized by its mean µ and variance σ 2 .

Original: For multi-dimensional predictions, we treat each dimension independently, and only predict the marginal mean and marginal variance for simplicity.
Revised: For multi-dimensional predictions, we treat each dimension independently, and only predict the marginal mean and marginal variance for simplicity.

Original: [Figure]

Original: Each ensemble member is trained independently by minimizing the negative log-likelihood,
Revised: Each ensemble member is trained independently by minimizing the negative log-likelihood,

Original: [Equation]
Revised: [Equation]

Revised: [Equation]

Original: We maintain N different dynamics models.
Revised: We maintain N different dynamics models.

Original: For each such model, we also compute M different policies; we use the soft actor-critic (SAC) (Haarnoja et al., 2018) method (see the appendix for details).
Revised: For each such model, we also compute M different policies; we use the soft actor-critic (SAC) (Haarnoja et al., 2018) method (see the appendix for details).

Original: The policy network π n,m is updated based on synthetic data D n,m M generated by dynamics model n and policy model m , as well as environmental data D E generated by interacting the real world dynamics with a sampled policy.
Revised: The policy network π n,m is updated based on synthetic data D n,m M generated by dynamics model n and policy model m , as well as environmental data D E generated by interacting the real world dynamics with a sampled policy.

Original: See Algorithm 1 for the pseudo-code.
Revised: See Algorithm 1 for the pseudo-code.

Original: 3. S AMPLING POLICIES
Revised: 3. S AMPLING POLICIES

Original: Ensemble Sampling.
Revised: Ensemble Sampling.

Original: Given the posterior distributions, it remains to specify the sampling approach for policies.
Revised: Given the posterior distributions, it remains to specify the sampling approach for policies.

Original: The simplest sampling strategy is uniform sampling at the beginning of each episode,
Revised: The simplest sampling strategy is uniform sampling at the beginning of each episode,

Original: [Equation]
Revised: [Equation]

Original:  In the case of bandits (where N = 1 , since there is no dynamics model), such a simple strategy has been shown to achieve a regret of ˜ O ( √ T + T (cid:112) A/M ) (Lu & Van Roy, 2017; Qin et al., 2022) for
Revised: [Equation] In the case of bandits (where N = 1 , since there is no transition model), such a simple strategy has been shown to achieve a regret of ˜ O ( √ T + T (cid:112) A/M ) (Lu & Van Roy, 2017; Qin et al., 2022) for

Original: T steps in Gaussian linear bandits, where M is the size of the ensemble and A is the number of arms.
Revised: T steps in Gaussian linear bandits, where M is the size of the ensemble and A is the number of arms.

Original: 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.
Revised: 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.

Original: Optimistic Ensemble Sampling.
Revised: Optimistic Ensemble Sampling.

Original: Unfortunately, ensemble sampling may overly explore some unpromising regions, as it treats each member in the ensemble model equally.
Revised: Unfortunately, ensemble sampling may overly explore some unpromising regions, as it treats each member in the ensemble model equally.

Original: This may lead to unnecessary or wasteful explorations.
Revised: This may lead to unnecessary or wasteful explorations.

Original: To reduce this, we propose an optimistic version of ensemble sampling which we call OPS-MBPO.
Revised: To reduce this, we propose an optimistic version of ensemble sampling which we call OPS-MBPO.

Original: 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).)
Revised: 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).)

Original: We then use this performance to determine the probability of choosing each member, thus gradually discarding unpromising ensemble members.
Revised: We then use this performance to determine the probability of choosing each member, thus gradually discarding unpromising ensemble members.

Original: More precisely, at the beginning of k th episode, we sample the policy from the following Boltzmann distribution, instead of uniformly at random,
Revised: More precisely, at the beginning of k th episode, we sample the policy from the following Boltzmann distribution, instead of uniformly at random,

Original: [Equation]
Revised: [Equation]

Original: where τ is a temperature term for controlling the level of optimism, and R E ( π i , l ) is the empirical cumulative reward of π i at the l th episode.
Revised: where τ is a temperature term for controlling the level of optimism, and R E ( π i , l ) is the empirical cumulative reward of π i at the l th episode.

Original: When τ → ∞ , we recover uniform ensemble sampling, which we call PS-MBPO.
Revised: When τ → ∞ , we recover uniform ensemble sampling, which we call PS-MBPO.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: Our empirical evaluation aims to 1) verify the effectiveness of our proposed methods; 2) offer a deeper understanding about the mechanisms that are key to the improved performance; and 3) provide additional ablation studies on non-key components.
Revised: Our empirical evaluation aims to 1) verify the effectiveness of our proposed methods; 2) offer a deeper understanding about the mechanisms that are key to the improved performance; and 3) provide additional ablation studies on non-key components.

Original: We start by introducing the experimental setup.
Revised: We start by introducing the experimental setup.

Original: 4. E XPERIMENTAL S ETUP
Revised: 4. E XPERIMENTAL S ETUP

Original: We consider seven tasks from OpenAI Gym (Brockman et al., 2016),
Revised: We consider seven tasks from OpenAI Gym (Brockman et al., 2016),

Original: Deepmind Control Suite (Tunyasuvunakool et al., 2020) and Metaworld (Yu et al., 2020).
Revised: Deepmind Control Suite (Tunyasuvunakool et al., 2020) and Metaworld (Yu et al., 2020).

Original: This includes four dense reward tasks ( ant , half-cheetah , walker2d and hopper ) where the agent receives an immediate reward at each step, and three sparse reward tasks ( ball-in-cup , cartpole-swingup and window-open-v2 ), where the agent receives a reward only if it finishes the corresponding task.
Revised: This includes four dense reward tasks ( ant , half-cheetah , walker2d and hopper ) where the agent receives an immediate reward at each step, and three sparse reward tasks ( ball-in-cup , cartpole-swingup and window-open-v2 ), where the agent receives a reward only if it finishes the corresponding task.

Original: For sparse reward tasks, efficiently exploring the environment is more crucial than in dense reward tasks.
Revised: For sparse reward tasks, efficiently exploring the environment is more crucial than in dense reward tasks.

Original: See Figure 3 for visualizations of the tasks.
Revised: See Figure 3 for visualizations of the tasks.

Original: For baseline methods, we consider the model-based approaches SLBO (Luo et al., 2019), PETS (Chua et al., 2018) and MBPO (Janner et al., 2019), and a model-free method SAC (Haarnoja et al., 2018).
Revised: For baseline methods, we consider the model-based approaches SLBO (Luo et al., 2019), PETS (Chua et al., 2018) and MBPO (Janner et al., 2019), and a model-free method SAC (Haarnoja et al., 2018).

Original: 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.
Revised: 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.

Original: To draw more robust conclusions, we repeat each experiment with 10 different random seeds, and report the mean and the standard error.
Revised: To draw more robust conclusions, we repeat each experiment with 10 different random seeds, and report the mean and the standard error.

Original: More details can be found in Appendix A and C.
Revised: More details can be found in Appendix A and C.

Original: 4. C OMPARISON WITH EXISTING METHODS
Revised: 4. C OMPARISON WITH EXISTING METHODS

Original: We report the results on the dense reward tasks in Figure 4.
Revised: We report the results on the dense reward tasks in Figure 4.

Original: We see that our (O)PS-MBPO outperforms the baselines in all four tasks, including the MBPO method.
Revised: We see that our (O)PS-MBPO outperforms the baselines in all four tasks, including the MBPO method.

Original: (We verified that our implementation of
Revised: (We verified that our implementation of

Original: MBPO is comparable, or superior, to the performance of the original implementation; see appendix for details.)
Revised: MBPO is comparable, or superior, to the performance of the original implementation; see appendix for details.)

Original: For example, on hopper , our method only requires roughly 40K iterations to reach an average reward around 3,500, whereas MBPO needs around 150K steps.
Revised: For example, on hopper , our method only requires roughly 40K iterations to reach an average reward around 3,500, whereas MBPO needs around 150K steps.

Original: We report the results on the sparse reward tasks in Figure 5.
Revised: We report the results on the sparse reward tasks in Figure 5.

Original: We first observe that both PS-MBPO and OPS-MBPO outperform MBPO across three tasks, and the improvement is more significant on Cartpole-swingup and Window-open-v2 .
Revised: We first observe that both PS-MBPO and OPS-MBPO outperform MBPO across three tasks, and the improvement is more significant on Cartpole-swingup and Window-open-v2 .

Original: In contrast to Figure 4, we also observe that
Revised: In contrast to Figure 4, we also observe that

Original: OPS-MBPO further improves PS-MBPO on these three tasks by a significant margin.
Revised: OPS-MBPO further improves PS-MBPO on these three tasks by a significant margin.

Original: This confirms the advantage of adopting an optimistic sampling strategy in sparse reward tasks.
Revised: This confirms the advantage of adopting an optimistic sampling strategy in sparse reward tasks.

Original:  A BLATION S TUDIES
Revised:  A BLATION S TUDIES

Original: In this section, we conduct a variety of ablation studies to better understand our proposed method.
Revised: In this section, we conduct a variety of ablation studies to better understand our proposed method.

Original: Additional experiments can be found in the appendix, including ablation study with forced exploration (Phan et al., 2019), random function prior network
Revised: Additional experiments can be found in the appendix, including ablation study with forced exploration (Phan et al., 2019), random function prior network

Original: (Osband et al., 2018) etc.
Revised: (Osband et al., 2018) etc.

Original: Does the gain come from posterior sampling or ensemble?
Revised: Does the gain come from posterior sampling or ensemble?

Original: To assess the importance of posterior sampling, we compare performance where we sample a policy from the posterior at each episode to the standard approach where we use the average policy, computed by averaging the distribution over actions across all ensemble members.
Revised: To assess the importance of posterior sampling, we compare performance where we sample a policy from the posterior at each episode to the standard approach where we use the average policy, computed by averaging the distribution over actions across all ensemble members.

Original: The results are presented in Figure 6 with N = 5 dynamics networks and M ∈ { 1 , 3 , 5 } policy networks.
Revised: The results are presented in Figure 6 with N = 5 dynamics networks and M ∈ { 1 , 3 , 5 } policy networks.

Original: We see that when we disable the sampling procedure, the performance drops significantly.
Revised: We see that when we disable the sampling procedure, the performance drops significantly.

Original: Another interesting observation is that, when with posterior sampling, the performance improves as we add more ensemble members (i e., increase M ).
Revised: Another interesting observation is that, when with posterior sampling, the performance improves as we add more ensemble members (i e., increase M ).

Original: By contrast, when posterior sampling is disabled , increasing M seems doesn’t improve the performance.
Revised: By contrast, when posterior sampling is disabled , increasing M seems doesn’t improve the performance.

Original: This confirms that posterior sampling is the main factor behind the the improved performance, not just the fact that we have a larger ensemble for both dynamics and policies.
Revised: This confirms that posterior sampling is the main factor behind the the improved performance, not just the fact that we have a larger ensemble for both dynamics and policies.

Original: Effect of N and M .
Revised: Effect of N and M .

Original: Since we are using the deep ensemble approximation, it is natural to wonder if we could get better performance by using a lager size of ensemble for dynamics models and policies.
Revised: Since we are using the deep ensemble approximation, it is natural to wonder if we could get better performance by using a lager size of ensemble for dynamics models and policies.

Original: 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 in { 1 , 2 , 3 , 4 , 5 } .
Revised: 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 in { 1 , 2 , 3 , 4 , 5 } .

Original: We see that increasing both N and M can improve the performance, and that both forms of uncertainty (uncertainty on polices and dynamics) seem to matter.
Revised: We see that increasing both N and M can improve the performance, and that both forms of uncertainty (uncertainty on polices and dynamics) seem to matter.

Original: Visualization of the state space.
Revised: Visualization of the state space.

Original: To gain insight into the exploration behavior, we project the high-dimensional states of each trajectories collected by PS-MBPO and MBPO into a two-dimensional space using Umap (McInnes et al., 2018).
Revised: To gain insight into the exploration behavior, we project the high-dimensional states of each trajectories collected by PS-MBPO and MBPO into a two-dimensional space using Umap (McInnes et al., 2018).

Original: The visualization of these embeddings can be found in
Revised: The visualization of these embeddings can be found in

Original: Figure 8.
Revised: Figure 8.

Original: We see that, at the initial phase, PS-MBPO (top left three figures) is more explorative than
Revised: We see that, at the initial phase, PS-MBPO (top left three figures) is more explorative than

Original: MBPO (top right three figures), and that this also leads to better final performance (see Figure 5).
Revised: MBPO (top right three figures), and that this also leads to better final performance (see Figure 5).

Original: 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.
Revised: 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.

Original: Visually, we observe that PS-MBPO can move the robot arm to cover a more diversified regions, whereas MBPO fails to do so and the two trajectories look very similar to each other (this is also reflected in the Umap visualization).
Revised: Visually, we observe that PS-MBPO can move the robot arm to cover a more diversified regions, whereas MBPO fails to do so and the two trajectories look very similar to each other (this is also reflected in the Umap visualization).

Original:  R ELATED W ORKS
Revised:  R ELATED W ORKS

Original: Model-based reinforcement learning (MBRL).
Revised: Model-based reinforcement learning (MBRL).

Original: Research in MBRL mainly concerns how to learn the dynamics model given the data, and how to use the learned model.
Revised: Research in MBRL mainly concerns how to learn the dynamics model given the data, and how to use the learned model.

Original: The most commonly adopted approach for model learning is using the L 2 loss for one-step transitions (Kurutach et al., 2018;
Revised: The most commonly adopted approach for model learning is using the L 2 loss for one-step transitions (Kurutach et al., 2018;

Original: Luo et 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.
Revised: Luo et 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.

Original: In addition to the one-step training,
Revised: In addition to the one-step training,

Original: Hafner et al.
Revised: Hafner et al.

Original: ; Asadi et al.
Revised: ; Asadi et al.

Original: Lutter et al.
Revised: Lutter et al.

Original: (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.
Revised: (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.

Original: 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.
Revised: 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.

Original: Given the learned model, there are several ways to utilize it.
Revised: Given the learned model, there are several ways to utilize it.

Original: Model predictive control (MPC) (Camacho & Alba, 2013) is a derivative-free optimization method that has been adopted in many prior works (Nagabandi et al., 2018;
Revised: Model predictive control (MPC) (Camacho & Alba, 2013) is a derivative-free optimization method that has been adopted in many prior works (Nagabandi et al., 2018;

Original: Chua et al., 2018; Hafner et al., 2019; Fan & Ming, 2021;
Revised: Chua et al., 2018; Hafner et al., 2019; Fan & Ming, 2021;

Original: Lutter et al., 2021).
Revised: Lutter et al., 2021).

Original: However, MPC is sensitive to the planning horizon and struggles with high-dimensional problems.
Revised: However, MPC is sensitive to the planning horizon and struggles with high-dimensional problems.

Original: As a mitigation, Kurutach et al.
Revised: As a mitigation, Kurutach et al.

Original: ; Luo et al.
Revised: ; Luo et al.

Original: ; Janner et al.
Revised: ; Janner et al.

Original: (2019) instead train a policy on top of the model for amortizing the planning cost.
Revised: (2019) instead train a policy on top of the model for amortizing the planning cost.

Original: Similarly, there are also works that utilize the model to facilitate learning of value functions (Feinberg et al., 2018; Buckman et al., 2018).
Revised: Similarly, there are also works that utilize the model to facilitate learning of value functions (Feinberg et al., 2018; Buckman et al., 2018).

Original: However, few of these works have investigated the exploration and exploitation tradeoff.
Revised: However, few of these works have investigated the exploration and exploitation tradeoff.

Original: Exploration and
Revised: Exploration and

Original: exploitation.
Revised: exploitation.

Original: Handling the exploration and exploitation tradeoff is the central problem in online learning.
Revised: Handling the exploration and exploitation tradeoff is the central problem in online learning.

Original: 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-samplingbased (Strens, 2000; Osband et al., 2013; Osband & Van Roy, 2014;
Revised: 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-samplingbased (Strens, 2000; Osband et al., 2013; Osband & Van Roy, 2014;

Original: Osband et al., 2018; Fan &
Revised: Osband et al., 2018; Fan &

Original: Ming, 2021); and 3) information-directed sampling (Russo & Van Roy, 2014; Nikolov et al., 2019) approaches.
Revised: Ming, 2021); and 3) information-directed sampling (Russo & Van Roy, 2014; Nikolov et al., 2019) approaches.

Original: Among them, optimism-based methods needs one to construct the confidence set that contains the target model/policy with high probability, which suffers from scalability issues (Osband & Van Roy, 2017); in addition this approach empirically performs worse than Thompson sampling (Chapelle & Li, 2011).
Revised: Among them, optimism-based methods needs one to construct the confidence set that contains the target model/policy with high probability, which suffers from scalability issues (Osband & Van Roy, 2017); in addition this approach empirically performs worse than Thompson sampling (Chapelle & Li, 2011).

Original: Information-directed sampling can be better than optimism-based methods and Thompson sampling, as it directly minimizes the “regret per information bit” (Russo &
Revised: Information-directed sampling can be better than optimism-based methods and Thompson sampling, as it directly minimizes the “regret per information bit” (Russo &

Original: Van Roy, 2014).
Revised: Van Roy, 2014).

Original: However, it relies on estimating the mutual information between random variables, which is especially difficult for high-dimensional continuous random variables (McAllester & Stratos, 2020).
Revised: However, it relies on estimating the mutual information between random variables, which is especially difficult for high-dimensional continuous random variables (McAllester & Stratos, 2020).

Original: Therefore, we focus on posterior sampling.
Revised: Therefore, we focus on posterior sampling.

Original: However, different from prior works, we study the effect of approximate inference in an RL setting.
Revised: However, different from prior works, we study the effect of approximate inference in an RL setting.

Original:  S UMMARY AND FUTURE WORK
Revised:  S UMMARY AND FUTURE WORK

Original: In this paper, we presented PS-MBPO and OPS-MBPO as two algorithms for efficient model-based reinforcement learning in complex environments.
Revised: In this paper, we presented PS-MBPO and OPS-MBPO as two algorithms for efficient model-based reinforcement learning in complex environments.

Original: We demonstrate that both PS-MBPO and OPSMBPO can greatly improve the sample efficiency in online reinforcement learning, and surpass various baseline methods by a large margin, especially on sparse reward tasks.
Revised: We demonstrate that both PS-MBPO and OPSMBPO can greatly improve the sample efficiency in online reinforcement learning, and surpass various baseline methods by a large margin, especially on sparse reward tasks.

Original: We hope that, beyond our specific approach, our analysis can inspire future works to propose improved factorizations of the posterior over policies and models.
Revised: We hope that, beyond our specific approach, our analysis can inspire future works to propose improved factorizations of the posterior over policies and models.

Original: In the future, we would like to explore automatically adapting the value of λ in an online fashion.
Revised: In the future, we would like to explore automatically adapting the value of λ in an online fashion.

Original: (See Section B.2 for some preliminary results.)
Revised: (See Section B.2 for some preliminary results.)

Original: In addition, we would like to extend the result of
Revised: In addition, we would like to extend the result of

Original: Phan et al.
Revised: Phan et al.

Original: and prove a sublinear regret for PSRL under approximate inference.
Revised: and prove a sublinear regret for PSRL under approximate inference.

Original: It would also be interesting to explore epistemic neural networks (Osband et al., 2021) and transformers (Vaswani et al., 2017; M ¨ uller et al., 2021) as alternatives to deep ensembles for approximate posterior inference.
Revised: It would also be interesting to explore epistemic neural networks (Osband et al., 2021) and transformers (Vaswani et al., 2017; M ¨ uller et al., 2021) as alternatives to deep ensembles for approximate posterior inference.

Original: Lastly, making information directed sampling (Russo & Van Roy, 2014) practical for reinforcement learning problems is another promising direction to further improve the sample efficiency.
Revised: Lastly, making information directed sampling (Russo & Van Roy, 2014) practical for reinforcement learning problems is another promising direction to further improve the sample efficiency.

Original: Reproducibility Statement.
Revised: Reproducibility Statement.

Original: We will open source the code for reproducing the results in our paper.
Revised: We will open source the code for reproducing the results in our paper.

Original: 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.
Revised: 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.

Original: C ONTENTS 1 Introduction
Revised: C ONTENTS 1 Introduction

Original: 2 Problem statement and theoretical results 3 Method 3.
Revised: 2 Problem statement and theoretical results 3 Method 3.

Original: Posterior decomposition . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
Revised: Posterior decomposition . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

Original: Proposed algorithm .
Revised: Proposed algorithm .

Original: Sampling policies . .
Revised: Sampling policies . .

Original: 4 Experiments 4.
Revised: 4 Experiments 4.

Original: Experimental Setup . .
Revised: Experimental Setup . .

Original: Comparison with existing methods . . . . . . . . . . . . . . . . . . . . . . . . . .
Revised: Comparison with existing methods . . . . . . . . . . . . . . . . . . . . . . . . . .

Original: Ablation Studies . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
Revised: Ablation Studies . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

Original: 5 Related Works
Revised: Works

Original:  Summary 
Revised: 6 Summary and future work

Original: and future work

Original: A Extended Backgrounds
Revised: A Extended Backgrounds

Original: A.1 Dynamics model . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
Revised: A.1 Dynamics model . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

Original: A.2 Soft actor-critic (SAC) . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
Revised: A.2 Soft actor-critic (SAC) . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

Original: A.3 Revisiting MBPO, PETS and ME-TRPO . . . . . . . . . . . . . . . . . . . . . . .
Revised: A.3 Revisiting MBPO, PETS and ME-TRPO . . . . . . . . . . . . . . . . . . . . . . .

Original: A.4 Posterior sampling reinforcement learning . .
Revised: A.4 Posterior sampling reinforcement learning . .

Original: B Additional Results for (O)PS-MBPO
Revised: B Additional Results for (O)PS-MBPO

Original: B.1 Visualization of the optimistic weights . . . . . . . . . . . . . . . . . . . . . . . .
Revised: B.1 Visualization of the optimistic weights . . . . . . . . . . . . . . . . . . . . . . . .

Original: B
Revised: B.

Original: .2 How does the temperature term τ affect the performance
Revised: How does the temperature term τ affect the performance? . . . . . . . . . . . . .

Original: B.3 How does the schedule of λ affect the performance
Revised: B .3 How does the schedule of λ affect the performance? . . . . . . . . . . . . . . . .

Original: C Additional Details about the Algorithm and Experiments
Revised: C Additional Details about the Algorithm and Experiments

Original: C.1  Implementation Details . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
Revised: C.1 Algorithm Details . . . C.2 Implementation Details . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

Original: C.2 Task Details . ..
Revised: C.3 Task Details . .

Original:  Exploration
Revised: D Forced Exploration

Original: E Random Function
Revised: E Random Function

Original: Prior
Revised: Prior

Original: F Additional Umap Visualizations
Revised: F Additional Umap Visualizations

Original: G Proofs
Revised: G Proofs

Original: G.1 Proof of Theorem 1 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
Revised: G. Proof of Theorem 1 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

Original: G.2 Supporting Lemmas . .
Revised: [Table] Supporting [Table]

Original: A E XTENDED B ACKGROUNDS A. D YNAMICS MODEL
Revised: A E XTENDED B ACKGROUNDS A. D YNAMICS MODEL

Original: We use deep ensemble for fitting the environment dynamics.
Revised: We use deep ensemble for fitting the environment dynamics.

Original: For each network in the ensemble f θ , it takes a whitened state and action pair as input, and predicts the residual of the next state as well as the reward, i.e.,
Revised: For each network in the ensemble f θ , it takes a whitened state and action pair as input, and predicts the residual of the next state as well as the reward, i.e.,

Original: [Equation]
Revised: [Equation]

Original: where µ s , µ a are the empirical mean of the states and actions, σ s and σ a are the empirical standard deviation of them.
Revised: where µ s , µ a are the empirical mean of the states and actions, σ s and σ a are the empirical standard deviation of them.

Original: Then, the predictions of next state and reward will be
Revised: Then, the predictions of next state and reward will be

Original: [Equation]
Revised: [Equation]

Original: Below is our implementation of each individual neural network in JAX (Bradbury et al., 2018).
Revised: Below is our implementation of each individual neural network in JAX (Bradbury et al., 2018).

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Revised: [Equation]

Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]

Original: [Equation]

Original: [Table]
Revised: [Table]

Original: def __call__(self x): h = self mlp(x) mean, logvar  jnp split self mean logvar
Revised: def __call__(self x): h = self mlp(x) mean, logvar = jnp split self mean _logvar

Original: , 2, axis logvar = self max  jax nn logvar self min  nn min mean, logvar
Revised: , 2, axis logvar = self max logvar jax nn logvar self min logvar nn min mean, logvar

Original: A. S OFT ACTOR - CRITIC (SAC)
Revised: A. S OFT ACTOR - CRITIC (SAC)

Original: We use SAC for learning the policies.
Revised: We use SAC for learning the policies.

Original: In a highlevel, SAC is a maximum entropy RL algorithm, which typically optimizes the following objective,
Revised: In a highlevel, SAC is a maximum entropy RL algorithm, which typically optimizes the following objective,

Original: [Equation]
Revised: [Equation]

Original: As a result, maximum entropy RL algorithm will favor those policies that not only optimize for the reward, but also has a large entropy.
Revised: As a result, maximum entropy RL algorithm will favor those policies that not only optimize for the reward, but also has a large entropy.

Original: This can in turn improve the robustness of the optimized policy.
Revised: This can in turn improve the robustness of the optimized policy.

Original: As for SAC, it searches the policy by iteratively solving the policy evaluation and policy improvement steps.
Revised: As for SAC, it searches the policy by iteratively solving the policy evaluation and policy improvement steps.

Original: [Equation]
Revised: [Equation]

Original: In the practical implementation of SAC
Revised: In the practical implementation of SAC

Original: [Equation]
Revised: [Equation]

Original: it uses a separate function approximator for the state value to stabilize the training.
Revised: it uses a separate function approximator for the state value to stabilize the training.

Original: Specifically, there are three components in SAC, a parameterized state value function V ψ ( s ) , a soft Q-function Q θ ( s , a )
Revised: Specifically, there are three components in SAC, a parameterized state value function V ψ ( s ) , a soft Q-function Q θ ( s , a )

Original: [Equation]
Revised: [Equation]

Original: and a policy π ϕ ( a | s ) .
Revised: and a policy π ϕ ( a | s ) .

Original: The objectives for each component are
Revised: The objectives for each component are

Original: [Equation]
Revised: [Equation]

Original: where Z θ ( · ) is a normalizing constant
Revised: where Z θ ( · ) is a normalizing constant

Original: [Equation]
Revised: [Equation]

Original: and ˆ Q ( s , a ) is defined as
Revised: and ˆ Q ( s , a ) is defined as

Original: [Equation]
Revised: [Equation]

Original: Additionally, ¯ ψ is the exponentially moving average of the weights of the value network, and J π ( ϕ ) can be optimized with reparameterization trick under Gaussian case, which can further reduces the variance of the gradient estimator and hence stabilizes the training.
Revised: Additionally, ¯ ψ is the exponentially moving average of the weights of the value network, and J π ( ϕ ) can be optimized with reparameterization trick under Gaussian case, which can further reduces the variance of the gradient estimator and hence stabilizes the training.

Original: We adopt the SAC implementation from Acme (Hoffman et al., 2020).
Revised: We adopt the SAC implementation from Acme (Hoffman et al., 2020).

Original: A. R EVISITING MBPO, PETS AND ME-TRPO
Revised: A. R EVISITING MBPO, PETS AND ME-TRPO

Original: Popular model-based reinforcement learning algorithms such as ME-TRPO (Kurutach et al., 2018),
Revised: Popular model-based reinforcement learning algorithms such as ME-TRPO (Kurutach et al., 2018),

Original: PETs (Chua et al., 2018) and MBPO (Janner et al., 2019) typically repeat the following three steps: 1) train a dynamics model (or an ensemble of models)
Revised: PETs (Chua et al., 2018) and MBPO (Janner et al., 2019) typically repeat the following three steps: 1) train a dynamics model (or an ensemble of models)

Original: q ( M|D E ) ; 2) train/extract a policy π ⋆ from the learned model; 3) collect data from the environment with the policy.
Revised: q ( M|D E ) ; 2) train/extract a policy π ⋆ from the learned model; 3) collect data from the environment with the policy.

Original: Consequently, their policy is (approximately) equivalent to the one obtained by solving
Revised: Consequently, their policy is (approximately) equivalent to the one obtained by solving

Original: [Equation]
Revised: [Equation]

Original: where the posterior of the model M is approximated by an ensemble of neural networks, Π is the search space of policies, and the cumulative reward R M ( π ) for an episode of length H of a policy π under dynamics model M is defined as
Revised: where the posterior of the model M is approximated by an ensemble of neural networks, Π is the search space of policies, and the cumulative reward R M ( π ) for an episode of length H of a policy π under dynamics model M is defined as

Original: [Equation]
Revised: [Equation]

Original: where s t +1 ∼ p M ( s | s t , a t ) and a t ∼ π ( a | s t ) .
Revised: where s t +1 ∼ p M ( s | s t , a t ) and a t ∼ π ( a | s t ) .

Original: [Equation]
Revised: [Equation]

Original: However, the above strategy only accounts for exploitation, so will lead to low data efficiency.
Revised: However, the above strategy only accounts for exploitation, so will lead to low data efficiency.

Original: A. P OSTERIOR SAMPLING REINFORCEMENT LEARNING
Revised: A. P OSTERIOR SAMPLING REINFORCEMENT LEARNING

Original: The idea of PSRL is introduced by Strens (2000).
Revised: The idea of PSRL is introduced by Strens (2000).

Original: The first regret bound ˜ O ( HS √ AT ) for PSRL is proved by Osband et al.
Revised: The first regret bound ˜ O ( HS √ AT ) for PSRL is proved by Osband et al.

Original: 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.
Revised: 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.

Original: In Osband & Van Roy (2017), the bound was improved to ˜ O
Revised: In Osband & Van Roy (2017), the bound was improved to ˜ O

Original: ( H √ SAT ) .
Revised: ( H √ SAT ) .

Original: For the continuous case, Osband &
Revised: For the continuous case, Osband &

Original: Van Roy (2014) provides the first regret bound ˜ O ( √ d K d E T ) based on the eluder dimension d E and
Revised: Van Roy (2014) provides the first regret bound ˜ O ( √ d K d E T ) based on the eluder dimension d E and

Original: Kolmogorov dimension d K .
Revised: Kolmogorov dimension d K .

Original: More recently, Fan & Ming (2021) study the regret bound for PSRL under Gaussian process assumption, and obtain a regret bound of
Revised: More recently, Fan & Ming (2021) study the regret bound for PSRL under Gaussian process assumption, and obtain a regret bound of

Original: ˜ O
Revised: ˜ O

Original: ( H 3 / 2 d √ T ) .
Revised: ( H 3 / 2 d √ T ) .

Revised: In addition to the bound on Bayesian regret, there is 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 a regret bound of order ˜

Revised: O ( √ T ) .

Original: Nevertheless, all these regret bounds are derived under exact Thompson sampling.
Revised: Nevertheless, all these regret bounds are derived under exact Thompson sampling.

Original: B A DDITIONAL R ESULTS FOR (O)PS-MBPO B. V ISUALIZATION OF THE OPTIMISTIC WEIGHTS
Revised: B A DDITIONAL R ESULTS FOR (O)PS-MBPO B. V ISUALIZATION OF THE OPTIMISTIC WEIGHTS

Original: For OPS-MBPO, we will maintain the weights for each policy throughout the entire process of online learning.
Revised: For OPS-MBPO, we will maintain the weights for each policy throughout the entire process of online learning.

Original: To investigate how those weights evolve, we plot the weights of each policy in Figureand Figure 10, which covers one dense reward task and one sparse reward task.
Revised: To investigate how those weights evolve, we plot the weights of each policy in Figureand Figure 10, which covers one dense reward task and one sparse reward task.

Original: The leftmost figure corresponds to the weights in the initial phase, which will change more rapidly than the later phases.
Revised: The leftmost figure corresponds to the weights in the initial phase, which will change more rapidly than the later phases.

Original: 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.
Revised: 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.

Original: More interestingly, the reward curve in the rightmost figure is also consistent with the pattern in the optimistic weights curve.
Revised: More interestingly, the reward curve in the rightmost figure is also consistent with the pattern in the optimistic weights curve.

Original: B. H OW DOES THE TEMPERATURE TERM τ AFFECT THE PERFORMANCE ?
Revised: B. H OW DOES THE TEMPERATURE TERM τ AFFECT THE PERFORMANCE ?

Original: We further study how does the temperature term will affect the performance on both dense reward and sparse reward tasks.
Revised: We further study how does the temperature term will affect the performance on both dense reward and sparse reward tasks.

Original: The results can be found in Figure 11 and Figure 12.
Revised: The results can be found in Figure 11 and Figure 12.

Original: We observe that the temperature term will affect the convergence speed of the reward on most of the tasks.
Revised: We observe that the temperature term will affect the convergence speed of the reward on most of the tasks.

Original: For some of the tasks, such as Ant , Hopper and Walker2d , it will also affect the converged reward slightly.
Revised: For some of the tasks, such as Ant , Hopper and Walker2d , it will also affect the converged reward slightly.

Original: In general, we recommend the temperature term to be around five times of the best averaged episodic reward that can be achieved
Revised: In general, we recommend the temperature term to be around five times of the best averaged episodic reward that can be achieved

Original: . B. H OW DOES THE SCHEDULE OF λ AFFECT THE PERFORMANCE ?
Revised: . B. H OW DOES THE SCHEDULE OF λ AFFECT THE PERFORMANCE ?

Original: 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.
Revised: 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.

Original: We consider two schemes for adjusting λ , i.e., 1) a constant schedule; and 2) a linear schedule.
Revised: We consider two schemes for adjusting λ , i.e., 1) a constant schedule; and 2) a linear schedule.

Original: For constant schedule, we fix the value of λ throughout training, whereas for linear schedule, we decrease the value of λ from 1 tolinearly.
Revised: For constant schedule, we fix the value of λ throughout training, whereas for linear schedule, we decrease the value of λ from 1 tolinearly.

Original: The rightmost figure in Figure 13 visualizes the difference between these two schedules.
Revised: The rightmost figure in Figure 13 visualizes the difference between these two schedules.

Original: To make them comparable, we ensure that the areas under both curves are of the same size, so that the total amount of real-world data is the same.
Revised: To make them comparable, we ensure that the areas under both curves are of the same size, so that the total amount of real-world data is the same.

Original: The comparison on three tasks are shown in the left three figures of Figure 13.
Revised: The comparison on three tasks are shown in the left three figures of Figure 13.

Original: We observe that the linear schedule has very little effect on the dense reward tasks, though slightly improves the the final reward in Hopper .
Revised: We observe that the linear schedule has very little effect on the dense reward tasks, though slightly improves the the final reward in Hopper .

Original: For Cartpole-swingup , the performance of linear schedule improves faster than constant schedule, but achieves similar rewards in the end.
Revised: For Cartpole-swingup , the performance of linear schedule improves faster than constant schedule, but achieves similar rewards in the end.

Original: 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 λ based on the model’s validation loss.
Revised: 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 λ based on the model’s validation loss.

Original: For simplicity, we recommend to use a constant schedule in practice.
Revised: For simplicity, we recommend to use a constant schedule in practice.

Original: For sparse reward tasks the  range of λ can be { reward
Revised: For sparse reward tasks, the search range of λ can be { 0 , 0 . 1 , 0 . 3 , 0 . 5 , 0 .

Revised: In addition to the grid search, we believe it’s also possible to adapt the value of λ online.

Original: C A DDITIONAL D ETAILS ABOUT THE A LGORITHM AND E XPERIMENTS Algorithm 2 PS-MBPO (abstract formulation)
Revised: We can cast the the problem of choosing the optimal value of λ as a bandit problem. The high-level idea of the algorithm is: 1) Initialize a set of possible values for λ , and treat each value of λ as an arm in bandit. 2) Apply any no-regret learning algorithms for solving it, e g., explore-then-commit

Revised: However,

Revised: we haven’t test this algorithm yet,

Revised: and

Revised: it would be

Revised: interesting

Revised: as

Revised: a future extension.

Revised: C A DDITIONAL D ETAILS ABOUT THE A LGORITHM AND E XPERIMENTS

Revised: Algorithm 2 PS-MBPO (abstract formulation)

Original: Require: Prior distributions q ( M ) , q ( π ) and tuning hyperparameter λ .
Revised: Require: Prior distributions q ( M ) , q ( π ) and tuning hyperparameter λ .

Original: Require: Initialize an empty dataset D E for storing data collected from the environment.
Revised: Require : Initialize an empty dataset D E for storing data collected from the environment.

Original: 1: for K episodes do 2:
Revised: 1: for K episodes do 2:

Original: Fit the posterior of the policy q ( π |D E , λ ) on data D E using equation 13.
Revised: Fit the posterior of the policy q ( π |D E , λ ) on data D E using equation 13.

Original: Sample a policy π k ∼ q ( π |D E , λ ) from the posterior distribution.
Revised: Sample a policy π k ∼ q ( π |D E , λ ) from the posterior distribution.

Original: 4: for H steps do 5:
Revised: 4: for H steps do 5:

Original: Run the policy π k in the environment and add the collected data to D E .
Revised: Run the policy π k in the environment and add the collected data to D E .

Original: 6: end for 7: end for
Revised: 6: end for 7: end for

Revised: [Figure]

Original: C. I MPLEMENTATION D ETAILS
Revised: C. A LGORITHM D ETAILS

Revised: In Algorithm 1, we approximate the posterior of MDPs and policies, i.e., q ( M|D E ) and q ( π |M , D E , λ ) using deep ensemble, which can be regarded as a finite particle approximation to the posterior.

Revised: Specifically, q ( M|D E ) is approximated by {M ˆ θ n } Nn =1 and q ( π |M ˆ θ n , D E , λ ) is approximated by { π ˆ ϕ n,m

Revised: } Mm =1 for all n ∈ [ N ] , where both M ˆ θ n and π ˆ ϕ n,m are implemented using a multi-layer perceptron (MLP) with parameters ˆ θ n trained on D E and ˆ ϕ n,m trained on the mixed dataset λ D E + (1 − λ )

Revised: D n,m M , respectively.

Revised: By mixed dataset λ D E + (1 − λ ) D n,m M , we mean that for each data point in the training batch, it is with probability of λ being sampled from the real data D E and probability of 1 − λ from the fictitious data D n,m M .

Revised: C. I MPLEMENTATION D ETAILS

Original: In this section, we provide the additional details about our algorithm and experiments.
Revised: In this section, we provide the additional details about our algorithm and experiments.

Original: We provide a detailed description of our approach in Algorithm 1 and a visual illustration about its difference with MBPO in Figure 14.
Revised: We provide a detailed description of our approach in Algorithm 1 and a visual illustration about its difference with MBPO in Figure 14.

Original: In terms of the hyperparameters, our choice of them are mostly the same as the ones adopted in MBPO (Janner et al., 2019) and Pineda et al.[Figure]
Revised: In terms of the hyperparameters, our choice of them are mostly the same as the ones adopted in MBPO (Janner et al., 2019) and Pineda et al.

Original:  for Ant , Halfcheetah , Hopper , Walker2D and Cartpole-swingup , and Eysenbach et al.
Revised: (2021) for Ant , Halfcheetah , Hopper , Walker2D and Cartpole-swingup , and Eysenbach et al.

Original: (2021) for Window-open-v2 , which are sufficiently optimized by the authors for MBPO.
Revised: (2021) for Window-open-v2 , which are sufficiently optimized by the authors for MBPO.

Original: Specifically, the hyperparameters of MBPO are directly adopted from https://github.com/ facebookresearch/mbrl-lib for dense reward tasks.
Revised: Specifically, the hyperparameters of MBPO are directly adopted from https://github.com/ facebookresearch/mbrl-lib for dense reward tasks.

Original: For sparse reward tasks, the hyperparameters are adopted from Eysenbach et al.
Revised: For sparse reward tasks, the hyperparameters are adopted from Eysenbach et al.

Original: The hyperparameters for our method on each task are reported in Table 1.
Revised: The hyperparameters for our method on each task are reported in Table 1.

Original: We will also release our code for reproducing all the experiments.
Revised: We will also release our code for reproducing all the experiments.

Original: Next, we introduce the details about each tasks.
Revised: Next, we introduce the details about each tasks.

Original: C. T ASK D ETAILS
Revised: C. T ASK D ETAILS

Original: Ant, Halfcheetah, Hopper and Walker2D. These tasks are taken from the official Github repository of MBPO (Janner et al., 2019), https://github.com/jannerm/mbpo .
Revised: Ant, Halfcheetah, Hopper and Walker2D. These tasks are taken from the official Github repository of MBPO (Janner et al., 2019), https://github.com/jannerm/mbpo .

Original: Ball-in-Cup and Cartpole-swingup.
Revised: Ball-in-Cup and Cartpole-swingup.

Original: These tasks are taken from the deepmind control suite (Tunyasuvunakool et al., 2020).
Revised: These tasks are taken from the deepmind control suite (Tunyasuvunakool et al., 2020).

Original: More details can be found in the Github repository at https: //github.com/deepmind/dm_control/ .
Revised: More details can be found in the Github repository at https: //github.com/deepmind/dm_control/ .

Original: Window-open-v2.
Revised: Window-open-v2.

Original: This task is based on the original Window-open-v2 in Metaworld benchmakr (Yu et al., 2020).
Revised: This task is based on the original Window-open-v2 in Metaworld benchmakr (Yu et al., 2020).

Original: The sparse reward is 0 only if the window is within 3 units of the open position, and − 10 for all other positions. 
Revised: The sparse reward is 0 only if the window is within 3 units of the open position, and − 10 for all other positions. [Table] [Figure] [Table]

Revised: [Figure]

Original:  D F ORCED E XPLORATION
Revised: [Table] [Figure] [Table] D F ORCED E XPLORATION

Original: Forced exploration is proposed in Phan et al.
Revised: Forced exploration is proposed in Phan et al.

Original: (2019) to improve approximate Thompson sampling for bandit problems.
Revised: (2019) to improve approximate Thompson sampling for bandit problems.

Original: Without properly dealing with the approximate inference error, there will be an extra term in the regret that is linear in T , regardless how small the error is.
Revised: Without properly dealing with the approximate inference error, there will be an extra term in the regret that is linear in T , regardless how small the error is.

Original: In their paper, they use the α -divergence for measuring the approximate inference error, defined as
Revised: In their paper, they use the α -divergence for measuring the approximate inference error, defined as

Original: [Equation]
Revised: [Equation]

Original: [Equation]

Original: [Figure]

Original: [Table]

Original: [Figure]

Original: [Table]

Original: [Figure]

Original: [Table] [Figure] [Table] The α -divergence can capture many divergences, including forward KL divergence ( α → 1 ), backward KL divergence ( α → 0 ), Hellinger distance ( α = 0 . 5 ) and χ 2 divergence ( α = 2 ).
Revised: The α -divergence can capture many divergences, including forward KL divergence ( α → 1 ), backward KL divergence ( α → 0 ), Hellinger distance ( α = 0 . 5 ) and χ 2 divergence ( α = 2 ).

Original: Different inference methods will give error guarantee measured by α -divergence with different α .
Revised: Different inference methods will give error guarantee measured by α -divergence with different α .

Original: We are interested in the error guarantee under the reverse KL-divergence, i.e., α = 0 , as the ensemble sampling (Lu & Van Roy, 2017) provides error guarantees under the reverse Kl-divergence.
Revised: We are interested in the error guarantee under the reverse KL-divergence, i.e., α = 0 , as the ensemble sampling (Lu & Van Roy, 2017) provides error guarantees under the reverse Kl-divergence.

Original: In Phan et al.
Revised: In Phan et al.

Original: (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 α -divergence with α ≤ 0 .
Revised: (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 α -divergence with α ≤ 0 .

Original: The reverse KL-divergence falls in this case.
Revised: The reverse KL-divergence falls in this case.

Original: Specifically, the forced exploration is a simple method, that maintains a probability of random exploration.
Revised: Specifically, the forced exploration is a simple method, that maintains a probability of random exploration.

Original: This probability decays as t , the online steps, grows.
Revised: This probability decays as t , the online steps, grows.

Original: 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.
Revised: 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.

Original: In our experiments, we consider the following exploration rate
Revised: In our experiments, we consider the following exploration rate

Original: [Equation]
Revised: [Equation]

Original: where k is the index for the episode, and τ is the hyperparameter for controlling the frequency of forced exploration.
Revised: where k is the index for the episode, and τ is the hyperparameter for controlling the frequency of forced exploration.

Original: As k increases, the random exploration probability will decrease.
Revised: As k increases, the random exploration probability will decrease.

Original: In our experiments, we consider τ ∈ { 1 , 5 , 10 } .
Revised: In our experiments, we consider τ ∈ { 1 , 5 , 10 } .

Original: All the other settings are the same as our experiments in the main paper.
Revised: All the other settings are the same as our experiments in the main paper.

Original: The results are presented in Figure 16.
Revised: The results are presented in Figure 16.

Original: We observe that forced exploration is mostly not helpful in our experiments, except for the Hopper task.
Revised: We observe that forced exploration is mostly not helpful in our experiments, except for the Hopper task.

Original: Moreover, increasing τ usually make the performance even worse.
Revised: Moreover, increasing τ usually make the performance even worse.

Original: 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.
Revised: 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.

Original: We leave the theoretical analysis as a future work.
Revised: We leave the theoretical analysis as a future work.

Original: E R ANDOM F UNCTION P RIOR
Revised: E R ANDOM F UNCTION P RIOR

Original: The random function prior (RFP) is proposed in Osband et al.
Revised: The random function prior (RFP) is proposed in Osband et al.

Original: (2018) for improving the uncertainty estimation.
Revised: (2018) for improving the uncertainty estimation.

Original: The prior network are chosen for modelling the uncertainty that does not come from the observed data.
Revised: The prior network are chosen for modelling the uncertainty that does not come from the observed data.

Original: The RFPs can also be viewed as a regularization in the output space.
Revised: The RFPs can also be viewed as a regularization in the output space.

Original: In contrast to weight space regularization, RFP makes it easier to incorporate different property (e.g., periodicity) of the function to be learned as a prior information.
Revised: In contrast to weight space regularization, RFP makes it easier to incorporate different property (e.g., periodicity) of the function to be learned as a prior information.

Original: More importantly, when using deep ensemble, incorporating the RFP is fairly simple.
Revised: More importantly, when using deep ensemble, incorporating the RFP is fairly simple.

Original: It modifies the original training objective ℓ ( f θ , D ) by adding an extra regularization term,
Revised: It modifies the original training objective ℓ ( f θ , D ) by adding an extra regularization term,

Original: [Equation]
Revised: [Equation]

Original: (34) where β is a scaling term for adjusting the effect of the prior, f θ 0 is the prior network which is held fixed during training.
Revised: where β is a scaling term for adjusting the effect of the prior, f θ 0 is the prior network which is held fixed during training.

Original: Hence, we also conduct experiments with RFPs in our experiments to investigate how does the RFPs will affect the learning of dynamics models.
Revised: Hence, we also conduct experiments with RFPs in our experiments to investigate how does the RFPs will affect the learning of dynamics models.

Original: We vary the value of β in { 0 . 1 , 0 . 3 , 1 } .
Revised: We vary the value of β in { 0 . 1 , 0 . 3 , 1 } .

Original: The results are reported in Figure 17.
Revised: The results are reported in Figure 17.

Original: Firstly, by properly choosing the value of β , RFPs slightly improve the performance on Hopper , and don’t affect the performance a lot on Walker2D .
Revised: Firstly, by properly choosing the value of β , RFPs slightly improve the performance on Hopper , and don’t affect the performance a lot on Walker2D .

Original: However, for Window-open-v2 , RFPs will hurt the performance a lot.
Revised: However, for Window-open-v2 , RFPs will hurt the performance a lot.

Original: We conjecture that this might because our choice of the prior function on Window-open-v is not suitable for the task, i.e., the reward is sparse in Window-open-v2 , but the RFPs don’t induce sparsity on the predictions.
Revised: We conjecture that this might because our choice of the prior function on Window-open-v is not suitable for the task, i.e., the reward is sparse in Window-open-v2 , but the RFPs don’t induce sparsity on the predictions.

Original: Nevertheless, one interesting observation is that both forced exploration and RFPs seem to help on
Revised: Nevertheless, one interesting observation is that both forced exploration and RFPs seem to help on

Original: Hopper , and their overall pattern on three tasks is a bit consistent.
Revised: Hopper , and their overall pattern on three tasks is a bit consistent.

Original: Therefore, it would be interesting to figure out if there is a deep connection between the forced exploration and RFPs.
Revised: Therefore, it would be interesting to figure out if there is a deep connection between the forced exploration and RFPs.

Original: F A DDITIONAL U MAP V ISUALIZATIONS
Revised: F A DDITIONAL U MAP V ISUALIZATIONS

Original: We provide the Umap visualization of the state embeddings of PS-MBPO and MBPO during training in Figure 18 and Figure 19.
Revised: We provide the Umap visualization of the state embeddings of PS-MBPO and MBPO during training in Figure 18 and Figure 19.

Original: We observe that PS-MBPO will mostly cover a more broad range of the embedding space, and its pattern also evolves more rapidly than MBPO.
Revised: We observe that PS-MBPO will mostly cover a more broad range of the embedding space, and its pattern also evolves more rapidly than MBPO.

Original: G P ROOFS
Revised: G P ROOFS

Original: In this section, we present the proof of Theorem 1.
Revised: In this section, we present the proof of Theorem 1.

Original: 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.
Revised: 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.

Original: G. P ROOF OF T HEOREM
Revised: G. P ROOF OF T HEOREM

Original: Theorem 1 For K episodes, the Bayesian regret of posterior sampling reinforcement learning algorithm A with any approximate posterior distribution q k at episode k is upper bounded by
Revised: Theorem 1 For K episodes, the Bayesian regret of posterior sampling reinforcement learning algorithm A with any approximate posterior distribution q k at episode k is upper bounded by

Original: [Equation]
Revised: [Equation]

Original: where H ( π ⋆ ) is the entropy of the prior distribution of  polices,
Revised: where H ( π ⋆ ) is the entropy of the prior distribution of optimal polices,

Revised: i.e., p ( π ) = (cid:82) δ

Revised: [Equation]

Original: i.e., p ( π ) = (cid:82) δ ( π |M ) p ( M ) d M ,and C is some problem-dependent constant.
Revised: π |M ) p [Equation] M ) d M , C is some problem-dependent constant and d KL ( ·| · ) is the KL-divergence.

Original: Proof: Recall the definition of Bayesian regret,
Revised: Proof: Recall the definition of Bayesian regret,

Original: [Equation]
Revised: [Equation]

Original: Let’s denote history at the beginning of episode k as H k .
Revised: Let’s denote history at the beginning of episode k as H k .

Original: Then, we can rewrite the Bayesian regret as
Revised: Then, we can rewrite the Bayesian regret as

Original: BayesianRegret

Original: [Equation]
Revised: [Equation]

Original: [Equation]

Original: By doing so, we can bound each term E [∆ k | H k ] separately.
Revised: By doing so, we can bound each term E [∆ k | H k ] separately.

Original: For convenience, we define E k [∆ k ] :=
Revised: For convenience, we define E k [∆ k ] :=

Original: E [∆ k | H k ] .
Revised: E [∆ k | H k ] .

Original: Then, by Lemma 1, we can further decompose it into,
Revised: Then, by Lemma 1, we can further decompose it into,

Original: [Equation]
Revised: [Equation]

Original: (37) where

Original: [Equation]

Original: 38) and
Revised: ) where

Original: [Equation]
Revised: [Equation]

Original: Then
Revised: and

Original: [Equation]
Revised: [Equation]

Original:  it remains to bound (cid:80)
Revised: Then, it remains to bound (cid:80)

Original: Kk =1 E [ G k ] and (cid:80)
Revised: Kk =1 E [ G k ] and (cid:80)

Original: Kk =1 E [ D k ] .
Revised: Kk =1 E [ D k ] .

Original: By Lemma [Equation] we can bound the sum of expectation of D k by
Revised: By Lemma 2, we can bound the sum of expectation of D k by

Original: [Equation]
Revised: [Equation]

Original: By Lemma [Equation] the upper bound for the sum of the expectation of G k is
Revised: By Lemma 3, the upper bound for the sum of the expectation of G k is

Original: [Equation]
Revised: [Equation]

Original: Hence

Original: [Equation]
Revised: [Equation]

Original:  the term D k captures the regret incurred by the approximate inference error [Equation] and G k captures the standard regret for Thompson sampling, which is of order ˜ O
Revised: Hence, the term D k captures the regret incurred by the approximate inference error, and G k captures the standard regret for Thompson sampling, which is of order ˜ O

Original: ( √ K ) .
Revised: ( √ K ) .

Original: By combining them together, we finally arrive at the upper bound of the Bayesian regret 
Revised: By combining them together, we finally arrive at the upper bound of the Bayesian regret BayesianRegret

Original: [Equation]
Revised: [Equation]

Original: G. S UPPORTING L EMMAS
Revised: G. S UPPORTING L EMMAS

Original: Lemma 

Revised: Lemma

Original: For each time k
Revised: For each time k

Original: = 1 , ..., K , we have
Revised: = 1 , ..., K , we have

Original: [Equation]
Revised: [Equation]

Original: where
Revised: where

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: and
Revised: and

Original: [Equation]
Revised: [Equation]

Original: Proof: Conditioning on the history H k , we can write the conditional Bayesian regret as
Revised: Proof: Conditioning on the history H k , we can write the conditional Bayesian regret as

Original: [Equation]
Revised: [Equation]

Original: where the second equality holds because the value function is independent of the instantiation of the policy π k when given the history H k .
Revised: where the second equality holds because the value function is independent of the instantiation of the policy π k when given the history H k .

Original: □
Revised: □

Original: Lemma 

Revised: Lemma

Original: For any k = 1 , ..., K , we have
Revised: For any k = 1 , ..., K , we have

Original: [Equation]
Revised: [Equation]

Original: Proof: Recall D k ,
Revised: Proof: Recall D k ,

Original: [Equation]
Revised: [Equation]

Original: By using the Cauchy-Schwarz inequality
Revised: By using the Cauchy-Schwarz inequality

Original: [Equation]
Revised: [Equation]

Original: we have
Revised: we have

Original: D k ≤
Revised: D k ≤

Original: [Equation]
Revised: [Equation]

Original: By the definition of Hellinger distance d H ( ·| · ) between two random variables
Revised: By the definition of Hellinger distance d H ( ·| · ) between two random variables

Original: [Equation]
Revised: [Equation]

Original: we have
Revised: we have

Original: D k
Revised: D k

Original: [Equation]
Revised: [Equation]

Original: Since [ d H ( ·| · )] 2 ≤ d KL ( ·| · ) and V Mπ, 1 is a bounded random variable with HR max as its upper bound, we have
Revised: Since [ d H ( ·| · )] 2 ≤ d KL ( ·| · ) and V Mπ, 1 is a bounded random variable with HR max as its upper bound, we have

Original: [Equation]
Revised: [Equation]

Original: Hence, we have
Revised: Hence, we have

Original: [Equation]
Revised: [Equation]

Original: 55) □
Revised: 55) □

Original: Lemma 

Revised: Lemma

Original: For each k = 1 , ..., K , we have
Revised: For each k = 1 , ..., K , we have

Original: [Equation]
Revised: [Equation]

Original: Proof: Recall the definition of G k ,
Revised: Proof: Recall the definition of G k ,

Original: [Equation]
Revised: [Equation]

Original: Since V Mπ, 1 (here, we drop the dependency on s 1 for clearness) is a bounded random variable, and more specifically, it’s
Revised: Since V Mπ, 1 (here, we drop the dependency on s 1 for clearness) is a bounded random variable, and more specifically, it’s

Original: (( HR max ) / 2) -sub-Gaussian.
Revised: (( HR max ) / 2) -sub-Gaussian.

Original: Hence, by Lemma 4, the following holds,
Revised: Hence, by Lemma 4, the following holds,

Original: [Equation]
Revised: [Equation]

Original: This gives us that
Revised: This gives us that

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: Then, we can further rewrite the KL-divergence using the conditional mutual information I k ( · ; · ) (i.e., conditioning on the history H k ),
Revised: Then, we can further rewrite the KL-divergence using the conditional mutual information I k ( · ; · ) (i.e., conditioning on the history H k ),

Original: [Equation]
Revised: [Equation]

Original: When conditioning on the history H k
Revised: When conditioning on the history H k

Original: [Equation]
Revised: [Equation]

Original: the optimal policy π ⋆ and M is independent of the π k , hence we have
Revised: the optimal policy π ⋆ and M is independent of the π k , hence we have

Original: [Equation]
Revised: [Equation]

Original: By the fact that π ⋆ is jointly independent of V Mπ k , 1 and π k when conditioning on the history H k
Revised: By the fact that π ⋆ is jointly independent of V Mπ k , 1 and π k when conditioning on the history H k

Original: [Equation]
Revised: [Equation]

Original: hence we have
Revised: hence we have

Original: [Equation]
Revised: [Equation]

Original: By the chain rule of mutual information
Revised: By the chain rule of mutual information

Original: [Equation]
Revised: [Equation]

Original: we finally get
Revised: we finally get

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: Now
Revised: Now

Original: [Equation]
Revised: [Equation]

Original: let’s define the following function g k and C ,
Revised: let’s define the following function g k and C ,

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: Thus
Revised: Thus

Original: [Equation]
Revised: [Equation]

Original: we further have
Revised: we further have

Original: [Equation]
Revised: [Equation]

Original: On the other hand
Revised: On the other hand

Original: [Equation]
Revised: [Equation]

Original: we can rewrite G k as
Revised: we can rewrite G k as

Original: [Equation]
Revised: [Equation]

Original: By rearranging the terms
Revised: By rearranging the terms

Original: [Equation]
Revised: [Equation]

Original: we get
Revised: we get

Original: [Equation]
Revised: [Equation]

Original: Hence,
Revised: Hence,

Original: [Equation]
Revised: [Equation]

Original: Hence, we have
Revised: Hence, we have

Original: [Equation]
Revised: [Equation]

Revised: where the last inequality holds because of the chain rule of mutual information

Revised: [Equation]

Original: Remark 

Revised: Remark

Original: When the number of policies | Π | is finite, and the value function V Mπ, 1 is linear with its parameter lives in R d , then C can be upper bounded by d , i [Equation] e., C ≤ d .
Revised: When the number of policies | Π | is finite, and the value function V Mπ, 1 is linear with its parameter lives in R d , then C can be upper bounded by d , i.e., C ≤ d .

Original: Lemma 4 (Russo & Van Roy (2016))
Revised: Lemma 4 (Russo & Van Roy (2016))

Original: Suppose that there is a H k -measurable random variable η , such that for each π ∈ Π , V Mπ, 1 is a η -sub-Gaussian random variable when conditioned on H k , then for every π, π ′ ∈ Π , the following holds with probability 1 ,
Revised: Suppose that there is a H k -measurable random variable η , such that for each π ∈ Π , V Mπ, 1 is a η -sub-Gaussian random variable when conditioned on H k , then for every π, π ′ ∈ Π , the following holds with probability 1 ,

Original: [Equation]
Revised: [Equation]