MICo: Improved representations via sampling-based state similarity for Markov decision processes
Authors: Pablo Samuel Castro
Tyler Kastner
Prakash Panangaden
Mark Rowland
Venue: NeurIPS 2021 Poster (2021)
Keywords: Reinforcement Learning
Metrics
Deep Reinforcement Learning

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 analyses, 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.

MICo: Improved representations via sampling-based state similarity for Markov decision processes
Pablo Samuel Castro" Google Research, Brain Team
Tyler Kastner" McGill University
Prakash Panangaden McGill University
Mark Rowland DeepMind
Abstract
We present a new behavioural distance over the state space of Markov deci- sion process, and demonstrate the use of this distance shaping the learnt representations of deep reinforcement learning agents. While existing notions of state similarity are typically difficult to learn at scale due high computational cost and lack of sample-based algorithms, our newly-prope pwwly-pooposed 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.
effective means of
Introduction
The success of reinforcement learning (RL) algorithms large-scale, complex tasks depends on forming usefuu 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 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 the environment implicitly, typically via gradient descent the additional learning objectives. As such, while auxiliary tasks continue to play an important role improving the performance deep RL algorithms, our understanding of the effects of auxiliary tasks on representations in RL is still in its infancy.
In contrast the implicit representation shaping of auxiliary tasks, separate line of work on be- havioural 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 RI. agents [Gelada al., 2019, Zhang et al., 2021, Agarwal et al., 2021a]. However, in practice behavioural metrics are difficult to estimate from both statistical and computa- tional 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 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
Equaa contribution. Correspondence to Pablo Samuel Castro: pse google.com.
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

00.30
10M Human Normarzed
200
150
0.20
0.15
50
10
0.10
0.05
25
75
180
125
150
175
200
80
150
240
320
400
Number of Fromes (n millions
Number of Frrrrmes bx 1000)
Figure I: Interquantile Mean human normalized scores of all the agents and losses on the ALE suite (left) and on the DM-Control suite (right), both run with five independent seeds for each agent and environment. In both suites MICo provides a clear advantage.
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 al., 2018], in which the performance of 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 al., 2018], 8, 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 right).
Background
Before describing the details of our contributions, we give brief overview of the required background in reinforcement learning and bisimulation. We provide more extensive background in Appendix A.
Reinforcement learning. We consider Markov decision process (X,A, P,r) defined by finite state space X, finite action space A, transition kernel P: X x A (X), reward function X' A - R. and discount factor y [0, 1). For notational convenience we will write Pe and r for transitions and rewards, respectively. Policies are mappings from states to distributions (AA)* and induce a value function V X I defined via the recurrence: (*) == Eawa(e) (xx [1 + B [V* (xr')] In RL we are concerned with finding the optimal policy π = arg max (A)* V from interaction with sample trajectories with an MDP, without knowledge of P or (and sometimes not even X',, and the optimal value function V* induced by n²"
over actions:
State similarity and bisimulation metries. 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 z, E X to be close if their immediate rewards are similar, and the transition dynamics 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' the unique fixed- point the operator Tx : :M(X) - M(X). symmetric and satisfies the triangle inequality) is the set of pseudometrics on X. given by (x(d) (z, y) maxae.A(11 - rgl + yWa(PP (P P. Here, Wa is the Kantorovich distance (also known as the Wasserstein distance) over the set of distributions (X) with base distance d, defined by Wa(n, v) = inf x~ E[d(X, Y)). for all μr, E (X)), where the infimum is taken over all couplings of (X,Y) with the prescribed marginals [Villani, 2008].
where M(X)= = {d E [0. x
The mapping Tk is a 7-contraction on M(X) under the - norm [Ferns et al., 2011], and thus by standard contraction mapping arguments analogous those used study value iteration, 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:
VV(x) vx) < dz
Vx,y X.
Representation learning in RL. In large-scale enviromments, is infeasible to express value func- tions directly as vectors in X Instead, RL agents must approximate value functions more concise manner, by forming representation of the environment, that is, feature embedding Φb + - and predicting state-action values linearly from these features. Representation learning 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 side-effect of learning to solve auxiliary tasks.
Since bisimulation metries capture additional information about the MDP in addition that sum- marised 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 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.
3 Advantages and limitations of the bisimulation metric
The bisimulation metric d' is 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 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 opera- tor Tx is contraction mapping. The map Tx contracts at rate with respect to the L norm on M, and therefore obtaining an e-approximation of d under this norm requires Olog (1 (1/e)/ log (1/7)) applications of T'x to an initial pseudometric dg- The cost of each application of Tx is dominated by the computation of 2 ||| Wa distances for distributions over X, each costing ((12212.5) in theory
[Lee and Sidford, 2014], and ((x)))) in practice [Pele and Werman, 2009, Guo et al., 2020a, Peyré and Cuturi, 2019]. Thus, the overall practical cost is ( ||A| log(s)/
log(>)).
(ii) Bias under sampled transitions. Computing Tk requires access to the transition probability distributions Pe for each (zr, a) E X'x 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 (x) - (v) I dx, y) for arbitrary policies и E ПΙ; a simple example 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 computa- tional difficulty of computing this metric. Prop- erty (ii) illustrates the problems associated with attempting to move from operator-based com- putation to online, sampled-based computation of the metric (for example, when the environ- ment dynamics are unknown). Finally, prop- erty (iii) shows that even if the metric is com- putable exactly, the information it yields about the MDP may not be practically useful. Al- though π m-bisimulation (introduced by Castro
L
b
a
y
b
+
k
Figure 2: MDP illustrating that the upper bound for
any π not generally satisfied. Here, d' d' (I. y)
(1 - 7)- but for IT m(b/z)) = 1.T((al)) = 1, we
have |V" (x) V (y) = k(1 -
[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., 3.21]. Taken together, these three properties motivate the search for metric without these shortcomings, which can be used in combination with deep reinforcement learning.
3

---

The MICo distance
We now present 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, make several modifications to the operator Tи 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 and Pjj (ii) to deal with lack of connection to non-optimal policies, we consider an on-policy variant of the metric, pertaining chosen policy и ((A)x. This leads to the following definition.
the independent coupling;
Definition 4.1. Given π E ((AA), the MICo update operator Ths- : Rxxx 7 xxx* is:
(TT = 1r 2
IE s' V'-PT
(U(x',
for all U X x + R, with r = EaEA x ( and Pr = E x(a for all a E X.
As with the bisimulation operator, this can be thought of as encoding desired properties of notion of similarity between states in self-referential manner; the similarity of two states æ, E X should be determined by the similarity of the rewards and the similarity the states they lead to. Proposition 4.2. The operator T is contraction mapping on xxx with respect to the _ norm. Proof. See Appendix B.
The following corollary now follows immediately from Banach's 's point theorem and the com- pleteness of IR. xX under the 0 norm.
Corollary 4.3. The MICo operator Ty has a unique fixed point U RXXXX and repeated applica- tion of T3 any initial function U Xxx converges to U
Having defined new operator, and shown that it has corresponding fixed-point, there 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 about the underlying MDP?
4.1 Addressing the drawbacks of the bisimulation metric
We introduced the MICo distance as a means of overcoming some the shortcomings associated with the bisimulation metric, described in Section 3. In this section, provide 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 lifted MDP. This result crucial for much of the analysis that follows, so we describe the proof in full detail.
Lemma 4.4 (Lifted MDP). The MICo operator T5 is the Bellman evaluation operator for an auxiliary MDP.
Proof. Given the MDP specified by the tuple A, R), we construct an auxiliary MDP (,,A, R). by taking the state space to be X X2, the action space to be Á A= A2, R the transition dynamics to be given by Pla.)) ((x, v)) = P. (x) (v) for all (a, u), (u, u) E X², a, b E A, and the action-independent rewards to be R(z.,) = 1r - r.. for all x, E X. The Bellman evaluation operator 7 for this auxiliary MDP discount rate y under the policy ((a, b/x, v) x (blp))(bl))) is given by (for all U X and (I, E X' x X)):
(T*U)(E,)) Rlaw +?
Σ
plab ((x', (2,))
= 1r r
Σ ((())))((/)) ,/ -TTUUe,)
(x' (xx))
Remark 4.5. Ferns and Precup [2014] noted that the bisimulation metric can interpreted as the optimal value function in a related MDP, and that the functional Tx of Tx can be interpreted as 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 Wa distance in the definition of T 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 follows.
Proposition 4.6 (MICo computational complexity). The computational complexity of computing an e-approximation in L to the MICo metric is O(x'' log(e)/ log(m)))
Proof. Since, by Proposition 4.2, the operator 75 a 7-contraction under L we require ((log(e))log())) applications of the operator to obtain an e-approximation in L Each itera- tion of value iteration updates |2|2 table entries, and the cost of each update is ())x) leading to an overall cost of ((12]* log(s)/ (o))))
In contrast to the bisimulation metric, this represents computational saving of O(Ix1), which arises from the lack of a need to solve optimal transport problems over the state space in computing the MICo distance. There is further saving of ([-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 TMr 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 current estimate U² of the fixed point of Ty and a pair of observations (z, a, x'), F'D0) generated under ıT, we can define a new estimate Ue+1 via
U141(2,V) + (1 ()) ()r + >U1((',))) )'))
(3)
for all other state-pairs (ū, 5) A (a,, y). for some sequence of stepsizes (es(1, y) 1 1.2 0, (x, v) E 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 not satisfied.
and U4+1 (5,0) 1
5) - U.(5,)
Proposition 4.7. Suppose rewards depend only on state, and consider the sequence of esti- mates (U))20, with Uo initialised arbitrarily, and U241 updated from U, via a pair of transitions (I2, at, TT, x[). (yr, bt, v() as in Equation (3). If all state-pairs tuples are updated infinitely often, and stepsizes for these updates satisfy the Robbins-Mont Monno conditions. Then U, - U almost surely.
Proof. Under the assumptions of the proposition, the update described is exactly a TD(O) 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 π E (A) and states x. y E X,, we have (V*(2)) (F) Uz.
Proof. We apply coinductive argument [Kozen, 2007] to show that [V*(<> VT( S U(x, v) for all x. y E X, for some U E XXXX symmetric in its two arguments, then we also have (V*(x) Va [y) & (T, y) for all x, E X. Since the hypothesis holds for the constant function (zz v) 2 m ((z.a))((1 7), and Ty contracts around U", the conclusion then
5

---

follows. Therefore, suppose the hypothesis holds. Then we have
(= = r
PT(d')V(') x'
Σ
VEE
s 5 r..
E P(N)) -V-(())
w'y'EX
< lr ryl +7
Σ PECR) (TT(((())
x' x'Ey'Ex
By symmetry, VT (*(x) < (Z,, v). as required.
4.2 Diffuse metries
To characterize the nature of the fixed point UH, we introduce the notion of a diffuse metric.
Definition 4.9. Given a set X, function d : X x X 7 R is a diffuse metric if the following axioms hold: (i) d(x, y) 2 for any y E X: (ii) d(x.y) d(y. x) for any x. y. E X:: (iii)
d(z, y) I d(z, 2) + d(u, 2) Vx, y. 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-regularte also do not satisfy the axiom of zero self- elf-distance.
optimal transport distances [Cuturi, 2013], which
An example of diffuse metric is the Lukaszyk- Lukaszyk-Karoowsk distance [Lukaszyk, 2004], which is used in the MICo metric as the operator between the next-state distributions. Given a diffuse metric space (X' p). the Mikaszyk-Karmowsee distance dfx is diffuse metric on probability measures on X' given by (f,(u, = Ea--,y-><l(zz v)]. This example demonstrates the origin of the name diffuse metrics: the non-zero self distances arises from a point being spread across probability distribution. In terms of the Lukaszyk-Karmowak Karmowski distance, the MICo distance can be written as the fixed point U (xx y) = rr= r + (UK(UT) PT). This characterisation leads to the following result.
Proposition 4.10. The MICo distance is 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 (Ux))20 in X by Uo(2, = 0 for all v,y E X, and Uk+1 = T3 Ux for each 2 0. Recall that by Corollary 4.3 that Us - U We will show that each UR 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 Ug. For the inductive step, assume that for some i 2 (z, < U2(x, 2) + U2(2, y) for all x, y, 2 E X. Now for any x. y, 2 E X, we have
U2+i(E,v) Ir=h r.l xBx''Pg P' (U,(X',Y)) Y
S 1r=h r#l + Ir - ryl + TEx y P UX Z') + U Y')] = U1-1,1, 2) + Uk+1(2,8)
It is interesting to note that a state x E 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 (² x) > 0, and U U²(x, x) # U (y. v) for distinct states x, y E X.. See the appendix for further discussion of diffuse metrics and related constructions.
5 The MICo loss
The impetus of our work is the development of principled mechanisms for directly shaping the representations used by RI. agents as to improve their learning. In this section we present a novel loss based on the MICo update operator Ty, 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

---

Lancs (o(x).((v))
(u(x))
8 (xx
((s(a))
((x)
+lv)
(n v)
(V*()))) (7 (07(),)) (0%
Figure 3: Left: Illustration of network architecture for learning MICo; Right: The projection of MICo distances onto representation space.
care needs be taken in how these distances are learnt; indeed, traditional mechanisms for measuring distances between representations (e.g. Euclidean and cosine distances) are ccometrically-basee 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 w, respectively: (x,-) Ve (ow (note that this can be the critic in an actor-crie actor-critic algorithm such as SAC). We will refer to (.) as the representation of state x and aim make distances between representations match the MICo distance; we refer to deg as the value approximator. We define the parameterized representation distance, U.. as an approximant to U
(x, =o (Z( (z) + ( + (0( (z), (3))
2
where (())(a)) (U)) is the angle between vectors &w (1) and (v) and 3 is scalar (in our results we use 3 = 0.1 but present results with other values 3 in the appendix).
Based on Equation (2), our learning target is then TV (ra, x', Ty. v/) Ira - +ryl+ (2', v/, where a separate copy of the network parameters that synchronised with at infrequent intervals. This is 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
LMICo(w) =Bqu.ra.a" (onryy'" (T (ra, ,yy.n/.y///))))))) Uw (I,
where (xr, T,, x²) and (y, Ty, v/) are pairs of transitions sampled from the agent's replay buffer. We combine LMICO with the defferea-difference loss LTD of any RL agent as (1-a) aLvico. where a E (0, 1). Each sampled mini-batch is used for both MICo and TD losses. Figure (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 parametrised through Us defined above, which is constituted by several distinct terms. This appears to leave a question as to how the representations &w (x: and (y). as Euclidean vectors, are related to another when the MICo loss is minimised. Careful inspection of the form of Uw y) shows that the (scaled) angular distance between (x) and d.. can be recovered from U., by subtracting the learnt approximations right). We therefore define the reduced MICo distance IT"", which encodes the distances enforced between the representation vectors dhw and (w),, by:
the self-distances U (2x and U ("(y, y) (see Figure
())) (v)) ITU (x, ("( UU 1 U
In the following section we investigate the following two questions: (1) How informative of V is ? and (2) How useful are the features encountered by - 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.
5.1 Value bound gaps
Although Proposition 4.8 states that we have (V*(2) V*(()) S Ux, y), we do not, in general, have the same upper bound for I1 (z, y) as demonstrated by the following result.
7

---

Metric nu=
Lineat regreeegeeeee aggression, avg. error
Mow Wiat doie dag
2.000
PVF
9.066
8.004
0.002
0.000
Number
22
Number of features
105
Figure 4: Left: The gap between the difference in values and the various distances for Garnet MDPs with varying numbers of actions frepresented by circle sizes); Right: Average error when performing linear regression varying numbers of features in the four-rooms GridWorld, averaged over 10 runs; shaded areas represent 95% confidence intervals.
Lemma 5.1. There exists MDP with x; y X, and क E II where [V*(x) V (T(() I1 v).
Proof. Consider a single-action MDP with two states (x and y) where y is absorbing, x transitions with equal probability to and y. and a reward of 1 is received only upon taking action from state x. There only one policy for this MDP which yields the value function V(x) Rs 1.8 and Vyy) O.
The MICo distance gives U (x., .)) R 1.06, Uxx v) a 1.82, and (v, y) = 0. while the reduced MICo distance yields (x, x) I1U (y. 3) = O, and IU (z, y) ~ 1.29 < VV - (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. and the number of actions nA, Garnet(n.x: T2,)) is generated as follows: 1. The branching factor by,a of each transition pa is sampled uniformly from [1: nx)) 2. bz.a states picked uniformly randomly from X' and assigned a random value in [0, 1}: these values are then normalized to produce a proper distribution Pa.
3. Each r is sampled uniformly in [0, 1].
For each Garnet(nx nAA) we sample 100 stochastic policies {m.) and compute the average gap: I00| E. En.y d(I,, w) (x) V**( (v)l, where d stands for any of the considered metrics. Note
measuring the signed difference, we 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 ( and r-bisimulation. This suggests that the resulting representations remain informative of value similarities.
5.2 State features
In order to investigate the usefuleness of the representations produced by I1 we construct state features directly by using the computed distances to project the states into lower-dimensional space with the UMAP dimensionality reduction algorithm [McInnes et al., 20 We then apply linear regression of the true value function V against the features 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 r-bisimulation [Castro, 2020]. We present our results on the well-known four-rooms GridWorld [Sutton et al., 1999] Figure 4 (right) and provide results on more environments in the appendix. Despite the independent couplings, 110 performs on par with π r-bisimulation, which optimizes over all couplings.
6 Large-scale empirical evaluation
Having developed greater understanding of the properties inherent 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 the Dopamine library [Castro al., 2018]: DQN [Mnih et al., 2015], Rainbow [Hessel et al., 2018], QR-DQN [Dabney et al., 2018b].
Note that since UMAP expects metric, it is del-nnd with the diffuse metric U

---

and IQN [Dabney et al., 2018a], using mean squared error loss to minimize LID for DQN (as suggested by Obando-Ceron and Castro [2021]). Given the state-of-the-ar results demonstrated by the Munchausen-IQQN (M-IQN) agent [Vieillard et al., 2020], we also evaluated incorporating our loss into M-IQN, For all experiments we used the hyperparameter settings provided with Dopamine. We found that a value of a = 0.5 worked well with quantile- hasele agents (QR-DQN, IQN, and M-IQN), while a value of a = 0.01 worked well with DQN and Rainbow. We hypothesise that the difference in scale of the quantile, categorical, and non-distributional Ioss functions concerned leads to these distinct values of a 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 (left), using the interquantile metric (IQM), proposed by Agarwal et al. [2021b] as 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) [Hаarnoja 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 a 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 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.
7 Related Work
Bisimulation metrics were introduced for MDPs by Ferns et al. [2004], and have been extended in 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., 20,3a,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].
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 homomor- phisms [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] l performed exhaustive analysis of the continuity properties, relative to functions of interest in RL, of 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- regulaaieee 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.,
*Given that the authors of M- M-IQN had implemented their agent in TensorFlow (whereas our agents are in JAX), we have reimplemented M- M-IQN in JAX and run 5 independent runs (in contrast to the 3 run by Vieillard al. [2020].

---

2017, Franç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].
Conclusion
In this paper, have introduced the MICo distance, notion of state similarity that can be learnt at scale and from samples. We have studied the theoretical properties of MICo, and proposed 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 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 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 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.
Acknowledgements
The authors would like 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 thank the reviewers (both ICML'21 and NeurlPS'21) for helping make this paper better.
10

---

References
Alessandro Abate. Approximation metries based on probabilistic bisimulations for general state-space Markov processes: A survey. Electr. Notes Theor. Comput. Sci., 297:3-25, 2013.
Rishabh Agarwal, Marlos C. Machado, Pablo Samuel Castro, and Mare G. Bellemare. Contrastive behavioral similarity embeddings for generalization reinforcement learning. In International Conference on Learning Representations (ICLR), 2021a.
Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron Courville, and Mare G Bellemare. Deep reinforcement learning the edge of the statistical precipice. Advances in Neural Information Processing Systems (NeurIPS), 2021b.
Philip Amortila, Marc G Bellemare, Prakash Panangaden, and Doina Precup. Temporally extended metrics for Markov decision processes. In AAAI Workshop Safe AI, 2019.
T. W. Archibald, K. I. M. McKinnon, and L. C. Thomas. On the generation of Markov decision processes. The Journal of the Operational Research Society, 46(3):354-36) 1995.
Giorgio Bacci, Giovanni Bacei, Kim G Larsen, and Radu Mardare. Computing behavioral distances, compositionally. In International Symposium on Mathematical Foundations of Computer Science (MFCS), 2013a.
Giorgio Bacci, Giovanni Bacci, Kim G Larsen, and Radu Mardare. On-the-fly exact computation of bisimilarity distances. In International Conference on Tools and Algorithms for the Construction and Analysis of Systems (TACAS), 2013b.
André Barreto, Will Dabney, Rémi Munos, Jonathan J Hunt, Tom Schaul, Hado Van Hasselt, and David Silver. Successor features for transfer in reinforcement learning. In Advances in Neural Information Processing Systems (NIPS), 2017.
Mare G. Bellemare, Yavar Naddaf, Joel Veness, and Michael Bowling. The Arcade Learning Environment: An evaluation platform for general agents. Journal of Artificial Intelligence Research, 47:253-277 June 2013.
Marc G Bellemare, Will Dabney, and Rémi Munos. A distributional perspective reinforcement learning. International Conference on Machine Learning (ICML), 2017.
Dimitri P. Bertsekas and John N. Tsitsiklis. Meuro-Dynamic Programming. Athena Scientific, 1996.
Pablo Samuel Castro. Scalable methods for computing state similarity in deterministic Markov Decision Processes. In AAAI Conference on Artificial Intelligence, 2020.
Pablo Samuel Castro and Doina Precup. Using bisimulation for policy transfer in MDPs. In International Conference Autonomous Agents and Multiagent Systems (AAMAS), 2010.
Pablo Samuel Castro, Subhodeep Moitra, Carles Gelada, Saurabh Kumar, and Marc G. Bellemare. Dopamine: A Research Framework for Deep Reinforcement Learning. 2018.
Lenaic Chizat, Pierre Roussillon, Flavien Léger, François-Xavis Xavier Vialard, and Gabriel Peyré. Faster wasserstein distance estimation with the sinkhorn divergence. Advances in Neural Information Processing Systems (NeurIPS), 2020.
Gheorghe Comanici, Prakash Panangaden, and Doina Precup. On-the-fly algorithms for bisimulation metrics. In International Conference on Quantitative Evaluation of Systems (QEST), 2012.
Gheorghe Comanici, Doina Precup, and Prakash Panangaden. Basis refinement strategies for linear value function approximation in MDPs. In Advances in Neural Information Processing Systems (NIPS), 2015.
Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in Neural Information Processing Systems (NIPS), 2013.
Will Dabney, Georg Ostrovski, David Silver, and Rémi Munos. Implicit quantile networks for distributional reinforcement learning. In International Conference on Machine Learning (ICML), 10.
11

---

Will Dabney, Mark Rowland, Mare Bellemare, and Rémi Munos. Distributional reinforcement learning with quantile regression. In AAAI Conference on Artificial Intelligence, 2018b.
Peter Dayan. Improving generalization for temporal difference learning: The successor representation. Neural Comput., 5(4):613-62) July 1993.
Josée Desharnais, Vineet Gupta, Radhakrishnan Jagadeesan, and Prakash Panangaden. Metrics for labeled Markov systems. In International Conference on Concurrency Theory (CONCUR), 1999.
Josée Desharnais, Vineet Gupta, Radhakrishnan Jagadeesan, and Prakash Panangaden. A metric for labelled Markov processes. Theoretical Computer Science, 31 ((33)3323-354, June 2004.
Kilian Fatras, Younes Zine, Rémi Flamary, Rémi Gribonval, and Nicolas Courty, Learning with minibatch Wasserstein: asymptotic and gradient properties. In International Conference on Artificial Intelligence and Statistics (AISTATS), 2020.
Kilian Fatras, Younes Zine, Szymon Majewski, Rémi Flamary, Rémi Gribonval, and Nicolas Courty. Minibatch optimal transport distances; analysis and applications. arXiv, 2021.
William Fedus, Carles Gelada, Yoshua Bengio, Marc G Bellemare, and Hugo Larochelle. Hyperbolic discounting and learning over multiple horizons. arXiv, 2019.
Norm Ferns and Doina Precup. Bisimulation metrics are optimal value functions. In Conference on Uncertainty in Artificial Intelligence (UAI), 2014.
Norm Ferns, Prakash Panangaden, and Doina Precup. Metries for finite Markov decision processes. In Conference on Uncertainty in Artificial Intelligence (UAI), 2004.
Norm Ferns, Pablo Samuel Castro, Doina Precup, and Prakash Panangaden. Methods for comput- ing state similarity in Markov decision processes. Conference on Uncertainty in Artificial Intelligence (UAI), 2006.
Norm Ferns, Prakash Panangaden, and Doina Precup. Bisimulation metrics for continuous Markov decision processes. SIAM Journal on Computing, 1066)1662-1114, 2011.
Norman Ferns, Prakash Panangaden, and Doina Precup. Metrics for Markov decision processes with infinite state spaces. In Conference on Uncertainty in Artificial Intelligence (UAI), 2005.
Vincent Frangois-Lavet, Yoshua Bengio, Doina Precup. and Joelle Pineau. Combined reinforcement learning via abstract representations. In AAAI Conference on Artificial Intelligence, 2019.
Carles Gelada, Saurabh Kumar, Jacob Buckman, Ofir Nachum, and Mare G Bellemare. DeepMDP: Learning continuous latent space models for representation learning. In International Conference on Machine Learning (ICML), 2019.
Aude Genevay, Gabriel Peyré, and Marco Cuturi. Learning generative models with sinkhorn diver- gences. In International Conference Artificial Intelligence and Statistics (AISTATS), 2018.
Robert Givan, Thomas Dean, and Matthew Greig. Equivalence notions and model minimization in Markov decision processes. Artificial Intelligence, 147(1-2):163-223, 2003.
Wenshuo Guo, Nhat Ho, and Michael I. Jordan. Fast algorithms for computational optimal transport and Wasserstein barycenter. In International Conference on Artificial Intelligence and Statistics (AISTATS), 2020a.
Zhaohan Daniel Guo, Bernardo Avila Pires, Bilal Piot, Munos, and Mohammad Gheshlaghi Azar. Bootstrap latent- iitent-predictive representations for multitask reinforcement learning. International Conference on Machine Learning (ICMI.), 2020b.
Grill, Florent Altché, Rémi
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. International Conference Machine Learning (ICML), 2018.
12

---

Matteo Hessel, Joseph Modayil, Hado van Hasselt, Tom Schaul, Georg Ostrovski, Will Dabney, Dan Horgan, Bilal Piot, Mohammad Azar, and David Silver. Rainbow: Combining Improvements in Deep Reinforcement learning. In AAAI Conference on Artificial Intelligence, 2018.
Peter J. Huber. Robust Estimation of Location Parameter. The Annals of Mathematical Statistics, 35(1):73 - 101, 1964.
Max Jaderberg, Volodymyr Mnih, Wojciech Marian Czarnecki, Tom Schaul, Joel Z Leibo, David Silver, and Koray Kavukcuoglu. Reinforcement learning with unsupervised auxiliary tasks. In International Conference on Learning Representations (ICLR), 2017.
Nicholas K Jong and Peter Stone. State abstraction discovery from irrelevant state variables. In International Joint Conference on Artificial Intelligence (IJCAI), 2005.
Dexter Kozen. Coinductive proof principles for stochastic processes. Logical methods in computer science, 3, November 2007.
Kim G Larsen and Ame Skou. Bisimulation through probablistic testing. Information and Computa- tion, 1-21- 1991.
Charline Le Lan, Mare G. Bellemare, and Pablo Samuel Castro. Metrics and continuity in reinforce- ment learning. In AAAI Conference on Artificial Intelligence, 2021.
Yin Tat Lee and Aaron Sidford. Path finding methods for linear programming: Solving linear programs in Ö(yrank iterations and faster algorithms for maximum flow. In IEEE Annual Symposium on Foundations of Computer Science (FOCS), 2014.
Lihong Li, Thomas J Walsh, and Michael L. Littman. Towards unified theory of state abstraction for MDPs. In International Symposium on Artificial Intelligence and Mathematics (ISAIM), 2006.
Marlos C Machado, Marc G Bellemare, Erik Talvitie, Joel Veness, Matthew Hausknecht, and Michael Bowling. Revisiting the Arcade Learning Environment: Evaluation protocols and open problems for general agents, Journal of Artificial Intelligence Research, 61:523-566 2018.
Sridhar Mahadevan and Mauro Maggioni. Proto-value functions: A Laplacian framework for learning representation and control in Markov decision processes. Journal of Machine Learning Research, 8:2169-2231, December 2007.
Steve Matthews. Partial metric topology. Annals of the New York Academy of Sciences, 728(1): 183-197, 1994.
Andrew Kachites McCallum. Reinforcement Learning with Selective Perception and Hidden State. PhD thesis, The University of Rochester, 1996.
Leland McInnes, John Healy, Nathaniel Saul, and Lukas Grossberger. UMAP: Uniform manifold approximation and projection. The Journal of Open Source Software, 3(2)):861, 2018.
R. Milner. Communication and Concurrency. Prentice-Hall, 1989.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Marc G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, Stig Petersen, Charles Beattie, Amir Sadik, Ioannis Antonoglou, Helen King, Dharshan Kumaran, Daan Wierstra, Shane Legg, and Demis Hassabis. Human-level control through deep reinforcement learning. Nature, 518(7540):529-533 533, 2015.
Ted Moskovitz, Michael Arbel, Ferene Huszar, and Arthur Gretton. Efficient Wasserstein natural gradients for reinforcement learning. In International Conference on Learning Representations (ICI.),, 2021.
Johan S Obando- Ceron and Pablo Samuel Castro. Revisiting rainbow: Promoting more insightful and inclusive deep reinforcement learning research. In International Conference on Machine Learning (ICML), 2021.
13

---

Aldo Pacchiano, Jack Porker-Holdes Yunhao Tang, Krzysztof Choromanski, Anna Choromanska, and Michael Jordan. Learning to score behaviors for guided policy optimization. In International Conference on Machine Learning (ICML), 2020.
Ofir Pele and Michael Werman. Fast and robust earth mover's distances. In IEEE International Conference on Computer Vision (ICCV), 2009. 2009.
Gabriel Peyré and Marco Cuturi. Computational optimal transport. Foundations and Trends in Machine Learning, 11(5-6):355-600 607, 2019.
Bilal Piot, Matthieu Geist, and Olivier Pictquin. Difference of convex functions programming for reinforcement learning. In Advances in Neural Information Processing Systems (NIPS). 2014.
Balaraman Ravindran and Andrew G. Barto. SMDP homomorphisms: An algebracc approach to abstraction in semi-Markov decision processes. In International Joint Conference on Artificial Intelligence (IJCAT), 2003.
Sherry Shanshan Ruan, Gheorghe Comanici, Prakash Panangaden, and Doina Precup. Representation discovery for MDPs using bisimulation metrics. In AAAI Conference on Artificial Intelligence. 2015.
Anirban Santara, Rishabh Madan, Balaraman Ravindran, and Pabitra Mitra. ExTra: Transfer-guie exploration. In International Conference on Autonomous Agents and Multiagent Systems (AAMAS), 2019.
Richard S Sutton, Joseph Modayil, Michael Delp, Thomas Degris, Patrick M Pilarski, Adam White, and Doina Precup. Horde: A scalable real-time architecture for learning knowledge from unsuper- vised sensorimotor interaction. In International Conference on Autonomous Agents and Multiagent Systems (AAMAS), 2011.
R.S. Sutton, D. Precup, and S. Singh. Between MDPs and semi-MDPs: A framework for temporal abstraction in reinforcement learning. Artificial Intelligence, 112:181-211, 1999.
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmaleki, Josh Merel, Andrew Lefraneq, Timothy P. Lillicrap, and Martin A. Riedmiller. Deepmind control suite. arXiv, 2018
Jonathan Taylor. Lax probabilistic bisimulation. Master's 's thesis, McGill University, 2008.
Jonathan Taylor, Doina Precup, and Prakash Panagaden. Bounding performance loss in approximate
MDP homomorphisms. In Advances in Neural Information Processing Systems (NIPS). 2009.
Nino Vieillard, Olivier Pietquin, and Matthieu Geist. Munchausen reinforcement learning. In
Advances in Neural Information Processing Systems (NeurIPS), 2020.
Cédric Villani. Optimal Transport. Springer- Verlag Berlin Heidelberg, 2008
Denis Yarats, Amy Zhang, Ilya Kostrikov, Brandon Amos, Joelle Pineau, and Rob Fergus. Improving sample efficiency in model-free reinforcement learning from images. arXiv, 2019.
Amy Zhang, Rowan McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Invariant rep- resentations for reinforcement learning without reconstruction. In International Conference on Learning Representations (ICLR), 2021.
Szymon Lukaszyk. A new concept of probability metric and its applications in approximation of scattered data sets. Computational Mechanics, 33:299-304, 03 2004.
14


Reviews:
[{"review_id": "wu2pF40ADzv", "reviewer": "Reviewer_vaui", "summary": "The paper presents MICo - a behavioural metric to learn representations for deep RL agents and shows its effectiveness - both theoretically and empirically (on the Arcade Learning Environment benchmark).", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "## Strengths\n\n1. The paper tackles an interesting, relevant and impactful problem in RL.\n2. The paper is very well written. I am not much familiar with the convergence proof techniques etc but I did not have much trouble understanding the paper.\n3. The paper provides both theoretical and empirical results. The empirical results show that MICo can be combined with a large number of baselines.\n\n## Suggestions for improvement\n\n1. Figure 4 is very hard to parse. Maybe the lines could be thinner and \"....\"can be used in place of \"----\". Alternatively, consider moving some of the results to the appendix. Also consider adding error bars to show the variance.\n2. Why arent other bisimulation metric based methods reported in the experiment section? I understand they have additional assumptions baked in (line 82) but some of them should still be evaluated and reported results with.\n\n\nAn additional comment (Note that I am not holding this as a weakness for the paper): If the authors have resources, they should consider reporting the results over at-least 10 seeds. They should also consider releasing the metrics as a csv/json file. I expect this paper to be well-cited and several other future works to use this as a baseline. ", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Limitations And Societal Impact are addressed.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}, {"review_id": "qwBecbx7GKe", "reviewer": "Reviewer_apPQ", "summary": "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.", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "**Originality**\n\nThey suggest the new operator for defining the MICo distance. While it can be thought of as a simplified version of the bisimulation operator with the eliminated need for the Wasserstein metric, they describe how it differs from previous approaches, importantly addressing some issues of the bisimulation operator. Also, they suggest the learned approximation to the MICo distance for practical use.\nI think the proposed methods and derivations are meaningfully novel.\n\n**Quality**\n\nI find the theory up to Sec. 4 (including the MICo distance with its operator $T_M^\\pi$, equivalence to the Bellman operator in the auxiliary MDP, the online approximation of the solution, upper-bounding of the value function difference, etc.) technically sound.\n\nIn Sec. 5, I think the definition of $U_{\\omega}$ (L259) is somewhat abrupt. Despite the usefulness of $\\Pi U^\\pi$ shown throughout Sec. 5 and its connection to $U_{\\omega}$, I believe addressing why the parametrized approximation of $U^\\pi$ should be (or at least can reasonably be) the form at L259 will be a good addition to the paper.\n\nRegarding the experiments in Sec. 6, while it is extensive that the various RL algorithms with or without the MICo loss are evaluated on the large number of Atari environments, including some additional comparison with other baseline representation shaping methods would be great.\n\n**Clarity**\n\nThe paper is overall easy to follow and well-organized.\n\nTypo \"oppoosed\" at 328.\n\n**Significance**\n\nThe theoretical results in this work, despite providing the diffuse metric not a metric, seem to have a good chance to be used in future literatures, thanks to the nice properties it has (bounding on an arbitrary policy's value function gaps, the approximation schemes, and the reduced computational complexity).", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "I think the authors have fairly covered the limitations of their work.\nPlease see my main review, as well.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "6"}, "extra_text": {}}, {"review_id": "P3eANmUxIiy", "reviewer": "Reviewer_dAxX", "summary": "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.", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Overall, I think the paper is of high quality. The proposed method is well motivated and theoretically sound. The structure and writing are very clear. Nevertheless, I have some questions and hope the authors can help address.\n\n* The flow of the paper is to first identify the drawbacks of the bisimulation metrics and then propose MICo to resolve them. However, it seems to me that it is more like proposing a new behavioural distance that \"bypasses\" the drawbacks of the bisimulation metrics. \"Addressing the drawbacks\" sounds like MICo can achieve the exactly same results as bisimulation metrics without suffering the drawbacks. I guess this is not what the authors want to convey. I would like to see some more discussion on comparing MICo and bisimulation metrics. For example, with the advantages of MICo being said, what does it lose (or fail to capture) compared to bisimulation metrics? In addition to Section 5.2, it would be better to add some visualizations of the MICo and bisimulation metrics (on a small grid world env), so that readers may get a more intuitive understanding.\n* I do not quite agree with the description in Line 85 that explicitness means not using additional network parameters. Whether or not using additional network parameters is more on the implementation side.\n* The fixed point $U^\\pi$ is defined w.r.t. a policy $\\pi$. However, when learning together with TD loss (as in Atari experiments), the policy keep changing. I am wondering how to deal with such non-stationarity.\n\nPossible typos (please correct me if I am wrong):\n\n* Line 163: \"functional $T_K$ of $T_K$\"\n* Line 59: \"$d$ symmetric\", missing \"is\"?", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "6"}, "extra_text": {}}, {"review_id": "-Bu4d77s7p4", "reviewer": "Reviewer_PkRV", "summary": "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.\n", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "### Strengths\n- The paper is very relevant to the representation learning / (deep) RL community.\n- The proposed MICo distance is a natural extension of the existing on-policy bisimulation metric, which removes the restriction to only deterministic MDPs. In addition, the theoretical analyses and proofs in Section 4 and Section 5 are novel to the best of my knowledge.\n- I appreciate the paper's technicality level. The authors start by mentioning the inherent limitations of the bisimulation metric, propose the MICo distance, and then provide theoretical results to show how it overcomes these limitations. The results of the paper are significant, and the analyses and proofs are thorough and accurate.\n- The paper is very well written. I enjoyed reading the paper.\n\n### Weaknesses\n- In my opinion, the lack of connection to the optimal policy is actually a disadvantage of MICo against the bisimulation metric. Imagine when the observation contains a lot of rich but task-irrelevant information, using the bisimulation metric will result in representations that are informative of the value similarities under the *optimal* policy. In other words, the resulting representations can successfully capture task-relevant information while filtering out redundant information.\n- Using MICo as an additional loss only improves the performance of M-IQN and Rainbow. The results with and without MICo in other methods are fairly the same. Moreover, in some Atari games, using MICo makes the performance drop significantly. Why is this the case?\n- While the authors have experimented MICo on many Atari games and with many different RL methods, the paper is lacking comparisons with other distances. I would like to see at least the comparison between MICo and the bisimulation metric used in [1].\n\n### Other comments and questions\n- Between lines 259 and 260, how did the authors come up with this specific way of parameterizing the MICo distance? There must be other parameterizations that are positive and symmetric.\n- The first equation between lines 161 and 162, there should be a sum over all actions $(a,b)$.\n\nI am willing to raise my score if these concerns are addressed.\n\n[1] Zhang, Amy, et al. \"Learning invariant representations for reinforcement learning without reconstruction.\" arXiv preprint arXiv:2006.10742 (2020).", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors have adequately addressed the limitations and potential negative societal impact of their work.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "16"}, "extra_text": {}}]

Author rebuttal:

Meta-review: This seems like a solid paper that develops new approach to the interesting problem of inducing bias in learning through a notion of state similarity.  The notion overcomes some limitations of previous notions that were based on bisimulation concepts.  The reviewers are aligned in their ratings of the paper.
Decision: Accept (Poster)