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 paper conditions the value function on a representation of the policy. The representation can be based on a batch of state-action pairs or based on the policy filters. When conditioning on the representation of a new policy, the value function can better approximate the value of the new policy. Experiments show benefits on continuous control tasks.