F INE -T UNING O FFLINE R EINFORCEMENT L EARNING WITH M ODEL -B ASED P OLICY O PTIMIZATION
Anonymous authors
Paper under double-blind review
A BSTRACT
In ofﬂine reinforcement learning (RL), we attempt to learn a control policy from a ﬁxed dataset of environment interactions.
This setting has the potential beneﬁt of allowing us to learn effective policies without needing to collect additional interactive data, which can be expensive or dangerous in real-world systems.
However, traditional off-policy RL methods tend to perform poorly in this setting due to the distributional shift between the ﬁxed data set and the learned policy.
In particular, they tend to extrapolate optimistically and overestimate the action-values outside of the dataset distribution.
Recently, two major avenues have been explored to address this issue.
First, behavior-regularized methods that penalize actions that deviate from the demonstrated action distribution.
Second, uncertainty-aware model-based (MB) methods that discourage state-actions where the dynamics are uncertain.
In this work, we propose an algorithmic framework that consists of two stages.
In the ﬁrst stage, we train a policy using behavior-regularized model-free
RL on the ofﬂine dataset.
Then, a second stage where we ﬁne-tune the policy using our novel Model-Based Behavior-Regularized Policy Optimization (MB2PO) algorithm.
We demonstrate that for certain tasks and dataset distributions our conservative model-based ﬁne-tuning can greatly increase performance and allow the agent to generalize and outperform the demonstrated behavior.
We evaluate our method on a variety of the Gym-MuJoCo tasks in the D4RL benchmark and demonstrate that our method is competitive and in some cases superior to the state of the art for most of the evaluated tasks.
I NTRODUCTION
Deep reinforcement learning has recently been able to achieve impressive results in a variety of video games (Badia et al., 2020) and board games (Schrittwieser et al., 2020).
However, it has had limited success in complicated real-world tasks.
In contrast, deep supervised learning algorithms have been achieving extraordinary success in scaling to difﬁcult real-world datasets and tasks, especially in computer vision (Deng et al., 2009) and NLP (Rajpurkar et al., 2016).
The success of supervised learning algorithms can be attributed to the combination of deep neural networks and methods that can effectively scale with large corpora of varied data.
The previous successes of deep RL (Levine, 2016; Schrittwieser et al., 2020) seem to indicate that reinforcement learning can potentially scale with large active data exploration to solve speciﬁc tasks.
However, the ability to collect such large datasets online seems infeasible in many real-world applications such as automated driving or robotassisted surgery, due to the difﬁculty and inherent risks in collecting online exploratory data with an imperfect agent.
Existing off-policy RL algorithms can potentially leverage large, previously collected datasets, but they often struggle to learn effective policies without collecting their own online exploratory data (Agarwal et al., 2020).
These failures are often attributed to the Q-function poorly extrapolating to out-of-distribution actions, which leads to overly optimistic agents that largely over-estimate the values of unseen actions.
Because we train Q-functions using bootstrapping, these errors will often compound and lead to divergent Q-functions and unstable policy learning (Kumar et al., 2019).
Recently, there have been a variety of ofﬂine RL approaches that have attempted to address these issues.
Broadly, we group these approaches into two main categories based on how they address the extrapolation issue.
The ﬁrst set of approaches (Wu et al., 2019; Kumar et al., 2019) rely on behavior-regularization to limit the learned policy’s divergence from the perceived behavioral policy that collected the data.
These approaches discourage the agent from considering out-of-distribution actions in order to avoid erroneous extrapolation.
While these methods can often be effective when given some amount of expert demonstrations, they often seem too conservative and rarely outperform the best demonstrated behavior.
The second set of approaches (Yu et al., 2020; Kidambi et al., 2020) leverage uncertainty-aware
MB RL to learn a policy that is discouraged from taking state-action transitions where the learned model has low conﬁdence.
Thus, these methods allow a certain degree of extrapolation where the models are conﬁdent.
Because these methods tend to be less restrictive, they can generalize better than behavior-regularization methods and sometimes outperform the behavioral dataset.
However, this ﬂexibility also seems to make it harder for these methods to recover the expert policy when it is present in the dataset, and reduce their effectiveness when trained with a narrow distribution.
In this work, we develop an algorithmic framework that combines ideas from behavior-regularization and uncertainty-aware model-based learning.
Speciﬁcally, we ﬁrst train a policy using behaviorregularized model-free RL.
Then, we ﬁne-tune our results with our novel algorithm Model-Based Behavior-Regularized Policy Optimization (MB2PO).
We ﬁnd that our approach is able to combine the upside of these approaches and achieve competitive or superior results on most of the GymMuJoCo (Todorov et al., 2012) tasks in the D4RL (Fu et al., 2020) benchmark.
R ELATED W ORK
While there exist many off-policy RL methods that can learn to solve a large variety of complex control tasks and can scale with large amounts of online data collection, these methods often perform quite poorly when run completely ofﬂine without any online data collection.
Recently, there have been several methods that made progress in improving the capabilities of ofﬂine RL.
For a general overview of the ﬁeld of ofﬂine RL, we refer the reader to Levine et al.
Here we will discuss some recent works that are particularly relevant to our approach.
2. I MPROVING O FF - POLICY Q-L EARNING
Many of the recent advances in both discrete and continuous action off-policy deep RL can be attributed to improvements in stabilizing off-policy Q-learning and reducing overestimation due to erroneous extrapolation.
Some notable methods include target networks (Mnih et al., 2013), double Q-learning (DDQN) (van Hasselt et al., 2015), distributional RL (Bellemare et al., 2017;
Dabney et al., 2017), and variance reduction through invertible transforms (Pohlen et al., 2018).
In learning for continuous control, Fujimoto et al. (2018) introduced a conservative method that uses the minimum estimate of an ensemble of Q-networks as the target, which is often referred to as clipped double-Q-learning.
Agarwal et al.
(2020) demonstrated that Quantile Regression DDQN (Dabney et al., 2017) and other ensemble methods can be effective in certain discrete action ofﬂine
RL problems.
However, Agarwal et al. (2020) showed that when used naively, these methods do not perform well on complex continuous control tasks.
In our work, we incorporate the mentioned advances in off-policy Q-learning into our approach to stabilize performance and prevent potential divergence.
Additionally, the ofﬂine RL algorithm Conservative Q-learning (CQL) (Kumar et al., 2020) has attempted to address Q-learning’s overestimation issue on ofﬂine data directly by including a constraint term that discourages the agent from valuing an out-of-distribution action more than the demonstrated actions.
In our method, instead of using a constraint on the Q-values, we use a combination of behavior-regularized model-free RL and uncertainty-aware model-based RL to discourage erroneous extrapolation.
2. B EHAVIOR -R EGULARIZED M ODEL -F REE RL
A variety of recent ofﬂine RL approaches have incorporated constraints or penalties on the learned policy’s divergence from the empirical behavioral policy.
In particular, recent works have used both KL Divergence (Wu et al., 2019) and mean measure of divergence (MMD) (Kumar et al., 2019).
MMD is sometimes used over KL Divergence because MMD approximately constrains the learned policy to be in the support of the behavioral policy, which is less restricting than KL Divergence.
However, most behavior-regularization or policy-constraint methods require the behavioral policy to be represented explicitly in order to estimate these divergences or to enforce their policy constraint (Laroche et al., 2019).
In contrast, AWAC (Nair et al., 2020) or CRR (Wang et al., 2020) is able to incorporate a KL divergence constraint without explicitly representing the behavioral policy.
They do this by reformulating the policy-constrained RL optimization equations into a form that resembles behavioral cloning re-weighted by the exponential of the advantage.
Wang et al.
(2020) demonstrates that this method can effectively learn complex control tasks purely from ofﬂine data, and Nair et al.
(2020) demonstrate that performance can even be improved with further online data collection.
In this work, we demonstrate that these properties make AWAC work exceptionally well when used for initialization as well as when used for ﬁne-tuning with Model-Based Policy Optimization (MBPO)
(Janner et al., 2019).
2. U NCERTAINTY -A WARE M ODEL -B ASED RL
MB RL algorithms have several natural advantages for ofﬂine RL compared to model-free RL algorithms.
First, MB RL algorithms rely on supervised learning, which provide more robust gradient signals compared to bootstrapped learning and policy gradients.
Second, learning a dynamics model often provides strong task-independent supervision, which allows MB RL algorithms to learn from sub-optimal trajectories.
These beneﬁts make generalization easier, and can allow MB RL algorithms to surpass the performance of the demonstrated data.
In fact, in many environments, MB RL methods have already been effective in learning with ofﬂine or randomly collected datasets.
Additionally, there is a rich history of prior works that have explored robust solutions to MDPs with uncertain transition dynamics (Nilim & El Ghaoui, 2005; Iyengar, 2005).
However, it can be difﬁcult to scale these type of methods to high-dimensional continuous control tasks, especially when using deep neural networks.
Recently, incorporating uncertainty estimation techniques from supervised learning in MB RL has demonstrated further improvement in both online
(Chua et al., 2018) and ofﬂine deep RL.
In particular, two recent works, Model-Based Ofﬂine Policy Optimization (MOPO)
(Yu et al., 2020) and Model-Based Ofﬂine Reinforcement Learning (MoREL) (Kidambi et al., 2020), have demonstrated impressive results by incorporating uncertainty-aware MB RL with the Dyna (Sutton, 1991) style algorithm MBPO
(Janner et al., 2019).
Both methods use these models to create conservative MDPs that have a lower potential expected sum of rewards compared to the true MDP.
By performing policy optimization in the conservative MDP through MBPO they are able to learn a conservative policy that can outperform the demonstrated trajectories.
However, these methods can often fail to recover the expert policy even though it was demonstrated in the dataset.
We believe that this is largely due to a lack of effective methods for estimating epistemic uncertainty for neural network regression.
P RELIMINARIES
In RL, we assume our agent operates within a standard Markov decision process (MDP)
M = ( S , A , T, r, ρ 0 , γ ) , where S denotes the state space, A denotes the action space, T ( s (cid:48) | s, a ) represents the probabilistic transition dynamics, r is the reward function, ρ 0 is the initial state distribution, and γ ∈ (0 , 1) is the discount factor.
The objective in RL is to learn a policy π ( a | s ) that optimizes the expected discounted sum of rewards R π
= E π,T,ρ
[ (cid:80) ∞ t =
γ t r ( s t , a t )] .
In ofﬂine RL, we assume that during training we only have access to a ﬁxed dataset D β containing a set of tuples ( s, a, s (cid:48) , r ) of environment transitions and associated rewards.
We assume that the data was collected by a policy π β , which we call the behavioral policy.
Typically, when training with data not collected by your current policy π , we either use off-policy model-free algorithms or model-based algorithms.
The most common off-policy model-free algorithms are actor-critic algorithms that use policy iteration. Policy iteration involves alternating between policy evaluation and policy improvement in order to learn an effective policy.
In policy evaluation, these methods train a parametric Q-function by iteratively minimizing the temporal difference equation
[Equation]
In policy improvement, we update our parametric policy π to maximize our current Q-function
[Equation]
In MB RL, we attempt to learn a model ˆ T of the transition dynamics and a model ˆ r of the reward function.
With this learned model of the dynamics and reward function we can create a model
MDP ˆ
M = ( S , A , ˆ T, ˆ r, ρ 0 , γ ) to estimate the true underlying MDP M .
These methods tend to use either trajectory optimization or policy optimization in the model MDP to produce their policy.
M ODEL -B ASED B EHAVIOR -R EGULARIZED P OLICY O PTIMIZATION FOR O FFLINE F INE -T UNING
For many ofﬂine datasets, it could be much harder to learn an effective model of the MDP than to learn a reasonable policy.
This is especially the case when there is low variability or insufﬁcient coverage of the state and action space in the collected dataset, or in environments with complex observations, like images, or long horizons.
To overcome these issues, recent works (Yu et al., 2020;
Kidambi et al., 2020) have leveraged uncertainty estimation methods in order to construct conservative MDPs that use soft penalties or hard thresholds on model uncertainty to discourage deviating from the conﬁdent regions.
However, these methods rely on the efﬁcacy of ensemble-based neural network uncertainty estimation methods which currently are not particularly effective at estimating epistemic uncertainty in regression settings.
Therefore, we propose Model-Based BehaviorRegularized Policy Optimization (MB2PO).
In MB2PO, we likewise use uncertainty-aware models to perform ofﬂine MBPO, but use the behavior-regularized model-free algorithm AWAC (also known as CRR-exp) instead of SAC (Haarnoja et al., 2018) for policy optimization.
4. C ONSERVATIVE MBPO
In this work, we use MOPO (Yu et al., 2020) as a basis for our conservative MBPO, due to its simplicity and prior effective results on the D4RL benchmarks.
In MOPO, they construct a conservative
MDP by augmenting the reward function as follows
[Equation]
where ˆ r is the learned estimate of the reward and u is the estimated uncertainty for the model transition.
Note, that this general formulation for a conservative MDP has also been explored in other prior work such as (Ghavamzadeh et al., 2016).
Still, we speciﬁcally follow MOPO in using the maximum standard deviation across an ensemble of probabilistic dynamics models as our measure of uncertainty.
Therefore, we can decompose our Q-function in this conservative MDP as
[Equation]
where ˆ Q πr represents our estimate of the expected discounted sum of rewards in the real MDP and
Q πu represents our expected discounted sum of uncertainty penalties.
Now at convergence, if our policy π deviates from the behavioral policy π β that collected the data, then we expect for all states in the conservative MDP that
[Equation]
Thus, by plugging in our decomposition we get
[Equation]
While in theory, with well-calibrated uncertainty estimates and a proper tuning of λ , this should lead to only safe policy improvements over the behavioral policy, in practice it seems that MOPO is often unable to recover expert-level performance when it is provided in the ofﬂine dataset.
This is unsurprising given that it is hard to generate well-calibrated epistemic uncertainty estimates in regression settings, and there will inevitably be model errors that will lead to overestimated Qvalues.
To address these issues, we use policy constrained model-free RL in MB2PO.
In policy constrained model-free RL, we attempt to optimize the following policy objective
[Equation]
If we estimate both π and π β to be roughly univariate Gaussians with similar variances, then the KL constraint becomes an (cid:96) 2 constraint on the policy mean.
Because we expect our models to be locally accurate around the data, this constraint can help ensure that we stay in the effective region of the estimated MDP even if we have poorly calibrated uncertainty estimation.
Additionally, Janner et al.
(2019) demonstrated that the difference between the true expected returns J ( π ) and the expected returns ˆ J ( π ) of an MDP induced by an approximate model can be bounded by
[Equation]
where r max is the maximum reward, γ is the discount factor, (cid:15) m is a bound on the total variation distance (TVD) between the learned model and the true model, and (cid:15) π is a bound on the TVD between π and π β on the demonstrated states.
By Pinker’s inequality, bounding the KL divergence also bounds the TVD.
Therefore, by leveraging policy constraints in the policy optimization in
MBPO, we can reduce the gap in expected returns and improve the algorithm’s robustness to model errors.
4. B EHAVIOR -R EGULARIZED M ODEL -F REE RL WITH AWAC
For performing behavior-regularized policy optimization, we use AWAC (Nair et al., 2020), also known as CRR-exp (Wang et al., 2020) due to its impressive results in ofﬂine RL and its ability to be ﬁne-tuned with additional online data.
By enforcing the KKT conditions (Peng et al., 2019; Peters & Schaal, 2007; G ´ omez et al., 2014), we can derive an analytic solution to Equation 7, where the Lagrangian is
[Equation]
We can substitute A π ( s, a ) for Q π ( s, a ) because it does not affect the optimum and get the closedform solution
[Equation]
where Z ( s ) is the normalizing partition function.
In order to project this solution into our policy space, we update our parameters by minimizing D KL ( π ∗ (cid:107) π θ ) .
This leads to the following iterative update
[Equation]
We follow Wang et al.
and Peng et al.
and avoid estimating Z ( s ) and instead clamp the exponential term to be at most 20.
Additionally, one could adaptively learn α using dual gradient descent, but this would require us to explicitly model the behavioral policy π β .
Instead, we use a ﬁxed α for all of our results.
Additionally, the Q-function is updated off-policy using the Bellman equations as described in Equation 2 and the improvements from section 2.1.
One of the major beneﬁts of using AWAC with a ﬁxed α is that we can leverage behavior regularization in a principled manner without needing to explicitly represent the behavioral policy.
This is particularly important in 3 major cases: 1.
when there are not enough data to learn the behavioral policy; 2. when the data was collected by a variety of different policies or sources; 3. when the data was collected by a policy outside of your policy class, such as a human expert or a controller that leverages hidden state information.
Additionally, we can view AWAC as a reweighted behavioral cloning algorithm.
Unlike SAC (Haarnoja et al., 2018) and DDPG (Lillicrap et al., 2015), it does not rely on the reparametrization trick or gradients of your learned Q-function to perform policy updates.
This allows us to use a wider ranger of policy classes, which in this work we take advantage of by using a tanh squashed
GMM with 5 components.
We suspect that there are also some additional beneﬁts to not depending on the gradients of the learned Q-function, which might be particularly bad in ofﬂine settings, but leave further investigation to future work.
An important thing to note with AWAC is that we can inﬂuence the implicit behavioral penalty by controlling the source of the data we train with.
This holds, for example, if we perform a series of policy updates only using data collected by the previous policy iterate.
Then, we are implicitly performing a trust-region policy update like TRPO (Schulman et al., 2015) and PPO (Schulman et al., 2017) of the form
[Equation]
In fact, if we train on data collected by the last n policy iterates, then we are approximately constraining our policy to a weighted sum of the previous n policies π ( n ) k
[Equation]
and damping our learning process in the policy space.
In our work, we train with a ω ∈ [0 , 1] portion of the data from ofﬂine data collected by π β and a (1 − ω ) portion of the data collected online from the last n policy iterates in the conservative MDP deﬁned by our learned models.
Therefore, we are approximately optimizing the following objective
[Equation]
Therefore, by using AWAC as the policy optimization algorithm in MB2PO, we can easily perform behavior-regularized policy optimization with soft damped trust region updates in the conservative
MDP to reduce the effects of model errors and poor uncertainty estimation.
4. M ODEL -B ASED B EHAVIOR -R EGULARIZED P OLICY O PTIMIZATION
Train π θ , Q φ with AWAC with samples from D β if ﬁnetune then
Train an ensemble of N probabilistic dynamics [Equation] | s t , a t ) =
N ( µ iθ ( s t , a t ) , Σ iθ ( s t , a t )) }
Ni =1 on the data in D β for epoch k= 1, 2, ... do
Initialize empty replay buffer D k for 1, 2, ..., batchsize do
Sample state s 1 from D β for j = 1, 2, ..., h do a j ∼ π ( s j )
Uniformly sample ˆ T from { ˆ T i } Ni = s j +1 , r j ∼ ˆ T ( s j , a j ) ˜ r j = r j − λ max Ni =1 (cid:107) Σ
i ( s j , a j ) (cid:107) F
Add sample ( s j , a j , ˜ r j , s j =1 ) to D k end end
Draw ω portion of the samples from D β and the rest uniformly from {D k − i } 9 i =0 to train
AWAC end end We ﬁrst initialize our policy by training with AWAC solely on the ofﬂine data.
Next for ﬁne-tuning with MB2PO, we train an ensemble of probabilistic dynamics models represented by neural networks that output a diagonal Gaussian distribution over the next state and reward: { ˆ T iθ ( s t +1 , r | s t , a t ) =
N ( µ iθ ( s t , a t ) , Σ iθ ( s t , a t )) } Ni =1 .
We construct a conservative MDP that at every time step uses a randomly drawn dynamics model from { ˆ T i θ } M i =1 to determine the next state transition.
Additionally, we incorporate an penalty on the largest predicted standard deviation among the dynamics models as a practical means of penalizing both epistemic and aleatoric uncertainty.
Then, we alternate between collecting data with our current policy in the conservative MDP and updating our policy and Q-network using Equation 9 and Equation 2 respectively.
When collecting data in the conservative MDP, we collect h -length truncated trajectories starting from states in the original ofﬂine dataset.
By collecting data this way, we are able to collect a variety of imagined data without relying on long model rollouts, which would inevitably lead to compounding errors. When performing training updates
[Table]
we sample ω ∈ [0 , 1] of the data from the original dataset and the remaining 1 − ω uniformly from the last 100 policy iterates.
Our full algorithm is outlined in
Algorithm
[Table]
E XPERIMENTS
[Table]
In our experiments
[Table]
we aim to address two questions: (1) Is AWAC an effective initialization algorithm?
(2) Can we further improve performance by ﬁne-tuning with MB2PO?
We evaluate (1) by comparing AWAC to other state-of-the-art model-free ofﬂine RL algorithms.
In particular, we compare our results to BRAC-v (Wu et al., 2019), BEAR (Kumar et al., 2019) [Table]
and CQL (Kumar et al., 2020) on the Gym-MuJoCo tasks in the D4RL benchmark.
We evaluate (2) by ﬁne-tuning the policy and Q-function, after running AWAC for 500000 gradient steps, with MB2PO.
In addition to the model-free ofﬂine RL algoirthm above, we also compare these results to MOPO, which to the best of our knowledge is the state-of-the-art MB ofﬂine RL algorithm on the Gym-MuJoCo tasks in the D4RL benchmark. All of our hyperparameters for
AWAC and MB2PO are given in the appendix.
The Gym-MuJoCo tasks are a standard in evaluating modern deep RL algorithms.
The goal in these tasks is to learn to travel as far forward as possible within a set horizon on a variety of different robots.
The D4RL benchmark contains a variety of precollected datasets for the halfcheetah, walker2d, and hopper tasks.
For each robot task, there are 5 different provided datasets.
The ”-random” datasets contain 1 million samples collected from a randomly initialized policy.
The ”-medium” datasets contain 1 million samples collected from a RL policy partially trained to a performance of approximately 33 .
The ”-expert” datasets contain 1 million samples collected from a fully trained RL policy that reaches approximately 100 .
The ”-mixed” datasets contain all the data in the replay buffer from the partially trained ”medium” policy.
Finally, the ”-medium-expert” datasets are a combination of the ”-medium” and ”-expert” datasets.
An important thing to note is that all datasets besides the ”-mixed” datasets were collected with only 1 or 2 policies, and thus probably only cover a narrow part of the state-action distribution.
On the other hand, the ”-mixed” dataset contains the data collected by all of the policy iterates during an incomplete RL training run, and thus represents a much wider part of the state-action distribution.
Results in Table 1 demonstrate that AWAC on its own can get reasonable results on all the datasets and can approach state-of-the-art results on ”-expert” and ”-medium-expert” datasets.
Unlike the other behavior-regularized model-free methods, AWAC and CQL are able to get near or fully recover expert-level performance when trained on the ”medium-expert-” datasets.
This indicates that AWAC and CQL are more robust as there is less of a drop in performance compared to other methods when incorporating additional sub-optimal trajectories.
Next, we ﬁne-tune the trained AWAC policies with MB2PO.
For each task and dataset, we pretrain an ensemble of 5 probabilistic dynamics models for 100000 gradient steps on the behavioral dataset.
We then perform MB2PO for 500 iterations.
Each iteration consists of collecting 1000000 steps from h -length truncated trajectories in the conservative MDP, which should run in a few seconds on modern GPU hardware, followed by 1000 gradient steps.
Results in Table 1 demonstrate that our method is effective in improving the performance over
AWAC in 11 of the 15 tasks.
In particular, we ﬁnd that MB2PO signiﬁcantly improves the performance on all of the ”-mixed” datasets and even achieves state-of-the-art on ”walker2d-mixed” by a large margin.
These strong results in the ”-mixed” datasets demonstrate that our model-based ﬁne-tuning method can be especially beneﬁcial when there is sufﬁcient variation in the behavioral dataset.
Additionally, the noticeable improvement in some of the ”-medium” and ”-medium-expert” datasets demonstrate that our ﬁne-tuning can be effective even when the data was collected by one or two policies.
In the 4 cases where MB2PO ﬁne-tuning degrades performance, it is always negligible and never over 3 points. Our method also outperforms MOPO, the most direct comparison, in 8 out of the 12 comparable tasks.
These results demonstrate the beneﬁts of combining behavior-regularized model-free RL with uncertainty-aware MB RL as we are able to get the best of both worlds.
We are able to recover high-level performance when available in the dataset like AWAC, and we can still generalize and learn to outperform the best observed trajectory like MOPO.
Finally, our method is quite competitive with CQL as we beat it for 7 of the 15 tasks, and generally our results are quite comparable. However, our method signiﬁcantly outperforms CQL in all of the ”-mixed” datasets.
These results indicate that our method might be superior in situations where the dataset was collected by a variety of different policies.
C ONCLUSION
We proposed an algorithmic framework that leverages the beneﬁts of both behavior-regularized model-free methods and uncertainty-aware model-based methods.
We do this by ﬁrst training an initial policy with the ofﬂine model-free AWAC algorithm.
Then, we ﬁne-tune with our novel MB2PO algorithm.
We perform this by learning uncertainty-aware models that are used to create a conservative MDP.
Then, we continue to use AWAC to further update our policy and Q-function in this conservative MDP.
By using AWAC, we are able to perform policy optimization while implicitly constraining the learned policy’s KL divergence to the behavioral policy.
We demonstrate that this two-stage process allows us to get the best of both worlds between behavior-regularized model-free methods and uncertainty-aware model-based methods.
Speciﬁcally, the initial AWAC training allows us to often recover the best-performing behavior in the dataset, and, when possible, MB2PO ﬁne-tuning can allow us to generalize and outperform the demonstrated behavior.
We see four important directions of future work in order to extend the effectiveness and applicability of MB2PO: 1. developing a rigorous means of determining for what datasets MB2PO ﬁne-tuning can be effective; 2. improving MB RL and neural network uncertain estimation to increase the number of datasets where MB2PO can be effective; 3. better leveraging behavior-regularization in the policy optimization or the conservative MDP to improve MB2PO’s ability to recover expert behavior when available; 4. improving off-policy evaluation (Thomas et al., 2015) for neural network policies in order to facilitate ofﬂine hyperparameter tuning.  A PPENDIX
All methods were trained with the Adam optimizer(Kingma & Ba, 2014) 7. AWAC Our Q-networks and policies were represented with [256, 256, 256, 256] fully connected networks with relu hidden-activations. The policy was a 5-head tanh squashed GMM.
The Q-network outputted 100 quantiles and was trained with the following Q-learning improvement: Quantile Regression DQN (Dabney et al., 2017), Clipped DQN(Fujimoto et al., 2018), Double DQN(van Hasselt et al., 2015), and Invertible Transforms(Pohlen et al., 2018). The advantage was estimated withsamples. We ran AWAC for 500000 gradient steps. The rest of the parameters were taken from the original AWAC paper (Nair et al., 2020): α = 1 ., batch size 1000, policy weight decay 1 .e − 4 , policy and Q learning rate 3 .e − 4 , soft target network update 5 .e − 3 .
7. MB RL
Our probabilistic dynamics models were represented with [200, 200, 200, 200] fully connected networks with swish hidden-activations.
The models outputted a mean and variance for every state variable and the reward.
We trained the models for 100000 gradient steps.
The rest of the parameters are: batch size 256, weight decay 1 .e − 4 , learning rate 1 .e − 3 .
7. MB2PO
MB2PO used the same AWAC parameters as before, but changed α according to the table below.
We ran MB2PO for 500 iterations where each iteration consisted of collecting 1000000 samples from truncated h -length trajectories with the current GMM policy, then training for 1000 steps with ω of the data from real dataset and 1 − ω of the data from imagined model rollouts from the lastiterations.
We picked the better of α = 1 , ω
= { 0 . 01 , 0 . 05 } for the mixed datasets, and the better of α = { 1 , 2 } , ω = 0 .
8 for the other datasets.
We used h = 5 for several of the halfcheetah datasets and h = 1 for the rest because any larger h tended to cause the Q-values to explode.
[Table]