MICo: Improved representations via sampling-based state similarity for Markov decision processes
Year: 2022

Abstract:
We present a new behavioural distance over the state space of a Markov decision process, and demonstrate the use of this distance as an effective means of shaping the learnt representations of deep reinforcement learning agents. While existing notions of state similarity are typically difficult to learn at scale due to high computational cost and lack of sample-based algorithms, our newly-proposed distance addresses both of these issues. In addition to providing detailed theoretical analysis, we provide empirical evidence that learning this distance alongside the value function yields structured and informative representations, including strong results on the Arcade Learning Environment benchmark.

heading: 1 Introduction
text: The success of reinforcement learning (RL) algorithms in large-scale, complex tasks depends on forming useful representations of the environment with which the algorithms interact. Feature selection and feature learning has long been an important subdomain of RL, and with the advent of deep reinforcement learning there has been much recent interest in understanding and improving the representations learnt by RL agents.
Much of the work in representation learning has taken place from the perspective of auxiliary tasks [Jaderberg et al., 2017, Bellemare et al., 2017, Fedus et al., 2019]; in addition to the primary reinforcement learning task, the agent may attempt to predict and control additional aspects of the environment. Auxiliary tasks shape the agent’s representation of the environment implicitly, typically via gradient descent on the additional learning objectives. As such, while auxiliary tasks continue to play an important role in improving the performance of deep RL algorithms, our understanding of the effects of auxiliary tasks on representations in RL is still in its infancy.
In contrast to the implicit representation shaping of auxiliary tasks, a separate line of work on behavioural metrics, such as bisimulation metrics [Desharnais et al., 1999, 2004, Ferns et al., 2004, 2006], aims to capture structure in the environment by learning a metric measuring behavioral similarity between states. Recent works have successfully used behavioural metrics to shape the representations of deep RL agents [Gelada et al., 2019, Zhang et al., 2021, Agarwal et al., 2021a]. However, in practice behavioural metrics are difficult to estimate from both statistical and computational perspectives, and these works either rely on specific assumptions about transition dynamics to make the estimation tractable, and as such can only be applied to limited classes of environments, or are applied to more general classes of environments not covered by theoretical guarantees.
The principal objective of this work is to develop new measures of behavioral similarity that avoid the statistical and computational difficulties described above, and simultaneously capture richer information about the environment. We introduce the MICo (Matching under Independent Couplings) distance, and develop the theory around its computation and estimation, making comparisons with existing metrics on the basis of computational and statistical efficiency. We demonstrate the usefulness ∗Equal contribution. Correspondence to Pablo Samuel Castro: psc@google.com.
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
of the representations that MICo yields, both through empirical evaluations in small problems (where we can compute them exactly) as well as in two large benchmark suites: (1) the Arcade Learning Environment [Bellemare et al., 2013, Machado et al., 2018], in which the performance of a wide variety of existing value-based deep RL agents is improved by directly shaping representations via the MICo distance (see Figure 1, left), and (2) the DM-Control suite [Tassa et al., 2018], in which we demonstrate it can improve the performance of both Soft Actor-Critic [Haarnoja et al., 2018] and the recently introduced DBC [Zhang et al., 2021] (see Figure 1, right).

