Stable Offline Value Function Learning with Bisimulation-based Representations
Authors: Brahma S Pavse
Yudong Chen
Qiaomin Xie
Josiah P. Hanna
Venue: ICML 2025 poster (2025)
Keywords: reinforcement learning
representation learning
off-policy
offline policy evaluation
bisimulations
stability
value function learning
abstractions
Primary area: reinforcement_learning->batchoffline
TL;DR: We propose a bisimulation-based metric learning algorithm to learn state-action representations for accurate offline policy evaluation (OPE) and establish the stability properties of the learned representations theoretically and empirically.

Abstract: In reinforcement learning, offline value function learning is the procedure of using an offline dataset to estimate the expected discounted return from each state when taking actions according to a fixed target policy. The stability of this procedure, i.e., whether it converges to its fixed-point, critically depends on the representations of the state-action pairs. Poorly learned representations can make value function learning unstable, or even divergent. Therefore, it is critical to stabilize value function learning by explicitly shaping the state-action representations. Recently, the class of bisimulation-based algorithms have shown promise in shaping representations for control. However, it is still unclear if this class of methods can \emph{stabilize} value function learning. In this work, we investigate this question and answer it affirmatively. We introduce a bisimulation-based algorithm called kernel representations for offline policy evaluation (\textsc{krope}). \textsc{krope} uses a kernel to shape state-action representations such that state-action pairs that have similar immediate rewards and lead to similar next state-action pairs under the target policy also have similar representations. We show that \textsc{krope}: 1) learns stable representations and 2) leads to lower value error than baselines. Our analysis provides new theoretical insight into the stability properties of bisimulation-based methods and suggests that practitioners can use these methods to improve the stability and accuracy of offline evaluation of reinforcement learning agents.

Stable Offline Value Function Learning with Bisimulation-based Representations
Brahma S. Pavse' Yudong Chen
Qiaomin Xie Josiah P. Hanna
Abstract
(Sutton & Barto, 2018). While value function learning algo- rithms such as lemporal-differeee successful, they can be unreliable. In particular, the deadly triad, i.e., the combination of off- off-policy updates, function approximation, and bootstrapping, can make TD-based meth- ods diverge (Sutton & Barto, 2018: Tsitsiklis Van Roy, 1997; Baird, 1995; Hasselt et al., 2018). Function approx- imation is critical component of value function learning since it determines the representations of state-action pairs, which in turn defines the space of expressible value func- tions. Depending on how this value function space is rep- resented, value function learning algorithms may diverge (Ghosh & Bellemare, 2020). That is, the value function learning algorithm may not converge to its fixed-point, or may even diverge away from it. In this work, we explicitly learn state-action representations to improve the stability of offline value function learning.
learning (TD) have been
In reinforcement learning, offline value func- tion learning is the procedure of using an offline dataset to estimate the expected discounted re- turn from each state when taking actions accord- ing to fixed target policy. The stability of this procedure, i.e., whether it converges to its fixed- point, critically depends on the representations of the state-action pairs. Poorly learned represen- tations can make value function learning unsta- ble, or even divergent. Therefore, it is critical to stabilize value function learning by explicitly shaping the state-action representations. Recently, algorithms have
the class of bisimulation-bar shown promise in shaping representations for con- trol. However, it is still unclear if this class of methods can stabilize value function learning. In this work, we investigate this question and an- swer affirmatively. We introduce a bisimulation- based algorithm called kernel representations for offline policy evaluation (KROPE). KROPE uses kernel to shape acatinaction representations such that state-action pairs that have similar immediate rewards and lead to similar next state-action pairs under the target policy also have similar represen- tations. We show that KROPE: 1) learns stable rep- resentations and 2) leads to lower value error than baselines. Our analysis provides new theoretical insight into the stability properties of bisimulation- based methods and suggests that practitioners can use these methods to improve the stability and accuracy of offline evaluation of reinforcement learning agents.
denote d values
& d are low
is high
N'
S
BISIM
Figure I. The figure illustrates the native state-action represen- tations X' and π r-bisimulation representations Q, - x- bisimulation algorithms use distance function that captures differences between state-action pairs based on immediate rewards and differences of next state-action pairs under shape their representations. Ultimately, the goal of π r-bisimulation methods is to learn representations such that state-actions pairs with similar action-valus values under и also have similar representations. The func- tion d outputs low values within the blue (and orange) state-actions but high values between blue and orange state-actions. Therefore, the blue (and orange) state-actions have similar representations, but different representations between the distinct colors.
1. Introduction
Learning the value function of a policy is a critical com- ponent of many reinforcement learning (RI.) algorithms
Uniieersity of Wisconsin Madison. Correspondence to: Brahma Payse <paveee@wa wii..
In seeking such representations, we turn to r-bisimulation algorithms. These algorithms define a metric to capture behavioral similarity between state-action pairs such that similarity is based on immediate rewards received and the similarity of next state-action pairs visited by æ (Castro,
Proceedings of the 4nnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Stable Offline Value Function Learning with Bisimulation- based Representations
2020). The algorithms then use this metric to learn represen- tations such that state-action pairs that are similar under this metric have similar representations. Ultimately, the goal of T r-bisimmlattin methods is to learn representations such that state-actions pairs with similar values under æ also have sim- ilar representations (see Figure 1). While these algorithms have shown promise in improving the expected return of RL algorithms, it remains unclear whether they contribute to stability (Castro al., 2023; Zhang al., 2021; Castro et al., 2021). In this paper, we aim to understand whether the x- bisimulation-baser representations stabilize value function learning.
2. Background
In this section, we present our problem setup and discuss prior work.
2.1. Problem Setup and Notation
We consider the infinite- Markov decision pro- cess (MDP) framework (Puterman, 2014), M do), where S is the state space. A is the ac- tion space, r : S x A + [-1. is the deterministic reward function, P S x A A(S) is the transition dynamics function, 7 E 10. is the discount factor, and do E A(S) is the initial state distribution, where A(X) represents the set of all probability distributions over a set X. We refer to the joint state-action space as X == S x A. The agent acting according to policy x S + ((A) in the MDP generates trajectory: So, Ao. Ra, S1, A,, where So ~ do, A, ~ r(-[S,), R2 == r(S,; and S411 P(-|S,, At) for t 1 0.
(S,A, r. P. 7, 7,
In this work, we focus on offline value function learning. Given fixed, offline dataset generated by unknown and pos- sibly multiple behavior policies, the goal is estimate the value function of a fixed, target policy. Towards establishing the stability properties of r-bisimulation representations, we introduce kernel representations for offline policy evaluation (KROPE). KROPE defines a kernel that captures similarity between state-action pairs based on immediate rewards re- ceived and similarity of next state-action pairs under the target policy. It then shapes the state-action representations such that state-action pairs that are similar according to this kernel have similar representations. We use KROPE as the representative algorithm for the class of bisimulation- based representation learning algorithms to investigate the following question:
We define the action-value function of a policy π for a given state-action pair as q (s.a) a) == E, Er "',(S, A,) = 8, Ag = a). i.e., the expected discounted return when starting from state s with initial action a and then fol- lowing policy a. The Bellman evaluation operator T Rx R is defined as (T*))ss == r(8,a) + A')], Vf E IR. Accordingly, the action-value function satisfies the Bellman equation, i.e.. (a:)))((s.000
q = r(s. + I9("",
Can Iisimulation-basel representation learning stabilize offline value function learning?
It will be convenient to consider the matrix notation equiv- alents of the above functions. Since a policy π induces Markov chain on X, can denote the transition ma- trix of this Markov chain by pa E Rg x|x|, Here, each entry PT (i.j) is the probability of transitioning from state- actions to Similarly, we have the action-value function q"' E RR and reward vector r E æ! where the en- try qn (i) and r(i) are the expected discounted return from state- action under π and reward received at state-action i respectively.
Through theoretical and empirical analysis, we answer this question affirmatively and make the following contributions:
1. We introduce kernel representations for offline pol- icy evaluation (KROPE) for stable and accurate offline value function learning (Section 3).
2. We prove that KROPE'S representations stabilize least- squares policy evaluation (LSPE), popular value func- tion learning algorithm (Sections 3.2).
In this work, we study the representations of the state-action space. We use á : S x A + Rd to denote the state- action representations, which maps state-action pairs into a d-dimensional Euclidean space. We denote the matrix of all the state-action features as 4 g| where each row is the state-action feature d(s,a) a) IRd for state-action pair (s. When dealing with the offline dataset D, 4's dimensions are [D] x d, where (D) is the number state- actions in the dataset D. Note that & can be the native state-action features of the MDP, or the output of some rep- resentation learning algorithm, or the penultimate features of the action-valus function when using a neural network. Throughout this paper, we will view @ as an encoder or state-action abstraction (Li et al., 2006). Note that the state-
3. We prove that KROPE representations are Bellman com- plete, another indication of stability (Sections 3.3).
4. We empirically validate that KROPE representations im- prove the stability and accuracy of offline value func- tion learning algorithms on 10/13 11 offline datasets and against 7 baselines (Section 4).
5. We empirically analyze the sensitivity of the KROPE learning procedure under the deadly triad. These exper- iments shed light on when representation learning may be easier than value function learning (Section 4.4).
2

---

Stable Offline Value Function Learning with Bisimulation based Representations
action abstraction view enables to view di a state-action aggregator from the space of state-actions X' to the space of state-action groups X
two-stage approach of learning a linear function on top of fixed representations is called the linear evaluation protocol (Chang et al., 2022; Farebrother et al., 2023; 2024; Grill et al., 2020; He et al., 2020). This protocol enables us to cleanly analyze the nature of the learned representations within the context of cell-understood value function learn- ing algorithms such as LSPE. In Appendix A, we include the pseudo-code for LSPE.
2.2. Offline Policy Evaluation and Value Function Learning
In offline policy evaluation (OPE), the goal is to evaluate fixed target policy, Te, using fixed dataset of m transition tuples D == {(5s, as, s/, r=)] In this work, we evalu- ate Te by estimating the action-value function qd using D. Crucially, D may have been generated by a set of un- known behavior policies that are different from Te. which means that simply averaging the discounted returns in 'D will produce inconsistent estimate of qmr In our theoreti- cal results, we make the standard coverage assumption that Vs E S, Va E Tf (als) > 0, then the state-action pair (s, a) has non-zero probability of appearing in D (Sutton & Barto, 2018; Precup et al., 2000).
2.3. Stable Representations and Con-Sonsisy
We define stability of LSPE and related TD-methods follow- in (Ghosh & Bellemare, 2020):
Definition 1 (Stability). LSPE is said to be stable if for any initial 0g R4, lime-+oo 0g Oysen when Og is updated according to Equation (1).
When determining the stability of LSPE, we have the follow- ing proposition from prior work:
We measure the accuracy of the value function estimate with the mean squared value error (MSVE). Let qf= be the estimate returned by value function learning method using D. The MSVE of this estimate is defined as MSV == ((s.A)>D ((TT (S, A) -00= continuous state-action spaces, where it is impossible to compute g=' for all state-actions, we adopt a common eval- uation procedure from the OPE literature of measuring the MSE across only the initial state-action distribution, i.e.,
Proposition (Asadi al. 2024; Wang et al. 2021). LSPE is stable if and only if the spectral radius of ( p=) i.e. its maximum absolute eigenvalue, is less than 1.
xT=(S (S, A)) In environments with
Therefore, the stability of LSPE. largely depends on the rep- resentations & and the distribution shift between the data distribution of D and T.. In this work, we study the stability of LSPE. for a fixed distribution of D and learn $ If a given & stabilizes LSPE, we say & is stable representation.
M == As [(qT (So, Au) -qT (So. )1] For this procedure, we assume access to do (Voloshin et al., 2021; Fu et al., 2021). While in practice nTT is unknown, it is standard for the sake of empirical analysis to estimate -== by executing unbiased Monte Carlo rollouts of ne or com- puting 9Th exactly using dynamic programming in tabular environments (Voloshin et al., 2021; Fu et al., 2021).
In addition to stability, we also want the state-action fea- tures to be such that state-action features that are close in the representation space also have similar gT= values (Lyle et al., 2022; Lan et al., 2021). In this work, we call this property qT' *o*-connistency since the learned representations are consistent with the qT values.
Least-Squares Policy Evaluation Least-squares policy evaluation (LSPE) is a value function learning algorithm, which models the action-value function as a linear function: " (s, a)= == d(s, a) 8, where f E Rd (Nedic & Bertsekas, 2003). LSPE iteratively learns 0 with the following updates per iteration step t:
2.4. Related Works
In this section, we discuss the most relevant prior literature on OPE and representation learning.
Representations for Offline RI. and OPE. There are sev- eral works that have shown shaping representations can be effective for offline RL (Yang & Nachum, 2021; Islam et al., 2023; Nachum & Yang, 2024; Zang et al.. 2023a; Arora et al., 2020; Uehara et al., 2022; Chen & Jiang, 2019; Pavse & Hanna, 2023b). Ghosh & Bellemare (2020) presented theoretical understanding of how various representations can stabilize TD learning. However, they did not discuss bisimulation- based representations. Kumar et al. (2022); Ma et al. (2024); He et al. (2024) promote the stability of bD-based methods by increasing the rank of the represen- tations to prevent representation collapse. However, as we show in Section 4, these types of representations can still lead to inaccurate OPE. On the other hand, KROPE mit-
0++1 + (En[0T " "ED,,, [ ( (
(1)
where the expectations are taken with respect to the ran- domness of the dataset D and Te. Note that E[ T +) is the feature covariance matrix. Assuming LSPE converges, it will converge to the same fixed-point as TD(O) (Szepesvari, 2010), which we denote ÓLSPE- In this work, we follow two-stage approach to applying LSPP: we first obtain the encoder & either through representation learning or using the native features of the MDP, and then feed the obtained along with D and Te as input to LSPE, which outputs T (Nedic & Bertsekas, 2003; Chang et al., 2022). This
3

---

Stable Offline Value Function Learning with Bisimulation- based Representations
igates representation collapse and leads to accurate OPE. Chang et al. (2022) introduced BCRL to learn Bellman com- plete representations for stable OPE. While in theory, BC representations are desirable, we found that BCRL is sen- sitive to hyperparameter tuning. In contrast, we show that KROPE is more robust to hyperparameter tuning. Pavse & representa-
algorithm to learn them. We defer the proofs to Appendix B.
3.1. KROPE Kernel
Prior works typically define a distance metric to capture the notion of r-bisimilarity between two states: the distance between states is in terms of the immediate rewards received and differences of next states under π (see Figure 1) (Castro, 2020). In this work, we follow Castro et al. (2023) and take a perspective of kernels instead of distances. We first define a kernel that captures π- --bisimilarity AT : : x X - R, but for pairs of state-actions and under Te!
Hanna (2023a) showed that bisimulation-base tions mitigate the divergence of FQE; however, they did not provide explanation for divergence mitigation. Our work provides theoretical insight into the stability properties of bisimulation-basel algorithms.
Bisimulation- Representation Learning. Recently, there has been lot of interest in r-bisimulation algorithms for better generalization (Ferns et al., 2004; 2011; Ferns & Precup, 2014 4; Castro, 2020; Zang et al., 2023b). These algorithms measure similarity between two state-action pairs based on immediate rewards received and the sim- ilarity of next state-action pairs visited by т. These algo- rithms first define a distance metric that captures this x- bisimilarity, and then use this metric to learn representations such that π r-bisimular states have similar representations (Castro et al., 2021; Castro, 2020; Zhang et al., 2021; Cas- tro et al., 2023; Chen & Pan, 2022; Kemertas & Jepson, 2022). Castro (2020) introduced a r-bisimulation learning algorithm but assumed that the transition dynamics are deter- ministic. Zhang et al. (2021) introduced a bisimulation algo- rithm that allowed for Gaussian dynamics and demonstrated its effectiveness in discarding distracting features for control. Castro et al. (2021; 2023) introduced bisimulation-base based algorithms for control that allow for stochastic transition dy- namics. To the best of our knowledge, there is still a gap in the literature as to whether the general class of these learned representations stabilize offline value function learning. To address this gap, we introduce KROPE, an algorithm that leverages Castro et al. (2023)'s kernel perspective similarity metrics. We first explicitly define kernel that captures the relationship between state-action features in latent space in terms of the x a-bisimulation re- lation. This definition allows us to immediately establish stability since we can now analyze the spectral properties of offline value function learning algorithms. The proofs for KROPE'S basic theoretical properties (Section 3.1) follow those by Castro et al. (2023). Our related theoret- ical results (Sections 3.2 and 3.3) and empirical analysis (Section 4) are novel to this work.
h= (52,a1;82,000 02) = k1 ($1, 01: $2, 02)
+ (172 (-181, a1), p
(2)
where ha 01:81,01:8,,000 02)
1 (r(s),a)--((0a0
and ly(k* ( (1111, (PFr))))) p (1(s))2))) a2))
E ( 03 p (2) [k*T (s), d/ 8/2 d/3}-
Here, 4 measures short-term similarity based on rewards received and 12 measures long-term similarity between probability distributions by measuring similarity between samples of the distributions according to he (Castro et al., 2023). We refer to h as the KROPE kernel.
Remarks on kTx, We now discuss the trade-offs of k*t While km is different from prior work, it benefits and suffers from the same advantages and disadvantages discussed in Castro et al. (2023).
By definition 22 is a function of independently coupling (s of) and (s²2 ab) (Castro et al., 2021; 2023). The ad- vantage of this independence is that it enables a prac- tical algorithm and is flexible as it make no assump- tions on the nature of the transition dynamies of psc, While independent coupling does raise a complication with stochastic dynamics and/or stochastic policies since self- similarity may be lower than similarity between different state- actions, i.e., ATT (TT a1; $1, 01) < hT a1 82, 42) for some (5), a1), (2),a2) E X. he still induces a valid .- rr-biiimulation-aaaa based metric. To see its validity, con- sider the induced distance function by h:=e, i.e., Var, y E X, , dxRopn v) ;= k-T (I,z)+ (v. 2 (x, v). The metric dEROPE satisfies continuity (see Lemma 3 in Ap- pendix B.I), which is an important property for metrics to satisfy (Lan et al., 2021). Lemma 3 states that the abso- lute action-value difference between any two state-action pairs under मe upper-boundee by dxRorE plus an addi- tive constant, where the additive constant arises due the independent coupling. Thus far, prior work has removed the additive constant by assuming deterministic or Gaus- sian dynamics (Castro, 2020; Zhang et al., 2021). However, these assumptions are practically more restrictive. While one might expect the existence of the additive constant to hurt performance, prior work and our work (see Section 4)
r-bisimulation-ben
3. Kernel Representations for Offline Policy Evaluation
We now present our bisimulation- based representation learn- ing algorithm, kernel representations for OPE (KROPE). We present the desired KROPE kernel, prove stability properties of KROPE representations, and present a practical learning
4

---

Stable Offline Value Function Learning with Bisimulation- based Representations
show that independent coupling actually yields strong per- formance (Pavse& & Hanna, 2023a; Castro al., 2023). Fur- thermore, in Section 3.2 we show that this property does not affect the stability properties of the learned representations. An interesting future direction will be to develop practically feasible approaches to eliminate the additive constant.
representations instead of KROPE representations, which are H r-bisimulation-liik representations. We say function class F is Bellman complete if it is complete under the Bellman operator: f S Vf E F. For instance, suppose F is the class of linear functions spanned by 4, F == {f E R² f == Φw). w E IR Then if T f. Vf E F is also a linear function within the span of 4, we say & is a Bellman complete representation. Bellman completeness alternative condition for stability and is typically assumed to ensure to data-efficient policy evaluation (Wang et al., 2021; Szepesvári & Munos, 2005; Chang et al., 2022).
Since Lemma and the associated contraction, metric space completeness, and fixed-point uniqueness properties are similar to Castro et al. (2023)'s kernel, we defer these results to Appendix B.I.
3.2. Stability of KROPE Representations
Recall, that the induced distance function due to k is
In the previous section, we defined the KROPE kernel. UI- timately. however, we are interested in representations that satisfy the relationship in Equation (2). We modify Equation (2) accordingly by giving aa some functional form in terms of state-action representations. We do so with the dot product: (u, v) = uTu, U, Vu, E Rd, i.e., ke ($1, 01; $2, (2) (($), (522.22). With this setup, write Equation (2) in matrix notation and define the KROPE representations as follows:
Vz, y E X: dEROPE (2, y) == k (=,,(, (xxx
2k
We now present our second main result. It states that KROPE representations are Bellman complete under a reward func- tion injectivity assumption (see proof in Appendix B.2):
Theorem 2. Let é X X0 be the state- action abstraction induced by grouping state- actions 2, E X such that if dEROPE v) = 0, then x(x) = d(y), Vx, y E X. Then & is Bellman complete if the abstract reward function 1 : xx - (-1,1) is injective (i.e., distinct abstract rewards).
Definition 2 (KROPE Representations). Consider state- action representations & e xx that are embedded in Rd with keTe (81, (s1,a1; $2, a2) = ((s,,1)) ((2, 02). We say 4 is a KROPE representation if it satisfies the following:
Ep 20 1 = [ DDDD n. [1 (P")
(3)
Takeaway #1: Stability of Bisimulation- Representations
where each entry of K1 E R represents the short- term similarity, Kn, between every pair of state-actions, i.e., K₁ ($1, 01; 82; a2) == (r(s),a: (x)g.ag))
Under the theoretical assumptions made, KROPE representations avoid divergence of offline value function learning. They 1) induce non-expansive value function updates and 2) are Bellman complete.
I Finss Tam
Given this definition, we present our novel result proving the stability of KROPE representations:
Theorem 1. & is a KROPE representation as de- fined in Definition 2, then the spectral radius of (E[
3.4. KROPE Learning Algorithm
In this section, we present an algorithm that learns the KROPE representations from data. We include the pseudo- code of KROPE in Appendix A. The KROPE learning al- gorithm uses an encoder dw S x A - Rdd which is parameterized by weights w of a function approxima- tor. It then parameterizes the kernel with the dot product, i.e. (w a1; $2, a2) ;= (s,, T (52, a2) (see Equa- tion (3)). Finally, the algorithm then minimizes the follow- ing loss function, which is similar to how the value function
p) is less than 1. That is, & stabilizes LSPE.
By adopting the kernel perspective, Theorem 1, proved in Appendix B.2, tells us that r-bisimulationn-aased KROPE representations stabilize OPE with LSPE. Intuitively, they are stable since when ( radius is less than 1, each update to 0, in Equation (I) is non-expansive. That is, each update brings O2 closer to 0
x T -B)> p=* spectral
is learned in deep RL (Mnih et al., 2015):
LXROPE(W)
3.3. Connection to Bellman Completeness
ED ( (K0(01, - (s 01; 82, 01;82, o2))
)
27
In this section, we draw a novel connection between KROPE representations and Bellman completeness. Chen & Jiang (2019) proved similar result but focused on bisimulation
(4)
5

---

Stable Offline Value Function Learning with Bisimulation- Representations
where (51, 01: 82; 02)
(=($1-01) (rrpeaa))rrree)) ea)l
II) model-based, and IIII co-adaptation penalty-based. In class I, we consider 1) KROPE (ours), 2) deep bisimulation for control (DBC) (Zhang et al., 2021), 3) representations for OPE (ROPE) (Pavse & Hamna, 2023a; Castro et al., 2021); in class II, we consider bellman complete representation learn- ing (BCRL-EXP-NA) (Chang et al., 2022); and in class III, we consider 1) absolute DR3 regularizer (Kumar et al., 2022: Ma et al., 2024) and 2) BEER regularizer (He et al., 2024). In all cases, the penultimate layer features of FQE'S action- value encoder de are fed into LSPE for OPE. Note that since BCRI. was not designed as an auxiliary task (Chang et al., 2022), we evaluate it as a non-auxiliary (NA) task algorithm. We provide additional details on the baselines in Appendix C.
Fines Tamnl
T, (s a/; s/2, ab)], the state-action pairs are sampled from D, and w are weights of the target network that are periodically copied from w (Mnih et al., 2015). In this work, we use KROPE as an auxiliary task, which introduces only an auxiliary task weight as the additional hyperparameter. It is critical to note that since the learning algorithm a semi-gradient method, it may still diverge. Nevertheless, in Section 4 we show KROPE still improves the stability and accuracy of OPE compared to baselines.
4. Empirical Results
In this section, we present our empirical study designed to answer the following question: do KROPE representations lead to stable MSVE and low MSVE?
Domains We conduct our evaluation on variety of do- mains: 1) Garnet MDPs, which are class of tabular stochas- tic MDPs that are randomly generated given a fixed num- ber of states and actions (Archibald et al., 1995); 2) 4 DM Control enviromments: CartPoleSwingUp. CheetahRun, Fin- gerEasy, WalkerStand (Tassa et al., 2018); and 3) 9 D4RL datasets (Fu et al.. 2020; 2021). The first domain enables us to analyze the algorithms' performance across a wide range of stochastic tabular MDPs. The second and third set of do- mains test the algorithms in continuous higher-dimensional state-action environments.
4.1. Empirical Setup
In this section, we describe the main details of our empir- ical setup. For further details such as datasets, policies, hyperparameters, and evaluation protocol please refer to Appendix C.
Baselines Our primary representation learning baseline is fitted q-evaluation (FQE) (Le et al., 2019). FQE is the most fundamental deep RL OPE algorithm that learns representations of state-actions to predict the long-term performance of a policy. While FQE is typically used as an OPE algorithm, it can also be viewed as a value-predictive representation learning algorithm (Lehnert & Littman, 2020). More specifically, consider its loss function: "") (r(s,a) - ad)] - a))) where (s, a)= == (s, (T)ssa)TTorr T w and & = & as the neural network weights of an action-value neural network and w as the linear weights of the network applied on the output of the penultimate layer (s, a) of the neural network. Then minimizing this loss function shapes the representations de (s, a) to predict the expected future discounted return. As noted in Section 2, we follow the linear evaluation protocol where &e is shaped by different auxiliary tasks and is then used with LSPE for OPE since it helps us understand the properties of the representations within the context of a well-understood value function learning algorithm (Grill et al., 2020; Chang et al., 2022; Farebrother et al., 2024; Wang et al., 2021). We provide the pseudocode of this setup Appendix A. We also note that in Appendix C, we present results of performing OPE using FQE instead LSPE, and find that KROPE still reliably produces stable OPF estimates.
4.2. Analyzing Stability and C*=nSnnist Properties
In this set of experiments on the Garnet MDPs domain, we analyze the stability and gT consistency properties of the learned representations. We present the results in Figure 2. Our Garnet MDPs were generated with states and 5 actions, with a total of |.X|| = 40 state-actions, and each native (s, a) representation a 1-hot vector. In these experiments, the native representation is fed into linear encoder with bias component and no activation function. All algorithms are trained for 500 epochs and we report the results by evaluating the final learned representations for different latent dimensions
We view
Stability. Based on Theorem 1, a representation is stable if it induces a spectral radius of (E[$ 4))-1(>E x that is less than 1. In Figure 2(a), we present the fraction of runs that result in such representations. We find that up till d = 30, 100% of KROPE and BEER runs have spectral radius less than 1. We also find that BCRL-EXP-NA produces stable representations up till 40. At d = 50, all algorithms produce unstable representations. These results suggest that KROPF, BEER, and BCRI-EEP-NA are reliable in producing stable representations when projecting state-actions into low dimensions. When d 2 |x' the covariance matrix E/ T Φ} is more likely to be a singular matrix, which implies higher chance of instability.
p==
We consider the following three classes (and total 6) aux- iliary representation learning algorithms that are typically paired with FQE for stability: 1) iilimulation-mettt based,
6

