# Efficient Embedding of Semantic Similarity in Control Policies via Entangled
Bisimulation

Martin Bertran  Walter Talbott  Nitish Srivastava  Joshua Susskind

###### Abstract

Learning generalizeable policies from visual input in the presence of visual
distractions is a challenging problem in reinforcement learning. Recently,
there has been renewed interest in bisimulation metrics as a tool to address
this issue; these metrics can be used to learn representations that are, in
principle, invariant to irrelevant distractions by measuring behavioural
similarity between states. An accurate, unbiased, and scalable estimation of
these metrics has proved elusive in continuous state and action scenarios. We
propose entangled bisimulation, a bisimulation metric that allows the
specification of the distance function between states, and can be estimated
without bias in continuous state and action spaces. We show how entangled
bisimulation can meaningfully improve over previous methods on the Distracting
Control Suite (DCS), even when added on top of data augmentation techniques.

Machine Learning, ICML

]————^2_2#1

##  1 Introduction

Learning control policies from primarily visual input is an important task in
many real world applications (Lillicrap et al., 2015; Jaderberg et al., 2016;
Kalashnikov et al., 2018; Espeholt et al., 2018). While, in principle,
reinforcement learning (RL) algorithms can eventually learn a suitable control
policy with sufficient experience, poor sample efficiency and the time and
monetary cost of data collection still present significant hurdles in practice
(Duan et al., 2016; Henderson et al., 2018; Kaiser et al., 2019). Moreover,
these agents may additionally suffer from poor generalization to unseen, but
semantically equivalent, scenarios (Tachet et al., 2018; Song et al., 2019;
Cobbe et al., 2019), these distractions can also impact learning efficiency.

One promising approach to deal with distractors while also preserving task
relevant information is the use of bisimulation metrics (Ferns et al., 2011;
Ferns & Precup, 2014; Castro, 2020), where states that are indistinguishable
with respect to future reward trajectories are grouped together (Zhang et al.,
2020) (deep bisimulation for control, DBC), or states which produce the same
action sequence under the optimal policy are similarly grouped (Agarwal et
al., 2021) (policy similarity metrics, PSM). Bisimulation tackles
generalization by incorporating and measuring important invariances of the
policy and environment dynamics into the representation learning process.
Bisimulation metrics have rich theoretical underpinnings (Ferns et al., 2011;
Ferns & Precup, 2014), but efficient and unbiased estimation of this metric is
an active research area; with current research efforts aimed at stabilizing
this metric learning process in practice (Kemertas & Aumentado-Armstrong,
2021).

In this work, we provide a unifying view of currently used bisimulation
metrics, their empirical approximations, and provide a mathematically formal
way of producing unbiased sample estimates of bisimulation metrics in
continuous, stochastic environments for a wide variety of desired invariance
specifications. We show how this contribution can meaningfully improve results
on current approximations on the Distracting Control Suite benchmark (Stone et
al., 2021); as well as an experimental analysis on how bisimulation interacts
with existing augmentation techniques and agent capacity.

##  2 Related Work

The works in (Laskin et al., 2020; Kostrikov et al., 2020) make explicit use
of low-level data augmentation to improve the robustness of the learned policy
to predefined input perturbations. These works have proven to be very
effective at tackling these issues, since the perturbations are pre-specified
and can be generated without additional interaction with the environment,
improving sample efficiency. One of the drawbacks of these approaches is that
the invariances must be known a-priori.

Contrastive learning (Srinivas et al., 2020; Chen et al., 2020; Hjelm et al.,
2018) allows the algorithm designer to specify positive and negative matches
in representation space, and embed a similarity measure based on this by
maximizing agreement between positive matches, and minimizing it w.r.t
negative matches; contrastive learning has been shown to improve sample
efficiency. Observation prediction and reconstruction (Hafner et al., 2019;
Sekar et al., 2020) also provides a rich auxiliary training signal, but forces
the agent to model and reconstruct task-irrelevant distractors, which may be a
significant disadvantage on natural scenarios (Zhang et al., 2020; Agarwal et
al., 2021).

The work in (Ferns et al., 2004, 2011) proposed a bisimulation metric that
required the computation of an auxiliary adversarial policy, as well as the
computation of a $1-$Wasserstein coupling between state transition
distributions, this metric was shown to be a value function of an augmented
(bi-state) environment (Ferns & Precup, 2014). The work in (Castro, 2020)
addressed the overly pessimistic and computationally costly concern of
optimizing over an adversarial policy by proposing policy bisimulation, where
the policy-averaged rewards and transition functions are used instead.
Estimating the $1-$Wasserstein coupling remained a challenging issue, so the
work analyzed how to efficiently compute this quantity for deterministic state
transitions in Markov decision processes (MDP).

The works of (Agarwal et al., 2021; Zhang et al., 2020) (DBC, PSM) tackle the
issue of approximating policy bisimulation (or related notions) over partially
observable Markov decision processes (POMDP). Both approximations
unfortunately suffer from drawbacks where their proposed bisimulation
optimization objectives have nontrivial biases for estimating Wasserstein
distances, and provably lose important metric properties like state self
similarity. Although both have shown to be effective methods in practice, the
differences in implementation, methodology and architecture make comparisons
between them nontrivial. In this work, we will address these shortcomings, and
compare the corrected methods to examine how effective reward- (DBC) or
policy-based (PSM) bisimulation metrics are at dealing with semantically
equivalent distractions.

##  3 Contributions

  * •

We discuss and analyze policy-based bisimulation metrics for continuous action
and state spaces and their practical implementations and biases in a unifying
framework.

  * •

We provide an alternative bisimulation metric, entangled bisimulation, which
allows for simple, unbiased estimation for any distance measure (reward-based,
action-based, or otherwise) under mild assumptions.

  * •

We show the effectiveness of the proposed bisimulation measures on the DCS
benchmark, and how it improves performance beyond standard image augmentation.

##  4 Preliminaries

###  4.1 Bisimulation Metrics

Bisimulation is a principled way of comparing (latent) states and learning or
inducing invariances in a policy or environment; it is a distance metric over
state pairs. Here we first introduce bisimulation in the context of a finite-
time Markov decision process (MDP), described by the tuple
$\mathcal{M}\coloneqq(\mathcal{Z},\mathcal{A},\mathcal{P},\mathcal{R},\gamma)$,
where $\mathcal{Z}$ is the latent state space and $\mathcal{A}$ the action
space, the distribution $\mathcal{P}(\bm{z}^{\prime}\mid\bm{z},\bm{a})$
contains the probabilities of transitioning to state $\bm{z}^{\prime}$ after
performing action $\bm{a}$ on state $\bm{z}$, the agent then receives reward
$r=\mathcal{R}(\bm{z},\bm{a})\in\mathbb{R}$. The goal of the agent is to
maximize discounted rewards via policy maximization
$\max_{\pi}\mathbb{E}_{\mathcal{P}}[\sum_{t=1}^{\infty}\gamma^{t}R(\bm{z}_{t},\bm{a}_{t})]$,
where actions are sampled according to a learned policy
$\bm{a}\sim\pi(\cdot|\bm{z})$. We later show how bisimulation metrics can be
applied and efficiently computed in a partially observable Markov decision
process (POMDP) where our agent’s latent states $\bm{z}$ define a learned MDP.

Bisimulation metrics are specified by three key components:

  * •

An equivalence relation $E\in\mathcal{Z}\times\mathcal{Z}$ between states
specifying a desired invariance.

  * •

A metric transformation $\mathcal{F}:\mathbb{M}\rightarrow\mathbb{M}$ that has
a fixed point on a metric
$d:\mathcal{Z}\times\mathcal{Z}\rightarrow\mathbb{R},\,d\in\mathbb{M}$
satisfying the equivalence relation (i.e.,
$F(d)=d,\,d(\bm{z},\bm{z^{\prime}})=0\iff(\bm{z},\bm{z^{\prime}})\in E$). We
call this metric the bisimulation metric.

  * •

For continuous state and action spaces, we additionally require a computable
pointwise estimate of the metric transformation,
$\widehat{\mathcal{F}}(d):\mathcal{Z}\times\mathcal{Z}\rightarrow\mathbb{R}$
used to learn the bisimulation metric. The methods discussed in this work make
use of an auxiliary objective to learn the bisimulation metric of the form

$J(\bm{z},\bm{z^{\prime}})\coloneqq
21[2{d(\bm{z},\bm{z^{\prime}})-\widehat{\mathcal{F}}(d)(\bm{z},\bm{z^{\prime}})}.$
(1)

As a motivating example, we consider $\pi$-bisimulation as proposed in
(Castro, 2020), which compares the state and reward dynamics introduced by the
policy $\pi$. We use this definition to show the above key concepts in
bisimulation. We then observe practical implementations of this bisimulation
metric and extract key insights used to design bisimulation metrics that
measure a wide variety of invariances, and how to ensure that the empirical
estimates used to learn these metrics in practice are easy to compute,
consistent, and unbiased.

Consider the $\pi$-bisimulation relation as proposed in (Castro, 2020), which
compares the state and reward dynamics introduced by the policy $\pi$. Letting
$R(\bm{z},\pi)\coloneqq\mathbb{E}_{\pi(\bm{a}|\bm{z})}R(\bm{z},\bm{a})$ and
$P(\bm{z}_{+}|\bm{z},\pi)\coloneqq\mathbb{E}_{\pi(\bm{a}|\bm{z})}P(\bm{z}_{+}|\bm{z},\bm{a})$
be the policy-averaged reward and state transitions functions respectively,
(Castro, 2020) defines a $\pi$-bisimulation relation as follows:

######  Definition 4.1.

(Castro, 2020) Given an MDP $\mathcal{M}$, an equivalence relation
$E^{\pi}\subseteq\mathcal{Z}\times\mathcal{Z}$ is a $\bm{\pi}$-bisimulation
relation if whenever $(\bm{z},\bm{z^{\prime}})\in E^{\pi}$ the following
properties hold

$\begin{array}[]{rl}R(\bm{z},\pi)&=R(\bm{z^{\prime}},\pi)\\\
P(\cdot|\bm{z},\pi)&=P(\cdot|\bm{z}^{\prime},\pi).\end{array}$ (2)

Further, we say that states $\bm{z},\bm{z^{\prime}}$ are $\bm{\pi}-$bisimilar
if $\bm{z},\bm{z^{\prime}}\in E^{\pi}$ for some $E^{\pi}$. We say that
$d^{\pi}\in\mathbb{M}:\mathcal{Z}\times\mathcal{Z}\rightarrow\mathbb{R}^{+}$
is a $\bm{\pi-}$bisimulation metric if
$d(\bm{z},\bm{z^{\prime}})=0\leftrightarrow(\bm{z},\bm{z^{\prime}})$ are
$\pi-$bisimilar. Note that this definition implies that two states are
identical if both states are expected to collect the same reward throughout
every step of the rollout policy $\pi$.

To find such a bisimulation metric, the work in (Castro, 2020) makes use of a
fixed point formulation using the $1-$Wasserstein metric,

$\begin{array}[]{cc}\mathcal{W}_{1}(P,Q;d)&\coloneqq\inf\limits_{\gamma\in\Gamma(P,Q)}\mathbb{E}_{\bm{z}_{+},\bm{z^{\prime}}_{+}\sim\gamma}[d(\bm{z}_{+},\bm{z^{\prime}}_{+})],\\\
\end{array}$ (3)

where $P,Q$ are two distributions over the same support, and $\Gamma(P,Q)$ is
the set of all couplings of $P,Q$111joint distribution
$\gamma(\bm{z}_{+},\bm{z^{\prime}}_{+})$ is a coupling of distributions
$P(\bm{z}_{+}),Q(\bm{z}^{\prime}_{+})$, denoted $\gamma\in\Gamma(P,Q)$, if
$\int\gamma(\bm{z}_{+},\bm{z^{\prime}}_{+})d\bm{z^{\prime}}_{+}=P(\bm{z}_{+})$
and
$\int\gamma(\bm{z}_{+},\bm{z}_{+})d\bm{z^{\prime}}_{+}=Q(\bm{z}^{\prime}_{+})$.

######  Theorem 4.2.

(Castro, 2020) Define the metric operator
$\mathcal{F}^{\pi}:\mathbb{M}\rightarrow\mathbb{M}$ as

$\begin{array}[]{rl}\mathcal{F}^{\pi}(d)(\bm{z},\bm{z^{\prime}})=&|R(\bm{z},\pi)-R(\bm{z^{\prime}},\pi)|\\\
&+c\mathcal{W}_{1}(P(\cdot|\bm{z},\pi),P(\cdot|\bm{z}^{\prime},\pi);d).\end{array}$
(4)

Then, with $c\in[0,1)$, $\mathcal{F}^{\pi}$ has at least a fixed point
$d_{\pi}$ which is a bisimulation metric.

The Wassertein metric is used to optimally compare differences in state
transitions, and plays a vital role in ensuring that states are self-similar
($d_{\pi}(\bm{z},\bm{z})=0\,\forall\,\bm{z}$). That is, bisimulation measures
differences in results attributed to the agent’s actions in the environment
and its current latent state, and optimally discounts the environment’s
subsequent stochastic transitions. Without this Wasserstein coupling,
distances between states would be inflated due to environment
stochasticity222for example, if both $P$ and $P^{\prime}$ were equally likely
to transition to the same pair of dissimilar states $\bm{z}_{1},\bm{z_{2}}$ we
would have $\mathbb{E}_{\bm{z}_{+},\bm{z}^{\prime}_{+}\sim P\times
P^{\prime}}d(\bm{z}_{+},\bm{z^{\prime}}_{+})=\frac{1}{2}d(\bm{z}_{1},\bm{z_{2}})$,
while $\mathcal{W}_{1}(P,P^{\prime};d)=0$, and, most crucially, the ordering
of distances between states would no longer encode potential behavioural
differences.

We now look at the propsed computable approximations of
$\mathcal{F}^{\pi}(d)(\bm{z},\bm{z^{\prime}})$. This poses a significant
challenge since it requires the computation of the policy-averaged state
transition $P(\bm{z}_{+}|\bm{z},\pi)$, and, most importantly,
$\mathcal{W}_{1}(P,Q;d)$ over a multivariate continuous distribution, which is
a notoriously hard estimation problem in general (Chizat et al., 2020).

Two recent approaches to embed $\pi-$bisimulation metrics are DBC (Zhang et
al., 2020) and PSM (Agarwal et al., 2021). We first analyze these approaches
before presenting our solution.

The work in (Zhang et al., 2020) (DBC) uses the following relaxation for
bisimulation estimation, which uses the $\mathcal{W}_{2}$ distance
$\begin{array}[]{rl}d_{DBC}(\bm{z},\bm{z}^{\prime})&\coloneqq||\bm{z}-\bm{z}^{\prime}||_{1}\\\
\mathcal{F}_{DBC}(d)(\bm{z},\bm{z^{\prime}})&\coloneqq|R(\bm{z},\pi)-R(\bm{z}^{\prime},\pi)|\\\
&+c\mathcal{W}_{\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}2}(P(\bm{z}_{+}|\bm{z},\pi),P(\bm{z}^{\prime}_{+}|\bm{z}^{\prime},\pi);\ell_{1})],\\\
\widehat{\mathcal{F}}_{DBC}(d)(\bm{z},\bm{z^{\prime}})&\coloneqq|R(\bm{z},\bm{a})-R(\bm{z}^{\prime},\bm{a}^{\prime})|\\\
&+c\mathcal{W}_{\color[rgb]{0,0,1}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,1}2}(P(\bm{z}_{+}|\bm{z},\bm{a}),P(\bm{z}^{\prime}_{+}|\bm{z}^{\prime},\bm{a}^{\prime});\ell_{1}),\\\
\bm{a},\bm{a^{\prime}}&\sim\pi(\cdot|\bm{z})\times\pi(\cdot|\bm{z}^{\prime}).\end{array}$
(5)

Where they define the bisimulation metric to be the $\ell_{1}$ distance
between states, learn the latent state dynamics $P$ to conform to
$d_{DBC}(\bm{z},\bm{z}^{\prime})=\mathcal{F}_{DBC}(d)(\bm{z},\bm{z^{\prime}})$,
and plug in single-sample estimates of
$R(\bm{z},\pi),R(\bm{z}^{\prime},\pi),P(\bm{z}_{+}|\bm{z},\bm{a}),$ and
$P(\bm{z}_{+}|\bm{z}^{\prime},\bm{a})$, with actions
$(\bm{a},\bm{a^{\prime}})$ independently sampled from each policy
$\bm{a},\bm{a^{\prime}}\sim\pi(\cdot|\bm{z})\times\pi(\cdot|\bm{z}^{\prime})$
to create their computable approximation
$\widehat{\mathcal{F}}_{DBC}(d)(\bm{z},\bm{z^{\prime}})$. The key insight is
that $\mathcal{W}_{2}(P,Q;\ell_{1})$ can be computed in closed form if $P,Q$
are multivariate Gaussian distributions, so the authors make use of this.

The specific choice of the distance function links the overall scale of latent
states to the environment reward, since doubling the reward function doubles
the average distance between states in the DBC formulation. A more pressing
issue is that the proposed relaxation does not satisfy self similarity. That
is, $\widehat{\mathcal{F}}_{DBC}(d_{DBC})(\bm{z},\bm{z})>0$, since

$\mathbb{E}_{\pi(\bm{a}|\bm{z})\times\pi(\bm{a}^{\prime}|\bm{z})}||R(\bm{z},\bm{a})-R(\bm{z},\bm{a}^{\prime})||_{1}>0,$
(6)

unless
$R(\bm{z},\bm{a})=R(\bm{z},\bm{a^{\prime}})\,\forall\bm{a},\bm{a^{\prime}}$ in
the support of $\pi(\cdot|\bm{z})$, and, likewise

$\mathbb{E}_{\bm{a},\bm{a^{\prime}}\sim\pi(\cdot|\bm{z})\times\pi(\cdot|\bm{z})}[\mathcal{W}_{2}(P(\cdot|\bm{z},\bm{a}),P(\cdot|\bm{z},\bm{a}^{\prime});||\cdot||_{1})]>0.$
(7)

The choice of the pointwise approximation $\widehat{\mathcal{F}}_{DBC}(d)$
means that there is no metric $d$ and transition function $P$ minimizing Eq. 1
that satisfies the fixed point relation.

PSM proposes a related bisimilarity metric where they compare the average
action of the optimal policy in place of average reward comparison333We
specifically discuss their formulation for continuous action and state spaces
$\begin{array}[]{rl}\mathcal{F}_{PSM}(d)(\bm{z},\bm{z^{\prime}})\coloneqq&||\mathbb{E}_{\pi(\bm{a}|\bm{z})}[\bm{a}]-\mathbb{E}_{\pi(\bm{a}^{\prime}|\bm{z}^{\prime})}[\bm{a}^{\prime}]||_{1}\\\
&+c\mathcal{W}_{1}(P(\cdot|\bm{z},\pi),P(\cdot|\bm{z}^{\prime},\pi);d),\\\
\widehat{\mathcal{F}}_{PSM}(d)(\bm{z},\bm{z^{\prime}})\coloneqq&||\mathbb{E}_{\pi(\bm{a}|\bm{z})}[\bm{a}]-\mathbb{E}_{\pi(\bm{a}^{\prime}|\bm{z}^{\prime})}[\bm{a}^{\prime}]||_{1}\\\
&+cd(\bm{z}_{+},\bm{z}^{\prime}_{+}),\\\
\bm{a},\bm{a^{\prime}}\sim&\pi(\cdot|\bm{z})\times\pi(\cdot|\bm{z}^{\prime}),\\\
\bm{z}_{+},\bm{z^{\prime}}_{+}\sim&P(\cdot|\bm{z},\bm{a})\times
P(\cdot|\bm{z}^{\prime},\bm{a}^{\prime}).\end{array}$ (8)

We note that the expected policy action
$\mathbb{E}_{\pi(\bm{a}|\bm{z})}[\bm{a}]$ is explicitly computed in most agent
architectures, so there is no estimation error with this part of the estimate.
The Wasserstein distance is replaced by an expectation over independent state
transition and policy distributions; this allows for easy computation, but
does not consider all potential couplings as in the Wasserstein distance.
Similarly to DBC, the proposed relaxation does not satisfy self-similarity,
that is, $\widehat{\mathcal{F}}_{PSM}(d_{PSM})(\bm{z},\bm{z})>0$ since
$\mathbb{E}_{\begin{subarray}{c}\bm{a},\bm{a^{\prime}}\sim\pi(\cdot|\bm{z})\times\pi(\cdot|\bm{z}^{\prime}),\\\
\bm{z}_{+},\bm{z^{\prime}}_{+}\sim P(\cdot|\bm{z},\bm{a})\times
P(\cdot|\bm{z}^{\prime},\bm{a}^{\prime})\end{subarray}}d_{PSM}(\bm{z}_{+},\bm{z}^{\prime}_{+})>0$
unless all independently reachable states $\bm{z}_{+},\bm{z^{\prime}}_{+}$
have $0$ distance. Since self-similarity is not satisfied, then we can
conclude that $d_{PSM}$ is not a metric.

###  4.2 Entangled Bisimulation

We introduce $\epsilon$-bisimulation (entangled bisimulation) which allows for
simple, principled computation in standard RL settings without approximation
bias. We show how this formulation relates to the DBC and PSM measures
discussed in the previous paragraphs with minimal adaptations, while still
showing the desireable properties of a proper bisimulation metric. To achieve
unbiased estimation, we first introduce entangled sampling.

######  Definition 4.3.

Given two multivariate distributions $P,Q$ defined over $\mathbb{R}^{n}$, and
let $F_{P}^{-1},F_{Q}^{-1}$ be the inverse sampling functions of $P$ and $Q$
respectively444e.g., for a two-dimensional vector $\\{x_{1},x_{2}\\}\sim P$ we
have
$F_{P}^{-1}(u_{1},u_{2})=F_{P_{x_{1}}}^{-1}(u_{1}),F_{P_{x_{2}|x_{1}}}^{-1}(u_{2})$
the entangled coupling of $P,Q$, denoted $\gamma^{\epsilon}(P,Q)$ is defined
as:

$\begin{array}[]{rl}x,y=&F_{P}^{-1}(\bm{\epsilon}),F_{Q}^{-1}(\bm{\epsilon}),\;\bm{\epsilon}\sim
U_{[0,1]}^{\otimes n},\\\ x,y&\stackrel{{\scriptstyle
d}}{{\sim}}\gamma^{\epsilon}(P,Q).\end{array}$ (9)

Note that for architectures using the reparametrization trick,555for example,
to take a continuous action sample in the $[-1,1]$ range with hyperbolic
tangent squashing $\bm{a}\sim\pi(\bm{a}\mid\bm{z})$, we compute state
dependent mean and standard deviations $\mu_{i}(\bm{z}),\sigma_{i}(\bm{z})$,
and sample
$a_{i}=\tanh(\mu_{i}(\bm{z})+\epsilon_{i}\sigma_{i}(\bm{z})),\epsilon_{i}\sim\mathcal{N}(0,1)$,
the entangled coupling amounts to reusing the noise vector between two
distributions. To improve readability, we use a shortened notation to denote
two couplings of interest, namely policy and state couplings, as:

$\begin{array}[]{rl}\gamma^{\epsilon}_{\pi}=&\gamma^{\epsilon}(\pi(\cdot|\bm{z}),\pi(\cdot|\bm{z}^{\prime})),\\\
\gamma^{\epsilon}_{P}=&\gamma^{\epsilon}(P(\cdot|\bm{z},\bm{a}),P(\cdot|\bm{z}^{\prime},\bm{a}^{\prime})).\end{array}$
(10)

Where we omit the state conditioning on the policy coupling, and the state
action conditioning on the state coupling. We define a family of equivalence
relations based on any state-action metric
$G:\\!(\mathcal{Z}\\!\times\\!\mathcal{A})^{2}\rightarrow\\!\mathbb{R}^{+}\\!,G\in\\!\mathbb{M}$:

######  Definition 4.4.

Given an MDP $\mathcal{M}$, an equivalence relation
$E^{\epsilon}\subseteq\mathcal{Z}\times\mathcal{Z}$ is a
$\bm{\epsilon}$-bisimulation relation (entangled bisimulation relation) w.r.t
policy $\pi$ and state-action metric $G$ if whenever
$(\bm{z},\bm{z^{\prime}})\in E^{\epsilon}$ the following properties hold
$\forall\bm{a},\bm{a}^{\prime}\in\textit{Supp}(\gamma^{\epsilon}_{\pi})$

$\begin{array}[]{rl}G(\bm{z},\bm{a},\bm{z}^{\prime},\bm{a}^{\prime})&=0\\\
P(\cdot|\bm{z},\bm{a})&=P(\cdot|\bm{z}^{\prime},\bm{a}^{\prime}).\end{array}$
(11)

Where we require $G$ to satisfy all core metric properties, that is

######  Definition 4.5.

A function $G:(\mathcal{Z}\times\mathcal{A})^{2}\rightarrow\mathbb{R}^{+}$ is
a state-action similarity metric if
$\forall\bm{z},\bm{z}^{\prime},\bm{z}^{\prime\prime},\bm{a},\bm{a}^{\prime},\bm{a}^{\prime\prime}\in\mathcal{Z}^{3}\times\mathcal{A}^{3}$,
$G$ satisfies  $\begin{array}[]{lll}\text{Non-
negativity}:&\\!\\!\\!\\!G(\bm{z},\bm{a},\bm{z}^{\prime},\bm{a}^{\prime})&\geq
0\\\ \text{Self-similarity}:&\\!\\!\\!\\!G(\bm{z},\bm{a},\bm{z},\bm{a})&=0\\\
\text{Symmetry}:&\\!\\!\\!\\!G(\bm{z},\bm{a},\bm{z}^{\prime},\bm{a}^{\prime})&=G(\bm{z}^{\prime},\bm{a}^{\prime},\bm{z},\bm{a})\\\
\text{Triangle}&\\!\\!\\!\\!G(\bm{z},\bm{a},\bm{z}^{\prime\prime},\bm{a}^{\prime\prime})&\leq
G(\bm{z},\bm{a},\bm{z}^{\prime},\bm{a}^{\prime})\\\
\text{Inequality}:&&+G(\bm{z}^{\prime},\bm{a}^{\prime},\bm{z}^{\prime\prime},\bm{a}^{\prime\prime}).\\\
\end{array}$ (12)

We can thus use any state-action similarity metric as a basis for an
equivalence class captured by $\epsilon$-bisimulation.We can now compute the
associated $\bm{\epsilon}$-bisimulation metric by using the following fixed
point theorem

######  Theorem 4.6.

Define the metric operator
$\mathcal{F}^{\epsilon}:\mathbb{M}\rightarrow\mathbb{M}$ as
$\begin{array}[]{rl}\mathcal{F}_{\epsilon}(d)(\bm{z},\bm{z^{\prime}})=&\mathop{\mathbb{E}}\limits_{\bm{a},\bm{a}^{\prime}\sim\gamma^{\epsilon}_{\pi}}[G(\bm{z},\bm{a},\bm{z}^{\prime},\bm{a}^{\prime})\\\
&+c\mathcal{W}_{1}(P(\bm{z}_{+}|\bm{z},\bm{a}),P(\bm{z}^{\prime}_{+}|\bm{z}^{\prime},\bm{a}^{\prime});d)].\end{array}$
(13)

Then, with $c\in[0,1)$, $\mathcal{F}_{\epsilon}$ has at least a fixed point
$d_{\epsilon}$ which is an$\epsilon$-bisimulation metric.

We observe that entangled bisimulation can build stricter equivalence
relations than policy-based bisimulation.

######  Lemma 4.7.

Let $E^{\pi}$ be the largest $\pi$-bisimulation relation, and let $E^{PSM}$
likewise be average action bisimulation relation in (Agarwal et al., 2021).
The entangled bisimulation relation $E^{\epsilon}$ satisfies

$\begin{array}[]{rl}E^{\epsilon}&\subseteq E^{\pi}\\\ E^{\epsilon}&\subseteq
E^{PSM}\\\ \end{array}$ (14)

for the choice respective choice of state-action metrics

$\begin{array}[]{rl}G(\bm{z},\bm{a},\bm{z}^{\prime},\bm{a}^{\prime})&=|R(\bm{z},\bm{a})-R(\bm{z}^{\prime},\bm{a}^{\prime})|,\\\
G(\bm{z},\bm{a},\bm{z}^{\prime},\bm{a}^{\prime})&=||\mathbb{E}_{\pi(\bm{a}|\bm{z})}[\bm{a}]-\mathbb{E}_{\pi(\bm{a}^{\prime}|\bm{z}^{\prime})}[\bm{a}^{\prime}]||_{1}.\end{array}$
(15)

Finally, to address estimation of the $\mathcal{W}_{1}$ metric, the following
theorem provides an entangled upper bound that still provably satisfies the
$\epsilon$-bisimulation relation; this upper bound can be computed as a simple
average. Further, when state transition dynamics also satisfy coordinate
independence (a common feature in RL architectures), and we restrict the
learned distance functions to be of the form
$d(\bm{z},\bm{z}^{\prime})=\sum_{i}w_{i}d_{i}(z_{i},z^{\prime}_{i})$, where
$w_{i}\geq 0$, and $d_{i}(\cdot,\cdot)$ convex, then the entangled upper bound
is exact.

######  Theorem 4.8.

Given an MDP $\mathcal{M}$, policy $\pi$, and state-action metric $G$. Define
the metric operator
$\mathcal{F}_{\bar{\epsilon}}\\!:\\!\mathbb{M}\\!\rightarrow\\!\mathbb{M}$ as
$\begin{array}[]{rl}\mathcal{F}_{\bar{\epsilon}}(d)(\bm{z},\bm{z^{\prime}})=&\mathop{\mathbb{E}}\limits_{\begin{subarray}{c}\bm{a},\bm{a^{\prime}}\sim\gamma^{\epsilon}_{\pi}\\\
\bm{z}_{+},\bm{z}^{\prime}_{+}\sim\gamma^{\epsilon}_{P}\end{subarray}}[G(\bm{z},\bm{a},\bm{z}^{\prime},\bm{a}^{\prime})+cd(\bm{z}_{+},\bm{z}^{\prime}_{+})].\end{array}$
(16)

Then, $\mathcal{F}_{\bar{\epsilon}}$ has at least a fixed point
$d_{\bar{\epsilon}}$ satisfying $d_{\bar{\epsilon}}\geq d_{\epsilon}$, and
$(\bm{z},\bm{z^{\prime}})\in E^{\epsilon}\rightarrow
d_{\bar{\epsilon}}(\bm{z},\bm{z^{\prime}})=0$, in particular,
$d_{\bar{\epsilon}}(\bm{z},\bm{z})=0,\,\forall\bm{z}\in\mathcal{Z}$.

Further, if the state transition function is coordinate independent
($P(\bm{z}_{+}\mid\bm{z},\bm{a})=\Pi_{i}P_{i}(z_{i,+}\mid\bm{z},\bm{a})\,\forall\bm{z}_{+},\bm{z},\bm{a}$),
and $d_{\bar{\epsilon}}$ is of the form

$d_{\bar{\epsilon}}(\bm{z},\bm{z}^{\prime})=\sum\limits_{\begin{subarray}{c}i=[n]\\\
j=[p]\end{subarray}}w_{i,j}|z_{i}-z^{\prime}_{i}|^{j},\;w_{i,j}\geq 0\forall
i,j,\;p>0,$ (17)

then the bound is tight, that is

$\begin{array}[]{rl}d_{\epsilon}(\bm{z},\bm{z}^{\prime})&=\mathop{\mathbb{E}}\limits_{\begin{subarray}{c}\bm{a},\bm{a^{\prime}}\sim\gamma^{\epsilon}_{\pi}\\\
\bm{z}_{+},\bm{z}^{\prime}_{+}\sim\gamma^{\epsilon}_{P}\end{subarray}}[G(\bm{z},\bm{a},\bm{z}^{\prime},\bm{a}^{\prime})+cd_{\epsilon}(\bm{z}_{+},\bm{z}^{\prime}_{+})].\end{array}$
(18)

Theorem 4.8 suggests the use of the sample estimate in Eq.18 to estimate
bisimilarity embeddings in the same way one would estimate a simple value
learning objective,

$\begin{array}[]{rl}\bm{a},\bm{a^{\prime}}\sim\gamma^{\epsilon}_{\pi},&\bm{z}_{+},\bm{z^{\prime}}_{+}\sim\gamma^{\epsilon}_{P},\\\
\widehat{\mathcal{F}}_{\epsilon}(d)(\bm{z},\bm{z^{\prime}})&=G(\bm{z},\bm{a},\bm{z}^{\prime},\bm{a}^{\prime})+cd(\bm{z}_{+},\bm{z}^{\prime}_{+})\\\
J_{\epsilon}(\bm{z},\bm{z}^{\prime}))&=||d(\bm{z},\bm{z}^{\prime})-\widehat{\mathcal{F}}_{\epsilon}(d)(\bm{z},\bm{z^{\prime}})||^{2}_{2}.\end{array}$
(19)

Unlike other approaches, which use actions and/or state transitions from the
replay buffer, we sample actions and states from our entangled distributions.
We note that sampling multiple latent states from the entangled policy is
computationally efficient since it requires a single feedforward pass through
the network; multi-sample estimates are likewise supported. Note that the
entangled policy
$\gamma^{\epsilon}(\pi(\cdot|\bm{z}),\pi(\cdot|\bm{z}^{\prime}))$ satisfies
$\gamma^{\epsilon}(\pi(\cdot|\bm{z}),\pi(\cdot|\bm{z})))(\bm{a},\bm{a^{\prime}})=\pi(\bm{a}|\bm{z})\delta_{\bm{a}=\bm{a}^{\prime}}$
and is smoothly varying with $\pi$. Entangled sampling thus ensures that

$\mathop{\mathbb{E}}\limits_{\bm{a},\bm{a}^{\prime}\sim\gamma^{\epsilon}_{\pi}}G(\bm{z},\bm{a},\bm{z},\bm{a}^{\prime})=0\,\forall\pi,\bm{z},G,$
(20)

this property is not shared by independent action sampling666for continuous
variables in general, entangled coupling is a one to one mapping between
distributions; for finite, discrete distributions with cardinality
$|\mathcal{A}|$, entangled coupling visits at most $2|\mathcal{A}|$ distinct
value pairs. Equation 20 shows that self-similarity is an intrinsic property
of this estimator, indeed, we observe that every sample of the distribution
satisfies $\widehat{\mathcal{F}}_{\epsilon}(d)(\bm{z},\bm{z})=0$.

##  5 Model description

Our bisimulation state regularization technique can be added as an extra
objective on top of any RL agent that has a latent transition model and a
continuous action policy that makes use of reparametrization. To benchmark
this regularization, we use SAC (Haarnoja et al., 2018) as the basis for our
experimental setup, with minor modifications described next. For some of the
experiments, we also consider the use of the data augmentation technique
proposed in DrQ (Kostrikov et al., 2020), a state of the art augmentation
technique in RL. The model components are:

$\begin{array}[]{rl}\text{Observation encoder}&\bm{z}_{t}=h_{\theta}(o_{t})\\\
\text{Reward predictor}&\hat{r}_{t}=R_{\theta}(\bm{z}_{t})\\\ \text{Inverse
dynamics model}&\hat{\bm{a}}_{t}=A_{\theta}(\bm{z}_{t},\bm{z}_{t-1})\\\
\text{Transition
model}&P_{\theta}(\hat{\bm{z}}_{t}|\bm{z}_{t-1},\bm{a}_{t-1})\\\
\text{Critics}&Q^{i}_{\theta}(\bm{z}_{t},\bm{a}_{t}),\;i\in\\{1,2\\}\\\
\text{Policy}&\pi_{\theta}(\bm{a}_{t}|\bm{z}_{t})\\\ \text{Bisimulation
distance}&d_{\theta}(\bm{z},\bm{z}^{\prime})\end{array}$

Algorithm 1 shows the core $\epsilon$-bisimulation computation, where we use
$\overline{\bm{z}}$ to denote the gradient stop operation applied to variable
$\bm{z}$, the full optimization loss for the model parameters is computed as

$J_{SAC}+J_{ID}+\beta J_{\epsilon}.$ (21)

Here $J_{SAC},J_{ID}$ and $J_{\epsilon}$ represent the standard SAC loss,
inverse dynamics loss, and the proposed $\epsilon-$bisimulation loss, with
parameter $\beta$ controlling the relative importance of the latter objective.
The full algorithm is shown in Appendix B.

Algorithm 1 Bisimulation Algorithm

0: Latent state batch $B=\\{\bm{z}\\}_{i=1}^{n}$, policy $\pi_{\theta}$,
latent transition model $P_{\theta}(\cdot\mid,\bm{z},\bm{a})$, bisimulation
distance function $d_{\theta}(\cdot,\cdot)$, similarity pseudometric $G$

Permute states

$B^{\prime}=\\{\bm{z}^{\prime}\\}=\text{Perm}(\\{\bm{z}\\})$

Sample noise variables

$\epsilon^{A}\sim\mathcal{N}(0,I^{|A|}),\epsilon^{Z}\sim\mathcal{N}(0,I^{|Z|})$

Compute tied actions

$\bm{a}=\pi_{\theta}(\cdot|\bm{z},\epsilon^{A}),\;\bm{a}^{\prime}=\pi_{\theta}(\cdot|\bm{z}^{\prime},\epsilon^{A})$

Compute tied latent transitions

$\bm{z}_{+}=P_{\theta}(\cdot|\bm{z},\bm{a},\epsilon^{Z}),\;\bm{z}^{\prime}_{+}=P_{\theta}(\cdot|\bm{z}^{\prime},\bm{a}^{\prime},\epsilon^{Z})$

Compute bisimulation target

$\hat{d}(\bm{z},\bm{z}^{\prime})=G(\bm{z},\bm{a},\bm{z}^{\prime},\bm{a}^{\prime})+cd_{\theta}(\bm{z}_{+},\bm{z}^{\prime}_{+})$

Compute bisimulation loss

$J_{\epsilon}(\bm{z},\bm{z}^{\prime})=\frac{1}{n}\sum\limits_{\bm{z^{\prime}},\bm{z}\in,B,B^{\prime}}||d_{\theta}(\bm{z},\bm{z}^{\prime})-\overline{\hat{d}(\bm{z},\bm{z}^{\prime})}||^{2}_{2}$

$J_{\epsilon}(\bm{z},\bm{z}^{\prime})$

##  6 Experiments and Results

In this section we benchmark our improvements in bisimulation learning on a
challenging control benchmark with distractors. The goal is to evaluate the
efficacy of the improved bisimulation metrics under a variety of environment
conditions and agent capacity, and to evaluate if bisimulation provides
additional benefits beyond those achieved with standard data augmentation. We
use 6 environments in the Distracting Control Suite (DCS): Ball in cup catch
(BiC), carpole swingup (CS), cheetah run (CR), finger spin (FS), reacher easy
(RE), and walker walk (WW). Similarly to (Agarwal et al., 2021), environments
use dynamic background distractions from the DAVIS 2017 dataset (Pont-Tuset et
al., 2017), with $2$ videos used for training and $30$ for evaluation to test
for out of distribution generalization. We also include an out-of-plane camera
motion to further increase task difficulty, an example of the observations
produced can be seen on Figure 1. Additional experiments on autonomous driving
using CARLA (Dosovitskiy et al., 2017) are provided in Appendix D.

(a)

(b)

(c)

(d)

(e)

(f)

Figure 1: DCS environments with distracting backgrounds and camera. The same
semantic situation (agent position relative to environment) is shown from
varying viewpoints and backgrounds. The trained agent needs to be robust to
these perturbations.

We test $6$ training setups: SAC as a baseline, DrQ, which adds low-level
visual augmentation, DrQ $+$DBC, a reward-based bisimulation measure on top of
DrQ, $\epsilon$-R, our entagled correction for reward-based bisimulation (
$G(\bm{z},\bm{a},\bm{z}^{\prime},\bm{a}^{\prime})=|R(\bm{z},\bm{a})-R(\bm{z}^{\prime},\bm{a}^{\prime})|$)
on top of DrQ, DrQ$+$PSM, a policy-based bisimulation measure, and
$\epsilon$-$\pi$, our entagled correction for policy-based bisimulation (
$G(\bm{z},\bm{a},\bm{z}^{\prime},\bm{a}^{\prime})=||\mathbb{E}_{\pi(\bm{a}|\bm{z})}[\bm{a}]-\mathbb{E}_{\pi(\bm{a}^{\prime}|\bm{z}^{\prime})}[\bm{a}^{\prime}]||_{1}$)
on top of DrQ. All methods share a common architecture and hyperparameters,
described in Appendix C; for bisimulation-based methods, we further tune the
bisimulation strength hyperparameter $\beta$, only the results for the best
performing hyperparameters are shown.

Table 1 shows the results with background distractions and camera movement,
data augmentation in the form of DrQ improves on the SAC baseline, and
bisimulation metrics improve on this even further, with $\epsilon$-$\pi$ and
$\epsilon$-R providing the best overall results. We observe that the largest
improvement provided by entanglement is seen on the DrQ+DBC to $\epsilon$-R
comparison, in some cases boosting performance over $3\times$ over the
unmodified algorithm; it also alleviates failiure cases as seen in BiC, FS,
and RE, where DrQ+DBC is a net loss over just the DrQ baseline. Some
environments (BiC, FS, WW) show significant improvement of the bisimulation
methods w.r.t. data augmentation, though none are adversely affected by the
inclusion of bisimulation. The training dynamics of these methods are shown in
Figure 2. Additionally, Figure 3 shows random observations, along with their
closest bisimulation matches across other episodes and backgrounds on a small
subset of 10 episodes acquired on the Walker Walk environment. Semantically
similar situations are consistently identified regardless of background or
camera position.

Table 1: Episodic reward on Distracting Control Suite with DAVIS backgrounds
and hard camera distractions, $2M$ environment steps. Mean and standard
deviation computed across 5 seeds.

Method BiC CS CR FS RE WW SAC $240\pm 9$ $133\pm 20$ $36\pm 3$ $80\pm 51$
$98\pm 9$ $64\pm 17$ DrQ $311\pm 181$ $229\pm 3$ $145\pm 14$ $293\pm 72$
$613\pm 33$ $142\pm 14$ DrQ+DBC $132\pm 39$ $207\pm 5$ $134\pm 12$ $144\pm 32$
$410\pm 83$ $150\pm 21$ $\bm{\epsilon}$-R $\bm{735\pm 37}$ $270\pm 22$ $146\pm
18$ $\bm{487\pm 33}$ $574\pm 42$ $\bm{472\pm 37}$ DrQ+PSE $563\pm 94$ $248\pm
20$ $161\pm 23$ $348\pm 83$ $\bm{680\pm 29}$ $379\pm 59$
$\bm{\epsilon}$-$\bm{\pi}$ $683\pm 103$ $\bm{280\pm 20}$ $\bm{175\pm 18}$
$354\pm 105$ $666\pm 57$ $364\pm 80$

(a)

(b)

(c)

(d)

(e)

(f)

Figure 2: Episodic reward on test environments in DCS benchmark with hard
camera distractions and $30$ background videos. Standard deviations are
collected across $5$ independent seeds. (a)  Figure 3: Sample observations,
alongside their closest bisimulation matches, acquired over 10 episodes on the
Walker Walk environment with hard backgrounds and camera distractions. We
acquired $5$ independent cameras, each with its own independent motion and
distracting background, to provide novel viewpoints for the bisimulation
comparisons.

To test if network capacity plays a major role in these results, we repeat
these experiments using a smaller architecture in Table 2. Here, all hidden
layers have their size reduced from $1000$ to $200$ units. The results show a
loss of performance across the board, suggesting that capacity is being
properly utilized on the larger architecture size. Bisimulation-based methods
still show the best overall performances, with reward-based bisimulation
methods taking a comparatively larger hit from the reduced network capacity.
The same environments that previously exhibited large performance differences
between DrQ and the bisimulation-based methods continue to do so. Figure 4
shows the training dynamics in the low-capacity scenario.

Table 2: Episodic reward on Distracting Control Suite with DAVIS backgrounds
and hard camera distractions, $2M$ environment steps, low capacity network.
Mean and standard deviation computed across 5 seeds.

Method BiC CS CR FS RE WW SAC $78\pm 43$ $130\pm 28$ $31\pm 7$ $74\pm 35$
$75\pm 10$ $63\pm 10$ DrQ $261\pm 206$ $231\pm 13$ $156\pm 19$ $341\pm 94$
$532\pm 52$ $150\pm 6$ DrQ+DBC $120\pm 21$ $211\pm 6$ $162\pm 17$ $276\pm 89$
$518\pm 80$ $147\pm 8$ $\bm{\epsilon}$-R $\bm{725\pm 62}$ $251\pm 9$ $124\pm
7$ $333\pm 37$ $\bm{628\pm 41}$ $277\pm 116$ DrQ+PSE $608\pm 36$ $247\pm 33$
$\bm{171\pm 40}$ $376\pm 121$ $560\pm 90$ $\bm{378\pm 51}$
$\bm{\epsilon}$-$\bm{\pi}$ $704\pm 64$ $\bm{261\pm 18}$ $117\pm 52$
$\bm{418\pm 97}$ $\bm{628\pm 80}$ $292\pm 75$

(a)

(b)

(c)

(d)

(e)

(f)

Figure 4: Episodic reward on test environments in DCS benchmark with hard
camera distractions and $30$ background videos from the DAVIS dataset. Low
capacity architecture. Standard deviations are collected across $5$
independent seeds. Table 3: Episodic reward on Distracting Control Suite with
DAVIS backgrounds and no camera distractions, $2M$ environment steps. Mean and
standard deviation computed across 5 seeds.

Method BiC CS CR FS RE WW $\bm{\epsilon}$-R $771\pm 9$ $\bm{367\pm 29}$
${239\pm 48}$ $419\pm 129$ $642\pm 35$ $\bm{541\pm 77}$ DrQ+PSE $673\pm 328$
$348\pm 51$ $\bm{305\pm 33}$ $\bm{573\pm 131}$ $\bm{703\pm 11}$ $411\pm 138$
$\bm{\epsilon}$-$\bm{\pi}$ $\bm{782\pm 84}$ $348\pm 51$ $249\pm 58$ $528\pm
77$ $696\pm 29$ $463\pm 103$

Out of the two distractions present in our benchmark, camera movement is the
one least addressed by low-level augmentation techniques like the random
cropping added in DrQ. For this reason, we remove camera distractions and
rerun the experiments, which are shown in Table 3. Without camera
distractions, all environments significantly improve their performance.

##  7 Discussion

This paper examines current methods for extracting task-relevant state in
environments that present significant task-irrelevant distractions. We focus
on bisimulation-based methods, their practical implementations, and some of
the consistent estimation errors present in current methods for continuous
action spaces. We propose entangled bisimulation, a new bisimulation framework
that is able to eliminate estimation bias for a broad class of user-specified
invariances, and is trivial to implement in many practical scenarios. We
enable the use of user-specified invariances throughout the representation
learning process. We empirically show how this improved estimation positively
impacts performance on a variety of environments for the common reward- and
policy-based invariances, and how bisimulation measures interact with standard
data augmentation techniques.

Our experiments show that certain, but not all, environments are improved by
the addition of bisimulation, even beyond standard data augmentation. No
environments see their performance negatively affected by this addition,
suggesting that this technique can be added into standard RL practice without
adverse effects. We also observed that, with proper baseline architectures,
reward- and policy-based bisimulation yields similar results after correcting
for estimation biases.

Future work should be aimed at understanding why certain environments show
large performance differences with bisimulation-based measures, while others
remain unaffected. A proper understanding of what properties of an environment
or a distraction are directly addressed by either data augmentation or
bisimulation is needed.

## References

  * Agarwal et al. (2021) Agarwal, R., Machado, M. C., Castro, P. S., and Bellemare, M. G.  Contrastive behavioral similarity embeddings for generalization in reinforcement learning.  _arXiv preprint arXiv:2101.05265_ , 2021.
  * Castro (2020) Castro, P. S.  Scalable methods for computing state similarity in deterministic markov decision processes.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 34, pp. 10069–10076, 2020.
  * Chen et al. (2020) Chen, T., Kornblith, S., Norouzi, M., and Hinton, G.  A simple framework for contrastive learning of visual representations.  In _International conference on machine learning_ , pp. 1597–1607. PMLR, 2020.
  * Chizat et al. (2020) Chizat, L., Roussillon, P., Léger, F., Vialard, F.-X., and Peyré, G.  Faster wasserstein distance estimation with the sinkhorn divergence.  _Advances in Neural Information Processing Systems_ , 33, 2020.
  * Cobbe et al. (2019) Cobbe, K., Klimov, O., Hesse, C., Kim, T., and Schulman, J.  Quantifying generalization in reinforcement learning.  In _International Conference on Machine Learning_ , pp. 1282–1289. PMLR, 2019.
  * Dosovitskiy et al. (2017) Dosovitskiy, A., Ros, G., Codevilla, F., Lopez, A., and Koltun, V.  Carla: An open urban driving simulator.  In _Conference on robot learning_ , pp. 1–16. PMLR, 2017.
  * Duan et al. (2016) Duan, Y., Schulman, J., Chen, X., Bartlett, P. L., Sutskever, I., and Abbeel, P.  Rl2: Fast reinforcement learning via slow reinforcement learning.  _arXiv preprint arXiv:1611.02779_ , 2016.
  * Espeholt et al. (2018) Espeholt, L., Soyer, H., Munos, R., Simonyan, K., Mnih, V., Ward, T., Doron, Y., Firoiu, V., Harley, T., Dunning, I., et al.  Impala: Scalable distributed deep-rl with importance weighted actor-learner architectures.  In _International Conference on Machine Learning_ , pp. 1407–1416. PMLR, 2018.
  * Ferns & Precup (2014) Ferns, N. and Precup, D.  Bisimulation metrics are optimal value functions.  In _UAI_ , pp. 210–219. Citeseer, 2014.
  * Ferns et al. (2004) Ferns, N., Panangaden, P., and Precup, D.  Metrics for finite markov decision processes.  In _UAI_ , volume 4, pp. 162–169, 2004.
  * Ferns et al. (2011) Ferns, N., Panangaden, P., and Precup, D.  Bisimulation metrics for continuous markov decision processes.  _SIAM Journal on Computing_ , 40(6):1662–1714, 2011.
  * Haarnoja et al. (2018) Haarnoja, T., Zhou, A., Abbeel, P., and Levine, S.  Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor.  In _International conference on machine learning_ , pp. 1861–1870. PMLR, 2018.
  * Hafner et al. (2019) Hafner, D., Lillicrap, T., Ba, J., and Norouzi, M.  Dream to control: Learning behaviors by latent imagination.  _arXiv preprint arXiv:1912.01603_ , 2019.
  * Henderson et al. (2018) Henderson, P., Islam, R., Bachman, P., Pineau, J., Precup, D., and Meger, D.  Deep reinforcement learning that matters.  In _Proceedings of the AAAI conference on artificial intelligence_ , volume 32, 2018.
  * Hjelm et al. (2018) Hjelm, R. D., Fedorov, A., Lavoie-Marchildon, S., Grewal, K., Bachman, P., Trischler, A., and Bengio, Y.  Learning deep representations by mutual information estimation and maximization.  _arXiv preprint arXiv:1808.06670_ , 2018.
  * Jaderberg et al. (2016) Jaderberg, M., Mnih, V., Czarnecki, W. M., Schaul, T., Leibo, J. Z., Silver, D., and Kavukcuoglu, K.  Reinforcement learning with unsupervised auxiliary tasks.  _arXiv preprint arXiv:1611.05397_ , 2016.
  * Kaiser et al. (2019) Kaiser, L., Babaeizadeh, M., Milos, P., Osinski, B., Campbell, R. H., Czechowski, K., Erhan, D., Finn, C., Kozakowski, P., Levine, S., et al.  Model-based reinforcement learning for atari.  _arXiv preprint arXiv:1903.00374_ , 2019.
  * Kalashnikov et al. (2018) Kalashnikov, D., Irpan, A., Pastor, P., Ibarz, J., Herzog, A., Jang, E., Quillen, D., Holly, E., Kalakrishnan, M., Vanhoucke, V., et al.  Qt-opt: Scalable deep reinforcement learning for vision-based robotic manipulation.  _arXiv preprint arXiv:1806.10293_ , 2018.
  * Kemertas & Aumentado-Armstrong (2021) Kemertas, M. and Aumentado-Armstrong, T.  Towards robust bisimulation metric learning.  _Advances in Neural Information Processing Systems_ , 34, 2021.
  * Kostrikov et al. (2020) Kostrikov, I., Yarats, D., and Fergus, R.  Image augmentation is all you need: Regularizing deep reinforcement learning from pixels.  _arXiv preprint arXiv:2004.13649_ , 2020.
  * Laskin et al. (2020) Laskin, M., Lee, K., Stooke, A., Pinto, L., Abbeel, P., and Srinivas, A.  Reinforcement learning with augmented data.  _arXiv preprint arXiv:2004.14990_ , 2020.
  * Lillicrap et al. (2015) Lillicrap, T. P., Hunt, J. J., Pritzel, A., Heess, N., Erez, T., Tassa, Y., Silver, D., and Wierstra, D.  Continuous control with deep reinforcement learning.  _arXiv preprint arXiv:1509.02971_ , 2015.
  * Pont-Tuset et al. (2017) Pont-Tuset, J., Perazzi, F., Caelles, S., Arbeláez, P., Sorkine-Hornung, A., and Van Gool, L.  The 2017 davis challenge on video object segmentation.  _arXiv preprint arXiv:1704.00675_ , 2017.
  * Sekar et al. (2020) Sekar, R., Rybkin, O., Daniilidis, K., Abbeel, P., Hafner, D., and Pathak, D.  Planning to explore via self-supervised world models.  In _International Conference on Machine Learning_ , pp. 8583–8592. PMLR, 2020.
  * Song et al. (2019) Song, X., Jiang, Y., Tu, S., Du, Y., and Neyshabur, B.  Observational overfitting in reinforcement learning.  _arXiv preprint arXiv:1912.02975_ , 2019.
  * Srinivas et al. (2020) Srinivas, A., Laskin, M., and Abbeel, P.  Curl: Contrastive unsupervised representations for reinforcement learning.  _arXiv preprint arXiv:2004.04136_ , 2020.
  * Stone et al. (2021) Stone, A., Ramirez, O., Konolige, K., and Jonschkowski, R.  The distracting control suite–a challenging benchmark for reinforcement learning from pixels.  _arXiv preprint arXiv:2101.02722_ , 2021.
  * Tachet et al. (2018) Tachet, R., Bachman, P., and van Seijen, H.  Learning invariances for policy generalization.  _arXiv preprint arXiv:1809.02591_ , 2018.
  * Zhang et al. (2020) Zhang, A., McAllister, R., Calandra, R., Gal, Y., and Levine, S.  Learning invariant representations for reinforcement learning without reconstruction.  _arXiv preprint arXiv:2006.10742_ , 2020.

##  Appendix A Theorem Proofs

Here we provide the proof for all theorems stated in this work.

#### Theorem 4.6.

Define the metric operator
$\mathcal{F}^{\epsilon}:\mathbb{M}\rightarrow\mathbb{M}$ as

$\mathcal{F}_{\epsilon}(d)(\bm{z},\bm{z^{\prime}})=\mathbb{E}_{\bm{a},\bm{a}^{\prime}\sim\gamma^{\epsilon}(\pi(\cdot|\bm{z}),\pi(\cdot|\bm{z}^{\prime}))}[G(\bm{z},\bm{a},\bm{z}^{\prime},\bm{a}^{\prime})+c\mathcal{W}_{1}(P(\bm{z}_{+}|\bm{z},\pi),P(\bm{z}^{\prime}_{+}|\bm{z}^{\prime},\pi);d)].$
(22)

Then, with $c\in[0,1)$, $\mathcal{F}_{\epsilon}$ has at least a fixed point
$d_{\epsilon}$ which is a bisimulation metric.

#### Proof.

We first observe that, since G is non-negative