heading: 2 Background
text: Before describing the details of our contributions, we give a brief overview of the required background in reinforcement learning and bisimulation. We provide more extensive background in Appendix A.
Reinforcement learning. We consider a Markov decision process (X ,A, γ, P, r) defined by a finite state space X , finite action space A, transition kernel P : X ×A →P(X ), reward function r : X × A → R, and discount factor γ ∈ [0, 1). For notational convenience we will write P ax and rax for transitions and rewards, respectively. Policies are mappings from states to distributions over actions: π ∈ P(A)X and induce a value function V π : X → R defined via the recurrence: V π(x) := Ea∼π(x) [ rax + γEx′∼Pax [V π(x′)] ] . In RL we are concerned with finding the optimal policy π∗ = arg maxπ∈P(A)X V π from interaction with sample trajectories with an MDP, without knowledge of P or r (and sometimes not even X ), and the optimal value function V ∗ induced by π∗. State similarity and bisimulation metrics. Various notions of similarity between states in MDPs have been considered in the RL literature, with applications in policy transfer, state aggregation, and representation learning. The bisimulation metric [Ferns et al., 2004] is of particular relevance for this paper, and defines state similarity in an MDP by declaring two states x, y ∈ X to be close if their immediate rewards are similar, and the transition dynamics at each state leads to next states which are also judged to be similar. This self-referential notion is mathematically formalised by defining the bisimulation metric d∼ as the unique fixed-point of the operator TK :M(X)→M(X), whereM(X) = {d ∈ [0,∞)X×X : d symmetric and satisfies the triangle inequality} is the set of pseudometrics on X , given by TK(d)(x, y) = maxa∈A[|rax − ray |+ γWd(P ax , P ay )]. Here, Wd is the Kantorovich distance (also known as the Wasserstein distance) over the set of distributions P(X ) with base distance d, defined by Wd(µ, ν) = infX∼µ,Y∼ν E[d(X,Y )], for all µ, ν ∈P(X ), where the infimum is taken over all couplings of (X,Y ) with the prescribed marginals [Villani, 2008].
The mapping TK is a γ-contraction onM(X) under the L∞ norm [Ferns et al., 2011], and thus by standard contraction mapping arguments analogous to those used to study value iteration, it has a unique fixed point, the bisimulation metric d∼. Ferns et al. [2004] show that this metric bounds differences in the optimal value function, hence its importance in RL:
|V ∗(x)− V ∗(y)| ≤ d∼(x, y) ∀x, y ∈ X . (1)
Representation learning in RL. In large-scale environments, it is infeasible to express value functions directly as vectors in RX×A. Instead, RL agents must approximate value functions in a more concise manner, by forming a representation of the environment, that is, a feature embedding φ : X → RM , and predicting state-action values linearly from these features. Representation learning is the problem of finding a useful representation φ. Increasingly, deep RL agents are equipped with
additional losses to aid representation learning. A common approach is to require the agent to make additional predictions (so-called auxilliary tasks) with its representation, typically with the aid of extra network parameters, with the intuition that an agent is more likely to learn useful features if it is required to solve many related tasks. We refer to such methods as implicit representation shaping, since improved representations are a side-effect of learning to solve auxiliary tasks.
Since bisimulation metrics capture additional information about the MDP in addition to that summarised in value functions, bisimulation metrics are a natural candidate for auxiliary tasks in deep reinforcement learning. Gelada et al. [2019], Agarwal et al. [2021a], and Zhang et al. [2021] introduce auxiliary tasks based on bisimulation metrics, but require additional assumptions on the underlying MDP in order for the metric to be learnt correctly (Lipschitz continuity, deterministic, and Gaussian transitions, respectively). The success of these approaches provides motivation in this paper to introduce a notion of state similarity applicable to arbitrary MDPs, without further restriction. Further, we learn this state similarity explicitly: that is, without the aid of any additional network parameters.

heading: 3 Advantages and limitations of the bisimulation metric
text: The bisimulation metric d∼ is a strong notion of distance on the state space of an MDP; it is useful in policy transfer through its bound on optimal value functions [Castro and Precup, 2010] and because it is so stringent, it gives good guarantees for state aggregations [Ferns et al., 2004, Li et al., 2006]. However, it has been difficult to use at scale and compute online, for a variety of reasons that we summarize below.
(i) Computational complexity. The metric can be computed via fixed-point iteration since the operator TK is a contraction mapping. The map TK contracts at rate γ with respect to the L∞ norm onM, and therefore obtaining an ε-approximation of d∼ under this norm requires O(log(1/ε)/ log(1/γ)) applications of TK to an initial pseudometric d0. The cost of each application of TK is dominated by the computation of |X |2|A|Wd distances for distributions over X , each costing Õ(|X |2.5) in theory [Lee and Sidford, 2014], and Õ(|X |3) in practice [Pele and Werman, 2009, Guo et al., 2020a, Peyré and Cuturi, 2019]. Thus, the overall practical cost is Õ(|X |5|A| log(ε)/ log(γ)). (ii) Bias under sampled transitions. Computing TK requires access to the transition probability distributions P ax for each (x, a) ∈ X×Awhich, as mentioned in Section 2, are typically not available; instead, stochastic approximations to the operator of interest are employed. Whilst there has been work in studying online, sample-based approximate computation of the bisimulation metric [Ferns et al., 2006, Comanici et al., 2012], these methods are generally biased, in contrast to sample-based estimation of standard RL operators.
(iii) Lack of connection to non-optimal policies. One of the principal behavioural characterisations of the bisimulation metric d∼ is the upper bound shown in Equation (1). However, in general we do not have |V π(x)− V π(y)| ≤ d∼(x, y) for arbitrary policies π ∈ Π; a simple example is illustrated in Figure 2. More generally, notions of state similarity that the bisimulation metric encodes may not be closely related to behavioural similarity under an arbitrary policy π. Thus, learning about d∼ may not in itself be useful for large-scale reinforcement learning agents.
Property (i) expresses the intrinsic computational difficulty of computing this metric. Property (ii) illustrates the problems associated with attempting to move from operator-based computation to online, sampled-based computation of the metric (for example, when the environment dynamics are unknown). Finally, property (iii) shows that even if the metric is computable exactly, the information it yields about the MDP may not be practically useful. Although π-bisimulation (introduced by Castro
[2020] and extended by Zhang et al. [2021]) addresses property (iii), their practical algorithms are limited to MDPs with deterministic transitions [Castro, 2020] or MDPs with Gaussian transition kernels [Zhang et al., 2021]. Taken together, these three properties motivate the search for a metric without these shortcomings, which can be used in combination with deep reinforcement learning.