---

Stable Offline Value Function Learning with Bisimulation- based Representations
learning curves).
1.0
1.00
FOE KAOPE
Roc of Ben Spegial Beegte
1.02
D63 AMS
REFR
0.8
In general, we find that KROPE representations lead to low and stable MSVE on 10/13 datasets. On the other hand, we find that the other auxiliary tasks inconsistently produce stable OPE estimates across all environments. Compared to the other bisimulation-bast baaed algorithms (DBC and ROPE), KROPE consistently achieves lower error. In all cases, DBC seems to be unreliable for OPE. ROPE, which was the previ- ous state-of-the bisimulation also achieves low error. While ROPE is competitive with KROPE, KROPE generally outperforms ROPE and has one less hyperparameter. The performance of the ABS-DR3 and BEER regularizer suggests that explicitly trying to increase the rank of the features of the penultimate layer may hurt sta- bility, and even if the OPE error stable, it can hurt accuracy. We also make similar observation for BCRL. However, in this case, we attribute poor performance to difficulty in opti- mizing the BCRL objective. In fact, in Figure 3(a), we will see that BCRL is sensitive to hyperparameter tuning. We also observe results consistent with a known result that BCRL- EXP-NA achieves lower error than BCRL--NA indicating the known result that exploration maximization of the covari- ance matrix helps produce stable representations (Chang et al., 2022). While FQE achieves low error on WalkerStand and CheetahRun, it is very unstable on the other datasets, which motivates the need to shape the representations for stable and accurate OPE. Note that in practice, KROPE may still diverge, as it did it in 3 cases, because the learning objective in Equation (4) is still a semi- coni-gradient learning algorithm (see discussion in Section 5 and Section 4.4).
1.00
0.96
0.4
sera
SCRL
KRore
0.96
(R3 485
0.94
Orthock
0.2
0.90
0.90
15
20
30
53
15
20
30
40
50
Gg dimension)
(p dimension!
(a) Spectral Radius
(b) Orthogonality
algorithm for OPE,
Figure 2. Evaluation of basic representation properties on Garnet MDPS with 40 kate-actions vs. output dimension d. Figure 2(a): Fractions of runs out of 30 trials that resulted in spectral radius of (E[$* ")) ()) to he less than 1; higher is hetter. Figure 2(b): Pearson correlation between orthogonality between pairs of latent features vs. their corresponding absolute q" action- value difference; higher is better. All results are averaged over 30 trials and the shaded region is the 95% confidence interval.
Con-iinsis,, We say that the representations have maintained qa c'==coonsst well when state-actions that have similar q=* values are close to each other in the representation space (Lan et al., 2021). We assess this property by first measuring the orthogonality: 1 (0)_))))))))0000))) (Wang et al., 2024) between every state- action pair, (1, and the absolute action-value difference: (st 01) - 9T (s2,a2)) [2)]. We then compute the Pearson correlation between these values for every pair and plot the correlation for each d in Figure 2(b). A correla- tion coefficient close to 1 indicates that the representations maintain consiiiityy well. We find that KROPF rep- resentations satisfy this property almost perfectly since it specifically tries to learn representations such that state- action pairs with similar values under Te are similar. We algorithms
[o(s)ar)] [2)))
11.00
11.00
Frec. of w/ MSVE <
Frat: of runs MSWE <
0.75
0.75
0.50
0.50
observe that the other cam-bisimulation-aaae typically have zero or even negative correlation. A nega- tive correlation indicates that state-actions with different action-values may have similar representations in latent space, which can complicate value function learning.
0.25
SCRL-CMPra
SCRLIS
062-485
0.25
BCAL-EX0-
FOE
DBC
0.00
0.00
(S)
MSVE
(a) CartPoleSwingUp
(b) WalkerStand
4.3. Offline Policy Evaluation
Figure 3. Hyperparameter sensitivity plots on CanPPleSSingUUp and WalkerStand. Results are over 20 trials for each hyperparame- ter configuration and shaded region is the 95% confidence interval. Larger area under the curve is better.
In this set of experiments, we evaluate the algorithms for OPE on 13 datasets: 4 DM control and 9 D4RL datasets. We also evaluate BCRL-NA, which is BCRL without the explo- ration maximization regularizer. To stabilize training for all algorithms, we use wide neural networks with layernorm (Gallici et al., 2025; Ota et al., 2024). Note that while wide networks and layernorm stabilize training, they may not lead to stable LSPE under the linear evaluation protocol. We report the final (normalized) squared value errors when the learned representations are used with LSPE for OPE. The results are presented in Table 1 (see Appendix C.3.1 for full
Hyperparameter Sensitivity. In OPE, hyperparameter tuning can be challenging since it may be infeasible to get access to ground truth performance of ne (Fu et al., 2021). Therefore, we prefer algorithms that are robust to hyperparameter tuning, i.e, they reliably produce accurate OPE estimates for a wide range of hyperparameters. In Fig- ure 3, we present the performance profile for each algorithm
7

---

Stable Offline Value Function Learning with Bisimulation-bus -aseed Representations
Algorithm
Dataset (DMC)
FQE
BCRL+EXP
BCRL
BEER
DR3
DBC
ROPE
KROPE (ours)
CariPoleSwingUp
Div.
2.0 + 1.6
2.2 0.8 Div.
0.9 0.0
Div.
0.2 0.1
0.0 + 0.0
CheetahRun
0.0 - 0.0
0.3 0.2
0.8
0.0 0.0
0.4
Div.
Div.
0.0 0.0
FingerEasy
Div.
0.6
0.8 0.2
Div.
0.0
Div.
0.100.0
0.6 - 0.0
WalkerStand
0.0 0.0
0.2
0.2 + 0.1 1.9 + 3.6
0.1 0.1
Div.
0.2 + 0.0
0.0 + 0.0
Dataset (D4RL.)
FQE
BCRL+EXP BCRL
BEER
DR3
DBC
ROPE
KROPE (ours)
cheetah random
0.9 + 0.0
Div.
Div.
0.9 + 0.0
0.9 + 0.0
0.9 + 0.0
1.0 0.0
1.0 + 0.0
cheetah medium
Div.
Div.
0.2 = 0.2
Div.
Div.
Div.
0.0 0.0
0.0 I 0.0
cheetah med-expert
Div.
0.2 0.1
0.3 = 0.1
Div.
Div.
Div.
0.1
0.0 - 0.0
hopper random
Div.
Div.
Div.
Div.
0.0
Div.
Div.
0.1 + 0.0
hopper medium
Div.
Div.
Div.
Div.
Div.
Div.
Div.
Div.
hopper med-expert
Div.
Div.
Div.
Div.
0.0
Div.
0.0 0.0
0.0 + 0.0
walker random
Div.
Div.
Div.
Div.
1.0
Div.
Div.
0.5 + 0.1
walker medium
Div.
Div.
Div.
Div.
Div.
Div.
Div.
Div.
walker med-expert
Div.
1.3
2.6 2.1
Div.
6.6 + 11.6
Div.
0.1 0.0
Div.
Table 1. Final normalized squared value error achieved by LSPE with different representation learning algorithms. Results are averaged over 20 trials and the variation is the 95% confidence interval. Lower and less erratic is better. Values rounded single place decimal. When algorithm's 's error is 2 10. we label it as diverged (Div.). Bolded and highlighted error indicates lowest error among baselines.
across all hyperparameter combinations and all trials (Agar- wal et al., 20.1). We tune the hyperparameters discussed in Appendix C.I. We find that 100% KROPE runs across all instances produce MSVE < 1, which is not the case with other algorithms. This result indicates the reliability us- ing KROPE for OPE over other algorithms including other bisimulation-baseel algorithms.
W3 2w3 is bad transition since updates made wg may move 2wy further away. When this transition is sampled at frequency that is different from the on-policy distribution, algorithms such as TD, LSPE, and FQE tend to diverge (Asadi et al., 2024). Similarly, for KROPE, which uses its weights to compute the latent representation of states, we would expect that pairs of transitions that lead to moving dot product targets are bad transition pairs.
Takeaway #2: Practical Offline Policy Evalua- tion with KROPE
To understand the stability of the learning procedures, design the following experiment when using KROPF as an auxiliary loss function to FQE. We first create an on-policy dataset D We then define off-policy datasets of the form D*, which consists of transitions starting from the speci- fied state 8, where = {w1, W2, W3, 2w3 }- Using these datasets, we then construct D Dun U Dwa and four D2 variations: 1) Di Dan U D 2) D
KROPE can increase the stability and accuracy of evaluation of offline RL agents.
4.4. Stability of the KROPE Learning Procedure
Den U D""2, Don U D2ws, There-
In this section, we compare the susceptibility of KROPE's and FQE's 's semi-gradiens learning to divergence. We refer the reader to Appendix C.1 for more details.
3) D Don D''ss, and 4) D2 fore, DT. D2 consist of mostly bad transitions while D DD consist of good transitions. Since KROPE sam- ples pairs of transitions, we investigate the consequences of pairing Dj with D2 variation. We pair each of the four D: datasets with D2 and train FQE+KROPE, where FQE is fed transition samples from D2 and KROPF is fed samples from D and the specific D2 variation.
We conduct our experiments on the Markov reward process (MRP) in Figure 4(a) (see Feng al. (2019)). The MRP consists of 4 non-terminal states, terminal state (the box), and only 1 action. The value function estimate is linear in the weights w = [W1, W2, wa]. so, starting clockwise from the left-most circle, the native features of the states are [1, O, 0), [0, 1, 0]. [0, 0, 2]. and [0, 0, 1]. In this setup, we say a transition is bad if the bootstrapping target a moving target for the current state. For example, the transition from
In Figure 4(b), we show the training loss of only KROPE to analyze when KROPE may diverge. We find that even though Dj consists of mostly bad transitions, KROPE sam-
8

