Weighted model estimation for offline model-based reinforcement learning
Year: 2021

Abstract:
This paper discusses model estimation in offline model-based reinforcement learning (MBRL), which is important for subsequent policy improvement using an estimated model. From the viewpoint of covariate shift, a natural idea is model estimation weighted by the ratio of the state-action distributions of offline data and real future data. However, estimating such a natural weight is one of the main challenges for off-policy evaluation, which is not easy to use. As an artificial alternative, this paper considers weighting with the state-action distribution ratio of offline data and simulated future data, which can be estimated relatively easily by standard density ratio estimation techniques for supervised learning. Based on the artificial weight, this paper defines a loss function for offline MBRL and presents an algorithm to optimize it. Weighting with the artificial weight is justified as evaluating an upper bound of the policy evaluation error. Numerical experiments demonstrate the effectiveness of weighting with the artificial weight.

heading: 1 Introduction
text: Reinforcement learning (RL) is a framework to learn a policy in an unknown environment [1]. Model-based RL (MBRL) is an approach that explicitly estimates a transition model and utilizes it to improve a policy. Compared to model-free RL, one major advantage of MBRL is the data efficiency [2–4], which is important for applications where the data collection is expensive, such as robotics or healthcare. This paper focuses on offline MBRL, i.e. MBRL that learns a policy from previously collected datasets.
One of the main challenges in offline MBRL is distribution shift [5]. The distribution of offline data is different from the distribution of data obtained in the future when applying an improved policy. This is a situation called covariate shift, where a model estimated using standard supervised learning techniques such as empirical risk minimization (ERM) may not generalize to test data. A major idea in MBRL is to estimate a model using such a standard method and to improve a policy using the estimated model in regions where the predictions are accurate [6–8]. A point to be improved here is model estimation without considering covariate shift. In supervised learning under covariate shift, importance-weighted ERM can obtain good predictive performance [9]. If a model in offline MBRL can be estimated in a similar way, then the estimated model will make more accurate predictions, resulting in increased policy improvement. Motivated by this, this paper tackles the issue of weighted model estimation considering covariate shift in offline MBRL.
From the viewpoint of covariate shift, a natural idea is weighting with the ratio of the state-action distributions of offline data and real future data, as the agent uses a model to predict real future data. For convenience, this paper calls the ratio “the natural weight.” However, estimating the natural weight is one of the main challenges for off-policy evaluation [5, 10]. Compared to density ratio estimation for supervised learning [11], the difficulty of off-policy evaluation is the inaccessibility of samples from the distribution of real future data.
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
As an alternative, this paper considers weighting with the state-action distribution ratio of offline data and “simulated” future data. This paper calls the ratio “the artificial weight,” in contrast to the natural weight. Simulated future data can be generated in offline simulation, unlike real future data. Since samples from both distributions are available, the artificial weight can be relatively easily obtained by standard density ratio estimation techniques for supervised learning. Based on the artificial weight, this paper defines a loss function for offline MBRL in Section 4. This paper presents an algorithm to optimize it in Section 5.
The question here is the validity of weighting with the artificial weight, as it may not be natural for covariate shift. This paper justifies it as evaluating an upper bound of the policy evaluation error in Section 4. This paper demonstrates that it works in practice in numerical experiments in Section 6. The above is the contribution of this paper.

