Represent Your Own Policies: Reinforcement Learning with Policy-extended Value Function Approximator

We study Policy-extended Value Function Approximator ( PeVFA ) in Reinforce-1 ment Learning ( RL ) , which extends conventional value function approximator2 ( VFA ) to take as input not only the state ( and action ) but also an explicit policy3 representation . Such an extension enables PeVFA to preserve values of multi-4 ple policies at the same time and brings an appealing characteristic , i.e. , value5 generalization among policies . We formally analyze the value generalization un-6 der Generalized Policy Iteration ( GPI ) . From theoretical and empirical lens , we7 show that generalized value estimates offered by PeVFA may have lower initial8 approximation error to true values of successive policies , which is expected to9 improve consecutive value approximation during GPI . Based on above clues , we10 introduce a new form of GPI with PeVFA which leverages the value generalization11 along policy improvement path . Moreover , we propose a representation learning12 framework for RL policy , providing several approaches to learn effective policy em-13 beddings from policy network parameters or state-action pairs . In our experiments,14 we evaluate the efficacy of value generalization offered by PeVFA and policy15 representation learning in several OpenAI Gym continuous control tasks . For a16 representative instance of algorithm implementation , Proximal Policy Optimization17 ( PPO ) re-implemented under the paradigm of GPI with PeVFA achieves about 40 % 18 performance improvement on its vanilla counterpart in most environments.19 1 Introduction20 Reinforcement Learning ( RL ) has been widely considered as a promising way to learn optimal21 policies in many decision-making problems [ 35 , 31 , 53 , 65 , 47 , 62 , 16 ] . One fundamental element of22 RL is value function which defines the long-term evaluation of a policy . With function approximation23 ( e.g. , deep neural networks ) , a value function approximator ( VFA ) is able to approximate the values24 of a policy under large and continuous state spaces . As commonly recognized , most RL algorithms25 can be described as Generalized Policy Iteration ( GPI ) [ 55 ] . As illustrated on the left of Figure 1,26 at each iteration the VFA is trained to approximate the true values of current policy ( i.e. , policy27 evaluation ) , regarding which the policy is further improved ( i.e. , policy improvement ) . The value28 function approximation error hinders the effectiveness of policy improvement and then the overall29 optimality of GPI [ 5 , 46 ] . Unfortunately , such errors are inevitable under function approximation . A30 large number of samples are usually required to ensure high-quality value estimates , resulting in the31 sample-inefficiency of deep RL algorithms . Therefore , this raises an urgent need for more efficient32 value approximation methods [ 61 , 4 , 12 , 25 ] .33 An intuitive idea to improve the efficiency value approximation is to leverage the knowledge on34 the values of previous encountered policies . However , a conventional VFA usually approximates35 the values of one policy and values learned from old policies are over-written gradually during36 Submitted to 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Do not distribute . the learning process . This means that the previously learned knowledge can not be preserved and37 utilized with one conventional VFA . Thus , such limitations prevent the potentials to leverage the38 previous knowledge for future learning . In this paper , we study Policy-extended Value Function39 Approximator ( PeVFA ) , which additionally takes an explicit policy representation as input in contrast40 to conventional VFA . Thanks to the policy representation input , PeVFA is able to approximate values41 for multiple policies and induces value generalization among policies . We formally analyze the42 generalization of approximate values among policies in a general form . From both theoretical and43 empirical lens , we show that the generalized value estimates can be closer to the true values of44 the successive policy , which can be beneficial to consecutive value approximation along the policy45 improvement path , called local generalization . Based on above clues , we introduce a new form46 of GPI with PeVFA ( the right of Figure 1 ) that leverages the local generalization to improve the47 efficiency of consecutive value approximation along the policy improvement path.48 One key point of GPI with PeVFA is the representation of policy since it determines how PeVFA gen-49 eralizes the values . For this , we propose a framework to learn effective low-dimensional embedding50 of RL policy . We use network parameters or state-action pairs as policy data and encode them into51 low-dimensional embeddings ; then the embeddings are trained to capture the effective information52 through contrastive learning and policy recovery . Finally , we evaluate the efficacy of GPI with PeVFA53 and our policy representations . In principle , GPI with PeVFA is general and can be implemented54 in different ways . As a practical instance , we re-implement Proximal Policy Optimization ( PPO ) 55 with PeVFA and propose PPO-PeVFA algorithm . Our experimental results on several OpenAI Gym56 continuous control tasks demonstrate the effectiveness of both value generalization offered by PeVFA57 and learned policy representations , with an about 40 % improvement in average returns achieved by58 our best variants on standard PPO in most tasks.59 We summarize our main contributions below . 1 ) We study the value generalization among policies60 induced by PeVFA . From both theoretical and empirical aspects , we shed the light on the situations61 where the generalization can be beneficial to the learning along policy improvement path . 2 ) We62 propose a framework for policy representation learning . To our knowledge , we make the first attempt63 to learn a low-dimensional embedding of over 10k network parameters for an RL policy . 3 ) We64 introduce GPI with PeVFA that leverages the value generalization in a general form . Our experimental65 results demonstrate the potential of PeVFA in deriving practical and more effective RL algorithms.66 2 Related Work67 Extensions of Conventional Value Function . Sutton et al . [ 56 ] propose General Value Functions68 ( GVFs ) as a general form of knowledge representation of rewards and arbitrary cumulants . Later,69 conventional value functions are extended to take extra inputs for different purposes of generalization.70 One notable work is Universal Value Function Approximator ( UVFA ) [ 45 ] , which is proposed to71 generalize values among different goals for goal-conditioned RL . UVFA is further developed in72 [ 1 , 37 , 9 ] and influences the occurrence of other value function extensions in context-based Meta-RL73 [ 43 , 29 ] , Hierarchical RL [ 64 ] and multiagent RL [ 19 , 14 ] and etc . Most of the above works study74 how to generalize the policy or value function among extrinsic factors , i.e. , environments , tasks and75 opponents ; while we mainly study the value generalization among policies along policy improvement76 path , an intrinsic learning process of the agent itself.77 Policy Embedding and Representation . Although not well studied , representation ( or embedding ) 78 learning for RL policies is involved in a few works [ 18 , 14 , 3 ] . The most common way to learn a79 policy representation is to extract from interaction experiences . As a representative , Grover et al . [ 14 ] 80 propose learning the representation of opponent policy from interaction trajectories with a generative81 policy recovery loss and a discriminative triplet loss . These losses are later adopted in [ 64 , 42 ] .82 Another straightforward idea is to represent policy parameters . Network Fingerprint [ 17 ] is such a83 differentiable representation that uses the concatenation of the vectors of action distribution outputted84 by policy network on a set of probing states . The probing state set is co-optimized along with the85 primary learning objective , which can be non-trivial especially when the dimensionality of the set is86 high . Besides , some early attempts in learning low-dimensional embedding of policy parameters are87 studies in Evolutionary Algorithms [ 13 , 44 ] , mainly with the help of VAE [ 23 ] . Our work introduce a88 learning framework of policy representation including both above two perspectives.89 PVN and PVFs . Recently , several works study the generalization among policy space . Harb et al.90 [ 17 ] propose Policy Evaluation Network ( PVN ) to directly approximate the distribution of policy91 π ’ s objective function J ( π ) = Eρ0 [ vπ ( s0 ) ] with initial state s0 ∼ ρ0 . PVN takes as input Network92 Fingerprint ( mentioned above ) of policy network . After training on a pre-collected set of policies , a93 random initialized policy can be optimized in a zero-shot manner with the policy gradients of PVN by94 backpropagting through the differentiable policy input . We call such gradients GTPI for short below.95 Similar ideas are later integrated with task-specific context learning in multi-task RL [ 42 ] , leveraging96 the generalization among policies and tasks for fast policy adaptation on new tasks . In PVN [ 17 ] ,97 as an early attempt , the generalization among policies is studied with small policy network and98 simple tasks ; besides , the most regular online learning setting is not studied . Concurrent to our work,99 Faccio and Schmidhuber [ 10 ] propose a class of Parameter-based Value Functions ( PVFs ) that take100 vectorized policy parameters as inputs . Based on PVFs , new policy gradient algorithms are introduced101 in the form of a combination of conventional policy gradients and GTPI ( i.e. , by backpropagating102 through policy parameters in PVFs ) . Except for zero-shot policy optimization as conducted in PVN,103 PVFs are also evaluated for online policy learning . Due to directly taking parameters as input , PVFs104 suffer from the curse of dimensionality when the number of parameters is high . Besides , GTPI can105 be non-trivial to rein since policy parameter space are complex and extrapolation generalization106 error can be large when the value function is only trained on finite policies ( usually much fewer than107 state-action samples ) thus further resulting in erroneous policy gradients.108 Our work differs with PVFs from several aspects . First , we make use of learned policy representation109 rather than policy network parameters . Second , we do not resort to GTPI for the policy update110 in our algorithms but focus on utilizing value generalization for more efficient value estimation in111 GPI . Furthermore , we shed the light on two important problems — how value generalization among112 policies can happen formally and whether it is beneficial to learning or not — which are neglected in113 in previous works from both theoretical and empirical lens.114 3 Policy-extended Value Function Approximator115 In this section , we propose Policy-extended Value Function Approximator ( PeVFA ) , an extension116 of conventional VFA that explicitly takes as input a policy representation . First , we introduce the117 formulation ( Sec . 3.1 ) , then we study value generalization among policies theoretically ( Sec . 3.2 ) 118 along with some empirical evidences ( Sec . 3.3 ) . Finally , we derive a new form of GPI ( Sec . 3.4 ) .119 3.1 Formulation120 Consider a Markov Decision Process ( MDP ) defined as 〈S , A , r , P , γ〉 where S is the state space , A121 is the action space , r is the ( bounded ) reward function , P is the transition function and γ ∈ [ 0 , 1 ) is122 the discount factor . A policy π ∈ P ( A ) |S| defines the distribution over all actions for each state . The123 goal of an RL agent is to find an optimal policy π∗ that maximizes the expected long-term discounted124 return . The state-value function vπ ( s ) is defined as the expected discounted return obtained through125 following the policy π from a state s : vπ ( s ) = Eπ [ ∑∞ t=0 γ trt+1|s0 = s ] for where rt+1 = r ( st , at ) .126 We use V π to denote the vectorized form of value function.127 In a general form , we define policy-extended value function V : S × Π→ R over state and policy128 space : V ( s , π ) = vπ ( s ) for all s ∈ S and π ∈ Π . In this paper , we focus on V ( s , π ) and policy-129 extended action-value function Q ( s , a , π ) can be obtained similarly . We use V ( π ) to denote the value130 vector for all states in the following . The key point is that PeVFA V is able to preserve the values of131 multiple policies . With function approximation , a PeVFA is expected to approximate the values of132 policies among policy space , i.e. , { V π } π∈Π and then enable value generalization among policies.133 Formally , given a function g : Π→ X ⊆ Rn that maps any policy π to an n-dimensional represen-134 tation χπ = g ( π ) ∈ X , a PeVFA Vθ with parameter θ ∈ Θ is to minimize the approximation error135 over all possible states and policies generally:136 Fµ , p , ρ ( θ , g , Π ) = ∑ π∈Π µ ( π ) ‖Vθ ( χπ ) − V π‖p , ρ , ( 1 ) where µ , ρ are distributions over policies and states respectively , ‖f‖p , ρ = ( ∫ s ρ ( ds ) |f ( s ) |p ) 1/p is137 ρ-weighted Lp-norm [ 26 , 46 ] for any f : S → R. The policy distribution µ of interest depends on138 the scenario where value generalization is considered . As illustrated in Figure 2 , we provide two139 value generalization scenarios . In the global generalization scenario , a uniform distribution over140 known policy set may be considered with a general purpose of value generalization for unknown141 policies . For the specific local generalization scenario along policy improvement path during GPI , a142 sophisticated distribution that adaptively weights recent policies more during the learning process143 may be more suitable in this case . In the following , we care more about the local generalization144 scenario and use uniform state distribution ρ and L2-norm for demonstration . The subscripts are145 omitted and we use ‖ · ‖ for clarity.146 3.2 Theoretical Analysis on Value Generalization among Policies147 In this part , we theoretically analyze the value generalization among policies induced by PeVFA . We148 start from a two-policy case and study whether the value approximation learned for one policy can be149 generalized to the other one . Later , we study the local generalization scenario ( Figure 2 ( b ) ) and shed150 the light on the superiority of PeVFA for GPI . All the proofs are provided in Appendix A.151 For the convenience of demonstration , we use an identical policy representation function , i.e. , χπ = π,152 and define the approximation loss of PeVFA Vθ for any policy π ∈ Π as fθ ( π ) = ‖Vθ ( π ) −V π‖ ≥ 0.153 We use the following definitions for a formal description of value approximation process with PeVFA154 and local property of loss function fθ that influences generalization [ 40 , 63 ] respectively:155 Definition 1 ( π-Value Approximation ) We define a value approximation process Pπ : Θ → Θ156 with PeVFA as a γ-contraction mapping on the approximation loss for policy π , i.e. , for θ̂ = Pπ ( θ ) ,157 we have fθ̂ ( π ) ≤ γfθ ( π ) where γ ∈ [ 0 , 1 ) .158 Definition 2 ( L-Continuity ) We call fθ is L-continuous at policy π if fθ is Lipschitz continuous at159 π with a constant L ∈ [ 0 , ∞ ) , i.e. , |fθ ( π ) − fθ ( π′ ) | ≤ L · d ( π , π′ ) for π′ ∈ Π with some distance160 metric d for policy space Π.161 With Definition 1 , the consecutive value approximation for the policies along policy improvement path162 during GPI can be described as : θ−1 Pπ0−−−→ θ0 Pπ1−−−→ θ1 Pπ2−−−→ . . . , as the green arrows illustrated in163 Figure 1 . One may refer to Appendix A.1 for a discussion on the rationality of the two definitions.164 To start our analysis , we first study the generalized value approximation loss in a two-policy case165 where only the value of policy π1 is approximated by PeVFA as below:166 Lemma 1 For θ Pπ1−−−→ θ̂ , if fθ̂ is L̂-continuous at π1 and fθ ( π1 ) ≤ fθ ( π2 ) , we have : fθ̂ ( π2 ) ≤167 γfθ ( π2 ) +M ( π1 , π2 , L̂ ) , whereM ( π1 , π2 , L̂ ) = L̂ · d ( π1 , π2 ) .168 Corollary 1 Pπ1 is γg-contraction ( γg ∈ [ 0 , 1 ) ) for π2 when fθ ( π2 ) > L̂·d ( π1 , π2 ) 1−γ .169 Lemma 1 shows that the post-Pπ1 approximation loss for π2 is upper bounded by a generalized170 contraction of prior loss plus a locality margin termM which is related to π1 , π2 and the locality171 property of fθ̂ . In general , the form ofM depends on the local property assumed . Some higher-172 order variants are provided in Appendix A.2 . For a step further , Corollary 1 reveals the condition173 where a contraction on value approximation loss for π2 is achieved when PeVFA is only trained to174 approximate the values of π1 . Concretely , such a condition is apt to reach with tighter contraction for175 policy π1 is , closer two policies , or smoother approximation loss function fθ̂.176 Then we consider the local generalization scenario as illustrated in Figure 2 ( b ) . For any iteration t177 of GPI , the values of current policy πt are approximated by PeVFA , followed by a improved policy178 πt+1 whose values are to be approximated in the next iteration . The value generalization from each179 πt and πt+1 can be similarly considered as the two-policy case . In addition to the former results , we180 shed the light on the value generalization loss of PeVFA along policy improvement path below:181 Lemma 2 For θ−1 Pπ0−−−→ θ0 Pπ1−−−→ θ1 Pπ2−−−→ . . . with γt for each Pπt , if fθt is L̂t-continuous at πt182 for any t ≥ 0 , we have fθt ( πt+1 ) ≤ γtfθt−1 ( πt ) +Mt , whereMt = Lt · d ( πt , πt+1 ) .183 Corollary 2 By induction , we have fθt ( πt+1 ) ≤ ∏t i=0 γtfθ−1 ( π0 ) + ∑t−1 i=0 ∏t j=i+1 γjMi +Mt .184 The above results indicate that the value generalization loss can be recursively bounded and has185 a upper bound formed by a repeated contraction on initial loss plus the accumulation of locality186 margins induced from each local generalization . An infinity-case discussion for Corollary 2 is in187 Appendix A.5 . The next question is whether PeVFA with value generalization among policies is188 preferable to the conventional VFA . To this end , we introduce a desirable condition which reveals the189 superiority of PeVFA during consecutive value approximation along the policy improvement path:190 Theorem 1 During θ−1 Pπ0−−−→ θ0 Pπ1−−−→ θ1 Pπ2−−−→ . . . , for any t ≥ 0 , if fθt ( πt ) + fθt ( πt+1 ) ≤191 ‖V πt − V πt+1‖ , then fθt ( πt+1 ) ≤ ‖Vθt ( πt ) − V πt+1‖.192 Theorem 1 shows that the generalized value estimates Vθt ( πt+1 ) can be closer to the true values of193 policy πt+1 than Vθt ( πt ) . Note that Vθt ( πt ) is the value approximation for πt which is equivalent194 to the counterpart Vφt for a conventional VFA as value generalization among policies does not195 exist . To consecutive value approximation along policy improvement path , this means that the value196 generalization of PeVFA has the potential to offer closer start points at each iteration . If such closer197 start points can often exist , we expect PeVFA to be preferable to conventional VFA since value198 approximation can be more efficient with PeVFA and it in turn facilitates the overall GPI process.199 However , the condition in Theorem 1 is not necessarily met in practice . Intuitively , it depends on the200 locality margins that may be related to function family and optimization method of PeVFA , as well201 as the scale of policy improvement . We leave these further theoretical investigations for future work.202 Instead , we empirically examine the existence of such desirable generalizations in the following.203 3.3 Empirical Evidences204 We empirically investigate the value generalization of PeVFA with didactic environments . In this205 section , PeVFA Vθ is parameterized by neural network and we use the concatenation of all weights206 and biases of the policy network as a straightforward representation χπ for each policy , called Raw207 Policy Representation ( RPR ) . Experimental details are provided in Appendix B.208 First , we demonstrate the global generalization ( illustrated in Figure 2 ( a ) ) in a continuous 2D Point209 Walker environment . We build the policy set Π with synthetic policies , each of which is a randomly210 initialized 2-layer tanh-activated neural network with 2 units for each layer . The size of Π is 20k and211 the behavioral diversity of synthetic policies is verified ( see Figure 7 ( b ) in Appendix ) . We divide Π212 into training set ( i.e. , known policies Π0 ) and testing set ( i.e. , unseen policies Π1 ) . We rollout the213 policies in the environment to collect trajectories , based on which we perform value approximation214 training . Our results show that a PeVFA trained on Π0 achieves reasonable generalization performance215 when evaluating on Π1 . The average losses on training and testing set are 1.782 and 2.071 over 6216 trials . Figure 3 ( a ) shows the value predictions for policies from training and testing set ( 100 for each ) .217 Next , we investigate the value generalization along policy improvement path , i.e. , local generalization218 as in Figure 2 ( b ) . We use a 2-layer 8-unit policy network trained by standard PPO algorithm [ 50 ] in219 MuJoCo continuous control tasks . Parallel to the conventional value network Vφ ( s ) ( i.e. , VFA ) in220 PPO , we set a PeVFA network Vθ ( s , χπ ) as a reference for the comparison on value approximation221 loss . Compared to Vφ , PeVFA Vθ ( s , χπ ) takes RPR as input and approximates the values of all222 historical policies ( { πi } ti=0 ) in addition . We compare the value approximation losses of Vφ ( red ) and223 Vθ ( blue ) before ( solid ) and after ( dashed ) updating with on-policy samples collected by the improved224 policy πt+1 at each iteration . Figure 3 ( b ) shows the results for InvertedPendulum-v1 and Ant-v1.225 Results for all 7 MuJoCo tasks can be found in Appendix B.2 . By comparing approximation losses226 before updating ( red and blue solid curves ) , we can observe that the approximation loss of Vθt ( χπt+1 ) 227 is almost consistently lower than that of Vφt . This means that the generalized value estimates228 offered by PeVFA are usually closer to the true values of πt+1 , demonstrating the consequence229 arrived in Theorem 1 . For the dashed curves , it shows that PeVFA Vθt+1 ( χπt+1 ) can achieve lower230 approximation loss for πt+1 than conventional VFA Vφt+1 after the same number of training with the231 same on-policy samples . The empirical evidence above indicates that PeVFA can be preferable to232 the conventional VFA for consecutive value approximation . The generalized value estimates along233 policy improvement path have the potential to expedite the process of GPI.234 3.4 Reinforcement Learning with PeVFA235 Based on the results above , we expect to leverage the value generalization of PeVFA to facilitate236 RL . In Algorithm 1 , we propose a general description of RL algorithm under the paradigm of237 GPI with PeVFA . For each iteration , the interaction experiences of current policy and the policy238 Algorithm 1 RL under the paradigm of GPI with PeVFA ( V ( s , χπ ) is used for demonstration ) 1 : Initialize policy π0 , policy representation model g , PeVFA V−1 and experience buffer D 2 : for iteration t = 0 , 1 , . . . do 3 : Rollout policy πt in the environment and obtain k trajectories Tt = { τi } ki=0 4 : Get representation χπt = g ( π ) for policy πt and add experiences ( χπt , Tt ) in buffer D 5 : if t % M = 0 then 6 : Update PeVFA Vt−1 ( s , χπi ) for previous policies with data { ( χπi , Ti ) } t−1i=0 7 : Update policy representation model g , e.g. , with approaches provided in Sec . 4 8 : end if 9 : Update PeVFA Vt−1 ( s , χπt ) for current policy χπt and set Vt ←− Vt−1 10 : Update πt w.r.t Vt ( s , χπt ) by policy improvement algorithm and set πt+1 ←− πt 11 : end for representation are stored in a buffer ( line 3-4 ) . At an interval of M iterations , PeVFA is trained via239 value approximation for previous policies with the stored data and the policy representation model240 is updated according to the method used ( line 5-8 ) . This part is unique to PeVFA for preservation241 and generalization of knowledge of historical policies . Next , value approximation for current policy242 is performed with PeVFA ( line 9 ) . A key difference here is that the generalized value estimates243 ( i.e. , Vt−1 ( χπt ) ) are used as start points . Afterwards , a successive policy is obtained from typical244 policy improvement ( line 10 ) . Algorithm 1 can be implemented in different ways and we propose an245 instance implemented based on PPO [ 50 ] in our experiments later . In the next section , we introduce246 our methods for policy representation learning.247 4 Policy Representation Learning248 To derive practical deep RL algorithms , one key point is policy representation , i.e. , a low-dimensional249 embedding of RL policy . Intuitively , policy representation influences the approximation and gener-250 alization of PeVFA . Thus , it is of interest to find an effective policy representation based on which251 the superiority of PeVFA can be leveraged to improve RL algorithms . To our knowledge , policy252 representation is not well studied and it remains unclear on how to obtain an effective representation253 for an RL policy in a general case in practice . In previous section , we demonstrate the effectiveness254 of using policy parameters as a naive representation when policy network is small , called RPR.255 However , a usual policy network may have large number of parameters , thus making it inefficient256 and even irrational to use RPR for approximation and generalization [ 17 , 10 ] . More generally , policy257 parameters of the policy we wish to represent may not be accessible.258 To this end , we propose a general framework of policy representation learning as illustrated in Figure259 4 . The first thing to consider is data source , i.e. , from which we can extract the information for an260 effective policy representation . Recall that the policy is a distribution over state and action space261 of high dimensionality . The features of such a distribution is not directly available . Therefore , we262 consider two kinds of data source below that indirectly contains the information of policies : 1 ) Surface263 Policy Representation ( SPR ) : The first data source is state-action pairs ( or trajectories [ 14 ] ) , since264 they reflect how policy may behave under such states . This data source is general since no explicit265 form of policy is assumed . In a geometric view , learning policy representation from state-action pairs266 can be viewed as capturing the features of policy via scattering sample points on the curved surface267 of policy distribution . 2 ) Origin Policy Representation ( OPR ) : The other data source is parameters of268 policy since they determine the underlying form of policy distribution . Such a data source is often269 available during the learning process of deep RL algorithms when policy is parameterized by neural270 networks . Generally , we consider a policy network to be an MLP with well represented state features271 ( e.g. , features extracted by CNN for pixels or by LSTM for sequences ) as input.272 The remaining question is how we extract the policy representation from the data sources mentioned273 above . As shown in Figure 4 , we use permutation-invariant ( PI ) transformations followed by an274 MLP to encode the data of policy π into an embedding χπ for both SPR and OPR . For SPR , each275 state-action pair of { ( si , ai ) } ki=1 is fed into a common MLP , followed by a Mean-Reduce operation276 on the outputted features across k. For OPR , we perform PI transformation ( similar as done for277 state-action pairs ) inner-layer weights and biases { ( wi , bi ) } hi=1 for each layer first , where h denotes278 the number of nodes in this layer and wi , bi is the income weight vector from previous layer and279 the bias of ith node ; then we concatenate encoding of layers and obtain the OPR . A illustrative280 description for the encoding of OPR is in Figure 12 of Appendix.281 To train the policy embedding χπ obtained above , the most straightforward way is to backpropagate282 the value approximation loss of PeVFA in an End-to-End ( E2E ) fashion as illustrated on the lower-283 right of Figure 4 . In addition , we provide two self-supervised training losses for both OPR and SPR,284 as illustrated on the upper-right of Figure 4 . The first one is an auxiliary loss ( AUX ) of policy recovery285 [ 14 ] , i.e. , to recover the action distributions of π from χπ under different states . To be specific,286 an auxiliary policy decoder π̄ ( ·|s , χπ ) is trained through behavioral cloning , formally to minimize287 cross-entropy objective LAUX = −E ( s , a ) [ log π̄ ( a|s , χπ ) ] . For the second one , we propose to train288 χπ by Contrastive Learning ( CL ) [ 54 , 51 ] : policies are encouraged to be close to similar ones ( i.e.,289 positive samples π+ ) , and to be apart from different ones ( i.e. , negative samples π− ) in representation290 space . For each policy , we construct positive samples by data augmentation on policy data , depending291 on SPR or OPR considered ; and different policies along the policy improvement path naturally292 provide negative samples for each other . Finally , the embedding χπ is optimized through minimizing293 the InfoNCE loss [ 41 ] below : LCL = −E ( π+ , { π− } ) [ log exp ( χTπWχπ+ ) exp ( χTπWχπ+ ) + ∑ π− exp ( χ T πWχπ− ) ] .294 Now , the training of policy representation model in Algorithm 1 can be performed with any com-295 bination of data sources and training losses provided above . A pseudo-code of the overall policy296 representation training framework and complete implementation details are provided in Appendix D.297 5 Experiments298 In this section , we conduct experimental study with focus on the following questions:299 Question 1 Can value generalization offered by PeVFA improve a deep RL algorithm in practice ? 300 Question 2 Can our proposed framework to learn effective policy representation ? 301 Our experiments are conducted in several OpenAI Gym continuous control tasks ( one from Box2D302 and five from MuJoCo ) [ 6 , 58 ] . All experimental details and curves can be found in Appendix B.303 Algorithm Implementation . We use PPO [ 50 ] as the basic algorithm and propose a representative304 implementation of Algorithm 1 , called PPO-PeVFA . PPO is a policy optimization algorithm that305 follows the paradigm of GPI ( Figure 1 , left ) . A value network Vφ ( s ) with parameters φ ( i.e.,306 conventional VFA ) is trained to approximate the value of current policy π ; while π is optimized with307 respect to a surrogate objective [ 48 ] using advantages calculated by Vφ and GAE [ 49 ] . Compared with308 original PPO , PPO-PeVFA makes use of a PeVFA network Vθ ( s , χπ ) with parameters θ rather than309 the conventional VFA Vφ ( s ) , and follows the training scheme as in Algorithm 1 . Note PPO-PeVFA310 uses the same policy optimization method as original PPO and only differs at value approximation.311 Baselines and Variants . Except for original PPO as a default baseline , we use another two baselines:312 1 ) PPO-PeVFA with randomly generated policy representation for each policy , denoted by Ran PR ; 313 2 ) PPO-PeVFA with Raw Policy Representation ( RPR ) , i.e. , use the vector of all parameters of policy314 network as representation as adopted in PVFs [ 10 ] . Our variants of PPO-PeVFA differ at the policy315 representation used . In total , we consider 6 variants denoted by the combination of the policy data316 choice ( i.e. , OPR , SPR ) and representation principle choice ( i.e. , E2E , CL , AUX ) .317 Experimental Details . For all baselines and variants , we use a normal-scale policy network with318 2 layers and 64 units for each layer , resulting in over 3k to 10k ( e.g. , Ant-v1 ) policy parameters319 depending on the environments . We do not assume the access to pre-collected policies . Thus the320 size of policy set increases from 1 ( i.e. , the initial policy ) during the learning process , to about 1k to321 2 for a single trial . The dimensionality of all kinds of policy representation expect for RPR is set322 to 64 . The buffer D maintains recent 200k steps of interaction experience and the policy data of323 corresponding policy . The number of interaction step of each trial is 1M for InvDouPend-v1 and324 LunarLander-v2 , 4M for Ant-v1 and 2M for the others.325 Results . The overall experimental results are summarized in Table 1 . In Figure 5 , we provide326 aggregated results across all environments expect for InvDouPend-v1 and LunarLander-v2 ( since327 most algorithms achieve near-optimal results ) , where all returns are normalized by the results of PPO328 in Table 1 . Full learning curves are omitted and can be found in Appendix F.2.329 To Question 1 . From Table 1 , we can find that both PPO-330 PeVFA w/ OPR ( E2E ) and PPO-PeVFA w/ SPR ( E2E ) 331 outperforms PPO in all 6 tasks , and achieve over 20 % 332 improvement in Figure 5 . This demonstrates the effec-333 tiveness of PeVFA . Moreover , the improvement is further334 enlarged ( to about 40 % ) by CL and AUX for both OPR335 and SPR . This indicates that the superiority of PeVFA can336 be further utilized with better policy representation that337 offers a more suitable space for value generalization.338 To Question 2 . In Table 1 , consistent degeneration is339 observed for PPO-PeVFA w/ Ran PR due to the nega-340 tive effects on generalization caused by the randomness341 and disorder of policy representation . This phenomenon342 seems to be more severe for PPO-PeVFA w/ RPR due343 to the complexity of high-dimensional parameter space.344 In contrast , the improvement achieved by our proposed345 PPO-PeVFA variants shows that effective policy repre-346 sentation can be learned from policy parameters ( OPR ) 347 and state-action pairs ( SPR ) though value approximation348 loss ( i.e. , E2E ) and further improved when additional self-349 supervised representation learning is involved as CL and350 AUX . Overall , OPR slightly outperforms SPR as CL does351 over AUX . We hypothesize that it is due to the stochas-352 ticity of state-action pairs which serve as inputs of SPR353 and training samples for AUX . This reveals the space for354 future improvement . In addition , we visualize the learned355 representation in Figure 6 . We can observe that policies356 from different trials are locally continuous and show dif-357 ferent modes of embedding trajectories due to random358 initialization and optimization ; while a global evolvement359 among trials emerges with respect to policy performance.360 6 Conclusion and Future Work361 In this paper , we propose Policy-extended Value Function Approximator ( PeVFA ) and study value362 generalization among policies . We propose a new form of GPI based on PeVFA which is potentially363 preferable to conventional VFA for value approximation . Moreover , we propose a general framework364 to learn low-dimensional embedding of RL policy . Our experiments demonstrate the effectiveness of365 the generalization characteristic of PeVFA and our proposed policy representation learning methods.366 Our work opens up some research directions on value generalization among policies and policy367 representation . A possible future study on the theory of value generalization among policies is to368 consider the interplay between approximation error , policy improvement and local generalization369 during GPI with PeVFA . Besides , analysis on influence factors of value generalization among policies370 ( e.g. , policy representation , architecture of PeVFA ) and other utilization of PeVFA are expected . For371 better policy representation , inspirations on OPR may be got from studies on Manifold Hypothesis of372 neural network ; the selection of more informative state-action pairs for SPR is also worth research.373 References374 [ 1 ] M. Andrychowicz , D. Crow , A. Ray , J. Schneider , R. Fong , P. Welinder , B. McGrew , J. Tobin , P. Abbeel,375 and W. Zaremba . Hindsight experience replay . In NeurIPS , pages 5048–5058 , 2017.376 [ 2 ] M. Andrychowicz , A. Raichuk , P. Stanczyk , M. Orsini , S. Girgin , R. Marinier , L. Hussenot , M. Geist,377 O. Pietquin , M. Michalski , S. Gelly , and O. Bachem . What matters in on-policy reinforcement learning ? A378 large-scale empirical study . CoRR , abs/2006.05990 , 2020.379 [ 3 ] I. Arnekvist , D. Kragic , and J . A. Stork . VPE : variational policy embedding for transfer reinforcement380 learning . In ICRA , pages 36–42 , 2019.381 [ 4 ] M. G. Bellemare , W. Dabney , and R. Munos . A distributional perspective on reinforcement learning . In382 ICML , volume 70 , pages 449–458 , 2017.383 [ 5 ] D. P. Bertsekas and J. N. Tsitsiklis . Neuro-dynamic programming , volume 3 of Optimization and neural384 computation series . 1996.385 [ 6 ] G. Brockman , V. Cheung , L. Pettersson , J. Schneider , J. Schulman , J. Tang , and W. Zaremba . Openai gym.386 CoRR , abs/1606.01540 , 2016.387 [ 7 ] T. Chen , S. Kornblith , M. Norouzi , and G. E. Hinton . A simple framework for contrastive learning of388 visual representations . CoRR , abs/2002.05709 , 2020.389 [ 8 ] L. Engstrom , A. Ilyas , S. Santurkar , D. Tsipras , F. Janoos , L. Rudolph , and A. Madry . Implementation390 matters in deep RL : A case study on PPO and TRPO . In ICLR , 2020.391 [ 9 ] B. Eysenbach , X. Geng , S. Levine , and R. R. Salakhutdinov . Rewriting history with inverse RL : hindsight392 inference for policy improvement . In NeurIPS , 2020.393 [ 10 ] F. Faccio and J. Schmidhuber . Parameter-based value functions . CoRR , abs/2006.09226 , 2020.394 [ 11 ] H. Fu , H. Tang , J. Hao , C. Chen , X. Feng , D. Li , and W. Liu . Towards effective context for meta-395 reinforcement learning : an approach based on contrastive learning . CoRR , abs/2009.13891 , 2020.396 [ 12 ] S. Fujimoto , H. v. Hoof , and D. Meger . Addressing function approximation error in actor-critic methods.397 In ICML , 2018.398 [ 13 ] A. Gaier , A. Asteroth , and J. Mouret . Discovering representations for black-box optimization . In GECCO,399 pages 103–111 , 2020.400 [ 14 ] A. Grover , M. Al-Shedivat , J. K. Gupta , Y. Burda , and H. Edwards . Learning policy representations in401 multiagent systems . In ICML , volume 80 , pages 1797–1806 , 2018.402 [ 15 ] T. Haarnoja , A. Zhou , P. Abbeel , and S. Levine . Soft actor-critic : Off-policy maximum entropy deep403 reinforcement learning with a stochastic actor . In ICML , pages 1856–1865 , 2018.404 [ 16 ] D. Hafner , T. P. Lillicrap , J. Ba , and M. Norouzi . Dream to control : Learning behaviors by latent405 imagination . In ICLR , 2020.406 [ 17 ] J. Harb , T. Schaul , D. Precup , and P. Bacon . Policy evaluation networks . CoRR , abs/2002.11833 , 2020.407 [ 18 ] K. Hausman , J. Tobias Springenberg , Z. Wang , N. Heess , and M. A. Riedmiller . Learning an embedding408 space for transferable robot skills . In ICLR , 2018.409 [ 19 ] H. He and J. L. Boyd-Graber . Opponent modeling in deep reinforcement learning . In ICML , volume 48,410 pages 1804–1813 , 2016.411 [ 20 ] K. He , H. Fan , Y. Wu , S. Xie , and R. B. Girshick . Momentum contrast for unsupervised visual representation412 learning . In CVPR , pages 9726–9735 , 2020.413 [ 21 ] M. Igl , G. Farquhar , J. Luketina , W. Boehmer , and S. Whiteson . The impact of non-stationarity on414 generalisation in deep reinforcement learning . CoRR , abs/2006.05826 , 2020.415 [ 22 ] N. S. Keskar , D. Mudigere , J. Nocedal , M. Smelyanskiy , and P. T. P. Tang . On large-batch training for416 deep learning : Generalization gap and sharp minima . In ICLR , 2017.417 [ 23 ] D. P. Kingma and M. Welling . Auto-encoding variational bayes . In ICLR , 2014.418 [ 24 ] I. Kostrikov , D. Yarats , and R. Fergus . Image augmentation is all you need : Regularizing deep reinforcement419 learning from pixels . CoRR , abs/2004.13649 , 2020.420 [ 25 ] A. Kuznetsov , P. Shvechikov , A. Grishin , and D. P. Vetrov . Controlling overestimation bias with truncated421 mixture of continuous distributional quantile critics . In ICML , volume 119 , pages 5556–5566 , 2020.422 [ 26 ] M. G. Lagoudakis and R. Parr . Least-squares policy iteration . J. Mach . Learn . Res. , 4:1107–1149 , 2003.423 [ 27 ] Q. Lan , Y. Pan , A. Fyshe , and M. White . Maxmin q-learning : Controlling the estimation bias of q-learning.424 In ICLR , 2020.425 [ 28 ] M. Laskin , K. Lee , A. Stooke , L. Pinto , P. Abbeel , and A. Srinivas . Reinforcement learning with augmented426 data . CoRR , abs/2004.14990 , 2020.427 [ 29 ] K. Lee , Y. Seo , S. Lee , H. Lee , and J. Shin . Context-aware dynamics model for generalization in428 model-based reinforcement learning . CoRR , abs/2005.06800 , 2020.429 [ 30 ] A . Levy , G. D. Konidaris , R. P. Jr. , and K. Saenko . Learning multi-level hierarchies with hindsight . In430 ICLR , 2019.431 [ 31 ] T. P. Lillicrap , J. J . Hunt , A. Pritzel , N. Heess , T. Erez , Y. Tassa , D. Silver , and D. Wierstra . Continuous432 control with deep reinforcement learning . In ICLR , 2015.433 [ 32 ] H. Liu , K. Simonyan , and Y. Yang . DARTS : differentiable architecture search . In ICLR , 2019.434 [ 33 ] R. Luo , F. Tian , T. Qin , E. Chen , and T. Liu . Neural architecture optimization . In NeurIPS , pages435 7827–7838 , 2018.436 [ 34 ] L. V. D. Maaten and G. E. Hinton . Visualizing data using t-sne . Journal of Machine Learning Research , 9:437 2579–2605 , 2008.438 [ 35 ] V. Mnih , K. Kavukcuoglu , D. Silver , A . A. Rusu , J. Veness , M. G. Bellemare , A. Graves , M. A. Riedmiller,439 A. Fidjeland , G. Ostrovski , S. Petersen , C. Beattie , A. Sadik , I. Antonoglou , H. King , D. Kumaran,440 D. Wierstra , S. Legg , and D. Hassabis . Human-level control through deep reinforcement learning . Nature,441 518 ( 7540 ) :529–533 , 2015.442 [ 36 ] V. Mnih , A. P. Badia , M. Mirza , A. Graves , T. P. Lillicrap , T. Harley , D. Silver , and K. Kavukcuoglu.443 Asynchronous methods for deep reinforcement learning . In ICML , 2016.444 [ 37 ] O. Nachum , S. Gu , H. Lee , and S. Levine . Data-efficient hierarchical reinforcement learning . CoRR,445 abs/1805.08296 , 2018.446 [ 38 ] O. Nachum , S. Gu , H. Lee , and S. Levine . Near-optimal representation learning for hierarchical reinforce-447 ment learning . In ICLR , 2019.448 [ 39 ] Y. E. Nesterov and B. T. Polyak . Cubic regularization of newton method and its global performance . Math.449 Program. , 108 ( 1 ) :177–205 , 2006.450 [ 40 ] R. Novak , Y. Bahri , D. A. Abolafia , J. Pennington , and J. Sohl-Dickstein . Sensitivity and generalization in451 neural networks : an empirical study . In ICLR , 2018.452 [ 41 ] A. Oord , Y. Li , and O. Vinyals . Representation learning with contrastive predictive coding . CoRR,453 abs/1807.03748 , 2018.454 [ 42 ] R. Raileanu , M. Goldstein , A. Szlam , and R. Fergus . Fast adaptation via policy-dynamics value functions.455 CoRR , abs/2007.02879 , 2020.456 [ 43 ] K. Rakelly , A. Zhou , C. Finn , S. Levine , and D. Quillen . Efficient off-policy meta-reinforcement learning457 via probabilistic context variables . In ICML , volume 97 , pages 5331–5340 , 2019.458 [ 44 ] N. Rakicevic , A. Cully , and P. Kormushev . Policy manifold search : Exploring the manifold hypothesis for459 diversity-based neuroevolution . CoRR , abs/2104.13424 , 2021.460 [ 45 ] T. Schaul , D. Horgan , K. Gregor , and D. Silver . Universal value function approximators . In ICML,461 volume 37 , pages 1312–1320 , 2015.462 [ 46 ] B. Scherrer , M. Ghavamzadeh , V. Gabillon , B. Lesner , and M. Geist . Approximate modified policy iteration463 and its application to the game of tetris . J. Mach . Learn . Res. , 16:1629–1676 , 2015.464 [ 47 ] J . S Schreck , C. W Coley , and K. JM Bishop . Learning retrosynthetic planning through simulated465 experience . ACS central science , 5 ( 6 ) :970–981 , 2019.466 [ 48 ] J. Schulman , S. Levine , P. Abbeel , M. I. Jordan , and P. Moritz . Trust region policy optimization . In ICML,467 pages 1889–1897 , 2015.468 [ 49 ] J. Schulman , P. Moritz , S. Levine , M. I. Jordan , and P. Abbeel . High-dimensional continuous control using469 generalized advantage estimation . In ICLR , 2016.470 [ 50 ] J. Schulman , F. Wolski , P. Dhariwal , A. Radford , and O. Klimov . Proximal policy optimization algorithms.471 CoRR , abs/1707.06347 , 2017.472 [ 51 ] M. Schwarzer , A. Anand , R. Goel , R. D. Hjelm , A. C. Courville , and P. Bachman . Data-efficient473 reinforcement learning with momentum predictive representations . CoRR , abs/2007.05929 , 2020.474 [ 52 ] D. Silver , G. Lever , N. Heess , T. Degris , D. Wierstra , and M. A. Riedmiller . Deterministic policy gradient475 algorithms . In ICML , pages 387–395 , 2014.476 [ 53 ] D. Silver , A. Huang , C. J. Maddison , A. Guez , L. Sifre , G. van den Driessche , J. Schrittwieser,477 I. Antonoglou , V. Panneershelvam , M. Lanctot , S. Dieleman , D. Grewe , J. Nham , N. Kalchbrenner,478 I. Sutskever , T. P. Lillicrap , M. Leach , K. Kavukcuoglu , T. Graepel , and D. Hassabis . Mastering the game479 of go with deep neural networks and tree search . Nature , 529 ( 7587 ) :484–489 , 2016.480 [ 54 ] A. Srinivas , M. Laskin , and P. Abbeel . CURL : contrastive unsupervised representations for reinforcement481 learning . CoRR , abs/2004.04136 , 2020.482 [ 55 ] R. S. Sutton and A. G. Barto . Reinforcement learning - an introduction . Adaptive computation and machine483 learning . MIT Press , 1998.484 [ 56 ] R. S. Sutton , J. Modayil , M. Delp , T. Degris , P. M. Pilarski , A . White , and D. Precup . Horde : a scalable485 real-time architecture for learning knowledge from unsupervised sensorimotor interaction . In AAMAS,486 pages 761–768 , 2011.487 [ 57 ] A. Tacchetti , H. F. Song , P. A. M. Mediano , V. Flores Zambaldi , J. Kramár , N. C. Rabinowitz , Th . Graepel,488 M. Botvinick , and P. W. Battaglia . Relational forward models for multi-agent learning . In ICLR , 2019.489 [ 58 ] E. Todorov , T. Erez , and Y. Tassa . Mujoco : A physics engine for model-based control . 2012 IEEE/RSJ490 International Conference on Intelligent Robots and Systems , pages 5026–5033 , 2012.491 [ 59 ] Y. H. Tsai , Y. Wu , R. Salakhutdinov , and L. Morency . Demystifying self-supervised learning : An492 information-theoretical framework . CoRR , abs/2006.05576 , 2020.493 [ 60 ] T. Unterthiner , D. Keysers , S. Gelly , O. Bousquet , and I. O. Tolstikhin . Predicting neural network accuracy494 from weights . CoRR , abs/2002.11448 , 2020.495 [ 61 ] H. v. Hasselt . Double q-learning . In J. D. Lafferty , C. K. I. Williams , J. Shawe-Taylor , R. S. Zemel , and496 A. Culotta , editors , NeurIPS , pages 2613–2621 , 2010.497 [ 62 ] O. Vinyals , I. Babuschkin , W. M. Czarnecki , M. Mathieu , A. Dudzik , J. Chung , D. H. Choi , R. Powell,498 T. Ewalds , P. Georgiev , J. Oh , D. Horgan , M. Kroiss , I. Danihelka , A. Huang , L. Sifre , T. Cai , J. P. Agapiou,499 M. Jaderberg , A. S. Vezhnevets , R. Leblond , T. Pohlen , V. Dalibard , D. Budden , Y. Sulsky , J. Molloy,500 T. L. Paine , C. Gulcehre , Z. Wang , T. Pfaff , Y. Wu , R. Ring , D. Yogatama , D. Wünsch , K. McKinney,501 O. Smith , T. Schaul , T. Lillicrap , K. Kavukcuoglu , D. Hassabis , C. Apps , and D. Silver . Grandmaster level502 in starcraft ii using multi-agent reinforcement learning . Nature , 575 ( 7782 ) :350–354 , 2019.503 [ 63 ] H. Wang , N. S. Keskar , C. Xiong , and R. Socher . Identifying generalization properties in neural networks.504 CoRR , abs/1809.07402 , 2018.505 [ 64 ] R. Wang , R. Yu , B . An , and Z. Rabinovich . I2hrl : Interactive influence-based hierarchical reinforcement506 learning . In IJCAI , pages 3131–3138 , 2020.507 [ 65 ] J . You , B. Liu , Z. Ying , V. S. Pande , and J. Leskovec . Graph convolutional policy network for goal-directed508 molecular graph generation . In NeurIPS 2018 , pages 6412–6422 , 2018.509 Checklist510 1 . For all authors ... 511 ( a ) Do the main claims made in the abstract and introduction accurately reflect the paper ’ s contribu-512 tions and scope ? [ Yes ] 513 ( b ) Did you describe the limitations of your work ? [ Yes ] See the future work in Sec . 6.514 ( c ) Did you discuss any potential negative societal impacts of your work ? [ No ] Our work is on515 general Reinforcement Learning study . No specific practical application is considered.516 ( d ) Have you read the ethics review guidelines and ensured that your paper conforms to them ? [ Yes ] 517 2 . If you are including theoretical results ... 518 ( a ) Did you state the full set of assumptions of all theoretical results ? [ Yes ] 519 ( b ) Did you include complete proofs of all theoretical results ? [ Yes ] 520 3 . If you ran experiments ... 521 ( a ) Did you include the code , data , and instructions needed to reproduce the main experimental522 results ( either in the supplemental material or as a URL ) ? [ No ] Our experimental environment523 are public and standard . All the information needed to reproduce our results is provided in the524 main body and appendix . Code will be available publicly soon.525 ( b ) Did you specify all the training details ( e.g. , data splits , hyperparameters , how they were chosen ) ? 526 [ Yes ] Partially in main body and all details can be found in the appendix document.527 ( c ) Did you report error bars ( e.g. , with respect to the random seed after running experiments528 multiple times ) ? [ Yes ] 529 ( d ) Did you include the total amount of compute and the type of resources used ( e.g. , type of GPUs,530 internal cluster , or cloud provider ) ? [ Yes ] 531 4 . If you are using existing assets ( e.g. , code , data , models ) or curating/releasing new assets ... 532 ( a ) If your work uses existing assets , did you cite the creators ? [ Yes ] 533 ( b ) Did you mention the license of the assets ? [ Yes ] We use a free education licence for students for534 MuJoCo.535 ( c ) Did you include any new assets either in the supplemental material or as a URL ? [ No ] 536 ( d ) Did you discuss whether and how consent was obtained from people whose data you ’ re us-537 ing/curating ? [ N/A ] 538 ( e ) Did you discuss whether the data you are using/curating contains personally identifiable informa-539 tion or offensive content ? [ N/A ] 540 5 . If you used crowdsourcing or conducted research with human subjects ... 541 ( a ) Did you include the full text of instructions given to participants and screenshots , if applicable ? 542 [ N/A ] 543 ( b ) Did you describe any potential participant risks , with links to Institutional Review Board ( IRB ) 544 approvals , if applicable ? [ N/A ] 545 ( c ) Did you include the estimated hourly wage paid to participants and the total amount spent on546 participant compensation ? [ N/A ] 547

Summary: Update: See the comments. Review score remains at 7.  The paper considers learning value function approximators that take a policy representation as an input to the value function to allow generalization across different policies. There are several other recent works of this type. The aspect differentiating the current work is the type of policy representation used that allowed the method to scale to policies with 10^5 parameters.  The paper considered two types of data sources for the policy representations: 1. State action pairs obtained by the policy. 2. The policy parameters.  Then, the policy representation is formed by passing the data through a multilayer neural network that embeds the data into a lower dimensional representation. The embedding is fed into the value function as an input, and the network creating the embedding is trained end-to-end to minimize the value function prediction error. Optionally they also added auxiliary losses such as contrastive learning.  Experiments were performed using PPO combined with their new value function and compared to the standard PPO as well as variants taking a random policy representation and also a policy representation with no emebedding. They performed experiments on 6 MuJoCo tasks, and generally the new methods outperformed the baselines.  tSNE plots showed that the policies with similar representations had a similar performance at the task (when comparing the achieved rewards). Moreover, experiments showed that the value prediction error was smaller for the new method compared to a standard value function method, and also that the value function was accurate for a test set of new policies not encountered during training.