Title: How to Fine-tune the Model: Unified Model Shift and Model Bias Policy Optimization
Venue: NeurIPS 2023
Accepted: yes
Keywords: model-based reinforcement learning

model shift

model bias

fine-tuning

performance difference bound
Abstract: Designing and deriving effective model-based reinforcement learning (MBRL) algorithms with a performance improvement guarantee is challenging, mainly attributed to the high coupling between model learning and policy optimization. Many prior methods that rely on return discrepancy to guide model learning ignore the impacts of model shift, which can lead to performance deterioration due to excessive model updates. Other methods use performance difference bound to explicitly consider model shift. However, these methods rely on a fixed threshold to constrain model shift, resulting in a heavy dependence on the threshold and a lack of adaptability during the training process. In this paper, we theoretically derive an optimization objective that can unify model shift and model bias and then formulate a fine-tuning process. This process adaptively adjusts the model updates to get a performance improvement guarantee while avoiding model overfitting. Based on these, we develop a straightforward algorithm USB-PO (Unified model Shift and model Bias Policy Optimization). Empirical results show that USB-PO achieves state-of-the-art performance on several challenging benchmark tasks.

How to Fine-tune the Model: Unified Model Shift and Model Bias Policy Optimization

Abstract

Designing and deriving effective model-based reinforcement learning (MBRL) algorithms with a performance improvement guarantee is challenging, mainly attributed to the high coupling between model learning and policy optimization. Many prior methods that rely on return discrepancy to guide model learning ignore the impacts of model shift, which can lead to performance deterioration due to excessive model updates. Other methods use performance difference bound to explicitly consider model shift. However, these methods rely on a fixed threshold to constrain model shift, resulting in a heavy dependence on the threshold and a lack of adaptability during the training process. In this paper, we theoretically derive an optimization objective that can unify model shift and model bias and then formulate a fine-tuning process. This process adaptively adjusts the model updates to get a performance improvement guarantee while avoiding model overfitting. Based on these, we develop a straightforward algorithm USB-PO 2 (Unified model Shift and model Bias Policy Optimization). Empirical results show that USB-PO achieves state-of-the-art performance on several challenging benchmark tasks.

1 Introduction

Nowadays, reinforcement learning (RL) has been gaining much traction in a wide variety of complicated decision-making tasks ranging from academia to industry [42, 9, 5, 13, 28]. Part of this is due to some remarkable model-free RL (MFRL) algorithms [40, 15, 11, 41, 18], which show desirable asymptotic performance. However, their applications are hindered by the bottleneck of sample efficiency. On the contrary, model-based RL (MBRL) algorithms, using a world model to generate the imaginary rollouts and then taking them for policy optimization [31, 20], have high sample efficiency while achieving similar asymptotic performance, thus becoming a compelling alternative in practical cases [39, 16, 52].

Typically, MBRL algorithms iterate between model learning and policy optimization. Hence the model quality is crucial for MBRL. Many prior methods [31, 20, 53, 38, 24] rely on return discrepancy to obtain model updates with a performance improvement guarantee. While having achieved comparable results, they only account for model bias in one iteration [20] but do not consider the impacts of model shift between two iterations [21], which can lead to performance deterioration due to excessive model updates. Although CMLO [21] explicitly considers model shift from the perspective of the performance difference bound, it only sets a fixed threshold to constrain the impacts of model shift and determines when the model should be updated accordingly. If this threshold is

set too low, the model bias of the following iteration will be large, which impairs the subsequent optimization process. If this threshold is set too high, the performance improvement can no longer be guaranteed. We remark that such an update is heavily dependent on the choice of this threshold and should be adjusted adaptively during the training process. Therefore, a smarter scheme is required to unify model shift and model bias, enabling adaptively adjusting their impacts to get a performance improvement guarantee.

In this paper, we theoretically derive an optimization objective that can unify model shift and model bias. Specifically, according to the performance difference bound, we propose to minimize the sum of two terms: the model shift term between the pre-update model and the post-update model and the model bias term of the post-update model, each of which is denoted by a second-order Wasserstein distance [47]. By minimizing this optimization objective after the model update via maximum likelihood estimation (MLE) [6, 20], we can tune the model to adaptively find appropriate updates to get a performance improvement guarantee.

Based on these, we develop a straightforward algorithm USB-PO (Unified model Shift and model Bias Policy Optimization). To the best of our knowledge, this is the first method that unifies model shift and model bias and adaptively fine-tunes the model updates during the training process. We evaluate USB-PO on several continuous control benchmark tasks. The results show that USB-PO has higher sample efficiency and better final performance than other state-of-the-art (SOTA) MBRL methods and yields promising asymptotic performance compared with the MFRL counterparts.

2 Related works

MBRL algorithms are promising candidates for real-world sequential decision-making problems due to their high sample efficiency. Existing studies can be divided into several categories [4, 45, 1, 34, 49, 8, 17] following their different usage of the model. Our work falls into the Dyna-style category [45, 44]. Specifically, after model learning, the model generates the imaginary rollouts into the replay buffer for subsequent policy optimization. Hence, both model learning and policy optimization have critical impacts on asymptotic performance.

Some previous algorithms focus on the policy optimization process. CMBAC [51] introduces conservatism to reduce overestimation of the action value function, and ME-TRPO [23] imposes constraints on the policy to get reliable updates within the trust region. Our work is oriented towards model learning, which is orthogonal to these algorithms. Hence, we are inclined to propose a generic algorithm similar to [24, 20, 21] that can be plugged into many SOTA MFRL algorithms [15, 29], rather than just proposing for a specific policy optimization algorithm.

A key issue in model learning is model bias, which refers to the error between the model and the real environment [20]. As the imaginary rollout horizon increases, the impacts of model bias accumulate rapidly, leading to compounding error and unreliable transitions. To mitigate this problem, several effective methods have been proposed. The ensemble model technique [6, 25, 37] and the dropout method [12] are employed to prevent model overfitting. The uncertainty estimation techniques are used to adjust the rollout length [20, 30] or the transition weight [19, 38]. Furthermore, The multi-step techniques [2, 48] are applied to prevent direct input of the imaginary states. We follow the previous work [20] to use the combination of the ensemble model technique with short model rollouts to mitigate the compounding error.

Performance improvement guarantee is a core concern in both MFRL and MBRL theoretic avenues. In MFRL, methods such as TRPO [41] and CPI [22] choose to optimize the performance difference bound, whilst most of the previous work in MBRL [31, 20, 53, 38, 24] choose to optimize the difference of expected return under the model and that of the real environment, which is termed return discrepancy. However, return discrepancy ignores model shift between two consecutive iterations compared to the performance difference bound under the MBRL setting, which can lead to performance deterioration due to excessive model updates. Although some recent methods have also employed performance difference bound to construct theoretical proofs, they still suffer from certain limitations. OPC [10] designs an algorithm to optimize on-policy model error, but it is similar to return discrepancy in nature. DPI [43] uses dual updates to improve sample efficiency but tries to restrict policy updates within the trust region, thus inhibiting exploration. CMLO [21] relies on a fixed threshold to constrain the impacts of model shift, resulting in a heavy dependence on the threshold and a lack of adaptability during the training process. Hence, we try to unify model shift

and model bias to form a novel optimization problem, adaptively fine-tuning the model updates to get a performance improvement guarantee. Still, some prior work [7, 36, 54] choose to consider regret bound, among which [54] also reduce the impacts of the model changing dramatically between successive iterations. Instead of unifying model shift and model bias, they choose to realize dual optimization by considering maximizing the expectation of the model value rather than that of the single model as a sub-process. Different from [54, 43, 27] that use dual optimization to train the policy, we devise an extra phase to fine-tune the model.

