Proper Value Equivalence

1 Introduction . It has long been argued that , in order for reinforcement learning ( RL ) agents to solve truly complex tasks , they must build a model of the environment that allows for counterfactual reasoning [ 29 ] . Since representing the world in all its complexity is a hopeless endeavor , especially under capacity constraints , the agent must be able to ignore aspects of the environment that are irrelevant for its purposes . This is the premise behind the value equivalence ( VE ) principle , which provides a formalism for focusing on the aspects of the environment that are crucial for value-based planning [ 17 ] . VE distinguishes models based on a set of policies and a set of real-valued scalar functions of state ( henceforth , just functions ) . Roughly , a model is said to be VE to the environment if the Bellman operators it induces for the policies yield the same result as the environment ’ s Bellman operators when applied to the functions . The policies and functions thus become a “ language ” to specify which parts of the environment a model should capture . As the number of policies and functions increase the requirements on the model become more stringent , which is to say that the class of VE models shrinks . In the limit , the VE class collapses to a single point corresponding to a perfect model . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Although this result is reassuring , in practice we want to stop short of collapsing—after all , at this point the agent is no longer ignoring irrelevant aspects of the environment . A fundamental question is thus how to select the smallest sets of policies and functions such that a resulting VE model is sufficient for planning . In this paper we take an important additional step in this direction : we show that the VE principle can be formulated with respect to value functions only . This result drastically reduces the space of functions that must be considered by VE , as in general only a small fraction of the set of all functions will qualify as value functions in a given environment . Since every policy has an associated value function , this new formulation of VE removes the need for selecting functions , only requiring policies . We name our new formulation proper value equivalence ( PVE ) to emphasize its explicit use of value functions . PVE has several desirable properties . Unlike with VE , the class of PVE models does not collapse to a singleton in the limit . This means that , even if all value functions are used , we generally end up with multiple PVE models—which can be beneficial if some of these are easier to learn or represent than others . Crucially , all of these models are sufficient for planning , meaning that they will yield an optimal policy despite the fact that they may ignore many aspects of the environment . Finally , we make more precise Grimm et al. ’ s [ 17 ] suggestion that the VE principle may help explain the good empirical performance of several modern algorithms [ 38 , 33 , 24 , 12 , 30 ] . Specifically , we show that , with mild assumptions , minimizing the loss of the MuZero algorithm [ 31 ] can be understood as minimizing a PVE error . We then leverage this connection to suggest a modification to MuZero and show a small but significant improvement in the Atari Learning Environment [ 3 ] . 2 Background . The agent ’ s interaction with the environment will be modeled as a Markov decision process ( MDP ) M≡ 〈S , A , r , p , γ〉 , where S and A are the state and action spaces , r ( s , a ) is the expected reward following taking a from s , p ( s′|s , a ) is the transition kernel and γ ∈ [ 0 , 1 ) is a discount factor [ 27 ] . A policy is a mapping π : S 7→ P ( A ) , where P ( A ) is the space of probability distributions over A ; we define ≡ { π |π : S 7→ P ( A ) } as the set of all possible policies . A policy π is deterministic if π ( a|s ) > 0 for only one action a per state s. A policy ’ s value function is defined as vπ ( s ) ≡ Eπ [ ∞∑ i=0 γir ( St+i , At+i ) |St = s ] , ( 1 ) where Eπ [ · ] denotes expectation over the trajectories induced by π and the random variables St and At indicate the state occupied and the action selected by the agent at time step t. The agent ’ s goal is to find a policy π ∈ that maximizes the value of every state [ 36 , 37 ] . Usually , a crucial step to carry out this search is to compute the value function of candidate policies . This process can be cast in terms of the policy ’ s Bellman operator : Tπ [ v ] ( s ) ≡ EA∼π ( ·|s ) , S′∼p ( ·|s , A ) [ r ( s , A ) + γv ( S′ ) ] , ( 2 ) where v is any function in the space V ≡ { f | f : S 7→ R } . It is known that limn→∞ T nπ v = vπ , that is , starting from any v ∈ V , the repeated application of Tπ will eventually converge to vπ . Since in RL the agent does not know p and r , it can not apply ( 2 ) directly . One solution is to learn a model m̃ ≡ ( r̃ , p̃ ) and use it to compute ( 2 ) with p and r replaced by p̃ and r̃ [ 36 ] . We denote the set of all models as M. The value equivalence principle defines a model as value equivalent ( VE ) to the environment m∗ ≡ ( r , p ) with respect to a set of policies Π and a set of functions V if it produces the same Bellman updates as m∗ when using Π and V [ 17 ] . Classes of such models are expressed as follows : M ( Π , V ) ≡ { m̃ ∈M : T̃πv = Tπv ∀π ∈ Π , v ∈ V } ( 3 ) whereM⊆ M is a class of models , T̃π denotes one application of the Bellman operator induced by model m̃ and policy π to function v , and Tπ is environment ’ s Bellman operator for π. Grimm et al . [ 17 ] showed that the VE principle can be used to learn models that disregard aspects of the environment which are not related to the task of interest.1 Classical approaches to model learning 1A related approach is taken in value-aware model learning [ 11 ] which minimizes the discrepancy between the Bellman optimality operators induced by the model and the environment . do not take the eventual use of the model into account , potentially modeling irrelevant aspects of the environment . Accordingly , Grimm et al . [ 17 ] have shown that , under the same capacity constraints , models learned using VE can outperform their classical counterparts . 3 Proper value equivalence . One can define a spectrum of VE classes corresponding to different numbers of applications of the Bellman operator . We define an order-k VE class as : Mk ( Π , V ) ≡ { m̃ ∈M : T̃ kπ v = T kπ v ∀π ∈ Π , v ∈ V } ( 4 ) where T̃ kπ v denotes k applications of T̃π to v. Under our generalized definition of VE , Grimm et al . [ 17 ] studied order-one VE classes of the formM1 ( Π , V ) . They have shown thatM1 ( , V ) either contains only the environment or is empty . This is not generally true for k > 1 . The limiting behavior of order-k value equivalent classes can be described as follows Proposition 1 . Let V be a set of functions such that if v ∈ V then Tπv ∈ V for all π ∈ Π . Then , for k , K ∈ Z+ such that k divides K , it follows that : ( i ) For anyM⊆ M and any Π ⊆ , we have thatMk ( Π , V ) ⊆MK ( Π , V ) . ( ii ) If Π is non-empty and V contains at least one constant function , then there exist environments such that Mk ( Π , V ) ⊂ MK ( Π , V ) . We defer all proofs of theoretical results to Appendix A.2 . Based on Proposition 1 we can relate different VE model classes according to the greatest common divisor of their respective orders ; specifically , two classesMk ( Π , V ) andMK ( Π , V ) will intersect atMgcd ( k , K ) ( Π , V ) ( Figure 1 ) . Proposition 1 also implies that , in contrast to order-one VE classes , higher order VE classes potentially include multiple models , even if VE is defined with respect to all policies and all functions V. In addition , the size of a VE class can not decrease as we increase its order from k to a multiple of k ( and in some cases it will strictly increase ) . This invites the question of what happens in the limit as we keep increasing the VE order . To answer this question , we introduce a crucial concept for this paper : Definition 1 . ( Proper value equivalence ) . Given a set of policies Π ⊆ , let M∞ ( Π ) = lim k→∞ Mk ( Π , V ) = { m̃ ∈M : ṽπ = vπ ∀π ∈ Π } , ( 5 ) where ṽπ and vπ are the value functions of π induced by model m̃ and the environment . We say that each m̃ ∈M∞ ( Π ) is a proper value equivalent model to the environment with respect to Π . Because the process of repeatedly applying a policy ’ s Bellman operator to a function converges to the same fixed point regardless of the function , in an order-∞ VE class the set Π uniquely determines the set V . This reduces the problem of defining Π and V to defining the former only . Also , since all functions in an order-∞ VE are value functions , we call it proper VE or PVE . It is easy to show that Proposition 1 is valid for any k ∈ Z+ when K = ∞ ( Corollary 2 in Appendix A.2 ) . Thus , in some sense , M∞ is the “ biggest ” VE class . It is also possible to define this special VE class in terms of any other : Proposition 2 . For any Π ⊆ and any k ∈ Z+ it follows that M∞ ( Π ) = ⋂ π∈Π Mk ( { π } , { vπ } ) , ( 6 ) where vπ is the value of policy π in the environment . We thus have two equivalent ways to describe the class of models which are PVE with respect to a set of policies Π . The first , given in ( 5 ) , is the order-∞ limit of value equivalence with respect to Π and the set of all functions V. The second , given in ( 6 ) , is the intersection of the classes of models that are order-k VE with respect to the singleton policies in Π and their respective value functions . This latter form is valid for any k , and will underpin our practical algorithmic instantiations of PVE . Setting k = 1 in Proposition 2 we see that PVE can be written in terms of order-one VE . This means thatM∞ inherits many of the topological properties ofM1 shown by Grimm et al . [ 17 ] . Specifically , we know thatM′∞ ( Π ) ⊆M∞ ( Π ) ifM′ ⊆M and also thatM∞ ( Π′ ) ⊆M∞ ( Π ) when Π ⊆ Π′ ( these directly follow from Grimm et al. ’ s [ 17 ] Properties 1 and 3 respectively ) . Proposition 2 also sheds further light into the relation between PVE and order-k VE more generally . Let Π be a set of policies and Vπ their value functions . Then , for any k ∈ Z+ , we have that Mk ( Π , Vπ ) = ⋂ π∈Π ⋂ v∈Vπ Mk ( { π } , { v } ) ⊆ ⋂ π∈Π Mk ( { π } , { vπ } ) =M∞ ( Π ) , ( 7 ) which is another way to say thatM∞ is , in some sense , the largest among all the VE classes . The reason why the size of VE classes is important is that it directly reflects the main motivation behind the VE principle . VE ’ s premise is that models should be constructed taking into account their eventual use : if some aspects of the environment are irrelevant for value-based planning , it should not matter whether a model captures them or not . This means that all models that only differ with respect to these irrelevant aspects but are otherwise correct qualify as valid VE solutions . A larger VE class generally means that more irrelevant aspects of the environment are being ignored by the agent . We now make this intuition more concrete by showing how irrelevant aspects of the environment that are eventually captured by order-one VE are always ignored by PVE : Proposition 3 . Let Π ⊆ . If the environment state can be factored as S = X × Y where |Y| > 1 and vπ ( s ) = vπ ( ( x , y ) ) = vπ ( x ) for all π ∈ Π , then M1 ( Π , V ) ⊂ M∞ ( Π ) . Note that the subset relation appearing in Proposition 3 is strict . We can think of the variable ‘ y ’ appearing in Proposition 3 as superfluous features that do not influence the RL task , like the background of an image or any other sensory data that is irrelevant to the agent ’ s goal . A model is free to assign arbitrary dynamics to such irrelevant aspects of the state without affecting planning performance . Since order-one VE eventually pins down a model that describes everything about the environment , one would expect the size ofM∞ relative toM1 to increase as more superfluous features are added . Indeed , in our proof of Proposition 3 we construct a set of models inM∞ ( ) which are in one-to-one correspondence with Y , confirming this intuition ( see Appendix A.2 ) . Proper value equivalence yields models that are sufficient for optimal planning In general PVE does not collapse to a single model even in the limit of Π = . At first this may cause the impression that one is left with the extra burden of selecting one among the PVE models . However , it can be shown that no such choice needs to be made : Proposition 4 . An optimal policy for any m̃ ∈M∞ ( ) is also an optimal policy in the environment . According to Proposition 1 any model m̃ ∈M∞ ( ) used for planning will yield an optimal policy for the environment . In fact , in the spirit of ignoring as many aspects of the environment as possible , we can define an even larger PVE class by focusing on deterministic policies only : Corollary 1 . Let det be the set of all deterministic policies . An optimal policy for any m̃ ∈ M∞ ( det ) is also optimal in the environment . Given that bothM∞ ( ) andM∞ ( det ) are sufficient for optimal planning , one may wonder if these classes are in fact the same . The following result states that the class of PVE models with respect to deterministic policies can be strictly larger than its counterpart defined with respect to all policies : Proposition 5 . There exist environments and model classes for whichM∞ ( ) ⊂M∞ ( det ) . Figure 2 illustrates Proposition 5 with an example of environment and a model m̃ such that m̃ ∈ M∞ ( det ) but m̃ /∈M∞ ( ) . To conclude our discussion on models that are sufficient for optimal planning , we argue that , in the absence of additional information about the environment or the agent , M∞ ( det ) is in fact the largest possible VE class that is guaranteed to yield optimal performance . To see why this is so , suppose we remove a single deterministic policy from det and pick an arbitrary model m̃ ∈M∞ ( det − { π } ) . Let ṽπ be the value function of π according to the model m̃ . Because π is not included in the set of policies used to enforce PVE , ṽπ may not coincide with vπ , the actual value function of π according to the environment . Now , if π happens to be the only optimal policy in the environment and ṽπ is not the optimal value function of m̃ , the policy returned by this model will clearly be sub-optimal .

Summary: The paper proposes an extension to the value equivalence (VE) principle in MBRL called Proper VE (PVE). The original VE states that a model is equivalent to the true model if the induced Bellman operators coincide for a set of policies and value functions. PVE is motivated by the observation that a repeated application of the Bellman operator will eventually converge to a value function of a given policy for a chosen model. PVE defines a model to be equivalent to the true if, for a set of policies, it results in the same value functions as in the original MDP. While PVE assumes an infinite Bellman operator application (called order-∞), the paper studies the relation between non-asymptotic versions of PVE.  For example, the authors prove that order-k PVE models are also order-K PVE models if k divides K and demonstrate the result empirically. The paper shows that PVE models are sufficient for planning, i.e. a policy that is optimal for a PVE model is optimal in the true MDP. Finally, the authors propose a PVE-inspired loss for model learning and show that MuZero optimizes an upper bound on the loss corresponding to a behavior policy. Inspired by the observation, the authors extend to the MuZero loss with additional terms corresponding to previous policies and show an increase in the median score on the Atari 57 benchmark.