heading: 2 Related Works
text: Most of offline MBRL researches [7, 8, 12] discuss how to learn a policy using a learned model. In contrast, this paper studies how to estimate a model considering covariate shift in MBRL, which depends on a policy to be applied. In principle, these can be combined as an algorithm that repeats between estimating a model and learning a policy, like Algorithm 3.
Importance-sampling-based methods estimate the natural weight, namely, the distribution ratio of offline data and real future data on a trajectory or state(-action) space [13–15, 10]. The advantage is unbiasedness, while the disadvantage is high variance. Model-based methods, including this paper, evaluate policies utilizing models trained using data. Model-based methods have the opposite advantage and disadvantage of importance-sampling-based methods. Doubly robust methods combine importance-sampling-based and model-based methods [16–18]. In principle, incorporating the idea of this paper as a model-based component is possible.
VAML [19, 20] defines a state-action-wise loss function based on the expected error of Bellman operator and learns a model to optimize a weighted sum. Similar ideas are found in [21–23]. The main focus is the state-action-wise loss function, while there are few studies about how to weight it is. The few studies, e.g. GAMPS [24], use the natural weight. In contrast, this paper simply uses the negative log likelihood loss function for a state-action pair and discusses the artificial weight, which is relatively easy to estimate.
This paper is similar to VMBPO [25] in that an EM-style algorithm updates a model and a policy using a universal loss function. The differences is that this paper considers covariate shift but does not use the idea of RL as inference [26]. Formulation that explicitly incorporates RL as inference is also an interesting topic.
Another approach is learning a model robust against covariate shift by reducing the difference between the distributions of offline data and future data on a feature space [22, 27, 28]. Since covariate shift measured in a feature space is encouraged to decrease, policy improvement increasing covariate shift is also discouraged. That is, such an idea implicitly limits policy improvement. This paper differs in that it is trying to utilize covariate shift.

heading: 3 Preliminaries
text: This paper discusses a discounted infinite horizon MDP [29]. Let s be a state. Let a be an action. Let ρ(s) be an initial state distribution. Let P(s′|s, a) be a transition probability function. Let r(s, a) be a reward function. Let γ be a discount factor. Let π be a policy. The value function is
V πP (s 0) = Eat∼π(·|st),st+1∼P(·|st,at) [ ∞∑ t=0 γtr(st, at) ∣∣∣∣∣ s0 ] .
The expected return is ηπP = Es0∼ρ [V πP (s)]. The state distribution is
dπP(s) = (1− γ) ∞∑ t=0 γtPr ( st = s ∣∣ ρ, π,P) . The state-action distribution is dπP(s, a) = d π P(s)π(a|s). Note that ηπP = E(s,a)∼dπP [ r(s,a) 1−γ ].
Let P∗ denote the true transition probability, which is unknown to the agent. Let Pθ denote the transition probability model parameterized by θ. Let D = {(sn, an, s′n)}Nn=1 be offline data, which is obtained from P∗. Let dDP∗ be the distribution of the input part of D, i.e. {(sn, an)} N n=1. Let D̂πθ be simulation data generated given (Pθ, π) with state initialization with probability 1−γ, which is equal to sampling (s, a) from dπPθ . Let cθ(s, a) = Es′∼P∗(·|s,a) [− lnPθ(s
′|s, a)] be the cross entropy between P∗(·|s, a) and Pθ(·|s, a). Let h(s, a) = Es′∼P∗(·|s,a) [− lnP∗(s′|s, a)] be the self-entropy of P∗(·|s, a).

heading: 4 Loss Function
text: This paper discusses the theory for the case where the reward function is known. For the unknown reward function, if the prediction of a reward model is bounded, the same loss functions can be derived, see Appendix.

heading: 4.1 Policy Evaluation
text: The goal of policy evaluation is estimating ηπP∗ when given π. For a model-based setting, the goal reduces to finding θ decreasing |ηπP∗ − η π Pθ |, which is bounded as
|ηπP∗ − η π Pθ | ≤
γE(s,a)∼dπPθ [∣∣∑ s′ (P∗(s′|s, a)− Pθ(s′|s, a))V πP∗(s ′) ∣∣]
1− γ
≤ BE(s,a)∼dπPθ [||P∗(·|s, a)− Pθ(·|s, a)||1]√
2 ≤ B √ E(s,a)∼dπPθ [cθ(s, a)− h(s, a)]
= B √ E(s,a)∼dDP∗ [w π θ (s, a)cθ(s, a)]− E(s,a)∼dDP∗ [w π θ (s, a)h(s, a)]
≤ B √ E(s,a)∼dDP∗ [w π θ (s, a)cθ(s, a)]− hmin, (1)
where B = γmaxs,a |r(s,a)| √ 2
(1−γ)2 , w π θ (s, a) =
dπPθ (s,a) dDP∗ (s,a) , and hmin = mins,a h(s, a). The first inequality
is the Telescoping lemma [21]. The second inequality holds from Holder’s inequality and ||V πP ||∞ ≤ maxs,a |r(s,a)|
1−γ . The third inequality holds from Pinsker’s inequality and Jensen’s inequality. When trying to optimize the third inequality, E(s,a)∼dDP∗ [w π θ (s, a)h(s, a)] is difficult to handle, as h(s, a) is unknown. The fourth inequality (1) excludes the difficult part from the optimization.
This paper considers decreasing the policy evaluation error by minimizing Equation (1). What is to be minimized is
Lπ(θ) = E(s,a)∼dDP∗ [w π θ (s, a)cθ(s, a)] = E(s,a)∼dDP∗ ,s′∼P∗(·|s,a) [−w π θ (s, a) lnPθ(s′|s, a)] .
The empirical average of Lπ(θ) is
L̂π(θ) = − 1 N N∑ n=1 wπθ (sn, an) lnPθ(s′n|sn, an), (2)
Equation (2) is a loss function weighted with wπθ (s, a), which is the artificial weight, namely, the distribution ratio of offline data D and simulated future data D̂πθ . Thus, weighting with the artificial weight is justified as evaluating an upper bound of the policy evaluation error.

heading: 4.2 Policy optimization
text: The goal of policy optimization is finding π increasing ηπP∗ . In offline MBRL, the agent utilizes η π Pθ for the optimization. A lower bound of ηπP∗ is
ηπP∗ ≥ η π Pθ − |η π P∗ − η π Pθ | ≥ η π Pθ −B √ Lπ(θ)− hmin.
The first term evaluates the expected return by model Pθ, and the second penalizes the policy evaluation error.
Inspired from increasing ηπP∗ by maximizing the lower bound, this paper defines a generalization loss function for offline MBRL by
J(θ, π) = −ηπPθ +B ′ √ Lπ(θ)− hmin (3)
where B′ ∈ [0, B] is a tunable coefficient to scale the penalty for the policy evaluation error. The agent optimizes an estimate of Equation (3), where empirical averages of ηπPθ and L
π(θ) are weighted with the artificial weight. Similarly to the discussion in Section 4.1, weighting with the artificial ratio is justified as evaluating the expected return with the penalty for the policy evaluation error.
Remark. If using the Telescoping lemma with respect to dπP∗ instead of d π Pθ , the policy evaluation error is similarly bounded as
|ηπP∗ − η π Pθ | ≤ B √√√√E(s,a)∼dDP∗ [ dπP∗(s, a)
dDP∗(s, a) cθ(s, a)
] − hmin, (4)
where dπP∗ (s,a)
dDP∗ (s,a) is the natural weight. There are two advantages to consider Equation (1) rather
than Equation (4). The first is that the artificial weight can be more easily ontained, as described in Section 1. The second is that Equation (1) leads to a standard MDP optimization problem, as described in Section 5.2, while Equation (4) results in a more complicated optimization problem related to both dπPθ and d π P∗ .

heading: 5 Algorithm

heading: 5.1 Policy Evaluation

heading: 5.1.1 Full Version
text: This paper considers a gradient-based optimization algorithm. The gradient of Equation (2) is
∇θL̂π(θ) = − 1
N N∑ n=1 wπθ (sn, an) {∇θ lnPθ(s′n|sn, an) + lnPθ(s′n|sn, an)zπθ (sn)} ,
where zπθ (s) = ∇θ ln dπPθ (s), which can be estimated using an extension of LSDG [30], as described later. The artificial weight, wπθ , can be estimated using a density ratio estimation method given (D, D̂πPθ ). The gradient of Equation (2) can be estimated using them.
In practice, the computational efforts to estimate (wπθ , zθ) are non-negligible. This paper presents a method that iterates between estimating (wπθ , zθ) and updating θ locally. To locally update from θ
(i), this paper uses a local approximate loss function defined by
L̂π(θ; θ(i)) = − 1 N N∑ n=1 wπθ(i)(sn, an) { lnPθ(s′n|sn, an) + (θ − θ(i)) · ( lnPθ(i)(s′n|sn, an)zπθ(i)(sn) )} . (5) The gradient of Equation (5) with respect to θ can be obtained immediately, as (wπ
θ(i) , zπ θ(i) ) is fixed. For θ = θ(i), Equation (2) and Equation (5) are equal, and their gradients are also equal. That is, for θ = θ(i), the first-order necessary conditions for optimality are the same. In theory, if assuming that (wπ
θ(i) , zπ θ(i) ) is accurately estimated, the above-mentioned iterative method converges.
Importance-weighted ERM is consistent but can be unstable in practice [9], even when the weight is exactly known. Importance-weighted ERM weighted with wπ
θ(i) can be more unstable due to the
estimation and approximation errors of density ratio estimation of wπ θ(i) . To stabilize the estimation of θ, this paper replaces wπ
θ(i) (sn, an) in Equation (5) with [wπθ(i)(sn, an)] α, where α = 0 corresponds to ERM, and α = 1 corresponds to importance-weighted ERM weighted with wπ
θ(i) (sn, an). That is,
α is a hyperparameter controlling the trade-off between consistency and stability.
Algorithm 1 Weighted model estimation for policy evaluation (full version). 1: Input: D, π, α, and Pθ(0) . 2: for i = 0, 1, 2, · · · do 3: Generate D̂π
θ(i) , using simulation given (Pθ(i) , π).
4: Estimate wπ θ(i) , using density ratio estimation given (D, D̂π θ(i) ). 5: Estimate zπ
θ(i) , using extended-LSDG given (Pθ(i) , π).
6: Obtain Pθ(i+1) , maximizing Equation (5) given (D, wπθ(i) , z π θ(i) , α). 7: end for
Algorithm 2 Weighted model estimation for policy evaluation (simplified version). 1: Input: D, π, α, and Pθ(0) . 2: for i = 0, 1, 2, · · · do 3: Generate D̂π
θ(i) , using simulation given (Pθ(i) , π).
4: Estimate wπ θ(i) , using density ratio estimation given (D, D̂π θ(i) ). 5: Obtain Pθ(i+1) , maximizing Equation (7) given (D, wπθ(i) , α). 6: end for
Extending LSDG. LSDG [30] uses the gradient of ln dπP(s) with respect to policy parameters in model-free RL. This paper extends this idea to the gradient with respect to model parameters in MBRL. In a discounted MDP, dπPθ satisfies
dπPθ (s ′) = (1− γ)ρ(s′) + γ ∑ s,a dπPθ (s)π(a|s)Pθ(s ′|s, a).
The gradient with respect to the j-th parameter, θj , is
dπPθ (s ′)zπθj (s ′) = ∑ s,a dπPθ (s)π(a|s)Pθ(s ′|s, a) { γ∇θi lnPθ(s′|s, a) + γzπθj (s) } , (6)
where zπθj (s) = ∇θj ln d π Pθ (s). Equation (6) implies that z π θj is the value function satisfying a “forward” Bellman equation, where the transition probability function is Pθ, and the reward function is γ∇θi lnPθ(s′|s, a). To estimate zπθj , this paper extends LSDG by replacing the gradient of lnπ(a|s) with γ∇θj lnPθ(s′|s, a). The extended-LSDG estimates zπθj using simulation data generated using (Pθ, π). Since some off-policy evaluation methods estimate the discounted state distribution ratio as a solution to another forward Bellman equation [14, 15, 5], their extensions could also be used instead of LSDG.
Algorithm 1 summarizes how to optimize θ for policy evaluation, described above.

heading: 5.1.2 Simplified version
text: The computationally expensive part in Algorithm 1 is estimating zπ θ(i) . This requires estimating the same number of the value functions satisfying Equation (6) as the number of model parameters. This is unrealistic for neural networks with many parameters. As a practical simplification, this paper also considers replacing Equation (5) with
L̂π(θ; θ(i)) = − 1 N N∑ n=1 wπθ(i)(sn, an) lnPθ(s ′ n|sn, an). (7)
Algorithm 2 is the resulting simplified algorithm. The price of the simplification is the loss of convergence in theory, as the first-order necessary conditions are no longer the same. For this reason, the number of iterations should be specified. In practice, Algorithm 2 also works well, see Section 6.

heading: 5.2 Policy Optimization
text: This paper considers an EM-style optimization algorithm, where the E-step fixes π and optimizes θ, and the M-step fixes θ and optimizes π. Instead of optimizing Equation (3) directly, this paper uses a
Algorithm 3 Model-based offline policy optimization based on weighted model estimation.
1: Input: D, α, B′, Pθ(0) , and π(0). 2: for i = 0, 1, 2, · · · do 3: if i = 0 then 4: Estimate Pθ(i+1) , using ERM given D. 5: else 6: Compute b̂, given (B′,D,Pθ(i) , π(i)). 7: Obtain Pθ(i+1) , using E-step given (D, b̂, π(i), α). 8: end if 9: Estimate ĉθ(i+1) , using ERM given (D,Pθ(i+1)).
10: Compute b̂, given (B′,D,Pθ(i+1) , π(i)). 11: Update π(i+1), using M-step given (Pθ(i+1) , b̂, ĉθ(i+1)). 12: end for
surrogate function in the majorization minimization framework [31],
Jsurr(θ, π; θ(i), π(j)) = −ηπθ + b(Lπ(θ)− hmin) + C
= −E(s,a)∼dπPθ
[ r(s, a)
1− γ − b(cθ(s, a)− hmin)
] + C, (8)
where b = B ′
2 √ Lπ (j) (θ(i))−hmin and C = b(Lπ (j) (θ(i))− hmin). The advantages of using Equation (8)
are that the E-step reduces to a similar problem with Section 5.1 and that the M-step comes down to a MDP optimization problem, as described later.
Let b̂ = B ′
2 √ L̂π (j) (θ(i))−ĥmin be an estimate of b, where ĥmin = min(s,a,s′)∈D [− lnPθ(i)(s′|s, a)].
Note that L̂π (j) (θ(i)) is computed using Equation (2) given (D,Pθ(i) , wπ (j) θ(i) ), where wπ (j) θ(i) can be estimated using density ratio estimation given (D, D̂π(j) θ(i) ), and D̂π(j) θ(i)
can be generated using simulation given (Pθ(i) , π(j)). That is, b̂ can be computed when given (D,Pθ(i) , π(j)).
E-step. To update θ while fixing π = π(i), the E-step estimates Equation (8) by
− 1 N N∑ n=1 wπ (i) θ (sn, an) [ r(sn, an) (1− γ) + b̂ lnPθ(s′n|sn, an) ] , (9)
where the constant terms are ignored. Equation (9) can be minimized using Algorithm 1 replacing lnPθ(s′n|sn, an) with lnPθ(s′n|sn, an) + r(sn,an)
b̂(1−γ) , or Algorithm 2.
M-step. To update π while fixing θ = θ(i+1), the M-step estimates Equation (8) by
−E(s,a)∼dπP θ(i+1)
[ r(s, a)
1− γ − b̂(ĉθ(i+1)(s, a)− ĥmin)
] , (10)
where ĉθ(i+1) is a model of cθ(i+1) , and C is ignored. This paper trains ĉθ(i+1) using input data {(sn, an)}Nn=1 and output data {− lnPθ(i+1)(s′n|sn, an)}Nn=1. Note that Equation (10) is the expected return in a MDP, where the transition probability is Pθ(i+1) , and the reward function is
r(s, a)− (1− γ)b̂(ĉθ(i+1)(s, a)− ĥmin).
Here, ĉθ(i+1)(s, a)− ĥmin can be seen as an estimate of the modeling error, and (1− γ)b̂ can be seen as the coefficient to scale the penalty for the modeling error. Since this MDP is known to the agent, Equation (10) can be optimized using any MDP optimization approaches, including RL algorithms.
Algorithm 3 summarizes the above-mentioned procedure.
Remark. If α = 0, then the E-step reduces to ERM, and Lines 6-7 in Algorithm 3 are skipped. This case can be seen as an instantiation of MOPO [8], where ĉθ(i+1) − ĥmin and b̂ are the admissible error estimator and the penalty coefficient, respectively. This means that the M-step is essentially the same as MOPO, while the E-step is different from MOPO. In other words, MOPO can be seen applying only M-step to optimize the penalized return. Thus, Algorithm 3 improves MOPO in that it optimizes the penalized return with respect to both a model and a policy by the E-step and M-step.

heading: 6 Numerical Experiment

heading: 6.1 Pendulum
text: Firstly, to simply illustrate the effectiveness of weighting with the artificial weight, this paper studies policy evaluation on a pendulum environment. State s is the pair of angle and angular velocity, and action a is the torque. Offline data D is collected using a random policy. The initial state is around (π, 0), and the cost-free state is (0, 0). Target policy π is a controller that swing-up the pendulum to (0, 0) in the real environment. That is, real future data DπP∗ is the swing-up behavior towards (0, 0). Figure 1 shows the state part of D and DπP∗ . The outline of the implementation is as follows. The agent uses Pθ represented by two-layer neural networks with 8 units with tanh activation. The agent estimates wπθ using density ratio estimation via logistic regression [32, 11]. The agent estimates zπθ using an extension of LSDG [30]. For more details, see Appendix.
Unweighted vs weighted ERM. The discounted return of π averaged over 100 episodes in the real environment is −15.7. This paper regards this value as the expected return, namely, the goal of policy evaluation. When using the ERM method, the estimated value is −80.7. When using Algorithm 1, the estimated value is −17.8. That is, Algorithm 1 estimates the expected return more accurately. Below, this paper describes why Algorithm 1 obtains such a better estimate.
Figure 1 (c) shows simulated future data D̂πPθ when using the unweighted maximum likelihood method, a typical ERM method. The pendulum is predicted not to swing-up to (0, 0), unlike DπP∗ . This result can be described as follows. Based on the unweighted loss function, Pθ tries to represent D as much as possible. However, its expressiveness is insufficient for global prediction. As a result, Pθ does not predict the swing-up behavior.
Figure 1 (d) shows D̂πPθ when using Algorithm 1 with α = 1. The state part of D is colored based on the logarithm of wπθ (sn, an) when Algorithm 1 converges. The pendulum is predicted to swing-up to (0, 0), like DπP∗ . This situation can be described as follows. When Algorithm 1 converges, Line 3 in Algorithm 1 predicts the swing-up behavior shown in Figure 1. Using the swing-up behavior as D̂πPθ , Line 4 estimates w π θ (sn, an) shown in Figure 1. After estimating z π θ at Line 5, Line 6 estimates Pθ to predict more accurately where the weight is larger. Since the expressive capability of Pθ is collected near the larger weights, Pθ can predict the same swing-up behavior again at Line 4 in the next iteration, despite the lack of global expressiveness as mentioned in the paragraph above. This is what happens when Algorithm 1 converges.
Simplified version. Figure 1 (b) shows the lerning curves of Algorithm 1 and Algorithm 2, where θ is initialized by the learning result of the ERM method, and wπθ (sn, an) is normalized to satisfy 1 N ∑ n w π θ (sn, an) = 1. The results are averaged over 5 runs, and the errorbar is the standard deviation. The simplified version also reduces the weighted loss function, compared to the ERM method, i.e. the start of the iteration. This result suggests that, although the simplified version has the limitation of not converging, it can be used as better than ERM. The simplified version can also predict the swing-up behavior, similarly to the full version, see Appendix.

heading: 6.2 D4RL MuJoCo Benchmark
text: Secondly, to demonstrate weighting with the artificial weight on large-scale environments, this paper studies policy optimization on the D4RL Benchmark [33] based on the MuJoCo simulator [34]. This paper shows the results of Algorithm 3 with α = 0 and α = 0.2. This paper regards the case of α = 0 as a MBRL baseline, as it can be seen as an instantiation of MOPO [8]. Appendix discusses the result when using different values for α for the Walker2d-medium-expert dataset.
The outline of the implementation is as follows. The agent iterates the main loop of Algorithm 3 twice. The agent uses the simplified version of the E-step. The agent estimates Pr(s′, r|s, a) instead of P(s′|s, a), to handle the unknown reward function, as in [8]. The agent represents Pr(s′, r|s, a) using an ensemble of models, where each model is parametrized as 7-layer neural networks with 256 hidden units. The agent estimates wπθ in the same way as in Section 6.1. The agent estimates cθ using L2 norm minimization. For the M-step, the agent improves π using SAC [35], where samples are drawn from D ⋃ D̂πθ . This paper implements SAC by modifying the implementation code by [36]. To stabilize the M-step, the agent terminates a simulation episode when a state or action variable diverges. In addition, the agent clips a penalized reward so that the absolute value is not more than 10 times the maximum absolute value of the reward data. For more details, see Appendix.
Table 1 shows the normalized scores of Algorithm 3, where the results are averaged over 5 runs, and ± means the standard deviations. The major differences are found for Walker2d-medium-expert and Walker2d-medium datasets. The normalized scores of α = 0.2 is better for the former and worse for the latter. Below, this paper describes the reasons.
Walker2d-medium-expert dataset. The case of α = 0.2 moves forward without falling for all runs, while the case of α = 0 falls down for two runs. Figure 2 (a) shows the learning curves of real and simulated returns in the M-step at the second iteration for the seed for that run. The case of α = 0.2 obtains similar curves for real and simulation, while the case of α = 0 does not. This difference comes from the weighted model estimation. In the case of α = 0, the model for the first move is not sufficiently accurate, and this lead to falling at the first move. In the case of α = 0.2,
simulated future data includes predictions of the first move, and similar parts of offline data are weighted with large weights. Based on the weights, the estimated model more accurately predicts the first move. As a result, the case of α = 0.2 can learn to move forward without falling, based on the more accurate prediction. That is, Algorithm 3 improves the performance based on weighted model estimation. For the other runs, although the case of α = 0 moves forward without falling in a real environment, it obtains real and simulation returns that are relatively different, see Appendix.
Walker2d-medium dataset. The case of α = 0.2 falls down for two runs, while the case of α = 0 moves forward without falling. Figure 2 (b) shows the learning curves fo those runs. In Figure 2 (b), for both α = 0 and α = 0.2, the simulation curves do not seem to capture the real returns. The difference of the score seems due to the high variance of the real returns.

heading: 7 Limitations
text: This paper has limitations common to standard offline MBRL approaches, such as Markov assumption and reward design. The limitations specific to this paper are as follows. This paper derives Algorithm 1 based on the theory, but it cannot be applied to large-scale problems due to the large amount of computations. This paper also presents Algorithm 2 as a practical simplified version, but it has no convergence with respect to Equation (3), see Section 5.1.

heading: 8 Conclusion and Future Direction
text: This paper discusses weighted model estimation considering covariate shift in offline MBRL. The main idea of this paper is weighting with the artificial weight, which is relatively easy to estimate. Based on this idea, this paper defines a loss function for offline MBRL and presents algorithms to optimize it. This paper justifies this idea as evaluating the upper bound of the policy evaluation error. This paper demonstrates the effectiveness of this idea in numerical experiments. The main contribution of this paper is showing the validity of weighting with the artificial weight.
The propose algorithm shows clear improvement in Section 6.1, while it improves only one case in Section 6.2. One reason for this gap might be that Section 6.1 needs interplolation, while Section 6.2 requires extrapolation. Although interpolation and near extrapolation can be addressed by importanceweighting, far extrapolation would require additional idea, e.g. using a structured prior. Since further discussion is needed, we leave it a future work.
Another interesting direction is extending to Bayesian MBRL [38, 39]. In supervised learning under covariate shift, the posterior distribution can be defined using a weighted likelihood function [9]. In offline MBRL, the likelihood function weighted with the artificial weight is effective, as shown in this paper. Combining these ideas, a weighted posterior considering covariate shift will be introduced into Bayesian MBRL.
Acknowledgments and Disclosure of Funding
This work was partly supported by Collaborative Research for Evolutionary Mechanical System Technology with Mitsubishi Electric Corporation and by General Research Grant from SECOM Science and Technology Foundation. We thank all anonymous reviewers for their constructive comments.

Review summaries:
The paper presents a model-based policy optimization algorithm for offline reinforcement learning. To account for the mismatch between the available dataset and the experience used during training, the proposed method uses estimation of the density ratio of the state distributions induced by the estimated and the real model. This importance weight, called “artificial”, is in contrast with the one, referred to as “natural”, employed by other off-policy methods involving the densities for two different policies under a same model. The resulting model learning objective is inserted into an expectation-maximization algorithm, which is then evaluated in a standard offline RL benchmark.

This work studies how to estimate a learned model in offline model-based reinforcement learning (MBRL). The authors propose to use "the artificial weight," which is the ratio of offline data and simulated future data, to improve the model estimation. They further propose a new objective function for offline MBRL and an EM-style algorithm to optimize it. The authors evaluate their algorithm in several different tasks.

The paper addresses the problem of covariate shift in offline memory-based reinforcement learning by proposing a novel method for more accurate evaluation of the expected return of a given policy by means of simulation of a learned model of the real system that might not be entirely correct. Two algorithm that uses a weighting factor, called the artificial weight, are proposed for tuning the learned system's model to predict more accurately the parts of the state space visited by the policy being evaluated, and thus improve the estimate of the expected return of this policy. This improved estimate can further be used for policy improvement, as customary for RL algorithms. An empirical evaluation in simulated environments demonstrates the validity of the proposed approach, both for policy evaluation and for policy improvement.   Update: I appreciate the clarification about the evaluation task the authors are solving (pendulum swing-up, not stabilization). I also noticed that the other reviewers also found the advantages of the proposed algorithm not all that overwhelming. This, combined with the generally unclear presentation, makes the paper borderline. I am revising my rating to 6, just above borderline - it would still be good to accept this paper, if possible, and depending on the ratings of other papers.

This paper presents a model-based algorithm for offline reinforcement learning (RL) that attempts to address the covariate shift issue by re-weighting the model losses for different datapoints. Based on an upper bound of the policy evaluation error, a joint loss function is defined over both the policy and the model parameters, and the two are optimized in an alternating fashion. Several simplifications and tricks are described to make the algorithm computationally tractable.  On a simple pendulum environment, the proposed algorithm predicts the policy value much more accurately than a model based on empirical risk minimization (ERM). It also obtained good results on the MuJoCo tasks from the D4RL benchmark.