{'fraction_non_alphanumeric': 0.04592694344533411, 'fraction_numerical': 0.011260324976955194, 'mean_word_length': 4.52561669829222, 'pattern_counts': {'":': 0, '<': 6, '<?xml version=': 0, '>': 4, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 3, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
Towards Robust Bisimulation Metric Learning

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.


Introduction
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.
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 unconstrained 2 . 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.

Related Work
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 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].

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

Preliminaries
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, s 0 ∈ 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, a t ∈ A, according to a policy π(a t |s t ). As a result, the MDP transitions to the next state according to a transition distribution P(s t+1 |s t , a t ). The agent collects a scalar reward, r t = R(s t , a t ), from the environment according to a bounded reward function, R : S × A → [R min , R max ]. 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 γ t r t . 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 → R n , which is used to learn a policy of the form π(a|φ(s)).
The following (pseudo 3 )-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, (s i , s j ) ∈ 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: An earlier version of this metric for finite MDPs used separate weighting constants c R , c T ≥ 0 for the first and second terms respectively, and required that c R + c T ≤ 1 [13]. Here, when the weighting constant c T of the W 1 (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 c T ) 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].
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: where d π,φ (s i , s j ) := φ(s i ) − φ(s j ) q2 and they used (q 1 = 2, q 2 = 1). Notice that the recursion induced by this objective is different from prior metrics in three ways; (i) a W 2 metric was used instead of a W 1 metric since W 2 has a convenient closed form for Gaussian distributions when q 1 = 2, (ii) the distance function used for the bisimulation metric and the Wasserstein metric are different (L 1 and L 2 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 W p when necessary (see Assumption 1), (ii) study losses that use a matching metric, i.e., q 1 = q 2 = 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 W 2 term. The resulting gradient updates may be considered as approximate fixed-point iteration in the space of metrics: 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: where d π * is the on-policy bisimulation metric for the optimal policy π * .

Theoretical Analysis
In this section, we generalize prior value function bounds for bisimulation metrics to cases where arbitrary weighting constants c R , c T 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.

Value Function Bounds
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 (s i , s j ), the following bound holds for on-policy bisimulation metrics (see Definition 2): |V π (s i ) − V π (s j )|≤ d π (s i , s j ).
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 c T ∈ [0, 1), c R ∈ [0, ∞) and p ≥ 1, the following bisimulation metric exists and is unique: Remark 1. If p = 1, or both the environment and policy are deterministic, A1 holds.
For an on-policy bisimulation metric given by Eq. (6), for any c T ∈ [0, 1) and p ≥ 1, define γ = min(c T , γ). Given A1, the bisimulation distance between a pair of states upper-bounds the difference in their values: Note that Thm. 3 of [10] is a special case with c R = p = 1 and c T = γ. Suppose c T ≥ γ; 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.
This result generalizes previous performance bounds for V * to non-optimal policies. Previous bounds also assumed c T ≥ γ, while we generalize to arbitrary c T ∈ [0, 1). Here, the second term of the upper bound characterizes the penalty paid for "myopic" bisimulation (i.e., c T < γ) in VFA error guarantees. We further show in the next section that c T < γ may be desirable when approximate forward models are used (see Thm. 4). Indeed, in Appendix C, we not only connect large c T to high variance in embedding norms, but also find surprisingly strong results when empirically investigating the use of c T < γ. 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 c T .

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 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 compact 4 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  [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 c T = 0.5, Eq. (17) correctly predicts the ratio µ π bd /µ π rd 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 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.
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 : 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: 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, B c , with diameter given in Lemma 1, such that the following are true: 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, where E φ := d π,φ − d π ∞ is the metric learning error, E r := r π − r π ∞ is the reward approximation error, and E P := sup s∈S W 1 (d π )(P π (·|s), P π (·|s)) is the state transition model error.
See Appendix D for details, including a generalized version for c T ∈ [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 φ = E r = 0). Then, we observe that the choice of c T defines a trade-off between VFA error due to (a) forward model error E P , and (b) "myopic" bisimulation (see Thm. 2, c T < γ), where the BSM timescale is shorter than that of the discounted return.

On the Dangers of On-policy Bisimulation
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: 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 dispersion 6 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, (s i , s j ) sampled independently from ρ π . Then; 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): However, latent normalization immediately improves performance (top row). Further, combining it with IR or IR+ID improves performance on all tasks. Shading shows standard error over 10 seeds.
Clearly, a collapsed state encoder, φ * (s i ) = φ 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.

Intrinsic Rewards and Inverse Dynamics
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 r t = r I,t + r E,t , where r E,t is the extrinsic reward (i.e., the original environmental reward) and r I,t is the IR. Following prior work [8,36,43], we use the forward model error in the latent space to compute intrinsic rewards: 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 R max,I , so that r I,t := min(R max,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 g I : R n × R n → A be an inverse model, which predicts the action a t ∈ A that changed s t to s t+1 using the learned latent space: a t := g I (φ(s t ), φ(s t+1 )). This can be trained from observed transitions (s t , a t , s t+1 ), via J d (φ, g I ; η d ) := η d ||a t − a t || 1 /n a , 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

Experiments
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 N m dim(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 N m = 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.

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

A Wasserstein Distances
Definition 5 (Wasserstein metric [50]). Let d : X × X → [0, ∞) be a distance function and Ω the set of all joint distributions with marginals µ and λ over the space X; Definition 6 (Dual formulation of the Wasserstein metric [50]). Let d : X × X → [0, ∞) be a distance function, and µ and λ marginals over the space X; This dual formulation takes a simple form for p = 1: where Note that the 2-Wasserstein metric W 2 ( · 2 ) (or simply W 2 ) has a closed-form for Gaussians [35]: where · F denotes the Frobenius norm. We can observe that for point masses (i.e., Σ i , Σ j → 0), the 2-Wasserstein metric is equivalent to the Euclidean distance between the two points. Lemma 4 (p-Wasserstein Inequality [50]). For any two distributions µ, λ, if p ≤ q: Lemma 5 (Bounds on Wasserstein distances [41]). For any two distributions µ, λ over a space X, for all p ≥ 1: B Proofs Remark 1. If p = 1, or both the environment and policy are deterministic, A1 holds.
Proof. The existence proof is virtually identical to the proof of Thm. 3.12 of [14], except it discards max a∈A operations in favor of expectations under a policy π. We need to show that the following fixed-point update is a contraction: , P π (·|s j )), and invoke the Banach fixed-point theorem to show the existence of a unique metric.
Next, we consider the case where both P and π are deterministic, such that P π is a delta distribution. Observe that for point masses, W p (d)(δ(s i ), δ(s j )) = d(s i , s j ), due to Definition 5 of the Wasserstein metric. Then: Then, fixed point iterations that update the metric as d (n+1) (s i , s j ) ← F(d (n) )(s i , s j ) will converge for finite MDPs.
Lemma 6 (p-Wasserstein value difference bound). For an on-policy bisimulation metric given by Eq. (6), for any c T ∈ [γ, 1) and p ≥ 1, the bisimulation distance between a pair of states upper-bounds the difference in their values: Proof. The proof follows similarly to the proofs of Thm. 5.1 of [13] and Thm. 3 of [10]. We will prove by induction. Consider the following updates: π )(P π (·|s i ), P π (·|s j )). We need to show that the following holds for all n ∈ N: Then, Eq. (24) holds by taking a limit n → ∞. The base case holds since: In the general case: Notice that by the induction hypothesis, c R V (n) (s) is a 1-Lipschitz function with respect to the distance function d Then, since γ ≤ c T by assumption, c R γ c T V (n) (s) is also 1-Lipschitz. Using the dual form of the W 1 metric in Eq. (20): where the last inequality is due to Lemma 4.

Lemma 7 (Value function difference bound for different discount factors [37]
). Consider two otherwise identical MDPs with different discount factors γ 1 ≤ γ 2 , and a bounded reward function R ∈ [0, 1]. Let V π γ denote the value function for policy π given discount factor γ.
, ∀s ∈ S. For an on-policy bisimulation metric given by Eq. (6), for any c T ∈ [0, 1) and p ≥ 1, define γ = min(c T , γ). Given A1, the bisimulation distance between a pair of states upper-bounds the difference in their values: Proof. A bisimulation metric with c T ≤ γ can be viewed as approximating a value function for another MDP with γ = c T . We will make use of this view and apply Lemma 6 with Lemma 7 to derive the above relation.

Then,
where · ∞ is the supremum norm over S. Due to Lemma 6, c R V (s) is a 1-Lipschitz function with respect to the distance function d π . Then, since γ ≤ c T by assumption, c R γ c T V (s) is also 1-Lipschitz. Using the dual form of the W 1 metric: c R |r π (s) − r π (z)| + c T W 1 (d π )(P π (s |s), P π (s |z))dξ(z) Thus, taking the supremum on the LHS over the state space S: , ∀s ∈ S.
which implies F is a c T -contraction. It remains to prove that the distance is bounded. First, note that due to Lemma 5: Then, similarly to Lemma 1, Proof. In the on-policy case, the bound in Lemma 1 can be much tighter depending on the policy: 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): Proof. For point masses, W p (d)(δ(s i ), δ(s j )) = d(s i , s j ): Simply taking an expectation under ν π , due to the stationarity assumption: C Notes on Reward Scale, c R and c T In recent work [10,17,53], various forms of bisimulation metrics have been presented with different scaling constants; (c R = 1 − c, c T = c) as in Definition 1, and (c R = 1, c T = γ) as in Definition 2. Here, we aim to add clarity to the effect of these choices and how they relate to the reward scale. First, note that due to Lemma 1, setting c R = 1 − c T serves to ensure that d ∈ [0, 1] when the reward range is specified as (R max = 1, R min = 0) as in [13,14].
Corollary 1 (Policy-independent mean and variance bounds). Due to Lemma 1, the first two moments of the random variable given by bisimulation distance have the following bounds independently of any policy π: Conversely, if (c R = 1, c T = γ) as in Zhang et al. [53], the formulation allows large distances between embeddings since γ is commonly set to a value close to 1. Large pairwise distances imply large norms and high variance (see Corollary 1), which may cause instabilities in optimization (especially in the absence of norm constraints), considering the compactness conditions discussed in Sec. 3.2.2. Definition 7 (Variance of distances and reward differences). Given a stationary distribution ρ π over states, and ν π the distribution over pairs of states, (s i , s j ) sampled independently from ρ π : Proposition 1 (On-policy variance bound). Assume a deterministic environment. Given c T ∈ [0, √ 0.5), the variance of the optimal on-policy bisimulation distance for an objective of the form Eq. (3) can be upper-bounded as follows: while for all c T ∈ [0, 1), the bound in Eq. (31) applies.
Proof. Given a deterministic environment: When c T ≥ √ 0.5 (as in Zhang et al. [53]), the above bound is loose. However, c T < √ 0.5 provides a convenient upper bound: where the equality is due to Lemma 3.
Tighter bounds on the variance of the on-policy bisimulation metric may be important, since the statistics of r π undergo change throughout training between policy updates. Hence, it is desirable to remove the dependence of the bound from the µ π rd term with a choice of c T = 0.5, such that tighter bounds can be obtained. This choice renders the formulation more robust to changes in the scale of the expected rewards. The resulting bound for c T = 0.5 has a simpler form: (34) Indeed, in Figure 4, we show that such a choice can stabilize the DBC [53] algorithm significantly, resulting in higher overall performance.

D Value Bounds with Model Error
Our goal in this section is to characterize the errors induced by using approximate dynamics and an imperfect encoder, with respect to estimating both the ground-truth on-policy bisimulation metric, as well as preserving the value function with the encoded state.
For this section, we first remark on the difference between three forms of the PBSM, for some fixed policy π: • d π (s i , s j ) is the ground-truth bisimulation metric, defined as the fixed point of the operator defined in Assumption 1 (when it exists uniquely).
• d π (s i , s j ; P) is the fixed point of the same operator but with approximate dynamics (i.e., using P and r π instead of the true P and r π ; we leave out conditioning on r π for brevity).
• d π,φ (s i , s j ) := ||φ(s i ) − φ(s j )|| q is a non-negative function of states, dependent on an encoder function φ. For example, φ may be produced by a metric learning process, such as by stochastic minimization of the objective in Eq. 3.
Notice that, if we attempt to learn φ with a metric learning process based on approximate dynamics, then the best we can do is obtain d π,φ → d π , which can still have some irreducible error. On the other hand, even with perfect approximate dynamics, the metric learning process may be incomplete, meaning d π,φ = d π . We therefore hope to characterize the error into two types: dynamics approximation error and metric learning error.
Next, we define three types of model errors, relating to the quality of encoding and forward dynamics prediction (in terms of state and reward). Definition 8 (Model and Encoder Errors). The bisimulation distance approximation error E φ , transition probability error E P , and reward prediction error E r are given by where ||·|| ∞ is the supremum (or uniform) norm over states and d π is the fixed-point bisimulation metric with the W p distance defined by using P π and r π , instead of the true dynamics P π and r π . 7 Note that we have defined our forward dynamics model errors here irrespective of φ (i.e., it may be used by P and/or r π , or not). Our goal is to use these errors to bound the difference in the MDP value function induced by the approximate nature of the environmental model and state encoder. First, we consider how the model errors affect the optimal policy-dependent bisimulation distance that we can obtain given our approximate forward dynamics.
First, by the Wasserstein triangle inequality [11], as for the reward difference: Second, the convexity of d p implies that, Third, recall that when supp( P) ⊆ S, due to Lemma 5, we have: Then, the difference in distances can be bounded by: We can then plug these into the difference between the true and approximate policy-dependent bisimulation distances: where the second-last inequality follows by taking the supremum over states for both sides.
For the remainder of this section, we assume p = 1.
Corollary 2 (Bisimulation distance error with p = 1). Let p = 1, with the remaining conditions as in Lemma 9. Then Proof. When p = 1, we have a p = a 1 = 1, giving the expression above.
This bounds the error between the true on-policy bisimulation distance and the optimal approximate bisimulation distance (i.e., the best distance function we could hope to achieve with our encoder, given the error in our forward dynamics model). However, ultimately, we wish to bound the error in the value function in terms of d π,φ , not just d π (to take the error of the encoder φ into account, as well as that of the dynamics model). First, we can bound the true bisimulation distance in terms of the encoder and model error as follows: Lemma 10 (Bound on bisimulation distance with encoder error). Consider the same conditions as Corollary 2. Then Proof.
Thus, if we can relate d π to the value function, we can also do so for d π,φ , as a function of model error.
Finally, we look at bounding the difference in the state value function, using the approximate bisimulation distance defined through the learned encoder (i.e., our partitioning Z is defined via d π,φ ). Let be the aggregation radius in φ-space (meaning the maximum diameter with respect to d π,φ per partition subset, or equivalence class, is at most 2 ): Notice that bounds the maximal diameter of the partition cells with respect to the learned metric, using φ, rather than the ground truth bisimulation distance.
Proof. Following the proof of Lemma 8, we have that 8 where the last line used Lemma 10.
Rather than bound the ground-truth on-policy bisimulation distance d π , we have instead bound the estimated distance d π,φ , using both approximate predictive dynamics and error in the metric learning process.
Notice that as we shrink the size of the equivalence classes in the partition, so → 0, we get Φ → φ. This tells us that information about the value of a state is preserved by the encoder φ, as long as the error in the forward dynamics model and metric learning process is small. Further, in the low error case, if the vectors φ(s i ) and φ(s j ) are close, we are guaranteed they have similar value under π, with the difference growing only linearly with the error.
Proof. The proof follows similarly to the proof of Thm. 2. Suppose c T < γ: where the second and third inequalities are due to Thm. 4 and Lemma 7 respectively. For γ ≤ c T , we recover Thm. 4, hence, for all c T ∈ [0, 1) and s ∈ S: 8 Notice that using the Equation ( †) allows us to recover the value bound with model error from [53]: |V π (s) − V π (Φ(s))|≤ First, we consider modifications to the Cartpole-v0 task. In the standard version, the agent receives a constant reward of +1 at each time-step for as long as the cart-pole system is upright between [−θ term , θ term ] degrees. If the pole falls below this range, the episode terminates. To make the task more challenging, we introduce a second parameter θ rew θ term , such that the agent receives a reward only if the pole is between [−θ rew , θ rew ] degrees. We refer to this task as Sparse Cartpole and set θ term = 12 • with θ rew = 0.01 θ term . Additionally, we consider a noisy version of Sparse Cartpole to mimic distractors. In particular, we concatenate an N m dim(S)-dimensional vector sampled from an isotropic Gaussian to the state vector. The resulting task is referred to as Noisy Sparse Cartpole. Thus, the encoder φ must learn to embed functionally similar states in close proximity, despite the distractions, and maintain a well-behaved embedding, despite reward sparsity. While we provide sparse reward signals at training time to make the learning problem more difficult, we evaluate the resulting models in the standard environment based on θ term , since this provides a lower variance return signal and still allows us determine whether the task has been solved.

E.1.2 Noisy Mountain Car Additional Details
We next tested on the Noisy Mountain Car task [32], implemented as MountainCarContinuous-v0 in the OpenAI Gym [7], and modified to concatenate N m dim(S) noise dimensions to the observed state to simulate distraction. Briefly, the agent controls a car that should reach the top of a mountain, but has an engine of insufficient power to attain it directly. It must therefore learn a sequence of actions that build enough momentum to complete the task. The reward signal is highly uninformative: a small negative reward is given at every step, unless the task is solved, in which case a large positive reward is given. For N m > 0, this task has both noisy distractors and high sparsity, making it rather challenging.
Note that only methods with intrinsic reward were able to solve the task (see Fig. 2 and Fig. 5), and that DBC without normalization was also unable to complete it. We remark that all methods rely on the maximum policy entropy RL formulation, which is known to improve exploration [23, 24, 55]; nevertheless, our results suggest that curiosity-driven exploration, induced by intrinsic rewards based on predictive error, is at least complementary to such techniques.

E.1.3 Sparse Pendulum Task Details
For the Pendulum-v0 task, we implement similar modifications to those in SparseCartpole. The standard Pendulum task starts with a pole in downright position, and provides negative rewards proportional to |θ pend |, degrees away from upright position. Our SparsePendulum instead provides a reward of +1 only when the pendulum is between [−θ rew , θ rew ] degrees, where θ rew , and does not provide a reward otherwise. NoisySparsePendulum similarly concatenates to state vectors a noise vector of N m times the original dimensionality. Note that we evaluate with an environment with θ rew = 1 • , to reduce variance in the evaluation reward.
Results (see Fig. 6) show that our method performs comparably to DBC. In the presence of high distraction (N m = 10), DBC-based approaches can do much better than SAC, which is no longer able to solve the task.