3 Preliminaries

We consider a Markov Decision Process (MDP), defined by the tuple M = ( S , A , p, r, γ, ρ 0 ) . S and A denote the state space and action space respectively, and γ ∈ (0 , 1) denotes the discount factor. p ( s ′ | s, a ) denotes the dynamic transition distribution and we denote p M ∗ ( s ′ | s, a ) as that of the real environment. ρ 0 ( s ) denotes the initial state distribution and r ( s, a ) denotes the reward function. RL aims to find the optimal policy π ∗ that maximizes the expected return under the real environment M ∗ denoted by the value function V π M ∗ as Eq.(1):

MBRL algorithms aim to learn the dynamic transition distribution model, p M ( s ′ | s, a ) , by using samples collected from interaction with the real environment via supervised learning. We denote the expected return under the model M of the policy π as V π M and denote that under the real environment of the policy π derived from the model M as V π | M . Additionally, we assume that r ( s, a ) is unknown to the model M and the model will predict r M as the reward function. Besides, we denote M as a parameterized family of models and Π as a parameterized family of policies.

Let d π M ( s, a ) denote the normalized discounted visitation probability for ( s, a ) when starting at s 0 ∼ ρ 0 and following π under the model M . Let p π M,t ( s ) denote the probability of visiting s at timestep t given the policy π and the model M .

We define the total variation distance (TVD) estimator as D TV ( ·||· ) and the second-order Wasserstein distance estimator as W 2 ( · , · ) .

4 USB-PO framework

In this section, we demonstrate a detailed description of our proposed algorithmic framework. i.e., USB-PO. In Section 4.1, a meta-algorithm of the USB-PO framework is provided as a generic solution. In Section 4.2, we theoretically show how to unify model shift and model bias to get a performance improvement guarantee 3 . In Section 4.3, the practical algorithm is proposed to instantiate the USB-PO framework.

4.1 The overall algorithm

The general algorithmic framework of USB-PO is depicted in Algorithm 1, where the main difference compared to the existing MBRL algorithms is the two-phase model learning process, namely phase 1 and phase 2. Phase 1 uses traditional MLE loss to train the model, which may impair the performance by excessive model updates due to only considering the impacts of model bias. To mitigate this problem, we introduce phase 2 to further fine-tune the model updates, whose optimization objective is defined as Eq.(3).

Algorithm 1 Meta-Algorithm of the USB-PO Framework

1: Initialize the policy π and the learned model;

2: Initialize the environment replay buffer D and the model replay buffer D M ;

3: for each epoch do

4: Use π to interact with the real environment: D ← D ∪ { ( s, a, r, s ′ ) } ;

5: Backup the current learned model for future use and denote this backed-up model as p M 1 ;

6: Phase 1: use D to train the learned model with the supervision of MLE and denote this updated model as p M 2 ;

7: Phase 2: use Eq.(3) as optimization objective to further fine-tune p M 2 ;

9: Use D∪D M to train the policy π

8: Use p M 2 to generate the imaginary rollouts: D M ←D M ∪ { ( s M , a M , r M , s ′ M ) } ;

10: end for

Eq.(3) unifies the model shift term and the model bias term in the second-order Wasserstein distance form, namely W 2 ( p M 1 , p M 2 ) and W 2 ( p M 2 , p M ∗ ) , thus achieving adaptive adjustment of their impacts during the fine-tuning process. As demonstrated in Section 4.2 and Section 5.4, this is not equivalent to the traditional methods of limiting the magnitude of model updates, but rather beneficial to get a performance improvement guarantee.

4.2 Theoretical proof

Definition 1 (Performance Difference Bound) . Recalling that V π i | M i denotes the expected return under the real environment of the policy π i ∈ Π derived from the model M i ∈ M . The lower bound on the true return gap of π 1 and π 2 can be stated as,

By constantly increasing the value of C , the lower bound on the performance difference is guaranteed to be lifted, leading to performance improvement. Therefore, we try to take model shift and model bias into the formulation of C and maximize it to achieve our goal.

Theorem 1 (Performance Difference Bound Decomposition) . Let M i ∈ M be the evaluated model and π i ∈ Π be the policy derived from the model. The performance difference bound can be decomposed into three terms,

Obviously, compared to directly optimizing the return discrepancy of each iteration [20], the performance difference bound chooses to optimize the return discrepancy of two adjacent iterations, namely V π 2 | M 2 -V π 2 M 2 and V π 1 | M 1 -V π 1 M 1 respectively, and the expected return variation between these two iterations, namely V π 2 M 2 -V π 1 M 1 , demonstrating better rigorousness. For further discussion, we introduce the following theorem.

Theorem 2 (Return Bound) . Let R max denote the bound of the reward function, ϵ π denote max s D TV ( π 1 || π 2 ) and ϵ M 2 M 1 denote E ( s,a ) ∼ d π 1 M 1 [ D TV ( p M 1 || p M 2 )] . For two arbitrary policies π 1 , π 2 ∈ Π , the expected return under two arbitrary models M 1 , M 2 ∈ M can be bounded as,

By using Eq.(6), we can easily bound the decomposition terms of the performance difference bound in Eq.(5).

Theorem 3 (Decomposition TVD Bound) . Let ϵ π i M i denote E ( s,a ) ∼ d π i Mi [ D TV ( p M i || p M ∗ )] . Let M i ∈ M be the evaluated model and π i ∈ Π be the policy derived from the model. The decomposition terms can be bounded as,

Notably, the model shift term and the model bias term in TVD form, namely D TV ( p M 1 || p M 2 ) and D TV ( p M 2 || p M ∗ ) , are already present in the right-hand side of the Eq.(7). However, due to the unknown term π 2 , we can not sample from d π 2 M 2 . Thus, we need to make a further transformation about the Eq.(7).

Theorem 4 (Unified Model Shift and Model Bias Bound) . Let κ denote the constant 2 R max (1 -γ ) 2 and ∆ denotes E ( s,a ) ∼ d π 1 M 1 [ D TV ( p M 2 || p M ∗ )] -E ( s,a ) ∼ d π 2 M 2 [ D TV ( p M 2 || p M ∗ )] . Let M i ∈ M be the evaluated model and π i ∈ Π be the policy derived from the model. The unified model shift and model bias bound can be derived as,

The ∆ term in Eq.(8) is still intractable. However, the fact that it covers d π 1 M 1 and d π 2 M 2 reminds us that we may explore its relationship with the three TVD terms lying in the expectation of d π 1 M 1 .

Theorem 5 (| ∆ | Upper Bound) . Let M i ∈ M be the evaluated model and π i ∈ Π be the policy derived from the model. The term ∆ can be upper bounded as:

Under the online setting, it is assumed that the error caused by policy shift, namely ϵ π , compared to model bias has a relatively small scale [20, 53]. Therefore, we ignore the minor influence brought by the policy. Additionally, the terms lying in the expectation of d π 1 M 1 in Eq.(9) is the product of the model shift term and the model bias term in TVD form, both of which have the range [0 , 1] [26], making the upper bound of | ∆ | become a higher-order term compared to each of them alone. Thus, the critical element to lift the lower bound in Eq.(8) is the terms lying in the expectation of d π 1 M 1 . By maximizing these terms, namely minimizing D TV ( p M 1 || p M 2 ) + D TV ( p M 2 || p M ∗ ) 4 , we can achieve unifying model shift and model bias and adaptively fine-tuning the model updates to get a performance improvement guarantee.

To make it practically feasible, we make additional assumptions. Specifically, let ( X , Σ) be a measurable space and F is a set of functions mapping X to R that contains V π M . When V π M is L v -Lipschitz with respect to a norm || · || , the integral probability metric [33] of two arbitrary dynamic transition function M,M ′ ∈ M defined on X is as Eq.(10). Notice that if 1 ≤ p ≤ q , W p ( p M , p M ′ ) ≤ W q ( p M , p M ′ ) [32], and we can get W 1 ( p M , p M ′ ) ≤ W 2 ( p M , p M ′ ) . Hence, Eq.(3) can be applied as the optimization objective to get a performance improvement guarantee.

4.3 Practical implementation

We now instantiate Algorithm 1 by demonstrating an explicit approach. To better clarify, we would like to state the four design decisions in detail: (1) how to parametrize the model, (2) how to estimate the model bias term and model shift term in phase 2, (3) how to use the model to generate rollouts and (4) how to use the model rollouts to optimize the policy π .

Predictive Model. We use a bootstrap ensemble of dynamic models { p 1 M , ..., p B M } , whose elements are all probabilistic neural network, outputting the Gaussian distribution with diagonal covariance: p i M ( s t +1 , r t +1 | s t , a t ) = N ( µ i M ( s t , a t ) , Σ i M ( s t , a t )) . The probabilistic ensemble model can capture the aleatoric uncertainty arising from inherent stochasticities and the epistemic uncertainty corresponding to ambiguously determining the underlying system due to the lack of sufficient data [6]. In the field of MBRL, properly handling these uncertainties can achieve better asymptotical performance. Following the previous work [20], we select a model uniformly from the elite set to generate the transition at each step in the rollout process. In phase 1, the ensemble models are trained on shared but differently shuffled data, where the optimization objective is MLE [6, 20].

Model Shift and Model Bias Estimation. Recalling that the model shift term refers to W 2 ( p M 1 || p M 2 ) and the model bias term refers to W 2 ( p M 2 || p M ∗ ) . For arbitrary two Gaussian distributions p ∼ N ( µ 1 , Σ 1 ) and q ∼ N ( µ 2 , Σ 2 ) , the second-order Wasserstein distance between them is derived as Eq.(11) [35].

Let k 1 denote the index of the selected model from M 1 ensemble and k 2 denote that from M 2 ensemble. For the model shift term, we approximate it as W ( p k 1 , p k 2 )

̸

Model Rollout. Following the previous work [20], we use short branch rollouts to alleviate the impacts of compounding error.

2 M 1 M 2 . For the model bias term, ∑ B

Policy Optimization. Our work allows being plugged in many SOTA MFRL algorithms, e.g. SAC [15], TD3 [11], etc. Here we employ SAC as an example.

5 Experiment

Our experiments are designed to investigate four primary questions: (1) Whether we need to propose an algorithm similar to USB-PO to adaptively adjust the impacts of model shift? (2) How well does USB-PO perform on reinforcement learning benchmark tasks compared to SOTA MBRL and MFRL algorithms? (3) How to understand USB-PO? (4) Does USBPO have a similar performance on different learning rate settings in phase 2?

5.1 Necessity of USB-PO

Werecall that CMLO [21] sets a fixed threshold to constrain the impacts of model shift, which we argue CMLO is heavily dependent on this threshold. To valid our statement, we devise an experiment that sets three different thresholds for CMLO on Walker2d environment in MuJoCo [46]. As Figure 1 shows, the performance corresponding to the other two thresholds (1.0

-′-K

-glyph[triangleinv]-′-K

-∞-′-′-K

-∞-glyph[triangleinv]-′-K

-∈-′-′-K -∈-glyph[triangleinv]-′-K-glyph[owner]-′-′-K

-′

-∈-′-′-′

-glyph[triangle]-′-′-′

-W-glyph[turnstileright]-glyph[arrowbothv]-‖-glyph[ceilingright]-∇-∈-glyph[ceilingleft]-↖-glyph[subsetsqequal]-∈

-glyph[triangleinv]-↙-′

-∞-↙-′

-∫-glyph[unionsq]-glyph[ceilingright]-√-∫

and 5.0) is severely affected. Therefore, setting a fixed threshold to constrain is inappropriate, and a smarter way like USB-PO should be applied to adaptively adjust the impacts of model shift.

5.2 Comparison with baselines

To highlight our algorithm, we compare USB-PO against some SOTA MBRL and MFRL algorithms. The MBRL baselines include CMLO [21], which carefully chooses the threshold for different environments to constrain the impacts of model shift, ALM [14], which learns the representation, model, and policy into one objective, PDML [50], which uses historical policy sequence to aim the model training process, MBPO [20], which is the variant without our proposed model fine-tuning process and uses return discrepancy to get a performance improvement guarantee. The MFRL baselines include SAC [15], the SOTA MFRL algorithm in terms of asymptotic performance and PPO [40], which explores monotonic improvement under the model-free setting.

We evaluate USB-PO and these baselines on six MuJoCo [46] continuous control tasks in OpenAI Gym [3], covering Humanoid, Walker2d, Ant, HalfCheetah, Hopper and Inverted-Pendulum, with more details showing in the appendix. To be fair, we employ the standard 1000-step version of these tasks with the same environment settings.

Figure 2 shows the learning curves of all compared methods, together with the asymptotic performance. The results show that our algorithm is remarkably advanced over the MFRL algorithms

-glyph[turnstileright]-glyph[subsetsqequal]-glyph[ceilingright]-∇-glyph[turnstileright]-}-glyph[ceilingright]-←-∇-glyph[ceilingright]-glyph[unionsq]-glyph[intersectionsq]-∇-\

HalfCheetah-v2

Humanoid-v2

InvertedPendulum-v2

Hopper-v2

USB-PO(ours)

ALM

CMLO

SAC

PPO

PDML

MBPO

convergence

6000

6000

average return

12000

averagereturn

averagereturn

4000

4000

8000

4000

2000

2000

0

0

0

OK

50K

100K150K

200K

250K

OK

50K100K 150K 200K 250K 300K

OK

50K100K 150K 200K 250K 300K

steps

steps

steps

6000

900

average return

averagereturn

3000

4000

600

2000

2000

300

1000

0

0

0

OK

50K

100K150K

steps

200K

250K

OK

2K

5K

steps

7K

10K

12K

15K

OK

20K

40K

steps

60K

80K100K120K

regarding sample efficiency, along with asymptotic performance on par with the SOTA MFRL algorithm SAC. Compared to the MBRL baselines, our method achieves higher sample efficiency and better final performance. Notably, compared to CMLO, which requires a finely chosen threshold for each environment to constrain the impacts of model shift, our method utilizes the same learning rate for the fine-tuning process in all environments. This further validates unifying the model shift and the model bias to adjust the model updates adaptively is reasonable.

Computational Cost. We report our computational cost compared to MBPO [20] in Appendix D.4. Although USB-PO is a two-phase model training process, continuing to use the fine-tuned model for the next iteration has the potential to accelerate model convergence and then possibly reduce the training time.

5.3 How to understand USB-PO

In this section, we first design an experiment to illustrate the value magnitude of ∆ , the model shift term and the model bias term, validating whether using the second-order Wasserstein distance satisfies the prerequisites for getting a performance improvement guarantee, and then use more in-depth experiments to illustrate how USB-PO works and show its superiority.

Value Magnitude. We choose 2 challenging tasks in MuJoCo, HalfCheetah and Walker2d, and plot the value of ∆ , the model shift term and the model bias term during the training process respectively. To approximate the ∆ practically, we use the samples from the model replay

HalfCheetah-v2

Δ

model shift

model bias

0.12

0.08

value

0.08

value

0.04

0.04

0.00

0.00

0

200

400

600

0

200

400

600

training times

trainingtimes