---

Stable Offline Value Function Learning with Bisimulation- -aseed Representations
10a
W2
10
W2
W2
Training Loss
Training Loss
10a
2W3
10
2W3
Nofalring
W,
p-02
P-O.1
r-0
10a
10
Wh
r=0
0
109
109
200
400
200
400
D 0.9
Epochs
Epochs
(a) Divergence Counterexample
(b) KROPE MSE Loss
(c) FOE+KROPE MSE Loss
Figure 4. Figure 4(a): MRP counterexample designed to illustrate divergence; denotes the rewards and p denotes the probability of transition (Feng et al., 2019). Figures 4(b) and 4(c): KROPE training loss and KOE+ training loss epochs respectively when different datasets are paired with Dj results are averaged over 20 trials, shaded region is the 95% confidence interval, and lower is better.
ples good transitions from D or Dz",, its divergence is mitigated. However, if KROPE samples bad transitions from Ditss or D2'',, it diverges since the pairing of samples from (D", and (D1, D2 leads to KROPE chasing a mov- ing bootstrapped dot product target. Therefore, sampling pairs of bad transitions can make KROPE more likely to diverge, while sampling single bad transition (with good transition) can make it less likely to diverge.
ceptibility to divergence (Sutton & Barto, 2018). While we employed commonly-used techniques such as layer- norm and wide neural networks to mitigate divergence (Ota et al., 2024; Gallici et al., 2025), the consequences of a semi -raddent method may still exist. One potential rem- edy is to avoid semi-gradient learning. Feng et al. (2019) suggest to replace the fixed-point loss function of semi- gradient methods with an equivalent expression that lever- ages the Legendre-Fencchl transformation (Rockafellar & Wets, 1998). However, a drawback with this approach is that the new learning objective is a minimax procedure, which can be challenging to optimize in practice. In future work, we will explore the viability of this approach to design provably convergent version of KROPE.
In Figure 4(c), we plot the combined training loss of FOE+KROPE. As expected, when FQE uses only bad transi- tions from D,, it diverges (No Pairing). In fact, in all cases, the divergence is due to FQE even when the corresponding KROPE variation is not diverging (see D or Der in Fig- ure 4(b)). For the D or D2 variations, both FQE and KROPE diverge. Therefore, sampling single bad transitions makes FQE more likely to diverge (Asadi et al., 2024).
6. Conclusion
While we hand-designed these datasets, in general, we would expect that the probability of sampling pair of bad transitions less than that of sampling a single bad tran- sition'. These experiments show that while value function learning with FQE may diverge, representation learning with KROPE may not.
In this work, we tackled the problem of stabilizing offline value function learning reinforcement learning. We in- troduced bisimulation-basce representation learning algo- rithm, kernel representations for OPE (KROpE), that shapes the state-action representations to stabilize this procedure. Theoretically, we showed that KROPE representations are stable from two perspectives: 1) non-expansiveness, i.c.. they lead to value function learning updates that enable con- vergence to a fixed-point and 2) Bellman completeness, i.e., they satisfy a condition for data-efficient policy evaluation. Empirically, we showed that KROPE leads to more stable and accurate offline value function learning than baselines. We also demonstrated when representation learning with KROPE may be easier than value function learning with FQE. Our work showed that bisimulation-basee representation learning can improve the stability and accuracy of long-term performance evaluations of offline reinforcement learning agents.
Takeaway #3: KROPE Divergence
While FQE'S divergence due to sampling bad tran- sitions, KROPE's 's divergence is due to sampling bad pairs of transitions. Intuitively, since the probability of sampling a bad transition pair is less than that of sampling a single bad transition, KROPE training may be casier than FQE training.
5. Limitations and Future Work
In this section, we discuss limitations and future work. A shortcoming of KROPE's semi-gradient algorithm is its sus-
We expect this statement roughly hold true unless the data consists of mostly bad transitions; see Appendix C.1
9

---

