On Multi-objective Policy Optimization as a Tool for Reinforcement Learning: Case Studies in Offline RL and Finetuning

Many advances that have improved the robustness and efficiency of deep reinforcement learning ( RL ) algorithms can , in one way or another , be understood as introducing additional objectives or constraints in the policy optimization step . This includes ideas as far ranging as exploration bonuses , entropy regularization , and regularization toward teachers or data priors . Often , the task reward and auxiliary objectives are in conflict , and in this paper we argue that this makes it natural to treat these cases as instances of multi-objective ( MO ) optimization problems . We demonstrate how this perspective allows us to develop novel and more effective RL algorithms . In particular , we focus on offline RL and finetuning as case studies , and show that existing approaches can be understood as MO algorithms relying on linear scalarization . We hypothesize that replacing linear scalarization with a better algorithm can improve performance . We introduce Distillation of a Mixture of Experts ( DiME ) , a new MORL algorithm that outperforms linear scalarization and can be applied to these non-standard MO problems . We demonstrate that for offline RL , DiME leads to a simple new algorithm that outperforms state-of-the-art . For finetuning , we derive new algorithms that learn to outperform the teacher policy . 1 INTRODUCTION . Deep reinforcement learning ( RL ) algorithms have solved a number of challenging problems , including in games ( Mnih et al. , 2015 ; Silver et al. , 2016 ) , simulated continuous control ( Heess et al. , 2017 ; Peng et al. , 2018 ) , and robotics ( OpenAI et al. , 2018 ) . The standard RL setting appeals through its simplicity : an agent acts in the environment and can discover complex solutions simply by maximizing cumulative discounted reward . In practice , however , the situation is often more complicated . For instance , without a carefully crafted reward function or sophisticated exploration strategy , learning may require hundreds of millions of environment interactions , or may not be possible at all . A number of strategies have been developed to mitigate the shortcomings of the pure RL paradigm . These include strategies that regularize the final solution , for instance by maximizing auxiliary rewards ( Jaderberg et al. , 2017 ) or the entropy of the policy ( Mnih et al. , 2016 ; Haarnoja et al. , 2018 ) . Others facilitate exploration , for instance via demonstrations ( Brys et al. , 2015 ) , shaping rewards ( Ng et al. , 1999 ) , exploration bonuses ( Bellemare et al. , 2016 ; Pathak et al. , 2017 ) , or guidance from a pre-trained teacher policy , known as policy finetuning ( Xie et al. , 2021 ) . Strategies for reusing old environment interactions can in the extreme case enable agents to learn solely from a fixed dataset , known as offline RL ( Fu et al. , 2020 ; Gulcehre et al. , 2020 ) . In this case regularization is needed to ensure that the learned policy is well supported by the data distribution . These approaches may seem quite different at first glance , but they can all be understood as introducing additional objectives , or constraints , in the policy optimization step . Since reward maximization and the additional objective ( s ) are usually in conflict , these algorithmic approaches can also be understood as solving a multi-objective RL ( MORL ) problem . In practice , the solution to this MORL problem is usually implemented via a form of linear scalarization ( LS ) , i.e. , by taking a weighted sum of the objectives . However , LS suffers from a number of limitations : it is sensitive to the scale of the objectives ( Abdolmaleki et al. , 2020 ) and is limited in the solutions it can find ( Das & Dennis , 1997 ) . Our key insight is that an interpretation of algorithmic approaches , such as the above , in terms of MORL provides us with novel tools to solve fundamental challenges in RL . To demonstrate this , we focus on offline RL and finetuning as case studies . We show recent offline RL algorithms ( e.g. , CRR ( Wang et al. , 2020 ) , AWAC ( Nair et al. , 2020 ) ) can be understood as using LS to solve the MO problem . We hypothesize that by replacing LS with a more appropriate approach , we can obtain betterperforming algorithms . We thus propose a new MORL approach , called Distillation of a Mixture of Experts ( DiME ) , that can replace LS in this setting . We derive novel algorithms for finetuning and offline RL , based on DiME . In offline RL , this new algorithm outperforms state-of-the-art ( Sec . 6 ) . Our main contributions are as follows : • We propose that MORL can be used as a tool for tackling fundamental challenges in RL . We show that existing algorithms ( e.g. , for offline RL ) rely on LS to solve the MO problem . Based on this insight , we explore replacing LS with a new , improved MO algorithm . • By replacing LS , we derive a novel offline RL algorithm from the multi-objective perspective ( Sec . 5.2 ) , that is simple , intuitive and outperforms the state-of-the-art ( Sec . 6.2 ) . • In a similar manner , we also derive new finetuning algorithms ( Sec . 5.1 ) , and show that these algorithms can train policies that exceed the teacher ’ s performance ( Sec . 6.3 ) . • To replace LS , we introduce a new MORL approach , DiME . We support DiME by theoretical derivations ( Sec . 4.1 ) and show it outperforms LS on standard MORL tasks ( Sec . 6.1 ) . 2 RELATED WORK . Offline RL and Finetuning . Our case studies in this work will involve two related learning problems : offline ( or batch ) RL ( Levine et al. , 2020 ; Lange et al. , 2012 ) and policy finetuning ( Xie et al. , 2021 ) . Both leverage existing information , either in the form of a dataset or a teacher policy . In this work , we use proximity to the dataset or policy as an additional objective in a MORL problem . In offline RL , the goal is to learn a policy from a given , fixed dataset of transitions . While classical off-policy RL algorithms can be used , they perform worse than bespoke algorithms ( Fujimoto et al. , 2019 ) . This is likely due to the extrapolation errors incurred when learning a Q-function—naïve policy optimization could optimistically exploit such errors . Therefore , many existing approaches attempt to constrain the policy optimization to stay close to the samples present in the available dataset ( Peng et al. , 2019 ; Nair et al. , 2020 ; Siegel et al. , 2020 ; Wang et al. , 2020 ) . Finetuning seeks to improve upon an expert ( or teacher ) policy , which can be specified either nonparametrically via a dataset of transitions or as a parametric model . In contrast to offline RL , finetuning happens in an online setting , where the agent is able to interact with the environment . Existing approaches for finetuning minimize the cross-entropy between the teacher and agent policies ( Schmitt et al. , 2018 ) , based on distillation or mimicking ( Rusu et al. , 2015 ; Parisotto et al. , 2015 ) . Multi-Objective RL . MORL algorithms train policies for environments with multiple sources of reward . There are single-policy and multi-policy approaches . The former finds a policy that is optimal for a single preference trade-off across rewards . One common technique is to use the trade-off to combine the rewards into a single scalar reward , and then use standard RL to optimize this ( Roijers et al. , 2013 ) . Typically linear scalarization is used , but it is sensitive to reward scales and can not always find optimal solutions ( Das & Dennis , 1997 ) . Non-linear scalarizations exist ( Van Moffaert et al. , 2013 ; Golovin & Zhang , 2020 ) , but are hard to combine with value-based RL . Instead of scalarization , recent work on MO-MPO ( Abdolmaleki et al. , 2020 ) relies on constrained optimization . Multi-policy approaches seek to find a set of policies that cover the Pareto front . Various techniques exist , for instance repeatedly calling a single-policy approach with strategically-chosen trade-off settings ( Roijers et al. , 2014 ; Mossalam et al. , 2016 ; Zuluaga et al. , 2016 ) , simultaneously learning a set of policies by using a multi-objective variant of Q-learning ( Moffaert & Nowé , 2014 ; Reymond & Nowé , 2019 ; Yang et al. , 2019 ) , learning a manifold of policies in parameter space ( Parisi et al. , 2016 ; 2017 ) , or combining single-policy approaches with an overarching objective ( Xu et al. , 2020 ) . In contrast , here we consider objectives that are not only environment rewards , but also learning- or regularization-focused such as staying close to a behavioral prior . State-of-the-art MORL approaches , including MO-MPO , can not be readily applied . We thus introduce a new algorithm , DiME , which is on-par with state-of-the-art and can handle these non-standard objectives ( Sec . 4 ) . We use DiME as a vehicle to demonstrate the benefits of taking a multi-objective perspective toward challenges in RL . 3 BACKGROUND . We consider the RL problem defined by a Markov Decision Process ( MDP ) . An MDP consists of states s ∈ S , actions a ∈ A , an initial state distribution p ( s0 ) , transition probabilities p ( st+1|st , at ) , a reward function r ( s , a ) ∈ R , and a discount γ ∈ [ 0 , 1 ) . A policy πθ ( a|s ) is a state-conditional distribution over actions , parametrized by θ . Together with the transition probabilities , this gives rise to a state visitation distribution µ ( s ) . The action-value function is the expected return from choosing action a in state s and then following policy π : Qπ ( s , a ) = Eπ [ ∑∞ t=0 γ tr ( st , at ) |s0 = s , a0 = a ] . This can also be written as Qπ ( s , a ) = Es′∼p [ r ( s , a ) + γV π ( s′ ) ] , where V π ( s ) = Ea∼π [ Qπ ( s , a ) ] . In this paper , both the LS and DiME approaches to multi-objective RL follow a policy iteration scheme , repeatedly alternating between policy evaluation and policy improvement . In the policy evaluation step , at iteration i the current policy πθi is evaluated by training a separate value function Qik for every objective k. This requires bootstrapping 1 according to Q-decomposition ( Russell & Zimdars , 2003 ) . Any policy evaluation algorithm can be used to learn these Q-functions . In the policy improvement step , given the previous policy πi ( short for πθi when clear from context ) and associated Q-functions { Qk } Kk=1 , the aim is to improve the policy for a given state visitation distribution µ.2 3.1 POLICY IMPROVEMENT : SINGLE OBJECTIVE . DiME introduces a new way to perform policy improvement with respect to multiple objectives , so we will now describe the policy improvement step , first for the single-objective setting . Policy improvement can itself be decomposed into the following two alternating steps ( Ghosh et al. , 2020 ) . Improvement . This step finds a ( nonparametric ) policy q that improves with respect to the Qfunction , while staying close to the current iterate πi , leading to the Lagrangian objective function F ( q ; θi ) = Es∼µ a∼q [ Q ( s , a ) ] − η Es∼µ [ DKL ( q ( ·|s ) ‖πi ( ·|s ) ) ] , ( 1 ) with the known analytic solution q ( a|s ) ∝ πi ( a|s ) exp ( Q ( s , a ) /η ) ( Peters et al. , 2010 ) . Projection . This step projects the improved policy to the space of parametric policies by maximizing the following ( which simplifies to policy gradient if one uses a parametric q in the improvement step ) : J ( θ ) = −Es∼µ DKL ( q ( ·|s ) ‖πθ ( ·|s ) ) . ( 2 )

Summary: Modern day policy (and value) function optimization procedures routinely include terms beyond the typical scalar reward, such as to induce regularization, encourage exploration and perform discriminative feature learning. How must one go about balancing these objectives? The typical scheme is to take a weighted sum; this paper recasts that into a multi-objective optimization question, and proposes an objective that combines policies from various objectives in a weighted KL sense.