Stable Offline Value Function Learning with Bisimulation-based Representations

Thank you very much for your thoughtful comments and feedback. Thank you for acknowledging that the work was an interesting read with valuable contents.

Your comments are helpful in making our paper precise. We do believe, however, that these adjustments involve minor reframings/edits. We address your concerns below.

Claims and evidence

All 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.

Re: “’Therefore, it is critical to stabilize value function learning by explicitly shaping the state-action representations.’ is not supported by the experiments presented.”

That 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.

Re: “"In this work, we investigate how to explicitly learn state-action representations to stabilize value function learning." is misleading.”

The 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.

Re: “The fourth and fifth claim are slightly misleading."

We 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.

Re: “The Takeaway #1 is not true in general and thus, it needs clarification.”

We understand the point, we will rephrase the takeaway to explicitly say that the statement is true under theoretical assumptions.

Re: “Takeaway #2 is not true in general. At least from my understanding this is not given in general.”

Takeaway #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.

The 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.

We will accordingly discuss this earlier in the paper instead of waiting till Section 5 and modify the claims/takeways.

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.”

Thanks 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).

Deadly triad reference

Thanks for the reference. We will include this.

Formatting issues: ICML compliant, ylabels, line 190 format, FQE upper case

Thanks and we will update this and the graphs to be clearer.

Meaning of continuity

We 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.

Once 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.

---

[1] Patterson et al. 2024. Empirical Design in Reinforcement Learning.

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

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

Theorem 2 is dependent on a very strong assumption (that the reward function is injective)

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. 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].

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 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].

Regarding 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.

Note 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.

Definition 2 clarification

Yes, 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.

Other inner products/feature maps for KROPE kernel

This 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).

Deferring background in Sections 1/2 to Appendix and moving Appendix C.3 up

We 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.

---

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

[2] Castro. 2020. Scalable methods for computing state similarity in deterministic Markov Decision Processes

We thank the reviewer for acknowledging that our empirical results and theoretical results are rigorous, and that the paper is clear and well-organized.

Below we address your concerns.

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.”)

In 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.

OPE 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].

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.”)

This 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.

Re: image states (“We are interested in exploring the effectiveness of the method on the V-D4RL dataset, wherein states are depicted as images”)

This 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.

Re: using KROPE for offline control (“could facilitate learning an effective policy” and “demonstrate that the policy learned by the KROPE algorithm”)

Thanks 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.

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.

---

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

[2] Gottesman et al. 2018. Evaluating Reinforcement Learning Algorithms in Observational Health Setting

[3] Li et al. 2011. Unbiased offline evaluation of contextual-bandit- based news article recommendation algorithms

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

[5] Uehera et al. 2022. A Review of Off-Policy Evaluation in Reinforcement Learning

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

[7] Pavse et al. 2023. State-Action Similarity-Based Representations for Off-Policy Evaluation

Thank you for acknowledging that our empirical results and theoretical results are solid. Below we address the concerns you raise.

Use of Mean Absolute Error vs. MSE

This 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].

Numerical precision in Table 1

After 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.

Relations to Pavse et al. 2023

This 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.

Instead 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.

In more detail, the main differences are:

1. 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.
2. 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].

We will make these differences explicit in the Appendix in the camera ready where we discuss the different baselines

Including learning curves

Thanks for the suggestion. These are included in the Appendix (Figure 6 and 7) due to lack of space in the main paper.

KROPE does not always converge

Yes, 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.

Briefly, 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].

We 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).

Pseudo-code

Appendix (page 15) includes one, but we will make it clearer.

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.

---

[1] Voloshin et al. Empirical Study of Off-Policy Policy Evaluation for Reinforcement Learning. 2021 [2] https://github.com/google-research/google-research/blob/master/policy_eval/q_fitter.py#L101 [3] Chaudhari et al. 2024. Abstract Reward Processes: Leveraging State Abstraction for Consistent Off-Policy Evaluation [4] Liu et al. 2018. Breaking the Curse of Horizon: Infinite-Horizon Off-Policy Estimation [5] Thomas et al. 2016. Data-Efficient Off-Policy Policy Evaluation for Reinforcement Learning [6] Hanna et al. 2019. Importance Sampling Policy Evaluation with an Estimated Behavior Policy [7] Sachdeva et al. 2023. Off-Policy Evaluation for Large Action Spaces via Policy Convolution [8] Lehnart and Littman. 2020. Successor features combine elements of model-free and model-based reinforcement learning [9] Bhatt et al. 2024. CrossQ: Batch Normalization in Deep Reinforcement Learning for Greater Sample Efficiency and Simplicity [10] Feng et al. 2019. A Kernel Loss for Solving the Bellman Equation