Orthogonalized Estimation of Difference of Q-functions

Thank you for your thorough review and recognizing the strengths of the paper, especially the fact that it doesn’t introduce dependence on unstable inverse behavior policy weights, which we think has hampered adoption of orthogonal ML or double ML estimation in offline RL in the past.

We’d like to respond with clarifications to your concerns on weaknesses below. We have already included some of your great suggestions in the updated revision but have given them all careful consideration.

Weaknesses:

does not clearly highlight … In lines 271-273 and lines 373-377 we have highlighted where our analysis is similar to prior work, and where the novelty and technical challenges arose.

We wish to emphasize our work as a methodological contribution with rigorous theoretical guarantees. Our primary goal was to develop orthogonalized estimators that address statistical challenges and adapt to structure. The essence of our contribution lies in deriving a new estimand and orthogonal loss function.

While our analysis process aligns with the high-level framework of orthogonal ML, what is new to our analysis is confirming that the functional derivative of our loss function wrt to the nuisance functions is 0, as shown in Lemmas 4-5. We will add this to the exposition of Thm. 2. This leads to a distinct error decomposition, although we believe the analysis remains fundamentally similar. The analysis framework for all orthogonal statistical learning is similar: derive an estimator that is Neyman orthogonal to nuisance functions, i.e. its functional derivative wrt nuisance functions is 0. Then study the von Mises expansion of the estimator: establish that the second-order term is a higher-order remainder and leverage the error decomposition to conclude the nature of product-rate improvements. The last part appears in our new Assumption 5.

To clarify our objectives, we aimed to create a new estimand and estimator to tackle specific problems, such as applying MI regularization in empirical RL on the difference-of-Q functions. We chose to modify the estimand as an elegant solution for adapting to sparse structures, rather than developing various methods for different structural assumptions, as illustrated in our experiments. This required us to develop and analyze a novel approach for estimating the difference-of-Q functions, which we believe is of independent interest.

However, the challenges for our goals come up in establishing that orthogonal estimation implies convergent policy optimization, which is more difficult because each policy optimization step in principle defines a new set of nuisance functions. We do this with a mix of conceptual (our new estimator) and technical innovations (our induction analysis with margin) that are specific to our task. We explicitly designed our estimation framework to sidestep most of these challenges and moved this discussion of our conceptual innovations earlier in the related work, see lines 100-105 in revision, to make this clearer.

Thanks for the review and the helpful comments. We have made major changes to the presentation in the updated revision and believe these should address your comments re: presentation.

Please let us know if there is anything else we can clarify.

Responses to comments:

• 1. Yes, that is a great summary of the results (indeed we are presenting somewhat “standard” results in the realm of orthogonal ML). In lines 317-319 we had some discussion of the qualitative takeaways from the standard orthogonal estimation results but we realize this got buried in the presentation of the analysis. In the revision, we have overhauled the presentation along these lines by refactoring the product-error nuisance function estimate and focusing on presenting a simpler version of Theorem 3 which preserves the main qualitative takeaway. We also proceed under the well-specified setting for statements on rates in a new Assumption 6.
• 2. “Margin assumption … one can use to perhaps break this dependence” Ultimately this is what happens in the policy optimization result. At every timestep, we re-run the estimation approach to “peel off” estimation of the time-t difference-of-Q function. This prevents the use of the margin assumption from exponentiating convergence over multiple time-steps. On the other hand, we use an inductive argument to show that the dependence on estimated policy arises only in $Q$ functions and is higher-order via the margin assumption at every timestep. Ultimately we don’t pursue the tightest analysis because we ultimately seek to conclude that the error from policy-dependent nuisance functions is higher-order, rather than accumulating over the horizon. The analysis is sufficient to establish that our analysis translates to faster policy optimization, too.

Our innovations are moreso conceptual and technical innovation via the new estimation framework. These issues of policy-dependent nuisance functions have been around for a while in dynamic treatment regimes (see discussion in Zhang et al; they ultimately use a heuristic localization). Unlike Zhang et al, which incurs dependence on the estimated policy via the definition of the population estimand (they optimize the DR policy value estimate itself), the estimated future policy only affects our convergence of the $Q$ function in our time-t loss function. And, unlike Zhang et al, we consider greedy policies with respect to a potentially rich difference-of-Q function, further avoiding issues of policy-dependent nuisance functions. We will add discussion on this point in the appendix.

B. Zhang, A. A. Tsiatis, E. B. Laber, and M. Davidian. Robust estimation of optimal dynamic treatment regimes for sequential treatment decisions. Biometrika, 100(3):681-694, 2013.

• 3. “I don't think it is possible to estimate Q functions, essentially at all, without such assumptions in the offline RL setting, see e.g., [1]. So where does such an assumption come into the analysis?”

We certainly aren’t making any claims that it’s possible to estimate $Q$ functions without the structural assumptions that appear previously in the literature — the assumption comes into the analysis via the assumptions on product of nuisance function estimation rates. (In the revision we refactored this out into its own assumption, A5). If there is no structural condition like “bellman completeness”, then we expect $Q$ estimation to also fail, meaning A5 would not be satisfied. But because we only need the rate of convergence of $Q$ function, we are fairly agnostic about which version of Bellman completeness is (implicitly) implied via requiring consistent $Q$ estimation. Nor does estimation of a well-specified difference-of-Q function imply estimation of the policy value $E[V^\pi(s)]$ in general, so we certainly aren't circumventing any such lower bounds on policy value estimation.

We added additional discussion in lines 306-311 in the revision, summarizing this response.

Re other connections to offline RL theory – in lines 253-257 of the revision (also in the included submission) we discussed the connection to strong concentrability assumptions appearing in the offline RL literature. We do acknowledge that the offline RL literature has developed methods under weaker versions of this assumption; we view analogous improvements for our method as interesting directions for future work.

Lastly, our overall approach using orthogonal ML generally implies that offline RL theory results improving convergence results for $Q$ estimation would apply directly via the black-box nuisance estimation assumption (A5).

Thank you for your review and recognizing the significance of the contributions and extensive results! We have incorporated your suggestions for presentation into the revision and hope you find the revision satisfactory. In general these were all minor changes in exposition and presentation that did not change the substance of the analysis.

Please let us know if you have any remaining questions for discussion or if we can clarify anything else.

Weaknesses:

• “If requirement on nuisances rates are the same” Thanks for pointing this out, indeed they rely on a high-level assumption on quantification of the product error nuisance rates. Thanks for the suggestion, we have refactored this out as an additional Assumption, explained the assumption concisely, and now Thm. 3 repeatedly invokes the assumption for different timesteps/policies.
• Thm. 3 is long -- See response to all reviewers. We opted for presentation of the main high-level result in the main text via some simplifications/weaker rates.
• Definition/assumption/theorem -- we have added additional expositions throughout after introducing these technical statements.

Questions:

1. We don’t require them to increase, but this is related to the sequential nature of the estimation problems. If they are not increasing, it is not a problem, the final rate is just related to the slowest estimation error rate over all timesteps. Thanks for the feedback/your potential concern, to avoid such impressions, we have presented a version in the main text that in the end uses the slowest of all nuisance rates.

2. We’d like to first clarify that the “traditional methods” you mention based on “Inverse Propensity Scoring, or its doubly robust variant” have been used in general to estimate E[V^\pi(S_0)], i.e. the average policy value, and not our estimand, $Q(s,a)-Q(s,a_0)$.

This is a really crucial point that we want to be clear about, to the best of our knowledge there are no prior methods that leverage double-robustness to for MSE convergence to estimate the entire difference of Q function (except for the naive direct method which would simply difference two estimates of Q functions, again without double robustness). Please let us know if you had a specific prior paper in mind that does this.

So, our R-learner based method develops a novel estimation approach and does just that. The theoretical advantage includes the speedup of statistical rates of convergence in mean-squared error estimation of $Q(s,a)-Q(s,a_0),$ rather than the averaged Q function that is the policy value, $E[V^\pi(S_0)]$.

