Pareto Policy Pool for Model-based Offline Reinforcement Learning

1 INTRODUCTION . Offline reinforcement learning ( offline RL ) ( Levine et al. , 2020 ) or batch RL ( Lange et al. , 2012 ) can train an agent without interacting with the environment by instead using pre-collected experiences on other agents/policies . Recently , offline RL has been attracting growing interest due to the availability of large-scale datasets of diverse experiences in many RL applications , e.g. , autonomous driving ( Shin & Kim , 2019 ) , healthcare ( Yu et al. , 2019 ) , robot control ( Schulman et al. , 2016 ) , etc . However , RL algorithms developed for the online/interactive setting usually perform poorly in the offline setting ( Fujimoto et al. , 2019 ; Janner et al. , 2019 ) due to the data distribution shift caused by ( 1 ) the difference between the policy-in-training and the behavior policies used to collect the data ; and ( 2 ) the difference between the realistic environment in which we will deploy the policy and the environments used to collect the data . These differences can result in function approximation error and biased policy learning ( Levine et al. , 2020 ) . To address these challenges , model-based RL approaches ( Yu et al. , 2020 ; Kidambi et al. , 2020 ; Rafailov et al. , 2021 ; Yu et al. , 2021 ) firstly learn an environment model from a dataset of logged experiences using supervised learning and then conduct online RL interacting with the model . The learned environment model fully exploits the pre-collected experiences and can avoid/reduce the costly interactions with the realistic environment required by RL , hence improving the sample efficiency . That being said , due to the large size of state-action space , model-based offline RL approaches can still suffer from “ model exploitation ” ( Levine et al. , 2020 ) : when the dataset does not contain sufficient samples for some state-action pairs , the epistemic uncertainty of the model on these “ outof-distribution ” pairs can lead to a poor policy in offline RL . For example , an RL agent may easily exploit the model by repeatedly visiting the out-of-distribution states where the model erroneously issues higher rewards than the true environment dynamics , and thus the RL objective is biased by overly optimistic evaluation of the agent ’ s performance . A widely-studied strategy to tackle this problem is applying regularization to the reward function , which keeps the model uncertainty small when maximizing the reward issued by the model ( Yu et al. , 2020 ; Kidambi et al. , 2020 ; Rafailov et al. , 2021 ) . By carefully tuning the regularization weight , it can achieve preferable trade-off between the model reward and model uncertainty . However , it is usually challenging to tune the weight without access to the realistic environment that the learned policy will be deployed to . Moreover , even with access to the realistic environment , hyperparameter tuning methods such as Bayesian optimization ( Frazier , 2018 ) need to run multiple instances of offline RL and can be computationally expensive . In addition , the offline-trained policy may not be directly generalized to different environments . In this paper , we address a challenging problem , i.e. , how to balance between the model reward and model uncertainty during offline RL so it can produce diverse policies that can be adapted to different realistic environments ? Instead of formulating the problem as a single-objective optimization with regularization as in previous works , a primary contribution of this paper is to treat the model reward/return and model uncertainty as two separate objectives and develop an efficient bi-objective optimization method producing a pool of diverse policies on the Pareto front , which correspond to different trade-offs between the two objectives . Therefore , when deployed to a new realistic environment , we can accordingly choose the best policy from the pool . Our method is called “ Pareto policy pool ( P3 ) , ” and an example of the policies achieved by P3 is shown in Fig . 1 . When deployed to the realistic environment , Pareto policy 1 is overly optimal on the model return , so it runs fast at the beginning but quickly falls to the ground due to the aforementioned “ model exploitation ” . On the contrary , Pareto policy 4 favoring small model uncertainty is overly conservative and keeps standing still for 1000 time-steps because it avoids taking exploratory actions that potentially increase the uncertainty . As expected , Pareto policy 2 & 3 with the more balanced trade-off between the model return and uncertainty perform better and achieve higher returns in the test environment . These results imply that model-based offline RL ’ s performance significantly relies on the trade-off between model return and uncertainty in the optimization objectives . However , it is usually challenging or intractable to determine the optimal trade-off before deployment and control it during training . Moreover , for existing methods adopting a regularized single-objective , i.e. , scalarization ( Boyd & Vandenberghe , 2004 , Chapter 4.7 ) , even trying all possible regularization weights can not fully recover the Pareto front and thus can not guarantee to find the optimal trade-off . For example , in Fig . 1 , by running multiple instances with different regularization weights , MOPO ( Yu et al. , 2020 ) and MOReL ( Kidambi et al. , 2020 ) can only generate a few separated solutions , and it is difficult to find one with advantageous trade-off among them . In contrast , our method aims at efficiently generating a rich pool of diverse and representative policies covering the entire Pareto front without tuning the trade-off during training . Thereby , when deployed to a new realistic environment , the best policy for the new environment can be efficiently selected from the pool . Hence , P3 provides a simple and principal approach that addresses the two major challenges in model-based offline RL , i.e. , “ model exploitation ” and generalization to different unseen states in order to achieve high returns . Due to the complicated shape of the Pareto front that is unknown during training , finding a pool of diverse policies covering the whole Pareto front raises several non-trivial algorithmic challenges : ( 1 ) How to find policies located at different regions of the Pareto front associated with different levels of model return-uncertainty trade-off ? ( 2 ) How to avoid training each policy from scratch and reduce the computational cost linearly increasing with the number of policies ? Inspired by recent works in multi-objective optimization ( Cheng et al. , 2016 ; Ma et al. , 2020 ; Xu et al. , 2020 ) , we explore different regions of the Pareto front by generating multiple diverse reference vectors in the bi-objective space , each defining a constrained bi-objective optimization whose solution resides in a local region of the Pareto front . By solving these constrained bi-objective optimization problems , we can obtain a diverse set of policies covering the whole Pareto front . For solving each problem , we extend MGDA algorithm ( Désidéri , 2012 ) to be a two-stage gradient-based method that provably converges to the Pareto front region targeted by the reference vector . In order to achieve more policies on the Pareto front without linearly increasing the cost , we start from the previously obtained policies for initialization and explore their neighborhoods on the Pareto front by perturbing their corresponding reference vectors , resulting in a dense set of Pareto policies in each local region . In experiments , we evaluate P3 and compare it with several state-of-the-art model-based/free offline RL methods on the standard D4RL Gym benchmark ( Fu et al. , 2020 ) . P3 achieves the highest average-score over all datasets and significantly outperforms the baseline methods in 5 out of the 9 low/medium-quality datasets , showing the advantages of P3 on learning from non-expert experiences . We also present a thorough ablation study to identify the most important components in P3 . 2 RELATED WORK . Due to lack of space , we focus our discussion here on directly related works and present a more detailed overview of related work in Appendix A.2 . To address the aforementioned “ model exploitation ” , recent works rely on applying uncertainty regularization to the model return ( Yu et al. , 2020 ; Kidambi et al. , 2020 ; Rafailov et al. , 2021 ) which can be difficult and costly to tune the weight of regularization . In contrast , our work reframes the policy learning under the environment model as a bi-objective optimization problem and produces a diverse set of policies on the Pareto front performing different levels of model return-uncertainty trade-off , which provides flexibility to select the best policy in the inference stage . 3 PRELIMINARIES . We consider an episodic Markov Decision Process ( MDP ) M = { S , A , P , r , H , ρ0 } , where S is the state space , A is the space of actions , P is the transition probability : S ×A → ∆ ( S ) where ∆ ( S ) is a probability distribution over S , r : S ×A → R is the reward function so r ( sh , ah ) is the immediate reward for taking action ah at state sh , H is the horizon of the process , and ρ0 is the distribution of the initial state s0 . RL aims at learning a policy π : S → A maximizing the expected return in Eq . ( 1 ) , where π in this paper takes a deterministic action ah based on the state sh . max πθ Rρ0 ( πθ , M ) = Es0∼ρ0 , πθ [ H−1∑ h=0 r ( sh , ah ) ] . ( 1 ) In model-based offline RL , the agent instead interacts with an environment model rather than the realistic one . We train an environment model M̂ = { Ŝ , A , P̂ , r̂ , H , ρ̂0 } using a pre-collected dataset D , { ( sh , ah , sh+1 , rh ) |πb } of experiences by behavior policies , hand-designed controllers , or human demonstrators . By interacting with the model , online RL methods can learn in an offline manner . However , when D does not contain sufficient samples for some state-action pairs , the epistemic uncertainty of the model on these “ out-of-distribution ( OOD ) ” pairs can result in a poor policy . For example , an RL agent may easily exploit the model by repeatedly visiting the OOD states where the model erroneously issues r̂ higher than the true environment reward r and thus the RL objective is biased by overly optimistic evaluation of the agent ’ s performance . A widely-studied strategy to tackle this problem is applying regularization to r̂ , which keeps the model uncertainty small while maximizing the model reward ( Yu et al. , 2020 ; Kidambi et al. , 2020 ; Rafailov et al. , 2021 ) . A practical implementation of the regularized reward function is developed by ( Yu et al. , 2020 ) : r̃h = r̂h − λu ( ŝh , ah ) , where u ( ŝh , ah ) denotes the estimation of the model uncertainty at the state-action pair ( ŝh , ah ) and λ controls the trade-off between r̂ and u , which has to be carefully tuned in practice . Based on the regularized reward function , Yu et al . ( 2020 ) proposed a modified policy optimization objective : max π R̃ρ̂0 ( π , M̂ ) = Es0∼ρ̂0 , π [ H−1∑ h=0 ( r̂ ( ŝh , ah ) − λu ( ŝh , ah ) ) ] . ( 2 ) Despite being intuitive , this method ’ s performance is sensitive to the regularization weight λ ( Kidambi et al. , 2020 ; Yu et al. , 2021 ) , and tuning λ is usually challenging without access to the realistic environment that the learned policy will be deployed to . Moreover , even granted the access , hyperparameter tuning methods such as Bayesian optimization ( Frazier , 2018 ) require running many instances of offline RL , which can be computationally prohibitive . Instead of optimizing a single regularized objective , our “ Pareto policy pool ( P3 ) ” proposed later treats r̂ and u as two separate objectives and develop an efficient bi-objective optimization method that does not need to tune the trade-off deliberately but produces a pool of diverse policies , which are learned under different trade-offs between the two objectives . Thereby , when deployed to a realistic environment , the best policy can be chosen from the pool .

