How to Fine-tune the Model: Unified Model Shift and Model Bias Policy Optimization
Authors: Hai Zhang
Hang Yu
Junqiao Zhao
Di Zhang
Chang Huang
Hongtu Zhou
Xiao Zhang
Chen Ye
Venue: NeurIPS 2023 poster (2023)
Keywords: model-based reinforcement learning
model shift
model bias
fine-tuning
performance difference bound
TL;DR: We theoretically derive an optimization objective that can unify model shift and model bias and then formulate a fine-tuning process, adaptively adjusting model updates to get a performance improvement guarantee while avoiding model overfitting.

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
Hai Zhang Hang Yu Junqiao Zhao" Di Zhang Chang Huang Hongtu Zhou Xiao Zhang Chen Ye Department of Computer Science, Tongji University. Shanghai, China MOE Key Lab of Embedded System and Service Computing, Tongji University, Shanghai, China (zhangha inteii121101 12138, 2053881, zhaojunqiao, (eeejungia)@@tnnnee @tongjj i cn
Abstract
Designing and deriving effective model- reinforcement learning (MBRL) al- gorithms with a performance improvement guarantee challenging. mainly at- tributed the high coupling between model learning and policy optimization. Many prior methods that rely on return discrepancy to guide model learning ig- nore the impacts of model shift, which can lead 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. 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 performance improvement guarantee while avoiding model over- fitting. Based on these, we develop straightforward algorithm USB-PO" (Uni- fied model Shift and model Bias Policy Optimization). Empirical results show that USB-PO achieves state-of-the-ee performance on several challenging benchmark tasks.
this paper, we theoretically
Introduction
Nowadays, reinforcement learning (RL) has been gaining much traction in a wide variety of com- plicated decision-makinn tasks ranging from academia industry [2 9, due to some remarkable model-free RL (MFRL) algorithms [20 IT. IT IL [S) which show desir- able asymptotic performance. However, their applications are hindered by the bottleneck of sample efficiency. On the contrary, model--asee RI. (MBRL) algorithms, using world model to generate the imaginary rollouts and then taking them for policy optimization [G1], 20), have high sample effi- ciency while achieving similar asymptotic performance, thus becoming compelling alternative in practical cases
28] Part of this is
I6. 52].
Typically, MBRL algorithms iterate between model learning and policy optimization. Hence the model quality is crucial for MBRL. Many prior methods [G], 20, 53, 38, [3]] rely return discrep- ancy 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 consider the impacts of model shift between two iterations [[I]] which can lead to performance deterioration due excessive model updates. Although CMLO [[]]] explicitly considers model shift from the per- spective 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
Corresponding author
Code: http://github.coo/bbrrtt 121 38/Ufiif Model-Shift-am Shift-anddModel- Model-Biiss Laiss-Poicy--ppiiiiii
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