heading: 4 The MICo distance
text: We now present a new notion of distance for state similarity, which we refer to as MICo (Matching under Independent Couplings), designed to overcome the drawbacks described above.
Motivated by the drawbacks described in Section 3, we make several modifications to the operator TK introduced above: (i) in order to deal with the prohibitive cost of computing the Kantorovich distance, which optimizes over all coupling of the distributions P ax and P a y , we use the independent coupling; (ii) to deal with lack of connection to non-optimal policies, we consider an on-policy variant of the metric, pertaining to a chosen policy π ∈P(A)X . This leads us to the following definition. Definition 4.1. Given π ∈P(A)X , the MICo update operator TπM : RX×X → RX×X is:
(TπMU)(x, y) = |rπx − rπy |+ γEx′∼Pπx y′∼Pπy [U(x′, y′)] (2)
for all U : X × X → R, with rπx = ∑ a∈A π(a|x)rax and Pπx = ∑ a∈A π(a|x)P ax (·) for all x ∈ X .
As with the bisimulation operator, this can be thought of as encoding desired properties of a notion of similarity between states in a self-referential manner; the similarity of two states x, y ∈ X should be determined by the similarity of the rewards and the similarity of the states they lead to. Proposition 4.2. The operator TπM is a contraction mapping on RX×X with respect to the L∞ norm. Proof. See Appendix B.
The following corollary now follows immediately from Banach’s fixed-point theorem and the completeness of RX×X under the L∞ norm. Corollary 4.3. The MICo operator TπM has a unique fixed point Uπ ∈ RX×X , and repeated application of TπM to any initial function U ∈ RX×X converges to Uπ .
Having defined a new operator, and shown that it has a corresponding fixed-point, there are two questions to address: Does this new notion of distance overcome the drawbacks of the bisimulation metric described above; and what does this new object tell us about the underlying MDP?

heading: 4.1 Addressing the drawbacks of the bisimulation metric
text: We introduced the MICo distance as a means of overcoming some of the shortcomings associated with the bisimulation metric, described in Section 3. In this section, we provide a series of results that show that the newly-defined notion of distance addressess each of these shortcomings. The proofs of these results rely on the following lemma, connecting the MICo operator to a lifted MDP. This result is crucial for much of the analysis that follows, so we describe the proof in full detail. Lemma 4.4 (Lifted MDP). The MICo operator TπM is the Bellman evaluation operator for an auxiliary MDP.
Proof. Given the MDP specified by the tuple (X ,A, P,R), we construct an auxiliary MDP (X̃ , Ã, P̃ , R̃), by taking the state space to be X̃ = X 2, the action space to be Ã = A2, the transition dynamics to be given by P̃ (a,b)(u,v)((x, y)) = P a u (x)P b v (y) for all (x, y), (u, v) ∈ X 2, a, b ∈ A, and the action-independent rewards to be R̃(x,y) = |rπx − rπy | for all x, y ∈ X . The Bellman evaluation operator T̃ π̃ for this auxiliary MDP at discount rate γ under the policy π̃(a, b|x, y) = π(a|x)π(b|y) is given by (for all U ∈ RX×X and (x, y) ∈ X × X ):
(T̃ π̃U)(x, y) = R̃(x,y)+γ ∑
(x′,y′)∈X 2̃ P
(a,b) (x,y)((x ′, y′))π̃(a, b|x, y)U(x′, y′)
= |rπx − rπy |+ γ ∑
(x′,y′)∈X 2 Pπx (x ′)Pπy (y ′)U(x′, y′) = (TπMU)(x, y) .
Remark 4.5. Ferns and Precup [2014] noted that the bisimulation metric can be interpreted as the optimal value function in a related MDP, and that the functional TK of TK can be interpreted as a Bellman optimality operator. However, their proof was non-constructive, the related MDP being characterised via the solution of an optimal transport problem. In contrast, the connection described
above is constructive, and will be useful in understanding many of the theoretical properties of MICo. Ferns and Precup [2014] also note that the Wd distance in the definition of TK can be upper-bounded by taking a restricted class of couplings of the transition distributions. The MICo metric can be viewed as restricting the coupling class precisely to the singleton containing the independent coupling.
With Lemma 4.4 established, we can now address each of the points (i), (ii), and (iii) from Section 3.
(i) Computational complexity. The key result regarding the computational complexity of computing the MICo distance is as follows. Proposition 4.6 (MICo computational complexity). The computational complexity of computing an ε-approximation in L∞ to the MICo metric is O(|X |4 log(ε)/ log(γ)). Proof. Since, by Proposition 4.2, the operator TπM is a γ-contraction under L
∞, we require O(log(ε)/ log(γ)) applications of the operator to obtain an ε-approximation in L∞. Each iteration of value iteration updates |X |2 table entries, and the cost of each update is O(|X |2), leading to an overall cost of O(|X |4 log(ε)/ log(γ)).
In contrast to the bisimulation metric, this represents a computational saving of O(|X |), which arises from the lack of a need to solve optimal transport problems over the state space in computing the MICo distance. There is a further saving of O(|A|) that arises since MICo focuses on an individual policy π, and so does not require the max over actions in the bisimulation operator definition.
(ii) Online approximation. Due to the interpretation of the MICo operator TπM as the Bellman evaluation operator in an auxiliary MDP, established in Lemma 4.4, algorithms and associated proofs of correctness for computing the MICo distance online can be straightforwardly derived from standard online algorithms for policy evaluation. We describe a straightforward approach, based on the TD(0) algorithm, and also note that the wide range of online policy evaluation methods incorporating off-policy corrections and multi-step returns, as well as techniques for applying such methods at scale, may also be used.
Given a current estimate Ut of the fixed point of TπM and a pair of observations (x, a, r, x ′), (y, b, r̃, y′) generated under π, we can define a new estimate Ut+1 via
Ut+1(x, y)← (1− t(x, y))Ut(x, y) + t(x, y)(|r − r̃|+ γUt(x′, y′)) (3)
and Ut+1(x̃, ỹ) = Ut(x̃, ỹ) for all other state-pairs (x̃, ỹ) 6= (x, y), for some sequence of stepsizes { t(x, y) | t ≥ 0, (x, y) ∈ X 2}. Sufficient conditions for convergence of this algorithm can be deduced straightforwardly from corresponding conditions for TD(0). We state one such result below. An important caveat is that the correctness of this particular algorithm depends on rewards depending only on state; one can switch to state-action metrics if this hypothesis is not satisfied. Proposition 4.7. Suppose rewards depend only on state, and consider the sequence of estimates (Ut)t≥0, with U0 initialised arbitrarily, and Ut+1 updated from Ut via a pair of transitions (xt, at, rt, x ′ t), (yt, bt, r̃t, y ′ t) as in Equation (3). If all state-pairs tuples are updated infinitely often, and stepsizes for these updates satisfy the Robbins-Monro conditions. Then Ut → Uπ almost surely. Proof. Under the assumptions of the proposition, the update described is exactly a TD(0) update in the lifted MDP described in Lemma 4.4. We can therefore appeal to Proposition 4.5 of Bertsekas and Tsitsiklis [1996] to obtain the result.
Thus, in contrast to the Kantorovich metric, convergence to the exact MICo metric is possible with an online algorithm that uses sampled transitions.
(iii) Relationship to underlying policy. In contrast to the bisimulation metric, we have the following on-policy guarantee for the MICo metric. Proposition 4.8. For any π ∈P(A)X and states x, y ∈ X , we have |V π(x)− V π(y)| ≤ Uπ(x, y). Proof. We apply a coinductive argument [Kozen, 2007] to show that if |V π(x) − V π(y)| ≤ U(x, y) for all x, y ∈ X , for some U ∈ RX×X symmetric in its two arguments, then we also have |V π(x) − V π(y)| ≤ (TπMU)(x, y) for all x, y ∈ X . Since the hypothesis holds for the constant function U(x, y) = 2 maxz,a |r(z, a)|/(1 − γ), and TπM contracts around Uπ, the conclusion then
follows. Therefore, suppose the hypothesis holds. Then we have V π(x)− V π(y) = rπx − rπy + γ ∑ x′∈X Pπx (x ′)V (x′)− γ ∑ y′∈X Pπy (y ′)V (y′)
≤ |rπx − rπy |+ γ ∑
x′,y′∈X Pπx (x ′)Pπy (y ′)(V π(x′)− V π(y′))
≤ |rπx − rπy |+ γ ∑
x′,y′∈X Pπx (x ′)Pπy (y ′)U(x′, y′) = (TπMU)(x, y) .
By symmetry, V π(y)− V π(x) ≤ (TπMU)(x, y), as required.

