{'fraction_non_alphanumeric': 0.045069186167386636, 'fraction_numerical': 0.012636004527709117, 'mean_word_length': 4.338512763596005, 'pattern_counts': {'":': 0, '<': 1, '<?xml version=': 0, '>': 1, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
MICo: Improved representations via sampling-based state similarity for Markov decision processes

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.


Introduction
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 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).

Background
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 a x and r a x 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) := E a∼π(x) r a x + γE x ∼P a x [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 T K : M(X) → M(X), where M(X) = {d ∈ [0, ∞) X ×X : d symmetric and satisfies the triangle inequality} is the set of pseudometrics on X , given by Here, W d is the Kantorovich distance (also known as the Wasserstein distance) over the set of distributions P(X ) , for all µ, ν ∈ P(X ), where the infimum is taken over all couplings of (X, Y ) with the prescribed marginals [Villani, 2008].
The mapping T K is a γ-contraction on M(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: Representation learning in RL. In large-scale environments, it is infeasible to express value functions directly as vectors in R X ×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 → R M , 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 2 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.

Advantages and limitations of the bisimulation metric
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 T K is a contraction mapping. The map T K contracts at rate γ with respect to the L ∞ norm on M, and therefore obtaining an ε-approximation of d ∼ under this norm requires O(log(1/ε)/ log(1/γ)) applications of T K to an initial pseudometric d 0 . The cost of each application of T K is dominated by the computation of |X | 2 |A| W d distances for distributions over X , each costingÕ(|X | 2.5 ) in theory [Lee and Sidford, 2014], andÕ(|X | 3 ) in practice [Pele and Werman, 2009, Guo et al., 2020a, Peyré and Cuturi, 2019. Thus, the overall practical cost isÕ(|X | 5 |A| log(ε)/ log(γ)).
(ii) Bias under sampled transitions. Computing T K requires access to the transition probability distributions P a x for each (x, a) ∈ X ×A which, 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.

4 The MICo distance
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 T K 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 a x 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 : R X ×X → R X ×X is: for all U : X × X → R, with r π x = a∈A π(a|x)r a x and P π x = a∈A π(a|x)P a x (·) 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 R X ×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 R X ×X under the L ∞ norm. Corollary 4.3. The MICo operator T π M has a unique fixed point U π ∈ R X ×X , and repeated application of T π M to any initial function U ∈ R X ×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?

Addressing the drawbacks of the bisimulation metric
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 = A 2 , the transition dynamics to be given by P , 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 ∈ R X ×X and (x, y) ∈ X × X ): 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 T K of T K 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 W d distance in the definition of T K 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 U t 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 U t+1 via and U t+1 (x,ỹ) = U t (x,ỹ) for all other state-pairs (x,ỹ) = (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 (U t ) t≥0 , with U 0 initialised arbitrarily, and U t+1 updated from U t via a pair of transitions (x t , a t , r t , x t ), (y t , b t ,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 U t → 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] for all x, y ∈ X , for some U ∈ R X ×X symmetric in its two arguments, then we also have for all x, y ∈ X . Since the hypothesis holds for the constant function U (x, y) = 2 max z,a |r(z, a)|/(1 − γ), and T π M contracts around U π , the conclusion then follows. Therefore, suppose the hypothesis holds. Then we have

Diffuse metrics
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: 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 ρ . 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 | + d LK (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 (U k ) k≥0 in R X ×X by U 0 (x, y) = 0 for all x, y ∈ X , and U k+1 = T π M U k for each k ≥ 0. Recall that by Corollary 4.3 that U k → U π . We will show that each U k 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 U 0 . For the inductive step, assume that for some k ≥ 0, U k (x, y) ≤ U k (x, z) + U k (z, y) for all x, y, z ∈ X . Now for any x, y, z ∈ X , we have 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) = U π (y, y) for distinct states x, y ∈ X . See the appendix for further discussion of diffuse metrics and related constructions.

The MICo loss
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 π : 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 T Ū ω (r x , x , r y , y ) = |r x − r y | + γUω(x , y ), wherē ω 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 where x, r x , x and y, r y , y are pairs of transitions sampled from the agent's replay buffer. We can combine L MICo with the temporal-difference loss L TD of any RL agent as (1 − α)L TD + αL MICo , 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 L MICo 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: 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.

Value bound gaps
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. shaded areas represent 95% confidence intervals.
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 n X and the number of actions n A , Garnet(n X , n A ) is generated as follows: 1. The branching factor b x,a of each transition P a x is sampled uniformly from [1 : n X ]. 2. b x,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 a x .

Each r a
x is sampled uniformly in [0, 1]. For each Garnet(n X , n A ) we sample 100 stochastic policies {π i } and compute the average gap: 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.

State features
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 computeV π 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.  [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 L MICo 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.

Conclusion
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.

Supplementary Material: MICo: Improved representations via sampling-based state similarity for Markov decision processes A Extended background material
In this section we provide a more extensive background review.

A.1 Reinforcement learning
In this section we give a slightly more expansive overview of relevant key concepts in reinforcement learning, without the space constraints of the main paper. Denoting by P(S) the set of probability distributions on a set S, we define a Markov decision process (X , A, γ, P, r) as: • A finite state space X ; • A finite action space A; • A transition kernel P : X × A → P(X ); • A reward function r : X × A → R; • A discount factor γ ∈ [0, 1).
For notational convenience we introduce the notation P a x ∈ P(X ) for the next-state distribution given state-action pair (x, a), and r a x for the corresponding immediate reward. Policies are mappings from states to distributions over actions: π ∈ P(A) X and induce a value function V π : X → R defined via the recurrence: . It can be shown that this recurrence uniquely defines V π through a contraction mapping argument [Bertsekas and Tsitsiklis, 1996].
The control problem is concerned with finding the optimal policy π * = arg max It can be shown that while the optimisation problem above appears to have multiple objectives (one for each coordinate of V π , there is in fact a policy π * ∈ P(A) X that simultaneously maximises all coordinates of V π , and that this policy can be taken to be deterministic; that is, for each x ∈ X , π(·|x) ∈ P(A) attributes probability 1 to a single action. In reinforcement learning in particular, we are often interested in finding, or approximating, π * from direct interaction with the MDP in question via sample trajectories, without knowledge of P or r (and sometimes not even X ).

A.2 Metrics
A metric d on a set X is a function d : X × X → [0, ∞) respecting the following axioms for any x, y, z ∈ X:

A.3 State similarity and bisimulation metrics
Bisimulation is a fundamental notion of behavioural equivalence introduced by Park and Milner [Milner, 1989] in the early 1980s in the context of nondeterministic transition systems. The probabilistic analogue was introduced by Larsen and Skou [1991]. The notion of an equivalence relation is not suitable to capture the extent to which quantitative systems may resemble each other in behaviour. To provide a quantitative notion, bisimulation metrics were introduced by Desharnais et al. [1999,2004] in the context of probabilistic transition systems without rewards. In reinforcement learning the reward is an important ingredient, accordingly the bisimulation metric for states of MDPs was introduced by Ferns et al. [2004]. Much work followed this initial introduction of bisimulation metrics into RL, as described in the main paper. We briefly reviewed the bisimulation metric in Section 2, and now provide additional detail around some of the key associated mathematical concepts.
Central to the definition of the bisimulation metric is the operator T k : M(X ) → M(X ), defined over M(X ), the space of pseudometrics on X . Pseudometrics were explored in more detail in Section A.2. We now turn to the definition of the operator itself, given by for each d ∈ M(X ), and each x, y ∈ X . It can be verified that the function T K (d) : X × X → R satisfies the properties of a pseudometric, so under this definition T K does indeed map M(X ) into itself.
The other central mathematical concept underpinning the operator T K is the Wasserstein distance W d using base metric d. W d is formally a pseudometric over the set of probability distributions P(X ), defined as the solution to an optimisation problem. The problem specifically is formulated as finding an optimal coupling between the two input probability distributions that minimises a notion of transport cost associated with d. Mathematically, for two probability distributions µ, µ ∈ P(X ), we have Note that the pair of random variables (Z, Z ) attaining the minimum in the above expression will in general not be independent. That the minimum is actually attained in the above example in the case of a finite set X can be seen by expressing the optimisation problem as a linear program. Minima are obtained in much more general settings too; see Villani [2008].
Finally, the operator T K can be analysed in a similar way to standard operators in dynamic programming for reinforcement learning. It can be shown that it is a contraction mapping with respect to the L ∞ metric over M(X ), and that M(X ) is a complete metric space with respect to the same metric [Ferns et al., 2011]. Thus, by Banach's fixed point theorem, T K has a unique fixed point in M(X ), and repeated application of T K to any initial pseudometric will converge to this fixed point.

A.4 Further details on diffuse and partial metrics
The notion of a distance function having non-zero self distance was first introduced by Matthews [1994] who called it a partial metric. We define it below: Definition A.1. Given a set X , a function d : X × X → R is a partial metric if the following axioms hold: This definition was introduced to recover a proper metric from the distance function: that is, given a partial metric d, one is guaranteed thatd(x, y) = d(x, y) − 1 2 (d(x, x) + d(y, y)) is a proper metric. The above definition is still too stringent for the Łukaszyk-Karmowski distance (and hence MICo distance), since it fails axiom 4 as shown in the following counterexample. Example A.2. The Łukaszyk-Karmowski distance does not satisfy the modified triangle inequality: let X be [0, 1], and ρ be the Euclidean distance | · |. Let µ,ν be Dirac measures concentrated at 0 and 1, and let η be 1 2 (δ 0 + δ 1 ). Then one can calculate that d LK (ρ)(µ, ν) = 1, while d LK (ρ)(µ, η) + d LK (ρ)(ν, η) − d LK (ρ)(η, η) = 1/2, breaking the inequality.
This naturally leads us to the notion of diffuse metrics defined in the main paper.

B Proof of Proposition 4.2
Proposition 4.2. The operator T π M is a contraction mapping on R X ×X with respect to the L ∞ norm.
Proof. Let U, U ∈ R X ×X . Then note that for any x, y ∈ X , as required.

C Experimental details
We will first describe the regular network and training setup for these agents so as to facilitate the description of our loss.

C.1 Baseline network and loss description
The networks used by Dopamine for the ALE consist of 3 convolutional layers followed by two fully-connected layers (the output of the networks depends on the agent). We denote the output of the convolutional layers by φ ω with parameters ω, and the remaining fully connected layers by ψ ξ with parameters ξ. Thus, given an input state x (e.g. a stack of 4 Atari frames), the output of the network is Q ξ,ω (x, ·) = ψ ξ (φ ω (x)). Two copies of this network are maintained: an online network and a target network; we will denote the parameters of the target network byξ andω. During learning, the parameters of the online network are updated every 4 environment steps, while the target network parameters are synced with the online network parameters every 8000 environment steps. We refer to the loss used by the various agents considered as L TD ; for example, for DQN this would be: where D is a replay buffer with a capacity of 1M transitions, and ρ is the Huber loss.

C.2 MICo loss description
We will be applying the MICo loss to φ ω (x). As described in Section 5, we express the distance between two states as: where θ(φ ω (x), φω(y)) is the angle between vectors φ ω (x) and φω(y) and β is a scalar. Note that we are using the target network for the y representations; this was done for learning stability. We used β = 0.1 for the results in the main paper, but present some results with different values of β below.
In order to get a numerically stable operation, we implement the angular distance between representations φ ω (x) and φ ω (y) according to the calculations Based on Equation (2), our learning target is then (note the target network is used for both representations here): T Ū ω (r x , x , r y , y ) = |r x − r y | + γUω(x , y ) , and the loss is As mentioned in Section 5, we use the same mini-batch sampled for L TD for computing L MICo . Specifically, we follow the method introduced by Castro [2020] for constructing new matrices that allow us to compute the distances between all pairs of sampled states (see code for details on matrix operations). Our combined loss is then L α (ξ, ω) = (1 − α)L TD (ξ, ω) + αL MICo (ω).

C.3 Hyperparameters for soft actor-critic
We re-implemented the DBC algorithm from Zhang et al. [2021] on top of the Soft Actor-Critic algorithm [Haarnoja et al., 2018] provided by the Dopamine library [Castro et al., 2018]. We compared the following algorithms, using the same hyperparameters for all 4 : 3. DBC-Det: In the code provided by Zhang et al. [2021], the default setting was to assume deterministic transitions (which is an easier dynamics model to learn), so we decided to compare against this version as well. It is interesting to note that the performance is roughly the same as for DBC.

4.
MICo: This a modified version of SAC, adding the MICo loss to the output of the encoder. Note that the encoder output is the same one used by DBC for their dynamics and reward models.

5.
DBC+MICo: Instead of using the bisimulation loss of Zhang et al. [2021], which relies on the learned dynamics model, we use our MICo loss. We kept all other components untouched (so a dynamics and reward model were still being learned).
It is worth noting that some of the hyperparameters we used differ from those listed in the code provided by Zhang et al. [2021]; in our experiments they hyperparameters for all agents are based on the default SAC hyperparameters in the Dopamine library [Castro et al., 2018].
For the ALE experiments we used the "squaring" of the sampled batches introduced by Castro [2020] (where all pairs of sampled states are considered). However, the implementation provided by Zhang et al. [2021] instead created a copy of the sampled batch of transitions and shuffled them; we chose to follow this setup for the SAC-based experiments. Thus, while in the ALE experiments we are comparing m 2 pairs of states (where m is the batch size) at each training step, in the SAC-based experiments we are only comparing m pairs of states.
The aggregate results are displayed in Figure 1, and per-environment results in Figure 14 below.

D.1 Additional state feature results
The results shown in Figure 4 are on the well-known four-rooms GridWorld [Sutton et al., 1999]. We provide extra experiments in Figure 5.

D.2 Complete ALE experiments
We additionally provide complete results for all the agents in Figure 8, Figure 9, Figure 10, Figure 11, and Figure 12.

D.3 Sweep over α and β values
In Figure 13 we demonstrate the performance of the MICo loss when added to Rainbow over a number of different values of α and β. For each agent, we ran a similar hyperparameter sweep over α and β on the same six games displayed in Figure 13 to determine settings to be used in the full ALE experiments.

D.4 Complete DM-Control results
Full per-environment results are provided in Figure 14.

D.5 Compute time and infrastructure
For Figure 4 each run took approximately 10 minutes. For Figure 4 and Figure 5 the running time varied for each environment and per metric but a conservative estimate is 30 minutes per run. All GPU experiments were run on NVIDIA Tesla P100 GPUs. Each Atari game takes approximately 5 days (300 hours) to run for 200M frames.  The results for all games and agents are over 5 independent runs, and shaded regions report 75% confidence intervals.