Proper Value Equivalence

Abstract:
One of the main challenges in model-based reinforcement learning (RL) is to decide which aspects of the environment should be modeled. The value-equivalence (VE) principle proposes a simple answer to this question: a model should capture the aspects of the environment that are relevant for value-based planning. Technically, VE distinguishes models based on a set of policies and a set of functions: a model is said to be VE to the environment if the Bellman operators it induces for the policies yield the correct result when applied to the functions. As the number of policies and functions increase, the set of VE models shrinks, eventually collapsing to a single point corresponding to a perfect model. A fundamental question underlying the VE principle is thus how to select the smallest sets of policies and functions that are sufficient for planning. In this paper we take an important step towards answering this question. We start by generalizing the concept of VE to order-k counterparts defined with respect to k applications of the Bellman operator. This leads to a family of VE classes that increase in size as k → ∞. In the limit, all functions become value functions, and we have a special instantiation of VE which we call proper VE or simply PVE. Unlike VE, the PVE class may contain multiple models even in the limit when all value functions are used. Crucially, all 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. We construct a loss function for learning PVE models and argue that popular algorithms such as MuZero can be understood as minimizing an upper bound for this loss. We leverage this connection to propose a modification to MuZero and show that it can lead to improved performance in practice.

heading: 1 Introduction
text: 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].

heading: 2 Background
text: 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 cannot 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.

heading: 3 Proper value equivalence
text: 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 cannot 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.

heading: 4 Learning a proper value-equivalent model
text: Having established that we want to find a model inM∞( det), we now turn our attention to how this can be done in practice. Following Grimm et al. [17], given a finite set of policies Π and a finite set of functions V , we cast the search for a model m̃ ∈Mk(Π,V) as the minimization of deviations from (4):
`kΠ,V(m ∗, m̃) ≡ ∑ π∈Π ∑ v∈V ‖T kπ v − T̃ kπ v||, (8)
where T̃π are Bellman operators induced by m̃ and ‖·‖ is a norm.2 Note that setting k = ∞ in (8) yields a loss that requires computing m̃’s value function—which is impractical to do if m̃ is being repeatedly updated. Thankfully, by leveraging the connection between order-k VE and PVE given in Proposition 2, we can derive a practical PVE loss:
`kΠ,∞(m ∗, m̃) ≡ ∑ π∈Π ‖T kπ vπ − T̃ kπ vπ|| = ∑ π∈Π ‖vπ − T̃ kπ vπ||. (9)
Interestingly, given a set of policies Π, minimizing (9) for any k will result in a model m̃ ∈M∞(Π) (cf. Proposition 2). As we will discuss shortly, this property can be exploited to generate multiple loss functions that provide a richer learning signal in practical scenarios.
Contrasting loss functions (8) and (9) we observe an important fact: unlike with other order-k VE classes, PVE requires actual value functions to be enforced in practice. Since value functions require data and compute to be obtained, it is reasonable to ask whether the benefits of PVE justify the associated additional burden. Concretely, one may ask whether the sample transitions and computational effort spent in computing the value functions to be used with PVE would not be better invested in enforcing other forms of VE over arbitrary functions that can be readily obtained.
We argue that in many cases one does not have to choose between order-k VE and PVE. Valuebased RL algorithms usually compute value functions iteratively, generating a sequence of functions v1, v2, ... which will eventually converge to ṽπ for some π. A model-based algorithm that computes ṽπ in this way has to somehow interleave this process with the refinement of the model m̃. When it comes to VE, one extreme solution is to only use the final approximation ṽπ ≈ vπ in an attempt to enforce PVE through (9). It turns out that, as long as the sequence v1, v2, ... is approaching vπ, one can use all the functions in the sequence to enforce PVE with respect to π. Our argument is based on the following result:
2We can also impose VE with infinite sets of functions and policies by replacing the respective sums with integrals; in this case one may consider taking a supremum over VE terms to avoid situations where VE is not necessarily satisfied on measure 0 sets.
Proposition 6. For any π ∈ , v ∈ V and k, n ∈ Z+, we have that
‖vπ − T̃ kπ vπ‖∞ ≤ (γk + γn) ‖vπ − v‖∞︸ ︷︷ ︸ v + ‖T nπ v − T̃ kπ v‖∞︸ ︷︷ ︸ ve . (10)
Note that the left-hand side of (10) corresponds to one of the terms of the PVE loss (9) associated with a given π. This means that, instead of minimizing this quantity directly, one can minimize the upper-bound on the right-hand side of (10). The first term in this upper bound, v , is the conventional value-function approximation error that most value-based methods aim to minimize (either directly or indirectly). The second term, ve, is similar to the terms appearing in the order-k VE loss (8), except that here the number of applications of Tπ and of its approximation T̃π do not have to coincide. All the quantities appearing in ve are readily available or can be easily approximated using sample transitions [17]. Thus, ve can be used to refine the model m̃ using functions v that are not necessarily value functions. As v → vπ, two things happen. First, ve approaches one of the terms of the PVE loss (9) associated with policy π. Second, v vanishes. Interestingly, the importance of v also decreases with n and k, the number of times Tπ and T̃π are applied in ve, respectively. This makes sense: since T nπ v → vπ as n→∞ and, by definition, VE approaches PVE as k →∞, we have that ve approaches the left-hand side of (10) as both n and k grow.