heading: 4.2 Diffuse metrics
text: To characterize the nature of the fixed point Uπ , we introduce the notion of a diffuse metric. Definition 4.9. Given a set X , a function d : X × X → R is a diffuse metric if the following axioms hold: (i) d(x, y) ≥ 0 for any x, y ∈ X ; (ii) d(x, y) = d(y, x) for any x, y ∈ X ; (iii) d(x, y) ≤ d(x, z) + d(y, z) ∀x, y, z ∈ X .
These differ from the standard metric axioms in the first point: we no longer require that a point has zero self-distance, and two distinct points may have zero distance. Notions of this kind are increasingly common in machine learning as researchers develop more computationally tractable versions of distances, as with entropy-regularised optimal transport distances [Cuturi, 2013], which also do not satisfy the axiom of zero self-distance.
An example of a diffuse metric is the Łukaszyk–Karmowski distance [Łukaszyk, 2004], which is used in the MICo metric as the operator between the next-state distributions. Given a diffuse metric space (X , ρ), the Łukaszyk–Karmowski distance dρLK is a diffuse metric on probability measures on X given by dρLK(µ, ν) = Ex∼µ,y∼ν [ρ(x, y)]. This example demonstrates the origin of the name diffuse metrics: the non-zero self distances arises from a point being spread across a probability distribution. In terms of the Łukaszyk–Karmowski distance, the MICo distance can be written as the fixed point Uπ(x, y) = |rπx − rπy |+ dLK(Uπ)(Pπx , Pπy ). This characterisation leads to the following result. Proposition 4.10. The MICo distance is a diffuse metric. Proof. Non-negativity and symmetry of Uπ are clear, so it remains to check the triangle inequality. To do this, we define a sequence of iterates (Uk)k≥0 in RX×X by U0(x, y) = 0 for all x, y ∈ X , and Uk+1 = TπMUk for each k ≥ 0. Recall that by Corollary 4.3 that Uk → Uπ. We will show that each Uk satisfies the triangle inequality by induction. By taking limits on either side of the inequality, we will then recover that Uπ itself satisfies the triangle inequality. The base case of the inductive argument is clear from the choice of U0. For the inductive step, assume that for some k ≥ 0, Uk(x, y) ≤ Uk(x, z) + Uk(z, y) for all x, y, z ∈ X . Now for any x, y, z ∈ X , we have
Uk+1(x, y) = |rπx − rπy |+ γEX′∼Pπx ,Y ′∼Pπy [Uk(X ′, Y ′)]
≤ |rπx − rπz |+ |rπz − rπy |+ γEX′∼Pπx ,Y ′∼Pπy ,Z′∼Pπz [Uk(X ′, Z ′) + Uk(Z ′, Y ′)]
= Uk+1(x, z) + Uk+1(z, y) .
It is interesting to note that a state x ∈ X has zero self-distance iff the Markov chain induced by π initialised at x is deterministic, and the magnitude of a state’s self-distance is indicative of the amount of “dispersion” in the distribution. Hence, in general, we have Uπ(x, x) > 0, and Uπ(x, x) 6= Uπ(y, y) for distinct states x, y ∈ X . See the appendix for further discussion of diffuse metrics and related constructions.