$\begin{array}[]{l}A\coloneqq\\{(\bm{z},\bm{z^{\prime}}):\mathbb{E}_{\bm{a},\bm{a}^{\prime}\sim\gamma^{\epsilon}(\pi(\cdot|\bm{z}),\pi(\cdot|\bm{z}^{\prime}))}G(\bm{z},\bm{a},\bm{z}^{\prime},\bm{a}^{\prime})=0\\},\\\
B\coloneqq\\{(\bm{z},\bm{z^{\prime}}):G(\bm{z},\bm{a},\bm{z}^{\prime},\bm{a}^{\prime})=0,\forall\bm{a},\bm{a}^{\prime}\in\textit{Supp}(\gamma^{\epsilon}(\pi(\cdot|\bm{z}),\pi(\cdot|\bm{z}^{\prime})))\\\
(\bm{z},\bm{z^{\prime}})\in A\iff(\bm{z},\bm{z^{\prime}})\in B,\end{array}$
(23)

except on a zero-measure set in
$\gamma^{\epsilon}(\pi(\cdot|\bm{z}),\pi(\cdot|\bm{z}^{\prime}))$.

The rest of the proof largely mimics the proof of Theorem 4.5 in (Ferns et
al., 2004), and its adaptation in Theorem 2 in (Castro, 2020), which is
Theorem 1 in this paper, we largely use the same notation when appropriate,
and refer the reader to (Ferns et al., 2004) for the details ommited in this
proof. Lemma 4.1 in (Ferns et al., 2004) holds for definition 4.4. We make use
of the same distance ordering on $\mathbb{M}$ where $d\leq d^{\prime}\iff
d(\bm{z},\bm{z^{\prime}})\leq
d^{\prime}(\bm{z},\bm{z^{\prime}}),\,\forall\bm{z},\bm{z^{\prime}}\in\mathcal{Z}^{2}$,
this produces an $\omega$-cpo (an $\omega$ complete partial order) with bottom
$\perp$ the everywhere zero distance. Since we also use $\mathcal{W}_{1}$ to
measure state distribution distances in $\mathcal{F}^{\epsilon}$, then Lemma
4.4 in (Ferns et al., 2004) also applies (continuity of $\mathcal{W}_{1}$).

The remaining piece for the proof is showing that $\mathcal{F}^{\epsilon}$ is
continuous, meaning that for every $\omega$-chain $\\{d_{n}\\}$ (an increasing
sequence of metrics), then
$\mathcal{F}^{\epsilon}(\cup_{n}\\{d_{n}\\})=\cup_{n}\\{\mathcal{F}^{\epsilon}(d_{n})\\}$,
where
$\cup_{n}\\{d_{n}\\}(\bm{z},\bm{z^{\prime}})=\sup_{n\in\mathbb{N}}d_{n}(\bm{z},\bm{z^{\prime}})$.
This follows from

$\begin{array}[]{rl}\mathcal{F}^{\epsilon}(\cup_{n}\\{d_{n}\\})&=\mathbb{E}_{\bm{a},\bm{a}^{\prime}\sim\gamma^{\epsilon}(\pi(\cdot|\bm{z}),\pi(\cdot|\bm{z}^{\prime}))}[G(\bm{z},\bm{a},\bm{z}^{\prime},\bm{a}^{\prime})\\\
&+c\mathcal{W}_{1}(P(\bm{z}_{+}|\bm{z},\bm{a}),P(\bm{z}^{\prime}_{+}|\bm{z}^{\prime},\bm{a}^{\prime});\cup_{n}\\{d_{n}\\})],\\\
&=\sup_{n\in\mathbb{N}}\big{\\{}\mathbb{E}_{\bm{a},\bm{a}^{\prime}\sim\gamma^{\epsilon}(\pi(\cdot|\bm{z}),\pi(\cdot|\bm{z}^{\prime}))}[G(\bm{z},\bm{a},\bm{z}^{\prime},\bm{a}^{\prime})\\\
&+c\mathcal{W}_{1}(P(\bm{z}_{+}|\bm{z},\bm{a}),P(\bm{z}^{\prime}_{+}|\bm{z}^{\prime},\bm{a}^{\prime});d_{n})]\big{\\}},\\\
&=\sup_{n\in\mathbb{N}}\mathcal{F}^{\epsilon}(d_{n}),\\\ \end{array}$ (24)

Where the second equality follows by the continuity of $\mathcal{W}_{1}$
w.r.t. the metric $d_{n}$. The rest of the proof follows (Ferns et al., 2004).

#### Lemma 4.7.

Let $E^{\pi}$ be the largest $\pi$-bisimulation relation, and let $E^{PSM}$
likewise be the largest average-action bisimulation relation in (Agarwal et
al., 2021). The entangled bisimulation relation $E^{\epsilon}$ satisfies

$\begin{array}[]{rl}E^{\epsilon}&\subseteq E^{\pi}\\\ E^{\epsilon}&\subseteq
E^{PSM}\\\ \end{array}$ (25)

for the choice of state-action similarity metrics
$G(\bm{z},\bm{a},\bm{z}^{\prime},\bm{a}^{\prime})=||R(\bm{z},\bm{a})-R(\bm{z}^{\prime},\bm{a}^{\prime})||_{1}$
and
$G(\bm{z},\bm{a},\bm{z}^{\prime},\bm{a}^{\prime})=||\mathbb{E}_{\pi(\bm{a}|\bm{z})}[\bm{a}]-\mathbb{E}_{\pi(\bm{a}^{\prime}|\bm{z}^{\prime})}[\bm{a}^{\prime}]||_{1}$
respectively

#### Proof.

The proof is identical in both cases, so we focus on the reward-based case
with
$G(\bm{z},\bm{a},\bm{z}^{\prime},\bm{a}^{\prime})=||R(\bm{z},\bm{a})-R(\bm{z}^{\prime},\bm{a}^{\prime})||_{1}$.
From Jensen’s inequality we have

$\begin{array}[]{rl}\mathbb{E}_{{\bm{a},\bm{a^{\prime}}\sim\gamma^{\epsilon}(\pi(\cdot|\bm{z}),\pi(\cdot|\bm{z}^{\prime}))}}[G(\bm{z},\bm{a},\bm{z}^{\prime},\bm{a}^{\prime})]&\geq||\mathbb{E}_{\bm{a}\sim\pi(\cdot|\bm{z})}R(\bm{z},\bm{a})-\mathbb{E}_{\bm{a}^{\prime}\sim\pi(\cdot|\bm{z}^{\prime})}R(\bm{z}^{\prime},\bm{a}^{\prime})||_{1}\\\
&=||R(\bm{z},\pi)-R(\bm{z}^{\prime},\pi)||_{1}.\\\ \end{array}$ (26)

And, likewise, for any metric $d$ we have

$\begin{array}[]{rl}\mathbb{E}_{{\bm{a},\bm{a^{\prime}}\sim\gamma^{\epsilon}(\pi(\cdot|\bm{z}),\pi(\cdot|\bm{z}^{\prime}))}}\mathcal{W}_{1}(P(\bm{z}_{+}|\bm{z},\bm{a}),P(\bm{z}^{\prime}_{+}|\bm{z}^{\prime},\bm{a}^{\prime});d)&\geq\mathcal{W}_{1}(P(\bm{z}_{+}|\bm{z},\pi),P(\bm{z}^{\prime}_{+}|\bm{z}^{\prime},\pi);d)\end{array}$
(27)

Then, we observe that $\mathcal{F}^{\epsilon}(d)\geq\mathcal{F}^{\pi}(d)$, and
therefore $d^{\epsilon}\geq d^{\pi}$. From this we conclude that any state
pair $(\bm{z},\bm{z^{\prime}})\in E^{\epsilon}$ satisfies $0\leq
d^{\pi}(\bm{z},\bm{z^{\prime}})\leq d^{\epsilon}(\bm{z},\bm{z^{\prime}})=0$
and therefore $(\bm{z},\bm{z^{\prime}})\in E^{\pi}$.

#### Theorem 4.8.

Given an MDP $\mathcal{M}$, policy $\pi$, and a state-action similarity metric
$G$. Define the metric operator
$\mathcal{F}_{\bar{\epsilon}}:\mathbb{M}\rightarrow\mathbb{M}$ as

$\mathcal{F}_{\bar{\epsilon}}(d)(\bm{z},\bm{z^{\prime}})=\mathbb{E}_{\begin{subarray}{c}(\bm{a},\bm{a^{\prime}})\sim\gamma^{\epsilon}(\pi(\cdot|\bm{z}),\pi(\cdot|\bm{z}^{\prime}))\\\
(\bm{z}_{+},\bm{z}^{\prime}_{+})\sim\gamma^{\epsilon}(P(\cdot|\bm{z},\bm{a}),P(\cdot|\bm{z}^{\prime},\bm{a}^{\prime}))\end{subarray}}[G(\bm{z},\bm{a},\bm{z}^{\prime},\bm{a}^{\prime})+cd_{\phi}(\bm{z}_{+},\bm{z}^{\prime}_{+})].$
(28)

Then, $\mathcal{F}_{\bar{\epsilon}}$ has at least a fixed point
$d_{\bar{\epsilon}}$ satisfying $d_{\bar{\epsilon}}\geq d_{\epsilon}$, and
$(\bm{z},\bm{z^{\prime}})\in E^{\epsilon}\rightarrow
d_{\bar{\epsilon}}(\bm{z},\bm{z^{\prime}})=0$, in particular,
$d_{\bar{\epsilon}}(\bm{z},\bm{z})=0,\,\forall\bm{z}\in\mathcal{Z}$.

Further, if the state transition function is coordinate independent
($P(\bm{z}_{+}\mid\bm{z},\bm{a})=\Pi_{i}P_{i}(z_{i,+}\mid\bm{z},\bm{a})\,\forall\bm{z}_{+},\bm{z},\bm{a}$),
and $d_{\bar{\epsilon}}$ is of the form

$d_{\bar{\epsilon}}(\bm{z},\bm{z}^{\prime})=\sum\limits_{\begin{subarray}{c}i=[n]\\\
j=[p]\end{subarray}}w_{i,j}|z_{i}-z^{\prime}_{i}|^{j},\;w_{i,j}\geq 0\forall
i,j,\;p>0,$ (29)

then the bound is tight, that is

$\begin{array}[]{rl}d_{\epsilon}(\bm{z},\bm{z}^{\prime})&=d_{\bar{\epsilon}}(\bm{z},\bm{z}^{\prime})\\\
&=\mathbb{E}_{\begin{subarray}{c}(\bm{a},\bm{a^{\prime}})\sim\gamma^{\epsilon}(\pi(\cdot|\bm{z}),\pi(\cdot|\bm{z}^{\prime}))\\\
(\bm{z}_{+},\bm{z}^{\prime}_{+})\sim\gamma^{\epsilon}(P(\cdot|\bm{z},\bm{a}),P(\cdot|\bm{z}^{\prime},\bm{a}^{\prime}))\end{subarray}}[G(\bm{z},\bm{a},\bm{z}^{\prime},\bm{a}^{\prime})+cd_{\epsilon}(\bm{z}_{+},\bm{z}^{\prime}_{+})].\end{array}$
(30)

#### Proof.

The proof will be split in three parts, we first prove that
$d_{\bar{\epsilon}}$ exists and is an upper bound for $d_{\epsilon}$ on every
state pair, we then prove that $d_{\bar{\epsilon}}$ is zero on
$\epsilon$-bisimilar states, and, lastly, we prove that the bound is tight for
coordinate-independent transition functions if the distance function satisfies
the prescribed functional form. To reduce notation clutter throughout the
proof, we drop the state and action conditionals by noting
$P(\bm{z}_{+})\coloneqq P(\cdot|\bm{z},\bm{a})$, and, likewise,
$Q(\bm{z}_{+}^{\prime})\coloneqq P(\cdot|\bm{z}^{\prime},\bm{a}^{\prime})$.
Additionally, we denote
$\gamma^{\epsilon}_{P}\coloneqq\gamma^{\epsilon}(P(\cdot|\bm{z},\bm{a}),P(\cdot|\bm{z},\bm{a}^{\prime}))$
and
$\gamma^{\epsilon}_{\pi}\coloneqq\gamma^{\epsilon}(\pi(\cdot\mid\bm{z}),\pi(\cdot\mid\bm{z}^{\prime}))$.

#### Upper bound:

To prove the existence of the fixed point metric $d_{\bar{\epsilon}}$, we
again mimic the proof of Theorem 4.5 in (Ferns et al., 2004), only now we need
to prove that the operation
$T(d)\coloneqq\mathbb{E}_{\begin{subarray}{c}(\bm{a},\bm{a^{\prime}})\sim\gamma^{\epsilon}_{\pi}\\\
(\bm{z}_{+},\bm{z}^{\prime}_{+})\sim\gamma^{\epsilon}_{P}\end{subarray}}[d(\bm{z}_{+},\bm{z}^{\prime}_{+})]$
is continuous (the analog of Lemma 4.4 in that work).

This amounts to observing that, on an $\omega$-chain777Recall that an
$\omega$-chain satisfies $d_{i}\leq d_{j}$ for every $j>i$, and also
$\cup_{n}\\{d_{n}\\}(\bm{z},\bm{z^{\prime}})=\sup_{n\in\mathbb{N}}d_{n}(\bm{z},\bm{z^{\prime}})$,
we also define $d_{i}\leq d_{j}\iff d_{i}(\bm{z},\bm{z^{\prime}})\leq
d_{j}(\bm{z},\bm{z^{\prime}})\,\forall\,(\bm{z},\bm{z^{\prime}})\in\mathcal{Z}^{2}$
$\\{d_{n}\\}$,

$\begin{array}[]{rl}T(\cup_{n}\\{d_{n}\\})&=\mathbb{E}_{\begin{subarray}{c}(\bm{a},\bm{a^{\prime}})\sim\gamma^{\epsilon}_{\pi}\\\
(\bm{z}_{+},\bm{z}^{\prime}_{+})\sim\gamma^{\epsilon}_{P}\end{subarray}}[\cup_{n}\\{d_{n}\\}(\bm{z}_{+},\bm{z}^{\prime}_{+})],\\\
&=\sup_{n\in\mathbb{N}}\mathbb{E}_{\begin{subarray}{c}(\bm{a},\bm{a^{\prime}})\sim\gamma^{\epsilon}_{\pi}\\\
(\bm{z}_{+},\bm{z}^{\prime}_{+})\sim\gamma^{\epsilon}_{P}\end{subarray}}[d_{n}(\bm{z}_{+},\bm{z}^{\prime}_{+})]\\\
&=\sup_{n\in\mathbb{N}}T(d_{n}).\\\ \end{array}$ (31)

The rest of the proof of existence of a fixed point needs to show continuity
of $\mathcal{F}_{\bar{\epsilon}}$, which is nearly identical to the proof for
Theorem 4.6.

Now, to prove that $d_{\epsilon}\leq d_{\bar{\epsilon}}$, we recall that
$d_{\epsilon}=\cup_{n\in\mathbb{N}}\\{\mathcal{F}^{n}_{\epsilon}(\perp)\\}$
where $\perp$ is the everywhere-zero metric and
$\mathcal{F}^{n}_{\epsilon}(\perp)$ is the result of recursively applying
$\mathcal{F}_{\epsilon}$ to $\perp$ n times; and, likewise
$d_{\bar{\epsilon}}=\cup_{n\in\mathbb{N}}\\{\mathcal{F}^{n}_{\bar{\epsilon}}(\perp)\\}$.
From this, it suffices to prove that, if $d\leq d^{\prime}$ then
$\mathcal{F}_{\epsilon}(d)\leq\mathcal{F}_{\bar{\epsilon}}(d^{\prime})$, since
this would imply that
$\mathcal{F}^{n}_{\epsilon}(\perp)\leq\mathcal{F}^{n}_{\bar{\epsilon}}(\perp),\forall\,n\in\mathbb{N}$.

$\begin{array}[]{rl}\mathcal{F}_{\epsilon}(d)(\bm{z},\bm{z^{\prime}})&=\mathbb{E}_{\bm{a},\bm{a}^{\prime}\sim\gamma^{\epsilon}_{\pi}}[G(\bm{z},\bm{a},\bm{z}^{\prime},\bm{a}^{\prime})+c\mathcal{W}_{1}(P(\bm{z}_{+}|\bm{z},\pi),P(\bm{z}^{\prime}_{+}|\bm{z}^{\prime},\pi);d)],\\\
&\leq\mathbb{E}_{\bm{a},\bm{a}^{\prime}\sim\gamma^{\epsilon}_{\pi}}[G(\bm{z},\bm{a},\bm{z}^{\prime},\bm{a}^{\prime})+c\mathcal{W}_{1}(P(\bm{z}_{+}|\bm{z},\pi),P(\bm{z}^{\prime}_{+}|\bm{z}^{\prime},\pi);d^{\prime})],\\\
&\leq\mathbb{E}_{\bm{a},\bm{a}^{\prime}\sim\gamma^{\epsilon}_{\pi}}[G(\bm{z},\bm{a},\bm{z}^{\prime},\bm{a}^{\prime})+c\mathbb{E}_{\bm{z_{+}},\bm{z_{+}}^{\prime}\sim\gamma^{\epsilon}_{P}}[d^{\prime}(\bm{z_{+}},\bm{z_{+}}^{\prime})]],\\\
&=\mathbb{E}_{\begin{subarray}{c}\bm{a},\bm{a}^{\prime}\sim\gamma^{\epsilon}_{\pi}\\\
\bm{z_{+}},\bm{z_{+}}^{\prime}\sim\gamma^{\epsilon}_{P}\end{subarray}}[G(\bm{z},\bm{a},\bm{z}^{\prime},\bm{a}^{\prime})+cd^{\prime}(\bm{z_{+}},\bm{z_{+}}^{\prime})],\\\
&=\mathcal{F}_{\bar{\epsilon}}(d)(\bm{z},\bm{z^{\prime}}).\\\ \end{array}$
(32)

Here the first inequality derives from $d\leq d^{\prime}$ and the continuity
of $\mathcal{W}_{1}$, and the second inequality stems from
$\gamma^{\epsilon}_{P}\in\Gamma(P,Q)$ and the definition of $\mathcal{W}_{1}$.

####  $\epsilon$-similar states.

To prove that, if $\bm{z},\bm{z}^{\prime}\in E^{\epsilon}\rightarrow
d_{\bar{\epsilon}}(\bm{z},\bm{z}^{\prime})=0$, we again use that
$d_{\epsilon}=\cup_{n\in\mathbb{N}}\\{\mathcal{F}^{n}_{\epsilon}(\perp)\\}$.
Inductively, we have that
$\mathcal{F}^{0}_{\epsilon}(\perp)(\bm{z},\bm{z}^{\prime})=0\,\forall\bm{z},\bm{z}^{\prime}\in
E^{\epsilon}$(in fact, this is true everywhere), we now show that if
$\mathcal{F}^{j}_{\epsilon}(\perp)(\bm{z},\bm{z}^{\prime})=0,\forall\bm{z},\bm{z}^{\prime}\in
E^{\epsilon}$, then
$\mathcal{F}^{j+1}_{\epsilon}(\perp)(\bm{z},\bm{z}^{\prime})=0,\forall\bm{z},\bm{z}^{\prime}\in
E^{\epsilon}$. It suffices to observe that, under these conditions, if
$\bm{z},\bm{z}^{\prime}\in E^{\epsilon}$ and
$d:d(\bm{z},\bm{z}^{\prime})=0\forall\bm{z},\bm{z}^{\prime}\in E^{\epsilon}$
then

$\begin{array}[]{rl}\mathcal{F}_{\bar{\epsilon}}(d)(\bm{z},\bm{z}^{\prime})=&\mathbb{E}_{\begin{subarray}{c}\bm{a},\bm{a}^{\prime}\sim\gamma^{\epsilon}_{\pi}\\\
\bm{z_{+}},\bm{z_{+}}^{\prime}\sim\gamma^{\epsilon}_{P}\end{subarray}}[G(\bm{z},\bm{a},\bm{z}^{\prime},\bm{a}^{\prime})+cd(\bm{z}_{+},\bm{z}^{\prime}_{+})],\\\
&=0.\\\ \end{array}$ (33)

#### Optimality of upper bound under conditions:

Finally, let $\bm{z}_{\setminus i}=\\{z_{j}\\}_{j\neq i}$. If
$P(\bm{z}_{+}\mid\bm{z},\bm{a})=\Pi_{i}P_{i}(z_{i,+}\mid\bm{z},\bm{a})\,\forall\bm{z}_{+},\bm{z},\bm{a}$,
and we further restrict the function class $d(\bm{z},\bm{z}^{\prime})$ to be
coordinate-wise separable and convex, we can first show that expected
distances over couplings have a coordinate-factorized structure, that is, for
any coupling $\gamma\in\Gamma(P,Q)$, we have

$\begin{array}[]{rl}\mathop{\mathbb{E}}\limits_{\bm{z},\bm{z}^{\prime}\sim\gamma}[d(\bm{z},\bm{z}^{\prime})]&=\int
d(\bm{z},\bm{z}^{\prime})^{p}\gamma(\bm{z},\bm{z}^{\prime})d\bm{z}d\bm{z}^{\prime}\\\
&=\sum\limits_{i=1}^{n}\int
d_{i}(z_{i},z^{\prime}_{i})\gamma(\bm{z},\bm{z}^{\prime})d\bm{z}d\bm{z}^{\prime}\\\
&=\sum\limits_{i=1}^{n}\int
d_{i}(z_{i},z^{\prime}_{i})\gamma_{i}(z_{i},z^{\prime}_{i})\gamma(\bm{z}_{\setminus
i},\bm{z}^{\prime}_{\setminus
i}|z_{i},z^{\prime}_{i})d\bm{z}d\bm{z}^{\prime}\\\
&=\sum\limits_{i=1}^{n}\int_{z_{i},z^{\prime}_{i}}d_{i}(z_{i},z^{\prime}_{i})\gamma_{i}(z_{i},z^{\prime}_{i})\int_{\bm{z}_{\setminus
i},\bm{z}^{\prime}_{\setminus i}}\gamma(\bm{z}_{\setminus
i},\bm{z}^{\prime}_{\setminus i}|z_{i},z^{\prime}_{i})d\bm{z}_{\setminus
i}d\bm{z}^{\prime}_{\setminus i}dz_{i}dz^{\prime}_{i}\\\
&=\sum\limits_{i=1}^{n}\int
d_{i}(z_{i},z^{\prime}_{i})\gamma(z_{i},z^{\prime}_{i})dz_{i}dz^{\prime}_{i}\\\
&=\sum\limits_{i=1}^{n}\int
d_{i}(z_{i},z^{\prime}_{i})\gamma(z_{i},z^{\prime}_{i})\mathop{\Pi}\limits_{j\in[n]\setminus
i}\gamma(\hat{z}_{j},\hat{z}^{\prime}_{j})d\hat{\bm{z}}_{\setminus
i}d\hat{\bm{z}}_{\setminus i}dz_{i}dz^{\prime}_{i}\\\ &=\int
d(\hat{\bm{z}},\hat{\bm{z}}^{\prime})\mathop{\Pi}\limits_{i=1}^{n}\gamma(\hat{\bm{z}}_{i},\hat{\bm{z}}^{\prime}_{i})d\hat{\bm{z}}d\hat{\bm{z}}^{\prime}\\\
&=E_{\hat{\bm{z}},\hat{\bm{z}}^{\prime}\sim\mathop{\Pi}\limits_{i=1}^{n}\gamma_{i}}[d(\hat{\bm{z}},\hat{\bm{z}}^{\prime})]\\\
\end{array}$ (34)

In essence, the expected value of the coupling is the same as the one for the
coordinate-independent coupling. The preceding result only requires $d$ to
satisfy the coordinate independent assumption, also recall
$\mathcal{W}_{1}(P,Q;d)=\inf_{\gamma\in\Gamma(P,Q)}\mathop{\mathbb{E}}\limits_{\bm{z},\bm{z}^{\prime}\sim\gamma}[d(\bm{z},\bm{z}^{\prime})]$.

A well-known formula the p-Wasserstein distance between univariate variables
is

$W_{p}(P_{i},Q_{i})^{p}=\int_{0}^{1}|F_{i}^{-1}(u)-G_{i}^{-1}(u)|^{p}du.$ (35)

Where $F_{i}^{-1}(u)$ and $G_{i}^{-1}(u)$ are the inverse CDF’s of the i-th
coordinate of distributions $P(X)$ and $Q(Y)$ respectively.

From the closed form solution, we deduce that $\gamma^{\epsilon}(P,Q)$ is the
Wasserstein optimal coupling for any distance function of the form
$d(\bm{z},\bm{z}^{\prime})=\sum\limits_{\begin{subarray}{c}i=[n]\\\
j=[p]\end{subarray}}w_{i,j}|z_{i}-z^{\prime}_{i}|^{j},\;w_{i,j}\geq 0\forall
i,j,\;p>0$ if the transition distribution satisfies coordinate independence.

##  Appendix B Algorithm description

Here we provide a full description of the training algorithm, where an inverse
dynamics loss is added to the standard SAC objective, along our proposed
$\epsilon-$bisimulation loss.

Algorithm 2 Entangled SAC Algorithm

for Time $t=0$ to $\infty$ do

Sample action: $\bm{a}_{t}\sim\pi_{\theta}(h_{\theta}(o_{t}))$

Get environment transition: $o_{t+1},r_{t+1}\sim\text{Env}(\bm{a}_{t})$

Store transition: $D\leftarrow D\cup\\{o_{t},\bm{a}_{t},r_{t+1},o_{t+1}\\}$

Sample batch: $B\sim D$

Dynamics loss: $J_{P}=\frac{1}{n}\sum_{o_{t},o_{t+1},\bm{a}_{t}\in B}[\ln
P_{\theta}(\overline{h_{\theta}(o_{t+1})}|\bm{a}_{t},h_{\theta}(o_{t}))]$

Inverse dynamics loss: $J_{ID}=\frac{1}{n}\sum_{o_{t},o_{t+1},\bm{a}_{t}\in
B}||\bm{a}_{t}-A_{\theta}(h_{\theta}(o_{t}),h_{\theta}(o_{t+1}))||_{2}^{2}$

Reward loss: $J_{R}=\frac{1}{n}\sum_{o_{t},r_{t},\bm{a}_{t}\in
B}[||R_{\theta}(\bm{z}\sim
P_{\theta}(\bm{z}|\bm{a}_{t},h_{\theta}(o_{t})))-r_{t}||^{2}_{2}]$

Policy loss: $J_{\pi}$ {SAC policy update, Algorithm 3}

entangled bisimulation loss: $J_{\epsilon}$ {Entangled update, Algorithm 4}

Update parameters:
$\theta\leftarrow\theta+\eta\nabla_{\theta}(J_{P}+J_{ID}+J_{\pi}+J_{R}+J_{\pi}+\lambda
J_{\epsilon})$

end for

Algorithm 3 Policy Algorithm

0: data batch $B=\\{o_{t},o_{t+1},\bm{a}_{t},r_{t+1}\\}_{i=1}^{n}$, policy
$\pi_{\theta}$, critics $Q^{j},\,j=\\{1,2\\}$, target critics
$\hat{Q}^{j},\,j=\\{1,2\\}$, observation encoder $h_{\theta}$.

Get future state value:
$\overline{V}=\min\limits_{j=1,2}\hat{Q}^{j}_{\theta}(h_{\theta}(\bm{o_{t+1}}),\hat{\bm{a}}\sim\pi(\hat{\bm{a}}|h_{\theta}(\bm{o_{t+1}})))$

Critic loss: $J_{C}=\frac{1}{n}\sum_{o_{t},r_{t},\bm{a}_{t},o_{t+1}\in
B}\sum\limits_{j=1,2}||Q^{j}_{\theta}(h_{\theta}(\bm{o_{t}}),\bm{a}_{t})-r_{t+1}-\gamma\overline{V}||_{2}^{2}$

Actor loss: $J_{a}=\frac{1}{n}\sum_{o_{t},\bm{a}_{t}\in B}\
\overline{\alpha}\log\pi(\bm{a}_{t}|h_{\theta}(\bm{o_{t}}))-\min\limits_{j=1,2}\overline{Q}^{j}_{\theta}(h_{\theta}(\bm{o_{t}}),\bm{a}_{t})$

Train alpha:
$J_{\alpha}=\alpha\log\overline{\pi(\bm{a}_{t}|h_{\theta}(\bm{o_{t}}))}-\mathcal{H}(\bm{a}_{t}|h_{\theta}(\bm{o_{t}}))$

Update target critics:
$\hat{Q}^{j}_{\theta}=\tau{Q}^{j}_{\theta}+(1-\tau)\hat{Q}^{j}_{\theta}$

Policy loss $J_{\pi}=J_{a}+J_{C}$

Algorithm 4 Bisimulation Algorithm

0: Latent state batch $B=\\{\bm{z}\\}_{i=1}^{n}$, policy $\pi_{\theta}$,
latent transition model $P_{\theta}(\cdot\mid,\bm{z},\bm{a})$, bisimulation
distance function $d_{\theta}(\cdot,\cdot)$, similarity pseudometric $G$

Permute states $B^{\prime}=\\{\bm{z}^{\prime}\\}=\text{Perm}(\\{\bm{z}\\})$

Sample noise variables
$\epsilon^{A}\sim\mathcal{N}(0,I^{|A|}),\epsilon^{Z}\sim\mathcal{N}(0,I^{|Z|})$

Compute tied actions
$\bm{a}=\pi_{\theta}(\cdot|\bm{z};\epsilon^{A}),\;\bm{a}^{\prime}=\pi_{\theta}(\cdot|\bm{z}^{\prime};\epsilon^{A})$

Compute tied latent transitions
$\bm{z}_{+}=P_{\theta}(\cdot|\bm{z},\bm{a};\epsilon^{Z}),\;\bm{z}^{\prime}_{+}=P_{\theta}(\cdot|\bm{z}^{\prime},\bm{a}^{\prime};\epsilon^{Z})$

Bisimulation target
$\hat{d}(\bm{z},\bm{z}^{\prime})=G(\bm{z},\bm{a},\bm{z}^{\prime},\bm{a}^{\prime})+cd_{\theta}(\bm{z}_{+},\bm{z}^{\prime}_{+})$

bisimulation loss
$J_{\epsilon}(\bm{z},\bm{z}^{\prime})=\frac{1}{n}\sum\limits_{\bm{z^{\prime}},\bm{z}\in,B,B^{\prime}}||d_{\theta}(\bm{z},\bm{z}^{\prime})-\overline{\hat{d}(\bm{z},\bm{z}^{\prime})}||^{2}_{2}$

##  Appendix C Network Architecture and Hyperparameters

We use the a similar architecture to (Zhang et al., 2020; Agarwal et al.,
2021), with some key modifications to the encoder architecture. The encoder is
implemented using a convnet trunk with $4$ layers and ReLU activations, each
kernel has size $3$ with $32$ channels, stride $2$ on the first layer, and
stride $1$ everywhere else; the output of the convnet trunk is projected into
a $50$-dimensional vector with a fully-connected layer. Unlike (Zhang et al.,
2020; Agarwal et al., 2021), the encoder is fully shared between the actor and
the critic. Also unlike (Zhang et al., 2020; Agarwal et al., 2021), the
encoder is applied to each frame in the input framestack independently, and
the embeddings are then concatenated, yielding a $F\times 50$ embedding on a
framestack with $F$ repeats ($F=3$ for all DCS environments). The independent
embedding of each frame in the framestack is used to improve robustness on the
distracting camera movements, since these distractions break low-level pixel
correspondence in the same framestack.

The actor, critic, dynamics, inverse dynamics, and reward models are all
implemented with MLPs with two hidden layers with $1024$ neurons each ($200$
for low capacity experiments) and ReLU activations. The continuous actions of
the actor model are sampled via the reparametrization trick with a $\tanh$
nonlinearity, and the latent state prediction is parametrized as a coordinate-
independent Gaussian distribution.

Table 4 summarizes the hyperparameters used for the RL experiments; Table 5
summarizes the action repeats used for each DCS environment.

Parameter Name | Value
---|---
Replay buffer capacity | $100000$
Batch size | $512$
Discount $\gamma$ | $0.99$
Optimizer | Adam
Learning rate | $10^{-3}$
Critic target update frequency | $2$
Critic Q-function soft update rate $\tau$ | $0.01$
Actor update frequency | $2$
Actor log stddev bounds | $[-10,2]$
Temperature learning rate | $10^{-4}$
Initial Temperature | $0.1$
Entropy Target | $-\text{num actions}$
Image augmentation pad (DrQ) | $4$
Table 4: Architecture and training hyperparameters Environment | action repeats
---|---
Ball in cup Catch | 4
Cartpole Swingup | 8
Cheetah Run | 4
Finger Spin | 2
Walker Walk | 2
Reacher Easy | 4
Table 5: Action repeat parameter used in DCS environments

##  Appendix D Autonomous Driving Experiments

We replicate the autonomous driving experiments in (Zhang et al., 2020) using
the CARLA simulator (Dosovitskiy et al., 2017). The agent’s goal is to drive
as far as possible while avoiding collisions with other 20 vehicles in CARLA’s
Town04 highway map in 1000 time-steps, the reward function and environment
setup are identical to (Zhang et al., 2020). The agent has access to $5$
cameras placed on the car’s roof, each covering $60^{\circ}$, the images are
concatenated to create a single, $84\times 420$ pixel, $300^{\circ}$ view of
the environment, sample views are shown in Figure 5. Similarly to our previous
experiments, we repurpose the original training architecture, with a single
shared encoder for both actor and critic, consecutive frames are stacked into
a single observation, but these are processed separately, and their
corresponding embeddings are concatenated; see Appendix C for details.

(a)

(b)

Figure 5: Sample egocentric view in CARLA environment on Town04 highway map.
The highway layout is also shown

Table 6 summarizes the results, for these experiments DrQ was not used. Unlike
the results reported in (Zhang et al., 2020), we observe that the minor
adjustments we made to the architecture are enough to fully close the
performance gap between SAC and the bisimulation-based options. The lack of
result differentiation in Table 6 suggests that the original architecture was
not well suited for the task, and bisimulation acted as a helpful
regularization objective.

SAC | DBC |  $\bm{\epsilon}$-R | PSE |  $\bm{\epsilon}$-$\bm{\pi}$
---|---|---|---|---
$\bm{168\pm 9}$ | $160\pm 7$ | $157\pm 11$ | $158\pm 11$ | $148\pm 23$
Table 6: Episodic reward on CARLA’s Town04 highway map, $500k$ environment
steps. Mean and standard deviation are computed across 5 seeds.