A Unified Framework for Alternating Offline Model Training and Policy Learning
Authors: Shentao Yang
Shujian Zhang
Yihao Feng
Mingyuan Zhou
Venue: NeurIPS 2022 Accept (2022)
Keywords: Offline Reinforcement Learning
Model-based Reinforcement Learning
Mismatched Model Objectives

Abstract: In offline model-based reinforcement learning (offline MBRL), we learn a dynamic model from historically collected data, and subsequently utilize the learned model and fixed datasets for policy learning, without further interacting with the environment. Offline MBRL algorithms can improve the efficiency and stability of policy learning over the model-free algorithms. However, in most of the existing offline MBRL algorithms, the learning objectives for the dynamic models and the policies are isolated from each other. Such an objective mismatch may lead to inferior performance of the learned agents. In this paper, we address this issue by developing an iterative offline MBRL framework, where we maximize a lower bound of the true expected return, by alternating between dynamic-model training and policy learning. With the proposed unified model-policy learning framework, we achieve competitive performance on a wide range of continuous-control offline reinforcement learning datasets. Source code is released at https://github.com/Shentao-YANG/AMPL_NeurIPS2022.

Though such training paradigm is historically rooted [25, 26), this issue of objective A Unified Framework for Alternating Offline Model Training and Policy Learning
Shentao Yang', Shujian Zhang', Yihao Feng', Mingyuan Zhou'
'The University of Texas at Austin shentao yang@mccombs utexas. edu
"Salesforce Research
szhang190utexas. edu ri com mingyuan. zhou@mccombs utexas edu
Abstract
In offline model-based reinforcement learning (offline MBRL), we learn a dy- namic model from historically collected data, and subsequently utilize the learned model and fixed datasets for policy learning, without further interacting with the environment. Offline MBRL algorithms can improve the efficiency and sta- bility of policy learning over the model-free algorithms. However, in most of the existing offline MBRL algorithms. the learning objectives for the dynamic models and the policies are isolated from each other. Such an objective mis- match may lead to inferior performance of the learned agents. In this paper, we address this issue by developing an iterative offline MBRL framework, where we maximize a lower bound of the true expected return, by alternating between dynamic-model training and policy learning. With the proposed unified model- policy learning framework, we achieve competitive performance on wide range of aaniiuuous-cotree offline reinforcement learning datasets. Source code released at https://gitubb.com/Shetttte YANG/AMPL
/Ati/////////eeeeaoAAA/A///LLLL YANG/AMPPL_Nuu1P222222
Introduction
Offline reinforcement learning (offline RL) [1, 2]. where the agents are trained from static and pre-collected datasets, avoids direct interactions with the underlying real-world environment during the learning process. Unlike traditional online RL, whose success largely depends on simulator- -nde-- [e.g.. 3, 4], offline RL enables training policies for real-world applications, where is infeasible or even risky to collect online experimental data, such as robotics, advertisement, dialog systems [e.g., 5-8]. Though promising, remains challenging to train agents under the offline setting, due to the discrepancy between the distribution of the offline data and the state-action distribution induced by the current learning policy. With such a discrepancy, directly transferring standard online off- off-policy RL methods [e.g., 9-11] the offline setting tends to be problematic [12, 13], especially when the offline data cannot sufficiently cover the state-action space [14]. To tackle this issue in offline RL, recent works [e.g.. 15-17] propose to approximate the policy-induced state-action distribution by leveraging a learned dynamic model to draw imaginary rollouts. These additional synthetic rollouts help mitigate the distributional discrepancy and stabilize the policy-learing algorithms under the offline setting.
Most of the prior offline model-based RL (MBRL) methods [e.g.. 16, 18-21]. however, first pretrain a one-step forward dynamic model via maximum likelihood estimation (MLE) on the offline dataset, and then use the learned model to train the policy, without further improving the dynamic model during the policy learning process. As a result, the objective function used for model training (e.g.. MLE) and the objective of model utilization trained to be "simply a mimic the world." but is used to improve the performance of the learned policy
unrelated with each other. Specifically, the model is
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

