Stable Offline Value Function Learning with Bisimulation-based Representations

Large-scale environments, image, noisy states: This is a very interesting future direction for us to pursue! Thank you for the suggestion. While we do acknowledge that applying our ideas to image-based domains and states with noisy distractors would be very interesting, we believe that the lack of it does not detract any significance from the work since we still evaluate the algorithms in various continuous state-action space environments with a high variety in datasets. In our work, we investigated the performance of the various algorithms on challenging datasets (7 environments, 13 datasets) that are commonly used for offline policy evaluation such as various DM control environments and D4RL datasets (see Appendix). Given that the scope of the paper is focused on the stability properties of KROPE-like algorithms, we believe that the current paper suffices as a step toward that understanding. We also note that while additional steps may be needed to apply KROPE to the alternative settings, there is no fundamental bottleneck for adapting our kernel-based method to images or distractors since our kernel is defined on the feature space.

Discussion about the kernel and scaling: Below we make remarks relevant to our work but we also encourage the reviewer to see [4] for additional benefits of the kernel in terms of theoretical analysis and scaling to complex domains.

• Benefit of the kernel: The use of the kernel enabled us to concretely talk about positive definiteness when measuring similarity. Intuitively, this property gives insights into the rank of the feature matrices $\Phi$ involved, which gives us a sense of how distinct are the abstract/latent state-actions from each other. When designing KROPE and as done by [4], we are able to characterize the similarity measure (i.e., the KROPE kernel) as a positive (semi-)definite kernel. We used this property in proving the stability of KROPE representations (Theorem 1, see Appendix). Moreover, we also used this property to establish the relation to Bellman completeness (Theorem 2). In short, kernels give us a way to directly control the spectral properties of the dynamics. It may be possible to still prove these same properties through other means, but we found the tool of kernels to be helpful in our case.
• Scaling issues of kernel: There may be confusion on how the kernel is used, which casts doubts on its ability to scale. The kernel is applied on feature/vectors (and would not be applied directly to the raw image). Moreover, the kernel is not fixed but it is learned since it is a function of $\phi$ which is learned. In the feature space, the form of the kernel (linear dot product) is fixed (as is done by [4]), but the objective enables appropriate features to be learned such that the learned kernel captures the similarity relation we care about. Consequently, the use of kernels does not automatically preclude the use of our method in image-based domains. While additional steps may be needed to apply KROPE to the alternative settings, there is no fundamental bottleneck for adapting our kernel-based method to images or distractors since our kernel is defined on the feature space.

We now address specific points made in the review.

Weaknesses

1. Please see the general comment above of “Goal of our work” and “comparison to prior work”.
2. Please see above “Benefits of the kernel”.
3. Please see above “Large-scale environments”.
4. Please see “Goal of our work”, “comparison to prior work” (“Relations to dynamics model”).

Questions

1. This is a great question and we will clarify this in the camera ready. Theorem 2 draws a relation to Bellman completeness. Intuitively, Theorem 2 tells us that KROPE representations are expressive enough to predict the reward and latent features at the next step. While KROPE does not explicitly model the latent dynamics of the MDP, its features are expressive enough to capture them. We refer the reader to [1, 2, 3] on prior works connections of bisimulations, forward dynamics, and Bellman completeness, where they are all essentially different approaches for getting at the same idea. Please also see the general comment above of “Comparison to prior work” (“Relations to dynamics model”).
2. Thank you for the suggestion. Please see our response to “comparison to prior work” (“Other bisimulation methods”).
3. Please see comments on “Large-scale environments”.
4. This is an interesting direction to consider in future work. Given that the KROPE objective specifically tries to capture similarity between state-actions by considering short-term and long-term behavior, where this behavior is only a function of the reward function and transition dynamics, we also suspect that KROPE will discard irrelevant noise that does not impact the reward function and transition dynamics. That is, it will learn feature representations (and corresponding kernel) only relevant to capture similarity in action-values. Note that the kernel in KROPE is actually being learned. We do bias the form (using a linear dot product as done in prior work [4]), but the kernel, which is a function of the learnable representations, is being learned so we expect the features (and indirectly the kernel) to ignore irrelevant noise.
5. We would not apply the kernel to the images directly. We suspect that we would have a typical encoder network that projects the image into a vector, and then we would measure similarity on that vector. [4] applies this kernel idea to Atari and shows improvement, so we would expect similar results to hold. However, as mentioned earlier, our focus is showing the stability of these representations, which is a result that theoretically extends to images or state-based environments as long as the features satisfy the relationship in Theorem 1. Please also see above “Discussion about the kernel”.
6. Good question and this is related to your first question. Please see our response to Q1 above. We will make this connection clearer in the camera ready.
7. The primary goal is to show that if we take a kernel perspective (similarity and positive definiteness) to representations and if they satisfy the KROPE relationship (a bisimulation-based relation), the representations are stable and accurate for offline policy evaluation setting as we state in the introduction. This stability property of the metric-based algorithm was missing from the literature and our work addresses it. We investigated this specific stability-related question in our experiments on 13 datasets, 7 high dimensional state-action environments, and through Theorem 1. As an auxiliary result, we also prove Theorem 2 which draws a relation between KROPE and Bellman completeness, which is another important condition typically considered for stability. Please also see “Goal of our work” above.

[1] Successor Features Combine Elements of Model-Free and Model-based Reinforcement Learning. Lehnert et al.

[2] A Note on Loss Functions and Error Compounding in Model-based Reinforcement Learning. Jiang.

[3] Information-Theoretic Considerations in Batch Reinforcement Learning. Chen and Jiang.

[4] A Kernel Perspective on Behavioural Metrics for Markov Decision Processes. Castro et al.

[5] Scalable methods for computing state similarity in deterministic Markov Decision Processes. Castro et al.

[6] Learning Invariant Representations for Reinforcement Learning without Reconstruction. Zhang et al.

[7] MICo: Improved representations via sampling-based state similarity for Markov decision processes. Castro et al.

[8] Benchmarks for Deep Off-Policy Evaluation. Fu et al.

[9] Empirical Study of Off-Policy Policy Evaluation for Reinforcement Learning. Voloshin et al.

[10] A Review of Off-Policy Evaluation in Reinforcement Learning. Uehara et al.

[11] Learning Bellman Complete Representations for Offline Policy Evaluation. Chang et al.

Thank you for appreciating the clarity of our paper and our process of introducing the KROPE objective. We also thank the reviewer for their detailed questions, comments, and feedback. We first make some general statements regarding the overall comments and then address individual comments/questions below.

General comments

In order to clarify any possible confusion regarding the contributions and focus of our work, we make the following general comments.

Goal of our work: The primary objective of our work is to show that the general class of bisimulation-based metric learning can learn representations that are indeed useful for accurate and stable offline policy evaluation (Lemma 3 and Theorem 1). Our work provides further evidence that bisimulation-like metric learning algorithms have desirable properties. Our goal is not to show a bake-off that our bisimulation-based metric learning method is better than other algorithms within this class, nor is to say that KROPE is trying to learn the true underlying latent dynamics.

Importance of studying only policy evaluation (reference to comment: “I do not think the discussion of policy evaluation and stability here is enough to show that KROPE can learn good policies for control.”): We want to mention that we focus only on offline policy evaluation (OPE) only and do not make claims about policies for control. Furthermore, we want to emphasize that studying evaluation and prediction algorithms only is still a significant contribution since good prediction is often a precursor for understanding improved control algorithms. While control is important, OPE is important [8, 9, 10] due to: 1) its practical significance in AI safety and building trust worthy RL agents and 2) it is a clean problem for studying how to make better predictions without dealing with the exploration problem or a changing target policy. Of course, exploration and changing target policy are important but there is value in removing those aspects of the whole RL problem for study. As such, designing control algorithms for offline RL is not the focus of this work, but is a very interesting future direction.

Comparison to prior work: We want to address the following:

• “several prior works have done this already, whether implicitly or explicitly showing the policy evaluation measure”: It would be great if the reviewer could link us to the mentioned references so we can give a more thoughtful response on this issue. However, at a high-level: just because good control algorithms may hint at good policy evaluation, it does not mean we understand the policy evaluation algorithm well, nor is it clear if the evaluation algorithm has desirable properties such as stability. Our work attempts to understand this important piece from the view of stability and accuracy.
• Other bisimulation methods: As mentioned in the introduction, KROPE is a representative of the class of bisimulation-based metric learning algorithms. Our work provides evidence that KROPE-like algorithms can have favorable stability and accuracy properties. Furthermore, our aim is not to show our metric learning algorithm is better than another one. As such, we compare KROPE to other non-bisimulation based metric learning algorithms. While prior has introduced different bisimulation metric learning algorithms, they have typically been limited. For example, they either assume deterministic [5] or Gaussian transition dynamics [6], or are more difficult to analyze theoretically [7]. Castro et al. [4] later introduced a metric learning algorithm for general stochasticity in the dynamics, and we built upon it to illustrate our point on the stability and accuracy of this class of methods.
• Relations to dynamics model: This is a very good point. As you mention, it is nice to theoretically show the relationship between KROPE and algorithms that model forward dynamics. This is actually what Theorem 2 tells us. Bellman completeness essentially tells us how well the representation models the dynamics of the MDP in the latent/abstract MDP. Theorem 2 draws the connection between KROPE representations and modeling latent dynamics of the environment. We will make this point clear in the camera ready. Empirically, we compare against BCRL as the representative method for the class of methods that capture the dynamics of the environment since it was introduced as a strongly competitive algorithm specifically for our problem setting of offline policy evaluation [11].

Dear uF25, Please let us know if we have addressed your concerns. Your comments are valuable and we would like to clarify any potential misunderstandings to strengthen our submission. Thanks!

Dear uF25, to ensure there is ample time for discussion before the deadline it would be great to get your thoughts on our response to see if we have addressed your concerns. Thank you!

Thank you for appreciating our work, the rigor of our theoretical analysis, and thoroughness of our experiments.

We address your remarks below.

Weaknesses

In practice, it is common to assume that $\pi_e$ is available to us to evaluate it [3]. This is how algorithms such as FQE, BCRL, and KROPE typically work. The challenge, however, is knowing the value of $\pi_e$ since that is what we want to evaluate. In practice, this is unknown and KROPE does not depend on it. But we compute it for the sake of empirical analysis by running Monte Carlo rollouts to get an unbiased estimation of its performance (see Appendix C.1) and then computing the error between this value to the output of KROPE and LSPE.

Questions

1. Refer above. Typically, in OPE, we will always know the policy that we want to evaluate (i.e. $\pi_e$) [3].
2. In all cases, we do have an offline dataset, which is generated by a (collection) behavior policy (policies). We use this dataset for offline value prediction. We refer the reviewer to Appendix C.1 for details on how we generated the dataset.
3. This is an interesting question. In our work, we concatenate the state and actions and feed them into the encoder. This is similar to the BCRL paper [1]. However, you are right that it is also possible to have separate networks. One work that adopts this approach is [2], where they have separate networks but jointly train the encoders based on a homomorphism loss. We conjecture that having separate networks has benefits of capturing more precise symmetries between states and between actions, but it does make the optimization process more difficult since it introduces more hyperparameters and networks. It would be interesting to investigate the performance gains or losses of such an architecture. And no there will be marginal difference in applying KROPE to image-based vs. state-based states. For state-based environments, the state can be directly fed into KROPE. For image-based, we would typically use a CNN encoder to generate vector features, which can then be fed into KROPE.

[1] Learning Bellman Complete Representations for Offline Policy Evaluation. Chang et al.

[2] Continuous MDP Homomorphisms and Homomorphic Policy Gradient. Rezaei-Shoshtari et al.

[3] Benchmarks for Deep Off-Policy Evaluation. Fu et al.

Thank you for feedback, your kind words in appreciating the clarity of our work, and acknowledging the thoroughness in our experiments in demonstrating KROPE’s utility.

We address your comments below.

Weaknesses

Regarding the narrow scope: We want to emphasize that the stability of temporal difference learning algorithms is a fundamental problem in RL. Thus far, it was unclear whether bisimulation-based metric learning algorithms (instances of TD-like algorithms) produced representations that were indeed stable since prior work only focussed on generalization properties. Our work fills this gap, and through our theoretical and empirical analysis, shows that these algorithms do indeed learn stable and accurate representations for OPE.

Your comments regarding the clarity of the graphs make sense. We will update this for the camera ready. Thank you for the feedback.

Questions

• This completely makes sense. We will make this clear. To answer your questions briefly: a) the chosen notion of generalizability is quite common (see Section 4.2 in [1] and Figure 1 in [2]). Intuitively, if state-actions with similar representations had different action-values, then that would hurt performance since the action-value function would get confused on what it should output for identically looking inputs, so they should not have the same representation (example of bad generalization). b) Bellman completeness is commonly assumed in theoretical works for data-efficient policy evaluation [3]. Intuitively, it means that the representation captures the dynamics of the MDP perfectly in the latent MDP, which is naturally sufficient for accurate evaluation. c) It may not be reasonable in practice, but as shown in the experiments that does not detract from the utility of KROPE. Theorem 2 is to show that under some circumstances KROPE is related to this alternative notion of stability.
• Thanks for the suggestion. We will make this clearer. We refer the reviewer to Figure 1 in [4] for a nice illustration of the general idea of representation for value function learning. But briefly, yes, it is what you suggested of $q(s,a)$ vs. $q(\phi(s,a))$.
• Yes, since the policy is stochastic in those environments, Lemma 3 tells us that two state-actions with similar representations may have different action-values, which would hurt performance.
• While BCRL does optimize for condition number, it depends on a scalar regularizer coefficient ($\lambda$) [3]. So it depends on hyperparameter tuning. KROPE, on the other hand, is able to get favorable condition numbers without introducing any new objective and hyperparameters. We did a hyperparameter sweep for BCRL and reported the best result we got in terms of OPE error. At the expense of higher error for BCRL, it is possible to lower the condition number by increasing $\lambda$.
• We will make the description of FQE clearer in the camera ready.
• This makes sense too.
• Thanks for the suggested edits! We will revise these accordingly.

[1] Learning Dynamics and Generalization in Reinforcement Learning. Lyle et al.

[2] Metrics and continuity in reinforcement learning. Le Lan et al.

[3] Learning Bellman Complete Representations for Offline Policy Evaluation. Chang et al.

[4] On the Generalization of Representations in Reinforcement Learning. Le Lan et al.

We sincerely appreciate your very kind words and appreciate that you liked our work. We are glad that you found the paper clear in its contributions and thorough in the empirical and theoretical analysis. We agree with you that our work will be of interest to a general audience.

We address your remarks below.

Weakness

This is a valid remark. Thank you for the suggestion. As you point out, most prior work already focused only on the generalization properties of bisimulation-based methods. However, there was still a gap in the literature of whether this class of methods produced representations that were stable and accurate for OPE. As such, our focus was mostly on addressing this gap. Our work showed that these representations are indeed stable and accurate for OPE through theoretical and empirical analysis.

Questions

Thanks for all the editing suggestions. We will make these changes for the camera ready.

Regarding Q4: This is a fair point. We made the statement based on what we observed in those experiments, but we will clarify this general speculative point in the camera ready. Thank you.

Regarding Q5: This is a good question. Your suggestion of studying the difference between LHS and RHS is an interesting first step and worth investigating in future work. We think the actual analysis may be more complicated since one way to measure the LHS/RHS difference is the loss function, but as noted by prior work loss functions based on target networks and bootstrapping (such as the TD/FQE loss) tend to poorly correlate with metrics we actually care about such as value error [1].

Regarding Q8: Ah yes. We meant that it (gray line) produces a lower condition number than BCRL-NA and FQE (orange and yellow). We will clarify this in the camera ready.

[1] Why Should I Trust You, Bellman? The Bellman Error is a Poor Replacement for Value Error. Fujimoto et al.