Summary: The paper proposes a new model-based offline reinforcement learning, where a pool of Pareto optimal policies is trained on a model of the environment provided by offline RL data. First, a model of the environment is trained using supervised learning from the dataset of logged experiences and subsequently a set of Pareto optimal policies is trained along the dimension of return in the MDP model and uncertainty in the MDP model. The authors argue that providing a Pareto optimal set of policies is superior to prior methods that relied on regularization of the two objectives into a combined metric, as having a Pareto optimal set of policies enables one too explore a range of behavior that is inherent with trade-off between model return vs model uncertainty in offline model based RL.  The authors then describe their method, Pareto Policy Pool (P3), which addresses the bi-objective optimization problem two stage method. Initially, the authors initialize a set of reference vectors along the Pareto frontiers to ensure a diverse set of Pareto policies can be found. Subsequently, the authors apply "Local Pareto Extension" where first a given policy is corrected to a desired range along the objective dimensions provided by the reference vectors. Once the policies are within the desired range, the method applies MGDA based optimization to further improve the performance of the policy in a Pareto optimal way (ascension stage). In order to alleviate computational cost of training multiple policies, the authors apply a "Pareto extension" method to initialize new policies along the Pareto frontier.   The authors then perform experiments on D4RL based benchmarks and compare their method to other literature methods and claim outperformance, particularly in lower quality dataset settings. The authors also perform an ablation study pertaining to the various components of their method, as well as an analysis related to why it is challenging to obtain optimal trade-offs in model based offline RL.