set too low, the model bias of the following iteration will be large, which impairs the subsequent op- timization 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 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 [ 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 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 (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 [LL H5 7 following their different usage of the model. Our work falls into the Dyna-style 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 asymptotic performance.
II, 54L B
R.
category
Some previous algorithms focus on the policy optimization process. CMBAC [b_]]] introduces con- servatism to reduce overestimation of the action value function, and ME-TRPO [23]] imposes con- straints 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 generic algorithm similar to [24] 20, 21] that can be plugged into many SOTA MFRL algorithms [[LS] 2 rather than just proposing for specific policy optimization algorithm.
A key issue in model learning is model bias, which refers the error between the model and the real environment [20]. As the imaginary rollout horizon increases, the impacts of model bias accu- mulate rapidly, leading compounding error and unreliable transitions. To mitigate this problem. several effective methods have been proposed. The ensemble model technique [R. 25, [ and the dropout method [[2]] are employed to prevent model overfitting. The uncertainty estimation tech- niques are used to adjust the rollout length [20, 30] or the transition weight [19, [8]. Furthermore, The multi-step techniques [Di ES| 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 core concern in both MFRL and MBRL theoretic avenues. In MFRL, methods such TRPO [4I]] and CPI [22] choose optimize the performance differ- ence bound, whilst most of the previous work in MBRL. [11 2n b31 5% [e! choose to optimize the difference of expected return under the model and that of the real environment, which termed return discrepancy. However, return discrepancy ignores model shift between two consecutive it- erations 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 [I0]] designs an algorithm to optimize on- on-policy model error, but it is sim- ilar to return discrepancy in nature. DPI [EX]] uses dual updates to improve sample efficiency but tries to restrict policy updates within the trust region, thus inhibiting exploration. CMLO [2T]] relies on a fixed threshold to constrain the impacts of model shift, resulting a heavy dependence on the threshold and lack of adaptability during the training process. Hence, try to unify model shift
2

---

and model bias to form novel optimization problem, adaptively fine-tuning the model updates to get a performance improvement guarantee. Still, some prior work [[ 36. 54il choose consider re- gret bound, among which [544]] 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 [S4, EK 27 that use dual optimization to train the policy, we devise extra phase to fine-tune the model.
3 Preliminaries
We consider Markov Decision Process (MDP), defined by the tuple M AS, 7, Po). S and A denote the state space and action space respectively, and E (0, 1) denotes the discount factor. p(s'ls s a) denotes the dynamic transition distribution and we denote PM- (s'ls, a) as that of the real environment. Po(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 Vy- as Ell(l)):
π. = arg max VM- = (an (1)),
y (s ae) so, so ~ po(s)
(1)
MBRL algorithms aim to learn the dynamic transition distribution model, (s'(s's,, (8,8),,, 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 VM 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 rM as the reward function. Besides, we denote M as parameterized family of models and II as a parameterized family of policies.
Let d'r(s, a) denote the normalized discounted visitation probability for (s, when starting at 80 Po and following under the model M. Let ( denote the probability of visiting timestep & given the policy and the model M.
(8) (1 (1-7)2 (1--))
Σ I' (
We define the total variation distance (TVD) estimator as Drv(-1l.) and the second-order Wasser- stein distance estimator as ..,..).
4 USB-PO framework
In this section, we demonstrate a detailed description of our proposed algorithmic framework. i.e., USB-PO. In Section [. meta-algorithm of the USB-PO framework is provided generic solution. In Section 2 we theoretically show how to unify model shift and model bias to get per- formance improvement guarantee In Section B.3. the practical algorithm proposed to instantiate the USB-PO framework.
4.1
The overall algorithm
The general algorithmic framework of USB-PO is depicted in Algorithm MI where the main differ- ence compared to the existing MBRL, algorithms is the two-phase model learning process, namely phase and phase 2. Phase uses traditional MLE loss to train the model, which may impair the performance by excessive model updates due to only considering the impacts model bias. To mit- igate this problem, we introduce phase to further fine-tune the model updates, whose optimization objective is defined as (G.
arg min Jphase2 = Eog, [We(PM, Wa(PM,, Par-)] PMa
(3)
The detailed derivations of all the theorems in this section presented in the appendix.

---

Algorithm Meta-Algorithm of the USB-PO Framework
1: Initialize the policy and the learned model;
2: Initialize the environment replay buffer and the model replay buffer DMi
3: for each epoch do
4: Use π to interact with the real environment: D - U {(s, s.):
5: Backup the current learned model for future use and denote this backed-up model PM;
6: Phase I: use D to train the learned model with the supervision of MLE and denote this updated model PM
7: Phase 2: use (q) as optimization objective to further fine-tune PMai
8: Use РM, generate the imaginary rollouts: DM DM U {(M aM, sm)
9: Use D U MM to train the policy
end for
(q.(E) unifies the model shift term and the model bias term in the second-order Wasserstein distance form, namely (2 PMM,) and (2(PM,, WP(PMM-PM-) PMr- ). thus achieving adaptive adjustment of their im- Section 4.2 and Section 5.4 this is not
pacts during the fine-tuning process. As demonstrated equivalent to the traditional methods of limiting the magnitude of model updates, but rather benefi- cial to get performance improvement guarantee.
4.2 Theoretical proof
Definition 1 (Performance Difference Bound). Recalling that V=iM, denotes the expected return under the real environment of the policy ni derived from the model M4 E M. The lower bound on the true return gap of and T2 can be stated as,
VTalMa
y [M1 2
By constantly increasing the value of C, the lower bound the performance difference is guaran- teed to be lifted, leading to performance improvement. Therefore, we try to take model shift and model bias into the formulation of C and maximize
achieve our goal.
Theorem (Performance Difference Bound Decomposition). Let M; E M be the evaluated model and n. E II the policy derived from the model. The performance difference bound can be decom- posed into three terms,
V Mr V
(V=alMa
MT,) (VT
Mr,
(V
V
(5)
Obviously, compared to directly optimizing the return discrepancy of each iteration [20], the per- formance difference bound chooses optimize the return discrepancy of two adjacent iterations, namely /T VM, and |T VM, respectively, and the expected return variation between these two iterations. namely T M, demonstrating better rigorousness. For further discussion, we introduce the following theorem.
Theorem 2 (Return Bound), Let Rmar denote the bound of the reward function, 6x denote Mr
max, Drv (TT]I T2) and EM, denote B(saayad), [Drv (PM, [lpM,)]. For two arbitrary policies TI, T2 II, the expected return under two arbitrury models M1, M2 E -M can be bounded as,
M2
Mi 2
22Rma
(1 71-
(1 7)- M.
7
Ma
(6)
By using (6), can easily bound the decomposition terms of the performance difference bound in (5).
Theorem (Decomposition TVD Bound), Let ci, denote Eisaard [Drv (pM, ((pm-))), }), Let M;
M be the evaluated model and n. 11 be the policy derived from the model. The decomposition terms can be bounded as,
2 2Rmar To
2R 6
\ V >
(1 7)2 (EMs
-Ma
-Ma
(1 - 7)
(7)

---

Notably, the model shift term and the model bias term in TVD form, namely Drv (PM, llpMs) and Drv (Paz IPar.)). ) are already present the right-hand side of the Eq.(I. However, due to the unknown term T2, we can not sample from d Mag Thus, we need to make a further transformation about the (q.(W).
Theorem 4 (Unified Model Shift and Model Bias Bound). Let R. denote the constant (1-7)²
2Rmaz and
A denotes Bss,a)yaāh (pM, [|pM-]]] Bis,a)ddf [Drv (pM, [PM-))] Let M, M be the evaluated model and na E II be the policy derived from the model. The unified model shift and model bias bound can be derived as,
VTalMa VTIlAf
(8)
2 ((7(BB DDr (Pr,(P,,, - Drv (PM, lPMM)) ) - Drv (PMr(MM,llMm... )] A) - 6x)
The term in (g) is still intractable. However, the fact that it covers dii and d reminds us that we may explore its relationship with the three TVD terms lying in the expectation of dīi,-
Theorem 5 (IAI Upper Bound). Let M; E M be the evaluated model and m. E II be the policy derived from the model. The term A can be upper bounded as:
27 1-7 B(o.o)aa
2e, max DTV (PM 1-7 8,a
A S
max 5,a
DrV(PM,
(9)
Under the online setting, it is assumed that the error caused by policy shift, namely Cz. compared to model bias has relatively small scale [20, [33]. Therefore, we ignore the minor influence brought by the policy. Additionally, the terms lying in the expectation of dM, in (q. 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 IÅI become higher-order term compared to each of them alone. Thus, the critical element to lift the lower bound in (q.(S) is the terms lying in the expectation of di By maximizing these terms, namely minimizing Drv (PM, (lPM,) + Drv (PMa 1P ). we can achieve unifying model shift and model bias and adaptively fine-tuning the model updates to get performance improvement guarantee.
To make it practically feasible, we make additional assumptions. Specifically, let (X, Σ) be a measurable space and F is set of functions mapping X to R that contains Vāg- When V is LypLLipschiz with respect to a norm
||- Il. the integral probability metric [33]] of two arbitrary dy-
namic transition function M, M' E .M defined on X' is as Eq.(II Notice that if S p < q.
W Pwr)= S W₂ (PM PMr) [22), and we can get W1 (PM PM') S W2(PM PMr). Hence, Eq.(G)
can be applied as the optimization objective to get performance improvement guarantee.
Eer [f'')] Exeepy [(()]] = Rmaz Drv (PMllll)) = L,, W (PM.PMr) SEF
sup
(10)
1-7
4.3 Practical implementation
We now instantiate Algorithm 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 optimize the policy x.
Predictive Model. We use a bootstrap ensemble of dynamic models {phr..., P. ). whose ele- ments are all probabilistic neural network, outputting the Gaussian distribution with diagonal covari- ance: pls ($141, (tt1 |st, ax) = [(rrg a.), En at)). 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).
"Dv (pM, (pp->) ) is constant for the optimization of M2
5

---

Model Shift and Model Bias Estimation. Recalling that the model shift term refers to W2(PM, Ilpas,)) and the model bias term refers to W2(PM, (PM-). For arbitrary two Gaussian distri- butions ~ ((n En) and (u2,, E2), the second-order Wasserstein distance between them derived as (q.(II) [35]]
(2(p.)) = V Il - p2ll trace(>i) + + 2 (
(11)
Let K1 denote the index of the selected model from M, ensemble and l2 denote that from M2 ensemble. For the model shift term, we approximate it as W2(PM, PMr,). For the model bias term, we approximate it as -ーM 1b W2(PM,, Phaz).
Model Rollout. Following the previous work [20], use short branch rollouts alleviate the impacts of compounding error.
Policy Optimization. Our work allows being plugged in many SOTA MFRL algorithms, e.g. SAC [TS], TD3 [_]] etc. Here we employ SAC as an example.
Experiment
Our experiments are designed to investigate four primary questions: (1) Whether we need to propose algorithm similar to USB-PO to adaptively adjust the impacts of model shift? (2) How well does USB-PO perform on rein- forcement learning benchmark tasks compared to SOTA MBRL and MFRL algorithms? (3) How to understand USB-PO? (4) Does USB- PO have similar performance on different learning rate settings in phase 2?
Walker2d-12
Algoritim
4000
3.0 (CMLO)
pneegeereu retur
1.0
0
5.1 Necessity of USB-PO
0K
50K
100K
150K
200K
250K
300K
steps
We recall that CMLO PI sets a fixed threshold 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 Mu- JoCo HG. As Figure 0I shows, the performance corresponding to the other two thresholds (1.0
Figure 1: CMLO performance curves for differ- ent threshold settings over different random seeds, where 3.0 is the threshold recommended in the pa- per.
and 5.0) is severely affected. Therefore, setting fixed threshold to constrain is inappropriate, and smarter way like USB-PO should be applied 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 [2I]]. which carefully chooses the threshold for different envi- ronments 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 aim the model training process, MBPO [20]], which is the variant without our proposed model fine-tuning process and uses return discrepancy
get performance improvement guarantee. The MFRL baselines include SAC [[SS], the SOTA MFRL algorithm in terms of asymptotic performance and PPO [[]]. which explores monotonic improvement under the model-free setting.
We evaluate USB-PO and these baselines on six MuloCo [А6]] continuous control tasks in OpenAI Gym [G), covering Humanoid, Walker2d, Ant, HalfCheetah, Hopper and Inverted-Pendulung 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 shows the learning curves of all compared methods, together with the asymptotic perfor- mance. The results show that algorithm is remarkably advanced over the MFRL algorithms

---

HairCheetah-v2
Walker2d-v2
Humanoid- .2
8000
8000
12000
4000
4000
2000
2000
0
0K
50K
100K
150K
200K
250K
DK
50K
100K
150K
200K
250K
300K
0K
50K
100K
150K
200K
250K S0OK
stepa
steps
steps
Ant-v2
InvertedPendulummv2
Hopper-v2
8000
return
3000
800
2000
2000
1000
0
OK
100K
150K
250K
0K
2K
5K
7K steps
10K
12K
15K
0K
40K
80K
100K
120K
stops
USB-PO(ours)
ALM
CMLO
SAC PPO
PDML
MBPO
convergencc
Figure 2: Comparison against baselines on continuous control benchmarks. Solid curves refer to the mean performance of trials over different random seeds, and shaded area refers to the standard deviation of these trials. Dashed lines refer to the asymptotic performance of SAC (at 3M steps). Note that PDML is not evaluated in InvertedPendulum-
regarding sample efficiency, along with asymptotic performance on par with the SOTA MFRL algo- rithm 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 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 adjust the model updates adaptively is reasonable.
Computational Cost. We report our computational cost compared to MBPO [20] in Appendix D2.4L Although USB-PO is a two-phase model training process, continuing 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 A, the model shift term and the model bias term, validating whether using the second-order Wasserstein dis- tance satisfies the prerequisites for getting performance improvement guarantee, and then more in-depth experiments to illustrate how USB-PO works and show its superiority.
HalfCheetah-v2
Walker2d-v2
0.12
0.08
s 0.08
8 0.04
0.04
0.00
0.00
200
400
600
200
400
600
Irsining times
training times
model shift
model bias
Value Magnitude. We choose 2 challenging tasks MuJoCo, HalfCheetah and Walker2d, and plot the value of A, the model shift term and the model bias term during the training pro- cess respectively. To approximate the A practi- cally, we use the samples from the model replay buffer before the update of the model and the policy to approximate sampling from di), and use the samples generated by them after update approximate sampling from d di..
Figure 3: The value magnitude of the A, the model shift term and the model bias term during the training process in HalfCheetah and Walker2d.
As shown in Figure EL the magnitude of A generally oscillates in a slight manner around 0 while the magnitude of model shift term and the model bias term are much larger than A even in the converged cases. Therefore, is reasonable to ignore the impacts of A in (g) This further supports that minimizing Eq. (L) can get performance improvement guarantee..
7

---

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 Fig- ure BL All random seeds show the consistent result that at the beginning of training the sum of model bias and model shift large and thus the fine-tuning magnitude is relatively large. As the model 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 are reasonable), which further supports our theorem and the motivation this paper, i.e., actively adaptive adjustment of the model updates get performance improvement guarantee, rather than passively waiting for model shift to surpass given threshold to make a substantial update [2T]]. Based on the performance comparison in Figure E1 we argue that actively adaptive updates are more beneficial.
Walker2d-v2
2.0
Sead
Seed
Bead 2
oit Optimizatten objaclive
Seed
Seed 3
average prediction enroR
Beed
0.004
1.0
0.002
0.000
0.0
200
400
800
800
(a) training times
(b) training times
Figure 4: (a) the difference of the optimization objective value before and after the fine-tuning process during the training process over random seeds. (b) the average prediction error during the training process over these random seeds. The model is generally near convergence when training times reach 1K.
Walker2d-v2
0.006
Seed
difference of model shift
average enor
Seed2
Scod1
differance of optimization objective
0.004
0.002
0.000
0.000
10
20
30 0
200
400
600
800
(4) training times
(b) training times
Figure 5: (a) choose specific random seed to show the details of the first 30 training times, covering the difference of optimization objective value, the model shift term and the model bias term before and after the fine-tuning process. (b) the difference of average prediction error before and after the fine-tuning process during the training process over the previously used random seeds.
We further plot the difference of the model shift term, the model bias term and the average prediction error before and after the fine-tuning Ane-tuning process. As shown in Figure B (a), when the fine-tuning actually operates, the difference of the model shift term and the model bias term are both positive. As shown Figure the fine-tuning process has positive effect on reduction of the average prediction emror. 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.

---

Walkor2d-22
Walker2d-22
6000
Algonitim
USB-PO
7e-4
average retum
4000
average retumn
4000
56-4
3e-4
10-4
Algonthi with none model shift with model bias
(onginaa))
2000
2000
0
with
0K
50K
100K
150K
200K
250K
300K
0K
50K
100K
150K
200K
250K
300K
Sleps
steps
Figure The average return of different opti- mization objective variants over different random seeds.
Figure The average return of USB-PO and MBPO with different learning rates over differ- ent random seeds.
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 g! forming bi-verification.
To conclude, when USB-PO actively recognizes that an improper update happens, it performs pull- back 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 done, no actual operation will taken. According to performance comparison results in Figure A it is verified that the model quality plays determinant role in the Dyna-Q [[4]] algorithms.
Walkeer2dv22
HaliChectah.v2
6000
average return
average return
12000
4000
8000
2000
4000
1e-4
3e-4
1e-4
3e-4 leamnig rate
5e-4
learning rate
leamiing rate le-4
leaming rate 3e-4
leamiing rate 5e-4
MBPO maximum performance
Figure 8: Performance of different learning rates on Walker2d and HalfCheetah during the whole training process. Each learning rate is repeated several times over different random seeds and the dashed line refers to the maximum average returns of MBPO.
5.4 Ablation study
In this section, we design ablation experiments strengthen 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 [200], (2) with the model shift term, (3) with the model bias term. As shown in Figure G. only optimizing the model shift term results in drop in sample efficiency since the fine-tuning process makes M, tend to update towards Mj and only optimizing the model bias term leads to performance deterioration due to excessive model updates.
Not Equivalent to Limiting the Update Magnitude. to compare with USB-PO. As Figure W/ shows, USB-PO is more dominant from the perspective of the average return, which further strengthens that USB-PO is not equivalent limiting the update magnitude, but rather beneficial to get performance improvement guarantee.
set different learning rates MBPO

---

Learning Rate Performance Comparison. We set up an ablation experiment on the learning rate of the fine-tuning process, As shown in Figure unlike CMLO which is strongly dependent on carefully chosen threshold for each environment to constrain the impacts of model shift, USB-PO is less sensitive to the learning rate phase 2.
Discussion
In this paper, we propose novel MBRL algorithm called USB-PO, which unify model shift and model bias, enabling adaptive adjustment of their impacts to get a performance improvement guarantee. We further find that 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 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 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. 202I 104 No. 2020YFA071 1402). 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
[1] Brandon Amos al. "On the model-based stochastic value gradient for continuous reinforce- ment learning". In: Learning for Dynamics and Control. PMLR. 2021, pp. 6-20.
[2] Kavosh Asadi al. "Combating the compounding- error problem with multi-step model". In: arXiv preprint 19X.v:1905.13322 (2019).
[3] Greg Brockman et al. "Openai gym"" In: arXiv preprint
(2016).
[4] Jacob Buckman et al. "Sample--fficient reinforcement learning with stochastic ensemble value expansion"..
Advances neural information processing systems (2018).
[5] Annie Chen et al. "You Only Live Once: Single-Life Reinforcement Learning.. In: Advances in Neural Information Processing Systems 35 (2022), pp. 14784-14797.
[6]
Kurtland Chua al. "Deep reinforcement learning in handful of trials using probabilistic dynamics models". In: Advances in neural information processing systems 31 (2018). [7] Sebastian Curi, Felix Berkenkamp, and Andreas Krause. "Efficient model-based reinforce- ment learning through optimistic policy search and planning". : In: Advances Neural Infor- mation Processing Systems 33 (2020), pp. 14156-14170.
[8] Vladimir Feinberg al. "Model-basee value estimation for efficient model-free reinforce- ment learning". In: arXiv preprint 1803.0008000000 (2018).
[9]
Shuo Feng et al. "Dense reinforcement learning for safety validation of autonomous vehicles".
In: Nature 615.7953 (2023), pp. 620-627.
[10] Lukas Froehlich et al. "On-Policy Model Errors in Reinforcement Learning". In: Interna- tional Conference Learning Representations, 2022. URL: https openreview net. httt//ffrm?iddde ee
[11] 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.
[12] Yarin Gal, Jiri Hron, and Alex Kendall. "Concrete dropout". In: Advances in neural informa- tion processing systems 30 (2017).
10

---

[13] Chengqian Gao et al. "Value penalized Qllearning for recommender systems". In: Proceed- ings of the 45th International ACM SIGIR Conference on Research and Development in In- formation Retrieval. 2022, pp. 2008-2012.
[14] Raj Ghugare et al. "Simplifying Model-based RL: Learning Representations, sptce Models, and Policies with One Objective". In: The Eleventh International Conference on Learning Representations. 2023. URL: https obenreview. net forum id
MQcm qRxf Val
[15] 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.
[16] Danijar Hafner et al. "Mastering Atari with Discrete World Models". In: International Con- ference on Learning Representations. 2021. URL: https:: wwtp:ewwwwwwwwe....teeeeeee
[17] Nicolas Heess al. "Learning continuous control policies by stochastic value gradients". In: Advances in neural information processing systems 28 (2015).
[18] Matteo Hessel et al. "Rainbow: Combining improvements in deep reinforcement learning".
In: Proceedings of the AAAI conference on artificial intelligence. Vol. 32. 1, 2018.
[19] Wenzhen Huang et al. "Learning to Reweight Imaginary Transitions for Model-Based Re- inforcement Learning". In: Proceedings of the AAAI Conference Artificial Intelligence.
Vol. 35. 9. 2021, pp. 7848-7856.
[20] Michael Janner et al. "When to trust your model: Model-based policy optimization". In: Ad- vances in neural information processing systems 32 (2019).
[21] 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: http:::://ppeeeeviiee
net, wwwooummmd=aaaaaommm oVV
[22] Sham Kakade and John Langford. "Approximately optimal approximate reinforcement learn- ing" In: Proceedings of the Nineteenth International Conference on Machine Learning. 2002,
pp. 267-274.
[23] Thanard Kurutach et al. "Model-ensemble trust-region policy optimization". In: arXiv preprint arXiv: 10592 (2018).
[24] Hang et al. "Bidirectional model-based policy optimization". In: International Conference Machine Learning. PMLR. 2020, pp. 5618-5627.
[25] Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. "Simple and scalable pre- dictive uncertainty estimation using deep ensembles". In: Advances in neural information processing systems 30 (2017).
[26] David A Levin and Yuval Peres. Markov chains and mixing times. Vol. 107. American Math- ematical Soc., 2017.
[27] 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).
[28] Lanqing Li, Rui Yang, and Dijun Luo. "FOCAL: Efficient Fully-Offline Meta- Rent-Oeeinfrreement Learning via Distance Metric Learning and Behavior Regularization" In: International Con- ference on Learning Representations. 2021. URL:
[29] Timothy P Lillicrap et al. "Continuous control with deep reinforcement learning". In: arXiv preprint arXiv: nete:-50990291 (2015).
[30] Cong Lu et al. "Revisiting Design Choices in Offline Model Based Reinforcement Learn- ing". In: International Conference on Learning Representations. 2022. URL: https::/
openrevies eeriiewweee net www.nne//orrmmicee.... wor.mmid=z 9hxVhE40
[31] Yuping Luo et al. "Algorithmic framework for model-based deep reinforcement learning with theoretical guarantees". In: arXiv preprint /rXiv::807.000 /807.03878 (2018).
[32] Ester Mariucci and Markus ReiSS. "Wasserstein and total variation distance between marginals of Lévy processes". In: (2018).
[33] Alfred Müllrr. "Integral probability metrics and their generating classes of functions". In: Advances in applied probability 29.2 (1997), pp. 429-443.

---

[34] 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). IEFE. 201 8, pp. 7559-7566.
[35] 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.
[36] Ian Osband and Benjamin Van Roy. "Model-based reinforcement learning and the eluder dimension". In: Advances in Neural Information Processing Systems 27 (2014).
[37] Ian Osband et al. "Deep exploration via bootstrapped DQN". In: Advances in neural informa- tion processing systems 29 (2016).
[38] Feiyang Pan 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.
[39] 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.
[40] John Schulman et al. "Proximal policy optimization algorithms". In: arXiv preprint arXiv:1707.06347 (2017).
[41] John Schulman al. "Trust region policy optimization". In: International conference on ma- chine learning. PMLR. 2015, pp. 1889-1897.
[42] David Silver et "A general reinforcement learning algorithm that masters chess, shogi, and Go through self-play". In: Science 362.6419 (2018), pp. 1 1140-1144.
[43] Wen Sun et al. "Dual policy iteration". In: Advances in Neural Information Processing Sys-
tems 31 (2018).
[44] Richard S Sutton. "Dyna, an integrated architecture for learning, planning, and reacting". In: ACM Sigart Bulletin 2.4 (1991), pp. 160-163.
[45] Richard S Sutton. "Integrated architectures for learning, plamning, and reacting based ap- proximating dynamic programming" In: Machine learning proceedings 1990. Elsevier, 1990,
pp. 216-224.
[46] 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.
[47] Leonid Nisonovich Vaserstein. "Markov processes over denumerable products of spaces, de- scribing large systems of automata". In: Problemy Peredachi Informatsii 5.3 (1969), pp. 64- 72.
[48] Arun Venkatraman, Martial Hebert, and Bagnell. "Improving multi-step prediction of learned time series models". In: Proceedings of the AAAI Conference on Artificial Intelli- gence. Vol. 29. 2015.
[49] Tingwu Wang and Jimmy Ba. "Exploring Model-based Planning with Policy Networks". In: International Conference on Learning Representations. 2020. URL.: https
openreview net. net/torrmm e
[50] Xiyao Wang et al. "Live in the moment: Learning dynamics model adapted to evolving pol- icy". In: International Conference on Machine Learning. PMLR. 2023, pp. 36470-36493.
[51] Zhihai Wang et al. "Sample-Efficiee" Efficient Reinforcement Learning via Conservative Model-Basee Based Actor-Critic" In: Proceedings of the AAAI Conference on Artificial Intelligence. Vol. 36. 8. 2022. pp. 8612-8620.
[52] Jingda Wu, Zhiyu Huang, and Chen Lv. "Uncertainty-aware model-based reinforcement learning: methodology and application in autonomous driving". In: IEEE Transactions Intelligent Vehicles (2022).
[53] Zifan Wu al. "Plan To Predict: Learning Uncertainty-Foreseeinn Model for Model-Based Reinforcement Learning". In: Advances in Neural Information Processing Systems 35 (2022), pp. 15849-15861.
[54] Shenao Zhang. "Conservative Dual Policy Optimization for Efficient Model-Based Rein- forcement Learning". In: Advances Neural Information Processing Systems (2022),
pp. 25450-25463.
12

---

7 Appendix
A Proof Sketch
To better clarify theoretical results, we provide proof sketch here. Firstly, decompose the performance difference bound under the model-based setting into three terms (Theorem MI.. Sec- ondly, by means of using Return Bound (Theorem 2), we can bound these three terms individually (Theorem B) Then, we can do some transformation to get Unified Model Shift and Model Bias Bound (Theorem H), which bounds the model shift term and the model bias term in total variation form. However, due to the intractable property of A, we further explore the upper bound of I (Theorem 5D. finding that A be ignored. Finally, by the Integral Probability Metrics (Lemma R) 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
Difference Bound pait
Performance
Intergral Probability Metrics (Lemma
lipschitz continulty of function
Target total variation form
Target
Decomposition three terms (Theorem
Dyv (pan) (Px))) Dyv Pryip (par))
Return Bound (Theorem 2)
ignore
Unified Model Shift and Model Bias Bound (Thearem 4)
Decomposition- Individual bound (Theorem
Upper Bound Theorem 5)
transtormation
A
intractable
Figure 9: Theoretical sketch of USB-PO.
B Useful Lemmas
In this section, we provide some proof to support our theoretical analysis.
Lemma (Total variation Distance), Consider measurable space(S), Σ) and probability measures
P and Q are defined on (St )). The total variation distance between and Q defined as:
Drv (PII)=
sup PP(A) Q(A)) AEE
(12)
Eq. L can be equivalently written as:
Drv Prv((P11)) Σ
(((w)) - Q(()))
(13)
WER
Proof: The proof this lemma can found [26].
Lemma 2 (Total Variation Distance of Joint Distributions). Given distributions p(x, y) p(r)p(ulz) and q(2, y) ((x))(y)x)) the total variation distance between them can be bounded
Drv(p(a, 1) v)) < Drv (p()))))l((())) + max max Drv
(14)
13