buffer before the update of the model and the policy to approximate sampling from d π 1 M 1 and use the samples generated by them after update to approximate sampling from d π 2 M 2 .

As shown in Figure 3, the magnitude of ∆ generally oscillates in a slight manner around 0 while the magnitude of model shift term and the model bias term are much larger than ∆ even in the converged cases. Therefore, it is reasonable to ignore the impacts of ∆ in Eq.(8). This further supports that minimizing Eq.(3) can get a performance improvement guarantee.

Working Mechanism. To illustrate the significance of our method, we calculate the optimization objective value before and after the fine-tuning process and plot their difference, as shown in Figure 4. All random seeds show the consistent result that at the beginning of training the sum of model bias and model shift is large and thus the fine-tuning magnitude is relatively large. As the model is trained to converge, the fine-tuning process also gradually converges to ensure the stability of the model. Notice that in some cases the fine-tuning process does not choose to fine-tune the model updates actually (0 refers to no update, which means the updates in phase 1 are reasonable), which further supports our theorem and the motivation of this paper, i.e., actively adaptive adjustment of the model updates to get a performance improvement guarantee, rather than passively waiting for model shift to surpass a given threshold to make a substantial update [21]. Based on the performance comparison in Figure 2, we argue that actively adaptive updates are more beneficial.

2.0

Seed

Seed

of optimization objective

0.006

+

Seed3

Seed2

average prediction error

1.5

Seed2

Seed3

Seed5

Seed5

0.004

1.0

0.002

0.5

0.000

0

200

400

600

800

0.0

0

200

400

600

800

(a) training times

(b)training times

0.006

differenceofmodelbias

0.012

Seed

differenceofmodelshift

error

Seed2

differenceofoptimizationobjective

faverage prediction

0.008

Seed3

of values

0.004

Seed5

0.004

0.002

60.000

0.000

0

10

20

30

0

200

400

600

800

(a) training times

(b) training times

Wefurther plot the difference of the model shift term, the model bias term and the average prediction error before and after the fine-tuning process. As shown in Figure 5 (a), when the fine-tuning actually operates, the difference of the model shift term and the model bias term are both positive. As shown in Figure 5 (b), the fine-tuning process has a positive effect on the reduction of the average prediction error. These observations suggest that USB-PO has other positive effects during the fine-tuning process, i.e. potentially reducing both model shift and model bias and leading to model overfitting avoidance.

average return

6000

4000

2000

0

0K

50K

100K

150K

steps

0K

50K

100K

150K

200K

250K

300K

steps

0

2000

4000

average return

Algorithm

USB-PO

7e-4

5e-4

3e-4

1e-4

1e-3 (original)

Additionally, the change in the adjustment magnitude of the average prediction error indicates that the model can converge under the role of the fine-tuning process, which is consistent with Figure 4, forming bi-verification.

To conclude, when USB-PO actively recognizes that an improper update happens, it performs a pullback operation (the model fine-tuning process) to secure the performance improvement guarantee while avoiding model overfitting. In contrast, when USB-PO considers that a reasonable update is done, no actual operation will be taken. According to the performance comparison results in Figure 2, it is verified that the model quality plays a determinant role in the Dyna-Q [44] algorithms.

HalfCheetah-v2

learning rate 1e-4

learning rate 3e-4

learning rate 5e-4

MBPO maximum performance

6000

average return

average return

12000

4000

8000

2000

4000

1e-4

3e-4

5e-4

1e-4

3e-4

5e-4

learningrate

learningrate

5.4 Ablation study

In this section, we design 3 ablation experiments to strengthen our superiority.

Optimization Objective Variants. We set three variants for our optimization objective, covering (1) without the model shift term and the model bias term, which is equal to MBPO [20], (2) with the model shift term, (3) with the model bias term. As shown in Figure 6, only optimizing the model shift term results in a drop in sample efficiency since the fine-tuning process makes M 2 tend to update towards M 1 and only optimizing the model bias term leads to performance deterioration due to excessive model updates.

Not Equivalent to Limiting the Update Magnitude. We set different learning rates for MBPO to compare with USB-PO. As Figure 7 shows, USB-PO is more dominant from the perspective of the average return, which further strengthens that USB-PO is not equivalent to limiting the update magnitude, but rather beneficial to get a performance improvement guarantee.

200K

Algorithm with none

with model shift with model bias

with both

250K

300K

Learning Rate Performance Comparison. We set up an ablation experiment on the learning rate of the fine-tuning process. As shown in Figure 8, unlike CMLO which is strongly dependent on a carefully chosen threshold for each environment to constrain the impacts of model shift, USB-PO is less sensitive to the learning rate of phase 2.

6 Discussion

In this paper, we propose a novel MBRL algorithm called USB-PO, which can unify model shift and model bias, enabling adaptive adjustment of their impacts to get a performance improvement guarantee. We further find that our method can potentially reduce both model shift and model bias, leading to model overfitting avoidance. Empirical results on several challenging benchmark tasks validate the superiority of our algorithm and more in-depth experiments are conducted to demonstrate our mechanism. The limitation of our work is we do not investigate the change of model shift and model bias in this optimization problem theoretically. Therefore, one direction that merits further research is how to solidly interpret the variation of model shift and model bias in this fine-tuning process. Further, we aim to propose a general algorithmic framework in this paper. Due to the simple design of our model architecture, we find it difficult to capture some complex and practical locomotion. We hope to combine our adaptive adjusting technique with more advanced model architectures in the future.

Acknowledgments and Disclosure of Funding

This work is supported by the National Key Research and Development Program of China (No. 2021YFB2501104 No. 2020YFA0711402). We thank Michael Janner, Raj Ghugare, Zifan Wu, and Xiyao Wang for providing the baseline results and thank Tianying Ji for the worthy discussions.

References

Brandon Amos et al. 'On the model-based stochastic value gradient for continuous reinforcement learning'. In: Learning for Dynamics and Control . PMLR. 2021, pp. 6-20.

Kavosh Asadi et al. 'Combating the compounding-error problem with a multi-step model'. In: arXiv preprint arXiv:1905.13320 (2019).

Greg Brockman et al. 'Openai gym'. In: arXiv preprint arXiv:1606.01540 (2016).

Jacob Buckman et al. 'Sample-efficient reinforcement learning with stochastic ensemble value expansion'. In: Advances in neural information processing systems 31 (2018).

Annie Chen et al. 'You Only Live Once: Single-Life Reinforcement Learning'. In: Advances in Neural Information Processing Systems 35 (2022), pp. 14784-14797.

Kurtland Chua et al. 'Deep reinforcement learning in a handful of trials using probabilistic dynamics models'. In: Advances in neural information processing systems 31 (2018).

Sebastian Curi, Felix Berkenkamp, and Andreas Krause. 'Efficient model-based reinforcement learning through optimistic policy search and planning'. In: Advances in Neural Information Processing Systems 33 (2020), pp. 14156-14170.

Vladimir Feinberg et al. 'Model-based value estimation for efficient model-free reinforcement learning'. In: arXiv preprint arXiv:1803.00101 (2018).

Shuo Feng et al. 'Dense reinforcement learning for safety validation of autonomous vehicles'. In: Nature 615.7953 (2023), pp. 620-627.

Lukas Froehlich et al. 'On-Policy Model Errors in Reinforcement Learning'. In: International Conference on Learning Representations . 2022. URL: https://openreview. net/forum?id=81e1aeOt-sd .

Scott Fujimoto, Herke Hoof, and David Meger. 'Addressing function approximation error in actor-critic methods'. In: International conference on machine learning . PMLR. 2018, pp. 1587-1596.

Yarin Gal, Jiri Hron, and Alex Kendall. 'Concrete dropout'. In: Advances in neural information processing systems 30 (2017).

Chengqian Gao et al. 'Value penalized Q-learning for recommender systems'. In: Proceedings of the 45th International ACM SIGIR Conference on Research and Development in Information Retrieval . 2022, pp. 2008-2012.

Raj Ghugare et al. 'Simplifying Model-based RL: Learning Representations, Latent-space Models, and Policies with One Objective'. In: The Eleventh International Conference on Learning Representations . 2023. URL: https://openreview.net/forum?id= MQcmfgRxf7a .

Tuomas Haarnoja et al. 'Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor'. In: International conference on machine learning . PMLR. 2018, pp. 1861-1870.

Danijar Hafner et al. 'Mastering Atari with Discrete World Models'. In: International Conference on Learning Representations . 2021. URL: https://openreview.net/forum? id=0oabwyZbOu .

Nicolas Heess et al. 'Learning continuous control policies by stochastic value gradients'. In: Advances in neural information processing systems 28 (2015).

Matteo Hessel et al. 'Rainbow: Combining improvements in deep reinforcement learning'. In: Proceedings of the AAAI conference on artificial intelligence . Vol. 32. 1. 2018.

Wenzhen Huang et al. 'Learning to Reweight Imaginary Transitions for Model-Based Reinforcement Learning'. In: Proceedings of the AAAI Conference on Artificial Intelligence . Vol. 35. 9. 2021, pp. 7848-7856.

Michael Janner et al. 'When to trust your model: Model-based policy optimization'. In: Advances in neural information processing systems 32 (2019).

Tianying Ji et al. 'When to Update Your Model: Constrained Model-based Reinforcement Learning'. In: Advances in Neural Information Processing Systems . Ed. by Alice H. Oh et al. 2022. URL: https://openreview.net/forum?id=9a1oV7UunyP .

Sham Kakade and John Langford. 'Approximately optimal approximate reinforcement learning'. In: Proceedings of the Nineteenth International Conference on Machine Learning . 2002, pp. 267-274.

Thanard Kurutach et al. 'Model-ensemble trust-region policy optimization'. In: arXiv preprint arXiv:1802.10592 (2018).

Hang Lai et al. 'Bidirectional model-based policy optimization'. In: International Conference on Machine Learning . PMLR. 2020, pp. 5618-5627.

Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. 'Simple and scalable predictive uncertainty estimation using deep ensembles'. In: Advances in neural information processing systems 30 (2017).

David A Levin and Yuval Peres. Markov chains and mixing times . Vol. 107. American Mathematical Soc., 2017.

Sergey Levine and Pieter Abbeel. 'Learning neural network policies with guided policy search under unknown dynamics'. In: Advances in neural information processing systems 27 (2014).

Lanqing Li, Rui Yang, and Dijun Luo. 'FOCAL: Efficient Fully-Offline Meta-Reinforcement Learning via Distance Metric Learning and Behavior Regularization'. In: International Conference on Learning Representations . 2021. URL: https://openreview.net/forum? id=8cpHIfgY4Dj .

Timothy P Lillicrap et al. 'Continuous control with deep reinforcement learning'. In: arXiv preprint arXiv:1509.02971 (2015).

Cong Lu et al. 'Revisiting Design Choices in Offline Model Based Reinforcement Learning'. In: International Conference on Learning Representations . 2022. URL: https:// openreview.net/forum?id=zz9hXVhf40 .

Yuping Luo et al. 'Algorithmic framework for model-based deep reinforcement learning with theoretical guarantees'. In: arXiv preprint arXiv:1807.03858 (2018).

Ester Mariucci and Markus ReiSS. 'Wasserstein and total variation distance between marginals of Lévy processes'. In: (2018).

Alfred Müller. 'Integral probability metrics and their generating classes of functions'. In: Advances in applied probability 29.2 (1997), pp. 429-443.

Anusha Nagabandi et al. 'Neural network dynamics for model-based deep reinforcement learning with model-free fine-tuning'. In: 2018 IEEE international conference on robotics and automation (ICRA) . IEEE. 2018, pp. 7559-7566.

Ingram Olkin and Friedrich Pukelsheim. 'The distance between two random vectors with given dispersion matrices'. In: Linear Algebra and its Applications 48 (1982), pp. 257-263.

Ian Osband and Benjamin Van Roy. 'Model-based reinforcement learning and the eluder dimension'. In: Advances in Neural Information Processing Systems 27 (2014).

Ian Osband et al. 'Deep exploration via bootstrapped DQN'. In: Advances in neural information processing systems 29 (2016).

Feiyang Pan et al. 'Trust the model when it is confident: Masked model-based actor-critic'. In: Advances in neural information processing systems 33 (2020), pp. 10537-10546.

Julian Schrittwieser et al. 'Online and offline reinforcement learning by planning with a learned model'. In: Advances in Neural Information Processing Systems 34 (2021), pp. 27580-27591.

John Schulman et al. 'Proximal policy optimization algorithms'. In: arXiv preprint arXiv:1707.06347 (2017).

John Schulman et al. 'Trust region policy optimization'. In: International conference on machine learning . PMLR. 2015, pp. 1889-1897.

David Silver et al. 'A general reinforcement learning algorithm that masters chess, shogi, and Go through self-play'. In: Science 362.6419 (2018), pp. 1140-1144.

Wen Sun et al. 'Dual policy iteration'. In: Advances in Neural Information Processing Systems 31 (2018).

Richard S Sutton. 'Dyna, an integrated architecture for learning, planning, and reacting'. In: ACM Sigart Bulletin 2.4 (1991), pp. 160-163.

Richard S Sutton. 'Integrated architectures for learning, planning, and reacting based on approximating dynamic programming'. In: Machine learning proceedings 1990 . Elsevier, 1990, pp. 216-224.

Emanuel Todorov, Tom Erez, and Yuval Tassa. 'Mujoco: A physics engine for model-based control'. In: 2012 IEEE/RSJ international conference on intelligent robots and systems . IEEE. 2012, pp. 5026-5033.

Leonid Nisonovich Vaserstein. 'Markov processes over denumerable products of spaces, describing large systems of automata'. In: Problemy Peredachi Informatsii 5.3 (1969), pp. 6472.

Arun Venkatraman, Martial Hebert, and J Bagnell. 'Improving multi-step prediction of learned time series models'. In: Proceedings of the AAAI Conference on Artificial Intelligence . Vol. 29. 1. 2015.

Tingwu Wang and Jimmy Ba. 'Exploring Model-based Planning with Policy Networks'. In: International Conference on Learning Representations . 2020. URL: https : / / openreview.net/forum?id=H1exf64KwH .

Xiyao Wang et al. 'Live in the moment: Learning dynamics model adapted to evolving policy'. In: International Conference on Machine Learning . PMLR. 2023, pp. 36470-36493.

Zhihai Wang et al. 'Sample-Efficient Reinforcement Learning via Conservative Model-Based Actor-Critic'. In: Proceedings of the AAAI Conference on Artificial Intelligence . Vol. 36. 8. 2022, pp. 8612-8620.

Jingda Wu, Zhiyu Huang, and Chen Lv. 'Uncertainty-aware model-based reinforcement learning: methodology and application in autonomous driving'. In: IEEE Transactions on Intelligent Vehicles (2022).

Zifan Wu et al. 'Plan To Predict: Learning an Uncertainty-Foreseeing Model for Model-Based Reinforcement Learning'. In: Advances in Neural Information Processing Systems 35 (2022), pp. 15849-15861.

Shenao Zhang. 'Conservative Dual Policy Optimization for Efficient Model-Based Reinforcement Learning'. In: Advances in Neural Information Processing Systems 35 (2022), pp. 25450-25463.

