SlowGuess/NeurIPS_2023_OCR huggingface dataset:
where  $\alpha$  is the step size. The fixed-point of the measurement obtained using this expectile-based operator is denoted as  $G_{\tau}$ . Although the utilization of the expectile statistics is well established, its application for estimating bisimulation measurement is not particularly intuitive. In the following, we will show how expectile-based operator can be helpful in addressing the aforementioned issue. First, it is worth noting that when  $\tau = 1/2$ , this operator becomes the bisimulation expectation of the behavior policy, i.e.,  $\mathbb{E}_{\mu_{\pi_{\beta}}}[\hat{\epsilon}]$ . Next, we shall consider how this operator performs when  $\tau \to 1$ . We show that under certain assumptions, our method indeed approximates an "optimal" measurement in terms of the given dataset. We first prove a technical lemma stating that the update operator is still a contraction, and then prove a lemma relating different expectiles, finally we derive our main result regarding the "optimality" of our method.
Lemma 5. For any  $\tau \in [0,1)$ ,  $\mathcal{F}_{\tau}^{\pi}$  is a  $\gamma_{\tau}$ -contraction, where  $\gamma_{\tau} = 1 - 2\alpha(1 - \gamma)\min\{\tau,1 - \tau\} < 1$ .
Lemma 6. For any  $\tau, \tau' \in [0,1)$  with  $\tau' \geq \tau$ , and for all  $s_i, s_j \in S$  and any  $\alpha$ , we have  $G_{\tau'} \geq G_{\tau}$ .
Theorem 7. In deterministic MDP and fixed finite dataset, we have:
$$
\lim  _ {\tau \rightarrow 1} G _ {\tau} \left(s _ {i}, s _ {j}\right) = \max  _ {\substack {a _ {i} \in \mathcal {A}, a _ {j} \in \mathcal {A}\\s. t. \pi_ {\beta} \left(a _ {i} \mid s _ {i}\right) > 0, \pi_ {\beta} \left(a _ {j} \mid s _ {j}\right) > 0}} G _ {\sim} ^ {*} \left(\left(s _ {i}, a _ {i}\right), \left(s _ {j}, a _ {j}\right)\right). \tag{8}
$$
where  $G_{\sim}^{*}((s_{i},a_{i}),(s_{j},a_{j}))$  is a fixed-point measurement constrained to the dataset and defined on the state-action space  $\mathcal{S}\times \mathcal{A}$  as
$$
G_{\sim}^{*}((s_{i},a_{i}),(s_{j},a_{j})) = |r(s_{i},a_{i}) - r(s_{j},a_{j})| + \gamma \mathbb{E}_{\substack{s^{\prime}_{i}\sim T_{s_{i}}^{\pi_{\beta}}\\ s^{\prime}_{j}\sim T_{s_{j}}^{\pi_{\beta}}}}\left[ \max_{\substack{a^{\prime}_{i}\in \mathcal{A},a^{\prime}_{j}\in \mathcal{A}\\ s.t. \pi_{\beta}(a^{\prime}_{i}|s^{\prime}_{i}) > 0,\pi_{\beta}(a^{\prime}_{j}|s^{\prime}_{j}) > 0}}G_{\sim}^{*}((s^{\prime}_{i},a^{\prime}_{i}),(s^{\prime}_{j},a^{\prime}_{j}))\right].
$$
Intuitively,  $G_{\sim}((s_i,a_i),(s_j,a_j))$  can be interpreted as a state-action value function  $Q(\tilde{s},\tilde{a})$  in a lifted MDP  $\widetilde{M}$ , and  $G_{\sim}(s_i,s_j)$  as a state value function  $V(\tilde{s})$ . We defer the detailed explanation to Appendix E.
Theorem 7 illustrates that, as  $\tau \to 1$ , we are effectively approximating the maximum  $G_{\sim}((s_i,a_i),(s_j,a_j))$  over actions  $a_i', a_j'$  from the dataset. When we set  $\tau = 1$ , the expectile-based bisimulation operator achieves fully in-sample learning: we only consider state pairs that have corresponding actions in the dataset. For instance, only when we have  $(s_i', a_i') \in \mathcal{D}$  and  $(s_j', a_j') \in \mathcal{D}$ , can we apply the measurement of  $G_{\sim}^*$ . As such, by manipulating  $\tau$ , we balance a trade-off between minimizing the expected bisimulation residual (for  $\tau = 0.5$ ) and evaluating  $G_{\sim}^*((s_i,a_i),(s_j,a_j))$  solely on the dataset (for  $\tau = 1$ ), thereby sidestepping the failure case outlined in Proposition 4 in an implicit manner.
5.2 Reward Scaling
Most previous works [5, 54, 6, 52] have overlooked the impact of reward scaling in the bisimulation operator. To demonstrate its importance, we investigate a more general form of the bisimulation operator in Equation 2, given as:
$$
\mathcal {F} ^ {\pi} G \left(s _ {i}, s _ {j}\right) = c _ {r} \cdot \left| r _ {s _ {i}} ^ {\pi} - r _ {s _ {j}} ^ {\pi} \right| + c _ {k} \cdot \mathbb {E} _ {s _ {i} ^ {\prime}, s _ {j} ^ {\prime}} ^ {\pi} [ G \left(s _ {i} ^ {\prime}, s _ {j} ^ {\prime}\right) ]. \tag {9}
$$
We then can derive the following:
$$
\begin{array}{l} G _ {\sim} ^ {\pi} (s _ {i}, s _ {j}) = \mathcal {F} ^ {\pi} G _ {\sim} ^ {\pi} (s _ {i}, s _ {j}) = c _ {r} \cdot | r _ {s _ {i}} ^ {\pi} - r _ {s _ {j}} ^ {\pi} | + c _ {k} \cdot \mathbb {E} _ {s _ {i} ^ {\prime}, s _ {j} ^ {\prime}} ^ {\pi} [ G _ {\sim} ^ {\pi} (s _ {i} ^ {\prime}, s _ {j} ^ {\prime}) ] \\ \leq c _ {r} \cdot \left(R _ {\max } - R _ {\min }\right) + c _ {k} \cdot \mathbb {E} _ {s _ {i} ^ {\prime}, s _ {j} ^ {\prime}} ^ {\pi} \left[ G _ {\sim} ^ {\pi} \left(s _ {i} ^ {\prime}, s _ {j} ^ {\prime}\right) \right] \tag {10} \\ \leq c _ {r} \cdot (R _ {\max } - R _ {\min }) + c _ {k} \cdot \max  _ {s _ {i} ^ {\prime}, s _ {j} ^ {\prime}} G _ {\sim} ^ {\pi} (s _ {i} ^ {\prime}, s _ {j} ^ {\prime}). \\ \end{array}
$$
Figure 1: The effectiveness of Reward Scaling  $(RS)$  in SimSR on halfcheetah-medium-expert-v2, with results averaged on 3 random seeds. (Left) Effective Dimension [53] comparison: without  $RS$ , there is a significant reduction in the effective dimension, accompanied by a marked increase in instability as training progresses. (Right) Numerical value comparison of estimated bisimulation Bellman residual:  $\hat{\epsilon}$  is persistently greater than 0 in the absence of  $RS$ , which indicates that target  $G$  is invariably larger than  $G_{\phi}$ , suggesting that  $G_{\phi}$  does not achieve steady convergence.
Accordingly, we have  $G_{\sim}^{\pi}(s_i,s_j) \leq \frac{c_r \cdot (R_{\max} - R_{\min})}{1 - c_k}$ . Adopting the conventional settings of  $c_{r} = 1$  and  $c_{k} = \gamma$  as suggested in [6, 52], could possibly result in a relatively large upper bound of  $G_{\sim}^{\pi}$  between states. This is due to the common practice of setting  $\gamma$  at 0.99. However, when bisimulation operators are instantiated with bounded distances, e.g., cosine distance, such a setting may be unsuitable. Therefore, it becomes important to tighten the upper bound.
Besides, we can also derive the value bound between the ground truth value function and the approximated value function:
Theorem 8. (Value bound based on on-policy bisimulation measurements in terms of approximation error). Given an MDP  $\widetilde{\mathcal{M}}$  constructed by aggregating states in an  $\omega$ -neighborhood, and an encoder  $\phi$  that maps from states in the original MDP  $\mathcal{M}$  to these clusters, the value functions for the two MDPs are bounded as
$$
\left| V ^ {\pi} (s) - \widetilde {V} ^ {\pi} (\phi (s)) \right| \leq \frac {2 \omega + \hat {\Delta}}{c _ {r} (1 - \gamma)}. \tag {11}
$$
where  $\hat{\Delta} \coloneqq \|\hat{G}_{\sim}^{\pi} - \hat{G}_{\phi}^{\pi}\|_{\infty}$  is the approximation error.
In essence, Equation 10 and Theorem 8 reveal that: (i) there is a positive correlation between the reward scale  $c_{r}$  and the upper bound of the fixed-point  $G_{\sim}^{\pi}$ , and (ii) a larger reward scale  $c_{r}$  facilitates a more accurate approximation of the value function  $\widetilde{V}^{\pi}(\phi(s))$  to its ground-truth value  $V^{\pi}(s)$ . It is important to note that  $c_{r}$  also impacts the value of  $\hat{\Delta}$ , as depicted in Figure 7(Right) $^{4}$ . Therefore, it is crucial to first ensure the alignment with the instantiation of the bisimulation measurement, and then choose the largest possible  $c_{r}$  to minimize the value error. For instance, as the SimSR operator [52] uses the cosine distance,  $c_{k} = \gamma$  is predetermined. We should thus set  $c_{r} \in [0,1 - \gamma]$ , and apply min-max normalization to the reward function. This can make  $G_{\sim}^{\pi} \leq 1$  and therefore be consistent with the maximum value of 1 of the cosine distance. To achieve a tighter bound in Equation 11, we should then maximize the reward scale, setting  $c_{r}$  to  $1 - \gamma$ . Figure 7 illustrates the effectiveness of this reward scaling.
6.1 Performance Comparison in D4RL Benchmark
Implementation Details We analyze our proposed method on the D4RL benchmark [15] of OpenAI gym MuJoCo tasks [50] which includes a variety of datasets that have been commonly used in the
4Despite Figure 7(Right) depicting the approximate residual  $\hat{\epsilon}$ , we have drawn a connection between  $\epsilon_{\phi}^{\pi}$  and  $\Delta_{\phi}^{\pi}$  in the Appendix, which can reflect the possible situations for  $\hat{\Delta}$ .
Figure 2: Performance comparison on 12 D4RL tasks over 10 seeds with one standard error shaded in the default setting. For every seed, the average return is computed every 10,000 training steps, averaging over 10 episodes. The horizontal axis indicates the number of transitions trained on. The vertical axis indicates the normalized average return.
Timesteps
Figure 3: Bootstrapping distributions for uncertainty in IQM (i.e. inter-quartile mean) measurement on D4RL tasks (left) and visual D4RL tasks (right), following from the performance criterion in [2].
Offline RL community. To illustrate the effectiveness of our method, we implement it on top of two bisimulation-based approaches, MICO [6] and SimSR [52]. It is worth noting that there are two versions of SimSR depending on its use of a latent dynamics model: SimSRbasic follows the dynamics that the environment provides, and SimSR_full constructs latent dynamics for sampling successive latent states. We opt for SimSRbasic as our backbone, as it exhibits superior and more stable performance in the D4RL benchmark tasks compared to SimSR_full. Additionally, to explore the impact of bisimulation-based representation learning on the downstream performance of policy learning, we build these approaches on top of the Offline RL method TD3BC [16]. We examine three environments: halfcheetah, hopper, and walker2d, with four datasets per task: expert, medium-expert, medium-replay, and medium. We first pretrain the encoder during  $100k$  timesteps, then freeze it, pass the raw state through the frozen encoder to obtain the representations that serve as input for the Offline RL algorithm. Further details on the experiment setup are included in Appendix F.
Analysis Figure 2 illustrates the performance of two approaches and their variants in the D4RL tasks. We use EBS to represent the scheme of employing the expectile-based operator, while RS
denotes the reward scaling scheme. The latter includes both min-max reward normalization and penalization coefficient with  $(1 - \gamma)$  in the bisimulation operator. As discussed in Section 5.2, the role of reward scaling varies depending on the specific instantiation of  $G^5$ . We observe that without RS, SimSR almost fails in every dataset, which aligns with our understanding of the critical role reward scaling plays. The results also illustrate that EBS effectively enhances the downstream performance of the policy for both SimSR and MICo. It is noteworthy that in this experiment, we set  $\tau = 0.6$  for the expectile in SimSR and  $\tau = 0.7$  in MICo across all datasets, demonstrating the robustness of this hyperparameter. Regarding SimSR, when RS is applied (SimSR+RS), the performance is comparable to the TD3BC baseline, while the incorporation of the expectile-based operator (SimSR+RS+EBS) further enhances final performance and sample efficiency. Besides, we additionally present the IQM normalized return of all variants in Figure 3, illustrating our performance gains over the backbones. Further, we have also constructed an ablation study to investigate the impact of different settings of  $\tau$ , the results show that a suitable expectile  $\tau$  is crucial for control tasks. We present the corresponding results in Appendix E.
6.2 Performance Comparison in V-D4RL Benchmark
Implementation details We also evaluate our method on a visual observation setting of DMControl suite (DMC) tasks, V-D4RL benchmark [35]. Similar to the previous experiment, we add the proposed schemes on top of MICo and SimSR. In the experiments, we notice that the latent dynamics modeling can help to boost performance for the visual setting, hence we use SimSR_full as the backbone. Additionally, we also notice that MICo often gives really poor performance in the V-D4RL benchmark, while adding latent dynamics alleviates the issue. Therefore, we boost MICo with explicit dynamics modeling for a fair comparison. To compare the performance with the other representation approaches, we include 4 competitive representation learning approaches for Offline RL, including DRIML [38], HOMER [39],
CURL [31], and Inverse model [44]. Detailed descriptions of these approaches can be found in Appendix G.
Figure 4: Performance comparison on V-D4RL benchmark.
Analysis We evaluate all aforementioned approaches by integrating the pre-trained encoder from each into an Offline RL method  $\mathrm{DrQ + BC}$  [35], which combines data augmentation techniques with TD3BC. The results in Table 1 and Figure 4 illustrate the effectiveness of our proposed method, the numerical improvements are underlined with red upward arrows. Compared to the other baselines, while SimSR+RS+EBS does not achieve the highest score in all datasets, it achieves the best overall performance. Besides, our modifications on MICo and SimSR consistently show significant improvements. This indicates that our proposed method is not only applicable to raw-state inputs but also compatible with pixel-based observations.
Limitations and Future Work While  $\tau$  remains constant in our D4RL experiments, optimal performance may arise under different  $\tau$  settings, contingent on the specific attributes of the dataset. Therefore, to yield the best outcomes, one might need to set various  $\tau$  to identify the most suitable value. However, this process could consume substantial computational resources. Another area of potential study involves evaluating the effectiveness of our approach in off-policy settings, given that off-policy settings may also lead to similar failure cases.
<sup>5</sup>Since MICo does not necessitate a particular upper bound, RS may be harmful to its performance. Our experiments have substantiated this observation, leading us to exclude the MICo+RS results from Figure 2.
Table 1: Performance comparison with several other baselines on V-D4RL benchmark, averaged on 3 random seeds.
Dataset	CURL	DRIMLC	HOMER	ICM	MICo → MICo+EBS	SimSR → SimSR+RS+EBS
cheetah-run-medium	392	524	475	365	177 → 449 (↗272)	391 → 491(↗100)
walker-walk-medium	452	425	439	358	450 → 447 (—)	443 → 480(↗37)
cheetah-run-medium-replay	271	395	306	251	335 → 357 (↗22)	374 → 462(↗88)
walker-walk-medium-replay	265	235	283	167	207 → 240 (↗33)	197 → 240(↗43)
cheetah-run-medium-expert	348	403	383	280	282 → 341 (↗59)	360 → 547(↗187)
walker-walk-medium-expert	729	399	781	606	586 → 635(↗49)	755 → 845(↗90)
cheetah-run-expert	200	310	218	237	308 → 331(↗23)	409 → 454(↗45)
walker-walk-expert	769	427	686	850	370 → 447 (↗77)	578 → 580 (—)
total	3426	3118	3571	3114	2715 → 3253 (↗538)	3507 → 4043 (↗536)
Conclusion In this work, we highlight the effectiveness of the bisimulation operator over incomplete datasets and emphasize the crucial role of reward scaling in Offline settings. By employing the expectile operator in bisimulation, we manage to strike a balance between behavior measurement and greedy assignment of the measurement over datasets. We also propose a reward scaling strategy to reduce the risk of representation collapse in specific bisimulation-based measurements. Empirical studies show the effectiveness of our proposed modifications.
This work was partially supported by the NSFC under Grants 92270125 and 62276024, as well as the National Key R&D Program of China under Grant No.2022YFC3302101.
[52] Hongyu Zang, Xin Li, and Mingzhong Wang. Simsr: Simple distance-based state representations for deep reinforcement learning. In Thirty-Sixth AAAI Conference on Artificial Intelligence, AAAI 2022, Thirty-Fourth Conference on Innovative Applications of Artificial Intelligence, IAAI 2022, The Twelveth Symposium on Educational Advances in Artificial Intelligence, EAAI 2022 Virtual Event, February 22 - March 1, 2022, pages 8997-9005. AAAI Press, 2022.
[53] Hongyu Zang, Xin Li, Jie Yu, Chen Liu, Riashat Islam, Remi Tachet des Combes, and Romain Laroche. Behavior prior representation learning for offline reinforcement learning. In The Eleventh International Conference on Learning Representations, 2023.
[54] Amy Zhang, Rowan Thomas McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learning invariant representations for reinforcement learning without reconstruction. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021.
[55] Amy Zhang, Shagun Sodhani, Khimya Khetarpal, and Joelle Pineau. Multi-task reinforcement learning as a hidden-parameter block MDP. CoRR, abs/2007.07206, 2020.
[56] Amy Zhang, Shagun Sodhani, Khimya Khetarpal, and Joelle Pineau. Learning robust state abstractions for hidden-parameter block mdps. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021.
A Notation 17
B Algorithm 17
C Technical backgrounds 18
C.1 Bisimulation metric 18
C.2 MICo distance 19
C.3 SimSR operator 19
C.4 Lifted MDP 19
C.5 Expectile Regression 20
D Proof 20
D.1 Connection between bisimulation error and bisimulation Bellman residual 20
D.2 Thoerem 3 22
D.3 Proposition 4 22
D.4 Lemma 5 23
D.5 Lemma 6 24
D.6 Theorem 7 25
D.7 Theorem 8 25
E Understanding of Theorem 7 26
F Additional Experiments 27
F.1 Ablation Study - Value of Expectile 27
F.2 Ablation Study - Effectiveness of Reward Scaling 27
F.3 Case Study on MICo 28
G Additional Related Works 28
H Additional Discussion 29
H.1 The severity of the proposed problem 29
H.2 Suitability of different techniques 29
I Empirical estimation of bisimulation error 30
Table 2 summarizes our notation.
Table 2: Table of Notation.
Notation	Meaning	Notation	Meaning
M	MDP	M	Lifted MDP (auxiliary MDP)
S	state space	A	action space
T	transition function	r	reward function
γ	discount factor	π	policy of the agent
Vπ(s)	state value function given policy π	D	dataset
πβ	behavior policy	μβ(s)	state occupancy of the dataset
Fπ	on-policy bisimulation operator	gπ	π-bisimulation metric
D(·,·)	a specific distance	Gπ	fixed point of MICo and SimSR
φ	state encoder	Gφ(si, sj)	parameterized bisimulation measurement
Δφ	bisimulation error	εφ	bisimulation Bellman residual
μπ(s)	stationary distribution over states on policy π	μπ	the distribution over pairs of states
Eμπ[εφ]	expected on-policy bisimulation Bellman residual	Fπβ	behavior bisimulation operator
τ	expectile term	γτ	discount factor with expectile
Fπβ	behavior bisimulation operator with expectile	ˆ	estimated one-step residual
Gφ	bisimulation measurement parameterized by target encoder	G~(si, ai, sj, aj)	a measurement on state-action space
G*(si, ai, sj, aj)	maximum measurement constrained to dataset	cr	scale term of reward in bisimulation
ck	scale term of transition in bisimulation	Vπ(φ(s))	value function based on state encoder
ω	distance bound of aggregating neighbor	ˆΔ	approximation error of bisimulation measurement
We provide the algorithm in Algorithm 1, and a pytorch-like implementation build on top of SimSR in Algorithm 2.
Algorithm 1 Proposed Implementation
1: Stage 1 Preprocessing:
2: Min-Max reward normalization:  $\bar{r} = \frac{r - r_{\mathrm{min}}}{r_{\mathrm{max}} - r_{\mathrm{min}}}$
3: Stage 2 Pretraining the encoder:
4: Initialize encoder parameter  $\phi$ , expectile  $\tau$ , learning rate  $\alpha$ , discount factor  $\gamma$ .
5: for each gradient step do
6: Apply reward scaling when computing  $\hat{\epsilon}$ :
$$
\hat {\epsilon} = (1 - \gamma) | \bar {r} \left(s _ {i}, a _ {i}\right) - \bar {r} \left(s _ {j}, a _ {j}\right) | + \gamma G _ {\bar {\phi}} ^ {\pi_ {\beta}} \left(s _ {i} ^ {\prime}, s _ {j} ^ {\prime}\right) - G _ {\phi} ^ {\pi_ {\beta}} \left(s _ {i}, s _ {j}\right) \tag {12}
$$
7: Update encoder  $\phi$
$$
\phi \leftarrow \phi - 2 \alpha \mathbb {E} _ {a _ {i} \sim \pi_ {\beta} (\cdot | s _ {i}), a _ {j} \sim \pi_ {\beta} (\cdot | s _ {j})} [ \tau [ \hat {\epsilon} ] _ {+} + (1 - \tau) [ \hat {\epsilon} ] _ {-} ] \tag {13}
$$
9: Stage 3 Training value function and policy network:
10: Initialize value function parameter  $\psi$ , policy network parameter  $\theta$ , learning rate  $\lambda_{V}$  and  $\lambda_{\pi}$ .
11: for each gradient step do
12: Sample tuple  $(s, a, s', \bar{r})$  from dataset  $\mathcal{D}$
13: Encode the states to representation space:  $z = \phi (s),z^{\prime} = \phi (s^{\prime})$
14: Update value function with  $(z,a,z^{\prime},\bar{r})$
$$
\psi = \psi - \lambda_ {V} \nabla_ {\psi} \mathcal {L} _ {V} (\psi). \tag {14}
$$
15: Update policy network with  $(z,a,z^{\prime},\bar{r})$
$$
\theta = \theta - \lambda_ {\pi} \nabla_ {\theta} \mathcal {L} _ {\pi} (\theta). \tag {15}
$$
Algorithm 2 SimSR+RS+EBS Pseudocode, PyTorch-like
class ReplayBuffer(object):
    def __init__(self):
        ...
        selfreward_normalization()
        ...
        ...
        def reward_normalization(self):
            r_max = selfreward.max()
            r_min = selfreward.min()
            selfreward = (selfreward - r_min) / (r_max - r_min)
def compute_distance(features_a, features_b):
    similarity_matrix = torch/matmul_features_a, features_b.T)
    dis = 1-similarity_matrix
    return dis
def expectile_loss(diff, expectile):
    weight = torch.where(diff > 0, expectile, (1 - expectile))
    return weight * (diff ** 2)
# encoder: mlp, encoder network, the output is l2-normalized
# target Encoder: mlp, same as encoder, updated by EMA
# discount: discount factor  $\mathbb{S}\backslash \mathrm{gamma}\mathbb{S}$ 
# slope: expectile  $\mathbb{S}\backslash \mathrm{tau}\mathbb{S}$ 
def compute_ebs_loss(Encoder, target Encoder, replay_buffer, batch_size, discount, slope):
    observation, action, reward, discount, nextobservation = replay_buffer.sample(batch_size) # sample a batch of tuples from replay buffer
    latent_state = encoder(observation)
    latent_next_state = target Encoder(next Observation)
    r_diff = (1 - discount) * torch.abs(reward.T - reward)
    next_diff = compute_distance(latent_next_state, latent_next_state)
    z_diff = compute_distance(latent_state, latent_state)
    bisimilarity = r_diff + discount * next_diff
    encoder_loss = expectile_loss(bisimilaritydetach() - z_diff, slope)
    encoder_loss = encoder_loss.mean()
    return encoder_loss
C Technical backgrounds
C.1 Bisimulation metric
Bisimulation measures equivalence relations on MDPs with a recursive form: two states are deemed equivalent if they share the equivalent distributions over the next equivalent states and they have the same immediate reward [30, 20]. However, since bisimulation considers equivalence for all actions, including bad ones, it commonly results in "pessimistic" outcomes. Instead, [5] developed  $\pi$ -bisimulation which removes the requirement of considering each action and only needs to consider the actions induced by a policy  $\pi$ .
Definition 9. [5] Given an MDP  $\mathcal{M}$ , an equivalence relation  $E^{\pi} \subseteq S \times S$  is a  $\pi$ -bisimulation relation if whenever  $(\mathbf{s}, \mathbf{u}) \in E^{\pi}$  the following properties hold:
1.  $r(s,\pi) = r(u,\pi)$
2.  $\forall C\in S_{E^{\pi}},T(C|s,\pi) = T(C|u,\pi)$
where  $S_{E^{\pi}}$  is the state space  $S$  partitioned into equivalence classes defined by  $E^{\pi}$ . Two states  $s, u \in S$  are  $\pi$ -bisimilar if there exists a  $\pi$ -bisimulation relation  $E^{\pi}$  such that  $(s, u) \in E^{\pi}$ .
However,  $\pi$ -bisimulation is still too stringent to be applied at scale as  $\pi$ -bisimulation relation emphasizes the equivalence is a binary property: either two states are equivalent or not, thus becoming too sensitive to perturbations in the numerical values of the model parameters. The problem becomes even more prominent when deep frameworks are applied.
Thereafter, they proposed a  $\pi$ -bisimulation metric to leverage the absolute value between the immediate rewards w.r.t. two states and the 1-Wasserstein distance  $(\mathcal{W}_1)$  between the transition distributions conditioned on the two states and the policy  $\pi$  to formulate such measurement: