Towards Robust Bisimulation Metric Learning
Year: 2021

Abstract:
Learned representations in deep reinforcement learning (DRL) have to extract taskrelevant 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.

heading: 1 Introduction
text: Complex reinforcement learning (RL) problems require the agent to infer a 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 a 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; even more pertinent aspects, such as two mugs with slightly different patterns and shapes, should be treated nearly the same by the policy. In other words, many states are equivalent for a given task, with their differences being task-irrelevant distractors. Thus, a natural approach to generalization across environmental changes is constructing a representation that is invariant to such nuisance variables – effectively “grouping” such equivalent states together.
One method of obtaining such a representation uses the notion of a 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., in 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.
∗Equal contribution.
35th Conference on Neural Information Processing Systems (NeurIPS 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 a given task [10], resulting in a 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 a premature collapse of the embedding space, grouping states that should not be 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 unconstrained2. In 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 a norm constraint and (2) we prevent embedding collapse by altering the encoder training method.
Contributions We (i) generalize theoretical value function approximation bounds to a broader family of BSMs, with learned dynamics; (ii) analyze the BSM loss formulation and identify sources of potential embedding pathologies (explosion and collapse); (iii) show how to constrain representation space to obtain better embedding quality guarantees; (iv) devise theoretically-motivated training 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.

heading: 2 Related Work
text: Our work builds on two major aspects of DRL: representation learning, particularly task-specific encodings following the BSM, and methods for sparse reward environments, (e.g., intrinsic motivation).
Representation Learning for Control In the context of RL, there is a spectrum of representations from task-agnostic to task-specialized. One 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 so for model-free RL [33, 17, 30]. 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-supervised 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 homomorphisms [38, 39], previous efforts have exploited various structural regularities (e.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 -close states to 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 dynamics 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, we build on the use of task-specific bisimulation relations [19], which define equivalence classes of states based on rewards and transition probabilities. However, obtaining equivalence classes is brittle, as our estimates of these quantities may be inexact. Hence, Ferns et al. [13, 14, 15] instead turn to a 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 bisimulation-derived state aggregation. While DeepMDP [17] proves that its representation provides an upper bound on the bisimulation distance, DBC [53] directly imposes the bisimulation metric
2Embedding 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 uninformative 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 an extrinsic reward signal. This is often done by providing intrinsic motivation via a self-derived reward, resulting in “curiosity-driven” behaviour [8, 4]. Such approaches include surprise [1, 5] (where experiencing unexpected dynamics is rewarded), and empowerment [40, 20] (where the agent prefers states in which it has more control).
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 an intrinsic reward [36, 8, 43], thus augmenting the sparse extrinsic rewards and encouraging exploration. We also regularize the latent space by learning an inverse dynamics model on top of it, which does not rely on the extrinsic reward signal and has previously been used for distraction-robust representation learning in DRL [3, 36].

heading: 3 Technical Approach
text: 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 as a single point (embedding collapse). Based on our analysis, we propose an extension of deep bisimulation metric learning, with theoretically-motivated constraints on the optimization objective, including alterations to the forward dynamics model, intrinsic motivation, and inverse dynamics regularization.

heading: 3.1 Preliminaries
text: In this work, we consider a discounted Markov Decision Process (MDP) given by a tuple, 〈S,A,P, R, ρ0〉. At the beginning of each episode, an initial state, s0 ∈ S , is sampled from the initialstate distribution ρ0 over the state space S . Then, at each discrete time-step t ≥ 0, an agent takes an action, at ∈ A, according to a policy π(at|st). As a result, the MDP transitions to the next state according to a transition distribution P(st+1|st,at). The agent collects a scalar reward, rt = R(st,at), from the environment according to a bounded reward function, R : S × A → [Rmin, Rmax]. In infinite- and long-horizon settings, a discount factor, γ ∈ [0, 1), is used to calculate the agent’s discounted return in a given episode, G = ∑ t≥0 γ
trt. RL algorithms aim to find an optimal policy, π∗ := argmaxπ∈Π E[G], for a class of stationary policies Π. In high-dimensional, continuous state (or observation) spaces, this learning problem is rendered tractable via a state encoder (e.g., a neural network), φ : S → Rn, which is used to learn a policy of the form π(a|φ(s)). The following (pseudo3)-metric, based on the Wasserstein metric (see Appendix A for a review), is of particular relevance to this work. Distances are assigned to state pairs, (si, sj) ∈ S × 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 : S ×A → [0, 1] and c ∈ (0, 1) for continuous MDPs:
d(si, sj) = max a∈A
(1− c)|R(si,a)−R(sj ,a)| + cW1(d)(P(·|si,a),P(·|sj ,a)). (1)
An earlier version of this metric for finite MDPs used separate weighting constants cR, cT ≥ 0 for the first and second terms respectively, and required that cR + cT ≤ 1 [13]. Here, when the weighting
3Pseudo-metrics are a generalization of metrics that allow distinct points to have zero distance. For simplicity, we broadly use the term "metric" in the remainder of this paper at the cost of imprecision.
constant cT of the W1(d) term is in [0, 1), the RHS is a 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 a 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 1 in Appendix B. Notice that c (or cT ) determines a timescale for the bisimulation metric, weighting the importance of current versus future rewards, analogously to the discount factor γ. More recently, an on-policy bisimulation metric (also called π-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-policy bisimulation metric [10]). Given a fixed policy π, the following on-policy bisimulation metric exists and is unique:
dπ(si, sj) := |rπi − rπj | + γW1(dπ)(Pπ(·|si),Pπ(·|sj)), (2)
where rπi := Ea∼π[R(si,a)] and Pπ(·|si) := Ea∼π[P(·|si,a)].
Zhang et al. [53] proposed to learn a similar on-policy bisimulation metric directly in the embedding space via an MSE objective. They proposed an algorithm for jointly learning a policy π(a|φ(s)) with an on-policy bisimulation metric. Below, we define a generalized variant of their objective:
J(φ) := 1 2 E [( d̂π,φ(si, sj)− |rπi − rπj | − γW2(‖·‖q1)(P̂π(·|φ(si)), P̂π(·|φ(sj))) )2] , (3)
where d̂π,φ(si, sj) := ‖φ(si)− φ(sj)‖q2 and they used (q1 = 2, q2 = 1). Notice that the recursion induced by this objective is different from prior metrics in three ways; (i) a W2 metric was used instead of aW1 metric sinceW2 has a convenient closed form for Gaussian distributions when q1 = 2, (ii) the distance function used for the bisimulation metric and the Wasserstein metric are 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 a unique metric exists for all Wasserstein metrics Wp when necessary (see Assumption 1), (ii) study losses that use a matching metric, i.e., q1 = 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̂(d̂π(φn, P̂))(si, sj) := d̂π(si, sj ;φn + αn∇φJ(φn), P̂), (4)
where αn is a learning rate. However, in practice, π and P̂ may also be updated in training, which is of particular relevance when they have form π(a|φ(s)) and P̂(φ(s′)|φ(s),a) as in [53]. In the next section, we will discuss conditions under which joint updates to a policy π and a metric defined by state encodings φ may or may not converge in practical settings:
lim n→∞
F (n)(π, d̂π,φ) ? = dπ∗ , (5)
where dπ∗ is the on-policy bisimulation metric for the optimal policy π∗.

heading: 3.2 Theoretical Analysis
text: In this section, we generalize prior value function bounds for bisimulation metrics to cases where arbitrary weighting constants cR, cT are used, and further derive a value function approximation (VFA) bound for V π rather than V ∗ unlike prior work [13, 14, 53]. We then describe constraints on forward dynamics models for convergence to a unique metric in joint training, and derive VFA bounds as a 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.

heading: 3.2.1 Value Function Bounds
text: An important feature of bisimulation metrics is their relation to value functions since a provably tight connection implies guarantees in VFA. Similarly to previous bounds for policy-independent bisimulation metrics [13, 14], Castro [10] showed that given any two states (si, sj), the following bound holds for on-policy bisimulation metrics (see Definition 2): |V π(si)− V π(sj)|≤ dπ(si, sj). As discussed earlier, differently from previous approaches, Zhang et al. [53] used a 2-Wasserstein metric due to practical reasons. Here, in order to generalize previous value function bounds, we assume the existence of a unique bisimulation metric for p-Wasserstein metrics. Assumption 1 (A1, p-Wasserstein bisimulation metric). For a given cT ∈ [0, 1), cR ∈ [0,∞) and p ≥ 1, the following bisimulation metric exists and is unique:
dπ(si, sj) := cR|rπi − rπj | + cTWp(dπ)(Pπ(·|si),Pπ(·|sj)). (6)
Remark 1. If p = 1, or both the environment and policy are deterministic, A1 holds.
Theorem 1 (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 cT ∈ [0, 1) and p ≥ 1, define γ = min(cT , γ). Given A1, the bisimulation distance between a pair of states upper-bounds the difference in their values:
cR|V π(si)− V π(sj)| ≤ dπ(si, sj) + 2cR(γ − γ)
(1− γ)(1− cT ) , ∀(si, sj) ∈ S × S. (7)
Note that Thm. 3 of [10] is a special case with cR = p = 1 and cT = γ. Suppose cT ≥ γ; we observe that for a degenerate metric dπ = 0, the corresponding value function is a constant V π(s) = c. We will discuss the dangers posed by this relation in Section 3.2.3.
Theorem 2 (Generalized VFA bound). Let rewards be bounded as R ∈ [0, 1] and Φ : S → S̃ be a function mapping states to a finite partitioning S̃ such that Φ(si) = Φ(sj)⇒ dπ(si, sj) ≤ 2 , which produces an aggregated MDP 〈S̃,A, P̃, R̃, ρ̃0〉. For any cT ∈ [0, 1), let γ = min(cT , γ). Given A1,
|V π(s)− Ṽ π(Φ(s))| ≤ 2 cR(1− γ) + 2(γ − γ) (1− γ)(1− cT ) , ∀s ∈ S. (8)
This result generalizes previous performance bounds for V ∗ to non-optimal policies. Previous bounds also assumed cT ≥ γ, while we generalize to arbitrary cT ∈ [0, 1). Here, the second term of the upper bound characterizes the penalty paid for “myopic" bisimulation (i.e., cT < γ) in VFA error guarantees. We further show in the next section that cT < γ may be desirable when approximate forward models are used (see Thm. 4). Indeed, in Appendix C, we not only connect large cT to high variance in embedding norms, but also find surprisingly strong results when empirically investigating the use of cT < γ. We speculate that with sufficient modelling capacity and a critic V (φ) being trained based on γ, the metric space may hold information about the environment regarding both timescales γ and cT .

heading: 3.2.2 Bisimulation Metrics with Approximate Forward Dynamics
text: 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 a 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 on which the Banach fixed-point 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 compact4 state space was assumed to ensure completeness of met [14]. In practice, when using an approximate forward dynamics model P̂ : S×A →M(S ′)5, if compactness
4A metric space is compact if and only if it is totally bounded and complete [14]. 5M(X ) denotes the space of all probability distributions over X .
is not guaranteed, the same convergence guarantees do not apply. Thus, ideally, S ′ is a compact subset of S . Here, we formalize these restrictions on latent representations and the forward dynamics model to guarantee convergence to a unique metric. The following states a constraint on the BSM solution space given approximate dynamics. Lemma 1 (Diameter of S is bounded). Let d : S × S → [0,∞) be any bisimulation metric:
diam(S; d) := sup si,sj∈S×S
d(si, sj) ≤ cR
1− cT (Rmax −Rmin). (9)
The above lemma holds for all bisimulation metrics (on-policy or otherwise) that are based on exact dynamics, if said metrics exist. Let us also consider an 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 ×A →M(S ′), cT ∈ [0, 1), and cR ∈ [0,∞):
d̂π(si, sj) := cR|rπi − rπj | + cTW1(d̂π)(P̂π(·|si), P̂π(·|sj)). (10)
Next, we provide a sufficient condition for the existence of a unique metric d̂π based on approximate dynamics, which also satisfies the upper-bound of Lemma 1. Meeting this condition shrinks the solution space of metrics being searched to a set known to contain dπ . Theorem 3 (Boundedness condition for convergence). Assume S is compact. If the support of an approximate dynamics model P̂ , i.e., S ′ = supp(P̂), is a closed subset of S, then there exists a unique on-policy bisimulation metric d̂π of the form Eq. (10), and this metric is bounded:
supp(P̂) ⊆ S ⇒ diam(S; d̂π) ≤ cR
1− cT (Rmax −Rmin). (11)
If this condition is not satisfied at any point during training, the system may diverge. Indeed, we confirm empirically in Sec. 4 that the absence of a 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, Bc, with diameter given in Lemma 1, such that the following are true:
φ(s) ∈ Bc = {x ∈ Rn | ‖x‖q ≤ cR(Rmax −Rmin)
2(1− cT ) }, ∀s ∈ S (12)
P̂(·|φ(s),a) ∈M(Bc), ∀(s,a) ∈ 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. A natural question that arises from this view concerns the degree to which modelling errors affect VFA bounds. Theorem 4 (VFA bound in terms of model error). Consider the same conditions as in Theorem 2, except that cT ∈ [γ, 1), p = 1, and Φ(si) = Φ(sj)⇒ d̂π,φ(si, sj) = ‖φ(si)− φ(sj)‖q ≤ 2 ̂. Then:
|V π(s)− Ṽ π(Φ(s))|≤ 1 cR(1− γ)
( 2 ̂+ Eφ +
2cR 1− cT Er + 2cT 1− cT EP ) ,∀s ∈ S. (14)
where Eφ := ‖d̂π,φ − d̂π‖∞ is the metric learning error, Er := ‖r̂π − rπ‖∞ is the reward approximation error, and EP := sups∈SW1(dπ)(Pπ(·|s), P̂π(·|s)) is the state transition model error.
See Appendix D for details, including a generalized version for cT ∈ [0, 1), listed as Corollary 3. Consider an ideal case where ground-truth rewards rπ are available and the metric is learned perfectly (i.e., Eφ = Er = 0). Then, we observe that the choice of cT defines a trade-off between VFA error due to (a) forward model error EP , and (b) “myopic” bisimulation (see Thm. 2, cT < γ), where the BSM timescale is shorter than that of the discounted return.

heading: 3.2.3 On the Dangers of On-policy Bisimulation
text: Suppose that at training time, |rπi − rπj |= 0 for all pairs of states, e.g., during early training in a sparse reward setting. Then, unlike the policy-independent bisimulation metric, the on-policy formulation has a degenerate solution at diam(S; dπ) = 0, regardless of the structure of the underlying MDP. Lemma 2 (A reason for caution in on-policy bisimulation). On-policy bisimulation metrics of the form Eq. (6) have an upper bound determined by their policy:
diam(S; dπ) ≤ cR
1− cT sup i,j |rπi − rπj |. (15)
Due to policy-dependence, the target metric dπ changes with each policy update during joint training. As a result of this difficulty, convergence to a unique fixed point (i.e., a unique metric) for a 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 dπ∗ after the policy itself reaches a fixed point, namely, the optimal policy π∗. 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 π(a|φ(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 dispersion6 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 ρπ denote the stationary distribution over states, and νπ the distribution over pairs of states, (si, sj) sampled independently from ρπ . Then;
µπbd := E(si,sj)∼νπ [dπ(si, sj)] µ π rd := E(si,sj)∼νπ [|r π i − rπj |]. (16)
In the low dispersion case, µπrd ≈ 0, 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 low-dispersion rewards). Assume deterministic transitions and the existence of a stationary distribution ρπ over states. Given a bisimulation metric of the form Eq. (6):
µπbd = cR
1− cT µπrd. (17)
6Statistical dispersion is an umbrella term used to describe measures of variability or diversity (e.g., variance).
Clearly, a collapsed state encoder, φ∗(si) = φ0, loses all information about the state. These observations motivate us to extend the method with inverse dynamics-based regularization and intrinsic rewards based on forward prediction errors (see Sec 3.3), since they promote µπbd, µ π rd > 0 respectively. For empirical evidence of the relation in Lemma 3 at training-time, see Figure 1. Relationships between variances are also discussed in Appendix C.

heading: 3.3 Intrinsic Rewards and Inverse Dynamics
text: The core principle of DBC is to construct a representation such that a given φ(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: (1) using intrinsic rewards (IR), and (2) regularizing the latent space with inverse dynamics (ID) learning.
Curiosity-Driven 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 t, we redefine the reward signal to be rt = rI,t + rE,t, where rE,t is the extrinsic reward (i.e., the original environmental reward) and rI,t is the IR. Following prior work [8, 36, 43], we use the forward model error in the latent space to compute intrinsic rewards: r̃I,t := ηr||φ̂µ(st,at) − φ(st+1)||22/(2n), where φ̂µ(st,at) = Eφ(st+1)∼P̂(·|φ(st),at)[φ(st+1)] is the mean of the predicted distribution from the forward dynamics model and ηr > 0 is a hyperparameter controlling the intrinsic reward scale. We finally clamp to a fixed maximum Rmax,I , so that rI,t := min(Rmax,I , r̃I,t). 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-stationary, resulting in the target of the BSM learning changing as well; however, as the agent explores and can better predict its environment over time, the IR (and thus its influence on the metric learning process) should fade.
Latent Inverse Dynamics Prior We also consider regularizing the learned latent state space, by having it learn to predict the inverse dynamics of the world, following prior work [3, 36]. Let gI : Rn × Rn → A be an inverse model, which predicts the action at ∈ A that changed st to
st+1 using the learned latent space: ât := gI(φ(st), φ(st+1)). This can be trained from observed transitions (st,at, st+1), via Jd(φ, gI ; ηd) := ηd||at − ât||1/na, where ηd > 0 weighs the ID loss importance, and A = [−1, 1]na for continuous control tasks. This loss prevents removal of information that pertains to the agent’s ability to control the environment, but it is balanced against the loss driving φ to mimic the BSM. In this sense, we are placing a regularizing prior on φ via an auxiliary task, such that it adheres to the BSM requirement when possible, but 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).
Summary In this section, we (1) provided generalized bounds connecting the value function to the on-policy BSM (Sec. 3.2.1); (2) showed that the BSM satisfies a fixed latent diameter, and devised a normalization for enforcing 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 IR and ID to mitigate it (Sec. 3.2.3, 3.3).

heading: 4 Experiments
text: In this section, we seek answers to the following questions concerning our main hypotheses:
1. Do the embedding collapse and explosion issues predicted theoretically 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 OpenAI 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 3D robotics benchmarks from the DeepMind Control Suite [45].
Sparse Noisy Cartpole First, we modified the Cartpole-v0 task, in which an agent tries to keep a pole upright. To increase sparsity, we shrank the angular extent in which the pole must be to earn rewards to 1% of its standard value. Then, to mimic distractors, we concatenate an Nmdim(S)dimensional vector sampled from an isotropic Gaussian to the state vector. (See Appendix E.1.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 E.1.4 for results with Nm = 3.
Noisy Mountain Car We also consider the classic MountainCarContinuous-v0 task [32], in its continuous control form (see Appendix E.1.2 for details). Since the task was already highly sparse, we modified only the distraction aspect, as in 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 Pendulum-v0 task to have a 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 a set of 3D robotics tasks based on the MuJoCo physics simulator [47]. We include results from two DMC tasks, namely, cheetah-run (dim(S)=18,dim(A)=6) and cartpole-swingup-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 a 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.

heading: 5 Conclusion
text: 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, it 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.
Acknowledgments We thank Sven Dickinson, Allan Jepson, and Amir-massoud Farahmand for helpful discussions. The support of NSERC (CGSD3-534955-2019) and Samsung AI Research is gratefully acknowledged.

Review summaries:
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.

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.  

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.