7 Appendix

A Proof Sketch

To better clarify our theoretical results, we provide a proof sketch here. Firstly, we decompose the performance difference bound under the model-based setting into three terms (Theorem 1). Secondly, by means of using Return Bound (Theorem 2), we can bound these three terms individually (Theorem 3). Then, we can do some transformation to get Unified Model Shift and Model Bias Bound (Theorem 4), which bounds the model shift term and the model bias term in total variation form. However, due to the intractable property of ∆ , we further explore the upper bound of | ∆ | (Theorem 5), finding that ∆ can be ignored. Finally, by the Integral Probability Metrics (Lemma 3) and the property of the Wasserstein distance, we derive the target which bounds the model shift term and the model bias term in the Wasserstein distance form.

Definition

Performance

Intergral

Difference Bound

ProbabilityMetrics

lipschitzcontinuityofvaluefunction

V=2|M2 - V=1|M1 > C

(Lemma3)

Targetintotalvariationform

Target

Decomposition

DTv(pM llpM2） +DTv(pM2llPM*)

W2(PM,PM2）+W2(PM2，PM*)

ReturnBound

threeterms

(Theorem 2)

(Theorem 1)

ignore the [△l

Unified Model

Decomposition-

Shift and Model

[|UpperBound

individual bound

transformation

BiasBound

|intractable

(Theorem 5)

(Theorem3)

(Theorem 4)

B Useful Lemmas

In this section, we provide some proof to support our theoretical analysis.

Lemma1 (Total variation Distance) . Consider a measurable space (Ω , Σ) and probability measures P and Q are defined on (Ω , Σ) . The total variation distance between P and Q is defined as:

Eq. (12) can be equivalently written as:

Lemma 2 (Total Variation Distance of Joint Distributions) . Given two distributions p ( x, y ) = p ( x ) p ( y | x ) and q ( x, y ) = q ( x ) q ( y | x ) , the total variation distance between them can be bounded as:

Proof:

Lemma 3 (Integral Probability Metrics) . Consider a measurable space ( X , Σ) . The integral probability metric associated with a class F of real-valued functions on X is defined as

where P and Q are probability measures on X . We demonstrate the following special cases:

(a) If = f : f c , then we have

(b) If F is the set of L -Lipschitz function with a norm || · || , then we have

In our paper, to distinguish the dynamic transition function, we choose F to be the class covering V π M . Since the value function can converge to r max 1 -γ , it only needs to satisfy the L v -Lipschitz continuity and thus we can get r max 1 -γ D TV ( p M || p M ′ ) = L v W 1 ( p M , p M ′ ) for any arbitrary model M,M ′ .

C Missing Proof

Theorem 1 (Performance Difference Bound Decomposition) . Let M i ∈ M be the evaluated model and π i ∈ Π be the policy derived from the model. The performance difference bound can be decomposed into three terms,

Proof: Weintroduce two additional terms V π 1 M 1 and V π 2 M 2 that allow the performance difference bound objective to be divided into three operators based on the return bounds, which can be reformulated separately.

Theorem 2 (Return Bound) . Let R max denote the bound of the reward function, ϵ π denote max s D TV ( π 1 || π 2 ) and ϵ M 2 M 1 denote E ( s,a ) ∼ d π 1 M 1 [ D TV ( p M 1 || p M 2 )] . For two arbitrary policies π 1 , π 2 ∈ Π , the expected return under two arbitrary models M 1 , M 2 ∈ M can be bounded as,

Proof: We give the thorough proof referring to Lemma B.4 in MBPO [20] as follows.

According to the Lemma 2, we have:

Further we expand the first term:

Then move the result of Eq.(24) to Eq.(23), we can get:

Next, we move the result of Eq.(25) to Eq.(22), we can get:

Here, we first simplify the second term of the Eq.(26)

Then we bring this result back to Eq.(26) and the proof is complete.

Theorem 3 (Decomposition TVD Bound) . Let ϵ π i M i denote E ( s,a ) ∼ d π i Mi [ D TV ( p M i || p M ∗ )] . Let M i ∈ M be the evaluated model and π i ∈ Π be the policy derived from the model. The decomposition terms can be bounded as,

Proof: According to CMLO [21] and Eq.(21), the term V π 1 | M 1 -V π 1 M 1 can be approximated as -2 R max γ (1 -γ ) 2 ϵ π 1 M 1 , thus we only need to bound the remaining two terms.

For the term V π 2 | M 2 -V π 2 M 2 , we use Eq.(21) to bound it.

Similarly, for the term V π 2 M 2 -V π 1 M 1 , we can get:

We now combine these three bounds together and complete the proof.

Theorem 4 (Unified Model Shift and Model Bias Bound) . Let κ denote the constant 2 R max (1 -γ ) 2 and ∆ denotes E ( s,a ) ∼ d π 1 M 1 [ D TV ( p M 2 || p M ∗ )] -E ( s,a ) ∼ d π 2 M 2 [ D TV ( p M 2 || p M ∗ )] . Let M i ∈ M be the evaluated model and π i ∈ Π be the policy derived from the model. The unified model shift and model bias bound can be derived as,

Proof: Based on Eq.(29), we add a new term κ E ( s,a ) ∼ d π 1 M 1 [ D TV ( p M 2 || p M ∗ )] to reformulate the optimization objective.

Theorem 5 (| ∆ | Upper Bound) . Let M i ∈ M be the evaluated model and π i ∈ Π be the policy derived from the model. The term ∆ can be upper bounded as:

Proof: First, we combine these two terms.

Recalling that we get the result of the sum equation above in Eq.(25), and then we have:

D Experimental Details

D.1 Environment Setup

We evaluate the algorithm over a series of MuJoCo [46] continuous control benchmark tasks. To ensure fairness, we use the standard 1000-step version of all the environments. The details of the environment setup are from OpenAI Gym [3], as shown in Table 1.

D.2 Baseline implementation

MFRL Baselines. We use two state-of-the-art model-free algorithms, i.e. SAC [15] and PPO [40], to do baseline comparison. To demonstrate the final performance and sampling efficiency of our method, we train SAC for 3M steps, which is much more than MBRL algorithms. The hyperparameters are consistent with the author's settings.

MBRLBaselines. We use several state-of-the-art model-based algorithms to do baseline comparison, covering CMLO [21], MBPO [20], SLBO [31] and STEVE [4]. The implementation of CMLO is based on the opensource repo published by the author and all of the hyperparameters are set according to the paper [21]. Our algorithm USB-PO is implemented based on the opensource repo published by Janner who is the author of MBPO.

We present the final performance on six continuous benchmark tasks in Table 2. The results demonstrate that our algorithm achieves competitive performance compared to both MBRL and MFRL baselines over these tasks. Each result in the table shows the average and standard deviation on the maximum average returns among different random seeds and we choose 250K for HalfCheetah-v2, 300K for Walker2d-v2, 300K for Humanoid-v2, 250K for Ant-v2, 15K for Inverted-Pendulum-v2, 120K for Hopper-v2.

D.3 Hyperparameters

Our algorithm USB-PO is based on MBPO [20] and is implemented according to the opensource repo published by the MBPO author. Except for the learning rate in phase 2 of our USB-PO algorithm, the hyperparameters are completely identical to the MBPO settings for all environments. In all benchmark tasks, we set this learning rate to 1e-4.

D.4 Computing Infrastructure

In Table 3, we list our computing infrastructure and the computational time for training USB-PO on these six continuous benchmark tasks. Note that the time we report is the cost for 4 random seeds simultaneously on one graphics card. For Humanoid, only two random seeds can be run simultaneously because of the limitation of graphics memory.

E Comparison with Prior Works

In this section, we compare USB-PO with prior theoretical works to emphasize our contribution, as a complementary to the main paper. First, we give a summary and then show the details as follows. MBPO-Style does not consider model shift and CMLO-Style rely on a fixed threshold to constrain model shift. Our algorithm, USB-PO, adaptively adjusts the model updates in a unified manner (unify model shift and model bias) to get the performance improvement guarantee.

MBPO-Style [20, 41, 24, 53]. They use the return discrepancy bound V π | M ≥ V π M -C ( ϵ m , ϵ π ) to improve the lower bound on the performance under the real environment, i.e. as long as improving V π M by more than C ( ϵ m , ϵ π ) can guarantee improvement on V π | M . Obviously, This scheme is guaranteed under a fixed model and it does not consider the change in model dynamic during updates nor the performance variation concerning model shift. Even worse, if the model has some excessive updates, it is impractical to find a feasible solution to meet the improvement guarantee.

CMLO-Style [21]. They use the performance difference bound under the model-based setting V π 2 | M 2 -V π 1 | M 1 ≥ C to directly consider model shift and model bias. However, they finally derive a constrained lower-bound optimization problem and use a fixed threshold to constrain model shift, i.e. sup s ∈S ,a ∈A D TV ( P M 1 ( ·| s, a ) || P M 2 ( ·| s, a )) ≤ σ M 1 ,M 2 and determine when to update the model accordingly. Notably, we find that this fixed threshold plays a key role in the whole algorithm and needs to be carefully adjusted for each environment. If this threshold is set too low, the model bias of the following iteration will be large, which impairs the subsequent optimization process. If this threshold is set too high, the performance improvement can no longer be guaranteed. Additionally, using a fixed threshold during the whole training process makes the algorithm problematic to adjust adaptively.

USB-PO (Ours). Following CMLO-Style [21], we also use the performance difference bound under the model-based setting to directly consider model shift and model bias. Compared to relying on a fixed threshold to constrain model shift, we use a transformation to unify model shift and model bias into one formulation without the constraint (Theorem 4). Due to the intractable property of ∆ , we further explore the upper bound of | ∆ | , finding that ∆ can be ignored with respect to model shift and model bias alone (Theorem 5). Finally, the optimization objective we get can be used to fine-tune M 2 in a unified manner to adaptively adjust the model updates to get a performance improvement guarantee. Notably, our algorithm can use the same learning rate of Phase 2 and our algorithm is robust to this learning rate. To the best of our knowledge, this is the first method that unifies model shift and model bias and adaptively fine-tunes the model updates during the training process.

-H-glyph[turnstileright]-glyph[arrowbothv]-{-C-〈-glyph[ceilingright]-glyph[ceilingright]-glyph[unionsq]-glyph[turnstileright]-〈-↖-glyph[subsetsqequal]-∈

-W-glyph[turnstileright]-glyph[arrowbothv]-‖ -glyph[ceilingright]-∇-∈-glyph[ceilingleft]-↖-glyph[subsetsqequal]-∈

-H-glyph[intersectionsq]-glyph[arrowdblbothv]-glyph[turnstileright]-\-glyph[wreathproduct]-〉-glyph[ceilingleft]-↖-glyph[subsetsqequal]-∈

-A-\-glyph[unionsq]-↖-glyph[subsetsqequal]-∈

-I-\-glyph[subsetsqequal]-glyph[ceilingright]-∇-glyph[unionsq]-glyph[ceilingright]-glyph[ceilingleft]-P-glyph[ceilingright]-\-glyph[ceilingleft]-glyph[intersectionsq]-glyph[arrowbothv]-glyph[intersectionsq]-glyph[arrowdblbothv]-↖-glyph[subsetsqequal]-∈

-H-glyph[wreathproduct]-√-√-glyph[ceilingright]-∇ -↖-glyph[subsetsqequal]-∈

0.008

0.006

0.012

diff of values

0.004

fofvalues

0.008

0.004

0.002

0.000

0.000

0.000

0

10

20

30

0

10

20

30

0

10

20

30

training times

training times

training times

0.0002

0.0030

0.00002

ofvalues

0.0001

0.0000

0.00000

0.0000

0

10

20

30

0

10

20

30

0

10

20

30

training times

training times

training times

-glyph[ceilingleft]-〉-{-{-glyph[ceilingright]-∇-glyph[ceilingright]-\-glyph[floorright]-glyph[ceilingright]-←-glyph[wreathproduct]-{-←-glyph[arrowdblbothv]-glyph[wreathproduct]-glyph[ceilingleft]-glyph[ceilingright]-glyph[arrowbothv]-←-glyph[floorleft]-〉-glyph[turnstileright]-∫

-glyph[ceilingleft]-〉-{-{

glyph[ceilingright]

∇ glyph[ceilingright]

\-glyph[floorright]-glyph[ceilingright]-←-glyph[wreathproduct]-{-←-glyph[arrowdblbothv]-glyph[wreathproduct]-glyph[ceilingleft]-glyph[ceilingright]-glyph[arrowbothv]-←-∫-〈-〉-{-glyph[unionsq]

-glyph[ceilingleft]-〉-{-{-glyph[ceilingright]-∇-glyph[ceilingright]-\-glyph[floorright]-glyph[ceilingright]-←-glyph[wreathproduct]-{-←-glyph[wreathproduct]-√-glyph[unionsq]-〉-glyph[arrowdblbothv]-〉-‡-glyph[turnstileright]-glyph[unionsq]-〉-glyph[wreathproduct]-\-←-glyph[wreathproduct]-glyph[floorleft]-|-glyph[ceilingright]-glyph[floorright]-glyph[unionsq]-〉-glyph[subsetsqequal]-glyph[ceilingright]-←-

F Additional Experiment

F.1 Working Mechanism Extension

To illustrate that the ability of USB-PO to reduce both model shift and model bias potentially is not a coincidence that exists only in the Walker2d environment, we add experimental results in other MuJoCo [46] environments. As shown in Figure 10, when the fine-tuning actually operates, the difference of the model shift term and the model bias term among all of the benchmark tasks are generally both positive, further validating our superiority.

F.2 Ablation Study Extension

Here, we show the results of the ablation study on all of the MuJoCo benchmark tasks.

As shown in Figure 11, only optimizing the model shift term results in a drop in sample efficiency while only optimizing the model bias term leads to performance deterioration. Only fine-tuning the model updates in a unified manner can achieve excellent performance.

HalfCheetah-v2

Humanoid-v2

InvertedPendulum-v2

Hopper-v2

with none

with model shift

with model bias

with both

6000

12000

average retum

m

4000

8000

2000

2000

4000

Q

0

0

OK

50K

100K

150K

200K

250K

OK

50K

100K

150K

200K

250K

300K

OK

50K

100K

150K

200K

250K

300K

steps

steps

steps

6000

006

3000

4000

averagereturn

600

2000

2000

300

1000

0

0

OK

50K

100K

steps

150K

200K

250K

OK

2K

4K

6K

steps

8K

10K

12K

14K

OK

20K

40K

steps

60K

80K

100K

120K

F.3 Ensemble numbers

To illustrate the justification of the ensemble numbers we set, we conduct the ablation experiment on more ensemble numbers containing 3, 5, and 7. As Figure 12 shows, as the number of ensemble models goes up, the performance will be higher and more stable, but it will cost more time. To maintain the balance between performance and time, we finally set the value of this parameter to 7, which is recommended by the MBPO [20] original repo.

0K

50K

100K

150K

200K

250K

300K

steps

0

2000

4000

average return

Algorithm

3

5

7

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper investigates the problem of model shift in model-based Reinforcement Learning (RL). The author proposes a two-stage model learning method called USB-PO (Unified model Shift and model Bias Policy Optimization). The first stage is the same as previous work, updating based on the previous model using model replay buffer. The second stage is based on the previous model and the replay buffer, aiming to simultaneously reduce model shift and model bias after the model first stage update. The paper also provides a theoretical analysis for the proposed method, and the effectiveness of the method is demonstrated in experiments conducted on MuJoCo.\n\nSTRENGTHS:\n1. As far as I know, the method proposed in this paper is novel.\n2. The authors provide theoretical analysis and motivation for their method.\n3. The ablation study in the paper is rather comprehensive.\n\nWEAKNESSES:\n1. What is the basis for lines 133-135? I don't think such a conclusion can be derived from Definition 1 and Theorem 1.\n2. What is the relationship between Theorem 2 and 3? Both seem to be proving different bounds for the same term, what is the significance of this? Judging from the appendix, there appears to be a typo in Theorem 3.\n3. In the experiment section, the performance curve of the MBRL method CMLO drops much more than what is shown in the original paper. I would like to know the reason for this.\n4. The baselines STEVE and SLBO are too outdated, and they can no longer be considered as state-of-the-art (SOTA) baselines for MBRL.\n\nQUESTIONS:\n1. The two-stage model learning will inevitably lead to more time consumption. Can you specify how much the time consumption has increased?\n2. Please compare with some recent MBRL methods instead of STEVE and SLBO, such as PDML [1] and ALM [2].\n\n\nReference:\n\n[1]. Wang et. al. \"Live in the Moment: Learning Dynamics Model Adapted to Evolving Policy,\" in ICML 2023, https://arxiv.org/abs/2207.12141.\n\n[2]. Ghugare et. al. \"Simplifying Model-based RL: Learning Representations, Latent-space Models, and Policies with One Objective,\" in ICLR 2023, https://arxiv.org/abs/2209.08466.\n\nLIMITATIONS:\n1. The time efficiency needs to be reported.\n2. More baseline comparisons are required.\n3. It would be desirable to compare with more benchmarks, such as DMC and Metaworld, not just MuJoCo.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a new MBRL algorithm which adaptively adjusts the model update by considering the model shift (distance towards previous step update) and the model bias (towards true env model) simultaneously, which avoids model overfitting and shows better empirical performance compared to only considering either of them.\n\nSTRENGTHS:\n- This paper clearly discusses the drawbacks of previous MBRL approaches, either considering model bias only, which gives possible excessive model updates, or considering model shift only, which might be subject to overfitting to the previous model. A natural and neat idea is proposed by unifying the two aspects together to further fine-tune the model, which supported by strong empirical results. \n\n- The experiments section is thorough. The ablation study is interesting, which shows the performance of USB-PO is better than removing any of the model bias or shift term. Section 5.2 is also nice, which empirically verifies the assumptions in the method development, by removing the consideration of $\\Delta$, as well as the visualization on the optimization objective value to see how the update happens. These experiments facilitate the understanding the method.\n\nWEAKNESSES:\n- The biggest weakness of the paper is that it is very poorly written, the idea is simple but the way the paper organizes and writes make it very difficult to convey the main message. For example, Section 4.1 is too general and it is hard to get what it is trying to say. Algorithm 1 does not show any important message, for example, Line 115 says the optimization objective is MLE loss and the phase 2 is for finetuning, which is too high-level and it makes readers confused. \n\n- This might be correlated with point 1, the algorithm in Algorithm is not super clear to me. It seems Equation 3 is used to learning $M_2$, but the Algorithm 2 Line 2 says we get $M_2$ from $M_1$ update, this is confusing. Is Line 3 only for finetuning $M_2$ given the previously learned $M_1$, $M_2$? These are the main points of the paper and it is important to make things clear and precise.\n\nQUESTIONS:\n- For Line 195, could the $M_1$ ensembles be also used in the bias term estimate?\n- The experiment in Figure 6, is it possible to try more lr of the MBPO to showcase the advantage?\n- Another ablation study is to varying the number of ensembles for the models, and comment on the computational complexity.\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper considers model based RL (MBRL), focusing on the question of fine-tuning the model while learning the policy (the policy optimization).  While prior works treat these two aspects somewhat separately with tuning thresholds, the authors derive a single cost function that incorporates both.  The paper presents a lengthy derivation of the cost function and shows through inequalities and assumptions that it is a well grounded approach, and then numerical experiments validate the approach and compare to MBRL methods as well as model free (MFRL) methods.\n\nSTRENGTHS:\nThe paper is very clearly written, with good motivation and placement of the results in the state of the art.  The idea of combining the model refinement and policy learning is perhaps obvious; what is not obvious is how this can be carried out.  The paper does this theoretically, and then carries this through to a useful algorithm.  The derivations are clear and the assumptions needed to bring this to a numerical method are laid out.  \n\nTheorem 2 is nice because it bounds the expected return difference between and 2 models and their respective policies.  This should be useful in other learning contexts.\n\nThe meta-algorithm (Algorithm 1) puts the overall ideas across in a way that doesn’t require the reader to understand the subtleties of the proofs.\n\nWEAKNESSES:\nThe reviewer doesn't see any significant weaknesses in the paper.  Some clarifying questions are listed below.\n\nQUESTIONS:\nPerhaps it is worth saying more about the assumption after eqn (9), regarding policy shift being small wrt the model bias?  When would this not be true?\n\nDeriving eqn (10) relies on a Gaussian assumption.  Could you say why this is generally valid, or when it might not be so valid?  \n\nSimilarly in eqn (11) the Wasserstein distance is easily calculated for the Gaussian case.  So here the Gaussian assumption is carrying through?\n\nAfter Algorithm 2 statement, “following [18]”: What is the elite set of B models?  And you choose one of these with uniform probability? How large is B typically?\n\nThe “pull-back” interpretation is interesting.  That is what you call the no-update case?\n\nIn the appendix, Table 2, top 3 cases, the model-free works well, albeit with a lot of steps.  Can you compare the overall complexity of the MFRL and USB-PO for these cases, accounting for cost per step of each?  Also, is there a hybrid MFRL-MBRL approach?\n\nLIMITATIONS:\nThere are no negative social issues.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-06"}

comments: SUMMARY:
This paper studies the learning of the dynamics model in model-based reinforcement learning. The authors propose a novel method USB-PO that has provable properties.

STRENGTHS:
1. The paper studies an important problem in model-based RL, namely the model shift and model bias balance during model update without heavy dependence on the threshold and a lack of adaptability.
2. The algorithm is straightforward and the theoretical results also match the expectation.
3. The paper is clearly written and easy to follow.

WEAKNESSES:
1. In the pseudocode Algorithm 1, what is the difference between $M$ and $M_2$? Are they supposed to be the same?
2. Can the authors elaborate the MBRL algorithms that have heavy dependences on the threshold?
3. Can phase 1 and phase 2 be integrated to a single joint optimization objective? It would be easier to implement if the algorithm is a trust-region style algorithm that maintains a single model and regulates the updated and the last-iteration model. 
4. There are also previous MBRL works [1, 2, 3] that have a dual update procedure, in a very similar way to how the proposed method regulates the model updates. I recommend the authors to also discuss these related works in a later version of the manuscript.

[1] Zhang et al. Conservative Dual Policy Optimization for Efficient Model-Based Reinforcement Learning.\
[2] Sun et al. Dual Policy Iteration.\
[3] Levine et al. Learning neural network policies with guided policy search under unknown dynamics.

QUESTIONS:
See the weakness section above.

LIMITATIONS:
See the weakness section above.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-07