SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Hongyu Zang $^{1}$ , Xin Li $^{1,*}$ , Leiji Zhang $^{1}$ , Yang Liu $^{2}$ , Baigui Sun $^{2}$ , Riashat Islam $^{3}$ , Rémi Tachet des Combes $^{4}$ , Romain Laroche
<sup>1</sup> Beijing Institute of Technology, China <sup>2</sup> Alibaba Group, China
<sup>3</sup> McGill University, Mila, Quebec AI Institute, Canada <sup>4</sup> Wayve, UK
{zanghyu,xinli,ljzhang}@bit.edu.cn
{ly261666,baigui.sbg}@alibaba-inc.com
riashat.islam@mail.mcgill.ca
{remi.tachet,romain.laroche}@gmail.com
While bisimulation-based approaches hold promise for learning robust state representations for Reinforcement Learning (RL) tasks, their efficacy in offline RL tasks has not been up to par. In some instances, their performance has even significantly underperformed alternative methods. We aim to understand why bisimulation methods succeed in online settings, but falter in offline tasks. Our analysis reveals that missing transitions in the dataset are particularly harmful to the bisimulation principle, leading to ineffective estimation. We also shed light on the critical role of reward scaling in bounding the scale of bisimulation measurements and of the value error they induce. Based on these findings, we propose to apply the expectile operator for representation learning to our offline RL setting, which helps to prevent overfitting to incomplete data. Meanwhile, by introducing an appropriate reward scaling strategy, we avoid the risk of feature collapse in representation space. We implement these recommendations on two state-of-the-art bisimulation-based algorithms, MICo and SimSR, and demonstrate performance gains on two benchmark suites: D4RL and Visual D4RL. Codes are provided at https://github.com/zanghyu/Offline_Bisimulation.
Reinforcement learning (RL) algorithms often require a significant amount of data to achieve optimal performance [40, 48, 22]. In scenarios where collecting data is costly or impractical, Offline RL methods offer an attractive alternative by learning effective policies from previously collected data [29, 43, 32, 35, 16, 24]. However, capturing the complex structure of the environment from limited data remains a challenge for Offline RL [4]. This involves pre-training the state representation on offline data and then learning the policy upon the fixed representations [51, 47, 41, 53]. Though driven by various motivations, previous methods can be mainly categorized into two classes: i) implicitly shaping the agent's representation of the environment via prediction and control of some aspects of the environment through auxiliary tasks, e.g., maximizing the diversity of visited states [34, 10], exploring attentive contrastive learning on sub-trajectories [51], or capturing temporal information about the environment [47]; ii) utilizing behavioral metrics, such as bisimulation metrics [11, 13, 5], to capture complex structure in the environment by measuring the similarity of behavior on the
*Correspondence to Xin Li.
representations [52, 7]. The former methods have proven their effectiveness theoretically and empirically in Offline settings [41, 47, 51], while the adaptability of the latter approaches in the context of limited datasets remains unclear. This paper tackles this question.
Bisimulation-based approaches, as their name suggests, utilize the bisimulation metrics update operator to construct an auxiliary loss and learn robust state representations. These representations encapsulate the behavioral similarities between states by considering the difference between their rewards and dynamics. While the learned representations possess several desirable properties, such as smoothness [19], visual invariance [54, 1, 52], and task adaptation [56, 37, 46, 8], bisimulation-based objectives in most approaches are required to be coupled with the policy improvement procedure [54, 6, 52]. In Offline RL, pretraining state representations via bisimulation-based methods is supposed to be cast as a special case of on-policy bisimulation metric learning where the behavior policy is fixed so that good performance should ensue. However, multiple recent studies [51, 21] suggest that bisimulation-based algorithms yield significantly poorer results on Offline tasks compared to a variety of (self-)supervised objectives.
In this work, we highlight problems with using the bisimulation principle as an objective in Offline settings. We aim to provide a theoretical understanding of the performance gap in bisimulation-based approaches between online and offline settings: "why do bisimulation approaches perform well in Online RL tasks but tend to fail in Offline RL ones?" By establishing a connection between the Bellman and bisimulation operators, we uncover that missing transitions, which often occur in Offline settings, can cause the bisimulation principle to be compromised. This means that the bisimulation estimator can be ineffective in finite datasets. Moreover, we notice that the scale of the reward impacts the upper bounds of both the bisimulation measurement $^2$  fixed point and the value error. This scaling term, if not properly handled, can potentially lead to representation collapse.
To alleviate the aforementioned issues, we propose to learn state representations based on the expectile operator. With this asymmetric operator predicting expectiles of the representation distribution, we can achieve a balance between the behavior measurement and the greedy assignment of the measurement over the dataset. This results in a form of regularization over the bisimulation measurement, thus preventing overfitting to the incomplete data, and implicitly avoiding out-of-distribution estimation errors. Besides, by considering the specific properties of different bisimulation measurements, we investigate the representation collapse issue for the ones that are instantiated with bounded distances (e.g., cosine distance) and propose a way to scale rewards that reduces collapse. We integrate these improvements mainly on two bisimulation-based baselines, MICo [7] and SimSR [52], and show the effectiveness of the proposed modifications.
The primary contributions of this work are as follows:
- We investigate the potential harm of directly applying the bisimulation principle in Offline settings, prove that the bisimulation estimator can be ineffective in finite datasets, and emphasize the essential role of reward scaling.
- We propose theoretically motivated modifications on two representative bisimulation-based baselines, including an expectile-based operator and a tailored reward scaling strategy. These proposed changes are designed to address the challenges encountered when applying the bisimulation principle in offline settings.
- We demonstrate the superior performance our approach yields through an empirical study on two benchmark suites, D4RL [15] and Visual D4RL [35].
State representation learning in Offline RL Pretraining representations has been recently studied in Offline RL settings, where several studies presented its effectiveness [3, 47, 41, 25]. In this paradigm, we learn state representations on pre-collected datasets before value estimation or policy improvement steps are run. The learned representation can then be used for subsequent policy learning, either online or offline. Some typical auxiliary tasks for pretraining state representations include capturing the dynamical [42] and temporal [47] information of the environment, exploring
<sup>2</sup>Since some bisimulation-based approaches do not exactly use metrics but instead of pseudometrics, diffuse metrics or else, we will use the term "measurement" in the following.
attentive contrastive learning on sub-trajectories [51], or improving policy performance by applying data augmentations techniques to the pixel-based inputs [9, 35].
Bisimulation-based methods The pioneer works by [20, 33] aim to overcome the curse of dimensionality by defining equivalence relations between states to reduce system complexity. However, these approaches are impractical as they usually demand an exact match of transition distributions. To address this issue, [12, 14] propose a bisimulation metric to aggregate similar states. This metric quantifies the similarity between two states and serves as a distance measure to allow efficient state aggregation. Unfortunately, it remains computationally expensive as it requires a full enumeration of states. Later, [5] devise an on-policy bisimulation metric for policy evaluation, providing a scalable method for computing state similarity. Building upon this, [54] develop a metric to learn state representations by modeling the latent dynamic transition as Gaussian. [6] further investigate the independent couple sampling strategy to reduce the computational complexity of representation learning, whereas [52] propose to learn state representations built on the cosine distance to alleviate a representation collapse issue. Despite the promising results obtained, one of the major remaining challenges in this paradigm is its dependency on coupling state representation learning with policy training. This is not always suitable for Offline settings, given that obtaining on-policy reward and transition differences is infeasible due to our inability to gather additional agent-environment interactions. To adapt bisimulation-based approaches to Offline settings, one solution is to consider the policy over the dataset as a specific behavior policy, and then apply the bisimulation principle on it to learn state representations in a pretraining stage, thus disentangling policy training from bisimulation-based learning. Notably, although there exist recent studies [51, 42] investigating the potential of bisimulation-based methods to pretrain state representations, it has not yielded satisfactory results yet [51].
3.1 Offline RL
We consider the standard Markov decision process (MDP) framework, in which the environment is given by a tuple  $\mathcal{M} = (\mathcal{S},\mathcal{A},T,r,\gamma)$ , with state space  $\mathcal{S}$ , action space  $\mathcal{A}$ , transition function  $T$  that decides the next state  $s^{\prime}\sim T(\cdot |s,a)$ , reward function  $r(s,a)$  bounded by  $[R_{\min},R_{\max}]$ , and a discount factor  $\gamma \in [0,1)$ . The agent in state  $s\in S$  selects an action  $a\in \mathcal{A}$  according to its policy, mapping states to a probability distribution over actions:  $a\sim \pi (\cdot |s)$ . We make use of the state value function  $V^{\pi}(s) = \mathbb{E}_{\mathcal{M},\pi}\left[\sum_{t = 0}^{\infty}\gamma^{t}r\left(s_{t},a_{t}\right)\mid s_{0} = s\right]$  to describe the long term discounted reward of policy  $\pi$  starting at state  $s$ . In the sequel, we use  $T_s^a$  and  $r_s^a$  to denote  $T(\cdot |s,a)$  and  $r(s,a)$ , respectively. In Offline RL, we are given a fixed dataset of environment interactions that include  $N$  transition samples, i.e.  $\mathcal{D} = \{s_i,a_i,s_i',r_i\}_{i = 1}^N$ . We assume that the dataset  $\mathcal{D}$  is composed of trajectories generated i.i.d. under the control of a behavior policy  $\pi_{\beta}$ , whose state occupancy is denoted by  $\mu_{\beta}(s)$ .
3.2 Bisimulation-based Update Operator
The concept of bisimulation is used to establish equivalence relations on states. This is done recursively by considering two states as equivalent if they have the same distribution over state transitions and the same immediate reward [30, 20]. Since bisimulation considers worst-case differences between states, it commonly results in "pessimistic" outcomes. To address this limitation, the  $\pi$ -bisimulation metric was proposed in [5]. This new metric only considers actions induced by a given policy  $\pi$  rather than all actions when measuring the behavior distance between states:
Theorem 1. [5] Let  $\mathbb{M}$  be the set of all measurements on  $S$ . Define  $\mathcal{F}^{\pi}:\mathbb{M}\to \mathbb{M}$  by
$$
\mathcal {F} ^ {\pi} (g) \left(s _ {i}, s _ {j}\right) = \left| r _ {s _ {i}} ^ {\pi} - r _ {s _ {j}} ^ {\pi} \right| + \gamma \mathcal {W} (g) \left(T _ {s _ {i}} ^ {\pi}, T _ {s _ {j}} ^ {\pi}\right) \tag {1}
$$
where  $s_i, s_j \in S$ ,  $r_{s_i}^\pi = \sum_{a \in \mathcal{A}} \pi(a|s_i)r_{s_i}^a$ ,  $T_{s_i}^\pi = \sum_{a \in \mathcal{A}} \pi(a|s_i)T_{s_i}^a$ , and  $\mathcal{W}(g)$  is the Wasserstein distance with cost function  $g$  between distributions. Then  $\mathcal{F}^\pi$  has a least fixed point  $g_{\sim}^\pi$ , and  $g_{\sim}^\pi$  is a  $\pi$ -bisimulation metric.
Although it is feasible to compute the behavior difference measurement  $g_{\sim}^{\pi}$  by applying the operator  $\mathcal{F}^{\pi}$  iteratively (which is guaranteed to converge to a fixed point since  $\mathcal{F}^{\pi}$  is a contraction), this approach comes at a high computational complexity due to the Wasserstein distance on the right-hand
side of the equation. To tackle this issue, MICo [6] proposed using an independent couple sampling strategy instead of optimizing the overall coupling of the distributions  $T_{s_i}^{\pi}$  and  $T_{s_j}^{\pi}$ , resulting in a novel measurement to evaluate the difference between states. Additionally, SimSR [52] further explored the potentiality of combining the cosine distance with bisimulation-based measurements to learn state representations. Both works can be generalized as:
$$
\mathcal {F} ^ {\pi} G ^ {\pi} \left(s _ {i}, s _ {j}\right) = \left| r _ {s _ {i}} ^ {\pi} - r _ {s _ {j}} ^ {\pi} \right| + \gamma_ {\substack {\mathbb {E} _ {s _ {i} ^ {\prime}} \sim T _ {s _ {i}} ^ {\pi} \\ s _ {j} ^ {\prime} \sim T _ {s _ {j}} ^ {\pi}}} \left[ G ^ {\pi} \left(s _ {i} ^ {\prime}, s _ {j} ^ {\prime}\right) \right], \tag{2}
$$
and  $\mathcal{F}^{\pi}$  has a least fixed point  $G_{\sim}^{\pi^3}$ . The instantiation of  $G$  varies in different approaches [6, 52]. For example, in SimSR [52], the cosine distance is used to instantiate  $G$  on the embedding space, and the dynamics difference is computed by the cosine distance between the next-state pair  $(s_i', s_j')$  sampled from a transition model of the environment. A more detailed description can be found in Appendix C.
Lemma 2. [6] (Lifted MDP) The bisimulation-based update operator  $\mathcal{F}^{\pi}$  for  $\mathcal{M}$  is the Bellman evaluation operator for a specific lifted MDP.
Due to this interpretation of the bisimulation-based update operator as the Bellman evaluation operator in a lifted MDP, we can derive certain conclusions about bisimulation by drawing inspiration from policy evaluation methods. In the next section, we will borrow analytical ideas from [17] to prove that the bisimulation-based objective may be ineffective for finite datasets. We summarize all notations in Appendix A and provide all proofs in Appendix D.
4 Ineffective Bisimulation Estimators in Finite Datasets
The high-level idea of bisimulation-based state representation learning is to learn state embeddings such that when states are projected onto the embedding space, their behavioral similarity is maintained. We denote our parameterized state encoder by  $\phi : S \to \mathbb{R}^n$  and a distance  $D(\cdot, \cdot)$  in the embedding space  $\mathbb{R}^n$  by  $G_{\phi}^{\pi}(s_i, s_j) \doteq D(\phi(s_i), \phi(s_j))$ . For instance,  $D(\cdot, \cdot)$  may be the Lukaszyk-Karmowski distance [6] or the cosine distance [52]. To avoid unnecessary confusion, we defer implementation details to Section 5.
When considering bisimulation-based state representations, the goal is to acquire stable state representations under policy  $\pi$  via the measurement  $G_{\sim}^{\pi}$ . The primary focus is usually to minimize a loss over the bisimulation error, denoted by  $\Delta_{\phi}^{\pi}$ , which measures the distance between the approximation  $G_{\phi}^{\pi}$  and the fixed point  $G_{\sim}^{\pi}$ :
$$
\Delta_ {\phi} ^ {\pi} \left(s _ {i}, s _ {j}\right) := \left| G _ {\phi} ^ {\pi} \left(s _ {i}, s _ {j}\right) - G _ {\sim} ^ {\pi} \left(s _ {i}, s _ {j}\right) \right|. \tag {3}
$$
However, since the fixed point  $G_{\sim}^{\pi}$  is unobtainable without full knowledge of the underlying MDP, this approximation error is often unknown. Recall that in Lemma 2, we have shown that we can connect a bisimulation-based update operator to a lifted MDP. Taking inspiration from Bellman evaluation for the value function, we define the bisimulation Bellman residual  $\epsilon_{\phi}^{\pi}$  as:
$$
\epsilon_ {\phi} ^ {\pi} \left(s _ {i}, s _ {j}\right) := \left| G _ {\phi} ^ {\pi} \left(s _ {i}, s _ {j}\right) - \mathcal {F} ^ {\pi} G _ {\phi} ^ {\pi} \left(s _ {i}, s _ {j}\right) \right|. \tag {4}
$$
Then, we can connect the bisimulation Bellman residual with the bisimulation error by the following:
Theorem 3. (Bisimulation error upper-bound). Let  $\mu_{\pi}(s)$  denote the stationary distribution over states, let  $\mu_{\pi}(\cdot ,\cdot)$  denote the joint distribution over synchronized pairs of states  $(s_i,s_j)$  sampled independently from  $\mu_{\pi}(\cdot)$ . For any state pair  $(s_i,s_j)\in S\times S$ , the bisimulation error  $\Delta_{\phi}^{\pi}(s_i,s_j)$  can be upper-bounded by a sum of expected bisimulation Bellman residuals  $\epsilon_{\phi}^{\pi}$ :
$$
\Delta_ {\phi} ^ {\pi} \left(s _ {i}, s _ {j}\right) \leq \frac {1}{1 - \gamma} \mathbb {E} _ {\left(s _ {i} ^ {\prime}, s _ {j} ^ {\prime}\right) \sim \mu_ {\pi}} \left[ \epsilon_ {\phi} ^ {\pi} \left(s _ {i} ^ {\prime}, s _ {j} ^ {\prime}\right) \right]. \tag {5}
$$
Thereafter, the bisimulation Bellman residual is used as a surrogate objective to approximate the fixed point  $G_{\sim}^{\pi}$  when learning our state representation. Indeed, the minimization of the bisimulation Bellman residual objective over all pairs  $(s_i', s_j') \sim \mu_{\pi}$  leads to the minimization of the corresponding bisimulation error. This ensures that if the expected on-policy bisimulation Bellman residual (i.e.,
3For readability, we will conflate the notations  $G^{\pi}$  and  $G^{\pi}(x,y)$  , they are the same if not specified
$\mathbb{E}_{\mu_{\pi}}[\epsilon_{\phi}^{\pi}]$ , and we will use the term "expected bisimulation residual" in following) minimization objective is zero, then the bisimulation error must be zero for the state pairs under the same policy. However, when the dataset is limited, rather than an infinite transition set covering the whole MDP, minimizing the expected bisimulation residual will no longer be sufficient to guarantee a zero bisimulation error.
Proposition 4. (The expected bisimulation residual is not sufficient over incomplete datasets). If there exists states  $s_i'$  and  $s_j'$  not contained in dataset  $\mathcal{D}$ , where the occupancy  $\mu_{\pi}(s_i'|s_i, a_i) > 0$  and  $\mu_{\pi}(s_j'|s_j, a_j) > 0$  for some  $(s_i, s_j) \sim \mu_{\pi}$ , then there exists a bisimulation measurement  $G_{\phi}^{\pi}$  and a constant  $C > 0$  such that
- For all  $(\hat{s}_i, \hat{s}_j) \in \mathcal{D}$ , the bisimulation Bellman residual  $\epsilon_{\phi}^{\pi}(\hat{s}_i, \hat{s}_j) = 0$ .
- There exists  $(s_i, s_j) \in \mathcal{D}$ , such that the bisimulation error  $\Delta_{\phi}^{\pi}(s_i, s_j) = C$ .
As an example, if we only have  $(s_i,a_i,r,s_i')$  and  $(s_j,a_j,r,s_j')$  in a dataset, where both rewards equal to zero for state  $s_i$  and  $s_j$ , and if we choose  $G_{\phi}^{\pi}(s_i,s_j) = C$ , and  $G_{\phi}^{\pi}(s_i',s_j') = \frac{1}{\gamma} C$ , then the bisimulation Bellman residual is  $\epsilon_{\phi}^{\pi}(s_i,s_j) = 0$ , while the bisimulation error  $\Delta_{\phi}^{\pi} = G_{\phi}^{\pi}(s_i,s_j) - 0 = C$  is strictly positive. Note that this failure case does not involve modifying the environment in an extremely adversarial manner, it simply occurs when we are required to estimate the representation of states with subsequent states that are missing from the dataset. Since the distance between the missing states can be arbitrarily large as they are out-of-distribution, directly minimizing the Bellman bisimulation error could achieve the minimal Bellman bisimulation error over the dataset, while not necessarily improving the state representation.
In the context of Offline RL, since the dataset is finite, bisimulation-based representation learning ought to be conceptualized as a pretraining process over the behavior policy  $\pi_{\beta}$  of the dataset  $\mathcal{D}$ . However, the failure case above indicates that applying the bisimulation operator  $\mathcal{F}^{\pi_{\beta}}$  and minimizing the associated Bellman bisimulation error does not necessarily ensure the sufficiency of the learned representation for downstream tasks. Ideally, if we had access to the fixed-point measurement  $G_{\sim}^{\pi_{\beta}}$ , then we could directly minimize the error between the approximation  $G$  and the fixed-point  $G_{\sim}^{\pi_{\beta}}$ . However, given the static and incomplete nature of the dataset, acquiring the fixed-point  $G_{\sim}^{\pi_{\beta}}$  explicitly is not feasible. From another perspective, the failure stems from out-of-distribution estimation errors. Assuming we could estimate the bisimulation exclusively with in-sample learning, this issue could be intuitively mitigated. As such, we resort to expectile regression as a regularizer, allowing us to circumvent the need for out-of-sample / unseen state pairs.
In this section, we describe how we adapt existing bisimulation-based representation approaches to offline RL. We use the expectile-based operator to learn state representations that optimize the behavior measurement over the dataset, while avoiding overfitting to the incomplete data. In addition, we analyze the impact of reward scaling and propose as a consequence to normalize the reward difference in the bisimulation Bellman residual in order to satisfy the specific nature of different instantiations of the bisimulation measurement while keeping a lower value error. The pseudo-code of our method is shown in Algorithms in Appendix B.
5.1 Expectile-based Bisimulation Operator
The efficacy of expectile regression in achieving in-sample learning has already been demonstrated in previous research [28, 36]. Consequently, we will first describe our proposed expectile-based operator, and subsequently show how expectile regression can effectively address the aforementioned challenge. Specifically, we consider the update operator as follows:
$$
\begin{array}{l} \left(\mathcal{F}_{\tau}^{\pi_{\beta}}G_{\phi}^{\pi_{\beta}}\right)(s_{i},s_{j}):= \operatorname *{arg  min}_{G_{\phi}^{\pi_{\beta}}}\mathbb{E}_{a_{i}\sim \pi_{\beta}(\cdot |s_{i}),a_{j}\sim \pi_{\beta}(\cdot |s_{j})}\left[\tau [\hat{\epsilon} ]_{+}^{2} + (1 - \tau)[-\hat{\epsilon} ]_{+}^{2}\right], \\ \hat {\epsilon} = \mathbb {E} _ {s _ {i} ^ {\prime} \sim T _ {s _ {i}} ^ {\pi_ {\beta}}} \underset {s _ {j} ^ {\prime} \sim T _ {s _ {j}} ^ {\pi_ {\beta}}} {\left[ \underbrace {| r (s _ {i} , a _ {i}) - r (s _ {j} , a _ {j}) | + \gamma G _ {\phi} ^ {\pi_ {\beta}} \left(s _ {i} ^ {\prime} , s _ {j} ^ {\prime}\right)} _ {\text {t a r g e t} G} - G _ {\phi} ^ {\pi_ {\beta}} \left(s _ {i}, s _ {j}\right) \right],} \tag {6} \\ \end{array}
$$
where  $\hat{\epsilon}$  is the estimated one-step bisimulation Bellman residual,  $\pi_{\beta}$  is the behavior policy,  $G_{\bar{\phi}}$  is the target encoder, updated using an exponential moving average, and  $[\cdot]_{+} = \max(\cdot, 0)$ . Since the expectile operator in Equation 6 does not have a closed-form solution, in practice, we minimize it through gradient descent steps:
$$
G _ {\phi} ^ {\pi_ {\beta}} \left(s _ {i}, s _ {j}\right) \leftarrow G _ {\phi} ^ {\pi_ {\beta}} \left(s _ {i}, s _ {j}\right) - 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 {7}
$$