You do bring up an interesting point. One could ask, what is the closest version of traditional methods that estimates $Q(s,a)-Q(s,a_0)$? We do expect that recent approaches based on “pseudo-outcome regression” could be used to regress upon the doubly-robust score Jiang and Li 2016 / Thomas and Brunskill 2015, for example (or the naive IPW score). The R-learner based approach can be understood as a variance-weighted version of such an approach. To our knowledge however, this extension has not been fully developed and tested either. Even then, comparing our method to this hypothetical unpublished extension, our approach based on orthogonal estimation has the benefits of: 1) minimizing a variance-weighted version thereof and 2) the inverse propensity weights therefore do not appear directly in the loss function, reducing optimization issues which can be magnified when using stochastic-gradient based methods with importance sampling weights, as is common in more practical implementations of offline Q-learning.

So overall, our approach is different from traditional such DR methods in OPE (though related), and we expect it also improves upon natural extensions of traditional DR estimation to $Q(s,a)-Q(s,a_0)$ (which again, have not appeared in the prior literature).

Dear Reviewer zoM5,

Thank you for recognizing our paper's technical soundness and contributions. We would like to ensure we have fully addressed your concerns in the revision. Specifically, we have focused on:

1. Improving presentation by restructuring Theorem 3, consolidating nuisance rates, and providing more detailed explanations for assumptions.
2. Clarifying the theoretical advantages of our R-learning method, particularly how it differs from traditional IPW/DR methods by estimating the entire difference of Q function rather than just policy value estimation.

Have these revisions adequately addressed your concerns? As we still have time before the discussion period ends, we would be happy to clarify any of these aspects, especially regarding the technical content or the theoretical advantages of our approach.

Best,

Authors

Dear zoM5,

Thank you for your valuable feedback and support again! We are delighted that our revisions have addressed your concerns with the increased score, particularly regarding the presentation and theoretical distinctions of our approach. Your insights have been instrumental in enhancing both the clarity and technical depth of our work.

Best,

Authors

Responses to Questions:

• 1. on A1: See our prior clarifications. Assumption 1 holds in the numerical examples by construction; future states $S_{t+1}$ and actions $A_t$ depend only on observed states $S_t$.

• 2 We think our prior clarifications under Weaknesses should clear this up to. Assumption 1 is about the general suitability of MDPs for the underlying environment rather than POMDP, and is actually unrelated to the type of structural assumptions made in exogenous state MDP.

The kinds of structural assumptions made in Dietterich 2018 relate to additional conditional independences defined among certain dimensions of the state variable, dimensions of the next state variable, and actions. These structural assumptions are more restrictive and pertain to substructures of the state variable. For example, stated as a conditional independence restriction between endogenous variables $s_{t+1}^\rho$ (a sub-partition or representation of the state) and exogenous variables $s_t^{\rho_c}$, the exogenous state MDP posits that $s_{t+1}^{\rho_c} \perp s_t^{\rho} \mid s_t^{\rho_c}, a_t$. The reward-filtered model instead posits that $s_{t+1}^\rho$ is endogenous and $s_t^{\rho_c}$ is exogenous if $s_{t+1}^\rho \perp s_t^{\rho_c} \mid s_t^\rho, a_t$.

Regarding reward/value estimation, how is the proposed method compared to, e.g., Pan & Scholkopf, 2024, and others?

We mentioned this briefly in the related text but will inline the direct comparison. Pan & Scholkopf focused on ``direct” advantage estimation which elegantly avoids estimating future $Q$ functions; the drawback is that they require imposing nontraditional constraints on the advantage function, infinitely many nonconvex equality constraints, which is much harder than nonconvex unconstrained optimization. Their loss function is also just different.

Overall the philosophies and goals are almost opposite and complementary: we precisely want to build on prior $Q$ function estimation and bring in additional information via the behavior policy to achieve faster rates of statistical convergence. We clarified in lines 86-88. We have updated in the revision a more concise version of this comparison.

Regarding differences to value estimation: our method takes as black-box input estimates of $Q$ functions and is agnostic to which specific method is used; summarized in lines 191-195. Therefore, the convergence rates of prior methods relate to whether the product-error nuisance function rates (Assumption 5 in the revision, prev. in theorem statements) is satisfied.

• 3. Re: $alpha, \delta_0$ in Assumption 5 (margin, now Assumption 7 in the revision):

