SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Qitong Gao* Ge Gao† Juncheng Dong* Vahid Tarokh* Min Chi† Miroslav Pajic*
Off-policy evaluation (OPE) is important for closing the gap between offline training and evaluation of reinforcement learning (RL), by estimating performance and/or rank of target (evaluation) policies using offline trajectories only. It can improve the safety and efficiency of data collection and policy testing procedures in situations where online deployments are expensive, such as healthcare. However, existing OPE methods fall short in estimating human feedback (HF) signals, as HF may be conditioned over multiple underlying factors and is only sparsely available; as opposed to the agent-defined environmental rewards (used in policy optimization), which are usually determined over parametric functions or distributions. Consequently, the nature of HF signals makes extrapolating accurate OPE estimations to be challenging. To resolve this, we introduce an OPE for HF (OPEHF) framework that revives existing OPE methods in order to accurately evaluate the HF signals. Specifically, we develop an immediate human reward (IHR) reconstruction approach, regularized by environmental knowledge distilled in a latent space that captures the underlying dynamics of state transitions as well as issuing HF signals. Our approach has been tested over two real-world experiments, adaptive in-vivo neurostimulation and intelligent tutoring, as well as in a simulation environment (visual Q&A). Results show that our approach significantly improves the performance toward estimating HF signals accurately, compared to directly applying (variants of) existing OPE methods.
Off-policy evaluation (OPE) aims to estimate the performance of reinforcement learning (RL) policies using only a fixed set of offline trajectories [61], i.e., without online deployments. It is considered to be a critical step in closing the gap between offline RL training and evaluation, for environments and systems where online data collection is expensive or unsafe. Specifically, OPE facilitates not only offline evaluation of the safety and efficacy of the policies ahead of online deployment, but policy selection as well; this allows one to maximize the efficiency when online data collection is possible, by identifying and deploying the policies that are more likely to result in higher returns. OPE has been used in various application domains including healthcare [68, 53, 23, 22], robotics [15, 18, 24], intelligent tutoring [64, 45, 17], recommendation systems [50, 43].
The majority of existing OPE methods focus on evaluating the policies' performance defined over the environmental reward functions which are mainly designed for use in policy optimization (training). However, as an increasing number of offline RL frameworks are developed for human-involved systems [64, 45, 1, 48, 16], existing OPE methods lack the ability to estimate how human users would evaluate the policies, e.g., ratings provided by patients (on a scale of 1-10) over the procedure facilitated by automated surgical robots; as human feedback (HF) can be noisy and conditioned over various confounders that could be difficult to be captured explicitly [53, 7, 44]. For example, patient satisfaction over a specific diabetes therapy may vary across the cohort, depending on many subjective
*Duke University. Durham, NC, USA. Contact: {qitong.gao, miroslav.pajic}@duke.edu.
†North Carolina State University. Raleigh, NC, USA.
factors, such as personal preferences and activity level of the day, while participating in the therapy, in addition to the physiological signals (e.g., blood sugar level, body weight) that are more commonly used as the sources for determining environmental rewards toward policy optimization [70, 33, 21, 19]. Moreover, the environmental rewards are sometimes discrete to ensure optimality of the learned policies [67], which further reduces its correlation against HF signals.
In this work, we introduce the OPE for human feedback (OPEHF) framework that revives existing OPE approaches in the context of evaluating HF from offline data. Specifically, we consider the challenging scenario where the HF signal is only provided at the end of each episode - i.e., no per-step HF signals, referred to as immediate human rewards (IHRs) below, are provided - benchmarking the common real-world situations where the participants are allowed to rate the procedures only at the end of the study. The goal is set to estimate the end-of-episode HF signals, also referred to as human returns, over the target (evaluation) policies, using a fixed set of offline trajectories collected over some behavioral policies. To facilitate OPEHF, we introduce an approach that first maps the human return back to the sequence of IHRs, over the horizon, for each trajectory. Specifically, this follows from optimizing over an objective that consists of a necessary condition where the cumulative discounted sum of IHRs should equal the human return, as well as a regularization term that limits the discrepancy of the reconstructed IHRs over state-action pairs that are determined similar over a latent representation space into which environmental transitions and rewards are encoded. At last, this allows for the use of any existing OPE methods to process the offline trajectories with reconstructed IHRs and estimate human returns under target policies.
Our main contributions are tri-fold. (i) We introduce a novel OPEHF framework that revives existing OPE methods toward accurately estimating highly sparse HF signals (provided only at the end of each episode) from offline trajectories, through IHRs reconstruction. (ii) Our approach does not require the environmental rewards and the HF signals to be strongly correlated, benefiting from the design where both signals are encoded to a latent space regularizing the objective for reconstructions of IHRs, which is justified empirically over real-world experiments. (iii) Two real-world experiments, i.e., adaptive in-vivo neurostimulation for the treatment of Parkinson's disease and intelligent tutoring for computer science students in colleges, as well as one simulation environment (i.e., visual Q&A), facilitated the thorough evaluation of our approach; various degrees of correlations between the environment rewards and HF signals existed across the environments, as well as the varied coverage of the state-action space provided by offline data over sub-optimal behavioral policies, imposing different levels of challenges for OPEHF.
2 Off-Policy Evaluation for Human Feedback (OPEHF)
In this section, we introduce an OPEHF framework that allows for the use of existing OPE methods to estimate the human returns that are available only at the end of each episode, with IHRs remaining unknown. This is in contrast to the goal of classic OPE that only estimates the environmental returns following the user-defined reward function used in the policy optimization phase. A brief overview of existing OPE methods can be found in Appendix C.
We first formulate the human-involved MDP (HMDP), which is a tuple  $\mathcal{M} = (\mathcal{S},\mathcal{A},\mathcal{P},R,R^{\mathcal{H}},s_0,\gamma)$  where  $\mathcal{S}$  is the set of states,  $\mathcal{A}$  the set of actions,  $\mathcal{P}:S\times \mathcal{A}\to S$  is the transition distribution usually captured by probabilities  $p(s_{t}|s_{t - 1},a_{t - 1})$ $R:S\times \mathcal{A}\rightarrow \mathbb{R}$  is the environmental reward function,  $R^{\mathcal{H}}(r^{\mathcal{H}}|s,a)$  is the human reward distribution from which the IHR  $r_t^\mathcal{H}\sim R^\mathcal{H}(\cdot |s_t,a_t)$  are sampled,  $s_0$  is the initial state sampled from the initial state distribution  $p(s_0)$ , and  $\gamma \in [0,1)$  is the discounting factor. Note that we set the IHRs to be determined probabilistically, as opposed to the environmental rewards  $r_t = R(s_t,a_t)$  that are deterministic; this is due to the fact that many underlying factors may affect the feedback provided by humans [53, 7, 44], as we have also observed while performing human-involved experiments (see Appendix D). Finally, the agent interacts with the MDP following some policy  $\pi (a|s)$  that defines the probabilities of taking action  $a$  at state  $s$ .
In this work, we make the following assumption over  $R$  and  $R^{\mathcal{H}}$ .
Assumption 1 (Unknown IHRs). We assume that the immediate environmental reward function  $R$  is known and  $R(s, a)$  can be obtained for any state-action pairs in  $S \times \mathcal{A}$ . Moreover, the IHR distribution  $R^{\mathcal{H}}$  is assumed to be unknown, i.e.,  $r^{\mathcal{H}} \sim R^{\mathcal{H}}(\cdot | s, a)$  are unobservable, for all  $(s, a) \in S \times \mathcal{A}$ .
Instead, the cumulative human return  $G_{0:T}^{\mathcal{H}}$ , defined over  $R^{\mathcal{H}}$ , is given at the end of each trajectory, i.e.,  $G_{0:T}^{\mathcal{H}} = \sum_{t=0}^{T} \gamma^{t} r_{t}^{\mathcal{H}}$ , with  $T$  being the horizon and  $r_{t}^{\mathcal{H}} \sim R^{\mathcal{H}}(\cdot | s_{t}, a_{t})$ .
The assumption above follows the fact that human feedback (HF) is not available until the end of each episode, as opposed to immediate rewards that can be defined over the environment and evaluated for any  $(s_t, a_t)$  pairs at any time. This is especially true in environments such as healthcare where the clinical treatment outcome is not foreseeable until a therapeutic cycle is completed, or in intelligent tutoring where the overall gain from students over a semester is mostly reflected by the final grades. Note that although the setup can be generalized to the scenario where HF can be sparsely obtained over the horizon, we believe that issuing the HF only at the end of each trajectory leads to a more challenging setup for OPE. Consequently, the goal of OPEHF can be formulated as follows.
Problem 1 (Objective of OPEHF). Given offline trajectories collected by some behavioral policy  $\beta$ ,  $\rho^{\beta} = \{\tau^{(0)}, \tau^{(1)}, \ldots, \tau^{(N-1)}|a_{t} \sim \beta(a_{t}|s_{t})\}$ , with  $\tau^{(i)} = [(s_{0}^{(i)}, a_{0}^{(i)}, r_{0}^{(i)}, r_{0}^{\mathcal{H}(i)}, s_{1}^{(i)}), \ldots, (s_{T-1}^{(i)}, a_{T-1}^{(i)}, r_{T-1}^{(i)}, r_{T-1}^{\mathcal{H}(i)}, s_{T}^{(i)}), G_{0:T}^{\mathcal{H}(i)}]$  being a single trajectory,  $N$  the total number of offline trajectories, and  $r_{t}^{\mathcal{H}}$ s being unknown, the objective is to estimate the expected total human return over the unknown state-action visitation distribution  $\rho^{\pi}$  of the target (evaluation) policy  $\pi$ , i.e.,  $\mathbb{E}_{(s,a) \sim \rho^{\pi}, r^{\mathcal{H}} \sim R^{\mathcal{H}}} \left[ \sum_{t=0}^{T} \gamma^{t} r_{t}^{\mathcal{H}} \right]$ .
2.2 Reconstruction of IHRs for OPEHF
We emphasize that the human returns are only issued at the end of each episode, with IHRs remaining unknown. One can set all IHRs from  $t = 0$  to  $t = T - 2$  to be zeros (i.e.,  $r_{0:T - 2}^{\mathcal{H}} = 0$ ), and rescale the cumulative human return to be the IHR at the last step (i.e.,  $r_{T - 1}^{\mathcal{H}} = G_{0:T}^{\mathcal{H}} / \gamma^{T - 1}$ ), to allow the use of existing OPE methods toward OPEHF (Problem 1). However, the sparsity over  $r^{\mathcal{H}}$ 's here may impose difficulties for OPE to estimate the human returns accurately over the target policies. For OPEHF, we start by showing that for the per-decision importance sampling (PDIS) method – a variance-reduction variant of the importance sample (IS) family of OPE methods [61] – if IHRs were to be available, they could reduce the variance in the estimation compared to the rescale approach above.
Recall that the PDIS estimator follows  $\hat{G}_{PDIS}^{\pi} = \frac{1}{N}\sum_{i=1}^{N-1}\sum_{t=0}^{T-1}\gamma^{t}\omega_{0:t}^{(i)}r_{t}^{\mathcal{H}(\rangle)}$ , where  $\omega_{0:t}^{(i)} = \prod_{k=0}^{t}\frac{\pi(a_{k}^{(i)}|s_{k}^{(i)})}{\beta(a_{k}^{(i)}|s_{k}^{(i)})}$  are the PDIS weights for offline trajectory  $\tau^{(i)}$ . Moreover, the estimator of the rescale approach above is  $\hat{G}_{Rescale}^{\pi} = \frac{1}{N}\sum_{i=1}^{N-1}\omega_{0:T-1}^{(i)}G_{0:T}^{\mathcal{H}(i)}$ , which is equivalent to the vanilla IS estimator [61, 72]. We now show the variance reduction property of  $\hat{G}_{PDIS}^{\pi}$  in the context of OPEHF.
Proposition 1. Assume that (i)  $\mathbb{E}[r_t^{\mathcal{H}}] \geq 0$ , and (ii) given the horizon  $T$ , consider any  $1 \leq t + 1 \leq k \leq T$  of any offline trajectory  $\tau$ ,  $\omega_{0:k}$  and  $r_t^{\mathcal{H}}\omega_{0:k}$  are positively correlated. Then,  $\mathbb{V}(\hat{G}_{PDIS}^{\pi}) \leq \mathbb{V}(\hat{G}_{Rescale}^{\pi})$ , with  $\mathbb{V}(\cdot)$  representing the variance.
The proof can be found in Appendix A. Assumption (i) can be easily satisfied in the real world, as HF signals are usually quantified as positive values, e.g., ratings (1-10) provided by participants. Assumption (ii) is most likely to be satisfied when the target policies do not visit low-return regions substantially [46], which is a pre-requisite for testing RL policies in human-involved environments as initial screening are usually required to filter the ones that could potentially pose risks to participants [57].
Besides IS, doubly robust (DR) [71, 34, 69, 12] and fitted Q-evaluation (FQE) [40] methods require learning value functions. Sparsity of rewards (following the rescale approach above) in the offline dataset may lead to poorly learned value functions [74], considering that the offline data in OPE is usually fixed (i.e., no new samples can be added), and are often generated by behavioral policies that are sub-optimal, which results in limited coverage of the state-action space. Limited availabilities of environment-policy interactions (e.g., clinical trials) further reduce the scale of the exploration and therefore limit the information that can be leveraged toward obtaining accurate value function approximations.
Reconstruction of IHRs. To address this challenge, our approach aims to project the end-of-episode human returns back to each environmental step, i.e., to learn a mapping  $f_{\theta}(\tau ,G_{0:T}^{\mathcal{H}}):(\mathcal{S}\times \mathcal{A})^{T}\times$
3We call it the rescale approach instead of vanilla IS as the idea behind also generalizes to non-IS methods.
Figure 1: (Left) Architecture of the variational latent model with human returns (VLM-H). (Mid) Illustration of the clustering behavior in the latent space using  $t$ -SNE visualization [73], where the encoded state-action pairs (output by the encoder of VLM-H) are in general clustered together if they are generated by policies with similar human returns (shown in the legend at the top left). (Right) Diagram summarizing the pipeline of the OPEHF framework.
$\mathbb{R}\to \mathbb{R}^T$  , parameterized by  $\theta$  , that maximizes the sum of log-likelihood of the estimated IHRs,  $[\hat{r}_0^{\mathcal{H}},\dots ,\hat{r}_{T - 1}^{\mathcal{H}}]^{\intercal}\sim f_{\theta}(\tau ,G_{0:T}^{\mathcal{H}})$  , following  $\max_{\theta}\frac{1}{N}\sum_{i = 0}^{N - 1}\sum_{t = 0}^{T - 1}\log p(\hat{r}_t^{\mathcal{H}} = r_t^{\mathcal{H}(i)}|\theta ,\tau^{(i)},G_{0:T}^{\mathcal{H}(i)})$  where  $G_{0:T}^{\mathcal{H}(i)}$  and  $r_t^{\mathcal{H}(i)}$  s are respectively the human return and IHRs (unknown) of the  $i$  -th trajectory in the offline dataset  $\rho^\beta$  , and  $N$  is the total number of trajectories in  $\rho^{\beta}$  . Given that the objective above is intractable due to unknown  $r_t^{\mathcal{H}(i)}$  's, we introduce a surrogate objective
$$
\max  _ {\theta} \frac {1}{N} \sum_ {i = 0} ^ {N - 1} \left[ \log p \left(\sum_ {t = 0} ^ {T - 1} \gamma^ {t} \hat {r} _ {t} ^ {\mathcal {H}} = G _ {0: T} ^ {\mathcal {H} (i)} \mid \theta , \tau^ {(i)}, G _ {0: T} ^ {\mathcal {H} (i)}\right) - C \cdot \mathcal {L} _ {r e g u} \left(\hat {r} _ {0: T - 1} ^ {\mathcal {H}} \mid \theta , \tau^ {(i)}, G _ {0: T} ^ {\mathcal {H} (i)}\right) \right]. \tag {1}
$$
Here, the first term is a necessary condition for  $\hat{r}_t^{\mathcal{H}}$ s to be valid for estimating  $r_t^{\mathcal{H}}$ s, as they should sum to  $G_{0:T}^{\mathcal{H}}$ . Since many solutions may exist if one only optimizes over the first term, the second term  $\mathcal{L}_{regu}$  serves as a regularization that imposes constraints on  $r_t^{\mathcal{H}}$ s to follow the properties specific to their corresponding state-action pairs; e.g.,  $(s,a)$  pairs that are similar to each other in a representation space, defined over the state-action visitation space, tend to yield similar immediate rewards [18].
The detailed regularization technique is introduced in sub-section below. Practically, we choose  $f_{\theta}$  to be a bi-directional long-short term memory (LSTM) [32], since the reconstruction of IHRs can leverage information from both previous and subsequent steps as provided in the offline trajectories.
2.3 Reconstruction of IHRs over Latent Representations (RILR) for OPEHF
Now, we introduce the regularization technique for the reconstruction of IHRs, i.e., reconstructing IHRs over latent representations (RILR). Specifically, we leverage the representations captured by variational auto-encoders (VAEs) [35], learned over  $\rho^{\beta}$ , to regularize the reconstructed IHRs,  $\hat{r}_t^{\mathcal{H}}$ .
VAEs have been adapted toward learning a compact latent space over offline state-action visitations, facilitating both offline policy optimization [42, 81, 65, 27, 26, 28] and OPE [18]. In this work, we specifically consider building on the variational latent model (VLM) proposed in [18] since it is originally proposed to facilitate OPE, as opposed to others that mainly use knowledge captured in the latent space to improve sample efficiency for policy optimization. Moreover, the VLM has shown to be effective for learning an expressive representation space, where the encoded state-action pairs are clustered well in the latent space, as measured by the difference over the returns of the policies from which the state-action pairs are sampled; see Figure 1 (mid) which uses  $t$ -SNE to visualize the encoded state-action pairs in trajectories collected from a visual Q&A environment (Appendix E).
Note that VLM originally does not account for HF signals (neither  $r_t^{\mathcal{H}}$ 's nor  $G_{0:T}^{\mathcal{H}}$ 's), so we introduce the variational latent model with human returns (VLM-H) below, building on the architecture introduced in [18]. VLM-H consists of a prior  $p(z)$  over the latent variables  $z \in \mathcal{Z} \subset \mathbb{R}^L$ , with  $\mathcal{Z}$  representing the latent space and  $L$  the dimension, along with a variational encoder  $q_{\psi}(z_t | z_{t-1}, a_{t-1}, s_t)$ , a decoder  $p_{\phi}(z_t, s_t, r_{t-1} | z_{t-1}, a_{t-1})$  for generating per-step transitions (over both state-action and latent space), and a separate decoder  $p_{\phi}(G_{0:T}^{\mathcal{H}} | z_T)$  for the reconstruction of the human returns at the end of each episode. Note that encoders and decoders are parameterized by  $\psi$  and  $\phi$  respectively. The overall architecture is illustrated in Figure 1 (left).
Trajectory inference (encoding). VLM-H's encoder approximates the intractable posterior  $p(z_{t}|z_{t - 1},a_{t - 1},s_{t}) = \frac{p(z_{t - 1},a_{t - 1},z_{t},s_{t})}{\int_{z_{t}\in\mathcal{Z}}p(z_{t - 1},a_{t - 1},z_{t},s_{t})dz_{t}}$ , by avoiding to integrate over the unknown latent space a priori. The inference (or encoding) process can be decomposed as, i.e.,  $q_{\psi}(z_{0:T}|s_{0:T},a_{0:T - 1}) = q_{\psi}(z_0|s_0)\prod_{t = 1}^{T}q_{\psi}(z_t|z_{t - 1},a_{t - 1},s_t)$ ; here,  $q_{\psi}(z_0|s_0)$  encodes initial states  $s_0$  into latent variables  $z_0$ , and  $q_{\psi}(z_t|z_{t - 1},a_{t - 1},s_t)$  captures all subsequent environmental transitions in the latent space over  $z_{t}$ 's. In general, both  $q_{\psi}$ 's are represented as diagonal Gaussian distributions<sup>4</sup> with mean and variance determined by neural network  $\psi$ , as in [18, 42, 27, 26, 28].
Trajectory generation (decoding). The generative (or decoding) process follows, i.e.,  $p_{\phi}(z_{1:T}, s_{0:T}, r_{0:T-1}, G_{0:T}^{\mathcal{H}}|z_0, \pi) = p_{\phi}(G_{0:T}^{\mathcal{H}}|z_T) \cdot \prod_{t=1}^{T} p_{\phi}(z_t|z_{t-1}, a_{t-1}) p_{\phi}(r_{t-1}|z_t) \cdot \prod_{t=0}^{T} p_{\phi}(s_t|z_t)$ ; here,  $p_{\phi}(z_t|z_{t-1}, a_{t-1})$  enforces the transition of latent variables  $z_t$  over time,  $p_{\phi}(s_t|z_t)$  and  $p_{\phi}(r_{t-1}|z_t)$  are used to sample the states and immediate environmental rewards, while  $p_{\phi}(G_{0:T}^{\mathcal{H}}|z_T)$  generates the human return issued at the end of each episode. Note that here we still use the VLM-H to capture environmental rewards, allowing the VLM-H to formulate a latent space that captures as much information about the dynamics underlying the environment as possible. All  $p_{\phi}$ 's are represented as diagonal Gaussians with parameters determined by network  $\phi$ .
To train  $\phi$  and  $\psi$ , one can maximize the evidence lower bound (ELBO) of the joint log-likelihood  $\log p_{\phi}(s_{0:T}, r_{0:T-1}, G_{0:T}^{\mathcal{H}}|\phi, \psi, \rho^\beta)$ , i.e.,
$$
\begin{array}{l} \max _ {\psi , \phi} \mathbb {E} _ {q _ {\psi}} \left[ \log p _ {\phi} (G _ {0: T} ^ {\mathcal {H}} | z _ {T}) + \sum_ {t = 0} ^ {T} \log p _ {\phi} (s _ {t} | z _ {t}) + \sum_ {t = 1} ^ {T} \log p _ {\phi} (r _ {t - 1} | z _ {t}) \right. \\ \left. - K L \left(q _ {\psi} \left(z _ {0} \mid s _ {0}\right) \| p \left(z _ {0}\right)\right) - \sum_ {t = 1} ^ {T} K L \left(q _ {\psi} \left(z _ {t} \mid z _ {t - 1}, a _ {t - 1}, s _ {t}\right) \| p _ {\phi} \left(z _ {t} \mid z _ {t - 1}, a _ {t - 1}\right)\right) \right]; \tag {2} \\ \end{array}
$$
the first three terms are the log-likelihoods of reconstructing the human return, states, and environmental rewards, and the two terms that follow are Kullback-Leibler (KL) divergence [38] regularizing the inferred posterior  $q_{\psi}$ . Derivation of the ELBO can be found in Appendix B. In practice, if  $\phi$  and  $\psi$  are chosen to be recurrent networks, one can also regularize the hidden states of  $\phi$ ,  $\psi$  by including the additional regularization term introduced in [18].
Regularizing the reconstruction of IHRs. Existing works have shown that the latent space not only facilitates the generation of synthetic trajectories but demonstrated that the latent encodings of state-action pairs form clusters, over some measures in the latent space [73], if they are rolled out from policies that lead to similar returns [42, 18]. As a result, we regularize  $\hat{r}_t^{\mathcal{H}}$  following
$$
\min  _ {\theta} \mathcal {L} _ {r e g u} \left(\hat {r} _ {t} ^ {\mathcal {H}} \mid \theta , \psi , s _ {0: t} ^ {(i)}, a _ {0: t - 1} ^ {(i)}, G _ {0: T} ^ {\mathcal {H} (i)}\right) = \sum_ {j \in \mathcal {J}} - \log p \left(\hat {r} _ {t} ^ {\mathcal {H}} = (1 - \gamma) G _ {0: T} ^ {\mathcal {H} (j)} \mid \theta , \psi , s _ {0: t ^ {\prime}} ^ {(j)}, a _ {0: t ^ {\prime} - 1} ^ {(j)}, G _ {0: T} ^ {\mathcal {H} (j)}\right) \tag {3}
$$
for each step  $t$ ; here,  $(s_{0:t}^{(i)}, a_{0:t-1}^{(i)}) \in \tau^{(i)} \sim \rho^{\beta}$ ,  $\mathcal{J} = \{j_0, \ldots, j_{K-1}\}$  are the indices of offline trajectories that correspond to the latent encodings  $\{z_{t'}^{(j_k)} \sim q_\psi(\cdot | s_{0:t'}^{(j_k)}, a_{0:t'-1}^{(j_k)}) | j_k \in \mathcal{J}, t' \in [0, T-1]\}$  that are  $K$ -neighbours of the latent encoding  $z_t^{(i)}$  pertaining to  $(s_{0:t}^{(i)}, a_{0:t-1}^{(i)})$ , defined over some similarity/distance function  $d(|\cdot|)$ , following, i.e.,
$$
\min  _ {j _ {k} \in \mathcal {J}} \sum_ {k = 0} ^ {K - 1} d \left(z _ {t} ^ {(i)} \mid \mid z _ {t ^ {\prime}} ^ {(j _ {k})}\right), \quad \text {s . t .} z _ {t ^ {\prime}} ^ {(j _ {k})}, \text {s c o r r e s p o n d i n g} \left(s _ {0: t ^ {\prime}} ^ {(j _ {k})}, a _ {0: t ^ {\prime} - 1} ^ {(j _ {k})}\right) \in \tau^ {(j _ {k})} \sim \rho^ {\beta}. \tag {4}
$$
In practice, we choose  $d(\cdot ||\cdot)$  to follow stochastic neighbor embedding (SNE) similarities [73], as it has been shown effective for capturing Euclidean distances in high-dimensional space [75].