What About Taking Policy as Input of Value Function: Policy-extended Value Function Approximator

1 INTRODUCTION . Reinforcement learning ( RL ) has been widely considered as a promising way to learn optimal policies in many decision making problems ( Mnih et al. , 2015 ; Lillicrap et al. , 2015 ; Silver et al. , 2016 ; You et al. , 2018 ; Schreck et al. , 2019 ; Vinyals et al. , 2019 ; Hafner et al. , 2020 ) . Lying in the heart of RL is the value function which defines the long-term evaluation of a policy . With function approximation ( e.g. , deep neural networks ) , a value function approximator ( VFA ) is able to approximate the values of a policy under large and continuous state spaces . As commonly recognized , most RL algorithms can be described as Generalized Policy Iteration ( GPI ) ( Sutton & Barto , 1998 ) . As illustrated in the left of Figure 1 , at each iteration the VFA is trained to approximate the true values of current policy , regarding which the policy are improved . However , value approximation can never be perfect and its quality influences the effectiveness of policy improvement , thus raising a requirement for better value approximation ( v. Hasselt , 2010 ; Bellemare et al. , 2017 ; Fujimoto et al. , 2018 ) . Since a conventional VFA only approximates the values ( i.e. , knowledge ( Sutton et al. , 2011 ) ) for one policy , the knowledge learned from previously encountered policies is not preserved and utilized for future learning in an explicit way . For example in GPI , a conventional VFA can not track the values of the changing policy by itself and has no idea of the direction of value generalization when approximating the values of a new policy . In this paper , we propose Policy-extended Value Function Approximator ( PeVFA ) , which additionally takes an explicit policy representation as input in contrast to conventional VFA . PeVFA is able to preserve values for multiple policies and induces an appealing characteristic , i.e. , value generalization among policies . We study the formal generalization and contraction conditions on the value approximation error of PeVFA , focusing specifically on value generalization along the policy improvement path which we call local generalization . Based on both theoretical and empirical evidences , we propose a new form of GPI with PeVFA ( the right of Figure 1 ) which can benefit from the closer approximation distance induced by local value generalization under some conditions ; thus , GPI with PeVFA is expected to be more efficient in consecutive value approximation along the policy improvement path . Moreover , we propose a framework to learn effective policy representation for an RL policy from policy network parameters and state-action pairs alternatively , through contrastive learning and an auxiliary loss of action prediction . Finally , based on Proximal Policy Optimization ( PPO ) , we derive a practical RL algorithm PPO-PeVFA from the above methods . Our experimental results demonstrate the effectiveness of both value generalization offered by PeVFA and policy representation learning . Our main contributions are summarized as follows : • We propose PeVFA which improves generalization of values among policies and provide a theoretical analysis of generalization especially in local generalization scenario . • We propose a new form of GPI with PeVFA resulting in closer value approximation along the policy improvement path demonstrated through experiments . • To our knowledge , we are the first to learn a representation ( low-dimensional embedding ) for an RL policy from its network parameters ( i.e. , weights and biases ) . 2 BACKGROUND . 2.1 REINFORCEMENT LEARNING . We consider a Markov Decision Process ( MDP ) defined as 〈S , A , r , P , γ〉where S is the state space , A is the action space , r is the reward function , P is the transition function and γ ∈ [ 0 , 1 ) is the discount factor . The goal of an RL agent is to learn a policy π ∈ Π where π ( a|s ) is a distribution of action given state that maximizes the expected long-term discounted return . The state-value function vπ ( s ) is defined in terms of the expected discounted return obtained through following the policy π from a state s : vπ ( s ) = Eπ [ ∑∞ t=0 γ trt+1|s0 = s ] for all s ∈ S where rt+1 = r ( st , at ) . We use V π = vπ ( · ) to denote the vector of values for all possible states . Value function is determined by policy π and environment models ( i.e. , P and r ) . For a conventional value function , policy is modeled implicitly within a table or a function approximator , i.e. , a mapping from only state to value . One can refer to Appendix E.1 to see a more detailed description . 2.2 EXTENSIONS OF CONVENTIONAL VALUE FUNCTION . Schaul et al . ( 2015 ) introduced Universal Value Function Approximators ( UVFA ) that generalize values over goals in goal-conditioned RL . Similar ideas are also adopted for low-level learning of Hierarchical RL ( Nachum et al. , 2018 ) . Such extensions are also studied in more challenging RL problems , e.g. , opponent modeling ( He & Boyd-Graber , 2016 ; Grover et al. , 2018 ; Tacchetti et al. , 2019 ) , and context-based Meta-RL ( Rakelly et al. , 2019 ; Lee et al. , 2020 ) . General value function ( GVF ) in ( Sutton et al. , 2011 ) are proposed as a form of general knowledge representation through cumulants instead of rewards . In an unified view , each approach generalizes different aspects of the conventional VFA , focusing on different components of the vector form Bellman equation ( Sutton & Barto , 1998 ) expanded on as discussed in Appendix E.2 . Concurrent to our work , several works also study to take policy as an input of value functions . Harb et al . ( 2020 ) propose Policy Evaluation Networks ( PVN ) to approximate the objective function J ( π ) = Es0∼ρ0 [ vπ ( s0 ) ] of different policy π , where ρ0 is the initial state distribution . Later in ( Raileanu et al. , 2020 ) , Policy-Dynamics Value Function ( PDVF ) is proposed which takes both policy and task context as additional inputs , for the purpose of value generalization among policies and tasks so that to adapt quickly to new tasks . PDVF can be viewed as an integration of PVN and task-specific context learning ( Rakelly et al. , 2019 ; Zhou et al. , 2019a ) . Both PVN and PDVF conduct value approximation for a given collection of policies and then optimize policy with gradients through policy-specific inputs ( shorted as GTPI below ) of well trained value function variants in a zero-shot manner . We view this as a typical case of global generalization discussed further in Sec . 3 . In contrast , we focus more on a local generalization scenario and utilize value generalization to improve learning during standard GPI process with no prior policies given . Closely related to our work , Faccio & Schmidhuber ( 2020 ) propose a class of Parameter-based Value Functions ( PVFs ) which take policy parameters as inputs . Based on PVFs , new policy gradients are introduced in the form of a combination of conventional policy gradients plus GTPI ( i.e. , by backpropagating through policy parameters in PVFs ) . Beyond zero-shot policy optimization , PVFs also utilize value generalization of PVFs in the iterative learning process . Our work differs with PVFs at two aspects : first , we consider a general policy representation in our proposed PeVFA along with a framework of policy representation learning , in contrast to parsing the policy parameters directly ; second , we do not resort to GTPI for the policy update in our algorithms and only utilize value generalization for more efficient value estimation in GPI . Moreover , in these previous works , it is not clear how the value generalization among policies can be and whether it is beneficial for learning or not . In this paper , we study the condition of beneficial value generalization from both theoretical and empirical lens . 3 POLICY-EXTEND VALUE FUNCTION APPROXIMATOR . In this paper , we propose Policy-extended Value Function Approximator ( PeVFA ) , an extension of the conventional value function that explicitly takes policy ( representation ) as input . Formally , consider a function g : Π → X = Rn in a general form that maps any policy π to a n-dimensional representation χπ = g ( π ) ∈ X . The policy-extended value function V : S × X → R , defines the values over state and policy space : V ( s , χπ ) = Eπ [ ∞∑ t=0 γtrt+1|s0 = s ] , for all s ∈ S , π ∈ Π . ( 1 ) When only one policy π is considered , V ( s , χπ ) is equivalent in definition to a conventional value function vπ ( s ) . Note that if the policy πω is explicitly parameterized and its parameters are used as the representation , i.e. , χπω = ω , PeVFA is equivalent to Parameter-based State Value Function ( PSVF ) in PVF family ( Faccio & Schmidhuber , 2020 ) . Similarly , we can define policy-extended action-value function Q ( s , a , χπ ) . We use V ( s , χπ ) for demonstration in this paper . The key point is , V ( s , · ) is not defined for a specific policy and is able to preserve the values of multiple policies . With function approximation , a PeVFA is expected to approximate values among policy space , i.e. , { vπ ( s ) } π∈Π . Formally , with a finite parameter space Θ , we consider the optimal parameter θ∗ with minimum approximation error over all possible states and policies as : θ∗ = arg min θ∈Θ F ( θ , Π ) , where F ( θ , Π ) = max π∈Π ‖Vθ ( π ) − V π‖ = max π∈Π ‖Vθ ( · , π ) − vπ ( · ) ‖∞ , where F ( θ , Π ) is the overall approximation error of Vθ which is defined as the maximum L∞ norm of value vector distance among Π , as a typical metric commonly adopted in study on value estimation approximation and policy iteration ( Kakade & Langford , 2002 ; Munos , 2003 ; Lagoudakis & Parr , 2003 ) . Therefore , the learning process of a PeVFA Vθ is θ → θ∗ . As commonly recognized , a conventional value function approximator ( VFA ) can generalize the values to unseen states or state-action pairs after properly training on known states and state-action pairs under a specific policy . Beyond this , PeVFA provides an appealing characteristic that allows values to be generalized among policies , evaluating new policies with the knowledge of old ones . In the following , we study the characteristic of value generalization theoretically ( Sec . 3.1 ) and provide some empirical evidences ( Sec . 3.2 ) , from which we finally introduce a new form of GPI ( Sec . 3.3 ) . 3.1 VALUE GENERALIZATION AMONG POLICIES . Concretely , we introduce two scenarios of value generalization : global generalization and local generalization , which are closely related to many RL problems . An illustration is show in Figure 2 . In this section , we first focus on a general two-policy case and study whether the learned value approximation of one policy can generalize to that of the other one . We propose a common form of generalization approximation error ( Theorem 1 ) and analyze the condition of generalization contraction ( Corollary 1 ) . Further , we go to the two generalization scenarios as illustrated in Figure 2 and especially introduce the closer approximation distance along the policy improvement path induced by PeVFA ( Corollary 2 ) , from which a more efficient form of GPI can be derived ( Sec . 3.3 ) . Formally , consider two policies π1 and π2 and a PeVFA Vθ to approximate the values of them . For convenience of demonstration , we consider an identical representation space ( i.e. , X = Π ) and an infinite parameter space Θ where zero approximation error can be achieved , i.e. , F ( θ∗ , Π ) = 0 . We then define the approximation loss f of Vθ for each policy π ∈ Π as fθ ( π ) = ‖Vθ ( π ) −Vθ∗ ( π ) ‖ = ‖Vθ ( · , π ) − Vθ∗ ( · , π ) ‖∞ ≥ 0 , which measures the value distance of θ to the optimal parameter at some policy π . In the following , we analyze the generalization performance on value estimation of unlearned policy π2 when PeVFA Vθ is only trained to approximate the values of policy π1 . We first introduce the two assumptions below : Assumption 1 ( Contraction ) . The learning process P of value approximation for π1 is γcontraction , that is fθ̂ ( π1 ) ≤ γfθ ( π1 ) where γ ∈ ( 0 , 1 ] and θ P ( π1 ) −−−−→ θ̂ . θ and θ̂ are the parameters of PeVFA before and after some learning of π1 ’ s values . We use f and f̂ as abbreviations for fθ and fθ̂ below for the ease of demonstration . Recent works ( Keskar et al. , 2017 ; Novak et al. , 2018 ; Wang et al. , 2018 ) suggest that generalization performance is related to local properties of the model . We then assume following smoothness property for following analysis : Assumption 2 ( Smoothness ) . f̂ is ( or has ) Lipschitz continuous ( gradient/Hessian ) at π1 with Lipschitz constant L̂0 , e.g. , |f̂ ( π ) − f̂ ( π1 ) | ≤ L̂0 · d ( π , π1 ) for π ∈ Π with some metric space ( Π , d ) . With above two assumptions , next we derive a general form of an upper bound of the generalized value approximation for unlearned policy π2 ( i.e. , f̂ ( π2 ) ) ) as follows : Theorem 1 . Under Assumption 1 and 2 , when f ( π1 ) ≤ f ( π2 ) , we have the following bound : f̂ ( π2 ) ≤ γf ( π2 ) ︸ ︷︷ ︸ generalized contraction + M ( π1 , π2 , L̂ ) ︸ ︷︷ ︸ locality margin . ( 2 ) The form of M depends on the smoothness property , e.g. , M ( π1 , π2 , L̂ ) = L̂0 · d ( π , π1 ) when f̂ is Lipschitz continuous . See Appendix A.1 for more instances ofM when higher-order smoothness properties ( Nesterov & Polyak , 2006 ) are considered . Proof . See Appendix A.1 . The main idea is to chain the Lipschitz continuity upper bound ( Assumption 2 ) , the contraction upper bound ( Assumption 1 ) , and the inequality of f . Remark 1 . The case f ( π1 ) ≤ f ( π2 ) considered in Theorem 1 can usually exist since only π1 ’ s values are learned . Under such circumstances , we suggest that the complementary case f ( π1 ) > f ( π2 ) is acceptable since the approximation error of π2 is already lower than the trained one . Theorem 1 provides a generalization upper bound for f̂ ( π2 ) , as a generalized contraction on f ( π2 ) plus a locality margin termM which is related to π1 , π2 and the smoothness property L̂ . Further , we analyze the condition when a PeVFA can also obtain a contraction approximation for unlearned policy π2 though only trained on π1 , as below : Corollary 1 . Followed by Theorem 1 , consider f̂ is Lipschitz continuous and f ( π2 ) 6= 0 , we have , f̂ ( π2 ) ≤ γgf ( π2 ) where γg = γ + L̂0·d ( π1 , π2 ) f ( π2 ) . When f ( π2 ) ≥ L̂0·d ( π1 , π2 ) 1−γ , θ P ( π1 ) −−−−→ θ̂ is also a γg-contraction for value approximation of π2 , with γg ∈ ( 0 , 1 ] . Proof . ReplaceM with Lipschitz continuous upper bound and transform the RHS in Equation 2 , then let RHS not greater than f ( π2 ) . See Appendix A.2 for complete derivation . Remark 2 . From the generalization contraction condition provided in Corollary 1 , we can find that : as i. γ → 0 , or ii . d ( π1 , π2 ) → 0 , or iii . L̂0 → 0 , the contraction condition is easier to achieve ( or the contraction gets tighter ) , i.e. , the generalization on unlearned policy π2 is better . In the above , we discuss value generalization in a two-policy case , i.e. , from one learned policy to another unlearned policy . Global generalization in Figure 2 ( a ) is an extension scenario of the two-policy case where values generalize from a known policy set ( π ∈ Π0 ) of which the values are learned to unseen policies ( π′ ∈ Π1 ) . In this paper , we focus more on local generalization of PeVFA as shown in Figure 2 ( b ) , Recall the GPI form shown in the left of Figure 1 , at each iteration t , value generalization from current policy πt to improved policy πt+1 is exactly the two-policy case we discussed above . However , it is unclear how local generalization can impact the value estimation in GPI ( i , e. , policy evaluation ) . To this end , we propose the following Corollary to see a connection between local generalization of PeVFA and a closer value approximation distance : Corollary 2 . At iteration t in local generalization scenario of PeVFA ( Figure 2 ( b ) ) , if fθt ( πt ) + fθt ( πt+1 ) ≤ ‖Vθ∗ ( πt ) − Vθ∗ ( πt+1 ) ‖ , then fθt ( πt+1 ) ≤ ‖Vθt ( πt ) − Vθ∗ ( πt+1 ) ‖ . Proof . Proof can be obtained by applying Triangle Inequality . See Appendix A.3 . Corollary 2 indicates that local generalization of PeVFA can induce a more preferable start point ( Vθt ( πt+1 ) ) which is closer to the optimal approximation target ( Vθ∗ ( πt+1 ) ) than the conventional one ( V πtθt , equivalent to Vθt ( πt ) in definition ) for policy evaluation process at iteration t+ 1 . Remark 3 . With closer start points Vθt ( πt+1 ) , policy evaluation ( i.e. , minimize fθt ( πt+1 ) ) can be more efficient with PeVFA . One can assume an ideal case with perfect local generalization , where policy evaluation is no longer necessary and policy improvement can be performed consecutively .

Summary: The authors propose PeVFA: a value function able to evaluate the expected return of multiple policies. They do so by extending the conventional value function, allowing it to receive as input the parameter (or a representation) of the policy. The authors study the local generalization property of PeVFA, propose possible ways of encoding the policy parameters and compare traditional PPO with an extended version of PPO using PeVFA. While the idea of generalization among many policies is an interesting topic in RL, there are many theoretical and experimental issues that prevent acceptance. Moreover, the authors do not at all compare their approach to recent work which also uses value functions with policy parameters as input.