heading: 5 The MICo loss
text: The impetus of our work is the development of principled mechanisms for directly shaping the representations used by RL agents so as to improve their learning. In this section we present a novel loss based on the MICo update operator TπM given in Equation (2) that can be incorporated into any RL agent. Given the fact that MICo is a diffuse metric that can admit non-zero self-distances, special
care needs to be taken in how these distances are learnt; indeed, traditional mechanisms for measuring distances between representations (e.g. Euclidean and cosine distances) are geometrically-based and enforce zero self-distances.
We assume an RL agent learning an estimate Qξ,ω defined by the composition of two function approximators ψ and φ with parameters ξ and ω, respectively: Qξ,ω(x, ·) = ψξ(φω(x)) (note that this can be the critic in an actor-critic algorithm such as SAC). We will refer to φω(x) as the representation of state x and aim to make distances between representations match the MICo distance; we refer to ψξ as the value approximator. We define the parameterized representation distance, Uω, as an approximant to Uπ:
Uπ(x, y) ≈ Uω(x, y) := ‖φω(x)‖22 + ‖φω(y)‖22
2 + βθ(φω(x), φω(y))
where θ(φω(x), φω(y)) is the angle between vectors φω(x) and φω(y) and β is a scalar (in our results we use β = 0.1 but present results with other values of β in the appendix).
Based on Equation (2), our learning target is then TUω̄ (rx, x ′, ry, y ′) = |rx−ry|+γUω̄(x′, y′), where ω̄ is a separate copy of the network parameters that are synchronised with ω at infrequent intervals. This is a common practice that was introduced by Mnih et al. [2015] (and in fact, we use the same update schedule they propose). The loss for this learning target is
LMICo(ω) =E〈x,rx,x′〉,〈y,ry,y′〉 [( TUω̄ (rx, x ′, ry, y ′)− Uω(x, y) )2] where 〈x, rx, x′〉 and 〈y, ry, y′〉 are pairs of transitions sampled from the agent’s replay buffer. We can combine LMICo with the temporal-difference loss LTD of any RL agent as (1− α)LTD + αLMICo, where α ∈ (0, 1). Each sampled mini-batch is used for both MICo and TD losses. Figure 3 (left) illustrates the network architecture used for learning.
Although the loss LMICo is designed to learn the MICo diffuse metric Uπ, the values of the metric itself are parametrised through Uω defined above, which is constituted by several distinct terms. This appears to leave a question as to how the representations φω(x) and φω(y), as Euclidean vectors, are related to one another when the MICo loss is minimised. Careful inspection of the form of Uω(x, y) shows that the (scaled) angular distance between φω(x) and φω(y) can be recovered from Uω by subtracting the learnt approximations to the self-distances Uπ(x, x) and Uπ(y, y) (see Figure 3, right). We therefore define the reduced MICo distance ΠUπ, which encodes the distances enforced between the representation vectors φω(x) and φω(y), by:
βθ(φω(x), φω(y)) ≈ ΠUπ(x, y) = Uπ(x, y)− 1 2 Uπ(x, x)− 1 2 Uπ(y, y) .
In the following section we investigate the following two questions: (1) How informative of V π is ΠUπ?; and (2) How useful are the features encountered by ΠUπ for policy evaluation? We conduct these investigations on tabular environments where we can compute the metrics exactly, which helps clarify the behaviour of our loss when combined with deep networks in Section 6.

heading: 5.1 Value bound gaps
text: Although Proposition 4.8 states that we have |V π(x)− V π(y)| ≤ Uπ(x, y), we do not, in general, have the same upper bound for ΠUπ(x, y) as demonstrated by the following result.
Lemma 5.1. There exists an MDP with x, y ∈ X , and π ∈ Π where |V π(x)−V π(y)| > ΠUπ(x, y). Proof. Consider a single-action MDP with two states (x and y) where y is absorbing, x transitions with equal probability to x and y, and a reward of 1 is received only upon taking an action from state x. There is only one policy for this MDP which yields the value function V (x) ≈ 1.8 and V (y) = 0. The MICo distance gives U(x, x) ≈ 1.06, U(x, y) ≈ 1.82, and U(y, y) = 0, while the reduced MICo distance yields ΠU(x, x) = ΠU(y, y) = 0, and ΠU(x, y) ≈ 1.29 < |V (x)− V (y)| = 1.8.
Despite this negative result, it is worth evaluating how often in practice this inequality is violated and by how much, as this directly impacts the utility of this distance for learning representations.
To do so, we make use of Garnet MDPs, a class of randomly generated MDPs [Archibald et al., 1995, Piot et al., 2014]. Given a specified number of states nX and the number of actions nA, Garnet(nX , nA) is generated as follows: 1. The branching factor bx,a of each transition P ax is sampled uniformly from [1 : nX ]. 2. bx,a states are picked uniformly randomly from X and assigned a random value in [0, 1]; these values are then normalized to produce a proper distribution P ax . 3. Each rax is sampled uniformly in [0, 1].
For each Garnet(nX , nA) we sample 100 stochastic policies {πi} and compute the average gap: 1 100|X |2 ∑ i ∑ x,y d(x, y) − |V πi(x) − V πi(y)|, where d stands for any of the considered metrics. Note we are measuring the signed difference, as we are interested in the frequency with which the upper bound is violated. As seen in Figure 4 (left), our metric does on average provide an upper bound on the difference in values that is also tighter bound than those provided by Uπ and π-bisimulation. This suggests that the resulting representations remain informative of value similarities.

