Towards Robust Bisimulation Metric Learning
Authors: Mete Kemertas
Tristan Ty Aumentado-Armstrong
Venue: NeurIPS 2021 Poster (2021)
Keywords: reinforcement learning
bisimulation
state abstraction
sparse rewards
state similarity metrics
representation learning
continuous control

Abstract: Learned representations in deep reinforcement learning (DRL) have to extract task-relevant information from complex observations, balancing between robustness to distraction and informativeness to the policy. Such stable and rich representations, often learned via modern function approximation techniques, can enable practical application of the policy improvement theorem, even in high-dimensional continuous state-action spaces. Bisimulation metrics offer one solution to this representation learning problem, by collapsing functionally similar states together in representation space, which promotes invariance to noise and distractors. In this work, we generalize value function approximation bounds for on-policy bisimulation metrics to non-optimal policies and approximate environment dynamics. Our theoretical results help us identify embedding pathologies that may occur in practical use. In particular, we find that these issues stem from an underconstrained dynamics model and an unstable dependence of the embedding norm on the reward signal in environments with sparse rewards. Further, we propose a set of practical remedies: (i) a norm constraint on the representation space, and (ii) an extension of prior approaches with intrinsic rewards and latent space regularization. Finally, we provide evidence that the resulting method is not only more robust to sparse reward functions, but also able to solve challenging continuous control tasks with observational distractions, where prior methods fail.

Towards Robust Bisimulation Metric Learning
Mete Kemertas" Department of Computer Science University of Toronto kemertas@cs toronto. edu
Tristan Aumentado-AAmsstrong
Department Computer Science University of Toronto taumenêcs toronto. edu
Abstract
Learned representations in deep reinforcement learning (DRL) have to extract task- relevant information from complex observations, balancing between robustness to distraction and informativeness to the policy. Such stable and rich representa- tions, often learned via modern function approximation techniques, can enable practical application of the policy improvement theorem, even in high-dimensional continuous state-action spaces. Bisimulation metries offer solution this representation learning problem, by collapsing functionally similar states together in representation space, which promotes invariance noise and distractors. In this work, we generalize value function approximation bounds for on-policy bisim- ulation metrics to non-optimal policies and approximate environment dynamics. Our theoretical results help us identify embedding pathologies that may occur in practical use. In particular, we find that these issues stem from an underconstrained dynamics model and an unstable dependence of the embedding norm on the reward signal in environments with sparse rewards. Further, we propose set of practical remedies: (i) norm constraint on the representation space, and (ii) an extension of prior approaches with intrinsic rewards and latent space regularization, Finally, provide evidence that the resulting method is not only more robust to sparse reward functions, but also able to solve challenging continuous control tasks with observational distractions, where prior methods fail.
1 Introduction
Complex reinforcement learning (RL) problems require the agent to infer useful representation of the world state from observations. The utility of this representation can be measured by how readily it can be used to learn and enact a policy that solves specific task. As an example, consider a robot using visual perception to pour a cup of coffee. The clutter on the counter and the colour of the walls have little effect on the correct action; different patterns and shapes, should be treated nearly the same by the policy. In other words, many states equivalent for given task, with their differences being task-irrelevant distractors. Thus, natural approach to generalization across envirommental changes is constructing a representation that is invariant to such nuisance variables effectively "grouping" such equivalent states together.
more pertinent aspects, such two mugs with slightly
One method obtaining such a representation uses notion bisimulation metric (BSM) [13, 14). The goal is to abstract the states into a new metric space, which groups "behaviourally" similar states in a task-dependent manner. In particular, bisimilar states (i.e., this case, those close under the BSM) should yield similar stochastic reward sequences, given the same action sequence. In a recursive sense, this requires bisimilar states to have similar (i) immediate rewards and (ii) transition distributions in the BSM space (e.g., see Eq. 1). Thus, ideally, the resulting representation space contains the information necessary to help the policy maximize return and little else.
Equual contribution.
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

Recently, the Deep Bisimulation for Control (DBC) algorithm [53] showed how to map observations into a learned space that follows a policy-dependent (or on-policy) BSM (PBSM) for given task [10], resulting in powerful ability to ignore distractors within high-dimensional observation spaces. Using a learned dynamics model, DBC trains an encoder to abstract states to follow the estimated PBSM, using the aforementioned recursive formulation for metric learning.
However, this method can suffer from issues of robustness under certain circumstances. Conceptually, differentiating state encodings requires trajectories with different rewards. In other words, two states that only differ in rewards far in the future should still not be bisimilar. Yet, in practice, the encoder and policy are learned simultaneously as the agent explores; hence, in the case of uninformative (e.g., sparse or near-constant) rewards, it may incorrectly surmise that two states are bisimilar, as most trajectories will look very similar. This leads to premature collapse of the embedding space, grouping states that should not grouped. On the other hand, we can show that the formulation of the metric learning loss is susceptible to embedding explosion if the representation space is left unconstrained our work, we build upon the DBC model in an attempt to tackle both problems: (1) we address embedding explosion by stabilizing the state representation space via norm constraint and (2) we prevent embedding collapse by altering the encoder training method.
Contributions We (i) generalize theoretical value function approximation bounds to broader family of BSMs, with learned dynamics; (ii) analyze the BSM loss formulation and identify sources of potential embedding pathologies (explosion and collapse); (ii) show how to constrain representation space to obtain better embedding quality guarantees; (iv) devise decoretically-motiviee modifications based on intrinsic motivation and inverse dynamics regularization; and (v) show on a variety of tasks that our method is more robust to both sparsity and distractors in the environment.
training
2 Related Work
Our work builds on two major aspects of DRL: representation learning, particularly task-specific en- codings following the BSM, and methods for sparse reward environments, (e.g., intrinsic motivation).
Representation Learning for Control In the context of RI, there is a spectrum of representations from task-agnostic task-specialized. One task task-agnostic approach to learning latent representation spaces uses a reconstruction loss. In particular, some methods utilize this approach for model-based planning [27, 54, 25], while others do for model-free RL [33, 17, 30J. Such approaches often utilize generative models, which have shown use for learning via simulation, in addition to representation learning [22, 51]. Instead of fully reconstructing an input, other approaches use self-supe representation learning methods, especially contrastive learning algorithms [34, 28, 44].
While task-agnostic reconstruction is reusable and transferable, it can store unnecessary or distracting details. In contrast, state abstraction methods [31] can obtain compact state representations that ignore irrelevant information, by finding (approximately) equivalent aggregated MDPs that have (nearly) identical optimal policies. To learn such MDP bomomorphisms [38, 391, previous efforts have exploited various structural regularities (c.g.. equal rewards and transition probabilities [19, 46), state-action symmetries and equivariances [46, 48, 49]) and assumptions [6]. State similarity metric learning [29] is closely related, as one can aggregate e-close states produce a new MDP. For example, inter-state distances can be assigned based on the difference between policy distributions conditioned on respective states [2, 18]. Inverse dynamies prediction can also help encourage holding only information relevant to what the agent can control [3, 36). For model-based RL, prior work constructed equivalence classes of models [21], as well as other abstractions [12, 16), to save computation and memory.
In our work, build on the use of task-specific bisimulation relations [19), which define equivalence
classes of states based on rewards and transition probabilities. However, obtaining cquivalence classes brittle, as our estimates of these quantities may be inexact. Hence, Ferns al. [13, 14, 15] instead turn to bisimulation metric between states, which can vary smoothly as the rewards and transition probabilities change. More recently, Zhang et al. [53] devised Deep Bisimulation for Control (DBC), which applies a metric learning approach to enforce its representation to approximately follow iisimulation-derive state aggregation. While DeepMDP [17] proves that its representation provides an upper bound on the bisimulation distance, DBC [53] directly imposes the bisimulation metric
Embedding collapse and explosion issues also appear in the computer vision literature [26, 42, 52).

---

structure on the latent space. Herein, we extend prior analysis and methods for representation learning using PBSMs, with the goal of improving robustness and general performance. In particular, we focus on preventing embedding pathologies, in the context of uninformative reward signals and distraction.
Reinforcement Learning under Uninformative Rewards Many real-world tasks suffer from un- informative rewards, meaning that the agent receives sparse (mostly zero) or static (largely constant) rewards throughout its trajectories. For example, games like Montezuma's Revenge only provide reward signals after a large number of steps and the agent is likely not to receive one at all, in most cases, Thus, while common, this situation significantly increases the difficulty of the RL task, as the agent receives little signal about its progress. To deal with such situations, a common tactic is to encourage the agent to explore in the absence of extrinsic reward signal. This is often done by pro- viding intrinsic motivation via a self-derived reward, resulting in "curiosity-drive Such approaches include surprise [1, 5] (where experiencing unexpected dynamics rewarded), and empowerment [40, 20] (where the agent prefers states in which it has more control).
behaviour [8, 4].
In this work, we specifically target the case of uninformative rewards, showing that the DBC model is especially susceptible to instability or collapse for such tasks, and consider ways to ameliorate this issue. In particular, we utilize the forward prediction error as intrinsic reward [36, 8, 433, thus augmenting the sparse extrinsic rewards and encouraging exploration. We also regularize the latent space by learning an inverse dynamies model on top of it, which does not rely on the extrinsic reward signal and has previously been used for distraction-robud robust representation learning in DRL [3, 36].
3 Technical Approach
In this section, we first describe our notation and problem setting, and review prior definitions of bisimulation metrics and relevant metric learning objectives. Then, we consider theoretical connections to value functions and convergence properties, and show that these formulations may be susceptible to (i) instabilities in optimization due to embedding explosion (i.e., large norms), and (ii) convergence to trivial solutions where all states are represented collapse). Based on our analysis, we propose an extension of deep bisimulation metric learning, with poonetiically-motivttt constraints on the optimization objective, including alterations to the forward dynamics model, intrinsic motivation, and inverse dynamics regularization.
single point (embedding
3.1 Preliminaries
In this work, we consider a discounted Markov Decision Process (MDP) given by a tuple, (S, A. P, R. Po). At the beginning of each episode, an initial state, Su E S, is sampled from the initial- state distribution Po over the state space S. Then, at each discrete time-stept 2 O, an agent takes an action, a, E A, according to policy π [s.. a result, the MDP transitions to the next state ac- cording to a transition distribution P(( a.). The agent collects scalar reward, = R(s,, az). from the environment according to bounded reward function, R. S A -> [Rmin: Rmax) In infinite- and long-horizon settings, discount factor, 7 E 1), is used to calculate the agent's discounted return in a given episode, G = -.20 7'r, RL algorithms aim to find an optimal policy, n² == argmax = [[. for a class of stationary policies II. In nigh-dimensional, continuous state (or observation) spaces, this learning problem is rendered tractable via state encoder (e.g., a neural network), Φ : S > R", which is used to learn a policy of the form r(a (s)))
The following (pseudo²) (rreudo")-mettee meetri,, based on the Wasserstein metric (see Appendix A for review), is of particular relevance to this work. Distances are assigned to state pairs, S;) E S x S, according to a pessimistic measure [9] of how much the rewards collected in each state and the respective transition distributions differ. A distance of zero for a pair implies state aggregation, or bisimilarity.
Definition 1 (Bisimulation metric for continuous MDPs, Thm. 3.12 of [14]). The following metric exists and is unique, given R: A + [0, 1] and E (0, for continuous MDPs:
d(s s,) mmx 1 c) R(s,, a)- - R(s,. a)l + (W)(d))P(:ss,, a), P(-ls,, a)).
(1)
aE.A
An earlier version of this metric for finite MDPs used separate weighting constants CR, Or 2 0 for the first and second terms respectively, and required that CR + er - 1 [13]. Here, when the weighting
"Pseudo-metries are generalization of metries that allow distinct points to have zero distance. For simplicity. we broadly use the term "metric" in the remainder of this paper the cost of imprecision.
3

---

constant of the (d) term is in [0, 1), the RHS is contraction mapping, F(d) met + met, in the space of metrics. Then, the Banach fixed-point theorem can be applied to ensure the existence of unique metric, which also ensures convergence via fixed-point iteration for finite MDPs. For more details, we refer the reader to [13. 14] and the proof of Remark in Appendix B. Notice that e (or CT) determines timescale for the bisimulation metric, weighting the importance of current versus future rewards, analogously to the discount factor y, More recently, an on-policy bisimulation metric (also called 7-bisimulation) was proposed to circumvent the intractibility introduced by taking the max operation over high-dimensional action spaces (e.g., continuous control), as well as the inherent pessimism of the policy-independent form [10].
Definition 2 (On-polis) bisimulation metric [10]). Given a fixed policy n. the following on-policy
bisimulation metric exists and is unique:
(s. = 1 -151 +
(-(ss, PT))s))))))))))))))))))))))))P
where r7 := -aan [R(s a] and pm (-1s:) == Ea-a [P(-|s: a)].
Zhang et al. [53] proposed to learn a similar on on-policy bisimulation metric directly in the embedding space via MSE objective. They proposed an algorithm for jointly learning a policy π (alo(s)) with an on-policy bisimulation metric. Below, we define a generalized variant of their objective:
(10)=1pm LL J(0) = 2
(dr.(()))))) -r -r >Wa((I-ll (P' (-lo(s:)), (1((())))) (3)
where (s. 8,) = [(((s)) - -b(s))l and they used (4) = 2, = 1). Notice that the recursion induced by this objective is different from prior metrics in three ways; (i) a W22 metric was used instead of a Wi metric since W22 has a convenient closed form for Gaussian distributions when q 2, (ii) the distance function used for the bisimulation metric and the Wasserstein metric different (L1 and L2 respectively), and (iii) a forward dynamics model is used instead of the ground truth dynamics. While they may introduce practical benefits, these differences violate the conditions under which the existence of a unique bisimulation metric has been proven [10, 13, 14, 15]. Thus, we will (i) assume unique metric exists for all Wasserstein metrics Wp when necessary (see Assumption 1), (ii) study losses that use a matching metric, i.e., fi = Q2 = q. and (iii) introduce a constraint on forward models in Sec. 3.2.2.
Next, we note that they recommend the use of stop gradients for the W2 term The resulting gradient updates may be considered as approximate fixed-point iteration in the space of metrics:
F(ā (s: s,) = ār (s:, 8y; tne + on V V (on),,),
(4)
where on a learning rate. However, in practice, π and P may also be updated in training, which is of particular relevance when they have form =< (s)) and ((o(s))(((s)) a) as in [53]. In the next section, we will discuss conditions under which joint updates to a policy x and a metric defined by state encodings e may or may not converge practical settings:
lim F(n) or 2 dr.,
(5)
where dz is the on-policy bisimulation metric for the optimal policy x².
3.2 Theoretical Analysis
In this section, we generalize prior value function bounds for bisimulation metrics to cases where arbitrary weighting constants CR, Cr are used, and further derive value function approximation (VFA) bound for / rather than V* unlike prior work [13, 14, 533. We then describe constraints on forward dynamics models for convergence to unique metric in joint training, and derive VFA bounds as function of forward dynamics modelling error. Then, we discuss potential pitfalls in on-policy bisimulation metric learning, which lead us to connecting its weaknesses with sparse rewards. Our findings motivate our proposed remedies to the issues we identify in on-policy bisimulation; we recommend a particular norm constraint on the latent space, and motivate the use of inverse dynamics and intrinsic motivation outlined in Sec. 3.3. All proofs are provided in Appendix B.

---

3.2.1 Value Function Bounds
An important feature of bisimulation metrics is their relation to value functions since provably tight connection implies guarantees in VFA. Similarly to previous bounds for volicy-independent bisimulation metrics [13, 14], Castro [10] showed that given any two states (s;, s;), the following bound holds for on-policy bisimulation metrics (see Definition 2): [V*(s:)
V* [s))]= da (s, Sj).
As discussed earlier, differently from previous approaches, Zhang et al. [53] used 2- Wassseesteiin metric due to practical reasons. Here, in order to generalize previous value function bounds, we assume the existence of unique bisimulation metric for p- g-Wasserstein metrics.
Assumption 1 (A1, p- Wasserstein bisimulation metric). For given E [0, cR E [0, oo) and p 2 the following bisimulation metric exists and is unique:
(= 8,) = calr" - r3 + er W (dd(() P*)ss)
(6)
Remark 1. Ifp 1, or both the environment and policy are deterministic, Al holds.
Theorem (Generalized value difference bound). Let the reward function be bounded as R& [0. 1].
For an on-policy bisimulation metric given by Eq. (6). for any er E [0, and 2 1, define min(er,"), 7)- Given AI, the bisimulation distance between a pair of states upper-bounds the difference in their values:
2c V(s,, ESXS.
cnlV V S (=))),,)) + -1- (1- ccr)
Note that Thm. 3 of [10] is special case with en p and cT = 7. Suppose er 2 y; we observe that for degenerate metric d O, the corresponding value function a constant V V* (s) = c. We will discuss the dangers posed by this relation in Section 3.2.3.
Theorem (Generalized VFA bound). Let rewards be bounded as R E [0, 1] and $ S S be function mapping states to finite partitioning S such that (ss)) = (sy)) = d. (8;, 83) S 2e, which produces aggregated MDP (S, A, For any eT E 1). let - min( (TT, 7). Given Al,
2( 7) CR(1 7) (1-7) (1cr))
2c
(V*(s) - (+)) S
ss E S.
(8)
This result generalizes previous performance bounds for V* to non-optimal policies. Previous bounds also assumed cr 2 7. while generalize arbitrary cr 1). Here, the second term of the upper bound characterizes the penalty paid for "myopic" bisimulation (i.c., CT < 7) in VFA error guarantees. We further show in the next section that cr < y may be desirable when approximate forward models are used (see Thm. 4). Indeed, in Appendix C, we not only connect large er to high variance in embedding norms, but also find surprisingly strong results when empirically investigating the use of cT < 7. We speculate that with sufficient modelling capacity and a critic V()) being trained based y, the metric space may hold information about the environment regarding both timescales and cr.
3.2.2 Bisimulation Metrics with Approximate Forward Dynamics
We next examine theoretical properties of the on-policy bisimulation distance, which turn out to constrain properties of the approximate transition model. This motivates an additional architectural constraint, resulting in empirical improvements over prior work.
Metrics prior to [53] were defined based on ground-truth dynamics, which enabled the application of the Banach fixed-point theorem to show the existence of unique metric. However, the Banach fixed-point theorem requires that a contraction mapping is applied on complete metric spaces. In this case, the metric space which the Banach fixed-poiie theorem is applied is itself a space of metrics met over states. For that reason, when bisimulation metrics were generalized from finite to continuous MDPs, a compact' state space was assumed to ensure completeness of met [14]. In practice, when using an approximate forward dynamics model P. S A + M(S'),, if compactness
AA metric space is compact if and only if is totally bounded and complete [14]. 5 M(2) denotes the space of all probability distributions over X.
5

---

Effect Cartaple
Moan Bsimu
Reward Robio
286
6-
0-071
ortetrry form NarTuR
125
Eoc Prewert
100
TBC vor (Mcenkure)
500
30000
20000
330.0
40020 20000
5000 30000 1,000 20303
Figure I: Theoretical properties under the SparseCartpole task. (Left) Embedding explosion occurs for DBC [53] for the sparse 2D-Cartpole task, while our constraints discussed in Sec. 3.2.2 keep the model stable. Our approach achieves maximum returns at evaluation, while DBC diverges. (Right) Given cr 0.5. Eq. (17) correctly predicts the ratio / pd between bisimulation distance and difference of rewards at convergence. Dashed lines indicate analytically calculated targets, while solid lines correspond to mini-batch estimates of the ratio during training.
is not guaranteed, the same convergence guarantees do not apply. Thus, ideally, S' is compact subset S. Here, we formalize these restrictions on latent representations and the forward dynamics model to guarantee convergence to unique metrie. The following states constraint on the BSM solution space given approximate dynamics.
Lemma (Diameter of bounded). Let d - oo) be any bisimulation metric:
diam(S:d) = sup d(s., S,8,
CR
(Rmax
Rmmin.
1 Cr
The above lemma holds for all bisimulation metrics (on-policy or otherwise) that are based on exact dynamics, if said metrics exist. Let also consider on-policy BSM with imperfect dynamics, such that it may not necessarily satisfy Lemma 1.
Definition 3 (On-policy bisimulation metric with approximate dynamics), Given an approximate dynamics model P: S x A + M(S),, er E [0, 1), and eR E [0, oc): P (-1sg))
(= = cRIr r3] e
(10)
Next, we provide sufficient condition for the existence of unique metric d, based approximate dynamics, which also satisfies the upper-bout bound of Lemma 1. Meeting this condition shrinks the solution space of metrics being searched to set known to contain d_.
Theorem (Boundedness condition for convergence). Assume S is compact. If the support of an approximate dynamics model P, i.e., S' supp(P), is closed subset of S, then there exists unique on-policy bisimulation metric d of the form Eg. (10), and this metric bounded:
supp(p) C diam(S;
- -(Rmax Rmii).
CR
(11)
this condition not satisfied at any point during training, the system may diverge. Indeed, we confirm empirically in Sec. 4 that the absence of norm constraint on the forward dynamics model may result in embedding explosion with practical consequences (see Fig. 1). Such explosions are due to compactness violations of the approximate dynamics (e.g.. if the predictions always increase the embedding space diameter, the recursive nature of the metric can lead to runaway expansion). Luckily, the condition can be satisfied with ease, e.g.. by projecting larger vectors onto the surface of a closed ball, Bg, with diameter given in Lemma 1, such that the following are true:
(bs) E B (x RD ||x||, S Cm(Rmax - Rmia))
}, E
(12)
2(1 - 2(1-CT)
((-10(s),a M(B.), v(s A.
(13)

---

We find that the constraint applied here is mild yet effective, as evidenced by significantly improved performance and stability when the constraints are active (see Fig. 1).
Clearly, a necessary condition for d - d is an error-free dynamics model. natural question that arises from this view concerns the degree to which modelling errors affect VFA bounds.
Theorem (VFA bound in terms of model error). Consider the same conditions as in Theorem 2. except that Cr E h,, 1), p 1, and ((si) = ((8,) = (8, 85) = (1x(s;) - +((8))11 22. Then:
(S)
cn(1 7) (2 Eo 1 CT
2CR Er
2et Ep Aa
(14)
CT
where Eo :- -=1idz,s - 1oo is the metric learning error, E, := -l r oo is the reward approxi- mation error, and Ep == supses Wi (d=) (-1s), (-]ss) is the state transition model error.
See Appendix D) for details, including generalized version for er E [0, 1), listed as Corollary 3. Consider an ideal case where ground-trutt rewards r" are available and the metric learned perfectly (i.e., Eg- & 0). Then, we observe that the choice of eT defines trade-off between VFA error due to (a) forward model error Ep, and (b) "myopic" bisimulation (see Thm. 2, CT < 7), where the BSM timescale is shorter than that of the discounted return.
3.2.3 On the Dangers of On-policy Bisimulation
Suppose that at training time, InT - -r 0 for all pairs of states, e.g. during early training in a sparse reward setting. Then, unlike the poliey-independen bisimulation metric, the on-policy formulation has a degenerate solution at diam(S: S5 dz) = 0, regardless of the structure of the underlying MDP. Lemma (A reason for caution in on on-oolcc bisimulation). On-policy bisimulation metrics of the form Eq. (6) have an upper bound determined by their policy:
diam(S: S
CR
sup r - 5
(15)
1 - CT 4
Due to policy-dependence, the target metric di changes with each policy update during joint training. As result of this difficulty, convergence to unique fixed point (i.e., unique metric) for learning algorithm was previously guaranteed with a strong assumption: "a policy that is continuously improving", as in Thm. 1 of [53]. Informally, if the policy is assumed to be continuously improving, in the worst case, the metric will have a fixed point dg- after the policy itself reaches a fixed point, namely, the optimal policy n². However, this assumption may be too strong to guarantee convergence especially for continuous MDPs, as the policy learning process depends heavily on the encoder.
Specifically, to prove their Thm. 1, Zhang et al. [53] relied on the policy improvement theorem. For continuous MDPs, in practice, a policy π (s)) is learned via non-convex optimization (e.g., policy gradients, VFA), rather than the vanilla policy improvement algorithm. Thus, if the bisimulation metric is degenerate, a continuously improving policy cannot be guaranteed. As we will show, on-policy bisimulation metrics can in fact obstruct policy search in some cases (e.g., sparse rewards, low dispersion rewards) by yielding a collapsed or exploded metric space, which is unable to approximate the value function. Specifically, we can define and relate measures of (i) collapse in the embedding space, and (ii) statistical dispersion of rewards under the current policy.
Definition 4 (Measuring collapse and sparse rewards). Let pa denote the stationary distribution over states, and n the distribution over pairs of states, (s;, sampled independently from p" Then;
Pod = B( (S2, sy)]
p = E, [17 --
(16)
In the low dispersion case, pid O, meaning the current reward signal under is uninformative. However, this can have dire consequences for our embedding, as shown in the lemma below.
Lemma 3 (Relating collapse and lowddispersion rewards). Assume deterministic transitions and the existence of a stationary distribution over states. Given a bisimulation metric of the form Eg. (6);
CR
pbd =
Prd
(17)
1- CT
"Statistical dispersion is umbrella term used to describe measures of variability diversity (e.g., variance).
7

---

1000
1000 30000
20100
140
2.5
20090
Figure 2: Results on the modified Gym tasks: Cartpole (top row) and Mountain Car (bottom row). The left, middle, and right columns have O, dSm(S), and 2dim(S) distractor noise dimensions respectively. We show the Soft Actor-Critic (SAC) and Deep Bisimulation for Control (DBC) baselines, along with modifications using embedding normalization (normed), intrinsic rewards (IR), and inverse dynamics (ID). DBC struggles under all conditions of sparsity; while SAC handles it better (top-left), it cannot deal with high distraction (top-right) or more extreme sparsity (bottom row). However, latent normalization immediately improves performance (top row). Further, combining it with IR IR+ID improves performance on all tasks. Shading shows standard error over 10 seeds.
Clearly, a collapsed state encoder, o² (s;) Φu. loses all information about the state. These observations motivate us to extend the method with inverse dynamics-basee based regularization and intrinsic rewards based on forward prediction errors (see Sec 3.3), since they promote Pia Prd > 0 respectively. For empirical evidence of the relation in Lemma at training- time, see Figure 1. Relationships between variances are also discussed in Appendix C.
3.3 Intrinsic Rewards and Inverse Dynamics
The core principle DBC is to construct representation such that given 6(s) relates to other latent states via the PBSM, thus providing robustness to distractors, but also ensuring the latent state holds sufficient information to maximize return. As already discussed, however, uninformative rewards can induce metric learning issues for DBC, causing the embedding to hold insufficient information to solve the task. We consider two approaches to improving upon this representational deficiency: (I) using intrinsic rewards (IR), and (2) regularizing the latent space with inverse dynamics (ID) learning,
suriosity-Driver Intrinsic Rewards One technique for encouraging exploration in sparse-reward environments uses intrinsic motivation via self-generated rewards, based on a notion of curiosity [8, 4]. In such cases, the agent rewards itself for entering unpredictable areas of the environment: in particular, at every time step i, we redefine the reward signal to be rr TT,t + TEat where f,, the extrinsic reward (i.e., the original environmental reward) and rT, is the IR. Following prior work [8, 36, 43], we use the forward model error in the latent space to compute intrinsic rewards:
FI, 7, 1 (St. az) ((s(+1)))))))))))))))(()) 112/ (ll)/(2n) where (st ar) B ) P(Pl-l(ss)))))))))))- [o(.))(((())))) [o(s is the mean of the predicted distribution from the forward dynamics model and 17, 0 is hyper- parameter controlling the intrinsic reward scale. We finally clamp to fixed maximum Riaa, so
We already train P for the DBC loss, so the additional computational cost is limited. Finally, note that IRs cause the reward signal to become non- non-stationary resulting in the target of the BSM learning changing as well; however, the agent explores and can better predict its environment over time, the IR (and thus its influence on the metric learning process) should fade.
that T,E = main(Rmax.Jii
Latent Inverse Dynamics Prior We also consider regularizing the learned latent state space, by having learn to predict the inverse dynamics of the world, following prior work [3, 36). Let A be an inverse model, which predicts the action a, E A that changed to
91 1
R
8

---

cartocl wallsup spare
0BC
TCE
DBC
422-
293
230
Thaning 200
Traiing
Cheeogh run (wtt disraction))
Chitpete madt oparte (AIR
270
54G 1000
350
Teal Cuaaaae saae
150
0.2 2
Figure 3: Performance on 3D robotics tasks from DMC [45] with random seeds. Leftmost column illustrates the training setup for the cheetah-run task without (top) and with (bottom) natural video distractions. (Top-left) Our method significantly improves upon DBC [53] and is slightly below SAC [23]. (Top-right) DBC struggles under sparse rewards, while our approach remains robust to sparsity. (Bottom-left) Our approach significantly outperforms DBC and is more sample-efficient than SAC. (Bottom-right) DBC makes no progress under both sparsity and distractions, and SAC begins learning later due to sparsity.
8:41 using the learned latent space: ā = gr(((()))))))))))))))))()))))( ((()) ()). This can be trained from observed transitions (Se, a,, 11 ). via Ja(Ø, na) nal a. à. / where fa 0 weighs the ID loss importance, and A = [-1, 1|" for continuous control tasks. This loss prevents removal of information that pertains the loss driving é to mimic the BSM. In this sense, we are placing regularizing prior on á via an auxiliary task, such that it adheres the BSM requirement when possible, still transmits useful information in the sparse reward case instead of collapsing. Yet, since only actions need be predicted, distractor aspects of the observation that the agent cannot influence are naturally ignored. Hence, we expect less disruption to the BSM encoding than other auxiliary tasks (e.g., reconstruction).
the agent's s ubility control the environment, but it is balanced against
Summary In this section, we (1) provided generalized bounds connecting the value function the on-policy BSM (Sec. 3.2. 1); (2) showed that the BSM satisfies fixed latent diameter, and devised normalization for enforeing this property during learning (Sec. 3.2.2), which not only prevents embedding explosion but improves performance (Fig. 1); and (3) found the PBSM is susceptible to embedding collapse with sparse rewards, and suggested and ID to mitigate it (Sec. 3.2.3, 3.3).
Experiments
In this section, we seek answers the following questions concerning our main hypotheses:
Do the embedding collapse and explosion issues predicted thcoretically occur in practice?
2. Do our contributions address these problems?
3. Does our proposed approach preserve the noise-invariance property of bisimulation?
4. How do our proposed improvements interact with each other?
5. How does our method perform compared to prior work, particularly with sparse rewards?
To that end, we experiment on several altered classic control tasks from OpenAl Gym [7] by (i) sparsifying the reward signal, and (ii) augmenting the environment state with noisy dimensions, to simulate distractions. We also perform larger scale experiments on two challenging vision-based robotics benchmarks from the DeepMind Control Suite [45].

---

Sparse Noisy Cartpole First, we modified the Cartpole-vo task, in which an agent tries to keep a pole upright. To increase sparsity, we shrank the angular extent which the pole must be to earn rewards to 1 1% of its standard value. Then, to mimic distractors, we concatenate an (S) dimensional vector sampled from an isotropic Gaussian to the state vector. (See Appendix 1. for more details.) The performance of the respective models, as well as other baselines are shown in Fig. 2, top row. While DBC struggles on this task due to the sparsity, the Soft Actor-Critic (SAC, [23]) baseline, which does not use BSMs, is able to solve the problem (top-left inset). Our proposed modifications address the embedding explosion problem and perform on-par with the SAC baseline, both when combined and separately. Furthermore, when distractors are introduced (top-middle and top-right insets), SAC also fails, while our approach combining intrinsic rewards (IR) and robust metric learning still solves the task. See also Appendix 1.1.4 for results with Nn. = 3.
Noisy Mountain Car We also consider the classic MountainCo rCont i nuous-vO task [32], in its continuous control form (see Appendix 1.2 for details). Since the task was already highly sparse, we modified only the distraction aspect, as the Cartpole scenario. As shown in Fig. 2, all methods without IR are unable to solve the task and are stuck with rewards close to zero. However, DBC with IR is unstable and unable to solve the task either. Only using the normalization allows the agent to succeed; furthermore, the inverse dynamics (ID) regularization improves convergence speed and stability. This trend continues at even higher distraction levels as well (see Appendix E.1.4).
Sparse Noisy Pendulum We also modified the continuous Pendul um-vO task to have higher degree of sparsity and distraction. Due to space constraints, we relegate details to Appendix E.1.3. We find that our method performs on par with DBC, but outperforms SAC in the presence of distractions.
DeepMind Control Suite The DeepMind Control Suite (DMC) [45] contains set of 3D robotics tasks based on the MuJoCo physics simulator [47]. We include results from two DMC tasks, namely, cheetah-
(dim (S) 118, dim((A)=6) and cartpole wwwsssswinggsiieeee sparse (dim(S)-4, dim (A)==1), shown in Figure 3. The former is a common benchmark where Zhang et al. [53] showed that their method performed sub-par in the absence of distractors. The latter is task we select due to reward sparsity. Similarly to [53], we train for each task with and without natural video distractions, as illustrated in Figure 3. In all cases, our approach performs significantly better than the DBC baseline [53]. Without distractors, our method performs competitively with SAC [23], but when distractions are introduced, our approach is slightly more sample-efficient With these experiments, we verify that our improvements carry over to larger-scale tasks where learning is over raw-pixels.
5 Conclusion
Limitations and Impact One shortcoming of our approach is the lack of a principled way to set hyper-parameters for IR and ID, which was done empirically. Indeed, our use of forward model error and ID regularization itself may not be optimal. Sometimes, we observe that the embeddings all attain norms close to the maximum allowed value, suggesting alternative approaches to normalization may be helpful. Finally, IR and ID incur an additional computational cost to training. In terms of broader impact, while our work is foundational RL, focuses on removing distractions from internal representations. However, this culling can remove information that the agent deems unrelated to the task, potentially leading to unsafe decisions in unseen scenarios (e.g.. for mission critical systems).
Discussion In this work, we investigated the behavior of on-policy deep bisimulation metric learning approaches, which construct efficient neural representations that are invariant to noise and distractors, without reconstructing the raw input [10, 53]. We identified embedding collapse and explosion as potential failure modes of this method via theoretical analysis, and highlighted that it may be especially susceptible to failure in sparse reward settings. Our experiments confirmed the dangers of these failure modes and showed that our proposed remedies address the issue. In particular, we enforced a norm constraint on state representations, and incorporated intrinsic motivation and latent regularization in our technique. The resulting approach preserves the noise-invariance property of bisimulation metrics while comparing favorably against strong baselines [23, 53] on altered versions of classic control tasks, which we rendered more challenging by sparsifying the rewards and synthesizing distractors, as well as harder tasks with visual observations. Future work includes investigating alternative ways to improve embedding regularization and reward informativeness, as well as more realistic 3D robotics benchmarks, with sparse reward structures and heavy distraction.
10

---

We thank Sven Dickinson, Allan Jepson, and amir-massous Farahmand for helpful discussions. The support of NSERC (CGSD3-539955-019) and Samsung Al Research is gratefully acknowledged.
Acknowledgments
References
[1] Joshua Achiam and Shankar Sastry. Surprise-based intrinsic motivation for deep reinforcement learning. arXiv preprint arXiv: arXiv:1703.01332 2017.
[2] Rishabh Agarwal, Marlos C. Machado, Pablo Samuel Castro, and Mare G. Bellemare. Con- trastive behavioral similarity embeddings for generalization in reinforcement learning. In International Conference Learning Representations, 2021.
[3] Pulkit Agrawal, Ashvin Nair, Pieter Abbeel, Jitendra Malik, and Sergey Levine. Learning to poke by poking: experiential learning of intuitive physics. arXiv preprint arXiv: arXiv: 1606.0041, 2016.
[4] Arthur Aubret, Lactitia Matignon, and Salima Hassas. A survey on intrinsic motivation in reinforcement learning. arXiv preprint arXiv: (90.00900.06766 2019.
[5] Marc G Bellemare, Sriram Srinivasan, Georg Ostrovski, Tom Schaul, David Saxton, and Remi Munos. Unifying count-based exploration and intrinsic motivation. arXiv preprint arXiv:1606.014 08688, 2016.
[6] Ondrej Biza and Robert Platt. Online abstraction with MDP homomorphisms for deep learning. In Proceedings of the 18th International Conference on Autonomous Agents and MulriAgent Systems, 2019.
[7] Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. OpenAl gym. arXiv preprint arXiv: 160.00/66000 2016.
[8] Yuri Burda, Harri Edwards, Deepak Pathak, Amos Storkey, Trevor Darrell, and Alexci A Efros. Large-scale study of curiosity-driven learning. arXiv preprint arXiv:1808.0433 2018.
[9] Pablo Castro and Doina Precup. Using bisimulation for policy transfer MDPs. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 24, 2010.
[10] Pablo Samuel Castro. Scalable methods for computing state similarity in deterministic Markov decision processes. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pages 10069-10076, 2020.
[11] Philippe Clement and Wolfgang Desch, An elementary proof of the triangle inequality for the Wasserstein metric. Proceedings of the American Mathematical Society, 136(1):333-33) 2008
[12] Dane Corneil, Wulfram Gerstner, and Johanni Brea. Efficient model-based deep reinforcement learning with variational state tabulation. In International Conference Machine Learning, pages 1058. PMLR, 2018
[13] Norm Ferns, Prakash Panangaden, and Doina Precup. Metries for finite Markov decision processes. In Proceedings of the 20th Conference on Uncertainty in Artificial Intelligence, UAT '04, page 162-169, Arlington, Virginia, USA, 2004. AUAI Press.
[14] Norm Ferns, Prakash Panangaden, and Doina Precup. Bisimulation metrics for continuous Markov decision processes. SIAM J. Comput., 40(6): 1662-1664111 December 2011.
[15] Norm Ferns and Doina Precup. Bisimulation metrics are optimal value functions. In Proceedings of the Thirtieth Conference Uncertainty in Artificial Intelligence. UAI' 14, page 210-219, Arlington, Virginia, USA, 2014. AUAI Press.
[16] Vincent François- Yoshua Bengio, Doina Precup, and Joelle Pineau. Combined reinforce- ment learning via abstract representations. Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pages 3582-3589, 2019.

---

[17] Carles Gelada, Saurabh Kumar, Jacob Buckman, Ofir Nachum, and Mare G Bellemare. Deep- MDP: Learning continuous latent space models for representation learning. In International Conference Machine Learning, pages 2170-2179. PMLR, 2019.
[18] Dibya Ghosh, Abhishek Gupta, and Sergey Levine. Learning actionable representations with noal-conditioned policies. arXiv preprint arXiv:1811.07819, 2018.
[19] Robert Givan, Thomas Dean, and Matthew Greig. Equivalence notions and model minimization in markov decision processes. Artificial Intelligence, 147(1):163-223, 2003, Planning with Uncertainty and Incomplete Information.
[20] Karol Gregor, Danilo Jimenez Rezende, and Daan Wierstra. Variational intrinsic control. arXiv preprint arXiv:1611.07507 2016.
[21] Christopher Grimm, Andre Barreto, Satinder Singh, and David Silver. The value equivalence principle for model-based reinforcement learning. Advances in Neural Information Processing Systems, 33, 2020.
[22] David Ha and Jürgen Schmidhuber. World models. arXiv preprint arXiv: 10/20,0022)))
[23] Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off- policy maximum entropy deep reinforcement learning with a stochastic actor. In International Conference Machine Learning, pages 1861-1870. PMLR, 2018.
[24] Tuomas Haarnoja, Aurick Zhou, Kristian Hartikainen, George Tucker, Sehoon Ha, Jie Tan, Vikash Kumar, Henry Zhu, Abhishek Gupta, Pieter Abbeel, et al. Soft actor-critic algorithms and applications. arXiv preprint arXiv:1812.05900 2018.
[25] Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In International Conference on Muchine Leurning, pages 2555-2565. PMLR, 2019.
[26] Alexander Hermans, Lucas Beyer, and Bastian Leibe. defense of the triplet loss for person re-identification. arXiv preprint 3rXiv:1703.0773: 2017.
[27] Thanard Kurutach, Aviv Tamar, Ge Yang, Stuart Russell, and Pieter Abbeel. Learning plannable representations with causal InfoGAN. arXiv preprint /80.:0990 2018.
[28] Michael Laskin, Aravind Srinivas, and Picter Abbeel. Curl: Contrastive unsupervised represen- tations for reinforcement learning. In International Conference on Machine Learning, pages 1639-5650. PMLR, 2020.
[29] Charline Le Lan, Mare Bellemare, and Pablo Samuel Castro. Metrics and continuity in reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pages 8261-8269, 2021.
[30] Alex Lee, Anusha Nagabandi, Pieter Abbeel, and Sergey Levine. Stochastic latent actor-critic: Deep reinforcement learning with latent variable model. arXiv preprint 1xX.v:1907.00993 2019.
[31] Lihong Li, Thomas J. Walsh, and Michael L. Littman. Towards a unified theory of state abstraction for MDPs. In In Proceedings the Ninth International Symposium on Artificial Intelligence and Mathematics, pages 531-539, 2006.
[32] Andrew William Moore. Efficient memory- -baed learning for robot control. 1990.
[33] Ashvin Nair, Vitchyr Pong, Murtaza Dalal, Shikhar Bahl, Steven Lin, and Sergey Levine. Visual reinforcement learning with imagined goals. arXiv preprint 3rXiv:1807.04744 2018.
[34] Masashi Okada and Tadahiro Taniguchi. Dreaming: Model-based reinforcement learning by latent imagination without reconstruction. arXiv preprint arXiv:2007.145355 2020.
[35] Ingram Olkin and Friedrich Pukelsheim. The distance between two random vectors with given dispersion matrices. Linear Algebra and its Applications, 48:257-263, 1982.
12

---

[36] Deepak Pathak, Pulkit Agrawal, Alexei A Efros, and Trevor Darrell. Curiosity -divve exploration by self-supervised prediction. In International Conference on Machine Learning, pages 2778-
2787. PMLR, 2017.
[37] Marek Petrik and Bruno Scherrer. Biasing approximate dynamic programming with lower discount factor. In D. Koller, D. Schuurmans, Y. Bengio, and L. Bottou, editors, Advances in Neural Information Processing Systems, volume 21. Curran Associates, Inc., 2009.
[38] Balaraman Ravindran. An algebraic approach to abstraction in reinforcement learning. PhD thesis, 2004.
[39] Balaraman Ravindran and Andrew G Barto. Approximate bomomorphisms: A framework for non-exact minimization in Markov decision processes. 2004.
[40] Christoph Salge, Cornelius Glackin, and Daniel Polani. Empowerment-an introduction. In Guided Self- Organization: Inception, pages 67-114. Springer, 2014.
[41] Filippo Santambrogio. Optimal transport for applied mathematicians. Birkäuser, NY, 55(58- 63):94, 2015.
[42] Florian Schroff, Dmitry Kalenichenko, and James Philbin. FaceNet: A unified embedding for face recognition and clustering. In 2015 IEEE Conference on Computer Vision and Puttern Recognition (CVPR), pages 815 823, 2015.
[43] Bradly C Stadie, Sergey Levine, and Pieter Abbeel. Incentivizing exploration in reinforcement learning with deep predictive models. arXiv preprint arXiv: 00004 2015.
[44] Adam Stooke, Kimin Lee, Pieter Abbecl, and Michael Laskin. Decoupling representation learning from reinforcement learning, arXiv preprint wrXiv:2099.083199 2020.
[45] Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, et al. DeepMind control suite. arXiv preprint arXiv: 0000100000000 2018
[46] Jonathan Taylor, Doina Precup, and Prakash Panagaden. Bounding performance loss in approxi- mate MDP homomorphisms. Advances in Neural Information Processing Systems, 21:1649- 1656, 2008
[47] Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: physics engine for model-based control. In 2012 IEEE/RSJ International Conference Intelligent Robots and Systems, pages 0023. IEEE, 2012.
[48] Elise van der Pol, Thomas Kipf, Frans A Oliehoek, and Max Welling. Plannable approximations MDP homomorphisms: Equivariance under actions. arXiv preprint arXiv:2002. 1963, 2020.
[49] Elise van der Pol, Daniel Worrall, Herke van Hoof, Frans Olichoek, and Max Welling. MDP homomorphic networks: Group symmetries in reinforcement learning. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, editors, Advances Neural Information Processing Systems, volume 33, pages 4199 4210. Curran Associates, Inc., 2020.
[50] Cédric Villani. Optimal transport: old and new, volume 338. Springer Science & Business Media, 2008
[51] Manuel Watter, Jost Tobias Springenberg, Joschka Boedecker, and Martin Riedmiller. Embed control: locally linear latent dynamics model for control from raw images. In Proceedings of the 28th International Conference on Neural Information Processing Systems, 2015.
[52] Chao- Yuan Wu, R. Mammatha, Alexander J. Smola, and Philipp Krähenbühl. Sampling matters in deep embedding learning. In 2017 IEEE International Conference on Computer Vision (ICCV), pages 2859-2867, 2017.
[53] Amy Zhang, Rowan Thomas McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Invariant representations for reinforcement learning without reconstruction. In International Conference Learning Representations, 2021.
13

---

[54] Marvin Zhang, Sharad Vikram, Laura Smith, Pieter Abbeel, Matthew Johnson, and Sergey Levine. Solar: Deep structured representations for model-based reinforcement learning. In International Conference on Machine Learning, pages 7444-7453. PMLR, 2019.
[55] Brian D Ziebart. Modeling purposeful adaptive behavior with the principle of maximum causal entropy. 2010.
14


Reviews:
[{"review_id": "mzhYCQxRL-C", "reviewer": "Reviewer_mSYo", "summary": "This paper considers the problem of learning good representations for reinforcement learning problems. It approaches this from the perspective of bisimulation metrics, extending earlier work on on-policy bisimulation for control (DBC). The paper identifies issues with DBC that also appear in other types of metric learning, and includes a theoretical analysis of why this happens. To remedy the issues the paper proposes normalizing the representations and including an extension to DBC by including an inverse dynamics model and intrinsic rewards.", "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\n* This paper is well written and answers many of the questions that arise while reading it immediately. \n* The analysis of certain issues in DBC is very nice to have. The theoretical section was (mostly) easy to follow and was informative. \n* The theory section has a clear added value in analysing the issues with DBC and proposing improvements. For example, compact spaces are assumed for continuous MDPs and this assumption is violated in DBC. This motivates the authors to propose normalizing the representations to a closed ball. Additionally, there is a degenerate solution when rewards are sparse, which motivates the authors to propose intrinsic rewards to prevent this from occurring. \n* An ablation study is included (though it is incomplete)\n* The paper does a good job of reviewing the related work, only missing a few related references (included below). \n\n### Weaknesses:\n* The choice of inverse dynamics as a regularizer is not well-motivated. This choice is not grounded in the same way that the rest of the paper is. For example, why choose an inverse dynamics model and not a different self-supervised loss? Why not a different combination of approaches? The motivation given for inverse dynamics could also be applied to reconstructions or contrastive coding. Additionally, there are no experimental results showing that an inverse dynamics model is best; as such this choice is the weakest part of the paper.\n* As there is no experimental comparison to other additions to DBC (e.g. different self-supervised learning approaches), it is unclear why  inverse dynamics is the preferred method to solving the issues with DBC.\n* The paper focuses heavily on one method; DBC, and as such is somewhat limited in scope.\n* The paper, while well-written, is not fully self-contained - certain terms are used without being defined. A few examples: \n  * What is a ‘compact subset of $\\mathcal{S}$’? \n  * What is a diameter of a state space? \n  * What is sup? The supremum? \n  * What is supp? Also the supremum? \n  * What is $\\mathcal{H}$ in line 259?\n  * What is a low dispersion reward?\n* According to the paper, DBC suffers from common metric learning problems, but there are no references included to other metric learning papers that exhibit similar issues.\n* If the intrinsic reward is based on the forward model error, it should not be useful under state collapse, as the error would be 0. \n* The non stationarity introduced by IR might be an issue in more complex environments.\n* Line 294: Claim: “Hence, we expect less disruption to the BSM encoding than other auxiliary tasks (e.g. reconstruction)”. Reconstruction is not the only possible auxiliary task. What happens when using a better auxiliary loss, such as action-dependent contrastive coding?\n* Ablation study is incomplete -- DBC-ID and DBC-normed-ID are not included.\n* “(Left) Our method significantly improves upon DBC [36] and is on par with SAC” this looks like it isn’t true, as SAC achieves a higher performance in the end.\n* The method isn’t evaluated on the same benchmarks as DBC, with comparable levels of distraction (e.g. natural video). As a result, it is unclear what the performance improvement would be in those settings.\n\n\n### Questions:\n* “two states that only differ in rewards far in the future should still not be bisimilar” - Standard bisimilarity is defined over immediate, not future rewards. Could you elaborate?\n* The assumption that the reward is normalized between 0 and 1 -- is this taken into account in the experiments?\n* The sentences 193-195 are unclear, could you clarify?\n* Eq. 11: x \\in R^d is this d overloaded?\n* Eq. 11: what is B_c? The ball with radius c? It looks like c is not defined.\n* Eq. 12: Does M(B_c) mean the space of probability distributions over the ball with radius c?\n* Lemma 3 assumes a stationary distribution over states, but this distribution changes as the policy changes. How does this influence the theory/results?\n* Figure 1: what is magnitude? The magnitude of the embedding?\n\n### Minor comments:\n* Below are some suggested references on state representation learning that could be relevant to include in related work.\n* Plots are hard to read in black and white\n* While the motivation for requiring more informative rewards is clear (as mentioned in strengths), it is not motivated why this takes the form of intrinsic rewards, as compared to e.g. reward propagation. \n* SAC does not fail in sparse-cartpole with N_m=1, it just takes longer to converge\n\n\n### Overall:\n\nThe score is mostly based on the following:\n* The choice of inverse dynamics (vs action-dependent contrastive coding or another type of auxiliary loss) should be properly motivated, either theoretically or experimentally.\n* Missing references should be included\n* Evaluation should be performed on the same benchmarks as DBC, e.g. with natural video distraction.\n\nI would raise the score if these are addressed.\n\n#### Missing references:\n* Biza et al, \"Online Abstraction with MDP Homomorphisms for Deep Learning\"\n* Grimm et al, \"The value equivalence principle for model-based reinforcement learning\"\n* Le Lan et al, \"Metrics and continuity in reinforcement learning\"\n* Ravindran & Barto, \"Approximate Homomorphisms: A framework for non-exact minimization in Markov Decision Processes\"\n* Taylor et al, \"Bounding Performance Loss in Approximate MDP Homomorphisms\"\n* Van der Pol et al, \"Plannable Approximations to MDP homomorphisms: equivariance under actions\"\n\n#### Suggested references\n* Agarwal et al, \"Contrastive Behavioral Similarity Embeddings for Generalization in Reinforcement Learning\"\n* Corneil et al,  “Efficient model-based deep reinforcement learning with variational state tabulation”\n* Dadashi et al, Offline Reinforcement Learning with Pseudometric Learning\n* Francois-Lavet et al, “Combined reinforcement learning via abstract representations”\n* Watter et al,  “Embed to control: A locally linear latent dynamics model for control from raw images”\n\n#### Update\nThank you for the rebuttal, and for the update on the experiments. I think adding the new experiments to the paper will make the paper more convincing. \nWhen reviewing, I only had a few issues with this paper: the lack of motivation for the choice of inverse dynamics (as opposed to another auxiliary loss), the missing evaluations on benchmarks with natural video distractions (for a fair comparison to DBC), the somewhat limited scope (due to the focus on DBC), some missing references, and the paper not being fully self-contained.\nIn the author response, the authors motivate their choice of inverse dynamics, clarify that their results can be useful for bisimulation metric learning in general, promise to add the references and amend the paper where it is unclear, and in a second reply give an update on benchmarks with natural video distractions, with positive results. \n\nThe authors thus address my original concerns. However, reviewer pEeB raised concerns about the theoretical analysis that are important to address. I have therefore decided keep my original score.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "#### Limitations: \nYes, the authors discuss limitations of their work.\n#### Societal: \nYes, the authors discuss potential societal impact. A minor suggestion would be to add a comment on how improving RL (as this paper aims to do!) could impact society, e.g. in terms of labor markets or autonomous weapons.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "VXCRwpoCGT2", "reviewer": "Reviewer_pEeB", "summary": "The paper analyzes properties of bisimulation metrics, identifies potential issues in learning with sparse rewards, and proposes methods to fix the issues.", "questions": "", "limitations": "", "rating": 5, "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": "Learning good latent representation is important for RL problems with high dimensional observations, and recent approaches based on bisimulation metrics have shown promising results. This paper analyzes the existence of unique on-policy bisimulation metrics and their connection to value function bounds. Based on the properties, the paper proposes a normalization method to fix some embedding issues of prior methods in sparse reward environments, and it further proposes improvements using intrinsic rewards and inverse dynamics.\n\nThe ideas to analyze on-policy bisumlation metrics is very reasonable, and the proposed methods seem to be both theoretical and practical appealing, but the paper (particular Section 3) is very difficult to read, and many key concepts are either missing or hard to follow. Below are some issues with the paper in the current status.\n\nAs the paper claimed, one of the key contributions is the convergence condition of bisimulation metric in learning. However, there is never a formal statement of what is the \"sufficient condition for the existence of a unique bisimulation metric\". This condition seems to be mentioned after Theorem 3, but the discussion is very confusing. From the statement of Theorem 3, \"the sufficient conditions\" seem to refer to cT<1, p=1, and the support of the approximate dynamics is closed. But later in (11) the main condition seems to refer to the boundedness of the metric. The boundedness of the unique metric is mentioned in Theorem 3, but as a property instead of a sufficient condition. After checking the proof of Theorem 3, the boundedness is indeed not used in proving the existence of the metric. So why does the boundedness of metric becomes the sufficient condition? Does the boundedness have any relationship to the closedness of the support? Even more confusing, in all discussions after Section 3 the condition is referred to \"embedding normalization\" without any definition nor explanation. Does normalization mean the requirements (11)? How does the normalization actually work in learning with parameterized embedding functions?\n\nThere are some missing definitions and typos which highly affect readability. There is no definition of the diameter before Lemma 1, no definition for the different forms of metrics used in Theorem 4 besides in the appendix. In the paragraph before Theorem 3, one metric d with hat but another one without hat. What is the metric d appearing in (11)? There is no indication on what the sup in (14) over on which doesn't seem to be over policies.\n\nThe idea of investigating the danger of on-policy bisimulation with sparse rewards in 3.2.3 is very interesting, but the discussion is very unclear and may need further explanation. The section claims that the challenges of degeneration is only an issue with on-policy bisimulation, but there is no comparison with policy-independent bisimulation. If policy-independent bisimulation doesn't have the issue, then why prior methods don't work? The paper also claims that this degeneration might be the reason for poor performance of DBC and the proposed methods fix the degeneration, but the simulation results actually show exploding behaviors instead of collapsing. How does the analysis explains exploding behaviors of DBC? The proposed intrinsic rewards and inverse dynamics methods seem reasonable, but there does not seem to have any connection to the analysis in 3.2.3 and why these methods can resolve the degeneration issue.\n\nAs said above, I like the ideas of the paper but would not recommend to accept it given the issues mentioned above. I will consider raising the score if the paper can be revised to address the above points.\n\n\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Parts of the limitations are mentioned in the paper, but most of the analysis are only on the existence of the bisimulation metric, and there is a big gap for practical algorithms to actually converge to the unique metric even if it exists.", "review_text": "", "extra_scores": {"ethical_concerns": "No.", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "15 hours"}, "extra_text": {}}, {"review_id": "V8UcgHcomxH", "reviewer": "Reviewer_XWno", "summary": "This paper investigates learning better representations for deep reinforcement learning algorithms using bisimulation metric (BSM), such that the representations are more efficient and robust from distraction (e.g. noise). The proposed method improves the Deep Bisimulation for Control (DBC) algorithm for tasks with uninformative rewards and the authors show their method is competitive in those environments.\n\n", "questions": "", "limitations": "", "rating": 6, "confidence": 1, "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": "This paper focuses on learning more robust and efficient representations for deep reinforcement learning algorithms. The analysis and resulting method build on previous work of the Deep Bisimulation for Control (DBC) algorithm but improves for environments with sparse rewards in particular. The authors analyzed theoretically embedding collapse and explosion could happen with sparse rewards and proposed remedies to fix them. The authors showed that the resulting algorithm was not only better with sparse rewards but also more stable when the observations are noisy. The topic is definitely closely related and significant to the community. Although it builds on DBC, the analysis and proposed algorithm are novel to the best of my knowledge. The paper is also well-written and easy to follow in general. Some questions: 1.can you talk more about embedding collapsing and exploration? 2.the experiments are conducted on sparse cartpole and continuous mountain cars, which are relatively low-dimension, maybe evaluate and compare the methods in MuJuCo tasks with sparse rewards? 3. fig.3, why the proposed method is a bit worse or almost the same as sac? could you explain more? especially for cartpole-swingup-sparse, why with IR it is still a bit worse than SAC?", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "See the main review.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "12"}, "extra_text": {}}, {"review_id": "SZKIGUy4SwH", "reviewer": "Reviewer_FWDg", "summary": "In this paper, the author proposed a robust bisimulation metric that could combat the collapse and explosion issues. The authors analyzed the cause of these issues theoretically and further proposed  a set of remedies including a norm constraint on the representation space, and intrinsic rewards and latent space regularization. The authors did experiments in the gym environment to verify the proposed methods.", "questions": "", "limitations": "", "rating": 7, "confidence": 2, "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": "Pros:\n1. The authors modify the previous metrics and provide a generalized metric. The analysis of the new metric compared with the old one is solid.\n2. In the experiments, the newly proposed method shows better robustness against sparsity compared with DBC.\n3. The authors further show the policy-dependence may cause caution collapse in on-policy bisimulation.\n\nCons:\n1. The experiments have simple/similar dynamics. It would be of beneficial for readers to know whether the approach would still gain obvious edges when the dynamics of the task is complex, such as a sparse version of the spider or walker in gym. Now the main challenge seems only coming from the dimension of the noise.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "n/a", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "10"}, "extra_text": {}}]

Author rebuttal:

Meta-review: Three knowledgable reviewers recommended acceptance of the paper (2x accept, 1x weak accept) and one reviewer recommended (weak) rejection of the paper. The authors addressed most of the reviewers' concerns in their rebuttal but some concerns were not resolved. In the discussion about the paper we came to the conclusion that the paper can provide several interesting insights but needs to address several of the reviewer concerns the camera ready version. I am therefore recommending acceptance of the paper and at the same time strongly advise the authors to carefully adjust the paper to address the remaining reviewers' concerns (in particular reviewer pEeB's concerns regarding clarity and theoretical statements).
Decision: Accept (Poster)