mismatch in the model training and model utilization has been identified as problematic in recent works [23, 24]. In offline MBRL, this issue exacerbated. since the learned model can hardly be globally accurate, due to the limited amount of offline data and the complexity of the control tasks.
Motivated by the objective- mismatch issue, develop an iterative offline MBRL method, alternating between training the dynamic model and the policy to maximize lower bound of the true expected return. This lower bound, leading to a weighted MLE objective for the dynamic-model training, is relaxed to a tractable regularized objective for the policy learning. To train the dynamic model by the proposed objective, we need to estimate the marginal importance weights (MIW) between the offline-dee data distribution and the stationary state-action distribution of the current policy [27, 28]. This estimation tends to be unstable by standard approaches [e.g., 29, 30], which require saddle-poin optimization. Instead, we propose a simple yet stable tixed-point-siee which can be directly incorporated into our alternating training framework. With these considerations, our method, offline Alternating Model-Policy Learning (AMPI.), performs competitively on wide range of continuous-contro. offline RI. datasets in the D4RL benchmark [31]. These empirical results and ablation study show the efficacy of our proposed algorithmic designs.
method for MIW estimation,
2 Background
Markov decision process and offline RL. A Markov decision process (MDP) is denoted by M
(S, A, P.r, y, Mu), where & is the state space, A the action space, P (s' as.a): $ x A + [0. 1] the environmental dynamic, r(s, a) 8 x A - [-Tmany T'mar] the reward function, E [0, the discount factor, and (o(s): : - [0, 1] the initial state-distribution.
For any policy π (a s), we denote its state-action distribution at timestep & 2 as (s, a) & Pr (st - 8, 0r
80 ~ poo ae T,, 8141 ~ P, 14 2 0). The (discounted) stationary state-action distribution of is denoted di,, a
7) Ex= 7' (s.
Denote (=
(s, a) ET Er y'r (sr at) su s, ag a] as the action- value function of policy under the dynamic P. The goal of RL is to find policy æ maximizing the expected return
( P) - (1 Bn (a) [ (s, a)] ((a))adf [n(s,a)].
In offline RL, the policy π and critic Q. are typically approximated by parametric functions T and Qe, respectively, with parameters Φ and O. The critic Qe is trained by the Bellman backup
arg ming ((e.o,,e")> [Qo(s..)) a) (r(s, a) + (-
a*) (Qo (s', 66',0 =)
(2)
where Qe is the target network [12, 13]. The actor To is trained in the policy improvement step by arg max
a (11>) [Qo (s,a),,
(3)
where Deny denotes the offline dataset drawn from dh.. [2, 32], with ne being the behavior policy.
Offline model-based RL. In offline model-based RL algorithms, the true environmental dynamic P is typically approximated by parametric function P (s') s, a) in some function class P. With
the offline dataset Denv. P is trained via the MLE [15, 16, 18] as
arg mrg PeP E ((s,,,") log P (s' 8, s,
(4)
Similarly, the reward function can be approximated by parametric model if assumed unknown. With P and 7, the true MDP M can be approximated by M = A. 7. Ho). We further define d", (s. a) the stationary state-action distribution induced by on P* (or MDP M), and d a) (s,
as that on the learned dynamic P (or MDP M). We approximate d by simulating T on M for short horizon h starting from state E Deuv. as in prior work [e.g.. 16. 18, 19, 21]. The resulting transitions are stored in replay buffer Dinodel, constructed similar to the off off-policy RL [33, 9]. To better approximate d' sampling from Denv Eqs. (2) and (3) is commonly replaced by sampling from the augmented dataset D f'Denv + (1 - J)Droudel J E [0, 1), denoting sampling from Donv and D modde) with probabilities and f. respectively. We follow Yu al. [18] to use f = 0.5.
2

---

3 Offline alternating mooecy learning
Our goal is to derive the objectives for both dynamic-model training and policy learning from principled perspective. A natural idea is to build a tractable lower bound for P"), the expected return of the policy under the true dynamic P and then alternate between training the policy and the dynamic model P maximize this lower bound. Indeed, we can construct a lower bound as
(5) where J(x, P) the expected return of policy a under the learned model P. From the right hand side (RHS) of Eq. (5), if the policy evaluation error [J P*) J(T, P)] is small, J P) will be good proxy for the true expected return J (TT,, P*). We can empirically estimate J P) via P and
(T,, 2 [T, P) [J P*) - [ P)).
Further, if tractable upper bound for [J(xT, P*) [TT, P) can be constructed, it can serve unified training objective for both dynamic model P and policy и. We can then alternate between optimizing the dynamic model P and the policy æ to maximize the lower bound of J(T, , P*), i.e. simultaneously minimizing the upper bound of the evaluation error [J( P*) J (TT, P)J. This gives algorithm for model and policy learning.
us an iterative, muximiization-maxiiiaaaaa
The following theorem indicates a tractable upper bound for [J(T,, P*) [Tz, P)}, which can be subsequently relaxed for model training and policy learning.
Theorem 1. Let P be the true dynamic and P be the approximate dynamic model. Suppose the reward function (r(8,a) a)l S Tiaxx then have
J(Z,, J(T P) s with DP( P) A Be [( (P*(d' 80) P(d (,,)) 1+)]
V2(1- Da(M,, P).
7. -mmax
where To is the behavior policy, d is the offline-data distribution, and w(s a) (s,a)
d (s.a) is
the marginal importance weight (MIW) between the offline-dan distribution and the stationary state-action distribution of the policy π [27, 29].
Detailed proof of Theorem can be found in Appendix B.2.
The KL term in (P P) indicates the following two principles for model and policy learning:
For the dynamic model, s' ~ P(- s, a) should be close the true next state 3' (*((1 a). with (s, a) pairs drawn from the stationary state-action distribution of the policy æ. Since we cannot directly draw samples from df we reweight the offline data with w(s a). This leads to a weighted KI. minimization objective for the model training.
* For the policy म, the KI. term indicates a regularization term, that the tuple (s',a') from the joint conditional distribution P(s/ s, a))(aa) |&') should be close to the tuple (8',ā") a') from p (3') 8, (a) (ā' 8"). ,²) is simply sample from the offline dataset.
Based on the above observations, we can fixed and train the dynamic model P by minimizing D=(P", P) w.r.f P. Similarly, we can fix the dynamic model P and learn a policy æ to maximize the lower bound of J(T,, P*). This alternating training scheme provides a unified approach for model and policy learning. In the following sections, discuss how optimize the dynamic model P. the policy ग. and the MIW w under our alternating training farmework.
3.1 Dynamic model training
Expanding the KL term in (x(P", we have
Dx (P",PP ,)
Bis.a,s a'))aa [w(s, a) (log P* (s' 18, a)) + log mo(a' (s') log x (a'Is'))
)adg [u(s, a) log P(s' 18,
3

---

where the tuple (s, a, a', a') is simply two consecutive state-action pairs in the offline dataset. Further,
if the policy is fixed, the term ① is constant w.r.t. P. Thus, given the MIW wi can optimize P by minimizing the following loss
{(P) - - (s,a,,))) (sss a) log P(s' s,a)
(6)
which is an MLE objective weighted by w(s a). We discuss how to estimate w(8, a) in Section 3.3.
3.2 Policy learning
The lower bound for Jmm P*) implied by Theorem is
J [ P)
Da <2(1->)
T-Tmax
where J(H,,) P) can be estimated via the action-value function similar to standard offline MBRL algorithms [e.g., 16, 18, 19]. Thus, when the dynamic model P is fixed, the main difficulty is to , P) for the policy и.
estimate the regularizer (=
When the policy is Gaussian, direct estimation of (P P) is possible. Empirically, however, it is helpful learn the policy и in the class of implicit distribution. which richer distribution class and can better maximize the action-value function. Specifically, given noise distribution Pa (z), action a = TT 2) with 2 Pa (+), where कΦ deterministic network.
Unfortunately, we can not directly estimate the KL term Dx(P", P) if v is an implicit policy, since we can only draw samples from but the density is unknown. A potential solution is to use the dual representation of KL. divergence KL(p || g) = supr Eupy Iog((E, [34], which can be estimated with samples from the distributions p and q. However, the exponential function therein makes the estimation unstable in practice [30]. We instead use the dual representation of the Jensen-Shannon divergence (JSD) to approximate (p(P*, P), which can be approximately minimized using the GAN structure [35, 36). Our framework can thus utilize the many stabilization techniques developed in the GAN community (Appendix E.2.1).
Besides, remove the MIW w(s, a) during the policy training since we do not observe its empirical benefits, which will be discussed in Section 4.2. Further applying the replacement of KL with JSD and ignoring the V: for numerical stability, we get approximated
regularization for policy m: DP(P" P) A JSD (s" (.a))a(d (s, || P(s' s, a,a) ss))d"", (s) s)) (8)
Informally speaking. Eq. (8) regularizes the imaginary rollouts of on P towards state-action pairs from the offline dataset. Intuitively, D, , P)) is more effective regularizer for policy training than the original Dx P). since (P* P) regularizes action choices at both and s'. Appendix B.3
discusses how we move from (P", P) (P", P) in detail.
3.3 Marginal importance weight training
number of methods have been recently proposed to estimate the marginal importance weight w [27, 29, 30]. These methods typically require solving complex saddle-poin optimization, casting doubts on their training stability especially when combined with policy learning on continuous-control offline MBRL problems. this section, we mimic the Bellman backup derive fixed-point-style method for estimating the MIW.
Denote the true MIW as w (s, d (s.a) (o,0) we have (s) = di,, (s,a)) a). Expanding
the RHS. vs',a', a', (s'.
(8' (d''
Σ n(a (s) (0001/118,88) a,) (a,))=(,, a). a)) (8. 7 )=(((( ( (9)
P.a
The derivation is deferred to Appendix B.4. Therefore, a "Bellman equation" for w(s' cisss'd) a²) is w(s', (a)) =Tw(s', (u(ss'aa aa),
Tw(s', a²) e 7 >E., (a' s')P'((s (8' ()) a) (s," (aa))+ + (1 - 7) 1s')
dp d

---

Here T can be viewed as the "Bellman operator" for w. The update iterate defined by T has the following convergence property, which proved in Appendix B.5.
Proposition 2. On finite state-action space, the current policy is close the behavior policy nb. then the iterate for w defined by T converges geometrically.
The assumption that π is close to T coincides with the regularization term in the policy-learning objective discussed in Section 3.2.
Unfortumately, the RHS of Eq. (9) is not estimable since we do not know the density values therein. We therefore multiply both sides of Eq. (9) by some test function and subsequently sum over (s', a') on both sides to get tractable objective that only requires samples from the offline dataset and the initial vitt-distrinution (4g- the difference between the left-hand side (LHS) and the RHS of Eq. (9). A potential choice is the action-value function the policy иT, due to some primal-dual relationship between the stationary state-action censity-(ration (s",,,, a) or (s, a)) and the action- value function [37-39]. A detailed discussion on the choice of the test function is provided in Appendix C.
desired to choose test function that can better distinguish
Practically use Of as the test function. Note that multiplying both sides of Eq. (9) by the same of does not undermine the convergence property, under mild conditions on Q. Mimicking the Bellman backup to sum over (s', a') on both sides, with the notation d d
(8, a,,,,ss) d (8, a) P* (s(8' ( a).
k(ww)
kg(w")
of ()) = B (na.d [w' dff bir
( (s,a) of + ? mmpatt) a
[of
(10)
d' 69)
Thus for given w, we can optimize by minimizing the difference between the RHS and the LHS of Eq. (10). For training stability, we use a target network w' w'(s, a) for the RHS, and the final objective for learning w
- l2 (2 (w/))²
(11)
where the target network w'(s, a) is soft-updated after each gradient step, motivated by Qo; and Tay-
Our proposed training method is closely related to VPM [40]. By using the MIW w(s, a) itself as the test function, VPM leverages the variational power iteration to train MIW iteratively. Instead, our approach uses the current action-valus function as the test function, motivated by the primal-dual relationship between the MIW and the action-value function in off- -fo-colcc evaluation. We compare the empirical performance of several alternative approaches in Section 4.2 and in Table 3 of Appendix A.
3.4 Practical implementation
In this section we briefly discuss some implementation details of our offline Alternating Model-Policy Learning (AMPL) method, whose main steps are in Algorithm 1. Further details are in Appendix E.
Dynamic model training. We adopt common practice in offline MBRL (e.g., 16, 18) to an ensemble of Gaussian probabilistic networks P(- s, a) and f(s, a) parameterize the stochastic transition and reward. We initialize the dynamic model by standard MLE training, and periodically update the model by minimizing Eq. (6).
Critic training.
We use the conservative target in the offline RI. literature [e.g.. 12, 13]:
@ (s. a) _ r(s.a) a) + TEu (-1 sn) [cming=1,2 Q0; (s', d') + (1 - c) max Qoy (s', aa)], (12)
where we set c = 0.75. With mini-batch B sampled from the augmented dataset D, both critic networks are trained as
Vj = 1,2, arg mino, [8]
1 E(s,)E Huber(Qo, (s, (s,a).((s.a)),
(13)
where the Huber loss Huber(-) is used in lieu of the classical MSE for training stability [41].
Estimating (P' P). In D= P), using the notations of GAN, we denote the sample from the left distribution of ISD in Eq. (8) by Birue (i.e,, "true" sample), and the sample from the right distribution of JSD by Brake (Le., "fake" sample).
5

---

Algorithm Main steps of AMPL.
Initialize: Dynamic model P and f. policy ne. crities Qo, and Qo discriminator Dep, MIW w.
Initialize P and via the MLE (Eq. (4)). for iteration E (1,..., sotal_iterations do then Estimate w(s,a) a) via Eq. (11); train P and by weighted MLE (Eqq (6)) with w(s, a).
if iteration % rondel_reerrin_periioe
end if
Rollout synthetic data with TΦ. P and f. and add the data to D model.
Sample mini-batch B D = =f + (1 ff)
Optimize Qe Qe via Eqs. (12) (13).
Train Dy to maximize Eq. (16).
Optimize T by Eq. (17).
end for
The "true" sample Berue consists of samples from Dony. The "fake" sample Biake formed by first sampling D, followed by a ~ TT(: |8), s' P(-1s,a),ad a' ~ aT( [-([11>) s'). Coneretely, the "fake" sample Biake, generator loss Lg)()), and the discriminator loss (1bb can be described as
(8, (id, a)]. (14) Cg(Ø)
Brake -
E(sr)) [log (1 Da (s, a))] (15)
Braten
(D(()) Bernel E(s.))BB...
[log Dwp a)] + Biskel E(s,a)BBatt [log (1 - Dw a))l, (16)
where (Ø) is the empirical policy-learning regularizer implied by (P' P).
Policy training. Since the constant multiplier in Eq. (s) unknown, treat it hyperparameter. Adding the regularization term of policy training, our policy optimization objective is
arg ming -
B > (eB [min,-1.3 Qo, (s, a)] + (g)
(17)
where the regularization coefficient A X' X'/Qaug with soft-updated Qaeg similar to Fujimoto and Gu [42]. We use X' = 10 across all datasets in our experiments (Section 4).
4
Experiments
In this section, we first evaluate our AMPL on continuous-control offline-RL datasets in Section 4.1. Then we conduct ablation study on Section 4.2 to understand the efficacy of some proposed designs.
4.1 Continuous-controd results
Our experiments are conducted on diverse set of datasets in the D4RL benchmark [31], ranging across the Gym-Mojoco, Maze2D, and Adroit domains therein. We use the latest version of the datasets, i.e., "v2" version for the Gym-Mojoco domain and "v1" version for the Maze2D and Adroit domains. Details of our dataset choice are discussed in Appendix E.2.
We compare our AMPL with two DICE methods - AlgaeDICE [43] and OptiDICE [44] - that directly utilize MIW to improve policy learning. We further consider three state-of-the-art (SOTA) offline model-free RI. algorithms: CQL [45], FisherBRC [46], and TD3+BC [42]: and three SOTA offline MBRL methods: MOPO [16], COMBO [18], and WMOPO [47]. Experimental details and hyperparameter settings are discussed in Appendix E.2. We run baseline methods using the official implementation under the recommended hyperparameters, except for AlgaeDICE, for which we use the offline version provided by Fu et al. [31]. Table I shows the mean and standard deviation of the results of AMPL and baselines over five random seeds.
As shown in Table 1, our AMPL performs comparably well and is relatively stable across the sixteen tested datasets. In particular, AMPL generally performs better than the baselines on the Maze2D and Adroit datasets. The Maze2D and Adroit datasets are considered as more challenging than the MuloCo datasets [31], since the Maze2D datasets are collected by non-Markovian policies, and the amount of data in the high-dimensional Adroit datasets is limited.

---

Table I: Normalized returns for experiments the D4RL tasks. Mean and standard deviation across five random seeds are reported. High average score and low average rank are desirable. We bold the best result over all methods and underline the best of the model-based methods if different. Here, "hcheetah" denotes "halfcheetah," "med" denotes "medium," "rep" denotes "replay," and "exp" denotes "expert."
Task Name
AlgaDICE
OpiDDIE
CQL
FisherBRC
TD3-BC
MOPO
COMBO
WMOPO
AMPL
maze la-agg
-22 06
101.7 +: 50.7
1.5 d
0.9 = 6.4
107.1 + 45.9
40.5 2.6
138.5 + 82.2
1.8
180.0 =: 39.3
mazeddmas
25 1 142
119.4 1 52.5
6.3 1
16.1 1. 30.3
61.1 45.5
12.5.183 18.3
103.9 1 42.1
12.7 1. 38.3
1072 4500
45.3 +
114.0 39.7
37.5
366 16.4
38.6 + 14.4
4544
112.1
56.8
-11.5
4.1.5.33..1
55.8
hcheetab- med
0,9 a: 0.7
42.0 a: 35
48.6 à: 0.2
47.8 a 0.3
48.0 4: 0.3
69.2 3.4
73.0 +: 3.6
72.0 1: 4.7
51.7 +: 0.4
walker2ddmed
83.0
.0.1
0.5
83.1
hopper-med
0.9.2.22
47.5
68.0 6.0
94.2 1. 4.5
59.1 1
41.8
22.6
99.7
58,9
medeeta
2.3
40.5 = 33
38.1 + 18.6
34.9 + 18.2
44.5 + 0.6
62.7 = 7.5
66.0 + 1.8
66.4 +: 4.9
44.6 = 0.7
walker22d med--rp
21.8
6.7
$1.5
hopper-mode med
1.5
28.1 à: 12.4
96.0 2:
94.4 + 1.9
58.9 d: 19.7
84.8 30,0
53.6 à: 29.5
93.8 4: 8.5
91.1 + 9.5
hehectah-medeee exp
2.9
66.7 2.28.8
53.8 1. 13.6
94.5 0.6
91.4 5.22
61.0 20.8
52.5 52.0
68.8 1. 37.5
90.2 1. 2.0
walker2ddmeddessp
-.1
79.4
110.3
110,4 0.6
1.1
hupper-mod-esp
2.2 1.9
$2.6 9.3
79.3 1: 20.5
110.0 4.0
98.7 1. 9.6
15.6-
63.7 1. 51.9
62.6 1. 26.4
76.0 1: 15.7
pen-human
-3.0.1 0.8
-0.9 1: 3.1
-1.5 2.8
-0.1 1:300
10.
-1.6:::2.
5.9 -1::
-2.1 1: 1.4
20.6 10.7
pen-cloned
-2.6 1.3
-0.8 1 5.1
39,5 23.6
1.
6.3 4.6
5.2 1. 10.2
23.2
-3.0 2.4
57.4
pen-exp
2.4
0.4 7.3
119.3 +: 158
2.2 22 1.1
104.2 + 40.6
35.4 2009
57.1 4:4..6
10.4 +: 20.9
138.3 :: 6.6
door-eap
15.
-0.3
Avenge Score
-1.4
39.6
49.7
61.7
26.9
52.1
44.2
83.8
Average Rank
8.5
5.5
4.1
4.3
3.9
6.4
4.8
2.6
On the Maze2D datasets, the behavioral-cloning (BC) style algorithms, such as FisherBRC, are likely to fail, since these methods use Markovian policy to approximate the non-Markovian behavior policy, Meanwhile, OptiDICE performs relatively well on the Maze2D datasets, mainly because it uses Gaussian mixtsre policy for BC to alleviate this approximation difficulty. This advanced BC policy still offers little help on the nigher-dimensional Adroit datasets, due the challenge of learning the behavior policy. Similarly. TD3+BC performs well on the relatively lower-dimensional Maze2D datasets, but does not work well on the Adroit domain, since the action-space MSE-regularizer in TD3+BC may be insufficient for the high-dimensional tasks. WMOPO also shows performance drop in these two challenging task domains. The environmental dynamics on these two domains are complex, and thus may not be accurately learned to support WMOPO's 's full-trajectory model rollouts, which are used to estimate its MIW. By contrast, our AMPL does not require BC or long model-rollouts, leading to stable performance across task domains. Moreover, AMPL shows generally better performance than MOPO and COMBO, which use fixed pretrained dynamic models without considering the mismatched model objectives.
Compared with the DICE-based methods AlgaeDICE and OptiDICE, our AMPL shows generally better performance. This may indicate that maximizing a lower bound of the true expected return can be a more effective framework than explicit nationary-diitributton correction, since maximizing the lower bound is more directly related to RL's goal of maximizing the policy performance.
200
150
150
100
100
50
50
50
0
0
GHOOKCE DEWDICE WPR
MW
VPM
GenDICE DACE
GeeDICE DE
WPR
Maze2D
(b) MuloCo
(c) Adroit
Figure 1: Scores of each method in the ablation study (Section 4.2) on each domain of datasets. The faded
bars show the scores each datasets within the stated domain, averaged over five seeds, where each color corresponds to dataset. The highlighted har shows the average score the entire domain. Lahel "Main" refers the results of our main method in Section 4.1. Detailed numbers are on Table
4.2 Ablation study
(a): Does the wiW-wtighted model (re (re)training help the performance?
7

---

VPM
GenDICE
DuaIDICE
Main
-10
-5
-20
-10
400
900
(a) VPM
(b) GenDICF
10
0.0
-15
100
300
400
500
800
800
900
100
200
400
500
600
700
800
900
(c) DuaIDICE
(d) Our method ("Main""
Figure 2: Distribution plots of log (MIW) of the entire dataset by method ("Main") and the three alternatives dataset during the training process. dash lines represent the quartiles of the distribution. Recall that the MIW w is retrained every 100 epochs. The r-axis represents the number of epochs {100, 200, 900}. and the y-axis the log (MIW) values. For clarity, we use the normalized MIWs whose
Section 4.2 (b). on malkkerd--meiim--te
the entire dataset is onc. We random seed two for this plot. Other random seeds have similar patterns. We note the different scales in y y-axis of these four plots.
To verify the effectiveness of our MIW-weighted model (re)training scheme, we compare our AMPL with its variant training the model only the beginning using MLE, i.e., No Weights (dubbed as NW). As shown in Figure 1, the performance of the NW variant is generally worse than the main method (dubbed as Main) on all three domains. The performance difference is especially significant on the Maze2D domain. In this domain, the policy is required to "stitch together collected subtrajectories to find the shortest path the evaluation goal" [31]. As result, the state-action distribution induced by good policy should be different from the distribution of the offline-data. Table in Appendix confirms the benefit of MIW-weighted model (re (re)training. Indeed, on the majority of datasets, the NW variant not only shows worse performance, but also has larger standard deviation relative to the mean score.
(b): How does the algorithm perform if we change our MIW estimation method to other approaches?
We compare our AMPL with its three variants where the MIW is instead estimated by the Variational Power Method [40], GENeralized stationary DIstribution Correction Estimation [30], and Dual stationary DIstribution Correction Estimation [29]. Three variants are simply dubbed VPM, GenDICÉ, and DuaIDICE. For numerical stability, the estimated MIW from these three methods is clipped into (10-5 500). Implementation details are provided in Appendix E.3.
As shown in Figure 1 and Table 2, the variants with these three alternative MIW estimation methods generally perform worse than our approach. A plausible explanation is that these methods can be unstable to provide good MIW estimates for the model training. Figure shows the distribution plots of log (MIW) of the entire dataset estimated by method and by the variants using the three alternative MIW estimation methods, on the wwwlker2d-medium-rrpee medium-replay" dataset. We train the MIW every 100 epochs and ignore the MIW initialization from the plots. Notice that the MIWs obtained by these four methods have very different scales, which indicates the instability the alternative methods compared with ours.
As shown on Figure 2b, the MIW' estimated by the GenDICE variant degenerate, in a sense that most of the MIWs are as O at the begining and are 1 later on, with long tails on the MIW distributions. This may explain the relatively bad performance of GenDICE. Compared with GenDICE, DuaIDICE provides more diverse MIW estimates, which leads to its relatively better performance. Unfortunately,

---

Figure 2c shows that, for the DuaIDICE variant, the distribution of MIWs gradually concentrates on very small and very large values, which indicates the degeneration of the MIWs. The VPM variant performs the best among these three alternatives, mainly because it provides relatively better MIW estimates the second half of the training process. However, Figure 2a shows that the MIWs from the VPM variant are poorly- poorl-diitributte at the beginning. though relatively better later on. Our method provides well-behaved MIW estimates in the whole training process, which may explain our better result. As shown by Figure 2d, the MIWs from our method are well-shaped and concentrate around the mean 1 over the entire training process, In general, the well-shaping of the estimated MIW is important, since the Cramer Raa lower bound of the mean- mean-square- square-erroo for OPE is related to the square of the density ratio [38, 48].
Besides, on the Maze2D domain, these three variants using the alternative MIW estimation methods perform generally better than the NW variant discussed in Question (a), aligning with the benefit of the MIW-weighted model (re trraiingg scheme. Generally speaking, incorporating the MIW can help model training in offline MBRL.
(c): What is the performance of a weighted regularizer for policy learning implied by Theorem 1?
We compare our AMPL with its variant where the policy regularizer is weighted by the MIW w(s, a), as suggested by Theorem 1. This Weighted Policy Regularizer variant is dubbed as WPR. As shown in both Figure 1 and Table 2, the WPR variant underperforms
main method. This is because when we estimate the regularization term in WPR, we incorporate weights into the minimax optimization of the policy and the discriminator, which may bring additional instability.
Table 3 of Appendix provides additional ablation study on the performance of several alternatives.
5 Related work
Offline MBRL. Most of the existing offline MBRL works focus on policy learning under a given model trained by MLE. These works typically constrain the learned policy avoid visiting regions where the discrepancy between the true and the learned dynamic is large, thus reducing the policy evaluation error of using the learned dynamic [16, 19, 49, 50]. Besides, some recent works [51, 52] also adopt a GAN-style stochastic policy for its flexibility. Rather than using fixed MLE-trained model, we derive an objective that trains both the policy and the dynamic model toward maximizing a lower bound of true expected return (simultaneously minimizing the policy evaluation error [r(z,, P*) JTT, P).. Several recent works also propose to enhance model training. such as training a reverse dynamic model to encourage conservative model- imaginations [53], learning balanced state-action representation for the model to mitigate the distributional discrepancy [54, 55], and using advanced dynamic- model architecture of the Traa-ffaanfrmerr [56-59] to achieve accurate predictions [60]. These methods are orthogonal to and may further improve the performance.
Objective mismatch in MBRL. Our AMPL is related to prior works on online and offline MBRL that mitigate the mismatched model objectives. In online (off-policy) MBRL, Lambert et al. [23] identify the mismatched objectives between the MLE model-training and the model's 's usage of improving the control performance. This paper, however, only proposes an impractical weighted learning loss that requires the optimal trajectory. Eysenbach et al. [24] design an objective to jointly optimize the policy and the model in online MBRL. Voelcker et al. [61] propose per-sample diagonal scaling matrix for training a deterministic model in online RL, which improves on Farahmand et al. [62] and Farahmand [63]. In offline MBRL, D' D'Oro et al. [64] propose weighting scheme to learn the transition model. This method requires the hehavior policy and estimates the trajectory-wise importance sampling weight, which is known suffer from the "curse of horizon" [27]. Thus, this method may not scale to complex high-dimensional offline RL tasks. Voloshin et al. [65] propose minimax objective for model learning, where the model is optimized over function classes of the MIW and the state-value function. Most similar to our work, Hishinuma and Senda [47] also use a Meightettted MLE objective for model training. They use different version of the MIW estimated via multiple full-trajectory rollouts in the learned model. This strategy may suffer from the inaccuracy of the learned model (66-68) and can be consuiin.. In this paper, we develop fixed point style method to train the MIW that only requires samples from the offline dataset and the initial distibiisti which can be more stable and efficient.
10 in the log (MIW) plots.
9

---

Off-policy evaluation (OPE). OPE [29, 30, 48, 69-74] is a well-studied problem with several lines of research. natural way is trajectory importance sampling (IS) [75]. This, however, suffers from its variance, which can grow exponentially with the horizon length [27, 76]. Marginal importance sampling methods [27] are developed to address this problem by estimating the IS ratio between two stationary distributions, leading to estimators whose variance are polynomial w.r.f horizon [77]. Though prior works can also be used to estimate the MIW, they may not be directly applicable to our framework, since they either require knowing the behavior policy [27, 37], need to assume the MIW lying the space RKHS [28]. Another approach to estimate the MIW is through saddle-poine optimization [29, 30, 39, 78]. which suffer from its complex optimization problem. In this paper, we derive an poiet--oit--tyle and sehavior-agnostee MIW estimation method, which is simple to train and does not require additional assumptions.
6 Conclusion
In this paper, we are motivated by the mismatched model objectives in offline MBRL to design an iterative algorithm that alternates between model training and policy optimization. Both the model and the policy are trained to maximize a lower bound of the true expected return. The proposed new algorithm performs competitively with several SOTA baselines.
Limitation. Our current approach requires training two additional networks (discriminator D& and MIW w), and the model is (re)trained periodically mitigate the objective mismatch, which brings additional computational cost in order to obtain better empirical results.
Future work. We plan to conduct theoretical analysis of our proposed framework, and investigate more efficient ways to unify model training and policy learning.
Acknowledgments
S. Yang. S. Zhang, and M. Zhou acknowledge the support of NSF IIS 1812699 and 2212418, and the Texas Advanced Computing Center (TACC) for providing HPC resources that have contributed to the research results reported within this paper.
References
[1] Sascha Lange, Thomas Gabel, and Martin Riedmiller. Batch Reinforcement Learning, pages 45-73. Springer Berlin Heidelberg, Berlin, Heidelberg, 2012. ISBN 978-3-642-27645-3, doi: 10. 10.10077978-33-62-2265----
[2] Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint arXiv:2005.01643, 2020.
[3] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A. Rusu, Joel Veness, Marc G. Bellemare, Alex Graves, Martin Riedmiller, Andreas K. Fidjeland, Georg Ostrovski, Stig Petersen, Charles Beattie, Amir Sadik, Ioannis Antonoglou, Helen King, Dharshan Kumaran, Daan Wierstra, Shane Legg. and Demis Hassabis. Human-level control through deep rein- forcement learning. Nature, 518(7540):529-533 February 2015. ISSN 00280836.
[4] David Silver, Aja Huang, Chris Maddison, Arthur Guez, Laurent Sifre, George van den Driessche, Julian Schrittwieser, Ioannis Antonoglou, Veda Panneershelvam, Marc Lanctot, Sander Dieleman, Dominik Grewe, John Nham, Nal Kalchbrenner, Ilya Sutskever, Timothy Lillicrap, Madeleine Leach, Koray Kavukcuoglu, Thore Graepel, and Demis Hassabis. Master- ing the Game of Go with Deep Neural Networks and Tree Search. Nature, 529(7587):484-449, January 2016. doi: 10. 1038/nature16996
[5] Lihong Li, Wei Chu, John Langford, and Xuanhui Wang. Unbiased offline evaluation of contextual-banddi-buusse news article recommendation algorithms. In Proceedings of the 4th
International Conference on Web Search and Data Mining (WSDM), pages 297-306, 2011.
[6] Ekim Yurtsever, Jacob Lambert, Alexander Carballo, and K. Takeda. Survey of Autonomous Driving: Common Practices and Emerging Technologies. IEEE Access, 8:58443-58469, 2020.
10

---

[7] Youngsoo Jang, Jongmin Lee, and Kee-Eung Kim. GPT-critic: Offline reinforcement learning for end-to-end task-oriented dialogue systems. In International Conference on Learning Representations, 2022.
[8] Cenying Yang, Yihao Feng, and Andrew Whinston. Dynamic pricing and information dis- closure for fresh produce: An artificial intelligence approach. Production and Operations Munagement, 31(1):155-171, 2022.
[9] T. Lillicrap, Jonathan J. Hunt, A. Pritzel, N. Heess, T. Erez, Yuval Tassa, D. Silver, and Daan Wierstra. Continuous Control with Deep Reinforcement Learning. CoRR, 1509.02971, 2016.
[10] Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft Actor- Off- Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor. In Jennifer Dy and Andreas Krause, editors, Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pages 1861-1870. PMLR, 10-15 Jul 2018.
[11] Scott Fujimoto, Herke van Hoof, and David Meger. Addressing Function Approximation Error in Actor-Critic Methods. In Jennifer Dy and Andreas Krause, editors, Proceedings of the 35th International Conference Machine Learning, volume 80 of Proceedings of Machine Learning Research, pages 1587-1596. PMLR, 10-15 Jul 2018.
[12] Scott Fujimoto, David Meger, and Doina Precup. Off-Policy Deep Reinforcement Learning without Exploration. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pages 2052-2066 2062. PMLR, 09-15 Jun 2019.
[13] Aviral Kumar, Justin Fu, Matthew Soh, George Tucker, and Sergey Levine. Stabilizing Off- Policy Q-Learning via Bootstrapping Error Reduction. In Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
[14] Rishabh Agarwal, Dale Schuurmans, and Mohammad Norouzi. An optimistic perspective on offline reinforcement learning. In International Conference on Machine Learning, pages 104-114. PMLR, 2020.
[15] Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine. When to Trust Your Model: Model-Based Policy Optimization. ArXiv, abs/1996.0882 2019.
[16] Tianhe Yu, Garrett Thomas, Lantao Yu, Stefano Ermon, James Y. Zou, Sergey Levine, Chelsea Finn, and Tengyu Ma. MOPO: Model-based Offline Policy Optimization. ArXiv, abs/2005.13239, 2020.
[17] Aravind Rajeswaran, Igor Mordatch, and Vikash Kumar. A Game Theoretic Framework for Model Based Reinforcement Learning. In International Conference on Machine Learning. 2020.
[18] Tianhe Yu, Aviral Kumar, Rafael Rafailov, Aravind Rajeswaran, Sergey Levine, and Chelsea Finn. COMBO: Conservative Offline Model-Based Policy Optimization. ArXiv, abs/2102.088633, 2021.
[19] Rahul Kidambi, Aravind Rajeswaran, Praneeth Netrapalli, and T. Joachims. MOReL Model- Based Offline Reinforcement Learning. ArXiv, 0as/2005.055555 2020.
[20] Phillip Swazinna, Steffen Udluft, and Thomas A. Runkler. Overcoming Model Bias for Robust Offline Deep Reinforcement Learning. Eng. Appl. Artif. Intell., 104:104366. 2021.
[21] Catherine Cang, Aravind Rajeswaran, P. Abbeel, and M. Laskin. Behavioral Priors and Dynamics Models: Improving Performance and Domain Transfer in Offline RL. ArXiv, Abs/2106.09119. 2021.
[22] Yuping Luo, Huazhe Xu, Yuanzhi Li, Yuandong Tian, Trevor Darrell, and Tengyu Ma. Algorith- mic Framework for Model-based Deep Reinforcement Learning with Theoretical Guarantees. In International Conference Learning Representations, 2019.

---

[23] Nathan Lambert, Brandon Amos, Omry Yadan, and Roberto Calandra. Objective mismatch in model-based reinforcement learning. arXiv preprint 20Xi::2002.04523,, 2020.
[24] Benjamin Eysenbach, Alexander Khazatsky, Sergey Levine, and Ruslan Salakhutdinov. Mismatched No More: Joint Model-Policy Optimization for Model-Based RL. ArXiv, abs/2110.02758, 2021.
[25] Richard S. Sutton. Dyna, an integrated architecture for learning, planning, and reacting. ACM Sigart Bulletin, 2: 1-110066,,, 1991.
[26] Dimitri P. Bertsekas. Dynamic Programming and Optimal Control, volume 1. Belmont, MA: Athena Scientific, 1995.
[27] Q. Liu, Lihong Li, Ziyang Tang, and Dengyong Zhou. Breaking the Curse of Horizon: Infinite-Hoizzz Off- Off-Policy Estimation. In Advances in neural information processing systems, 2018.
[28] A. Mousavi, Lihong Li, Qiang Liu, and Denny Zhou, Black-box Off-policy Estimation for Infinite-Horizon Reinforcement Learning. ArXiv, abs/2003. 1126, 2020.
[29] Ofrr Nachum, Yinlam Chow, Bo Dai, and Lihong Li. DualDICE: Behavior- Agnostic Estima- tion Discounted Stationary Distribution Corrections. In Advances neural information processing systems, 2019.
[30] Ruiyi Zhang, Bo Dai, Lihong Li, and D. Schuurmans. GenDICE: Generalized Offline Estima- tion of Stationary Values. ArXiv, abs/2002.09072, 2020.
[31] Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning, arXiv preprint arXiv:2004.07219, 2020.
[32] Justin Fu, Aviral Kumar, Matthew Soh, and Sergey Levine. Diagnosing Bottlenecks in Deep Q-learning Algorithms. In International Conference on Machine Learning, 2019.
[33] Long-Ji Lin. Self-Improving Reactive Agents Based on Reinforcement Learning, Planning and Teaching. Machine Learning, 8(3 41:293-32232 1992.
[34] Monroe D Donsker and SR Srinivasa Varadhan. Asymptotic evaluation of certain markov process expectations for large time. iv. Communications on pure and applied mathematics, 36 (2):183-212, 1983.
[35] Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative Adversarial Nets. Z. Ghahramani, M. Welling, C. Cortes, N. Lawrence, and K. Q. Weinberger, editors, Advances in Neural Information Processing Systems, volume 27. Curran Associates, Inc., 2014.
[36] Xingchao Liu, Chengyue Gong, Lemeng Wu, Shujian Zhang, Hao Su, and Qiang Liu. Fuse- dream: Training free text-to-image generation with improved clip+gan space optimization. arXiv preprint arXiv:21 2021.
[37] Ziyang Tang, Yihao Feng, Lihong Li, Dengyong Zhou, and Qiang Liu. Doubly robust bias reduction in infinite horizon off- off-policy estimation. In International Conference on Learning Representations, 2020.
[38] Nathan Kallus and Masatoshi Uehara. Double reinforcement learning for efficient off-policy evaluation in markov decision processes. J. Mach. Learn. Res., 21(167): 2020.
[39] Mengjiao Yang. Ofir Nachum, Bo Dai, Lihong Li, and Dale Schuurmans. Off-policy evaluation via the regularized lagrangian. Advances in Neural Information Processing Systems, 33:6551- 6561, 2020.
[40] Junfeng Wen, Bo Dai, Lihong Li, and Dale Schuurmans. Batch stationary distribution estimation. arXiv preprint arXiv:2003.00722. 2020.
[41] Peter J. Huber. Robust Estimation of Location Parameter. The Annals of Mathematical Statistics, 35(1):73 101, 1964.
12

---

[42] Scott Fujimoto and Shixiang Shane Gu. A Minimalist Approach to Offline Reinforcement Learning. ArXiv, (bss2206.06860)) 2021.
[43] Ofir Nachum, Bo Dai, Ilya Kostrikov, Yinlam Chow, Lihong Li, and D. Schuurmans. Al- gaeDICE: Policy Gradient from Arbitrary Experience. ArXiv,
2019.
[44] Jongmin Lee. Wonseok Jeon, Byung-Jun Lee, J. Pineau, and Kee-Eung Kim. OptiDICE: Offline Policy Optimization via Stationary Distribution Correction Estimation. ArXiv, abs/2106.10783 2021.
[45] Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative Q-Learning for Offline Reinforcement Learning. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 1179-1191. Curran Associates, Inc., 2020.
[46] Ilya Kostrikov, Jonathan Tompson, Rob Fergus, and Ofir Nachum. Offline Reinforcement Learning with Fisher Divergence Critic Regularization. In International Conference Machine Learning, 2021.
[47] Toru Hishinuma and Kei Senda. Weighted model estimation for offline model- reinforce- ment learning. In Advances in neural information processing systems, 2021.
[48] Nan Jiang and Lihong Li. Doubly robust off off-policy evaluation for reinforcement learning. In Proceedings of the 23rd International Conference on Machine Learning, pages 652-661,
2016.
[49] Rafael Rafailov, Tianhe Yu, Aravind Rajeswaran, and Chelsea Finn. Offline reinforcement learning from images with latent space models. In Learning for Dynamics and Control, pages 1154 166. PMLR, 2021.
[50] Cong Lu, Philip J. Ball, Jack Parker-Holder, Michael A. Osborne, and Stephen J. Roberts. Revisiting Design Choices in Model-Basee Offline Reinforcement Learning. ArXiv, abs/2110.001335 2021.
[51] Shentao Yang, Zhendong Wang, Huangjie Zheng, Yihao Feng, and Mingyuan Zhou. A regularized implicit policy for offline reinforcement learning, arXiv preprint arXiv:2202.09673 2022.
[52] Shentao Yang. Yihao Feng, Shujian Zhang, and Mingyuan Zhou Regularizing model-based policy stationary distribution to stabilize offline reinforcement learning. In International Conference on Machine Learning, pages 24980-25006. PMLR, 2022.
[53] Jianhao Wang, Wenzhe Li, Haozhe Jiang, Guangxiang Zhu, Siyuan Li, and Chongjie Zhang. Offline reinforcement learning with reverse model-based imagination. Advances in Neural Information Processing Systems, 34, 2021.
[54] Yao Liu, Omer Gottesman, Aniruddh Raghu, Matthieu Komorowski, Aldo A Faisal, Finale Doshi- Velez, and Emma Brunskill. Representation balancing mdps for off-policy policy evaluation. Advances Neural Information Processing Systems, 31, 2018.
[55] Byung-Jun Lee, Jongmin Lee, and Kee-Eung Kim. Representation Balancing Offline Model- based Reinforcement Learning. International Conference on Learning Representations, 2021.
[56] Alec Radford and Ilya Sutskever. Improving Language Understanding by Generative Pre- Training. In arxiv, 2018.
[57] Xinjie Fan, Shujian Zhang, Bo Chen, and Mingyuan Zhou. Bayesian attention modules. Advances in Neural Information Processing Systems, 33:16662-16376 2020.
[58] Shujian Zhang, Xinjie Fan, Bo Chen, and Mingyuan Zhou. Bayesian attention belief networks. In International Conference Machine Learning, pages 12413-12426. PMLR, 2021.
13

---

[59] Shujian Zhang, Chengyue Gong, Xingchao Liu, Pengcheng He, Weizhu Chen, and Mingyuan Zhou. Allsh: Active learning guided by local sensitivity and hardness. arXiv preprint (rXiv:220504498, 2022.
[60] Michael Janner, Qiyang Li, and Sergey Levine. Reinforcement Learning One Big Sequence Modeling Problem. ArXiv, abs/2106.02003 2021.
[61] Claas Voelcker, Victor Liao, Animesh Garg, and Amir-massoud Farahmand. Value Gradient weighted Model-Bare -ased Reinforcement Learning. arXiv preprint arXiv:2204.01464 2022.
[62] Amir-massoud Farahmand, Andre Barreto, and Daniel Nikovski. Value-aware loss function for model-based reinforcement learning. In Artificial Intelligence and Statistics, pages 1486-1494.
PMLR, 2017.
[63]
Amir-massoud Farahmand. Iterative value-aware model learning. Advances in Neural Infor- mation Processing Systems, 31, 2018.
[64] Pierluca D'Oro, Alberto Maria Metelli, Andrea Tirinzoni, Matteo Papini, and Marcello Restelli. Gradient-aware model-based policy search. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pages 3801-3808, 2020.
[65] Cameron Voloshin, Nan Jiang, and Yisong Yue. Minimax model learning. In International Conference on Artificial Intelligence and Statistics, pages 1612-1620. PMLR, 2021.
[66] Stéphane Ross and Andrew Bagnell. Agnostic System Identification for Model-Based Reinforcement Learning. International Conference on Machine Learning, abs/1203.1007 2012.
[67] Ignasi Clavera, Jonas Rothfuss, John Schulman, Yasuhiro Fujita, Tamim Asfour, and P. Abbeel. Model-Based Reinforcement Learning via Meta-Policy Optimization. ArXiv, abs/1809 809.05214 2018.
[68] Thanard Kurutach, Ignasi Clavera, Yan Duan, Aviv Tamar, and P. Abbeel. Model-Ensemble Trust-Region Policy Optimization. International Conference on Learning Representations, abs/1802. 10592, 2018
[69] Bo Dai, Ofir Nachum, Yinlam Chow, Lihong Li, Csaba Szepesvári, and Dale Schuurmans. Coindice: Off-policy confidence interval estimation. Advances in neural information process- ing systems, 33:9398-9411, 2020.
[70] Yihao Feng, Lihong Li, and Qiang Liu. A kernel loss for solving the bellman equation. Advances in Neural Information Processing Systems, 32, 2019.
[71] Yihao Feng, Tongzheng Ren, Ziyang Tang, and Qiang Liu. Accountable off-policy evaluation with kernel bellman statistics. In International Conference Machine Learning, pages 3102-3111. PMLR, 2020.
[72] Yihao Feng, Ziyang Tang, Na Zhang, and Qiang Liu. Non-asymptotic confidence intervals of off off-policy evaluation: Primal and dual bounds. arXiv preprint arXiv:2103.05741, 2021.
[73] Ziyang Tang, Yihao Feng, Na Zhang, Jian Peng, and Qiang Liu. Off-policy interval estimation with lipschitz value iteration. Advances in Neural Information Processing Systems, 33:7887- 7897, 2020.
[74] Yihao Feng. Offline Reinforcement Learning via an Optimization Lens. PhD thesis, The University of Texas at Austin, 2022.
[75] Doina Precup. Richard S Sutton, and Sanjoy Dasgupta. Off Off-policy cemporal-differeee learning with function approximation. In International Conference on Machine Learning, pages 417-424, 2001.
[76] Lihong Li, Rémi Munos, and Csaba Szepesvári. Toward minimax off-policy value estimation. In Artificial Intelligence and Statistics, pages 608-616. PMLR, 2015.
14

---

[77] Tengyang Xie, Yifei Ma, and Yu-Xiang Wang. Towards optimal off policy evaluation for rein- forcement learning with marginalized importance sampling. Advances Neural Information Processing Systems, 32, 2019.
[78] Ofir Nachum and Bo Dai. Reinforcement learning via fenchel-nockafellar duality, arXiv preprint arXiv:2001.01000 0166,, 2020.
[79] Rémi Munos. Error bounds for approximate policy iteration. In International Conference Machine Learning, volume 3, pages 560-567, 2003.
[80] Jinglin Chen and Nan Jiang. Information-Thenee Learning. ArXiv, abs/ 19950003600, 2019.
Considerations in Batch Reinforcement
[81] Masatoshi Uehara, Jiawei Huang, and Nan Jiang. Minimax weight and q-function learning for off-policy evaluation. In International Conference on Machine Learning, pages 9659-9668.
PMLR, 2020.
[82] Xi Chen, Ali Ghadirzadeh, Tianhe Yu, Yuan Gao, Jianhao Wang, Wenzhe Li, Bin Liang, Chelsea Finn, and Chongjie Zhang. catent-variable advantage-weighted policy optimization for offline rl. arXiv preprint (xXiv:2203.0894) 2022.
[83] Kurtland Chua, Roberto Calandra, Rowan McAllister, and Sergey Levine. Deep Reinforcement Learning in Handful of Trials using Probabilistic Dynamics Models. In Advances in neural information processing systems, 2018.
[84] Tatsuya Matsushima, Hiroki Furuta, Yutaka Matsuo, Ofrr Nachum, and Shixiang Shane Gu. Deployment- Efficient Reinforcement Learning via Model- Offline Optimization. International Conference on Learning Representations, abs/2006.0366 2021.
[85] H. Hasselt. Double Q-learning. In Advances in neural information processing systems, 2010.
[86] Chengyue Gong, Xingehao Liu, and Qiang Liu. Bi-objective Trade-off with Dynamic Barrier Gradient Descent. Advances in Neural Information Processing Systems, 34, 2021.
[87] Diederik P. Kingma and Jimmy Ba. Adam: A Method for Stochastic Optimization. In International Conference on Learning Representationss 2014.
[88] Yuguang Yue, Zhendong Wang, and Mingyuan Zhou. Implicit Distributional Reinforcement Learning. In Hugo Larochelle, Mare' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tier Tien Lin, editors, Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 12, 2020, virtual, 2020.
[89] Yifan Wu, George Tucker, and Ofir Nachum. Behavior regularized offline reinforcement learning, arXiv preprint arXiv: 2019.
[90] Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, M. Laskin, P. Abbeel, A. Srinivas, and Igor Mordatch. Decision Transformer: Reinforcement Learning via Sequence Modeling. ArXiv, abs/2106.01345, 2021.
[91] Ilya Kostrikov, Ashvin Nair, and Sergey Levine. Offline Reinforcement Learning with Implicit Q-Learning. ArXiv,
2021.
[92] Aravind Rajeswaran, Vikash Kumar, Abhishek Gupta, Schulman, Todorov, and Sergey Levine. Learning Complex Dexterous Manipulation with Deep Reinforcement Learning and Demonstrations. ArXiv, 10087, 2018.
[93] Tom White. Sampling generative networks. arXiv preprint
2016.
[94] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, loannis Antonoglou, Daan Wierstra, and Martin Riedmiller. Playing atari with deep reinforcement learning. arXiv preprint arXiv:1312.5600 2013.
15

---

[95] Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
[96] Alec Radford, Luke Metz, and Soumith Chintala. Unsupervised Representation Learning with Deep Convolutional Generative Adversarial Networks. CoRR, abs/1511 1.06434, 2016.
[97] Tim Salimans, I. Goodfellow, Wojciech Zaremba, Vicki Cheung, Alec Radford, and Xi Chen. Improved Techniques for Training GANs. In Advances in neural information processing systems, 2016.
[98] Ian Goodfellow. Nips 2016 tutorial: Generative adversarial networks. arXiv preprint arXiv:1701.00160, 2016.
[99] Huangjie Zheng and Mingyuan Zhou, Exploiting Chain Rule and Bayes' Theorem to Compare Probability Distributions. Advances in Neural Information Processing Systems, 34, 2021.
[100] Shujian Zhang, Xinjie Fan, Huangjie Zheng, Korawat Tanwisuth, and Mingyuan Zhou. Align- ment attention by matching key and query distributions. Advances Neural Information Processing Systems, 34, 2021.
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? [Yes] See Sections 3 and 4.
(b) Did you describe the limitations of your work? [Yes] See Sections and 6.
(c) Did you discuss any potential negative societal impacts of your work? [Yes] See Appendix F.
(d) Have you read the ethics review guidelines and ensured that your paper conforms them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all thcoretical results? [Yes] See Section and Appendix B.
(b) Did you include complete proofs of all theoretical results? [Yes] See Appendix
3. If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material or as a URL)? [Yes] See Section 3 and Appendix for descriptions. Source code is publicly released.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes] See Appendix E.2 and E.3.
(c) Did you report error bars (e.g., with respect the random seed after running experi- ments multiple times)? [Yes] See Tables 1, 2, 3, and 4.
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)? [Yes] See Appendix G.
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? [Yes]
(b) Did you mention the license of the assets? [Yes] See the GitHub pages of the assets.
(c) Did you include any new assets either in the supplemental material or URL? [No]
(d) Did you discuss whether and how consent was obtained from people whose data you't 're
using/curating? [No] We use public benckmarks.
(e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [No] The used benckmarks are widely accepted in the community.
5. If you used crowdsourcing conducted research with human subjects...
16

---

(a) Did you include the full text of instructions given to participants and screenshots, if applicable? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable? [N/A]
(c) Did you include the estimated hourly wage paid participants and the total amount spent on participant compensation? [N/A]
17


Reviews:
[{"review_id": "soNI0PesdiF", "reviewer": "Reviewer_A7rq", "summary": "In this work, the authors investigate a new approach to the problem of Model-Based Reinforcement Learning in the off-line RL setting.\nThe paper proposes a number of new ideas. First, instead of first learning a model and subsequently using it to obtain a policy, their approach alternates between updating the model and improving the policy. The MLE model learning is aware of the discrepancy between the stationary distribution of the behavior policy and of the learned one through the use of marginal importance weighting. These weights are learned using a novel Bellman backup operator. At the same time the authors propose a new regularization loss in their policy updates to account for the difference between the true transition probabilities and the ones produced by the learned model. This encourages the imaginary rollouts produced by the model and policy to stay close to the true distribution observed in the offline dataset. \n", "questions": "The authors should look at the work proposed in [1] as another simpler method where model and policy learning are coupled to achieve state-of-the-art performance. Is there any benefit in using transition models instead of value equivalent ones ? \n\nMinor suggestions:\nCite and clarify the idea of target networks. \nCite and clarify the use of a Huber loss in the Q-function update rule. \n\n[1] Schrittwieser J, Hubert TK, Mandhane A, Barekatain M, Antonoglou I, Silver D (2021) Online and Offline Reinforcement Learning by Planning with a Learned Model. In: Beygelzimer A, Dauphin Y, Liang P, Vaughan JW (eds) Advances in Neural Information Processing Systems", "limitations": "The authors have adequately addressed the limitations and potential negative societal impact of their work. ", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The authors propose a theoretically and empirically principled algorithm. Their method introduces a number of novel ideas, in a clear way, which are in the right direction of coupling model and policy learning in offline RL. \n\nA possible weakness would be the stability of the learning rules. \nThe model learning updates depend on the values of the importance weights. At the same time the importance weights are learned through a Bellman backup which depends on the Q values learned using the model. Moreover, the policy is regularized through a GAN mechanism which can further destabilize learning. The authors show empirically that their method is stable which removes some of these concerns. However, since there are many hyperparameters involved it is not clear how important these are in achieving stability. ", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "MtmHllJDdW3", "reviewer": "Reviewer_UQCr", "summary": "\n\nThe authors propose a unified framework for alternating model training and policy learning in the offline RL setting. The model and the policy are trained to maximize a lower bound of the actual return. The experiments are conducted in the D4RL benchmark and several SOTA baselines are taken for comparison.", "questions": "\n1. In Section 3.2, to learn a policy, the authors optimize a surrogate objective of Eq.(7) via estimate D_\\pi. However, If we have an approximation of D_\\pi, why don’t we optimize the model through D_\\pi too? Similarly, If we have an approximation of \\omega, why don’t we use \\omega to estimate D_pi directly for policy optimization? \n\n2. The results in Figure 2 seem great but the details of experiment settings are still unclear to me. Can you give a detailed description (which can be in the Appendix)?  For example, how to normalize the MIWs; What’s the difference among the 1~9 policies for testing.\n\n3. Can the authors provide more evidence to demonstrate the mechanism of the regularizer and MIW training techniques such as conducting experiments in some interpretable environments or recording some intermediate  metrics in some specially designed experiment like (b) in Section 4.2", "limitations": "\nI think the overall idea of the article is interesting but needs some refinements before publication. The authors may be able to further analyze from a theoretical perspective; Or, from an experimental perspective, I think the mechanism of regularization and omega estimation can be more fully experimentally verified. I will consider increasing the score if the authors improve the article in either of the aspect.\n", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "\nStrengths:\n\n1. The motivation for building a unified framework for alternating model training and policy learning is in time and reasonable. Although current offline model-based RL algorithms have learned policies considering the inaccuracy of the models, the objective of model learning is still separated from the offline model-based RL pipeline. It is valuable as a research topic to investigate; \n\n2. The article is overall well-written and easy-to-follow to me;\n\n3. The experiment results in the benchmark datasets seem great.\n\nWeaknesses:\n\nWith respect to the theoretical modeling,  there still have large gaps between Theorem 1 and the algorithm designs in Sec. 3.2, 3.3, and 3.4. And thus, I prefer to evaluate the value of the article with respect to the novelty of the proposed techniques. \n\nIn Section 3.2 and 3.2, the authors proposed a new regularizer for policy learning and a new technique for importance weight estimation.  Both of the techniques are interesting: the regularizer not only evaluates the difference between two policies but the effect difference of the policies in different models, which makes sense for formulating a unified training framework. Importance weight estimation is also a challenging task. The proposed method utilize a test function to identify and correct the weight estimation. However, there are not enough experiments and descriptions to demonstrate the actual mechanism and the relation of the proposed techniques (See details in the following).", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ["I don’t know"]}, "extra_text": {}}, {"review_id": "GjndpxTq67T", "reviewer": "Reviewer_3XmH", "summary": "This work proposes a new algorithm for offline model-based reinforcement learning (MBRL). Based on a lower bound of the cumulative reward objective (Eq.(5)), the proposed method alternatively optimizes the policy and the transition model using a fixed dataset. To bypass the difficulties in the objective, the part of the lower bound is further approximated by the Jensen-Shannon divergence (JSD) so that GAN training can be applied. In addition, the objective includes a marginal importance weight (MIW) that is also being estimated during training. Experiments on the D4RL benchmark show that the proposed method can achieve better performance than existing methods in the literature.", "questions": "Although it is nice to have a unified objective, there are several questions/concerns regarding the method.\n\n1. The stability of the proposed method becomes a question when it alternates between learning several models (policy, action-value function, MIW, transition kernel, and the GAN discriminator). Each of them requires some sort of stability tricks, so the overall procedure seems very fragile as a whole.\n\n2. L121: Why implicit policy is preferred? Gaussian or Beta policies are common in the literature.\n\n3. When optimizing (13), 1. When optimizing (13), does some of the experience in the augmented dataset have fake rewards? Algorithm 1 does not mention anything about the learning of \\hat{r} and how it is used.\n\n4. The learning of MIW is interesting. Eq.(9) in the current paper is Eq.(4) in Wen et al., (2020) [44] when expanded. Such connection should be discussed and referenced. The main difference is that the current paper uses a test function. However, the motivation for using the action-value function as the test function is unsubstantiated and requires further discussion/evaluation. It seems more natural to use the reward function instead when it is motivated by the primal-dual relation (Eq.(1)), because we care more about learning accurate MIW. In this way, the LHS of (10) will be exactly the cumulative reward objective.\n\n5. Experiments\n- It would be more readable if the best model-based method is italicized/underlined instead of also bold in Table 1\n- How are the MIWs normalized in Fig.2? It is surprising that the median of DualDICE can be so far away from 1 after normalization.\n- The description in Sec.4.2(c) is unclear. The first sentence suggests that the regularization uses KL as in Thm.1, but the second sentence immediately suggests using JSD instead, which seems to be the “Main” variant. Specifically, how is WPR different from Main?\n\nMinor\n- L63: \\pi_b is presumably the behaviour policy?\n- L163: does not undermine the convergence property (under mild conditions on Q).", "limitations": "The authors have sufficiently addressed the limitations.", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n- A unified objective for training both the model and the policy\n- Solid experiments\n\nWeaknesses\n- Lacking theoretical justification\n- Some of the approximation steps are questionable", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "7jW_P7u3xl4", "reviewer": "Reviewer_dzbb", "summary": "The paper focuses on the problem of offline reinforcement learning, where the authors propose a novel model based algorithm. Specifically, the authors start from the general framework of bounding the difference between the return on the true model and a learned model, and then derive the objective for model learning and policy improvement. For the learning dynamics model, the authors propose a maximum likelihood objective weighted by the importance ratio of the policy state action visitation distribution and the dataset distribution. Specifically, the authors introduce a Bellman equation-like objective to estimate this important ratio. For policy improvement, the authors adopt an actor-critic style algorithm and penalizes the divergence between the policy and the dataset during actor updates. The divergence value is estimated using a generative adversarial network.\n\nThe authors evaluate the proposed method in the D4RL benchmark, and show that the proposed method outperforms prior model based and model free methods.\n", "questions": "Prior work [1] suggests that some baseline methods used in this paper can be tuned much better. I wonder if the authors could include these improvements in the baseline comparison as well as incorporate them into the proposed method if applicable.\n\n\nReferences\n\n[1] Lu, Cong, et al. \"Revisiting design choices in offline model based reinforcement learning.\" International Conference on Learning Representations. 2021.\n", "limitations": "The authors have sufficiently addressed the limitations and societal impacts.\n", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Overall I think the paper presents an interesting idea of doing model based RL from offline dataset. However I do have some concerns.\n\nPros\n\nThe empirical results of the paper are pretty strong. From table 1, it seems that the proposed method achieves better performance than pior method in most of the tasks from D4RL, which is a clear indication that the proposed method provides performance benefits. I also find the ablation study quite useful in showing that the MIW weighted model learning is important.\n\nThe paper is well written and the presentation is very clear. Building on top of the general framework of model based RL, the derivations and analysis of the proposed method are intuitive and easy to follow. \n\n\nCons\n\nIt seems to me that the practical implementation of the algorithm is somewhat disconnected from the derivation of the algorithm, so it is hard to tell if the method works because of the reasons justified by theoretical analysis or because of the choices made in practical implementation. For example, the authors start from the standard approach to bound the gap between the return between the true and learned dynamics model (theorem 1) and propose to estimate the importance ratio (MIW). However, the derived method for estimating MIW is not implementable since it uses quantities not available to the algorithm, so the authors propose to replace the unknown density with the learned Q function. It is not clear to me why this is a good choice and hence I believe that more theoretical analysis is necessary.\n\n\nDuring policy training, the authors used a GAN to estimate the divergence, which is one of many ways to estimate the divergence or quantify the uncertainty. Prior work [1] suggests that the performance difference of model based offline RL algorithms mainly come from these practical choices. Hence I believe the paper could benefit from a comparison of different ways of penalizing the policy during policy improvement. \n\n\nWhile the paper presents an interesting idea and the empirical results are strong, I do believe that some more analysis and justification is needed for the proposed method. Therefore, I’m neutral between acceptance and rejection for this paper. I highly encourage the authors to include them during the author response period.\n\n\n\nReferences\n\n[1] Lu, Cong, et al. \"Revisiting design choices in offline model based reinforcement learning.\" International Conference on Learning Representations. 2021.\n\n\n----\nThe authors' response addressed most of my concerns about the paper and therefore I'm recommending acceptance. I've increased my rating from 5 to 6.\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept