Weighted model estimation for offline model-based reinforcement learning

1 Introduction . 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 . 2 Related Works . 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 . 3 Preliminaries . 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 ) . 4 Loss Function . 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 . 4.1 Policy Evaluation . 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 . 4.2 Policy optimization . 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∗ .

Summary: 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.