---

Proof:
1
Drv(p(a, v)ll((,,)) v)) - Σ
1p(z, v) - q(2,)))
2.9
1
Σ
1 -00)g((()))())))
2
2.9
1
Σ
- ((<))(y/2)
2
2.8
(15)
1
$ 2
p(z)lp(ylz) (pc) -((x))((())))
2.8
1
Σ p(x))p(yl=)) [(v1=)]
Σ 1p(x) g(x)
2
2
2.8
x
Ex-p(s) Drv (pw.l + Drv (p(z))la(z))
s Drv (p(x))la(<)) + max Drv (rr(l(
Lemma (Integral Probability Metrics). Consider a measurable space E). The integral proba- bility metric associated with class F of real-valued functions on X is defined
d)(P,Q) sup fEF
[ [[xp] - EEx=olf(Y)]
(16)
where P and Q are probability measures on X. We demonstrate the following special cases:
(a) Ff (f: 1lfll S c). then we have
dy(P, = cDl(((P()))
(17)
(b) F is the set of L- funscciiouuunctonn with norm || I|, then we have
dy Q) = (P,Q),
(18)
In our paper, to distinguish the dynamic transition function, we choose F to be the class cover- ing Vğ. Since the value function can converge to I- continuity and thus we can get T-7 Drv (PM (ppM)))))))))))) = L.W, (PMs) for any arbitrary model M'..
Emaue, it only needs satisfy the Ly-Lisschitz
C Missing Proof
Theorem (Performance Difference Bound Decomposition). Let M. M be the evaluated model and ne E TI be the policy derived from the model. The performance difference bound can be decom- posed into three terms,
Ma - | (V=alMa V (V
VM, (V
(19)
M
Proof: We introduce two additional terms VM and V that allow the performance difference bound objective to be divided into three operators based on the return bounds, which can be reformulated separately.
M2
VxilMs
V=2lMM2 V.LlM + (Vil Vin + (V V M2
(20)
(VV=llM22 VMz. (VTIMM Vin,) ) (V V
Theorem 2 (Return Bound). Let Rman denote the bound of the reward function, denote max, (TV (12) and es denote Bleaaydd (Drv (pM, [|PM))] For two arbitrary policies
Mr
Ti T2 II, the expected return under two arbitrury models M1, M2 E M can be bounded as,
T M2
V 2 -2R
6n
Ma
(21)
(1-7)²
(1 71- MM
14

---

Proof: We give the thorough proof referring to Lemma B.4 in MBPO | as follows.
VM
(Pi (8. - Piin, (s. ())))))(s.)) a)
M2
t=0
5,a
2- Rmaz
Pi,Ms (8, - Piarr, (s,
(22)
5,a
22
y DTV (8. (8.a))
1=0
According to the Lemma 2 we have:
(rv(pE (s, (s..))l..M (s.))) S Drv (P.M. (s)llp (s)) + max Drv (1)))))
*
(23)
Drv (P (s) (s)) + c
Further we expand the first term:
DTV (Pi,ss, (8) (s)1l17,s, (8))
1
M (s)
P,M,, (s)
2
PM; (0|8 10 L, (6') - -PM (8|6 1')p 1. (6')
2
1
PMM (s|s (s.,,) - (s
(5|s 1
1. (s')
2
1
Σ
P M (8 (0)) (s|b) - 1 (s|61) + PNO (5|s (s1s')lPT_L.M) (6) -Pi-i.Ms (8')
2
E. (0) E PM (s|8') PMA, (s|5 DTV (6')llz(-1,M (5'))
-1M,, (0)) Σ (s|ss) - (s|s')1
(24)
2
Bylnp
1,00 (0)) Σ1 Σ
Pst (8,a)'') - (s,
2
1
PMM (s, als') -
2
5l mp 1,M,
8,
Σ
Ey Bann 7L 1M
(0)) Drv (p as 1PP2, (s,ass))
Eginp (oy]er + Earn, [Drv (pM, (s|s', aa)lll, (s)s), a))]]
f==1
= leg +
Ey arpp a (6 Drv (sls', (s|s/, a))
Then move the result of Eq. (24) to (Q. we can get:
DrV Op.t, (s. (s, a)) S (t+1)er Σ E Baa aapiii Drv (s)s', a))
(25)
15

---

Next, we move the result of Eq. (25) to Eq. (22)). we can get:
MM
M2
2 -2
7'(1 +
IJET +
E aa- (s ( (rv (ols',a)llpm, (s|s' ())s))))))
(26)
t=0
= -mam (
Ex
7 E. papp (r Drv(p,, (s|s', alllla a)))
71-
Here, we first simplify the second term of the Eq.(26)
B .mppl 1.m, (d) Drv (pM, (
7
L ampil (d',a) (s)s's)))
?
1=0
(27)
7
(1 7)
> LE s' (v) Drv (PM, (s|s', a)llp. (sls', a))
(1 7]2
==
7
Eg
anndd (s) Drv (pM (s|s', ()llp (sls,
(1 7)2
7
(1 7)2
Mr
Then we bring this result back to (2. and the proof is complete.
V
Mi 2
2 Rmaa
E
Y
M2
(28)
M2
71- (1 - 7)2 Ms
Theorem 3 (Decomposition TVD Bound). Let M,, denote (a (Drv (PM, lPM M. Let M, e M be the evaluated model and ne E II be the policy derived from the model. The decomposition terms can be bounded as.
V V >
2R 2Rman
7)2 (c
VMh can be approximated as Ma
Maa
2R C (1 7)2
(29)
Proof: According to CMLO [21]] and (21), the term I IM
2Rmaar Ti
(1->)² ', thus we only need bound the remaining two terms.
For the term VTalMa_ Viga we use (21)) to bound it.
max Drv (Tel (2)
VealM22 VIT 2 22
+ ?
(1 7)2 Mire
(1-
(30)
2R R.max Y
(1
7)²
CM/
Similarly, for the term
Ma
V M,
we can get:
V
V
2R
6
My) Mi
(31)
7)2 (1 7)2
We now combine these three bounds together and complete the proof.
/= = (V<alM² VMh))) - -(V Mr
Vin1) + M2
1/21
Mi
2Rmar m2 + 2Rmaa
(1 -7) <N²
(1 - -7) (//
22Raa
(1 - 7)2
En
(1 - 7)= Mi)
7
(32)
2Rmax
2R 2Rmazer cx (1 72
(cNi.
-M2 -c2)
(1 7)2
16

---

Theorem 4 (Unified Model Shift and Model Bias Bound). Let K denote the constant (1-4)² and A denotes a D (pMa [PMM)]] Eis,ayddi (Dv (pM, [PM-)]. Let M; M be the evaluated model and IT be the policy derived from the model. The unified model shift and model bias bound can be derived as,
V=alM2 - VTilM1
(rr((,,,lpm Drv (PM (pM. + + - -xx) (33)
2 (()(B(s.a)m)))
Proof: Based on (29)) we add new term A (aaa)aad (Drv(PM, [PM-)] to reformulate the optimization objective.
Theorem 5 (IAI Upper Bound). Let Mi M be the evaluated model and m. IT be the policy derived from the model. The term A can be upper bounded as:
27 (1 (Pa, max Drv (Pw(PMallPM-))) )]+ 1 max Drv (PMslaal)))
26
<
(34)
Proof: First, we combine these two terms.
A (s.a)aas)) DDrv (pM, (|pM-/-))))) ( (Pasl-))))
= (1 7)
(Pi,M, (s, a) - Pi,Ms, (s, CoDDrv(pMsf
(pm a)ll (6 (s'[s,)))
t=0
8,0
The (8, - Pi,M, (s) (35)
$ (1 7)
max s,a
Drv (pm, (s
Σ
Σ
1=0
8.a
2(1 - 7) max Drv (pM, (PM-)
7 Drv (FT, (s, a)llp, (8,
s.a
1=0
Recalling that we get the result of the sum equation above in Eq. (25), and then have:
A < 2(1 7) max Drv (pMallpas- )(
(11-7)² 7)2 + (1 7)2 Mi
En
7
8.a
27 E
2n
eannd "1 DDr (pms, (PNz) max Drv (pMr, 11m- a.a
+
max Dry (ppsslllps
7
1 1- 8a
(36)
D Experimental Details
D.1 Environment Setup
We evaluate the algorithm series of MuloCo [+6]] 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 MI
Table I: The general outline of the MuloCo environment.
Environment- Version
State Dim
Action Dim
Termination
Ant-v2
27
obs[0]<0.2 obs[0] 1.0
HalfCheetah-v2
17
Hopper-v2
obs[1] 2 0.2 or obs(0) < 0.7
Humanoid-v2
45
17
obs[0] 1.0 or obs[0] 2.0
InvertedPendulun-v2
4
obs > 0.2 or obs| -0.2
obs[0] > 2.0 or obs[0] 0.8 or
Walker2d-v2
17
obs[1] > 1.0 or obs
-1.0
17

---

D.2 Baseline implementation
MFRL Baselines. We use two state-of-the-arr model-free algorithms, i.e. SAC [MS]] and PPO [A0]] 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.
MBRL Baselines. We use several state-of-the-ar model-based algorithms to do baseline compari- son, covering CMLO [2I]], MBPO [20]] SLBO [+T]] and STEVE [B]] The implementation of CMLO based on the opensource repo published by the author and all of the hyperparameters are set ac- cording to the paper [ Our algorithm USB-PO is implemented based on the opensource repo
published by Janner who the author of MBPO.
We present the final performance on six continuous benchmark tasks in Table The results demon- strate that algorithm achieves competitive performance compared 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--2 120K for Hopper-v2.
Table 2: The final performance on six continuous benchmark tasks.
HalfCheetah
Humanoid
Walker2d
STEVE SLBO
12406.29:458.08
4318.32::83.66
1109.23+
1915.47++338.777 73
459.46
3107.93++:887.09
MBRL
MBPO
25755.67::94.54
5546.77-:21.72
4582.06=67.44
CMLO
10143.55=193.82
5577.01
4807.60
USB-PO
15105.91
5973.75
5691.622:162.57
MFRL @ 3M steps)
SAC
15012
6207
5879
Ant
InvertedPendulum
Hopper
STEVE
45.672+45.66
778.54::265.51
1131.61 623.61
SLBO
707.79+-218.8
793.24++334.90
893.68++233.21
MBRL
MBPO
4926. 118888
0.00000=0..00
3436.00= 120.72
CMLO
$123.71
1000.00=:0.00
3495,41 177.022
USB-PO
619..86:111.06
1000.00-10.00
3694.22
MFRL @3M steps)
SAC 5934
1000
3610
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 our USB-PO algo- rithm, the hyperparameters are completely identical to the MBPO settings for all environments. In all benchmark tasks, we set this learning rate le-4.
D.4 Computing Infrastructure
In Table B. we list our computing infrastructure and the computational time for training USB-PO on these six continuous benchmark tasks. Note that the time report is the cost for random seeds simultaneously on one graphics card. For Humanoid, only two random seeds can be run simultaneously because of the limitation of graphics memory.
Table 3: Computing infrastructure and the computational time for each benchmark task compared MBPO, where the time unit d denotes day and h denotes hour.
HalfCheetah
Humanoid
Walker2d Ant InvertedPendulum
Hopper
CPU
AMD EPYC 7B12
64-Core Processor
GPU
NVIDIA
2080Ti
MBPO times
2.46d
1.64d
1.75d
2.88d
3.43h
17.81h
USB-PO times
2.29d
1.51d
1.65d
2.91d
3.42h
18.28h
18

---

E
Comparison with Prior Works
In this section, we compare USB-PO with prior theoretical works to emphasize our contribution, as 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 fixed threshold to constrain model shift. Our algorithm, USB-PO, adaptively adjusts the model updates in unified manner (unify model shift and model bias) to get the performance improvement guarantee.
MBPO-Style [20, 91, 24, 53] They the return discrepancy bound V* 2 VTT CC(m, (>) to improve the lower bound on the performance under the real environment, i.e. as long as improving V by more than C'(Ems 6x) 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 feasible solution to meet the improvement guarantec.
CMI.O-Style [2]]]. They use the performance difference bound under the model-based setting VTalMa Vm (M 2 C' to directly consider model shift and model bias. However, they finally de- rive a constrained lower-bound optimization problem and use a fixed threshold to constrain model shift, i.e. sup,eS.aE.4 Dry (PM, (-|s, PM, (+,,(-s,, a)) I MM2 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 be carefully adjusted for each environment. If this threshold is 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. Addition- ally, using fixed threshold during the whole training process makes the algorithm problematic to adjust adaptively.
|
USB-PO (Ours). Following CMLO-Style [M]] we also use the performance difference bound under the model-based setting to directly consider model shift and model bias. Compared 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 R. Due the intractable property of A, we further explore the upper bound of |A), finding that A can be ignored with respect to model shift and model bias alone (Theorem HD). Finally, the optimization objective we get can be used to fine-tune M2 in unified manner to adaptively adjust the model updates to get 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.
HalfCheetahvv2
Walker2d-v2
Humanoid-v2
0.008
0.006
0.012
0.006
voluee
0.004
0.008
0.004
of
0.002
0.004
0.000
0.000
10
20
30
10
20
30
0
10
20
30
training times
Irsining times
training times
Ant--2
InvertedFenddllum--22
Hopper-v2
0.0002
0,0030
vsluea
0. 0.00002
of
0.0015
0.0000
0,00000
0.0000
20
30
10
20
30
10
30
training times
training times
training times
difference model blas
difference of model shift
difference of optimization objective value
Figure 10: We choose specific random seed to show the details of the first 30 training times all benchmark tasks, covering the difference of optimization objective value, the model shift term and the model bias term before and after the fine-tuning process.
19

---

Additional Experiment
F1
Working Mechanism Extension
To illustrate that the ability of USB-PO to reduce both model shift and model bias potentially is not coincidence that exists only in the Walker2d environment, add experimental results in other MuJoCo [[6] environments. As shown in Figure IM when the fine- tine-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 M 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
WalkerZd-v2
Humanoid-v2
0000
12003
2003
200K
250K
60K
100
190K
2004
200K
100K
InvertedPendulumm-2
Hopper-v2
100
2000
300
shia
26
25C
68
66
25
120
10K
14k
6CK
100
100
with nane
with model shift
with model bias
with both
Figure 11: Optimization Objective Variants on all MuJoCo benchmark tasks.
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, and 7. Figure IP 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 which is recommended by the MBPO [20]] original repo
20

---

Walker2d-v2
Algorithm
4000
average return
7
2000
OK
50K
100K
150K steps
200K
250K
300K
Figure 12: Ablation experimental results of the ensemble model numbers.
21


Reviews:
[{"review_id": "ptEV7JSeoY", "reviewer": "Reviewer_C4Te", "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.", "questions": "See the weakness section above.", "limitations": "See the weakness section above.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "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.\n2. The algorithm is straightforward and the theoretical results also match the expectation.\n3. 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?\n2. Can the authors elaborate the MBRL algorithms that have heavy dependences on the threshold?\n3. 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. \n4. 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.\n\n[1] Zhang et al. Conservative Dual Policy Optimization for Efficient Model-Based Reinforcement Learning.\\\n[2] Sun et al. Dual Policy Iteration.\\\n[3] Levine et al. Learning neural network policies with guided policy search under unknown dynamics.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We greatly appreciate the reviewer’s insightful comments on our paper. The relevant literature mentioned in your comment helps us better polish the paper! Our specific responses to the questions raised by the reviewer are as follows. \n\n___\n\n## 1. The difference between $M$ and $M_2$\n\n$M$ and $M_2$ are different. The purpose of Algorithm 1 is to put the overall ideas that don’t require the reader to understand the following details. Thus, in Algorithm 1 we do not distinguish the model backed up before training with MLE ($M_1$) and the model after training with MLE ($M_2$), but universally use $M$ to denote them. Later, to better illustrate our method theoretically, we denote $M_1$ and $M_2$ separately. Precisely, $M_2$ **denotes the model that needs to be fine-tuned after the training process by MLE in each iteration**. \n\n**To prevent misunderstanding, we revise the Algorithm pseudo-code and Section 4.1. Please refer to point 4 and 5 in global response. Also, we add a schematic diagram to illustrate USB-PO better. Please refer to point 8 in global response .**\n\n___\n\n## 2. Elaborate the algorithm that has heavy dependence on the threshold\nCMLO sets a threshold for model shift to satisfy the monotonic performance improvement guarantee. We argue that this threshold plays a determinant role in the whole algorithm.  ‘’Theorem 4.6 (Refined Bound with Constraint)’’ in CMLO stated\n$$\nV^{\\pi_2|M_2}-V^{\\pi_1|M_1} \\ge \\kappa (\\mathbb{E}\\_{s,a\\sim d^{\\pi_1}}D\\_{TV}[P(\\cdot|s,a)||P_{M_1}(\\cdot|s,a)] - \\mathbb{E}\\_{s,a\\sim d^{\\pi_2}}D\\_{TV}[P(\\cdot|s,a)||P_{M_2}(\\cdot|s,a)])   - \\frac{\\gamma}{1-\\gamma}L(2\\sigma_{M_1,M_2}) - \\epsilon_{opt}\n$$\n, where $\\sigma_{M_1,M_2}$ denotes the threshold to constrain model shift. Obviously, **to meet the monotonic performance improvement guarantee, the right-hand side of the upper inequality must be greater than 0**. Thus, this threshold needs to be carefully designed, as we stated in **line 33-37**. To further verify our description, we **devise an experiment** that sets three different thresholds for CMLO on Walker2d environment (1.0 for the lower threshold, 3.0 for the paper recommended and 5.0 for the higher threshold). As Figure 3 in global response pdf shows, the performance corresponding to the other two thresholds (1.0 and 5.0) is severely affected.\n\n___\n\n## 3. Trust-Region like updates\nNice view! We've thought about this idea. **However, this would introduce a new problem in determining the weights to balance the MLE updates (model bias) and trust region constraints (model shift), falling into the same situation as CMLO**. Though our method seems to be a little more complicated, it is **straightforward and effective** under our theoretical framework. Please see **Appendix C.4**, compared to MBPO, USB-PO does not introduce excessive computational resource cost. However, we believe that this idea may work to further optimize USB-PO with some technical support.\n\n___\n\n## 4. Related works\nThank you again for pointing out our omission regarding the relevant literature! According to your advice, we read the relevant papers covering **CPI [21], DPI [41], GPS with unknown dynamics [26] and CDPO [51]** then we update our related works as follows (changes are bolded). \n\nPerformance improvement guarantee is a core concern in both MFRL and MBRL theoretic avenues. In MFRL, methods such as TRPO [39] **and CPI [21]** choose to optimize the performance difference bound, whilst most of the previous work in MBRL [29, 19, 50, 36, 23] 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 [11] designs an algorithm to optimize on-policy model error, but it is similar to return discrepancy in nature. **DPI [41] uses dual updates to improve sample efficiency but tries to restrict policy updates within the trust region, thus inhibiting exploration.** CMLO [20] 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, 34, 51] choose to consider regret bound, among which [51] 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 [51, 41, 26] that use dual optimization to train the policy, we devise an extra phase to fine-tune the model.**\n\n[21]   Sham Kakade et al. “Approximately optimal approximate reinforcement learning”. In: Proceedings of the Nineteenth International Conference on Machine Learning. 2002, 358\npp. 267–274.\n\n[41] Wen Sun et al. “Dual policy iteration”. In: Advances in Neural Information Processing Systems 31 (2018).\n\n[26] Sergey Levine et al. “Learning neural network policies with guided policy search under unknown dynamics”. In: Advances in neural information processing systems 27 (2014).\n\n[51] Shenao Zhang. “Conservative Dual Policy Optimization for Efficient Model-Based Reinforcement Learning”. In: Advances in Neural Information Processing Systems 35 (2022), pp. 25450–25463.\n\n___\n\nWe hope the reviewer can consider raising the score if we resolved the reviewer’s concerns and we would be happy to have further discussion."}, {"review_id": "ZKAe0gYQDC", "reviewer": "Reviewer_dkWj", "summary": "The 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. ", "questions": "Perhaps 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\n", "limitations": "There are no negative social issues. ", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "The 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", "weaknesses": "The reviewer doesn't see any significant weaknesses in the paper.  Some clarifying questions are listed below.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank the reviewer for the high recognition of our work! Our specific responses to the questions raised by the reviewer are as follows:\n\n___\n\n## 1. Policy shift wrt the model bias\n· According to the experiment in MBPO, when the amount of data provided by the policy for training the model is very small, the policy shift can not be ignored.\n\n· In my view, as the model gradually converges, model bias will drop to a very small amount, which may make the policy shift unignorable. Further research on this topic may improve performance and save computational resources.\n___\n\n## 2. The detail of deriving Eq.(10)\nDeriving Eq.(10) does not require the assumption of the Gaussian distribution. Please refer to **line 167-170**, the derivation of Eq.(10) **only needs the assumption that $V^\\pi_M$ is** $L_v$**-Lipschitz** and please refer to **Appendix A Lemma 3** for the corresponding details. Though satisfying the Lipschitz continuity is actually difficult, **this occurred frequently in previous theoretical work to get better convergence properties**. Based on the above statement, deriving Eq.(3) also does not require the assumption of Gaussian distribution, except that a closed-form expression can be further obtained with the help of the Gaussian distribution.\n\n**To prevent misunderstanding, we make revisions on line 171-173. Please refer to point 6 in global response.** \n___\n\n## 3. Validity of the Gaussian distribution assumption\nIn MBRL, it is normal to assume that the model is subject to the Gaussian distribution. Since PETS [6] demonstrated that **Gaussian ensemble models can prevent model overfitting and can better capture model uncertainty** (please refer to the details in **line 180-186**), the subsequent methods (MBPO P2P CMLO M2AC...) have all used this approach. But if you're considering specifically that samples generated by the policy for training the model are internally unbalanced, using the Gaussian distribution has obvious weaknesses.\n\n[6] 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).\n\n___\n\n## 4. The details of elite set\nThe elite set first appeared as a trick in the source code published by the authors of MBPO. Each time the model is trained, a portion of the samples are set aside for holdout validation. Then, the algorithm is ranked from highest to lowest based on the validation error, and the $K$ lowest ranked is set as the elite set ($K<B$). **When generating rollouts, randomly selecting a model from the elite set rather than the original ensemble set gives better performance**. We set the same parameters as MBPO, namely the size of ensemble models is 7 and the size of the elite set is 5.\n\n___\n\n## 5. The pullback interpretation\nThanks again for your confirmation of the “pull-back” explanation! However, this is not for the no-update case. The no-update case means the fine-tuning optimization objective (Eq.(3)) basically does not change the updates of $M_2$ and it can be considered that the updates generated by MLE are appropriate for Eq.(3). When the MLE generates updates with large model shift which may impair the performance, the algorithm can fine-tune (**pull**) the model updates to get the performance improvement (**back**). We call this process of securing the performance improvement guarantee as “pull-back”.\n\n## 6. Overall complexity\nSorry, we can not analyze the overall complexity theoretically within the limited time, thus we choose to report the time cost as below and we leave this theoretical work as future work.\n\n| Env | USB-PO | SAC |\n| :-----:|:----: | :----: |\n| HalfCheetah | 2.29day(250K) | 12.53h(3M) |\n| Humanoid | 1.51day(300K) | 12.83h(3M) |\n| Walker2d | 1.65day(300K) | 12.55h(3M) |\n\nAs the above table shows, though SAC needs a lot of steps to get asymptotic performance, it costs much less time compared to USB-PO, because USB-PO costs much more time to train the model in each step. **However, in some sample-dangerous scenarios, we have to sacrifice the time complexity (MFRL) to pursue the sample complexity (MBRL)**. As for the hybrid algorithms, as far as we know, when the model is used for generating rollouts, this approach is generally recognized as a model-based algorithm.\n___\n\nIf the reviewer has more questions, please let us know and we'll be happy to answer them!"}, {"review_id": "vJzSgWobNc", "reviewer": "Reviewer_VRuX", "summary": "This 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. ", "questions": "- 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 ", "limitations": "N/A", "rating": 3, "confidence": 3, "soundness": 3, "presentation": 1, "contribution": 3, "strengths": "- 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", "weaknesses": "- 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\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for identifying our technical contributions. The valuable comments help us improve our submission! The reviewer’s primary concern seems to be how the paper is written and organized. Our specific responses to the concerns raised by the reviewer are as follows:\n\n___\n\n## 1. 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. / The algorithm in Algorithm is not super clear.\n\nWe notice a strong correlation between the two points mentioned by the reviewer in the Weaknesses, so **we make a unified answer here**. After carefully checking the reviewer’s comments and reading the statements in our paper, we find that Section 4.1 is indeed not very clear in its description and fails to highlight the main points. Also, the notation $p_M$ in Algorithm 1 and the notation $p_{M_1}$ and $p_{M_2}$ in Section 4.1 do have some confusion for the understanding. Therefore, we make revisions to clarify Section 4.1. Furthermore, we integrate Algorithm 2 into Algorithm 1 to unify the notation, removing the confusing notation $p_M$ and removing the original Algorithm 2 from the main paper. **We invite the reviewer to refer to point 4 and point 5 in global response . Moreover, we add a schematic diagram in global response pdf to further describe our method. Please refer to point 8 in global response.**\n\nHere we provide additional clarifications to address the reviewer’s confusion. Algorithm 1 wants to emphasize the difference from the traditional MBRL algorithms, namely the extra fine-tuning process (phase 2). $M_1$ denotes the model backed up before training by MLE while $M_2$ denotes the learned model after training by MLE. **Due to the possible performance drop, we devise phase 2 to further fine-tune $M_2$. As for $M_1$, it is used for computing the model-shift term in Eq.(3).**\n\n___\n\n## 2. could the $M_1$ ensembles be also used in the bias term estimate?\nModel bias refers to the error between the learned model and the real environment. According to previous work, the real environment is usually estimated using all models not selected in the ensemble models. Since **the model bias of $p_{M_2}$ is to be estimated in Eq.(3)**, the calculation only involves the $M_2$ ensemble rather than using the $M_1$ ensemble.\n\nHowever, **it may be possible that using the** $M_1$ **ensemble could lead to more accurate estimates, but this is not our key point and requires further research to validate.**\n\n___\n\n## 3. Figure 6: more lr of the MBPO to showcase the advantage\nWe conduct the ablation experiment on more learning rates containing 7e-4, 5e-4, 3e-4, 1e-4 and 1e-3 (the original MBPO) to compare with USB-PO, further strengthening the advantage of USB-PO. **Please refer to Figure 4 in global response pdf.**\n\n___\n\n## 4. Another ablation study is to varying the number of ensembles for the models, and comment on the computational complexity.\nWe conduct the ablation experiment on more ensemble numbers containing 3, 5 and 7. We find that **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 original repo. **We invite the reviewer to refer to Figure 5 and Table 1 in global response pdf.**\n\n___\n\nIf we have addressed the reviewer's concerns, we expect the reviewer to consider improving the score. If the reviewer has any additional questions or comments, we would be happy to discuss them further."}, {"review_id": "wPIlLjox7o", "reviewer": "Reviewer_eAwQ", "summary": "This 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.", "questions": "1. 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.", "limitations": "1. 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.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. 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.", "weaknesses": "1. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to appreciate the reviewer for the recognition and valuable comments! Our specific responses to the questions raised by the reviewer are as follows:\n\n___\n\n## 1. The basis for lines 133-135\nAccording to the previous work MBPO, the return discrepancy scheme refers to the difference between the return under the model and that in the real environment. Under our notation, the return discrepancy is denoted as $V^{\\pi_i|M_i}-V^{\\pi_i}\\_{M_i}$. According to Theorem 1, we make some revisions on line 133-135 to make it clear. **Please refer to point 2 in global response.**\n\n___\n\n## 2. The relationship between Theorem 2 and 3\nSorry, we do have a typo in Theorem 3. We correct the formula in Theorem 3 (Eq.(7)). **Please refer to point 3 in global response.**\n\n___\n\n## 3. The performance curve of the MBRL method CMLO drops much more than what is shown in the original paper.\nWe at first ran a lot of random seeds for CMLO when verifying its performance, and found that in very occasional cases CMLO would have extremely high performance (e.g. ant for 8088, walker2d for 7553). We do not doubt the performance of CMLO, we just think that the authors just select such special cases in their plots under limited runs. **For a fair comparison, we exclude these special cases from our plots.** Moreover, the HalfCheetah and Ant environments are only plotted up to 250K in our experiment, which may seem to have lower performance, but is actually compatible with CMLO, **as we take greater smoothing operation**.\n\n___\n\n## 4. More baseline comparisons\nWe are sorry for being unable to run PDML because the repo released by the author misses the key source-code file, **so we contact the author to get the results in their original paper**. **Please refer to Figure 2 in global response pdf.** USB-PO still exhibits SOTA asymptotic performance and sample efficiency, although not as good as PDML in Humanoid environments. We are not claiming to beat all SOTA algorithms with sophisticated designs as PDML needs extra computational resources to compute policy shift ($ξ_{\\pi_i}$) and to memorize the policy sequence. Although the growth on time reported in the PDML is small compared to the MBPO, **our method USB-PO has the potential to shorten the time consumption, please refer to the time cost in Appendix C.4**. As for the ALM, the training for its classifier may introduce instability, which impairs the performance.\n\n___\n\n## 5. Can you specify how much the time consumption has increased? / Time efficiency needs to be reported.\nAccording to Section 5.2, although USB-PO is a two-stage model training process, **continuing to use the fine-tuned model for the next iteration has the potential to accelerate model convergence** (the fine-tuning process reduces model bias) and then possibly reduce the training time. **Please refer to the time cost we reported compared to MBPO in Appendix C.4**. Training time is reduced on Humanoid, HalfCheetah, and Walker2d.\n\nIn addition, we make some revisions to describe the computational cost. **Please refer to point 7 in global response.**\n\n___\n\n## 6. It would be desirable to compare with more benchmarks, such as DMC and Metaworld, not just MuJoCo.\nSorry, we don't give experimental results on DMC or Metaworld within a limited time for the following three reasons:\n\n· Previous methods including MBPO, M2AC, P2P, BMPO, CMLO, and so on have not been experimented on DMC or Metaworld, so DMC and Metaworld are unfamiliar to us.\n\n· After investigating DMC and Metaworld, we made an attempt and found that DMC or Metaworld is quite different from MuJoCo, so both baseline and USB-PO need to tune the parameters carefully.\n\n· MBRL algorithms generally consume a long time to train once, e.g. almost 3 days for MBPO and 6 days for P2P on HalfCheetah. \n\nThough we don’t give specific results, we set this as future work.\n\n___\n\nWe expect the reviewer could increase the score if we address the issues the reviewer raised. Of course, if there are more questions, we are willing to further discuss them with the reviewer."}]

Author rebuttal:
We are very grateful to all the reviewers for the valuable feedback, helping us improve the paper further! 

We revise the paper and add suggested experiments according to the reviewer’s comments. The detailed revisions are as follows. The additional figures and a table are in the pdf file. 

___

## 1. Revisions of related work
We revise the line 75-87: Due to the character limit, we put the revision of the related work in the reply to the reviewer C4Te.

___

## 2. Revisions of the statement of Theorem 1
We revise the line 133-135 as: Obviously, compared to directly optimizing the return discrepancy of each iteration **[19]**, the performance difference bound chooses to optimize the return discrepancy of two adjacent iterations, **namely** $V^{\pi_2|M_2}-V^{\pi_2}\_{M_2}$ **and** $V^{\pi_1|M_1}-V^{\pi_1}\_{M_1}$ **respectively**, and the expected return variation between these two iterations, **namely** $V^{\pi_2}\_{M_2} - V^{\pi_1}\_{M_1}$, demonstrating better rigorousness.

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

___

## 3. Revisions of Eq.(7) in Theorem 3
We revise the typo as below: 

$V^{\pi_2|M_2} - V^{\pi_1|M_1} \ge \frac{2R\_{max}\gamma}{(1-\gamma)^2}(\epsilon^{\pi_1}\_{M_1} - \epsilon^{\pi_2}\_{M_2} - \epsilon^{M_2}\_{M_1}) - \frac{2R_{max}\epsilon\_\pi}{(1-\gamma)^2}$

___

## 4. Revisions of Section 4.1
We revise the line 114-122 as: 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)**. 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.3, this is not equivalent to the traditional methods of limiting the magnitude of model updates, but rather beneficial to get a performance improvement guarantee.

___

## 5. Revisions of Algorithm pseudo-code
To avoid confusion of notation between $p_M$ and $p_{M_1},p_{M_2}$, we integrate Algorithm 2 into Algorithm 1 and remove Algorithm 2 from the main paper: 

| *Algorithm 1* Meta-Algorithm of the USB-PO Framework |
| :----|
| 1: Initialize the policy $\pi$ and the learned model |
| 2: Initialize the environment replay buffer $\mathcal{D}$ and the model replay buffer $\mathcal{D}_M$ |
| 3: **for** each epoch **do** |
| 4: $\qquad$ Use $\pi$ to interact with the real environment: $\mathcal{D}\leftarrow  \mathcal{D}\cup\{(s,a,r,s')\}$ | 
| 5: $\qquad$ Backup the current learned model for future use and denote this backed-up model as $p_{M_1}$ | 
| 6: $\qquad$ Phase 1: use $\mathcal{D}$ to train the learned model with the supervision of MLE and denote this updated model as $p_{M_2}$ | 
| 7: $\qquad$ Phase 2: use Eq.(3) as optimization objective to further fine-tune $p_{M_2}$ | 
| 8: $\qquad$ Use $p_{M_2}$ to generate the imaginary rollouts: $\mathcal{D}_M\leftarrow  \mathcal{D}_M\cup\{(s_M,a_M,r_M,s_M')\}$ | 
| 9: $\qquad$ Use $\mathcal{D}\cup\mathcal{D}_M$ to train the policy $\pi$ | 
| 10: **end for** |

___

## 6. Revisions of the statement of Eq.(3)
To prevent misunderstanding, we revise the line 171-173 as: **Hence, Eq.(3) can be applied as the optimization objective to get a performance improvement guarantee.**

___


## 7. Revisions of time cost
We add the statement of time cost after line 225: **Computational Cost.  
We report our computational cost compared to MBPO in Appendix C.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.**

___

## 8. An added simple schematic diagram
To better illustrate USB-PO, we add an additional schematic diagram. Please see Figure 1 in global response pdf. We will put this into the appendix later.

___

The above is a summary of the revisions. For questions raised by each reviewer, please refer to the rebuttal to the specific reviewer.

Meta-review: All the reviewers were positive about this paper, except Reviewer VRuX, who found the paper to be presented not very clearly. The AC strongly recommends the authors to take their comments into account while revising the paper, but leans towards acceptance due to the positive response to the contribution made in the paper.

That said, I do personally think that the paper is quite lacking in its evaluation -- while the authors in response to the reviewers and the AC argue that DMC tasks are of no need, I would like to note that experiments on Gym tasks with small improvements were an accepted paradigm in 2018-2019 (when the SAC and MBPO papers came out), and since then several RL papers study a broader spectrum of tasks. In the context of the method prpposed in this paper, I believe that the same approach could be instantiated with Dreamer or image-based MBPO, such as with SLAC, and at least some experiments should be added to the camera-ready version of the paper. 

With regards to the point made by the authors, I agree that while the Gym domains are sufficient to validate the theory, I think viewing this paper as a theory paper primarily puts it below the acceptance threshold. In contrast, if we consider this as an empirical paper, then empirical contributions in and off themselves should be strong enough, which seems to not be the case in the submission.
Decision: Accept (poster)