heading: An extended example: MuZero through the lens of value equivalence
text: Grimm et al. [17] suggested that the VE principle might help to explain the empirical success of recent RL algorithms like Value Iteration Networks, the Predictron, Value Prediction Networks, TreeQN, and MuZero [38, 33, 24, 12, 31]. In this section we investigate this hypothesis further and describe a possible way to interpret one of these algorithms, MuZero, through the lens of VE. We acknowledge that the derivation that follows abstracts away many details of MuZero and involves a few approximations of its mechanics, but we believe it captures and explains the algorithm’s essence.
MuZero is a model-based RL algorithm that achieved state-of-the-art performance across both board games, such as Chess and Go, and Atari 2600 games [31]. The model m̃ in MuZero is trained on sequences of states, actions and rewards resulting from executing a “behavior policy” in the environment: st:t+n+K , at:t+n+K , rt:t+n+K where n and K are hyperparameters of the agent which will be explained shortly. The agent produces an “agent state” z0t from st and subsequently generates z1:Kt by using its model to predict the next K agent states following actions at:t+K . The agent also maintains reward and value function estimates as a function of agent states, which we denote r̃(z) and v(z) respectively. A variant3 of MuZero’s per-state model loss can thus be expressed as:
`µ(st) = K∑ k=0 (Vt+k − v(zkt ))2 + (rt+k − r̃(zkt )))2 (11)
where Vt+k = rt+k + · · · + γn−1rt+k+n−1 + γnvtarg(z0t+k+n). The term vtarg is a value target produced by Monte-Carlo tree search (MCTS, [7]). Because the behavior policy is itself computed via MCTS, we have that vtarg ≈ v; for simplicity we will assume that vtarg = v and only use v. In what follows we show, subject to a modest smoothness assumption, that minimizing MuZero’s loss with respect to its behavior policy, π, also minimizes a corresponding PVE loss. Put precisely:
C · Edπ [`µ(St)] ≥ ( `K{π},∞(m ∗, m̃) )2
(12)
for some C > 0, where dπ is a stationary distribution. We proceed by combining two derivations: a lower-bound on Edπ [`µ(St)] in (15), and an upper-bound on (`K{π},∞(m ∗, m̃))2 in (17).
As a preliminary step we note that `µ(st) and `K{π},∞(m ∗, m̃) are expressed in terms of samples and expectations respectively. We note the following connection between these quantities:
E[rt+k|st] = Pkπ [rπ](st), E[Vt+k|st] = PkπT nπ [vπ](st), E[r̃(zkt )|st] = P̃kπ [r̃π](st), E[v(zkt )|st] = P̃kπ [vπ](st),
(13)
3In reality MuZero uses a categorical representation for its value and reward functions and minimizes them using a cross-entropy objective. We argue that this choice is not essential to its underlying ideas and use scalar representations with a squared loss to simplify our analysis.
where Pkπ is the k-step environment transition operator under policy π: Pkπ [x](st) = E[x(St+k)|st,m∗, π], rπ(s) = EA∼π[r(s,A)] and P̃kπ and r̃π are the corresponding quantities using the model instead of the environment. The above expectations are taken with respect to the environment or model and π as appropriate. We now derive our lower-bounds on Edπ [`µ(St)]:
Edπ [`µ(St)] = Edπ [ K∑ k=0 E[ (Vt+k − v(zkt ))2 | St ] + K∑ k=0 E[ (rt+k − r̃(zkt ))2 | St ] ]
≥ Edπ [ K∑ k=0 (E[Vt+k | St ]− E[ v(zkt ) | St ])2 + K∑ k=0 (E[ rt+k | St ]− E[ r̃(zkt ) | St ])2 ]
= K∑ k=0 Edπ [ (PkπT nπ v(St)− P̃kπv(St))2 ] + K∑ k=0 Edπ [ (Pkπrπ(St)− P̃kπ r̃π(St))2 ] (14)
where we apply the tower-property, Jensen’s inequality and the identities in (13). We write the expression using norms and drop all terms except k ∈ {0,K} in the first sum to obtain:
Edπ [`µ(St)] ≥ ‖T nπ v − v‖2dπ + ‖P K π T nπ v − P̃Kπ v‖2dπ + K∑ k=0 ‖Pkπrπ − P̃kπ r̃π‖2dπ (15)
recalling that ‖x− y‖2dπ = Edπ [(x(St)− y(St)) 2]. To derive an upper-bound for (`K{π},∞(m ∗, m̃))2 we assume that the error in value estimation is smooth in the sense that there is some g > 0 (independent of v) such that ‖v − vπ‖∞ < g · ‖v − vπ‖dπ . We can then use a modified version of (10) for the dπ-weighted `2-norm (see Appendix A.2), plugging in n+K and K:
‖vπ − T̃ Kπ vπ‖dπ ≤ (gγK + γn+K)‖vπ − v‖dπ + ‖T K+nπ v − T̃ Kπ v‖dπ
≤ γK(g + γn)‖vπ − v‖dπ + ‖PKπ T nπ v − P̃Kπ v‖dπ + K∑ k=0 ‖Pkπrπ − P̃kπ r̃π‖dπ
≤ γK (g + γ n)
(1− γn) ‖T nπ v − v‖dπ + ‖PKπ T nπ v − P̃Kπ v‖dπ + K∑ k=0 ‖Pkπrπ − P̃kπ r̃π‖dπ ,
(16) from here we can square both sides and apply Jensen’s inequality,
‖vπ − T̃ Kπ vπ‖2dπ ≤ ab‖T π n v − v‖2dπ + b‖P K π T nπ v − P̃Kπ v‖2dπ + b K∑ k=0 ‖Pkπrπ − P̃kπ r̃π‖2dπ , (17)
where a = γK(g + γn)(1− γn)−1 and b = a+K + 2. Combining (17) and (15) we obtain:
ab · Edπ [`µ(St)] ≥ ‖vπ − T̃ Kπ vπ‖2dπ = ( `K{π},∞(m ∗, m̃) )2 , (18)
thus minimizing MuZero’s loss minimizes a squared PVE loss with respect to a single policy.

heading: 5 Experiments
text: We first provide results from tabular experiments on a stochastic version of the Four Rooms domain which serve to corroborate our theoretical claims. Then, we present results from experiments across the full Atari 57 benchmark [3] showcasing that the insights from studying PVE and its relationship to MuZero can provide a benefit in practice at scale. See Appendix A.3 for a full account of our experimental procedure.
In Section 3 we described the topological relationships between order-k and PVE classes. This is summarized by Proposition 1, which shows that, for appropriately defined V and Π,Mk ⊆MK ifK is a multiple of k. We illustrate this property empirically by randomly initializing a set of models and then using (8) (or (9) for the limiting case of k =∞) to iteratively update them towardsMk( ,V), with k ∈ {1, 30, 40, 50, 60,∞}. We take the vectors representing these models and project them onto their first two principal components in order to visualise their paths through learning. The results are
shown on the top row of Figure 3. In accordance with the theory, we see that the space of converged models, represented with the brightest yellow regions, grows with k. This trend is summarised in the rightmost plot, which shows the diameter of the scatter plots for each k. In the bottom row of Figure 3 we use color to show the value that the optimal policy of each model achieves in the true environment. As predicted by our theory, the space of models that are sufficient for optimal planning also grows with k.
Model classes containing many models with optimal planning performance are particularly advantageous when the set of models that an agent can represent is restricted, since the larger the set of suitable models the greater the chance of an overlap between this set and the set of models representable by the agent. Proposition 4 and Corollary 1 comparedM∞( ) andM∞( det), showing that, although any model in either class is sufficient for planning, M∞( ) ⊆ M∞( det). This suggests that it might be better to learn a model inM∞( det) when the agent has limited capacity.
We illustrate that this is indeed the case in Figure 4b. We progressively restrict the space of models that the agent can represent and attempt to learn models in eitherM∞( ) orM∞( det). Indeed, we find that the larger class,M∞( det), yields superior planning performance as agent capacity decreases.
Given their importance, we provide intuition on the ways that individual PVE models differ from the environment. In Figure 4a we compare trajectories starting at the same initial state (denoted by a red-circle) from the optimal environment policy in both the environment and in a randomly sampled model fromM∞( ). In the PVE model there are numerous diagonal transitions not permitted by the environment. Note that while the PVE model has very different dynamics than the environment, these differences must “balance out”, as it still has the same values under any policy as the environment.
In Section 4 we showed that minimizing Muzero’s loss function is analogous to minimizing an upper-bound on a PVE loss (9) with respect to the agent’s current policy π—which corresponds to finding a model inM∞(Π) where Π = {π}. Note that our guarantee on the performance of PVE models (Corollary 1) holds when Π contains all deterministic policies. While it is not feasible to enforce Π = det, we can use previously seen policies by augmenting the MuZero algorithm with a buffer of past policies and their approximate value functions (we do so by periodically storing the corresponding parameters). We can then add an additional loss to MuZero with the form of the original value loss but using the past value functions. We still use sampled rewards to construct value targets, but use the stored policies to compute off-policy corrections using V-trace [9].
To test this proposal we use an on policy (i.e., without a replay buffer) implementation of MuZero run for 500M frames (as opposed to 20B frames in the online result of [31]) on the Atari 57 benchmark and find that using our additional loss yields an advantage in the human normalised median performance shown in Figure 5. MuZero’s data efficiency can also be improved with the aid of a replay buffer of trajectories from past policies [32], which may also capture some of the advantages of expanding the set of policies used for PVE.

heading: 6 Related work
text: Our work is closely related to the value-aware model learning (VAML, IterVAML, [11, 10]) which learns models to minimize the discrepancy between their own Bellman optimality operators and the environment’s on the optimal value function—an a priori unknown quantity. To handle this VAML specifies a family of potential value functions and minimizes the worst-case discrepancy across them, whereas IterVAML minimizes the discrepancy with respect to model’s current estimate of the value function in a value iteration inspired scheme. PVE and the VAML family are complementary works, with VAML addressing its induced optimization problems and PVE addressing its induced model classes. Both, however, advocate for learning models with their eventual use in mind—a view that is aligned with many criticisms of the maximum likelihood objective for model learning [11, 21, 2, 23]).
It is worth mentioning the relationship between PVE and TD models [35] which, for a given policy, defines any R ∈ R|S| and P ∈ R|S|×|S| with limk→∞ P k = 0 as a valid model if V = R + P>V where V ∈ R|S| represents vπ . Clearly all models inM∞({π}) are valid models, however, since P is not necessarily a transition matrix, the converse does not hold. While TD models are restricted to prediction rather than control, their generality warrants further inquiry.
Order-k and PVE model classes form equivalences between MDPs and thus can be situated among other equivalence notions which can be formulated as state-aggregations [8, 26, 25, 16, 28, 13, 34, 39, 5, 40]. As pointed out by Grimm et al. [17], the interaction between arbitrary state-aggregation
and models can be captured with special cases of order-one VE. Our extension of higher-order VEs potentially offers the possibility of “blending” existing notions of state aggregation with PVE.
A notable instance of state-aggregation is bisimulation [22], which uses a relation to aggregate states that have the same immediate rewards and transition dynamics into other aggregated states. Bisimulation metrics [13] provide smooth measures of how closely pairs of states satisfy bisimulation relations. These concepts have become increasingly popular in deep reinforcement learning where they are used to guide the learning of effective representations [43, 42, 15, 1]. Although both bisimulation and PVE provide direction for learning internal aspects of an agent, they are fundamentally different in their purview—bisimulation concerns representation learning, while PVE concerns the learning of models given a representation of state.
Beyond bisimulation, representation learning has a wide literature [41, 20, 6, 14, 4] including several modern works which explicitly study the conjunction of model learning with state representation [44, 42, 15]. These are further complemented by efforts to learn state representations and models jointly in the service of value-based planning [12, 33, 24, 18, 31, 38].

heading: 7 Conclusion and future work
text: We extended the value equivalence principle by defining a spectrum of order-k VE sets in which models induce the same k-step Bellman operators as the environment. We then explored the topology of the resulting equivalence classes and defined the limiting class when k →∞ as PVE. If a model is PVE to the environment with respect to a set of policies Π, then the value functions of all policies in Π are the same in the environment and the model. The fact that PVE classes can be defined using only a set of policies eliminates the need for specifying a set of functions to induce VE—resolving a fundamental issue left open by Grimm et al. [17]. Importantly, we showed that being PVE with respect to all deterministic policies is sufficient for a model to plan optimally in the environment. In the absence of additional information, this is the largest possible VE class that yields optimal planning. On the practical side, we showed how the MuZero algorithm can be understood as minimizing an upper bound on a PVE loss, and leveraged this insight to improve the algorithm’s performance.
Though our efforts have advanced the understanding of value equivalence and proven useful algorithmically, there is still work to be done in developing a VE theory whose assumptions hold in practice. This remaining work can be broadly grouped into two areas (1) understanding the role of approximation in VE and (2) establishing performance guarantees for VE models with arbitrary sets of policies and functions. We leave these as future work.

heading: Acknowledgements
text: We thank Angelos Filos and Sonya Kotov for many thought-provoking discussions. Christopher Grimm’s work was made possible by the support of the Lifelong Learning Machines (L2M) grant from the Defense Advanced Research Projects Agency. Any opinions, findings, conclusions, or recommendations expressed here are those of the authors and do not necessarily reflect the views of the sponsors.