heading: 5.2 State features
text: In order to investigate the usefuleness of the representations produced by ΠUπ, we construct state features directly by using the computed distances to project the states into a lower-dimensional space with the UMAP dimensionality reduction algorithm [McInnes et al., 2018]2. We then apply linear regression of the true value function V π against the features to compute V̂ π and measure the average error across the state space. As baselines we compare against random features (RF), Proto Value Functions (PVF) [Mahadevan and Maggioni, 2007], and the features produced by π-bisimulation [Castro, 2020]. We present our results on the well-known four-rooms GridWorld [Sutton et al., 1999] in Figure 4 (right) and provide results on more environments in the appendix. Despite the independent couplings, ΠUπ performs on par with π-bisimulation, which optimizes over all couplings.

heading: 6 Large-scale empirical evaluation
text: Having developed a greater understanding of the properties inherent to the representations produced by the MICo loss, we evaluate it on the Arcade Learning Environment [Bellemare et al., 2013]. We added the MICo loss to all the JAX agents provided in the Dopamine library [Castro et al., 2018]: DQN [Mnih et al., 2015], Rainbow [Hessel et al., 2018], QR-DQN [Dabney et al., 2018b],
2Note that since UMAP expects a metric, it is ill-defined with the diffuse metric Uπ .
and IQN [Dabney et al., 2018a], using mean squared error loss to minimize LTD for DQN (as suggested by Obando-Ceron and Castro [2021]). Given the state-of-the-art results demonstrated by the Munchausen-IQN (M-IQN) agent [Vieillard et al., 2020], we also evaluated incorporating our loss into M-IQN.3 For all experiments we used the hyperparameter settings provided with Dopamine. We found that a value of α = 0.5 worked well with quantile-based agents (QR-DQN, IQN, and M-IQN), while a value of α = 0.01 worked well with DQN and Rainbow. We hypothesise that the difference in scale of the quantile, categorical, and non-distributional loss functions concerned leads to these distinct values of α performing well. We found it important to use the Huber loss [Huber, 1964] to minimize LMICo as this emphasizes greater accuracy for smaller distances as oppoosed to larger distances. We experimented using the MSE loss but found that larger distances tended to overwhelm the optimization process, thereby degrading performance.
We evaluated on all 60 Atari 2600 games over 5 seeds and report the results in Figure 1 (left), using the interquantile metric (IQM), proposed by Agarwal et al. [2021b] as a more robust and reliable alternative to mean and median (which are reported in Figure 6). The fact that the MICo loss provides consistent improvements over a wide range of baseline agents of varying complexity suggests that the MICo loss can help learn better representations for control.
Additionally, we evaluated the MICo loss on twelve of the DM-Control suite from pixels environments [Tassa et al., 2018]. As a base agent we used Soft Actor-Critic (SAC) [Haarnoja et al., 2018] with the convolutional auto-encoder described by Yarats et al. [2019]. We applied the MICo loss on the output of the auto-encoder (with α = 1e − 5) and maintained all other parameters untouched. Recently, Zhang et al. [2021] introduced DBC, which learns a dynamics and reward model on the output of the auto-encoder; their bisimulation loss uses the learned dynamics model in the computation of the Kantorovich distance between the next state transitions. We consider two variants of their algorithm: one which learns a stochastic dynamics model (DBC), and one which learns a deterministic dynamics model (DBC-Det). We replaced their bisimulation loss with the MICo loss (which, importantly, does not require a dynamics model) and kept all other parameters untouched. As Figure 1 illustrates, the best performance is achieved with SAC augmented with the MICo loss; additionally, replacing the bisimulation loss of DBC with the MICo loss is able to recover the performance of DBC to match that of SAC.
Additional details and results are provided in the appendix.

