L IVE IN THE M OMENT : L EARNING D YNAMICS M ODEL A DAPTED TO E VOLVING P OLICY
Anonymous authors
Paper under double-blind review
A BSTRACT
Model-based reinforcement learning (RL) often achieves higher sample efficiency in practice than model-free RL by learning a dynamics model to generate samples for policy learning.
Previous works learn a dynamics model that fits under the empirical state-action visitation distribution for all historical policies, i.e., the sample replay buffer.
However, in this paper, we observe that fitting the dynamics model under the distribution for all historical policies does not necessarily benefit model prediction for the current policy since the policy in use is constantly evolving over time.
The evolving policy during training will cause state-action visitation distribution shifts.
We theoretically analyze how this distribution shift over historical policies affects the model learning and model rollouts.
We then propose a novel dynamics model learning method, named Policy-adapted Dynamics
Model Learning (PDML) .
PDML dynamically adjusts the historical policy mixture distribution to ensure the learned model can continually adapt to the state-action visitation distribution of the evolving policy.
Experiments on a range of continuous control environments in MuJoCo show that PDML achieves significant improvement in sample efficiency and higher asymptotic performance combined with the state-of-the-art model-based RL methods.
I NTRODUCTION
Recent years have witnessed great successes of Reinforcement Learning (RL) in many complex decision-making tasks, such as robotics (Polydoros & Nalpantidis, 2017;
Yang et al., 2022) and chess games (Silver et al., 2016; Schrittwieser et al., 2020).
Among RL methods, a wide range of works in model-free RL (Schulman et al., 2015; Lillicrap et al., 2016; Haarnoja et al., 2018; Fujimoto et al., 2018; Hu et al., 2021) have shown promising performance.
However, model-free methods can be impractical for real-world scenarios (Dulac-Arnold et al., 2021) since massive samples from the real environment are required for policy training, resulting in low sample efficiency.
Model-based RL is considered one of the solutions to improve sample efficiency.
Most of the model-based RL algorithms first use supervised learning techniques to learn a dynamics model based on the samples obtained from the real environment, and then use this learned dynamics model to generate massive samples to derive a policy (Luo et al., 2018; Janner et al., 2019).
Therefore, it is crucial to learn a dynamics model which can accurately simulate the underlying transition dynamics of the real environment since the policy is trained based on the model-generated samples.
If the learned dynamics has a high prediction error, the model-generated samples will be biased, and the policy induced by these samples will be sub-optimal.
To reduce the model prediction error and learn an accurate dynamics model, some advanced architectures such as model ensemble (Kurutach et al., 2018; Chua et al., 2018) and multi-step model (Asadi et al., 2019) have been proposed to improve the multi-step prediction accuracy of the learned dynamics model.
Besides, the idea of a generative adversarial network (GAN) (Goodfellow et al., 2014) is used to design the training process of a dynamics model (Shen et al., 2020;
Eysenbach et al., 2021) to reduce the distribution mismatch between model-generated samples and real samples.
Those previous works mentioned above aim to learn a dynamics model that can fit all historical policies.
To be precise, when training the dynamics model, they randomly select the training data from the real samples obtained by all historical policies in the replay buffer.
This learned dynamics model needs to adapt to the state-action visitation distribution of all historical policies to obtain a dynamics model that predicts transitions accurately under different policies.
However, since we only use the current newest policy to interact with the learned model to generate samples for policy learning during model rollouts, learning such a dynamics model that fits under (highly likely sub-optimal) historical policies may be unnecessary.
Due to the state-action visitation distribution shift during policy updating, the state-action pairs visited by historical policies may not appear in the state-action visitation distribution of the current policy, and vice versa.
Thus, learning these samples may not benefit model rollouts.
Besides, in many complex tasks, it is hard to predict all samples from all historical policies due to limited model capacity (Abbas et al., 2020), and as shown later in our paper, trying to learn every sample from historical policies can even hurt the accuracy when predicting the transitions induced by the current policy.
Therefore, there is an objective mismatch between model learning and model rollouts — model learning tries to fit samples from state-action visitation distribution for all historical policies, whereas model rollouts require accurate prediction of the transitions induced by the current policy.
In this paper, we investigate how to learn an accurate dynamics model for model rollouts based on existing samples.
(a) To begin with, we confirm through experiments that although the dynamics model learned by the previous methods has a low overall prediction error on all transitions obtained by historical policies, its prediction error for the current newest policy can still be very high.
This leads to inaccurate model-generated samples which can hurt the sample efficiency and asymptotic performance of the policy.
(b) We then derive an upper bound of the expected performance gap between the model rollouts and real environment rollouts.
According to this upper bound, we analyze how the distribution of historical policies affects model learning and model rollouts.
The theoretical result suggests that the historical policy distribution used for model learning should be more inclined towards policies that are closer to the current policy rather than a uniform distribution over all historical policies to ensure the model prediction accuracy for model rollouts.
(c) Motivated by this insight, we propose a novel dynamics model learning method named Policy-adapted Dynamics Model Learning (PDML) .
Instead of learning a dynamics model that fits under a uniform mixture of all historical policies, PDML adjusts the historical policy distribution by reducing the total variation distance between the historical policy mixture and the current policy, then learns a policy-adapted dynamics model according to this adjusted historical policy distribution.
(d) We conduct systematic and extensive experiments on a range of continuous control benchmark MuJoCo environments (Todorov et al., 2012).
Experimental results show that PDML significantly improves the sample efficiency and asymptotic performance of the state-of-the-art model-based RL methods.
Summary of contributions: (1) Through detailed experimental results, we establish that learning a dynamics model that fits a uniform mixture of all historical policies may not be accurate enough for model rollouts.
(2) We propose an upper bound of an expected performance gap between the model rollouts and the real environment rollouts, and theoretically analyze how the distribution over historical policies affects model learning and model rollouts.
(3) We propose Policy-adapted
Dynamics Model Learning (PDML) , which dynamically adjusts the distribution over the historical policy sequence and allows the learned model to continuously adapt to the evolving policy.
Experimental results on a range of MuJoCo environments demonstrate that PDML can achieve significant improvement in sample efficiency and higher asymptotic performance combined with the state-of-the-art model-based RL methods.
B ACKGROUND 2. P RELIMINARIES
Reinforcement learning.
Consider a Markov Decision Process (MDP) defined by the tuple ( S , A , T, r, γ ) , where S is the state space, A is the action space, and T ( s ′ | s, a ) is the transition dynamics in the real world.
The reward function is denoted as r ( s, a ) and γ is the discount factor.
Reinforcement learning aims to find an optimal policy π which can maximize the expected sum of discounted rewards
[Equation]
In model-based RL, the transition dynamics T in the real world is unknown, and we aim to construct a model ˆ T ( s ′ | s, a ) of transition dynamics and use it to improve the policy.
In this paper, we concentrate on the Dyna-style (Sutton, 1990) model-based RL, which uses the learned dynamics model to generate samples and train the policy.
Policy mixture.
During policy learning, we consider the historical policies at iteration step k as a historical policy sequence Π k = { π 1 , π 2 , ..., π k } .
For each policy in the policy sequence, we denote its state-action visitation distribution as ρ π i ( s, a ) , and the policy mixture distribution over the policy sequence as w k = [ w k 1 . . . , w kk ] .
Then the state-action visitation distribution of the policy mixture π mix ,k
= (Π k , w k ) is ρ π mix ,k
( s, a ) = (cid:80) ki
=1 w ki ρ π i
( s, a ) (Hazan et al., 2019;
Zhang et al., 2021).
2. D YNAMICS M ODEL L EARNING IN M ODEL - BASED RL
Learning a dynamics model is the most crucial part of model-based RL since the ground-truth transition dynamics is unknown and the policy must be updated based on the samples generated by the learned dynamics model.
Previous works learn the dynamics model by randomly selecting training data from the samples obtained by the historical policy sequence Π k , which means the distribution of policy mixture is a random distribution: w ki = 1 k .
The learned dynamics model is trained based on the following state-action visitation distribution
[Equation]
This model tries to fit all the samples obtained by sampling the state-action visitation distribution corresponding to all policies in the historical policy sequence, so the learned dynamics model is (hopefully) able to predict the transition for any state-action input.
However, as shown in Figure 1a and 1b, since the policy is constantly evolving, the state-action visitation distribution of historical policies may have a huge shift from the current policy.
There is little overlap between the state-action visitation distribution of policies at different environment steps.
The state-action pairs visited by historical policies may not appear in the state-action visitation distribution of the current policy.
During model rollouts, we only use the current policy to interact with the learned dynamics model to generate samples.
Thus, learning these samples may not benefit model rollouts.
When the model capacity is not large enough, learning these samples may even be detrimental to the learning of the samples collected by the current policy.
[Figure]
in We conduct an experiment using a state-of-the-art model-based RL method called MBPO (Janner et al., 2019) on four MuJoCo (Todorov et al., 2012) environments HalfCheetah, Hopper, Walker2d, and Ant.
MBPO first trains a model based on the real samples and then uses the model to roll out multiple samples for policy learning.
The architecture of the dynamics model is a 4-layer neural network with a hidden size of 200, which is a very common architecture used in many recent modelbased methods (Yao et al., 2021; Froehlich et al., 2022; Li et al., 2022).
We present the overall error curves and the current error curves during learning steps on HalfCheetah and Hopper in Figure 1c and 1d.
Here the overall error means the model prediction error for all historical policies during training.
It is evaluated on an evaluation dataset which contains 1000 × N samples from the real environment.
N is the number of historical policies in the historical policy sequence.
The current error is the model prediction error for the current policy, which is evaluated using L2 error on the 1000 samples obtained by the current policy from the real environment.
The error curves for more environments can be found in Appendix F.1.
From Figure 1c and 1d, we observe that there is a gap between the overall error and the current error.
This means although the agent can learn a dynamics model which is good enough for all samples obtained by historical policies, this is at the expense of the prediction accuracy for the samples induced by current policy.
Since we only use the current policy during model rollouts, this will lead to inaccurate model-generated samples and misleading policy learning.
Therefore, learning a dynamics model that adapts the state-action visitation distribution for all historical policies, in other words, a random historical policy mixture distribution used for model learning, is not the most efficient way for model-based RL (especially for task-specific problems).
In the next section, we will analyze how the policy mixture distribution affects the performance of model-based RL.
P ERFORMANCE G AP I NFLUENCED BY P OLICY M IXTURE D ISTRIBUTION
In this section, we provide a theoretical analysis of how the policy mixture distribution affects the performance of model-based RL.
First, we derive a theorem that upper bounds the performance gap between the real environment rollouts and the model rollouts under any current policy π .
Theorem 3.1 Given the historical policy mixture π mix ,k
=
(Π k , w k ) at iteration step k , we denote ξ ρ i
= D TV ( ρ πT ( s, a ) ||
ρ π i T ( s, a ))
as the state-action visitation distribution shift and ξ π i
=
E s ∼ v π mix ˆ T [ D TV ( π ( a | s )
||
π i ( a | s ))] as the policy distribution shift between the historical policy π i and current policy π respectively, where v π mix ˆ T is the state visitation distribution of the policy mixture under the learned dynamics model ˆ T .
r max is the maximum reward the policy can get from the real environment, γ is the discount factor, and Vol( S ) is the volume of state space.
Then the performance gap between the real environment rollout J ( π, T ) and the model rollout J ( π, ˆ T ) can be bounded as:
[Equation]
Proof.
See Appendix D.
□
Remarks.
(1) The first term is about model prediction error .
This term suggests that the model needs to adapt to the state-action visitation distribution of the current policy to reduce the model prediction error, since this term is the expectation of prediction error of the learned dynamics model ˆ T under the current policy state-action visitation distribution ρ πT .
The second term shows the effect of
the
policy mixture distribution on
model
rollout
This item
contains two distribution shifts:
(2a) state-
action visitation distribution shift ξ ρ i and
(2b
) policy distribution shift ξ π i between the historical policy and current policy.
It should be noted that ξ ρ i is induced by ξ π i , so it is reasonable to believe that a historical policy with a larger ξ π i will have a larger ξ ρ i .
Both ξ ρ i and ξ π i are fixed since historical policies and the current policy are immutable during model learning and model rollout.
Therefore, to reduce this term, we can only adjust the policy mixture distribution w k .
Since the distribution shift varies across
historical
policies and the current
policy, it is obvious that the random distribution w ki =
k is not
the best choice.
The last term is related to the model sample buffer, which is used for policy learning .
To maximize sample utilization, the model-generated samples obtained by the historical policies will be maintained in the model sample buffer until they are replaced by the new samples generated by the current policy.
Therefore, the distribution of simulated samples in the model buffer is not exactly the same as the simulated sample distribution of the current policy, but is often mixed with the simulated sample distribution of the historical policies.
This makes it necessary to adjust the sample distribution in the model sample buffer to make it close to the simulated sample distribution of the current policy during the policy learning process.
This has been studied in many model-based and model-free methods (Schaul et al., 2016; Liu et al., 2021; Huang et al., 2021; Mu et al., 2021) and is out of the scope of this paper, and we focus on reducing the first two terms related to model learning.
The first two items on the right-hand side of Equation (3) provide useful insights on model learning.
This first term points out the goal of model learning: to make accurate predictions for the current policy.
The second item further demonstrates that to achieve this goal, we should adjust the policy mixture distribution to reduce the distribution shift between the historical policy mixture and the current policy.
According to the second term, we have the following proposition.
Proposition 3.2 The performance gap can be reduced if the weight w ki of each policy π i in the historical policy sequence Π k is negatively related to state action visitation distribution shift ξ ρ i and the policy distribution shift ξ π i between the historical policy π i and current policy π instead of an average weight w ki = 1 k .
The proof is in Appendix E. Proposition 3.2 illustrates how we should adjust the policy distribution to help the learned dynamics model adapt to the current policy.
This naturally motivates our method, which is described in the next section.
P OLICY - ADAPTED D YNAMICS M ODEL L EARNING
In this section, we introduce our model learning method called Policy-adapted Dynamics Model
Learning (PDML).
PDML is designed to reduce the model prediction error during model rollouts, and it contains two parts.
The first part is adjusting the policy mixture distribution into a nonuniform distribution, and the second part is learning the dynamics model based on this non-uniform distribution.
The pseudo-code is in Algorithm 1.
Algorithm 1 Policy-adapted Dynamics Model Learning (PDML)
Require: current policy proportion hyperparameter α , interaction epochs I 1:
Initialize historical policy sequence k ← 0 , Π k ← ∅ 2: for I epochs do 3:
Interact with the environment using current policy π c , add samples into real sample buffer D e 4:
Add current policy π c into historical policy sequence: π k ← π c , Π k ← { Π k − 1 , π k } 5:
Adjust the historical policy mixture distribution w k =
[ w k 1 , . . . , w kk ] via Equation (4) and (5) 6:
Normalize w k ← w k / ∥ w k ∥ 7:
Sample a training data batch of ( s n , a n , r, s n +1 ) from D e according to w k 8:
Train dynamics model ˆ T θ via Equation (7), use current policy π c to perform model rollouts 9: k ← k +10: end for
4. P OLICY M IXTURE D ISTRIBUTION A DJUSTMENT
In this section, we introduce a mechanism to adjust the policy mixture distribution.
According to our Theorem 3.1, to minimize the performance gap, one may set the weight of the policy with the smallest ξ ρ i and ξ π i to be 1 and the weights of other policies in the historical policy sequence to be 0.
However, this is not the best approach in practice since each policy can only interact with the environment for very few steps in model-based RL.
This means each policy can provide very limited samples for model learning.
If we only use a small number of samples from just one policy, it is difficult to learn accurate transition dynamics for the current policy.
Weights design for historical policies.
In order to maximize the use of limited samples to estimate the transition dynamics, inspired by Proposition 3.2, we design the weight of each policy in the historical policy sequence Π k
= { π 1 , π 2 , ..., π k } except for the current policy π c (i.e., π k ∈ Π k ) as follows:
[Equation]
where ξ π i is the policy distribution shift between historical policy π ki and the current policy π c ; it is also one of the distribution shifts in the second term of Equation (3).
We use [ k − 1] := { 1 , . . .
, k − 1 } to denote the integers from 1 to k − 1 .
We only use the policy distribution shift ξ π i (and not the state-action visitation distribution shift ξ ρ i ) because estimating the state-action visitation distribution shift using limited real samples is difficult, and thus the estimation may be inaccurate.
Besides, as mentioned in the remarks of Theorem 3.1, state-action visitation distribution is induced by the policy, so it is reasonable to believe that a historical policy with a larger ξ π i will have a larger ξ ρ i .
Weight design for the current policy.
In model-based RL, the current policy becomes a historical policy after interacting with the environment and is added to the historical policy sequence (see
Algorithm 1).
The total variation distance between the current policy and itself is 0, so Equation (4) cannot be used to calculate the weight of the current policy.
For the weight of the current policy w kk , we use the following equation:
[Equation]
where α is a hyperparameter to control the proportion of the weight of the current policy to the total weight over the historical policy sequence.
Equation (5) ensures that the weight of the current policy w kk is always the largest in the historical policy sequence.
Before each model learning iteration, we adjust the policy mixture distribution according to Equation (4) and Equation (5) and normalize the weights w k =
[ w k 1 , ..., w kk ] to make sure they sum to 1.
The details are illustrated in Algorithm 1.
Estimation of the policy distribution shift ξ π i ∀ i ∈ [ k − 1] .
Given a state s n , we define the output of policy π i as a multivariate Gaussian distribution N ( µ π ni , Σ π ni ) .
In order to make the empirical estimation more accurate, we use each historical policy to traverse all N samples in the real sample buffer and output the action distribution corresponding to each state.
Then we use the inequality between KL divergence and total variation distance to estimate ξ π i :
[Equation]
Novelty of PDML compared to prioritized experience replay proposed in model-free RL.
In model-free RL, prioritized experience replay methods only need to consider how to improve the policy based on existing samples.
Therefore, it is only necessary to select the sample that can bring the greatest improvement to the policy, and a weighting is designed for each sample.
In model-based RL, the policy is learned based on model-generated samples, and the accuracy of these model-generated samples determines the sub-optimality of the policy.
Thus, in the model-learning part, we focus on the model prediction accuracy.
Our theoretical analysis shows that we should consider whether the state-action visitation distribution that generates the samples is close to the current policy when reweighting samples.
Although a sample can bring a great improvement to the current policy (the
TD value is high), if this sample is not in the state-action visitation distribution of the current policy, this sample will not be encountered during model rollouts.
Then learning this sample will not bring any benefit to model learning and policy learning.
Therefore, we reweight the state-action visitation distribution that generates a batch of samples according to ξ π i , rather than a single sample as in model-free RL.
4. D YNAMICS M ODEL L EARNING
After adjusting the policy mixture distribution, we learn the dynamics model based on this adjusted distribution.
Although our method can be applied to learn any type of dynamic model, here we choose to use the current state-of-the-art structure probabilistic dynamics model ensemble Chua et al.
(2018): { ˆ T 1 θ , ..., ˆ T Bθ } .
θ is the parameters of each dynamics model in the ensemble, and B is the ensemble size.
Given a ( s n , a n ) pair as an input, the output ˆ T bθ of each network b in the ensemble is the Multivariate
Gaussian Distribution of the next state: ˆ T bθ ( s n +1 | s n , a n ) =
N ( µ bθ ( s n , a n ) , Σ bθ ( s n , a n ))
Before each model learning iteration, we sample the training data batch from the real sample buffer according to the adjusted policy mixture distribution w k , and train the dynamics model using maximum likelihood:
L ( θ ) =
N (cid:88)
[Equation]
During model rollouts, we use the current policy π c as the rollout policy and sample the initial states from the real sample buffer according to the adjusted policy mixture distribution w k .
E XPERIMENT
In this section, we will first compare our method with the previous state-of-the-art (including both model-free and model-based) baselines.
We demonstrate that after combining with SOTA modelbased method, PDML improves SOTA sample efficiency and SOTA asymptotic performance for
0k 50k 100k 150k 200k2 ,4 ,6 ,
Steps
R e t u r n
Walker2d 0k 30k 60k 90k 120k1 ,2 ,3 ,4 ,
Steps
Hopper 0k 75k 150k 225k 300k2 ,4 ,6 ,
Steps
Humanoid 0k 50k 100k 150k 200k0 . 1 . 5 ·
Steps
HalfCheetah
[Figure]
model-based RL.
Then we compare our method with three SOTA prioritized experience replay methods to indicate the advantage of our distribution adjustment method for model learning.
Lastly, we conduct a systematic ablation study to analyze the model errors of PDML.
5. C OMPARISON WITH S TATE - OF - THE - ARTS
In this section, we compare our method with several previous state-of-the-art (SOTA) baselines.
For model-based methods, we choose MBPO (Janner et al., 2019), AMPO (Shen et al., 2020), and
VaGraM (Voelcker et al., 2022).
MBPO is the SOTA model-based method, and our method is combined with MBPO for the model learning part.
We name our method PDML-MBPO and we provide the pseudo code in Appendix A. AMPO is another SOTA model-based method that uses unsupervised model adaptation during model learning to reduce the prediction error.
VaGraM is a
SOTA value equivalence model-based method.
Instead of accurately learning each dimension in the dynamics, it aims to learn the dimensions which impact policy learning most.
In other words, this method also learns a locally accurate model.
Both AMPO and VaGraM are implemented based on MBPO.
PDML-MBPO, AMPO, and VaGraM share the same model architecture and policy part; only the model learning part is different.
For model-free methods, we compare with two methods.
The first one is SAC (Haarnoja et al., 2018), which is the policy part of all model-based and model-free baselines we used and is one of the SOTA model-free methods.
The second one is REDQ (Chen et al., 2020), which improves the Update-To-Data (UTD) ratio of the model-free method and achieves higher sample efficiency than SAC.
The implementation details of our method are in Appendix G.1.
We conduct experiment on four complex MoJoCo-v2 (Todorov et al., 2012) environments, the performance curves are shown in Figure 2.
Results: (1) Improving SOTA sample efficiency.
PDML-MBPO outperforms all existing state-ofthe-art methods, including model-based and model-free, in sample efficiency in all four environments.
In Hopper, Walker2d, and Humanoid, PDML-MBPO achieves very impressive sample efficiency improvements, up to a 2 × improvement in Hopper and Humanoid compared to the SOTA model-based methods.
For example, our method using only 30k steps to achieve 3000 while other model-based methods need almost 60k steps.
Besides, its sample efficiency is also higher than REDQ which is modified for sample-efficient model-free RL.
(2) Improving SOTA asymptotic performance for
Model-based RL.
In addition, PDML-MBPO obtains significantly better asymptotic performance compared to other state-of-the-art model-based methods.
It is worth noting that the asymptotic performance of PDML-MBPO is very close to SAC in three environments (Hopper, Walker2d, and
Humanoid) and is even better than SAC occasionally.
Furthermore, our method achieves impressive improvement in the most complex environment Humanoid.
These indicate the effectiveness of our proposed model learning method.
Discussion of computational cost.
PDML requires saving all historical policies as well as computing their distances to the current policy for adjusting their weights (as shown in Equation 6).
This creates an additional memory overhead of storing historical policy networks ( k × policy network size) and an additional computational overhead of computing the distances, for each iteration k .
In PDML-MBPO, we observe storing historical policy networks costs a memory overhead of no more than k × 1 MB, compared with the high memory occupied by the model sample buffer, this cost is very small.
Besides,
Compared to MBPO, the training time of PDML-MBPO does not increase significantly, we provide the training time on four environments in Appendix F.8.
5. C OMPARISON WITH M ODEL - FREE E XPERIENCE R EPLAY M ETHODS
We compare with the other three prioritized experience replay methods in model-free RL to indicate the advantage of our PDML.
The first one is Prioritized Experience Replay (PER) (Schaul et al., 2016), which weighs the samples according to their TD-error.
The second method is RECALL (Goyal et al., 2018), which chooses the top k highest value sample.
They use this to recall the samples that induce the high-value trajectories and train the policy.
We implement this by choosing the top 25% highest Q value samples to train the model and as model rollout initial states.
The third method is
Model-augmented Prioritized Experience Replay (MaPER) (
Oh et al., 2022), which is an extension of PER using both TD-error and model prediction error to weight the samples for model learning.
[Figure]
The experiment results are shown in Figure 3a and 3b.
Our PDML significantly outperforms all three methods on both sample efficiency and asymptotic performance.
We believe these methods adjust the weights for each sample in the training data rather than each policy.
This will cause the samples belonging to the same stateaction visitation distribution to have different weights, and the samples with higher weights may not necessarily appear in the state-action visitation distribution of current policy.
Therefore, the learned model cannot be adapted to current policy’s state-action visitation distribution, and the model prediction error during model rollouts cannot be reduced.
In the model learning process, it is crucial to adapt to current policy’s state-action visitation distribution according to our theory.
This experiment result indicates our theory’s correctness and our method’s effectiveness.
[Figure]
5. M ODEL E RROR A NALYSIS
To further verify the impact of PDML, we compare the one-step prediction error and the compounding error of the policy-adapted model learned by PDML-MBPO and the original dynamics model learned by MBPO.
One-step prediction error.
As shown in Figure 4a, 4b and 4c, we evaluate the model prediction error for the current policy on Hopper, HalfCheetah, and Walker2d.
We evaluate the learned model every 1000 environment steps using L2 error on the samples obtained by the current policy from the real environment.
The error curves show that the one-step prediction error for the current policy of the policy-adapted model is much smaller than that of the original dynamics model, which means the model-generated samples of PDML-MBPO are more accurate than MBPO, so the policy induced by PDML-MBPO can perform better.
Compounding error.
We also compare the multi-step model rollouts compounding error of the policy-adapted model and the original dynamics model.
This directly determines the accuracy of the model-generated samples in each model rollout trajectory.
Figure 4d shows the compounding error curves of the policy-adapted model and the original dynamics model on Hopper.
We calculate the h -step compounding error as the difference between the state at each rollout step h in the model rollout trajectory and the real environment rollout trajectory using L2 error.
The results demonstrate that the policy-adapted model has much a smaller compounding error than the original dynamics model, which means the policy-adapted model has a more robust multi-step planning capability than the original dynamics model learned by MBPO.
R ELATED W ORK
Model adaptation.
Several adaptive control approaches (Sastry & Isidori, 1989; Pastor et al., 2011; Meier et al., 2016) aim to train a dynamics model that can adapt online.
However, scaling such methods to complex tasks is exponentially difficult.
Adaptive learning in the dynamics model has also been studied in inverse dynamics learning tasks.
A drifting Gaussian process (GP) keeps a history of a constant number of recently observed data points and updates its hyper-parameters at each time step (Meier & Schaal, 2016).
The drifting Gaussian process (GP) predicts the local dynamics errors to control the learning rate (Meier et al., 2016), resulting in more online hyperparameter learning and adaptive function approximator robustness.
Our method is different from these works that we learn a forward model which can always adapt to the evolving policy.
Some studies focus on an adaptive model predictive control for constrained linear systems (Tanaskovic et al., 2013) and guaranteeing safety, robustness, and convergence in a quadrotor helicopter testbed (Aswani et al., 2012).
Our work closely relates to a model adaptation in forward models from Fu et al.
; Nagabandi et al.
(2018a;b); Lee et al.
Guo et al.
These methods use meta-learning to train a dynamics model as a prior and then combine it with recent data to rapidly adapt to the new task.
However, these works are mainly about model transfer under different dynamics.
Different from their works, we study learning an accurate dynamics model for policy learning under a fixed transition dynamics, and we also provide theoretical analysis to motivate our method.
More related works about model-based
RL are provided in Appendix B.
Prioritized experience replay.
Another related line of work is prioritized experience replay in reinforcement learning.
This solves a classic issue in model-free RL.
Previous work Katharopoulos &
Fleuret (2018) claimed that emphasizing essential samples in the replay buffer can benefit off-policy
RL algorithms.
Prioritized Experience Replay (PER)
Schaul et al.
(2016) measured the importance of sample by temporal-difference (TD) error.
Based on this work, many methods are proposed to perform prioritized sampling.
Some methods (Brittain et al., 2019; Lee et al., 2019; Fujimoto et al., 2020; Jiang et al., 2021;
Liu et al., 2021; Lahire et al., 2021;
Oh et al., 2022) extend or explain PER from different perspectives, and others (Novati & Koumoutsakos, 2019; Fedus et al., 2020) propose to prioritize samples according to their age.
Our work is different from experience replay works in model-free RL in the following points: (1) In model learning, we re-weight the state-action visitation distribution that generates a batch of samples, rather than a single sample as in model-free RL.
During weighting, we use the distance between the policy distribution that each sample generated from and the policy distribution of the current policy as a metric, rather than how much improvement each sample can bring to the policy.
(3) We provide very detailed theoretical result to analyze how to reweight the samples for model learning.
C ONCLUSION AND D ISCUSSION
In this paper, we introduce a novel dynamics model learning method for model-based RL called
PDML, which learns a policy-adapted dynamics model based on a dynamically adjusted historical policy mixture distribution.
This policy-adapted dynamics model can continually adapt to the stateaction visitation distribution of the evolving policy.
This makes it more accurate than the previous dynamics model when making predictions during model rollouts.
We also provide theoretical analysis and experimental results to motivate our method.
After combining with the state-of-the-art model-based method MBPO, PDML achieves better asymptotic performance and higher sample efficiency than previous state-of-the-art model-based methods in MuJoCo.
We believe our work takes an important step toward more sample-efficient RL.
One limitation of our work is that the generalization ability of the policy-adapted dynamics model may not be strong enough because we focus on fitting the samples induced by the evolving policy to improve the convergence speed of the policy.
Therefore, our method is efficient for task-specific problems but may not perform well for some exploration-oriented tasks.
We leave this direction to future work.
Supplementary Material A P SEUDO C ODE OF PDML-MBPO
In Algorithm 2, we demonstrate the pseudo code of PDML-MBPO.
Algorithm 2 PDML-MBPO
Require: current policy proportion hyperparameter α , interaction epochs I , rollout horizon h 1: Initialize historical policy sequence k ← 0 , Π k ← ∅ 2: for I epochs do 3:
Interact with the environment using current policy π c , add samples into real sample buffer D e 4:
Add current policy π c into historical policy sequence: π k ← π c , Π k ← { Π k − 1 , π k } 5:
Adjust the historical policy mixture distribution w k =
[ w k 1 , . . . , w kk ] via Equation (4) and (5) 6:
Normalize w k ← w k / ∥ w k ∥ 7:
Sample a training data batch of ( s n , a n , r, s n +1 ) from D e according to w k 8:
Train dynamics model ˆ T θ via Equation (7) 9: for M model rollouts do 10:
Sample initial rollout states from real sample buffer D e according to w k 11:
Use current policy π c to perform h -step model rollouts, add model-generated samples into model sample buffer D m 12: end for 13: for G gradient updates do 14:
Update current policy π c using model-generated samples from model sample buffer D m 15: end for 16:
k ← k +17: end for
B A DDITIONAL R ELATED W ORK
Model-based reinforcement learning.
Model-based RL is proposed as a solution to reduce the sample complexity of model-free RL by learning a dynamics model.
Current model-based RL mainly focuses on better model learning and better model usage.
To learn a model with more accuracy, many model architectures have been proposed, such as linear models Parr et al.
; Sutton et al.
Kumar et al.
and nonparametric Gaussian processes Rasmussen & Kuss (2004)
; Deisenroth & Rasmussen (2011).
With the rapid development of deep learning, neural networks have become a popular choice of model architecture in recent years Kurutach et al.
; Chua et al.
Moreover, to reduce the model error, a multi-step model Asadi et al.
(2019) was designed to directly predict the transition of an action sequence input, and Shen et al.
(2020) used unsupervised model adaptation to reduce the potential data distribution mismatch.
For better model usage, Janner et al.
(2019) proved that short model rollouts could avoid the model error and improve the quality of model samples.
Based on this, Lai et al.
(2020) proposed a bidirectional model rollout scheme to avoid the model error further.
Furthermore, model disagreement was used to decide when to trust the model Pan et al.
and regularize the model samples Yu et al.
Besides, Luo et al.
(2018) provided a theoretical guarantee of monotone expected reward improvement of model-based RL.
Rajeswaran et al.
cast model-based RL as a game-theoretic framework by formulating the optimization of model and policy as a two-player game.
To save time tuning hyperparameters, Lai et al.
(2021) designed an automatic scheduling framework.
Abbas et al.
(2020) systematically studied how the model capacity affects the model-based methods.
Value-equivalence dynamics model.
Value-equivalence dynamics model has been noted by several authors in recent years.
Since learning an accurate dynamics model of the world remains challenging and often requires computationally costly and data-hungry models (Lovatto et al., 2020), Farahmand et al.
(2017) proposed value-aware model learning which aims to learn a value-equivalence model that induces the same Bellman operator as the real environment, rather than accurately predicting transitions.
However, they replaced the value function with the supremum over a function space, and it is difficult to find a supremum for a function space parameterized by complex function approximators like neural networks.
Based on this work, Farahmand (2018) proposed Iterative Value-Aware Model
Learning (IterVAML) which replaced the supremum over a value function space with the value function at current iteration.
Besides, Grimm et al.
(2020) introduced value equivalence principle and analysed how the space of possible solutions on model learning is impacted by the choice of policies and functions.
However, despite very detailed theoretical guarantees, there is still a performance gap between the value-equivalence dynamics model in the practical implementation and the model trained by the maximum likelihood estimate (Lovatto et al., 2020).
Eysenbach et al.
(2021) introduced a novel objective to jointly train the model and the policy.
Voelcker et al.
proposed Value-Gradient weighted Model loss (VaGraM) which approximated the value-aware model loss function with a
Taylor expansion of value function and achieved SOTA performance across all value-aware model learning methods.
Like our method, VaGraM also tries to learn a locally accurate dynamics model.
The difference is that our method aims to learn the samples that the current policy may encounter as accurately as possible, while VaGraM is to learn the dimensions in the state that can bring the greatest improvement to policy learning.
Experimental results demonstrate that our method outperforms
VaGraM in practice.
C U SEFUL LEMMA
Lemma C.1 Shen et al.
Assume the initial state distributions of the real dynamics T and the learned dynamics model ˆ T are the same.
For any state s ′ , assume F s ′ is a class of real-valued bounded measurable functions on state-action space, such that ˆ T ( s ′ |· , · ) : S ×
A → R is in F s ′ .
Then the gap between two different state visitation distributions v π 1 T ( s ′ ) and v π 2 ˆ T ( s ′ ) can be bounded as follows:
[Equation]
Proof.
For any state visitation distribution v πT
[Equation]
we have:
[Equation]
where v 0 is the probability of the initial state being the state s ′ .
Then the gap between two different state visitation distributions is:
[Equation]
10) □
D P ROOF OF MAIN THEOREM
Theorem D.1 Given the historical policy mixture π mix ,k
= (Π k , w k ) at iteration step k , we denote ξ ρ i
= D TV ( ρ πT ( s, a ) ||
ρ π i T ( s, a )) and ξ π i =
E s ∼ v π mix ˆ T [ D TV ( π ( a | s )
||
π i ( a | s ))] as the state-action visitation distribution shift and the policy distribution shift between the historical policy π i and current policy π respectively, where v π mix
ˆ T is the state visitation distribution of policy mixture under the learned dynamics model.
r max is the maximum reward the policy can get from the real environment, γ is the discount factor
[Equation]
and Vol( S ) is the volume of state space.
Then the performance gap between the real environment rollout J ( π, T ) and the model rollout J ( π, ˆ T ) can be bounded as follows:
[Equation]
Proof.
[Equation]
For term 1:
[Equation]
≤ r max
[Equation]
For the first term of last inequality in Eq.
[Equation]
according to Lemma
[Equation]
C.1 we have:
[Equation]
We use total variance distance as the F s ′ to measure the distance between ρ πT and ρ π mix ˆ T .
Suppose we can learn a dynamics model that can perfectly adapt the state-action visitation distribution of π mix , which means the difference between the model prediction and the environment next state s ′ is very small
[Equation]
and the state-action visitation density induced by the learned dynamics model ρ π mix ˆ T is approximately equal to ρ π mix T .
This assumption is required by many model-based RL methods
Voelcker et al.
Then Eq.
14 can be expressed as:
[Equation]
[Table]
[Equation]
Finally, based on Eq.
[Equation]
we get:
[Equation]
and the proof is completed
[Equation]
□
E P ROOF OF P ROPOSITION 3.
Proposition E.1 The performance gap can be reduced if the weight w ki of each policy π i in the historical policy sequence Π k is negatively related to state action visitation distribution shift
ξ ρ i and the policy distribution shift ξ π i between the historical policy π i and current policy π instead of an average weight w ki = 1 k :
[Equation] [Equation]
Proof.
Each policy π i in the historical policy sequence Π k corresponds to a distribution shift pair ( ξ ρ i , ξ π i ) , and these pairs form a distribution shift sequence { ( ξ ρ 1 , ξ π 1 ) , ( ξ ρ 2 , ξ π 2 ) , ......, ( ξ ρ k , ξ π k ) } , assuming that this sequence decreases as i increases (this is a reasonable assumption, because we can always arrange the historical policy sequence into a distribution shift decreasing sequence according to the magnitude of the shift).
As the weight of each policy is negatively related to state action visitation distribution shift ξ ρ
i and the policy distribution shift ξ π i , w ki increases with k .
Since k (cid:80)
[Equation]
i
[Equation]
1 k = 1 , there exists a k 0 that for all i > k 0 , w ki > 1 k .
Then we have:
[Equation]
where k 1 ∈ [ k 0 , k ]
[Equation]
where k 2 ∈ [0 , k 0 )
Based on these two equations:
[Equation]
Since distribution shift sequence { ( ξ ρ 1 , ξ π 1 ) , ( ξ ρ 2 , ξ π 2 ) , ......, ( ξ ρ k , ξ π k ) } decreases as i increases, and k 2 < k 1 , the first term will be less than 0.
Meanwhile, the second term will be equal to
because k (cid:80) i =1 w k i = k (cid:80) i = 1 k = 1 .
Therefore, we can get:
[Equation]
⩽
[Equation] [Equation]
The proof is finished.
□
Proposition 3.2 illustrate that after adjusting the policy mixture distribution according to the distribution shifts, the performance bound will be tighter than learning a global dynamics model ( w ki = 1 k ).
This provides a guidance for our proposed method, that the weight w ki of each policy π i in the historical policy sequence Π k should be negatively related to its state action visitation distribution shift ξ ρ
i and the policy distribution shift ξ π i .
F M ORE EXPERIMENTS F. M ORE E RROR C URVES FOR D YNAMICS M ODEL L EARNED BY MBPO
In this section, we provide the local error curves for global dynamics model in four MoJoCo environments: Hopper, HalfCheetah, Walker2d, and Humanoid.
The curves are shown in Figure 5.
[Figure]
Overall Error
Current Error
[Figure]
Overall Error
Current Error
[Figure]
Overall
[Figure]
Current
[Figure]
F. V ISUALIZATION OF S TATE -A CTION V ISITATION D ISTRIBUTION OF D IFFERENT H ISTORICAL P OLICIES
Due to the limited space of main paper
[Figure]
we provide detailed visualization of the state-action visitation distribution of policies under different environment steps in this section.
We conduct the experiment on HalfCheetah and Hopper, the results are shown in Figure 6 and Figure 7.
(a) in each figure is the comparison of different policies in the same figure, from (b) to (f) are the figures presenting the state-action visitation distribution of each policy individually.
We can see that the state-action visitation distribution of policies under different environment steps is very different.
[Figure]
[Figure]
F. V ISUALIZATION OF A DJUSTED P OLICY M IXTURE D ISTRIBUTION
To provide a further understanding of out method
[Figure]
we visualize the adjusted policy mixture distribution at different training steps on Humanoid in Figure 8.
We take Figure 8a as an example to explain the origin and meaning of the policy ID on the horizontal axis.
Each of policies interacts with the environment for 250 steps, so a 50k environment step has 200 historical policies.
The policy ID of 0 indicates the oldest policy.
The larger the ID, the newer the policy.
We can see that the policy mixture distribution is totally different at different training steps.
The weight of policy is not a simple exponentially decay or linearly decay, which indicates our proposed method is non-trivial
[Figure]
F. M ORE B ENCHMARK R ESULTS
Due to the limited space, we put more results on MuJoCo in this section.
Compared with Figure 2, we add the results on Ant and Pusher in Figure 9.
In Ant and Pusher, our method achieves a significant performance improvement compared with original MBPO.
Besides, our method also achieves comparable or better results than previous sample efficient model-based or model-free methods.
F. C OMPARISON WITH S IMPLE E XPONENTIALLY D ECAY P RIORITIZATION
To further demonstrate the effectiveness of our method, we compare with an exponentially decay method.
The weight of the historical policy exponentially decays as it lifetime increases.
To ensure a fair comparison, the weight of current policy is also compute using Eq.
The hyperparameter α of exponentially decay method is the same as PDML which is given in Appendix G.1.
We conduct the experiment on three MoJoCo environments: Hopper, Walker2d, and Humanoid.
The experiment results are given in Figure 10.
We can see that after using exponentially decay method, the performance in three environments is slightly improved, but it is much lower than PDML.
Besides, the model error of exponentially decay method is higher than PDML.
Combined with the analysis of distribution visualization in Appendix F.3, this further demonstrates that our method is non-trivial and effective.
F. A BLATION S TUDY OF PDML
As we described in Section 4, we use the adjusted policy mixture distribution for both model learning and sampling initial states for model rollouts. In this section, we provide the ablation study to show the impact of the adjusted policy mixture distribution in these two parts respectively. We conducted our experiments in Hopper and Walker2d, and the performance curves are shown in Figures 11a and 11b.
[Figure]
0k 30k 60k 90k 120k
[Figure]
0k 75k 150k
[Figure]
Humanoid
[Figure]
[Figure]
[Figure]
We find that using the adjusted policy mixture distribution only for model learning or model rollouts initial states sampling both improves the performance in Hopper and Walker2d compared to MBPO.
However, the improvement of only using the adjusted policy mixture distribution for model rollouts initial states sampling in Walker2d is not very significant.
Besides, the improvement of using the adjusted policy mixture distribution for model learning is better than using that for model rollouts initial states sampling, but both of them are worse than PDML.
This indicates two things.
First, model learning is more important than model rollouts initial states sampling, because even the initial state distribution obeys the state-action visitation distribution of current policy, the model-generated samples will still be inaccurate if the learned dynamics model is not accurate enough for the current policy.
Second, to achieve the best performance, sample distribution for model learning and sample distribution for model rollouts initial states should be synergistic; that is, the training data for training the dynamics model and the initial states of model rollouts should obey the same distribution, so that the model prediction error can be minimized.
F. A BLATION S TUDY OF C URRENT P OLICY P ROPORTION R ATE
We conducted experiments to explore the impact of current policy proportion rate on the performance of our method.
The α in Eq.
5 equals to current policy proportion rate divided by 1 minus current policy proportion rate.
As shown in Figure 11c, when the current policy proportion rate is small (0.02 and 0.1), the policy mixture distribution will not be too inclined to the current policy, so the model can learn a good transition dynamics.
When the current policy proportion rate is too large (0.3, 0.5, and 0.7), the learned dynamics capture information about the underlying transition too locally, resulting in performance decrease.
Therefore, we recommend that the selection of the current policy proportion rate should not be greater than 0.1.
In this section, we present the training time of PDML-MBPO and MBPO in four different environments.
As shown in Table 1, after using PDML, the training time doesn’t increase significantly.
In the most complex environment Humanoid, the training time for 300k steps increases by only one hour.
In other environments, the training time of PDML-MBPO is almost the same as that of MBPO.
[Figure]
Hopper
Humanoid
HalfCheetah
[Figure]
[Table]
[Figure]
[Table]
h
[Table]
h
[Table]
F. O NE - STEP E RROR IN F OUR E NVIRONMENTS
As an extension of Section 5.
[Figure]
we provide the one-step model prediction error curve in this section.
The results are shown in Figure 12.
[Figure]
G I MPLEMENTATION G. I MPLEMENTATION D ETAILS
We implement PDML-MBPO based on the PyTorch-version MBPO Liu et al.
We also set the ensemble size of PDML-MBPO to be the same as MBPO, which is 7.
The warm-up samples are collected through interaction with the real environment for 5000 steps using a randomly chosen policy.
After the warm-up, we train the dynamics model and update the lifetime weight every 250 interaction steps.
We set the current policy proportion to be 0.02 and α equals 0 .
One thing that needs to be noticed is the rollout horizon setting.
As introduced in MBPO Janner et al.
(2019), the rollout horizon should start at a short horizon and increase linearly with the interaction epoch.
[ a, b, x, y ] denotes a thresholded linear function, i.e . at epoch e , rollout horizon is h = min(max( x + e − a b − a ( y − x ) , x ) , y ) .
We set the rollout horizon to be the same as used in the MBPO paper, as shown in Table 2.
Other hyper-parameter settings are shown in Table 3.
For MBPO 1 , AMPO 2 , VaGraM 3 SAC 4 , and REDQ 5 , we use their open source implementations.
We evaluate PDML-MBPO and other baselines on four MuJoCo-v2 continuous control environments Todorov et al.
(2012) with a maximum horizon of 1000, including HalfCheetah, Hopper, Walker2d, and Humanoid.
For Humanoid, we use the modified version introduced by MBPO Janner et al.
All experiments are conducted using a single NVIDIA TITAN X Pascal GPU.
For the experiment of MaPER in Sec 5.2, we use their open-source code in the supplementary material on openreview 6 .
However, we find a bug in their code that comes from the PyTorch-version MBPO implementation, i.e., the same environment is used for policy training and policy evaluation. This [Table]
[Table] causes the rollout length to exceed the 1000-step limit during evaluation [Table] resulting in the performance of the policy being much higher than the 1000-step performance.
We fix this bug and conduct the experiment [Table] so the results of MaPER are lower than those reported in their paper
[Table]