Review summaries:
The paper "Proper Value Equivalence" addresses the question of how models for reinforcement learning can be obtained which are sufficient for planning without modelling irrelevant details of the environments. The authors build upon previous work by Grimm et al. and expand the notion of value equivalence to only require value functions actually obtainable by policies in the environment. They show several illustrative examples and proofs of why this notion of proper value equivalence alleviates some problems with previous work, such as not collapsing the space of admissible models down to the optimal model. This is very important in many cases where the true model might not be recoverable due to irreducible approximation errors or missing data.

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. 

This paper introduces a k-th order generalization of the Value Equivalence (VE) principle of Grimm et. al. 2016, generalizing that the candidate value functions are equal under multiple (rather than one) iterations of the Bellman operator for a candidate set of policies. Taking k -> infinity leads to the special case of Proper Value Equivalence (PVE), which only depends on a candidate set of policies. Several structural results and derived, and it is shown that models obtained from PVE (with the set of deterministic policies) have the same optimal policies as the true environment.  A loss function is introduced to learn k-VE models, connections with MuZero are drawn, and experimental simulations are carried out. 

This paper first extends value equivalency (VE) to k-step value equivalency, where the equivalency holds for k-step application of the Bellman operator. In the limit of k to infinity they then show this becomes ‘proper value equivalence’ (PVE), for which the policy uniquely defines the VE class. They also show that MuZero is a special case of the PVE class, and derive empirical experiments on PVE, also extending MuZero with a new loss.