heading: 7 Related Work
text: Bisimulation metrics were introduced for MDPs by Ferns et al. [2004], and have been extended in a number of directions [Ferns et al., 2005, 2006, Taylor, 2008, Taylor et al., 2009, Ferns et al., 2011, Comanici et al., 2012, Bacci et al., 2013a,b, Abate, 2013, Ferns and Precup, 2014, Castro, 2020], with applications including policy transfer [Castro and Precup, 2010, Santara et al., 2019], representation learning [Ruan et al., 2015, Comanici et al., 2015], and state aggregation [Li et al., 2006].
A range of other notions of similarity in MDPs have also been considered, such as action sequence equivalence [Givan et al., 2003], temporally extended metrics [Amortila et al., 2019], MDP homomorphisms [Ravindran and Barto, 2003], utile distinction [McCallum, 1996], and policy irrelevance [Jong and Stone, 2005], as well as notions of policy similarity [Pacchiano et al., 2020, Moskovitz et al., 2021]. Li et al. [2006] review different notions of similarity applied to state aggregation. Recently, Le Lan et al. [2021] performed an exhaustive analysis of the continuity properties, relative to functions of interest in RL, of a number of existing metrics in the literature.
The notion of zero self-distance, central to the diffuse metrics defined in this paper, is increasingly encountered in machine learning applications involving approximation of losses. Of particular note is entropy-regularised optimal transport [Cuturi, 2013] and related quantities [Genevay et al., 2018, Fatras et al., 2020, Chizat et al., 2020, Fatras et al., 2021].
More broadly, many approaches to representation learning in deep RL have been considered, such as those based on auxiliary tasks (see e.g. [Sutton et al., 2011, Jaderberg et al., 2017, Bellemare et al.,
3Given that the authors of M-IQN had implemented their agent in TensorFlow (whereas our agents are in JAX), we have reimplemented M-IQN in JAX and run 5 independent runs (in contrast to the 3 run by Vieillard et al. [2020].
2017, François-Lavet et al., 2019, Gelada et al., 2019, Guo et al., 2020b]), and other approaches such as successor features [Dayan, 1993, Barreto et al., 2017].

heading: 8 Conclusion
text: In this paper, we have introduced the MICo distance, a notion of state similarity that can be learnt at scale and from samples. We have studied the theoretical properties of MICo, and proposed a new loss to make the non-zero self-distances of this diffuse metric compatible with function approximation, combining it with a variety of deep RL agents to obtain strong performance on the Arcade Learning Environment. In contrast to auxiliary losses that implicitly shape an agent’s representation, MICo directly modifies the features learnt by a deep RL agent; our results indicate that this helps improve performance. To the best of our knowledge, this is the first time directly shaping the representation of RL agents has been successfully applied at scale. We believe this represents an interesting new approach to representation learning in RL; continuing to develop theory, algorithms and implementations for direct representation shaping in deep RL is an important and promising direction for future work.
Broader impact statement
This work lies in the realm of “foundational RL” in that it contributes to the fundamental understanding and development of reinforcement learning algorithms and theory. As such, despite us agreeing in the importance of this discussion, our work is quite far removed from ethical issues and potential societal consequences.

heading: 9 Acknowledgements
text: The authors would like to thank Gheorghe Comanici, Rishabh Agarwal, Nino Vieillard, and Matthieu Geist for their valuable feedback on the paper and experiments. Pablo Samuel Castro would like to thank Roman Novak and Jascha Sohl-Dickstein for their help in getting angular distances to work stably! Thanks to Hongyu Zang for pointing out that the x-axis labels for the SAC experiments needed to be fixed. Finally, the authors would like to thank the reviewers (both ICML’21 and NeurIPS’21) for helping make this paper better.

Review summaries:
The paper proposes a new behavioral distance over the state space of an MDP, which can be used to learn representations in deep RL tasks. The proposed distance has three advantages over the existing bisimulation metric: (1) it has better computational complexity; (2) it allows for online approximation; and (3) it has a close connection to the underlying policy. Empirical evidence on toy examples shows the informativeness of the representations learned by using the proposed distance, and large-scale experiments show improvements of deep RL agents when incorporating the proposed distance to learn representations. 

This work proposes a new formulation for quantifying the similarity of states based on rewards and next states. It points out three limitations of the bisimulation metric, and introduces the MICo distance. Given a specific policy, the MICo distance is defined with the MICo update operator, which the authors prove is a contraction mapping. They demonstrate some properties of it especially as a solution to the issues of the bisimulation metric. The authors suggest a specific form of approximation to the unique fixed point of the contraction mapping, so that it can be used for shaping state representations in environments with large state spaces. Empirically, the authors performed experiments to show the usefulness of a modified version of the MICo distance's learned approximation equipped with UMAP for predicting value functions and how the scores are improved with the state representations trained jointly with the MICo loss.

This paper introduces MICo, a new behavioural distance over the state space of an MDP. Compared to previous works on behavioural metrics such as bisimulation metrics, MICo is compatible with sample-based algorithms and therefore can be easily integrated with deep RL methods. Empirical results demonstrate that learning this distance improves performance on the Atari benchmark.