E.1.4 Plots with Additional Distraction
In Fig. 5, we show results for Sparse Cartpole (left) and Mountain Car (right), with an even higher level of distraction (N m = 3). On Sparse Cartpole, only DBC-normed-IR and DBC-normed-IR-ID were able to solve the task, with the latter being slightly more stable, while DBC-normed performed significantly better than any methods without normalization. On Mountain Car, all methods struggle to solve the task at this level of distraction; however, DBC-normed-IR-ID performs significantly better than the others, showing the utility of the inverse dynamics regularization.

E.1.5 Hyper-parameters and Architectures for OpenAI Gym Tasks
All models use the code generously released alongside the DBC paper [53] (CC-BY-NC 4.0 licensed), with default hyper-parameters and architectures according to the released code, unless otherwise specified. DBC and our modifications of it are all built on top of the Soft Actor-Critic implementation therein. Actor, critic, and encoder learning rates were all 0.001, using the Adam optimizer. A replay buffer of 50K was used, with a batch size for all training steps of 512. All encoders were implemented as multilayer perceptrons (MLPs) with four layers, except for the SparsePendulum task where two layers were used for all normalized approaches (unnormalized DBC still performed better with four). Encoder feature dimensionality was set to dim(φ(s)) = 50. When inverse dynamics prediction was needed, we implemented g I as an MLP with two hidden layers (of size 256 and 128), with ELU activations. Note that for the predictive approximate transition model P, we use a deterministic predictor. Following the implementation of [53], distances in reward and predicted encoded state, for the bisimulation metric, were computed with a Huber loss, defining the value of q as a function of the distance. However, embedding normalization was computed with the L 2 norm. A discount of γ = 0.99 was always used. In all tasks, rewards are bounded as R ∈ [0, 1].
We remark that all evaluation rewards in plots (per seed) are computed as averages over 10 episodes.
We used a maximum intrinsic reward clamping value of R max,I = 0.1 for all Gym tasks, and set η r and η d per task as in Table 1 below. Hyper-parameters not left at default (R max,I , η r , η d , and number of encoder layers and latent dimensionality) were set by searching over a small, manually defined set of values.

