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: This paper proposes a new algorithm for model-based offline reinforcement learning in which the learner attempts to optimally balance maximizing the reward under the learned model and minimizing the uncertainty of the model. The algorithm is motivated by the fact that offline RL methods suffer from distribution shift that often results in policies/predictions that are overly optimistic about highly uncertain areas of the state-action space. The authors’ proposed method attempts to address this problem by identifying a relevant collection of policies lying on the Pareto frontier that trades off these two criteria. The underlying supposition is that a well-performing policy should be among this collection of “non-dominated” policies on the frontier. Specialized subroutines are proposed in order to make sure the collection is diverse. Some theory is presented to validate certain aspects. Extensive experimental results support that the method performs well in relevant benchmarks compared to a number of state-of-the-art methods.