Stable Offline Value Function Learning with Iisimulation-baser Representations
Impact Statement
Castro, P. S., Kastner, T., Panangaden, P., and Row- land, M. Mico: Improved representations via sampling-based state similarity for markov deci- sion processes. In Ranzato, M., Beygelzimer, A., Dauphin, Y., Liang, P., and Vaughan, J. W. (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 30113-30126. Curran Associates, Inc., neurips.
Our work is largely focused on studying fundamental RL research questions, and thus we do not see any immediate negative societal impacts.
Acknowledgments
2021. URL https:: cc cccppaperrpm / paaper__ilees/ppppr///////1 fd0 Ea06s8ee00eebblc2216666000006000 Paper.. pdf.
The authors thank Abhinav Narayan Harish, Adam Labiosa, Andrew Wang, Anshuman Senapati. David Abel, Kyle Domico, Ishan Durugkar, Pablo Samuel Castro, Prabhat Nagarajan, Tengyang Xie, and the anonymous reviewers ICLR and ICML for feedback on our work, and thank Fanghui Liu for inspiring discussions on kernel methods. Y. Chen is partially supported by NSF CCF-2233152. Q. Xie is supported in part by National Science Foundation (NSF) grants CNS-1955997, EPCN-2339794 and EPCN-2432546.
6c Paper.
Castro, P. S., Kastner, T., Panangaden, P., and Rowland, M. A kernel perspective on behavioural metrics for markov decision processes. Transactions on Machine Learning Research, 2023. ISSN 2835-8856. URL https:/
openreview net. forumm pert Certification.
od -dd Ex-
J. Hamna acknowledges support from NSF (IIS (IS-2410981), American Family Insurance through research partnership Data Science Institute, and Sandia National Labs through University Partnership Award.
Chang, J., Wang, K., Kallus, N., and Sun, W. Learning Bell- man Complete Representations for Offline Policy Evalu- ation. In Proceedings of the 39th International Confer- ence on Machine Learning, pp. 2938 2971. PMLR, June 2022. URL //t:://rroooo http://oooedddceeoee ings mlr ppress/ v1 /262/6o chang22b. html. ISSN: 2640-3498.
with the University of Wisconsin
References
Agarwal, R., Schwarzer, M., Castro, P. S., Courville, A. C.,, and Bellemare, M. Deep reinforcement learning at the edge of the statistical precipice. Advances in Neural Information Processing Systems, 34, 2021.
Chen, J. and Jiang, N. Information-tthoorticc considerations in batch reinforcement learning. In International Conference Machine Learning. 2019. URL https: //api. sereiissematiiio www.annticccola....ommmm org /atp///rrpsspp 141460093.
Archibald, T. W., McKinnon, K. I. M., and Thomas, L. C. On the generation of markov decision processes. Jour- nal of the Operational Research Society, 46:354-361, 1995. URL https: //ap:p::///...../ https///ap.....aassaeeeeeeeeeeee .emanticschollll. CccCCCoop::::555::::
Chen, J. and Pan, S. Learning representations via robust behavioral metric for deep reinforcement learning. In Koyejo, S., Mohamed, S., Agarwal, A., Belgrave, D., Cho, K., and Oh, A. (eds.), Advances in Neural Informa- tion Processing Systems, volume 35, pp. 36654-36666. Curran Associates, Inc., 2022.
Arora, S., Du, S. S., Kakade, S., Luo, Y., and Saunshi, N. Provable representation learning for imitation learning via bi-level optimization. In Proceedings of the 37th In- lernational Conference on Machine Learning, 1CML JMLR.org, 2020.
Farebrother, J., Greaves, J., Agarwal, R., Lan, C. L., Goroshin, R., Castro, S., and Bellemare, M. G. Proto- value networks: Scaling representation learning with auxiliary tasks. In The Eleventh International Confer- ence Learning Representations, 2023. URL https: / //openrevien ew net. / forum? (aaerrummi==oooo.c.c doooooottottttr..
Asadi, K., Sabach, S., Liu, Y., Gottesman, O., and Fakoor, R. Td convergence: an optimization perspective. Pro- ceedings of the 37th International Conference on Neural Information Processing Systems, NIPS '23, Red Hook, NY, USA, 2024. Curran Associates Inc.
Farebrother, J., Orbay, J., Vuong, Q., Taiga, A. A., Chebotar, Y., Xiao, T., Irpan, A., Levine, S., Castro, P. S., Faust, A., Kumar, A., and Agarwal, R. Stop regressing: Train- ing value functions via classification for scalable deep RI. In Forty-firs International Conference on Machine Learning, 2024, URL https:ttp//tt openreview. .nt/ forum div.dddddd/////F....
Baird, L. C. Residual algorithms: Reinforcement learning with function approximation. In International Confer- ence on Machine Learning. 1995. URL https: app.. semant wew.ttiascolaa..o.oooooooo org wwwo//o//ooeeeeeeeoooo Corpus ID: 621595.
Castro, P. S. Scalable methods for computing state similarity in deterministic markov decision processes. In Proceed- ings of the Chirty-Fourdh AAA! Conference on Artificial Intelligence (AAAI-20), 2020.
Feng, Y., Li, L., and Liu, Q. A kernel loss for solving the bellman equation. In Wallach, H., Larochelle, H., Beygelzimer, A., d'Alché-Buc, E., Fox, E., and
10

---

Stable Offline Value Function Learning with Bisimulation-basee Representations
Garnett, R. (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc.,
Gretton, A., Borgwardt, K. M., Rasch, M. J., Schölkopf, B., and Smola, A. A kernel two-sample test. Jour- nal of Machine Learning Research, 13(25):723-73, 2012. URL http:// htpp://jmr...// httpmmt......aaaaeeeeeeeeeeeeeee www.ooggpgpp......... gretton12a htmmml.
2019. URL https://prcceed.ng..eeeee cc cc/paper_fiie /ttp//////aapp////////////e///// (tl0000oooccffcooceeeeeeeaeeeee0 pdf.
Grill, J.-B., Strub, F,, Altché, F., Tallec, C., Richemond, P. H., Buchatskaya, E., Doersch, C., Pires, B. A., Guo, Z. D., Azar, M. G., Piot, B., Kavukcuoglu, K., Munos, R., and Valko. M. Bootstrap your own latent a new approach to self-superviied learning. In Proceedings of the 34th International Conference on Neural Information Processing Systems, NIPS '20. Red Hook, NY, USA, 2020. Curran Associates Inc. ISBN 9781713829546
Ferns, N. and Precup. D. Bisimulation metrics are optimal value functions. In Proceedings of the Thirtieth Confer- ence on Uncertainty in Artificial Intelligence, UAI'14, 14, pp- 210-219, Arlington, Virginia, USA, 2014. AUAI Press. ISBN 9780974903910.
Ferns, N., Panangaden, P., and Precup, D. Metrics for finite markov decision processes. In Proceedings of the 20th Conference on Uncertainty in Artificial Intelligence, UAI '04, pp. 162-169, Arlington, Virginia, USA, 2004, AUAI Press. ISBN 0974903906.
Haamoja, T., Zhou, A., Abbeel, P., and Levine, Soft actor- critic: Off- maximum entropy deep reinforcement learning with stochastic actor. In Dy. J. and Krause, A. (eds.), Proceedings of the 35th International Confer- ence on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 1861-1870. PMLR, 10-15 Jul 2018. URL htthts::: http://pooeceiggeeeeeeeee mllr. press/ /v/0///88///// haarnojall.......... html.
Ferns, N., Panangaden, P., and Precup, D. Bisimulation metrics for continuous markov decision processes. SIAM Journal on Computing, 4066):1662-1714, 2011. doi: 10.1137/10080484X. URL https: //dot org/ 1137/10080484X,
Hasselt, H. V., Doron, Y., Strub, F., Hessel, M., Son- nerat, N., and Modayil, J. Deep reinforcement learn- ing and the deadly triad. ArXiv, abs/1812.02248, 2018. URL https: //tps://apii /api. semanticscholar: org / CorpooII::54445555
Fu, J.. Kumar, A., Nachum, O., Tucker, G., and Levine, S. D4rl: Datasets for deep data-driven reinforcement learning, 2020.
Fu, J., Norouzi, M., Nachum, O., Tucker, G., ziyu wang, Novikov, A., Yang, M., Zhang, M. R., Chen, Y., Kumar, A., Paduraru, C., Levine, S., and Paine, T. Benchmarks for deep off- off-policy evaluation. In International Confer- ence on Learning Representations, 2021. URL https: //openreview net wwwt/foruueeeeeee forum? idd
He, K. Fan, H., Wu, Y.. Xie, S., and Girshick, R. Mo- mentum contrast for unsupervised visual representation learning. In 2020 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 2020. doi: 10. 1.11099CVPR422202220000000
Fujimoto, S., Meger, D., Precup. D., Nachum, O., and Gu, S. S. Why should trust you, bellman? The Bellman error is poor replacement for value error. In Chaudhuri, K., Jegelka, S., Song, L., Szepesvari, C., Niu, G., and Sabato, S. (eds.), Proceedings of the 39th International Confer- ence on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 6918 09-46.33.. PMLR, 17-23 Jul 2022. URL httptttts///pcceeeeeeeeeee mmm..eeedinnnm.................. mlr. press/vle 62 ottp://1.../..ioooooo2....a..a.. o2w.. html.
He, Q., Zhou, T., Fang. M., and Maghsudi, S. Adaptive reg- ularization of representation rank as an implicit constraint of bellman equation. In The Twelfth International Confer- ence on Learning Representations, 2024. URL https: / //openreviee ew net/ forum? ii d=apXtol xDaJ.
Islam, R., Tomar, M., Lamb, A., Efroni, Y., Zang, H., Di- dolkar, A., Misra, D.,, Li, X., Van Seijen, H., Des Combes, R. T., and Langford, J. Principled offline rl in the pres- ence of rich exogenous information. In Proceedings of the 40th International Conference on Machine Learning, ICML'23, JMLR.org, 2023.
Gallici, M., Fellows, M., Ellis, B., Pou, B., Masmitja, I,, Foerster, J. N., and Martin, M. Simplifying deep temporal difference learning. In The Thirteenth In- ternational Conference on Learning Representations, 2025. URL httpt:tt///// fictps://operreve/ wwt.ee/wwweeeeeeeeeeeeeeeeeeeeee forum? id=7IzeLOkflu.
Kemertas, M. and Hamennadd-AAmstroon T. Towards robust bisimulation metric learning. In Ranzato, M., Beygelzimer, A., Dauphin, Y. N., Liang, P., and Vaughan, J. W. (eds.), Advances in Neural Infor- mation Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual, pp. 4764 4777, 2021. URL. https: http://pooccedeeee
Ghosh, D. and Bellemare, M. G. Representations for stable off-policy reinforcement learning. In Proceedings of the 37th International Conference on Machine Learning, ICML'20. Jor,, 2020.

---

Stable Offline Value Function Learning with Bisimulation-basee Representations
neurips. seeripgs.cc/paprr//22////ss 355fffe69226522deeddddd/do html.
Mnih, V., Kavukcuoglu, K., Silver, D., Rusu, A. A., Ve-
ness, J., Bellemare, M. G., Graves, A., Riedmiller, M., Fidjeland, A. K., Ostrovski, G., Petersen, S., Beattie, C., Sadik, A., Antonoglou, I., King, H., Kumaran, D., Wierstra, D., Legg, S., and Hassabis, D. Human-level control through deep reinforcement learning. Nature, 518 (7540):529-533, February 2015. ISSN 00280836. URL http://dx doi org/ 10. 1038 httpuunuuuuueeeeooeeee
Kemertas, M. and Jepson, A. D. Approximate pol- icy iteration with bisimulation metrics. Transactions on Machine Learning Research, 2022. ISSN 2835- 8856. URL
http://p...viie..eeeerr.rr.. id=Ii7UeHcCmO.
Nachum, O. and Yang, M. Provable representation learn- ing for imitation with contrastive fourier features. In Proceedings of the 35th International Conference on Neural Information Processing Systems, NIPS '21, Red Hook, NY, USA, 2024. Curran Associates Inc. ISBN 9781713845392
Kumar, A., Agarwal, R.. Ma, T., Courville, A. C., Tucker, G., and Levine, Dr3: Value-based deep reinforcement learning requires explicit regular- ization. In ICLR. OpenReview.net, 2022. URL
http:://b/l http:///bbp..ni-triee. uni uni-trier detde/d/d/bbb///eeeeee/ conf/ iclr/iclr2022 html #KumarAOCTL22
Nedic, A. and Bertsekas, D. Least squares policy evaluation algorithms with linear function approximation. Discrete Event Dynamic Systems, 13:79-110, 01 2003. doi: 10.
Lan, C. L.,, Bellemare, M. G., and Castro, P. S. Met- rics and continuity in reinforcement learning. In AAAI Conference on Artificial Intelligence, 2021. URL https://aph http://app..semnttssssctee www.nticscholar..o... org/ CorpusID:231749500
1023/A:1022112003333
Ota, K., Jha, D. K., and Kanezaki, A. A framework for training larger networks for deep reinforcement learning, Mach. Learn., 113(9):111--1199 June 2024. ISSN 0885- 6125, doi: 10. / //doi. org/ 10. 10t.//0099/111-10000 024-06477----00-00
Le. H. M., Voloshin, C., and Yue, Y. Batch Pol- icy Learning under Constraints. In International Conference on Machine Learning (ICML). arXiv, March 2019. URL http:// /tap:///axvv....//// www. / 1903. 08.0087733 738. 113.v:1903.88738 [cs, math, stat].
URL https:
Paulsen, V. and Raghupathi, M. An Introduction to the Theory of Reproducing Kernel Hilbert Spaces. Cambridge Studies in Advanced Mathematics. Cambridge University Press, 2016
Lehnert, L. and Littman, M. L. Successor features combine elements of model-free and model-based re- inforcement learning. Journal of Machine Learning Research, 21(196):1 2020. URL http://// fmmrr. org/papers/ wwtptr/////1/11 v21/19--111111 060.000..ee html.
Payse, B. S. and Hanna, J. State-Action Similarity- Basee Representations for Off-Policy Evaluation. In Proceed- ings of Advances in Neural Information Processing Sys- tems (NeurIPS), December 2023a.
Li, L., Walsh, T. J., and Littman, M. L. To- wards a unified theory of state abstraction for mdps. In International Symposium on Artifi- cial Intelligence and Mathematics, Al&Math 2006, Fort Lauderdale, Florida, USA, January 4-6, 2006, 2006. URL /tttt:///nyt.tt cs, www.uuasse............. edd/ /mmatto0/prrcceedinng//221....
Pavse, S. and Hanna, J. P. Scaling Marginalized Impor- tance Sampling to High-Dimensional State-Spaces via State Abstraction. In Proceedings of the AAAI Conference Artificial Intelligence, February 2023b.
Precup. D., Sutton, R. S., and Singh, S. P. Eligibility traces for off off-policy policy evaluation. In Proceedings of the Sev- enteenth International Conference on Machine Learning, ICML '00, pp. 759-766, San Francisco, CA, USA, 2000, Morgan Kaufmann Publishers Inc. ISBN 1558607072.
Lyle, C., Rowland, M., Dabney, W., Kwiatkowska, M., and Gal, Y. Learning dynamics and generalization in deep reinforcement learning. In Chaudhuri, K., Jegelka, S., Song, L., Szepesvari, C., Niu, G., and Sabato, S. (cds.), Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 14560-14581. PMLR,
Puterman, M. L. Markov decision processes: discrete
stochastic dynamic programming. John Wiley & Sons, 2014.
17-23 Jul 2022. URL https: httpp//tocceeg.....oo// mlr. press v 62 htt:/yylee2aa html.
Rockafellar, R. and Wets, R. J.B.. Variational Analysis. Springer Verlag, Heidelberg, Berlin, New York, 1998.
Ma, Y., Tang, H., Li, D., and Meng, Z. Reining generaliza- tion in offline reinforcement learning via representation distinction. In Proceedings of the 37th International Con- ference on Neural Information Processing Systems, NIPS 233 Red Hook, NY, USA, 2024. Curran Associates Inc.
Sutton, R. S. and Barto, A. G. Reinforcement Learn- ing: An Introduction. The MIT Press, second edition, 2018. URL http:/ //incompleteeeeee eompleteideas.... n book took/the-bool 2nd. html.
12

---

Stable Offline Value Function Learning with Bisimulation-base based Representations
Szepesvari, C. Algorithms for Reinforcement Learning. Mor- gan and Claypool Publishers, 2010. ISBN 1608454924.
Yang. M. and Nachum, O. Representation Matters: Offine Pretraining for Sequential Decision Making. In Proceedings of the 38th International Conference on Machine Learning, pp- 11784-11794, PMLR, July 2021. URI.tttps: wtwppprreeinn........... mw....mm........... htt.iga.........r.....m v139 1139/yang2hh.. html. ISSN: 2640-3498.
Szepesvári, C. and Munos, R. Finite time bounds for sam- pling based fitted value iteration. In Proceedings of the 22nd International Conference on Machine Learning, ICML '05, pp. 880-887, New York, NY, USA, 2005. As- sociation for Computing Machinery. ISBN 1595931805. doi: 10. 145/1 102351, 102355 102462. URL 1102462.
Zang, H., Li, X., Yu, J., Liu, C., Islam, R., des Combes, R. T., and Laroche, R. Behavior prior representation learn- ing for offline reinforcement learning. In The Eleventh International Conference on Learning Representations, 2023a. URL https: http////peeeeeeeeeeeeeeeeeeeeeee net/ forum? id=h04K9Bf4G2B.
org/10, 1tt.:1111111222
Tang. Y., Guo, Z. D., Richemond, H., Pires, B. A., Chan- dak, Y., Munos, R., Rowland, M., Azar, M. G., Lan, C. L., Lyle, C., György, A., Thakoor, S., Dabney, W., Piot, B.. Calandriello, D., and Valko, M. Understand- ing self-predictive learning for reinforcement learning. In Proceedings of the 40th International Conference on Machine Learning, ICMI:22. JMLR.org, 2023.
Zang, H., Li, X., Zhang, L., Liu, Y., Sun, B., Islam, R., des Combes, R. T., and Laroche, R. Understanding and addressing the pitfalls of bisimulation-basee representa- tions in offline reinforcement learning. In Thirty-sevente Conference Neural Information Processing Systems, 2023b. URL https: /owpnrvveewewweeeeee/eeeeeroe.we net forum? i id=sQyRQjun46.
Tassa, Y., Doron, Y., Muldal, A., Erez, T., Li, Y., de Las Casas, D., Budden, D., Abdolmaleki, A., Merel, J., Lefrancq, A., Lillicrap, T. P., and Riedmiller, M. A. 2018. Deepmind control suite. CoRR, 10000 00690.
Zhang, A., McAllister, R. T., Calandra, R., Gal, Y., and Levine, S. Learning invariant representations for re- inforcement learning without reconstruction. In In- ternational Conference on Learning Representations, 2021. URL https: /ttp:////peeeeeeeeeeeeeee net forum? 1 d== 2FCWDKREE..
URL http: wwwaaaavvvvv...... org/ abs
Tsitsiklis, J. and Van Roy, B. An analysis of temporal- difference learning with function approximation. IEEE Transactions on Automatic Control, 42(5):674-690, 1997. doi: 10. 58.84444
Uchara, M., Zhang, X., and Sun, W. Representation learn- ing for online and offline RL in low-rank MDPs. In International Conference Learning Representations, 2022. URL www.t:tt///c/..... /ttp::///weeeeeeeeee.oooo...o... n forum? id-J4iSIR9EhYO,
Voloshin, C., Le, H. M., Jiang, N., and Yue, Y. Empiri- cal study of off-policy policy evaluation for reinforce- ment learning. In Thirty-fifth Conference on Neural In- formation Processing Systems Datasets and Benchmarks Track (Round 1), 2021. URL hwtpthttp///// openreview. net/ forum? deaaa=KKKKK-----
Wang, H., Miahi, E., White, M. Machado. M. C., Abbas, Z., Kumaraswamy, R., Liu, V., and White, A. Investigating the properties of neural network representations in reinforcement learning. Artificial Intelligence, 330: 104100, 2024. ISSN 0004-3702. doi: http://dddiong///o....oo///antm. com/
URL
Wang, R., Wu, Y., Salakhutdinov, R., and Kakade, S. Insta- bilities of Offline RI with Pre-Trained Neural Representa- tion. In Proceedings of the 38th International Conference on Machine Learning, pp. 10948-10964 PMLR, July 2021. URL https: htt.attpee//roceaiiiteeoo http///eeeiggg...c//e/ppp www.nnmms......mmmm..m.........m press/
v1 1.39/ 9//wagg212 html. ISSN: 2640-3498.
13

---

Stable Offline Value Function Learning with Bisimulation- based Representations
A. Background
In this section, we present the theoretical background.
A.1. Bisimulation Metrics
In this section, we present background information on bisimulations and its associated metrics. Our proposed representation learning algorithm is a bisimulation- algorithm. Bisimulation abstractions are those under which two states with identical reward functions and that lead to identical groups of next states under any action are classified as similar (Ferns al.. 2004; 2011: Ferns Precup. 2014). Bisimulations are the strictest forms of abstractions. In practice, the exact bisimulation criterion is difficult to satisfy computationally and statistically. A more relaxed version of this notion is the r r-bisimmlatiin metrics. These metrics capture the similarity between two states such that two states are considered similar if they have identical expected reward functions and expected transitions to identical groups of next states under fixed policy t (Castro, 2020).
first give the definition of bisimulation.
Definition 3. (Li et al., 2006) An abstraction & : S + 8 over the state space S is a bisimulation for any action a and any abstract state d E S ø is such that for any two states $1, 82 E X, ((s1) = o(s2) implies that r(s1, a) = r(52, a) and) Ese P(s'|s), a) En' P(s' P(s||s2, $2
Below we define x- æ-bisimulations for state-actions instead of states:
Definition& 4. (Castro. 2020) An abstraction o X' + X over the state-action space X is a π T-bisimulation for fixed policy π for any two state-actions I, E X and abstract state-action xº E X2, & is such that (xxx) = <(y) implies that r(x) r(y) and Ev'E* PT (x'1z) Ex' P" (T(('ly)
The above definitions are based on exact groupings between state-actions. This strictness motivates the use of bisimulation and bibiiuullttion metrics, which we define below.
Theorem 3. (Ferns et al., 2004) Let M(S) be the space of bounded pseudometrics on the state-space S. Then define B: M(S + M(S) such that for each d E (M(S):
B(d)(s1. 82) max((r(s1,a) - + P(1ss,))) aEA
where W is Wasserstein distance between the two distributions under metric Then B has a unique fixed point, d', and d' bisimulation metric.
Similarly, we have the r-bisimulation metric:
Theorem (Castro, 2020) Let M(X) be the space of bounded pseudometrics on the state-action space X and a be a fixed
policy. Then define B M(x) > M(X) such that for each E MMS):
((d)(a, v) = + +W(()) P (-ly)),
where W is Wasserstein distance between the two distributions under metric d. Then B has a unique fixed point, d', and d' is a r-bisimulation metric.
Using the above metrics, prior works have introduced several representation learning algorithms to learn representations such that the distance between representations in latent space model the above distance metrics (Castro et al., 2021; ; 2023; Zhang et al., 2021; 1; Kemertas & Aumentado- Aumentado-Ammstrongg 2021; Pavse & Hanna, 2023a).
A.2. Reproducing Kernel Hilbert Spaces
Let X' be finite set and define a function & X' X' + R to be a positive semidefinite kernel if it is symmetric and positive semidefinite. We then have for any (I1, I2, Fn E X' and {ci, c,,... Cn E R:
Σ
Co (,, 2,) 22,20
14

---

Stable Offline Value Function Learning with Bisimulation- based Representations
Note that if the above inequality is strictly greater than zero whenever (cI...., Co has at least one nonzero, say the kernel is positive definite. Given kernel le on A' with the reproducing property, we can construct space of functions Ha referred to as reproducing kernel Hilbert space (RKHS with the following steps:
1. Construct a vector space of real-valued functions on A' of the form {k(a,,) )) 2}.
2. Equip this space with an inner product given by (k(z,) -), k(y, MH k(x,g).
3. Take the completion of the vector space with respect to the above inner product.
Our resulting vector space H4 is then an RKHS.
is often convenient to write (xx) == k( -) E He. which called the feature map and is an embedding of x in H4. One can also embed probability distributions into Ha. That is, Φ ((t)) - H.. which maps probability distributions over X to Hy. We define (p) Ex-=pl< (x))),, which is the mean embedding in H4 under p.
Given these embeddings in the Hilbert space, quantify the distances between elements in X' and P(x) in terms of the embeddings.
Definition 5. Given positive semidefintte kernel k; define Ph. as its induced distance:
Pk == [b(()) (x)l
By expanding the inner product, the squared distance can be written in terms of k:
p v) = k(,, x) + k(y.v) v) - 2k(2, v).
Similarly, we have distances on P(X) using 4
Definition 6. (Gretton et al., 2012) Let k be a kernel on X and & P(X) - Hy be as defined above. Then the Maximum Mean Discrepancy (MMD) is pseudo metric on (X)) defined by:
MMD(K))( = =10())) (T)
The core usage of the RKHS is to precisely characterize the nature of the KROPE kernel. In practice, we deal with neural network representations, which embedded in Euclidean space. Therefore, our goal is to learn representations in Euclidean space that approximate the properties of representations in the RKHS. For more details on the RKHS, we refer readers to (Castro et al., 2023) and (Gretton et al., 2012).
A.3. Algorithm Pseudocode
In this section, we present the pseudocode for LSPE and for our FQE auxiliary task with LSPE for OPE setup.
Algorithm LSPE
Input: policy evaluate me, batch D, fixed encoder function 6 A gd
2: Initialize &g E Rdd randomly.
3: Apply ó to D generate T.
4: for t= 0, 1, 2, ... T do
5: 0:41 +- ( -E T (r + PT 40.)]
6: end for
7: Return 0r
15

---

Stable Offline Value Function Learning with Bisimulation- based Representations
Algorithm 2 FQE + representation learning auxiliary task with LSPE for OPE
1: Input: policy to evaluate Ne. batch D, encoder parameters class Q, encoder function ø : x A RR action-value linear function Rd + R. E [0,
2: for epoch = 2, 3, Tdo
3: ((w) = (ax-Task D,Te) + (1 - (aa,)))) (r(s, a) + + a'))) q (de (s. (s))a))
{where the penultimate features are fed into an auxiliary representation learning algorithm such as KROPE,
DR3, BEER etc.}
4: We = arg minwes ((w)
5: Periodically run LSPE, 0 == LSPE(TTe) D, D.).)
6: Compute estimated action-values, == 22. O. {where og is applied to D to get )
7: end for
8: Return q = $2,0. 0. (Estimated action-value function of Tes -
B. Theoretical Results
In this section, we present the proofs of our main and supporting theoretical results. The first set of proofs in Section B. show that KROPE is valid operator. While new to our work, the proofs follow those by (Castro et al., 2023). The next set of proofs in Section B.2 prove the stability of KROPE representations and are novel to our work. For presentation purposes, will often be convenient to refer to state-action pair as æ E X instead (s, a).
B.1. KROPE Operator Validity
Given the definition of the KROPE kernel, we now present an operator that converges to kefi;: Definition (KROPE operator). Given a target policy Ne, the KROPE operator F x XXXX defined as follows: for each kernel X x X + R, Y(si, 02) X x X,
F (1,222)))) = ki(s1, al: a2) +7 EEs Pl [K(] ais 822. a22]
short-term similarity
long-term similarity
where s ((s))s),a1), P(sb)s2, az) (-(s(1)) (-(6)) and kn (11, S2, 02) a positive semidefiitt kernel.
(((1)a))-r(s20) Fiw Fann
We now present the proof' demonstrating the validity of the KROPE operator. All the proof's in this sub-section model those by (Castro et al., 2023). The primary difference is that our operator is for state-actions instead of states.
Lemma 1. Let KC(X) be the space of positive semidefinite kernels on X. The KROPE operator F is contraction with modulus r in lloe.
Proof. Let k2, l2 E K((X). We then have:
(kI)
FF (k2)
max
xx
max xxxx)
Exx. - Exx- YY)I
-T max
Ex (X' (2'X'
(xx))
< 1z b2
This completes the proof of the lemma.
Lemma 2. Let K(X) be the space of positive semidefinite kernels on X. Then the metric space (K(X)), || (1oo) is complete.
Proof. To show that KC( is complete it suffices to show that every Cauchy sequence [ hnz has a limiting point in K((Y). Since X' a finite, the space of function x'x'' is a mimee-dimensional vector space, which is complete under || I|
16

---

Stable Offline Value Function Learning with Bisimulation- -ased Representations
Thus, the limiting point & RXXX of the Cauchy sequence m22 lies in R Moreover, since we considering only positive semidefinite kernel elements in the Cauchy sequence and they uniformly converge RR-XX, k must also be positive semidefinite. Thus, ((X') is complete under || 1l.0.
Proposition 2. The KROPE operator FF has unique fixed point in K(X)). That is, there is unique kernel ka* E (()) satisfying
hTth (01 (11, a2) |r(s1, 01) r(82, 02)
+ + 166 P.aj. (6), a3; db)1
| I'max
Tmin
Proof. Due to Lemmas and 2, F is a contraction in a complete metric space. Therefore, by Banach's fixed point theorem, the unique fixed point ke exists.
Proposition The KROPE similarity metric dEROPE satisfies:
Var, y X, diRRor (I, v) |r(x) -ryy) (MM Pm (*u())))
Proof. To see this fact, we can write out the squared Hilbert space distance:
dxRore x. y) 1 (x) (y)
=I, +
(2.1) (x)
2k- (I,
||(x)) (rr)- + >/0 (-1z=), (T (-1x))) w + 7 (P== (+1y)): $(P (-)))()
-2>(0 (-1z)), + D(P (-M((-l))))mp
= Ir(x)) r(v)) + + (+|x)) pd pm (-1y)),
where the third line uses
(-, = 7 xxx ( (1)) (X1, = 7(0( (-1x), pr (+12))
This completes the proof.
Before presenting Lemma 3, we define the distance metric dEROPE: : X' x X + R induced by the KROPE kernel kT as follows:
Vx, E X: dEROPE (x. y) = 1= + xx (B.
22 (I.
Lemma 3. We have (x*)(x) 9dh (v)l I (RROPE(II.)) C, where C Inzo (AT-(x) A.. (v)) and (g-xx)
7M
Ex' ((=o))))()))))) [x xx""xx"-"a" xy (Po()))))) [lr(X)) r(Xx))]]
Proof. We will prove this with induction. We first define the relevant terms involved. We consider the sequences of functions (km)mas and {ym)m20. where ko, go = 0. Since T and T are contraction mappings, we know that limm-+oo km Te and limm- Im qT as T and T are applied respectively each iteration m. the mith application of the operators, have the corresponding kernel function kn along with its induced distance function (2, y) (x, + km (v. v) 2km (2, We will now prove the following for all m:
1 m
(m((x) - 9m S dm (x, v) +
7 "AA (x) (v)
(6)
2
A20
where A Te (z) Ex' (
[ [Ex; xx (x - r(X"))I
17

---

Stable Offline Value Function Learning with Bisimulation- Representations
The base case m 0 follows immediately since the LHS is zero while the RHS can be non-zero. We now assume the induction hypothesis in Equation (6) is true. We then consider iteration m 1:
Mm++ (x) 2m++1 (v)
r(x) (a r(v) -
4m (Y')]l
s |r(x) -r(v) + (1) 1m (X') 4m (Y))
dm (X', (X.)))) + 1 m
||(()) r(v)] +
Σ E (A (X') (
2
m++1
= -r(x)) - r(y)) + (1
dm (X', x +
(A n
+2)
2
where we have used the fact that E. (1) (Xx = A 1 We can then proceed from above as follows:
(x)).
= Ir(x) r(()l + - - dm (X', Σ ("(AT(()) A. (v))
<ir(x) - r(v)| ( (X' Y']]
+ E (
(1 r(x))
2 / (-18)
1 m+1
N (x) + (1)
a=1
1 m+1
[x()) r(v)l
(X', Y'] + Σ 7" (2)
dm+ (x, 1) + 1 m+1
Σ 7" "AA (x) 2
We thus have 19m+1 (2:) 1m+1(u)) (v) S dm+ (2, (r, + Σー 7" (A (x) + AF= which completes the proof.
Lemma 3 tells us that the KROPE state-actions that are close in latent space also have similar action-values upto constant C' : Z Im 7" (A (x) + A (v)). Intuitively, A (x) is the expected absolute reward difference between two trajectories at the nthh step after Te is rolled out from æ. If the transition dynamics and are deterministic, we have C = 0 (Castro, 2020; Zhang et al., 2021). Note that while the deterministic transition dynamics assumption is eliminated, the bound suggests that KROPE may hurt accuracy of qTT since when dxROPE y) = O, we get 1 (x) q" (2)l < C. This indicates that two state-actions that may have different action-values are considered the same under KROPE. This implies that while and y should have different representations, they actually may have the same representation.
B.2. KROPE Stability
In this section we present our main results. We present supporting theoretical results in Section B.2.1 and main theoretical results in Section B.3. To the best of our knowledge, even the supporting proofs in Section B.2.1 are new.
B.2.1. SUPPORTING THEORETICAL RESULTS
We present the following definitions that we refer to in our proofs.
Definition 8 (Bellman completeness (Chen & Jiang, 2019)). The function class F is said to be Bellman complete if E F, holds that T f E F. That is supjer inf Ilg T f ooo = O, where F C X operator.
+
b and T is the Bellman 11
Tun
Tan
Definition 9 (Piece-wise constant functions (Chen & Jiang, 2019)). Given a state-action abstraction d, let F X + Then E F is said be a piece-wise constant function Vx, y X where 6(x) Φ(v), have
Cain
Yous
f(x) = f())
18

---

Stable Offline Value Function Learning with Bisimulation- based Representations
Proposition 4. If a state-action abstraction function X X0 a mee bisimuulaaioo abstraction, then J is Bellman complete, that is, se y'ez I| T =o 0.
Proof. We first define ae-bisimulation (Castro, 2020). Note that (Castro, 2020) considered only state abstractions, while consider state-action abstractions. is considered π. mi-iisimulation abstraction if it induces mapping between X' and x such that for any z, E X' such that z, Φ(z), we have:
1. r(x()) = r(u)
2. yx E 20, Exveno (2(lx() = Ex'end p (x'l))
Given our Th -iiiimuuuatio abstraction function o, we can group state-actions actions according to its definition above. Once we have this grouping, according to Definition 9, ø induces a piece-wise constant (PWC) function class F Note that by definition of & we have:
Cy
max 31 (21 )
(r(z1) r(22)) =
Σ p pT (x'(x1) Ewe
Σ
p
0, Vo 00, x
6p
max Z1 (1:00(((1)) (
Once have d, consider the following to show Bellman completeness. Our proof closely follows the proof of Proposition
20 from (Chen & Jiang, 2019). First recall the definition of Bellman completeness from Definition 8: Vf E F, VT =- f E G. sup inf gng llg - T occ = 0. Given that the smallest value Vf E F. E G, supser geg llg T fl can take on is zero, we will prove our claim by showing that vʃ E F,VT"sf f E G. sup inf lig - T fl is upper-bounded by zero when Ø is a =.-bisimulation.
We will prove the upper bound by showing that there exists a function f _F such that - flloo S 0. which implies that inf fr I| Th fll 00 < 0.
We now construct such I' E Jº, We first define the following terms for given abstract state-action x E X2:
x+ = arg maxree--(ax) (T
(T ()()) and x- arg min (a=)n (a=>> (T fT)).. We can then define f' as follows:
f((x) = 1 ((++)))))++
(x++)) +
f)<=)) Yo ".
And since this holds true for Vr E xd, fi is piece-wise constant function. We can then upper bound || 1J"-Th. T= flloc as
19

---

Stable Offline Value Function Learning with Bisimulation-base Representations
follows:
((()) - (T (x)
1
$
((T (x=+)) + (xx-)) f((=)
2
= (1T (x++ ) - (T f)x-))
(r(#4) + B pro (st (x( r(x-) yE -poce (( )I)
2
2
Exg Pae
(x'+ - Ex Poo 118 ((-((x'_)]
(1)
2
Σ
p (x'1x_=)))
2
Z'EX
Σ
Σ sm (d'1))) - Σ (me (vclor))
2
EtE
ez'Ezo
Exx
f (20)
Σ pae (ro(((((0+)) Σ
pu pm (x'((s_)
(2)
2
2
F'EEU
x'Exv
7
Σ fm (x=0) (Pr(x*)x4+)) Pr(x
(3)
2
2 ||p((*((x4) - Pr(x*1x [((a=l=-)]] )11 (
13
(4)
K0
6p =0 0
where Pr denotes probability, (1) is due to 22022, 2 since PWC, (3) is due to Pr(x*/<:)- Eveose pd (1'l(r))) and (4) is due Hölder's, (1)())())) S
(=)))((()))((()))))) ((())))((2))) = is due to fT (T)
Similarly, we can show the other way around: (T f)(() (x) < 0 by giving the symmetric argument starting with (T)=f)(x) fz) < (T (z++ 2 (x-=) (T (a (x-)) )) Therefore, when ø is ne re-bisimulation, we have sup inf y II - T flloo = 0.
Lemma 4. Define the matrix K1 E R such that each entry is the short-term similarity, k2, between every pair of state-actions, i.e,, Ki (81, (1188, 82, a2) == 1 Finer- Tame |r(11, 01) r(82, u2)]. Then K1 is positive semidefinite matrix.
Proof. Proposition 2.21 from (Paulsen & Raghupathi, 2016) states that any kernel is positive semidefinite if it takes the form: k(a,b) min (a, b) where a, b [0,
First, recall that r(s, a) [-1, we then have each entry in the K, matrix of the following kernel form K (x. v) 1 - 1 I yl. We can then re-write kj as follows:
k (x,z) =
1
y
2
1
1+ min( -x, -y + min minfar y)
2
min
[1 xx, 1 v\ + min + x, 1 v)
k.
kp
That is.
(x = ka(a, y) + k6 y)-
Since x E [-1, 1), each term in the min function is non-negative. Thus, h2 and hú are both positive semidefinite kernels, which means kn is also positive semidefinite kernel. We then have that K₃ is a positive semidefinite matrix.
20

---

Stable Offline Value Function Learning with Bisimulation- based Representations
Lemma 5. Given finite X and kernel k defined on X, let K (k(a,, 9))x. E RI be the corresponding kernel matrix. If K is full-rank and MMD(k))pp.)) 0 for two probability distributions p and q on X, then p 4-
Proof. From (Gretton al., 2012), we have the definition of MMD between probability distributions p, given kernel k:
MMD(K)(p.g) = || (e )) Ex-ggBit [b( ))l
Now when MMD(k)(p.d) g) = 0, we have:
0 = | (kopppk(2,, [) Exng (k(c,, (k(
which implies
0 = (0 :]] Ex-alk(a, .)]l2
since all norms are equivalent in a mnite-dimensional Hilbert space. With p and viewed as vectors |! the above equality means
KKKyKKKa - Kqll2-
Hence, K(p- - 9) 0. Since K is full rank by assumption, we conclude that = q.
Lemma 6. Suppose we have a reproducing kernel k defined on the finite space X, which produces a reproducing kernel
Hilbert space (RKHS) Hy, with the induced distance function d such that d(x, y) = k(z, 2) + kly. 2k(x, y), Vz, X. When d(r, y) = 0, then k(xx)
k(y," ))
Proof. When d(z, v) = O, we have 2k(2r, v) = k: x) + k(y, v). Therefore, we the following equalities:
k(x, + k(y, y) = 2k(x, v)
k(0, k(æ, y) = k(I, 8) k(y. y)
(M(r, ,), (10 1)H, (k((r,) )) J, 1) PM = (M<(,, -), k(y, Hr, . (k(y -), ky,, 1)H.. (k(2, k(x, -) k(y, .)) (A(2,-) Myy. 3,k(
(1)
(2)
(k(2 -). k((, k(y, 1) k(z,,), -k(y, -)m
(3)
k(x,,) k)y,
where (1), (2), and (3) are is due to RKHS definition, linearity of inner product, and symmetry of inner product respectively.
Proposition 5. Let ,I,..., In E (0, 00) be distinct and strictly positive numbers. Let K E Ruxn be the matrix with entries Kij- = min((;;,) 2y}. Then K is a positive definite matrix.
Proof. By Proposition 2.21 in (Paulsen & Raghupathi, 2016), the matrix K is positive semidefinite, so only need to show that K is full rank. WLOG assume that 0 < x2 < ... < In. We prove by induction on n. The base case with n clearly holds. Suppose the claim holds for n numbers. Now consider n numbers. Let a = (a1...., an)) E R". suffices to show that Ka implies a 0. We write K block matrix form as
0
x2 Li
32
F1
0
K IIJn +
+
1
u2
Up
32 21
Fn
F1
where Jn is the n-by-n one matrix, 1 E 1 the all one vector, 0 E 1 the all zero vector, and u; E 1 2.....7. It follow that
Σ - a
0 = Ka
(x) En 0: 1 + Σn 2i=2 Uz4;
21

---

Stable Offline Value Function Learning with Bisimulation-bast basedd Representations
that is,
7
T1
0, i=1
0
(7)
R
Ø1
04
1
a =0
(8)
i=1
Plugging equation (7) into equation (8), we get
Σ 2
in
= 0. By the induction hypothesis, the (n - 1y-by-nn - 1) matrix
U2
Un min {
r1 i.j=2.....
has full rank since the (n 1) numbers 22 21....; In Φ1 are distinct and strictly positive. Therefore, we must have
Q2
Og Plugging back into equation (7) and using 1 >> we obtain 01 0.
B.3. Main KROPE Theoretical Results
We now present the main theoretical contributions of our work. Theorem I a KROPE representation as defined in Definition 2, then the spectral radius of [E[$ -1)/7>)))>)>> pd less than That is, stabilizes LSPE.
Proof. Recall from Definition 2, we have:
E Ki + (P J.
where K1 E x| such that each entry is the short-term similarity, ki, between every pair of state-actions i.c. Ki (81, (1; 82, a2) = (n(a).0))- ((a.a)))))aa)a)))
Taus
From this definition, we can proceed by left and right multiplying T and respectively to get:
[ T ΦΦ
R[
K1
rF p + b) $).
Notice that = E( T +) the feature covariance matrix and C == E p the appropriate substitutions, we get:
the cross-covariance matrix. By making
BB E( K +) +CC
We can then left and right multiply by B-1 and B to get the following where L:= -B-
1= B B-'E K + LIT.
7
Rearranging terms gives
I I- LL = B K- 7
From Lemma 4, we know that K2 is positive semidefinite, which means that B- B- T B is also positive semidefi- nite. Therefore, the eigenvalues of LHS above must also be greater than or equal to zero. Letting À be the eigenvalue of L, we know that that the following must hold:
X2
20 0
pT) is always less than 1. Thus, KROPE representations are stable. Finally, since KROPE representations are stable and due to Proposition 1, KROPE representations stabilize LSPE.
Since 7 < 1, the spectral radius of L = (E (>E
Theorem 2. Let Φ= : X - 2 be the state-action abstraction induced by grouping state-actions x, E X such that if dxROPE(x, 8) = 0, then Φ(x) = σ(y), Vx, E X. Then o is Bellman complete if the abstract reward function y ; X + (-1,1) injective (i.e,, distinct abstract rewards).
22

---

Stable Offline Value Function Learning with Bisimulation-base based Representations
Proof. Our proof strategy is to show that the abstraction function & due KROPE is me bisimuuuultion,, which implies it is Bellman complete due Proposition 4.
According to Propsition 3, dRRaPE (z. y) lr(x) - +
( (-l()) When dgRoPe v) = 0 for any two state-actions, it implies that r(x) = rly) and (M (P (-1z), P (-lu) = O.
+ MMD kT (P** (PP
For @ to be a Te bisiiuuuulatinn, we need Vx E X0, En' PT (z'1z) = En'er pm (x (y) to be true for any 2,y E X' such that x(x) = dly). While (MMD((*=)(P**)))) (-1x), pT (-1)) = O, it is possible that pd (-1x) + (-lu). However, as we will show, under the assumption that the abstract rewards r are distinct Vz E X we do have Vr E X Erew' P= = pa (x'ly)- Before we proceed, we make the following technical assumption on the reward function: r(a) E (-1,1), Vx E X. The exclusion of the rewards -1 and I allows us to use Proposition 5 to show that the KROPF kernel is positive definite instead of positive semi-definite.
Once we group state-actions 2, E X' together such that deRore (z. y) = O, we have the corresponding abstraction function 6: X + A Accordingly, & induces a Markov reward process, MM == (x .r,, ,r,,p, P 7) where r is the abstract reward function r X - (-1,1) and pd is the transition dynamics on the abstract MRP i.e. Po (-]a.)). We can also consider the abstract KROPE kernel, k (xx,, y² which measures the KROPE relation on X² Note that all these quantities are function of Te. We drop the notation for clarity. By this construction, we have:
r (x0) = r((x), Vx E xd k (a..)) k(2,,), Var E xd
Since all rewards are equal within x.
Lemma 6
Now, under the assumption that all abstract rewards (x9) are distinct Va E X², we have that the kernel matrix K E xxxxxx where each entry (x y² is positive definite. To see this fact, consider that:
1...30 20) + k (x ) (1
1k yd)]
(9)
where (1 y²)= =
(x²) r (1 From Lemma 4, we know that k is positive semidefinite. However,
under the assumption that all abstract rewards ~ are distinct, Proposition 5 tells us that ky is positive definite. Given that (x y²) (Equation (9)) is just a summation of positive definite kernels, hº is positive definite, which means K is positive definite.
We now consider when the MMD is zero Again, by construction. we have the following when MMD(K )(PT=)((=))))))))))))))))) pr
(x (-19) O. For clarity, we use k instead of km*
0 |X'' (110 [x(X', ))] - Ex pme (-1v) [M(X', )))ll
= ||x' (110) [x:(X,, ,))] - - (-1v) [X(X', ,]]l
|| Σ p (s' (xdd(()))('))) - Σ pse (1' ("c(x'ly)k(('' (1)k(('(,)))
x'EX
2'EX
Σ Σ
pT (c(o())))((
Σ
Σ
p. (2d (o)l)))((
xPEX E'
xPEX E'nn
Σ
A (0 ()) Σ
p - Σ
(x ..)
Σ
pi (2oll)
(1)
x' Z'E
a Ex
Z'Err
600 .) Pr(x)))))
Σ
* .)) Pr(r)) (v)
Pr denotes probability
x
where (1) is due to h (2*0 = k(z, -), Vx E x From above, we can see that the kernel and probability distributions are over X In matrix notation, we can write the above as follows where pd == Pr(((ar) and qd == Pr(-ly) are viewed as
23

---

Stable Offline Value Function Learning with Bisimulation- based Representations
probability distribution vectors in R 1
0 KK - 9b pd 9
since K is positive definite, from Lemma
We thus have Vx E X0; - P (2'lz) = Eweer (2'1l2 to be true for any 2, y E X such that o(x) = dly)
Given this condition holds true and r(x) vyy) E X such that <(x) <ly), @ is a me ne-bisimulation. From Proposition 4 we then have that & Bellman complete.
Remark on the injective reward assumption. The injective reward assumption simply means that each abstract state-action group will have a distinct associated reward from every other abstract state-action group. This assumption comes as a trade-off. Chen & Jiang (2019) proved that bisimulation abstractions are Bellman complete. Instead of assuming injective rewards, they assumed that two states were grouped together if each state's transition dynamics led next states that also grouped together (one of the conditions for exact bisimulations). This condition is also considered strict and inefficient to compute (Castro, 2020).
In our work, we relax this exact transition dynamics equality by considering independent couplings between next state distributions, thereby making the KROPE algorithm efficient to compute. However, the drawback of this relaxation is that preserving distinctness between abstract state-actions may be lost (see remarks on ka* Section 3.1). To ensure the distinetness between state-action abstractions, we assumed that the reward function is injective. This then allowed to show Bellman completeness in similar way that shown in Chen & Jiang (2019).
C. Empirical Details
In this section, we provide specific details on the empirical setup and additional results.
C.1. Empirical Setup
General Training Details. In all the continuous state-action experiments, we use a neural network with 1 layer and 1024 neurons using RELU activation function and layernorm to represent the encoder Φ X' - Rdd (Gallici et al., 2025). We use mini-batch gradient descent to train the network with mini-batch sizes of 2048 and for 500 epochs, where a single epoch is a pass over the full dataset. We use the Adam optimizer with learning rate (1e-5, 25 -e and weight decay 1e-2, The target network is updated with a hard update after every epoch. The output dimension is {1//4 x2/2, X/44), where | is the dimension of the original state-action space of the environment. All our results involve analyzing this learned o. Since FQE outputs a scalar, we add a linear layer on top of the d-dimensional vector to output a scalar. The entire network is then trained end-to-end. The discount factor is y 0.99. The auxiliary task weight with FQE for all representation learning algorithms is a = 0.1. When using LSPE for OPE, we invert the covariance matrix by computing the pseudoinverse.
In the tabular environments, we use a similar setup as above. The only changes are that we use a linear network with a bias component but no activation function and fix the learning rate to be 1e-3, For the experiment in Section 4.4, a 0.8. We refer the reader pseudo-code in Appendix A.
As noted earlier, measure OPE error by measuring MSVE. To ensure comparable and interpretable values, we normalize the MSVE by dividing with MAND == E(s.a)~D (S RAN A) - - A)) where GRAND is the action-value function of random-policy. Similarly, in the continuous state-action environments, we normalize by MSVE(gRAND] = Eso-do.Aoosss [[g RAND (So, Ag) q= (So, Ao)) Values less than one mean that the algorithm estimates the true performance of ग. better than a random policy.
Evaluation Protocol: OPE Error
Evaluation Protocol: Realizability Error. In tabular experiments, we normalize the realizability error. After solving the least-squares problem == - -== 112. where i =- arg min.. || - -== We divide by मे E, T (æx) and plot this value.
24

---

Stable Offline Value Function Learning with Bisimulation-base Representations
Pearson Correlation. The formula for the Pearson correlation used in the main experiments is:
Ei (r; x (1 5)
V 100E:-3334 Div (3/ Ø)²
where è and y are the means of all the x; and yi respectively.
Custom Datasets. We generated the datasets by first training policies in the environment using SAC (Haarnoja et al., 2018) and recording the trained policies during the course training. For each environment, we select 3 policies, where each contributes equally to generate a given dataset. We set Fe to be one of these policies. The expected discounted return the policies and datasets for each domain given in Table 2 (7 0.99). In all environments, Te zi (see Table 2). The values for the evaluation and behavior policies were computed by running each for 300 unbiased Monte Carlo rollouts, which was more than a sufficient amount for the estimate to converge. This process results in total of datasets, each of which consisted of 100K transitions.
Environments
We
"b
CariPoleSwingUp
50
20
5
FingerEasy
100
71
32
HalfCheetah
51
27
WalkerStand
90
55
40
Table 2. Policy values of the target policy and behavior policy on DM-control (Tassa et al., 2018).
D4RL. Datasets. Due to known discrepancy issues between newer environments of gym', we generat our datasets instead of using the publicly available ones. To generate the datasets, we use the publicly available policies For each domain, the expert (and target policy) was the 10th (last policy) from training. The medium (and behavior policy) was the 5th policy. We added noise of 0.1 to the policies. The values for the evaluation and behavior policies computed by running each for 300 unbiased Monte Carlo rollouts, which was more than a sufficient amount for the estimate to converge. We set y 0.99. We evaluate on the Cheetah, Walker, and Hopper domains. This generation process for three environments, led to datasets, each of which consisted of 100K transitions.
Toy Divergence Example We set 7 1. The sampled state-actions are fed into linear network encoder with bias component and no activation function which outputs a d = 3 representation. This representation is shaped by KROPE and this representation is then fed into linear function to output the scalar value for FQE. D is a dataset of 2000 on-policy transitions. D* is an oof-po dataset with 5000 transitions starting from the specified state s, where s = { [W1, W2. W3. 2wa), where the next state is sampled according to the transition dynamics of the MRP.
Regarding the statement of, "we would generally expect that the probability of sampling bad transition pair is less than that of sampling single bad transition", assume that the probability of choosing a bad transition is p and a good transition is 1 - p. Then sampling pair of bad transitions versus sampling single bad transition involves comparing the probabilities of p² and 2p(1 - p), respectively. In general, we would expect 2p(1 p) 2 p2, unless p 2 0.66, i.e., the dataset starts to overflow with bad transitions. Roughly, we would expect FQE to diverge when p 2 0.5, but we would expect KROPE to diverge when p 2 0.66, which suggests that KROPE less prone to divergence than FQE.
make the following simplistic and rough calculation. For a given dataset,
C.2. Baselines
We provide details of the 6 baselines in this section.
DBC. It is bisimulation-based metric learning algorithm that learns r-bisimilar representations by ensuring the I.1 distance between the latent representations approximates the r-bisimulation metric (Zhang et al., 2021). During the learning
Phttps: http://oooooooooooooooooooooo//a httpps: hittps://git //g thub. cww reseelleeeeeerre/deeeee
25

---

Stable Offline Value Function Learning with Bisimulation-base Representations
process, it also learns a model of the environment. Its additional hyperparameters the learning rate of model learning (which we to 1) and the output dimension of Φb
ROPE. It is a bisimulation-based metric learning algorithm that is off-policy evaluation variant of MICO (Castro et al., 2021; Pavse & Hanna, 2023a). It learns representations directly such that the r-bisimulation metric is satisfied and makes no assumptions on the transition dynamics. Its additional learning rate is the output dimension of Φ. Furthermore, we provide additional details on our experimental setup is different from that in (Pavse & Hanna, 2023a). The main differences are:
1. The original ROPE used ROPE a pre-training step, fixed the representations, and then fed them into FQE for OPE. In our case, we also pre-train the representations but with FQE as representation learning algorithm (for value predictive representations (Lehnert& & Littman, 2020)) along with other representation learning algorithms as auxiliary tasks, The fixed learned representations are then fed into LSPE for OPE.
2. The original ROPE encoder architecture had a TANH activation function on the output layer, which we effectively serves a clipping mechanism to the features, which is similar to how public implementations of FQE clip the return to avoid divergence".
We opted to our alternative setup for two reasons: 1) By using LSPE as the OPE algorithm, instead of FQE, we can precisely quantify the stability properties of the representations in terms of the spectral radius (Theorem 1), which is harder to do when using FQE as the OPE algorithm; 2) while valid architectural choice, for this current work, we viewed the use of the tanh function as obfuscating the true stability properties of the representations and so opted to avoid it. More practically, it is reasonable to use the TANH as part of the architecture.
BCRL. Unlike the other algorithms. BCRL is not used as an auxiliary loss with FQE (Chang al., 2022). We use the same learning rates mentioned above for (1e-b, 2e-5, 5) when training ø. As suggested by prior work, self- self-predictive algorithms such BCRL work well when network that outputs the predicted next state-action is trained at a faster rate (Tang et al., 2023; Grill et al., 2020). Accordingly, we set its learning rate to be 1e-1, For BCRL-EXP, which involves the log determinant regularizer, we set this coefficient to le BCRL's 's hyperparameters are: the learning rates for é, M, and the output dimension of Φo: and the log determinant coefficient (see Equation (10)).
DR3. The DR3 regularizer minimizes the total feature co-adaptation by adding the term ((s.,. dls To ohs's d') as an auxiliary task to the main FQE loss (Kumar et al., 2022). (Ma et al., 2024) introduced an improve- ment to this auxiliary loss by suggesting that the absolute value of the feature co-adaptation be minimized, i.e.. "( a' Jó(s, a) "'s's a')l. We use 0.1 as its auxiliary task weight. Absolute DR3's only hyperparameters are the auxiliary task weight a and the ø output dimension.
BEER. (He et al., 2024) introduced an alternative regularizer to DR3 rank regularizer since they suggested that the minimization of the unbounded feature co-adaptation can undermine performance. They introduced their bounded rank regularizer BEER (see Equation (12) in (He et al., 2024)). BEER introduces only the auxiliary task weight a as the additional hyperparameter.
KROPE. KROPE'S only hyperparameters are the output dimension of and the learning rate of the KROPE learning algorithm.
C.3. Additional Results
In this section, we include additional empirical results.
Realizability. In addition to stability, we also care about the realizability of क. We say &b is realizable representation if qT= E (pan($). where S2an( is the subspace of all expressible action-value functions with 4. Note that even if & is a realizable and stable representation, LSPE may not recover the TT= solution (Sutton & Barto, 2018).
basic criterion for learning nd is realizability. That is, we want e == || - gd I|g. where ii == arg min,.. || - q lg. to be low, In our experiments, we compute and plot it as a function of d in Figure 5(a). A critical message from our
httpps://gi www.hh.w..mmooooe www./oool/eeeooooooooooooooooooo weweeee-reeeerhcgeooeeeee e- aeweeeaarr/bbbbbb//attteeee
fitter..py
26

---

Stable Offline Value Function Learning with Bisimulation- based Representations
10
10
DR3 ABS
Realizaailit Eroo
109
20-3
10-4
10
10
20
30
d dimension)
(a) Realizability Error
Figure 5. Evaluation of basic representation properties on Garnet MDPS with 40 state-actions vs. output dimension d. Figure 5(a):
Realizability error; lower is better. All results averaged over 30 trials and the shaded region is the 95% confidence interval.
results is that stability and realizability do not always go hand-in-hand. While BCRL-EXP-NA has favorable spectral radius properties (Figure 2(a)), it has poor realizability, which will negatively affect its OPE accuracy. KROPE, on the other hand, has favorable stability and realizability properties up till d = 30. When d 2 40, the realizability error is 0 for all algorithms since the subspace spanned by & is large enough to contain the true action-value function (Ghosh & Bellemare, 2020). While the realizability error is for d 2 40, the representations can unstable (Figure 2(a)).
C.3.1. LEARNING GRAPHS FOR OFFLINE POLICY EVALUATION ON DMC AND D4RL DATASETS
In this section, we present the offline policy evaluation results on the DMC and D4RL datasets. We present the results in Figures 6 and 7.
Normalized Sq. Value Enor
Nonmalized Sq. Value Errar
101#
Normalized Vallue Emar
10th
Normalized Vallue Emar
10
10th
1n14
1014
1074
1017
ExP-na
EXP-na
BCRLEXP-na
BCRL-na
FQE
BCRL-na
FQE
BCRL-na
FQE
BCRL-na
FQE
KROPE
KROPE
BEER
DBC
ROPE
DBC
ROPE
DBC
ROPE
DBC
ROPE
10-
10
10
20
40
40
40
Epochs (x10)
Epachs (x10)
Epochs (x10)
Epochs (x10)
(a) CartPoleSwingUp
(b) ChectahRun
(c) FingerEasy
(d) WalkerStand
Figure Normalized squared value error achieved by LSPE when using particular representation vs. representation training epochs on the custom DMC datascts. LSPE estimates are computed every 10 epochs. Results are averaged over 20 trials and the shaded region is the 95% confidence interval. Lower and less erratic better.
On the DMC dataset, find that KROPE reliably produces stable and low MSVE during the full course of training.
On the D4RL dataset, we reach the similar conclusions: KROPE is effective in producing stable and accurate OPE estimates. However, in 3/9 instances, KROPE does diverge. This divergence is likely related to the discussion in Section 5 and Section 4.4. Recall that KROPE a semi- aami-gradient method, which does not optimize any objective function and is susceptible divergence (Feng et al., 2019; Sutton & Barto, 2018). So while KrOpf representations stabilize value function learning, KROPE'S learning algorithm may diverge and not converge KROPE representations. However, regardless of this result, KROPE does improve the stability and accuracy of FQE in all cases.
C.3.2. STABILITY-RELATTE ANALYSIS ON CUSTOM DATASETS
In this section, We include the remaining stability-related metric analysis that was deferred from the main paper.
Hyperparameter Sensitivity. In this subsection, we include all the remaining results related to the stability metrics for all environments.
27

---

Stable Offline Value Function Learning with Bisimulation- based Representations
Normalized Sq. Value Error
Normalized 5q Value Error
10
10th
1014
10
Normalized Value
1015
BCRL-EXP-na
DR3-ABS
DCRL-EXP-re
BCRL-EXP-ne
DR3-ABS
BCRI na
FOF
BCRL-na
FQE
BCRL-na
FOE
BEER
KROPE
BEER
КНОРЕ
DBC
ROPE
10a
10
100
20
40
20
20
40
Epochs (x10)
Epochs (x10)
Epochs (x10)
(a) Cheetah random
(b) Cheetah medium
(c) Cheetah medium-exee expert
Normalized Sq, Value Error
ue Error
Noraaizzed 5q. Error
1020
102>
10
10
BCRL-EXPMM BCRL-na
DR3-ABS
BCRL-EXP-na
DR3-ABS
BCRL-EXP-ma
DR3-ABS
FOE
BEER
KROPE
BEER
KROPE
BEER
KROPE
DBC
ROPE
DeC
ROPE
To
DBC
ROPE
Nomall:
10-
104
209
20 Epochs (x10)
20
40
20
40
Epochs (x10)
Epochs (x10)
(d) Hopper random
(e) Hopper medium
(f) Hopper medium-expert
1018
Normalized Sq. Value Error
1020
Normalized Sq. Value Enor
1020
Normalized Value
1014
10366
10
BCAL-EXP-nD
BCRL-EXP-na
DR3-ABS
BCRL-EXP-na
DR3-ABS
BCRL-na
FOF
BCKLina
HQE
1040
BEER
BEER
KROPE
BEER
KROPE
DR3-ABS
109
KROPE
10a
10
109
40
Epochs (x10)
Epochs (x10)
Epochs (x10)
Walker random
(h) Walker medium
(i) Walker medium
Figure 7. Normalized squared value error achieved by LSPE when using particular representation vs. representation training epochs on
the D4RL datasets. LSPE estimates are computed every epochs. Results are averaged over trials and the shaded region the 95% confidence interval. Lower and less erratic is better.
1..00
11.00
Frac. of MSVE
Frac. of runs w/ MSVE
0.75
0.75
0.50
0.50
0.25
BONL ma
POE
DK3-AB5
0.25
DCRL-ne
BORL-EXPas
FOE
DRS-BBS
BEER
KNorE
BEen
KROTE
DEC
ROPE
DBC
HOVE
0.00
0.00
MSVE
(a) CheetahRun
(b) FingerEasy
Figure 8. Hyperparameter sensitivity on different environments as a function of training epochs; larger area under the curve is better. All results averaged over 20 trials for each hyperparameter configuration and shaded region is the 95% confidence interval.
28

---

Stable Offline Value Function Learning with Bisimulation- based Representations
Bellman completeness. Another metric that is associated with stability is Bellman completeness (BC) (Chang et al., 2022; Wang al., 2021). We find that KROPE is approximately Bellman complete even though it does not explicitly optimize for it; this finding aligns with our Theorem 2. While BC is difficult to approximate, we can minimize the proxy metric introduced given in Equation (10) (Chang et al., 2022):
M ((s a)
( ( a'Il] ]
C(M. p) = ED
(10)
r(s,a)
where (p, M) E Rdddd, @ is fixed, and ((M, p) = indicates Bellman completeness. Given the final learned representation. we compute and report the BC error in Table 3. We find that KROPE is approximately Bellman complete even though it does not explicitly optimize for it; this finding aligns with our Theorem We note that BCRL-EXP is less Bellman complete since it also includes the exploratory objective in its loss function, which if maximized can reduce the Bellman completeness. While BCRL is more BC than BCRL-EXP, we found that is less BC in general. We attribute this finding due to the difficulty in explicitly optimizing the BCRL objective which involves multiple neural networks (M, p. Ø) and multiple loss functions on different scales (reward, seff-prediction, log determinant regularization losses). KROPE can achieve approximate Bellman completeness without these optimization- difficulties.
Algorithm
Domain
BCRL. + EXP BCRL
BEER
DR3
FQE
KROPF (ours)
CartPoleSwingUp
0.4 0.4:::0 0.1
0.2 0.200.1
0.1 + 0.0
0.0
0.1
0.0
0.0
CheetahRun
3.3 0.6
2.4
0.7 + 0.0
0.0
0.0
0.2 0.0
FingerEasy
1.3 0.6
0.7
0.0
137.0
4.4
0.9
0.0
0.2
WalkerStand
10.4 + 2.0
0.3 = 0.1 0.5 + 0.1 66.1 + 0.6
0.6 0.0
0.1 + 0.0
Table 3. Bellman completeness measure for all algorithms across all domains. Results averaged across 20 trials and the deviation shown is the 95% confidence interval. Values are rounded to the nearest single decimal.
C.3.3. USING FQE DIRECTLY FOR OPE
In our main empirical section (Section 4), we used FQE as representation learning algorithm custom datasets. We adopted the linear evaluation protocol, i.e., an approach of analyzing the penultimate features of the action-value function network and applied LSPE on top of these features for OPE. This protocol enabled us to better understand the nature of the learned features.
For the sake of completeness, we present results of FQE as an OPE algorithm where the action-value network is directly used estimate the performance of g.. We present the results in Figures 9 and 10. As done in Section 4, evaluate the performance of FQE and KROPE based on how they shape the penultimate features of the action-value network. However, when conducting OPE, we evaluate two variants: 1) using LSPE (-L) and 2) using the same end-to-end FQE action-value network (-E2B).
From Figure 9, we find that there are hyperparameter configurations that can outperform the KROPE variants. However, based on Figure 10, we find both KROPE variants are significantly more robust to hyperparameter tuning. This latter result suggests that KROPE does improve stability with respect to the hyperparameter sensitivity metric as well.
Regardless of FQE'S hyperparameter sensitivity, it is still interesting observe that when FQE is used as OPE algorithm, it produces reasonably accurate OPE estimates. It even outperforms the FOE+KROPE combination. The primary difference between using FQE for OPE vs. FQE features and LSPE for OPE is in how the last linear layer is trained. The former is trained by gradient descent while the latter is trained with the iterative LSPE algorithm on the fixed features. An interesting future direction will be to explore the learning dynamics of these two approaches.
On a related note, we point out that the training dynamics of FQR are still not well-understood. For example, (Fujimoto al., 2022) show that the FQE loss function poorly correlates with value error. That is, the FQE loss can be high but value error (and OPE error) can be low.
29

---

Stable Offline Value Function Learning with Bisimulation- based Representations
Normalized Sq. Value Emor
Normalized Value Error
1016
NE
109
1012
FOE-FEZE FOF-I KROPE-FEZE KROPE-L
10"
10-1
FQE-FEZE
109
FQE-L
109
KROPE-L
20
20
40
Epochs
Epochs (x10)
(a) CartPoleSwinUUp
(b) FingerFasy
Normalized lizee Sq. Value Error
1029
Normalized Sq. Value Error
10
1014
FQE-FEZE
10
FQE-FEZE
109
KROPE-FEZE KROPE-L
KROPE-FEZE
103
KROPE-L
10-1
104
20
40
20
40
Epochs (x10)
Epochs (x10)
(c) CheetahRun
WalkerStand
Figure 9. Normalized squared value achieved by LSPE (-1.) and FQE (-E2F) evaluated every 10 epochs of training. Results averaged over 20 trials shaded region the 95% confidence interval, Lower and less erratic better.
1.00
1.00
Frac. rurus w/ MSVE <
rums MSVE
0.50
0.50
FEE
FOE FEZE
0.25
ROEL
0.25
FOE
KROPE-FEDE
KADFE-PEDE
KROPE
KROPEL
0.00
0.00
MSVE (7)
MSVE (T)
(a) CartPoleSwingUp
(b) FingerEasy
1.00
1.000
of runs v/ MSVE
Frac. of runs w/ MSVE
0.75
0.75
0.50
0.50
PQE-FE2E
0.25
0.25
ROEL
0.00
0.00
MSVE
MSVE
(c) ChectahRun
(d) WalkerStand
Figure 10. Hyperparameter sensitivity different environments function of training epochs; larger under the curve is better. All results are averaged over 10 trials for each hyperparameter configuration and shaded region is the 95% confidence interval. We tuned the hyperparameters discussed in Appendix C.I. KROPE-FE2E overlaps with KROPE-L..
30


Reviews:
[{"review_id": "SdW7XEtbQa", "reviewer": "Reviewer_3GQD", "summary": "This paper addresses offline policy evaluation in offline RL, which involves estimating expected returns of state-action pairs under a fixed policy using offline datasets. Stability in this estimation process is critical for accurate evaluation. The authors propose KROPE, a new method combining bisimulation-based representation learning (ROPE) with kernel methods. KROPE constructs state-action representations such that pairs with similar immediate rewards and subsequent policy-induced states share similar representations. Experimental results suggest that KROPE enhances the stability of offline value function learning compared to baseline methods.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for acknowledging that our empirical results and theoretical results are solid. Below we address the concerns you raise.\n\n**Use of Mean Absolute Error vs. MSE**\n\nThis suggestion is valid since we understand the MAE may be more robust to outliers. However, it does not diminish the validity of the results since the MSE is a common metric used in the OPE literature, see this OPE benchmark paper [1]. The MSE is considered to be a valid metric due to its relation to the variance and bias of estimators (even though bias/variance may not be reported). We also refer the reviewer to many fundamental and prominent works that evaluate OPE algorithms based on the MSE [3,4,5,6,7].\n\n**Numerical precision in Table 1**\n\nAfter analyzing the raw data, we found that including additional precision did not alter the big picture. Moreover, we were tight on horizontal space, so we truncated it at 1 decimal.\n\n**Relations to Pavse et al. 2023**\n\nThis is a great point, which we will make clear in the paper. Below we state why we opted to use this setup and then explain the differences. Briefly, the cited ROPE paper did: 1) use FQE as the OPE algorithm and 2) used a tanh activation function on the last layer of the encoder.\n\nInstead of adopting Pavse et al.’s approach, we opted to use our alternative setup for two reasons: 1) By using LSPE as the OPE algorithm, instead of FQE, we can precisely quantify the stability properties of the representations in terms of the spectral radius (Theorem 1 and Section 4.2), which is harder to do when using FQE as the OPE algorithm; 2) while a valid architectural choice, for this current work, we viewed the use of the tanh function as obfuscating the true stability properties of the representations and so opted to avoid it. More practically, it is reasonable to use the tanh as part of the architecture.\n\nIn more detail, the main differences are:\n1. The original ROPE paper used ROPE as a pre-training step, fixed the representations, and then fed them into FQE for OPE. In our case, we also pre-train the representations but with FQE as a representation learning algorithm (for value predictive representation [8]) along with other representation learning algorithms as auxiliary tasks. The fixed learned representations are then fed into LSPE for OPE. \n2. The original ROPE encoder architecture (in the cited paper) had a tanh activation function on the output layer, which we effectively serves a clipping mechanism to the features [9], which is similar to how public implementations of FQE clip the return to avoid divergence [2].\n\nWe will make these differences explicit in the Appendix in the camera ready where we discuss the different baselines\n\n**Including learning curves**\n\nThanks for the suggestion. These are included in the Appendix (Figure 6 and 7) due to lack of space in the main paper.\n\n**KROPE does not always converge**\n\nYes, it is more accurate to say that KROPE improves the stability of OPE compared to other baselines, and we will reframe the paper as such. KROPE may not converge since KROPE relies on a semi-gradient learning algorithm. We discuss this limitation in Section 5 and provide insight into when this might occur in Appendix C.3.1. \n\nBriefly, it may be possible to leverage the Legendre-Fenchel transformation and replace the fixed-point loss function of semi-gradient methods with an equivalent expression that avoids semi-gradient learning. However, a drawback with this approach is that the new learning objective is a minimax procedure, which can be challenging to optimize in practice [10].\n\nWe will modify our claim to say that *KROPE improves the stability of OPE compared to 7 other baselines on 10/13 datasets.* (not just highlighted blue errors in table 1).\n\n**Pseudo-code**\n\nAppendix (page 15) includes one, but we will make it clearer.\n\nPlease let us know if we have addressed your concerns. If we have, we would greatly appreciate it if you could re-evaluate your review.\n\n---\n[1] Voloshin et al. Empirical Study of Off-Policy Policy Evaluation for Reinforcement Learning. 2021\n[2] https://github.com/google-research/google-research/blob/master/policy_eval/q_fitter.py#L101\n[3] Chaudhari et al. 2024. Abstract Reward Processes: Leveraging State Abstraction for Consistent Off-Policy Evaluation\n[4] Liu et al. 2018. Breaking the Curse of Horizon: Infinite-Horizon Off-Policy Estimation\n[5] Thomas et al. 2016. Data-Efficient Off-Policy Policy Evaluation for Reinforcement Learning\n[6] Hanna et al. 2019. Importance Sampling Policy Evaluation with an Estimated Behavior Policy\n[7] Sachdeva et al. 2023. Off-Policy Evaluation for Large Action Spaces via Policy Convolution\n[8] Lehnart and Littman. 2020. Successor features combine elements of model-free and model-based reinforcement learning\n[9] Bhatt et al. 2024. CrossQ: Batch Normalization in Deep Reinforcement Learning for Greater Sample Efficiency and Simplicity\n[10] Feng et al. 2019. A Kernel Loss for Solving the Bellman Equation", "claims_and_evidence": "The paper provides robust theoretical and experimental support for its claims, detailed in Sections 3 and 4, respectively. The theoretical part are solid.", "methods_and_evaluation_criteria": "While fundamentally sound, the methodology could be strengthened by addressing the evaluation score choice. The experimental section primarily presents squared value error, whereas the DOPE benchmark suggests using MAE for such evaluations (Fu et al., 2021). An explanation for this deviation or inclusion of MAE results would enhance the paper's credibility and relevance.\n\n----------\nFu et al., 2021.Benchmarks for Deep Off-Policy Evaluation. ICLR 2021", "theoretical_claims": "The theoretical framework is well-constructed with no apparent flaws in the claims or proofs provided.", "experimental_designs_or_analyses": "The presentation and clarity of results in Table 1 need substantial improvement for supporting the claims regarding KROPE's stability:\n1. Numerical results are currently presented with only one decimal place, making it difficult to discern clear differences among methods, e.g., 0.0 vs 0.0 or 0.1 vs 1.0. Higher numerical precision is necessary for meaningful comparisons.\n1. The absence of MAE results (as recommended by the DOPE benchmark) further complicates interpreting and validating the results.\n1. Notably, ROPE appears to diverge in the current experiments, contradicting convergence reported in the original ROPE paper (Pavse & Hanna, 2023a). This discrepancy raises concerns regarding experimental consistency and soundness. The authors should clearly explain this difference.\n\nLearning curves illustrating training stability should be included prominently in the main text, given the paper's central claim regarding stability.\n\nAdditionally, some learning curves presented in Figure 7 (Appendix) reveal that KROPE does not converge in certain environments (e.g., Walker and Hopper). Besides, the large squared value errors from certain baseline methods obscure meaningful comparisons with more stable methods. Addressing this visualization issue is recommended.\n\n\n--------\nPavse & Hanna, 2023a. State-action similarity-based representations for off-policy evaluation. NeurIPS 2023.", "relation_to_broader_scientific_literature": "NA", "essential_references_not_discussed": "NA", "other_strengths_and_weaknesses": "NA", "other_comments_or_suggestions": "1. The addition of algorithm pseudocode specific for KROPE method would enhance clarity.\n1. I would raise my score if the concerns related to the experimental results are thoroughly addressed.", "questions_for_authors": "1. Could the authors clarify whether any experimental settings differ from those used in the original ROPE paper?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "WBBbtFk5vk", "reviewer": "Reviewer_5dn2", "summary": "The paper introduces Kernel Representations for Offline Policy Evaluation (KROPE), a novel algorithm designed to stabilize offline value function learning in reinforcement learning. KROPE leverages π-bisimulation to shape state-action representations, ensuring that similar state-action pairs are represented consistently. This approach enhances convergence and reliability. The authors provide theoretical foundations that demonstrate KROPE's stability through non-expansiveness and Bellman completeness. Empirical results indicate that KROPE outperforms other baselines in terms of stability and accuracy.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for acknowledging that our empirical results and theoretical results are rigorous, and that the paper is clear and well-organized.\n\nBelow we address your concerns.\n\n**Re: reason to learn the value network (“the paper does not further elaborate on the advantages or potential applications of learning such a value network.”)**\n\nIn our context, our goal is to estimate the value of a policy. Our contribution is particularly relevant to the off-policy evaluation (OPE) literature [1, 4, 5]. In OPE, we want to use an offline dataset generated by different policies to estimate the performance of another target policy since deploying the target policy directly in the environment may be risky or costly. One way to evaluate the performance of this target policy is to compute its value function. Therefore, accurately estimating the value function for a target policy becomes important.\n\nOPE is particularly important in safety-critical tasks such as healthcare [2] or in situations where it may be monetarily costly to deploy a potentially poor performing policy such as in recommendation systems [3]. \n\n**Re: state-action input (“we are curious about whether states and actions should be concatenated prior to being fed into the network for representation, or if they should be processed by separate networks and concatenated subsequently.”)**\n\nThis is an interesting question. In our work, we concatenate the state-action pair and feed it directly into the network (i.e., we do not process them separately). This practice is fairly common (eg, [6, 7]). Further investigation, which is beyond the scope of this work, would be required to determine how this alternative approach performs for OPE.\n\n**Re: image states (“We are interested in exploring the effectiveness of the method on the V-D4RL dataset, wherein states are depicted as images”)**\n\nThis is an interesting future direction. In this present version of KROPE, we focussed on illustrating that KROPE improved the stability of OPE theoretically and on D4RL and deepmind control suite environments (across 13 datasets). A next step would be to apply these ideas to visual domains.\n\n**Re: using KROPE for offline control (“could facilitate learning an effective policy” and “demonstrate that the policy learned by the KROPE algorithm”)**\n\nThanks for this suggestion. We also expect KROPE to help in the offline control setting and it is an interesting direction to explore. In the current scope, however, we focussed on OPE. While control is interesting, studying OPE independently is important due to: 1) its practical significance in AI safety and building trust worthy RL agents and 2) prediction is a fundamental part of RL and is worth studying in isolation.\n\nPlease let us know if we have addressed your concerns. If we have, we would greatly appreciate it if you could re-evaluate your review.\n\n---\n[1] Voloshin et al. Empirical Study of Off-Policy Policy Evaluation for Reinforcement Learning. 2021\n\n[2] Gottesman et al. 2018. Evaluating Reinforcement Learning Algorithms in Observational Health Setting\n\n[3] Li et al. 2011. Unbiased offline evaluation of contextual-bandit- based news article recommendation algorithms\n\n[4] Fu et al. 2021. Benchmarks for Deep Off-Policy Evaluation\n\n[5] Uehera et al. 2022. A Review of Off-Policy Evaluation in Reinforcement Learning\n\n[6] Chang et al. 2022. Learning Bellman Complete Representations for Offline Policy Evaluation\n\n[7] Pavse et al. 2023. State-Action Similarity-Based Representations for Off-Policy Evaluation", "claims_and_evidence": "The claims made in the submission are supported by clear and convincing evidence.", "methods_and_evaluation_criteria": "The proposed methods and evaluation criteria do make sense for the problem.", "theoretical_claims": "Partially. We have carefully examined some theoretical proofs, including Lemma 1. However, due to the inconsistency between our research direction and other theoretical knowledge, we are unable to determine whether the theoretical proof is correct.", "experimental_designs_or_analyses": "We have confirmed the rationality of the experimental design in the revised paper. The experiment is mainly divided into two parts. The first part is to verify the stability and qπe consistency of the KROPE algorithm in the Garnet MDPs environment. The second part is to verify whether the KROPE algorithm can bring low and stable MSVE and whether the KROPE algorithm is sensitive to hyperparameters on 4 tasks in DMC and 9 datasets in D4RL.", "relation_to_broader_scientific_literature": "This paper focuses on learning a value network that accurately estimates the expected discounted return for each state. The author enhances the estimation accuracy by learning a more effective representation, setting this approach apart from previous studies like FQE [1]. However, the paper does not further elaborate on the advantages or potential applications of learning such a value network.\n[1] Le, Hoang, Cameron Voloshin, and Yisong Yue. \"Batch policy learning under constraints.\" International Conference on Machine Learning. PMLR, 2019.", "essential_references_not_discussed": "Related works that are essential to understanding the key contributions of the paper are currently cited/discussed in the paper.", "other_strengths_and_weaknesses": "Strengths:\n1.The author introduces a novel approach for evaluating the value function, with experimental results demonstrating its superior accuracy and stability in estimation compared to other baseline methods.\n2.The theoretical proof of this paper is rigorous, and the proof ideas are also clear.\nWeaknesses:\n1.This paper demonstrates that the KROPE algorithm can learn a value network that accurately estimates the expected discounted return for a given state. However, it does not further elaborate on the role or potential applications of this learned value network.", "other_comments_or_suggestions": "Overall, the paper is well-organized, with clear ideas and compelling theoretical and experimental evidence. However, the paper's starting point could be articulated more clearly. For instance, how do value networks further contribute to accurately predicting the expected discounted return from a given state?", "questions_for_authors": "1.We aim to gain insights into the particular implementation of representing the (s, a) pair. Specifically, we are curious about whether states and actions should be concatenated prior to being fed into the network for representation, or if they should be processed by separate networks and concatenated subsequently. \n2.Currently, this work has been focused on the D4RL dataset, which comprises physical states. We are interested in exploring the effectiveness of the method on the V-D4RL dataset, wherein states are depicted as images.\n3.We are curious about the potential applications of value networks when they can accurately estimate the expected discounted return for each state. We believe that such an accurate value network could facilitate learning an effective policy. In this context, we hope the author can provide a comparison between the performance of the policy learned through KROPE and classical offline RL algorithms (such as TD3-BC [1], CQL [2], etc.). If the author can provide the above experimental results and demonstrate that the policy learned by the KROPE algorithm performs comparably to or better than classical offline RL algorithms, we will consider increasing the paper's score.\n[1] Fujimoto, Scott, and Shixiang Shane Gu. \"A minimalist approach to offline reinforcement learning.\" Advances in neural information processing systems 34 (2021): 20132-20145.\n[2] Kumar, Aviral, et al. \"Conservative q-learning for offline reinforcement learning.\" Advances in neural information processing systems 33 (2020): 1179-1191.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {"ethical_review_concerns": "N/A"}}, {"review_id": "IBhj7ilUhF", "reviewer": "Reviewer_YH4r", "summary": "This paper introduces Kernel Representations for Offline Policy Evaluation (KROPE), a kernel-based representation learning algorithm based on bisimulation metric-like ideas. They study a class of representations which emerge as the solution to the representation learning loss, and prove that it has desirable theoretical properties, in particular being stable for off-policy value function learning and is Bellman-complete under additional assumptions. They then evaluate across a range of both tabular and larger-scale environments, and both validate their theoretical results and compare KROPE across a range of baselines for OPE.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for appreciating our work and the clarity of our writing. Thank you for mentioning the strengths of our empirical and theoretical work, especially Theorem 1 and our choice of experiments. We also appreciate your acknowledging the significance of our results and potential for further research\n\n**Theorem 2 is dependent on a very strong assumption (that the reward function is injective)**\n\nThe injective reward assumption simply means that each abstract state-action group will have a distinct associated reward from every other abstract state-action group. While we agree that it is strong, it comes as a tradeoff. Chen and Jiang [1] proved that bisimulation abstractions are Bellman complete. Instead of assuming injective rewards, they assumed that two states were grouped together if each state’s transition dynamics led to next states that are also grouped together (one of the conditions for exact bisimulations). This condition is also considered strict and inefficient to compute [2].\n\nIn our work, we relax this exact transition dynamics equality by considering independent couplings between next state distributions, thereby making the KROPE algorithm efficient to compute. However, the drawback is that preserving distinctness between abstract state-actions may be lost (see page 4 Section 3.1 on remarks on $k^{\\pi_e}$). To ensure the distinctness between state-action abstractions, we assumed that the reward function is injective. This then allowed us to show Bellman completeness in a similar way to that shown in Chen and Jiang [1].\n\nRegarding a weaker version: it may be possible to relax this assumption and consider the error induced in Bellman completeness. Ultimately, what is needed is the ability to preserve distinctiveness between state-action features/abstractions. Ensuring such a property with a weaker condition would be interesting to investigate.\n\nNote that this assumption is only for the BC proof. Theorem 1 does not make this assumption. We will include the rationale behind the injective reward assumption for Theorem 2 in the appendix.\n\n**Definition 2 clarification**\n\nYes, we will add the point that the inner product equals the kernel evaluated for those features to the camera ready. Regarding your other question, each row in the $\\Phi$ matrix corresponds to the feature vector for a state-action pair, so the dependence on $\\mathcal{D}$ means that the state-actions (i.e., the features) are sampled according to their appearance in the batch of data. \n\n**Other inner products/feature maps for KROPE kernel**\n\nThis is a really interesting question that is worth looking into. The features are shaped based on the PSD kernel used. Currently, Definition 2 defines a KROPE representation to be one that satisfies the relationship in Definition 2 in terms of linear dot products (which is a function of $k_1$, which is PSD (see Lemma 4 in Appendix)), which are easier to reason about. That said, Definition 2 can be broadened to just include short-term ($k_1$, which is PSD) and long-term similarity, where long-term similarity is determined by, say, a Gaussian kernel. This could also be a valid KROPE representation, but unclear if it will also be a stable representation as defined in our work (in terms of the spectral radius of the features).\n\n**Deferring background in Sections 1/2 to Appendix and moving Appendix C.3 up**\n\nWe will definitely re-review Sections 1 and 2 and see what we can defer to the Appendix. We do believe that it is important to be complete, even if slightly redundant for readers who may already be familiar with the background knowledge. If possible, we will then move Appendix C.3 to the main text since we also agree that Appendix C.3 brings useful insights regarding the potential instability of KROPE.\n\n--- \n[1] Chen and Jiang. 2019. Information-Theoretic Considerations in Batch Reinforcement Learning\n\n[2] Castro. 2020. Scalable methods for computing state similarity in deterministic Markov Decision Processes", "claims_and_evidence": "All claims made in the submission are supported by proper evidence.", "methods_and_evaluation_criteria": "The methods and evaluation done make sense for the problem at hand.", "theoretical_claims": "I checked all proofs, and did not find any issues.", "experimental_designs_or_analyses": "The experimental design and validity seems good to me (I only very quickly skimmed the code provided).", "relation_to_broader_scientific_literature": "In a sense the key contributions of this paper can be seen as extending the analysis/understanding of bisimulation-based representation learning to stability-type results, which I think is an important contribution on its own, and I expect it to lead to further research.", "essential_references_not_discussed": "I don't believe any essential references are not discussed.", "other_strengths_and_weaknesses": "**Strengths**\n- The application of stability analysis to bisimilarity metric-type representations is a nice perspective and I expect it to be built upon in the future. Additionally the proof of Theorem 1 appears rather general (I think it should apply with at least any reasonable choice of immediate similarity kernel).\n- The paper is well-written and easy to follow. \n- The choice of experiments in Section 4.2. nicely complement the theoretical results of section 3.\n\n**Weaknesses**\n- Theorem 2 is dependent on a very strong assumption (that the reward function $r^\\phi$ is injective), without any discussion around the assumption or what it may entail. I think that this assumption is violated in almost any non-contrived setting I can think of (from large-scale complex ones like Atari/MuJoCo to gridworlds, maze-based, CartPole, etc), which I believe limits the impact of the result.", "other_comments_or_suggestions": "- I quite like the introduction and discussion of the KROPE representation $\\Phi$ in Definition 2. To reassure the reader, can you add perhaps a minor which states that if $\\Phi$ is a KROPE representation, then the inner product between two state-action pairs is equal to the kernel evaluated at them?\n- It could strengthen the paper to shorten sections 1 & 2 and move some of Appendix C.3. to the main text. There is quite a bit of background/discussion before the novel contributions begin (halfway through page 4), and the main takeaways should be in the main body (takeaway #3 is currently in the appendix).\n- In section C.2., in the description of ROPE, minor nit: \"Its additional learning rate is the output dimension of $\\phi$.\"", "questions_for_authors": "- I'm a bit confused by Definition 2 -- in particular what does $\\mathbb{E}_{\\mathcal{D}} [\\Phi\\Phi^T]$ represent for a matrix $\\Phi \\in \\mathbb{R}^{|\\mathcal{X}| \\times d}$ (what depends on $\\mathcal{D}$)?\n- Similar to my first comment in the previous section -- is it possible for there to be a feature mapping $\\Phi$ such that $\\langle \\phi(x,a), \\phi(y,b)\\rangle = k^{\\pi_e}(x,a;y,b)$ but $\\Phi$ is not a KROPE representation?\n- Can there be a weaker statement similar to Theorem 2, without assuming the fact that $r^\\phi$ is injective (likely the statement would depend on $r^\\phi$, and the current theorem might appear as a special case).", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "2gX2lcHEa7", "reviewer": "Reviewer_KFHs", "summary": "The paper tackles the field of offline policy evaluation (OPE) and addresses methodology to find good state-action-pair representations. It introduces a kernel based state-action representation and gives theoretical properties for it. It then presents experimental results of the introduced KROPE method on different benchmarks.\n\n## update after rebuttal\nScore increased. See rebuttal comment.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for your thoughtful comments and feedback. Thank you for acknowledging that the work was an interesting read with valuable contents.\n\nYour comments are helpful in making our paper precise. We do believe, however, that these adjustments involve minor reframings/edits. We address your concerns below.\n\n**Claims and evidence**\n\nAll the concerns in this subsection are easily addressable since we acknowledge the shortcomings of KROPE explicitly in Section 5 and Appendix C.3.1. All changes are sentence-level clarifications and involve bringing discussions from Section 5 earlier in the paper.\n\n**Re: “’Therefore, it is critical to stabilize value function learning by explicitly shaping the state-action representations.’ is not supported by the experiments presented.”**\n\nThat statement is a general desired property we want from representation learning algorithms for offline policy evaluation. With regards to KROPE, we already discuss KROPE’s inability to stabilize in all settings in Section 5 and Appendix C.3.1. That said, we understand the concern and will discuss the limitation earlier in the paper.\n\n**Re: “\"In this work, we investigate how to explicitly learn state-action representations to stabilize value function learning.\" is misleading.”**\n\nThe statement means we show how one may go about learning KROPE representations. But we understand that this can be potentially confusing. We do not investigate different approaches to shape the representations. We propose only one way to do so, and we will clarify that.\n\n**Re: “The fourth and fifth claim are slightly misleading.\"**\n\nWe will adjust contribution 4 based on our discussion in Section 5 and Appendix C.3.1 to state that it does indeed improve the stability of OPE compared to 7 other baselines on 10/13 datasets (note: 10 does not include just highlighted blue errors in table 1). Regarding Contribution 5, we will try to move it to the main paper.\n\n**Re: “The Takeaway #1 is not true in general and thus, it needs clarification.”**\n\nWe understand the point, we will rephrase the takeaway to explicitly say that the statement is true under theoretical assumptions.\n\n**Re: “Takeaway #2 is not true in general. At least from my understanding this is not given in general.”**\n\nTakeaway #2 is true in an aggregate sense across datasets (10/13) and compared to the 7 other baselines. But we understand the point. We will state that KROPE improves the chances of stability for OPE compared to other representation learning baselines.\n\nThe more accurate characterization of our work is the following: *our theoretical results prove that if state-action features satisfy the KROPE relation (Definition 2), then they will lead to stable value function learning. Practically, since KROPE relies on a semi-gradient method (like DQN/fitted Q-evaluation; see Section 5 and Appendix C.3.1), the algorithm may still lead to divergence. Empirically, KROPE improves stability compared 7 other baselines, and leads to stable and accurate OPE in 10/13 cases (note: 10 does not include just highlighted blue errors in table 1). Therefore, KROPE improves upon other baselines in learning stable and accurate representations for OPE.*\n\nWe will accordingly discuss this earlier in the paper instead of waiting till Section 5 and modify the claims/takeways.\n\n**Re: “Table 1 is not informative. This needs thorough revision. Stating numbers in a table is a good idea in general, but the mean +/- the standard error should be the standard to do so.”**\n\nThanks for your feedback. However, the table does include information for statistical rigor. Each value is the MSE over 20 trials and the range shows the 95% confidence interval, which are all important measures of statistical rigor (see caption of Table 1) [1]. We also include learning curves in Figures 6/7 (appendix).\n\n**Deadly triad reference**\n\nThanks for the reference. We will include this.\n\n**Formatting issues: ICML compliant, ylabels, line 190 format, FQE upper case**\n\nThanks and we will update this and the graphs to be clearer.\n\n**Meaning of continuity**\n\nWe will add clarification on this and we refer the reviewer to the work of Le Lan et al. [2] for continuity of metrics in the context of RL (see their Figure 1 in Le Lan [2]). Briefly, ideally, states that have similar values are close to each other in feature space. This statement is true for general machine learning too: inputs with similar outputs should ideally be close to each other in feature space.\n\nOnce again thank you for making our work more precise. We believe the main concern on claims is easily addressable. Please let us know if we have addressed your concerns. If we have, we would greatly appreciate it if you could re-evaluate your review.\n\n---\n[1] Patterson et al. 2024. Empirical Design in Reinforcement Learning.\n\n[2] Le Lan et al. 2021. Metrics and continuity in reinforcement learning.", "claims_and_evidence": "- The statement in the abstract \"Therefore, it is critical to stabilize value function learning by explicitly shaping the state-action representations.\" is not supported by the experiments presented. It is shown in the paper that for some cases value function learning was successful without learning any state-action representation.\n\n- The statement \"In this work, we investigate how to explicitly learn state-action representations to stabilize value function learning.\" is misleading. The work introduces one way to do so, but does not investigate different approaches in my understanding. \n\n- There is a list of 5 contributions under the headline \"Can bisimulation-based representation learning stabilize offline value function learning?\". The fourth and fifth claim are slightly misleading. The fifth claim can not be part of the main paper since its contents are inside the appendix. The claim that KROPE representations can be successfully used in OPE is shown in the paper. To the best of my understanding the evidence is not convincing that KROPE representations always lead to more stable and accurate offline value function learning. \n\n- The Takeaway #1 is not true in general and thus, it needs clarification. \n\n- Takeaway #2 is not true in general. At least from my understanding this is not given in general.", "methods_and_evaluation_criteria": "The proposed methods and benchmarks seem reasonable for the scope of the work.", "theoretical_claims": "- All theoretical claims are based on the standard coverage assumption that gives a non-zero probability of state-action pairs appearing in the offline dataset for finite state spaces. \n- The LSPE stabilization looks fine to me.\n- The Bellman completeness proof looks fine as well.", "experimental_designs_or_analyses": "Table 1 is not informative. This needs thorough revision. Stating numbers in a table is a good idea in general, but the mean +/- the standard error should be the standard to do so.", "relation_to_broader_scientific_literature": "The paper contributes to the OPE literature.", "essential_references_not_discussed": "The reference to Van Hasselt et al., 2018 that coined the term \"deadly triad\" which is used throughout the paper is missing.", "other_strengths_and_weaknesses": "- The paper lacks clarity in some parts of the submission. \n\n- Formatting is not ICML compliant, e.g., the uppercased abbreviations. \n\n- General formatting needs streamlining. \n\n- In 209 right side it says that continuity is important. But why?\n\n- The ylabels in several plots are very hard to read and leave room for improvement. \n\n- Overall this work is an interesting read with valuable contents in need of a major revision and re-evalution of the made claims.", "other_comments_or_suggestions": "- Paragraph title \"Remarks...\" in 190 is weirdly formatted, as well as subsection title of 4.2\n- fitted *Q*-evaluation should have at least an uppercased Q", "questions_for_authors": "- Can you elaborate on my concerns regarding the claims made in the paper?\n\n- In the experiments in the appendix, you contradict the statement that KROPE always leads to more stable results, or do I missinterpretate the presented results?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper proposes KROPE, a method for offline policy evaluation via improved learned representations, based on prior work on bisimulation metrics. The authors provide theoretical results demonstrating the increased stability for least-squares policy evaluation when integrated with KROPE, that KROPE's representations are Bellman complete, and empirically demonstrate the stability and performance of KROPE.

Although there were some initial reservations on the impact and clarity of the work, the authors were able to clarify all issues during the rebuttal. All reviewers are now supportive of acceptance, as am I, after having reviewed the discussion and paper itself.
Decision: Accept (poster)