E.2 DeepMind Control Suite
For the DeepMind Control Suite [45], our encoder model architecture is identical to the open-source code repository released by [53]. Namely, a 3x3 convolutional layer with stride 2 is followed by another 3x3 convolution with stride 1 (both with 32 channels), before a fully-connected layer with 50-dimensional output and layer normalization. ReLU activations are used between neural layers. When an inverse dynamics model is used, we use the same architecture as those used for the OpenAI experiments (described in Appendix E.1.5), namely, a two-layer MLP. Differently from [53], to speed up training, we run 16 environments in parallel, all of which add experience tuples to a shared replay buffer. After every 16 environment steps (i.e., each parallel step), we apply 2 gradient updates. Our hyperparameters for inverse dynamics and intrinsic motivation are given in the last column of Table  1, and were selected by searching over a small, manually defined set of values.

E.3 Computational Resources and Timing
All training and evaluation was done on a small set of NVIDIA GPUs (GTX 1080 TI, Titan X, or RTX 2080 TI), less than 10 in total and shared with other users.
OpenAI Gym tasks were run with multiple seeds per GPU (up to the GPU memory limit) during training. In this parallel training context, which allowed us to complete a Gym task for all methods and seeds (for a single distraction level) within roughly a day on 2-4 GPUs, we obtain the following approximate timings (in seconds per training iteration