Fine-Tuning Offline Reinforcement Learning with Model-Based Policy Optimization

1 INTRODUCTION . 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 difficult 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 specific 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 difficulty 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 offline 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 first 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 confidence . Thus , these methods allow a certain degree of extrapolation where the models are confident . Because these methods tend to be less restrictive , they can generalize better than behavior-regularization methods and sometimes outperform the behavioral dataset . However , this flexibility 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 . Specifically , we first train a policy using behaviorregularized model-free RL . Then , we fine-tune our results with our novel algorithm Model-Based Behavior-Regularized Policy Optimization ( MB2PO ) . We find 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 . 2 RELATED WORK . 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 offline without any online data collection . Recently , there have been several methods that made progress in improving the capabilities of offline RL . For a general overview of the field of offline RL , we refer the reader to Levine et al . ( 2020 ) . Here we will discuss some recent works that are particularly relevant to our approach . 2.1 IMPROVING OFF-POLICY Q-LEARNING . 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 offline 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 offline RL algorithm Conservative Q-learning ( CQL ) ( Kumar et al. , 2020 ) has attempted to address Q-learning ’ s overestimation issue on offline 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.2 BEHAVIOR-REGULARIZED MODEL-FREE RL . A variety of recent offline 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 offline 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 fine-tuning with Model-Based Policy Optimization ( MBPO ) ( Janner et al. , 2019 ) . 2.3 UNCERTAINTY-AWARE MODEL-BASED RL . MB RL algorithms have several natural advantages for offline 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 benefits 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 offline 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 difficult 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 offline deep RL . In particular , two recent works , Model-Based Offline Policy Optimization ( MOPO ) ( Yu et al. , 2020 ) and Model-Based Offline 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 . 3 PRELIMINARIES . 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′|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 , ρ0 [ ∑∞ t=0 γ tr ( st , at ) ] . In offline RL , we assume that during training we only have access to a fixed dataset Dβ containing a set of tuples ( s , a , s′ , 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 Qπk+1 = arg min Q Es , a , s′∼D [ ( ( r ( s , a ) + γEa′∼π ( ·|s′ ) [ Qπk ( s′ , a′ ) ] ) −Qπ ( s , a ) ) 2 ] ( 1 ) In policy improvement , we update our parametric policy π to maximize our current Q-function πk+1 = arg max π Es∼D , a∼π ( ·|s ) [ Qπk ( s , a ) ] ( 2 ) 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 .

Summary: This paper proposes an improved offline RL (batch RL) algorithm combining the state-of-the-art behavior-regularization actor-critic method (Nair et al., 2020) with a model-based RL technique. N-trained probabilistic dynamics models generate fictitious trajectories with uncertainty-penalized rewards after pretraining the policy with the behavior-regularization solely on the offline data. Both these generated data and the original offline data are used for the further behavior-regularized actor-critic training. Numerical results showed that the proposed method outperformed recent offline model-free and model-based RL algorithms.