Let us explain more the role of $\alpha, \delta_0$: they quantify the strength of the assumption, i.e. how quickly the density falls away at the decision boundary. Importantly, they don’t need to be estimated in order to run the algorithm, they just quantify potential improvements in convergence rate. Typically an analyst would work with ``standard” structural assumptions on MDPs, such as linear etc; the margin assumption is common and other works have established $\alpha$ for common structural classes. In lines 332-334 of the revision, we described some values and clarify: for example, $\alpha=1$ for linear MDPs under a weak concentrability assumption, the margin assumption is always true if there is a nonzero gap between the value of optimal and sub-optimal actions.

• 4. Cross-fitting only requires re-estimation of the nuisance functions, i.e. twice or three-times as many function estimations on half or a third of the data. Since the datasize also reduces the overall impact on time is that experiments would take less than twice or three times as long (depending on whether cross-fitting for policy evaluation or optimization). The simple experiments run in a few minutes in parallel on a cluster with 24 parallel cores while the more complex experiments require ~20 minutes for a full comparison of all the methods.

Thanks for your review. We provide clarifications below.

Weaknesses:

• 1: Re other connections to offline RL theory and improvements to rates: our overall approach using orthogonal ML generally implies that offline RL theory results improving convergence results for $Q$ estimation would apply directly via the black-box nuisance estimation assumption (A5). Put simply, our results imply that just $n^{-¼}$ convergence for both $Q^{\pi^e}, \pi_b$ is needed to ensure $n^{-½}$ convergence of $Q(s,a’)-Q(s,a_0)$. Prior results would generally require $n^{-½}$ of $Q^\pi$ to achieve the same.

• 2: We think there is a simple misunderstanding regarding Assumption 1. Importantly, it does not restrict the actions from affecting the state variable. The way we wrote it was confusing since we were trying to avoid introducing potential outcome notation, so the formal statement via conditional independence requires introducing potential outcomes S_{t+1}(A_t); while this is standard in causal inference, it is not standard in RL.

We propose the best way forward instead is to revise Assumption 1 to the statement “Assumption 1: the underlying dynamics are from a Markov Decision Process on (S_t, A_t, R_t, …) and under the behavior policy, actions $A_t$ were taken with probability depending on observed states $S_t$ alone (and not on any unobserved states).”.

To summarize, all that Sequential Unconfoundedness in Assumption 1 represents conceptually is that underlying dynamics are Markovian and that the behavior policy was not acting in a POMDP. It does not restrict actions from affecting transitions. It is much less restrictive than your initial impression, although we understand why you got that impression.

We hope this improves your assessment of the paper, as what you perceived as a weakness was a simple misunderstanding we fixed in the updated revision. We hope that our fix introduces a clear conceptual statement. Importantly, we want to clarify that this was an oversight in our presentation, but what you thought this meant was not actually an assumption in the analysis or a restriction on the algorithm. It works for Markov decision processes.

In addition, the following might help improve the paper:

• 1. Fixed, thanks.
• 2. Fixed, thanks. We clarified that the "Misaligned exo-endo experiment" was just a small change from the prior DGP.

Dear Reviewer WGH1,

Thank you for your detailed review by mentioning our approach is novel and raising the insightful questions. We have worked to address the key points you raised, specifically:

1. We clarified that Assumption 1 is less restrictive than initially interpreted - it simply requires Markovian dynamics and that the behavior policy wasn't acting in a POMDP. We've revised its statement to be more precise and conceptually clear.
2. We explained how our approach connects to the literature:
   • Compared to Pan & Scholkopf (2024), we take a complementary approach by building on Q-function estimation
   • Our relationship to exogenous state MDPs (Dietterich 2018) involves different types of structural assumptions
   • The margin conditions in our assumptions quantify convergence rate improvements for common structural classes
3. Regarding implementation, we've clarified that cross-fitting's computational impact is moderate, with experiments taking between a few minutes to ~20 minutes on a parallel computing setup.

Combined with the updated draft with our point-wise response, have we sufficiently addressed your inquiries? We are open to further discussion if you have additional concerns or require further clarification.

Best,

Authors