Paper: Schrodinger Bridge to Bridge Generative Diffusion Method to Off-Policy Evaluation (ICLR-2024)

Reviewer comment (Theory):
What is $Q^*$ from Section 3.1 in the context of OPE? Is it the Q-function for some policy? And how exactly do we obtain $p_{\mathrm{data}}$ given $Q^*$ and $p_{\mathrm{prior}}$?

Author response:
> What is $Q*$  from Section 3.1 in the context of OPE? Is it the Q-function for some policy? And how exactly do we obtain $p_{\mathrm{data}}$  given $Q*$ and $p_{\mathrm{prior}}$

 $Q*$ from section 3.1 is not the Q-function for some policy, it is a path measure on $\mathcal{S}$ through time interval $[0,T]$ with marginal densities $Q_0^*=p_{{data}}$ and $Q_T^*=p_{\mathrm{prior}}$, which is the solution of the Schrodinger bridge problem (equation (3)). As a result, given $Q*$  and $p_{\mathrm{prior}}$ we can easily achieve $p_{\mathrm{data}}$ by integrating the path measure.

Reviewer comment (Writing):
Is "for $k$ in $1:K$" in Algorithm 1 a typo and what the authors actually meant is "for $k$ in $1:N$"?

Author response:
> Is for $k$ in $1:K$ in algorithm 1 a typo and what the authors actually meant for $k$ in $1:N$

Thank the reviewer for bringing this up. This is not a typo and $K$ represents the number of stages in the training procedure. In one stage, we train the forward and backward process by multiple gradient descent.  We need a repetitive training process, continuously updating the corresponding backward diffusion process through the well-trained forward diffusion process and vice versa.

Reviewer comment (Writing):
**Clarity of Sections 3 and 4**: These sections are challenging to follow and would benefit significantly from a simplification of notation and presentation to enhance readability and comprehension.

Author response:
> I had some trouble understanding how the DSB method described in Section 3.1 is applied to OPE. Section 3.2 is not very clear to me. It'd be nice if the authors could explain what $p_{\mathrm{prior}}$  is in the context of OPE as well as what $p_{\mathrm{obs}}$  is in Algorithm 1? And what is the relationship between $p_{\mathrm{prior}}$ and $p_{\mathrm{obs}}$?

$p_{prior}$ is the prior distribution in diffusion process, and usually it is a standard gaussian distribution. $p_{obs}$ is the same as $p_{data}$ and we apologize for using this term without defining it first. We shall make our notation more consistent in a revised paper.

Reviewer comment (Theory):
**Theoretical Analysis and Coverage Assumption**: The theoretical analysis is missing a discussion on the crucial coverage assumption, which necessitates that the ratio between the behavior and target policy be finite. This assumption is foundational for offline RL. Unfortunately, the paper does not provide assumptions or explanations on this aspect, leaving a gap in the theoretical groundwork.

Author response:
> Do you need the coverage assumption to guarantee the consistency of the estimated Markov transition kernel?

We do not need to additionally state the coverage assumption in the paper, since the coverage assumption is already contained in assumption 6 in our paper. For example, if $\mu$ is degenerate and the target policy $\pi$ is non-degenerate, there will be data points $(s,a)$ that are included in $\pi$'s support range but don't appear in the dataset $\mu$ generated (since they are not included in $\mu$'s support range), which we use to estimate our score. As a result, the estimation error
 
$\mathbb{E}_{q}[\| \mathbf{Z}(kh,\mathbf{X},(\theta,t,s,a))-\mathbf{Z} \|^2]$  

on these $(s,a)$ will be a constant instead of $\epsilon^2$, and assumption 6 does not hold. By the above argument, it requires that the support of $\pi$ to be a subset of the support of $\mu$ (that is, $\sup_{s,a}\frac{\pi(a|s)}{\mu(a|s)}<\infty$, i.e.the coverage assumption) for assumption 6 to hold.

Reviewer comment (Experiments):
**Inadequacy of Numerical Analysis**: The numerical analysis conducted seems insufficient. The paper explores only four benchmark environments with a fixed number of episodes. Additionally, the numerical experiments consider infinite horizons, whereas the paper focuses on fixed episodic settings. Moreover, the proposed estimator demonstrated larger errors in comparison to FQE or MB in two of the environments, questioning its practical effectiveness. To provide a more comprehensive evaluation, it would be beneficial to vary the sample size and offer a thorough comparison across different settings.

Author response:
> The numerical analysis conducted seems insufficient. ... To provide a more comprehensive evaluation, it would be beneficial to vary the sample size and offer a thorough comparison across different settings.

We thank the reviewer for bringing it up.  We have added experiments demonstrating the empirical performance of our estimator is consistent with our claim.  We refer to our rebuttal for reviewer VmGz for the details of the additional experiments.

Reviewer comment (Experiments):
The experimental evaluation is extremely limited

Author response:
We thank the reviewer for bringing it up. To further demonstrate the empirical behavior of our method, we carry out experiments investigating how error change with horizon and number of trajectories. The existing experiments is carried out on a benchmark proposed by Fu et.al.[1], which only provides settings for infinite dimensional MDP and the number of trajectories in the behavior dataset also cannot be changed. Therefore, we now utilize the benchmark in Voloshin et.al.[2] to verify our results.

This time, we do additional experiments on pixel Gridworld which is a pixel-based Grildworld with state being 64-dimensional. We compare it with Model-Based Method and Direct Method Regression meanwhile validating the empirical behavior is in favor of our major theoretical claim. These two baselines are provided in the open-source implementation of the second benchmark. We show the result in two graphs,  [graph1](https://ibb.co/3BtQvzv), [graph 2](https://ibb.co/bB5v9G7 ).

From the second plot, we can see a polynomial dependence on horizon of the absolute error, which is consistent with the result in our theoretical claim.

Reviewer comment (Writing):
The core of the contribution seems to be section 3.3, but there are a few elements that are not quite clear.

Author response:
We thank the reviewer for bringing those questions up. During our training, the time of the MDP $t$, action $a_t$,  current and next state $s_{t-1}$ and $s_t$ are all inputs of the diffusion direction $Z(r, \mathbf{X}_r,\mathbf{M},\phi)$ parameterized by a neural network with parameters $\phi$. We apologize for the typos and misleading expressions in our writing. We will revise this part for better readability.

Reviewer comment (Theory):
Why does the loss need to be masked?

Author response:
We thank the reviewer for bringing those questions up. During our training, the time of the MDP $t$, action $a_t$,  current and next state $s_{t-1}$ and $s_t$ are all inputs of the diffusion direction $Z(r, \mathbf{X}_r,\mathbf{M},\phi)$ parameterized by a neural network with parameters $\phi$. We apologize for the typos and misleading expressions in our writing. We will revise this part for better readability.

Reviewer comment (Theory):
What does it mean to add the time to the training parameters? Does this mean time gets gradient updates in every step? This seems like a mistake.

Author response:
We thank the reviewer for bringing those questions up. During our training, the time of the MDP $t$, action $a_t$,  current and next state $s_{t-1}$ and $s_t$ are all inputs of the diffusion direction $Z(r, \mathbf{X}_r,\mathbf{M},\phi)$ parameterized by a neural network with parameters $\phi$. We apologize for the typos and misleading expressions in our writing. We will revise this part for better readability.

Reviewer comment (Writing):
Overall, section 3.3 could do with a careful rewrite, in particular since this section seems to be the core part of the methodological innovation in the paper.

Author response:
We thank the reviewer for bringing those questions up. During our training, the time of the MDP $t$, action $a_t$,  current and next state $s_{t-1}$ and $s_t$ are all inputs of the diffusion direction $Z(r, \mathbf{X}_r,\mathbf{M},\phi)$ parameterized by a neural network with parameters $\phi$. We apologize for the typos and misleading expressions in our writing. We will revise this part for better readability.

Reviewer comment (Reproducibility):
Can the authors explain how the integral in (1) is evaluated in practice? While the paper discusses how to achieve the individual factors in the integrand, they do not discuss how they evaluate the integral in practice — Importance sampling? SMC?

Author response:
We thank the reviewer for bringing it up. The integral in (1) is evaluated using Monte Carlo method as we outlined in *Model-based OPE* of Algorithm 1. To evaluate the expected reward, we need to sample from the target policy that is given and the transition which is learned using the Schrodinger Bridge diffusion model on offline data. Then we calculate the reward based on the trajectories we sampled from the learned environment.

Reviewer comment (Theory):
While the paper is very rigorous in defining the assumptions of the theoretical analysis (which is a big plus!), it is not clear how feasible each of these is in practice. In particular, do the authors have any means of evaluating whether these requirements are indeed satisfied in their trained estimators?

Author response:
We thank the reviewer for pointing out this gap. Assumptions 2 to 5 are easily feasible and can be straightforwardly checked. Assumption 1 is the Lipschitz continuity of the score function, which is also a common assumption in diffusion modeling which covers a wide range of problems. Papers such as [3] and [4] use similar assumpition. Assumption 6 is also a common learning error assumption which can be obtained via standard stastical analysis, as is stated in [5], [6] and [7]. We will include these justifications in our revised paper.

Reviewer comment (Presentation):
Is the data in figure 1 identical to that in table 1?

Author response:
We thank the reviewer for bringing it up. To further demonstrate the empirical behavior of our method, we carry out experiments investigating how error change with horizon and number of trajectories. The existing experiments is carried out on a benchmark proposed by Fu et.al.[1], which only provides settings for infinite dimensional MDP and the number of trajectories in the behavior dataset also cannot be changed. Therefore, we now utilize the benchmark in Voloshin et.al.[2] to verify our results.

This time, we do additional experiments on pixel Gridworld which is a pixel-based Grildworld with state being 64-dimensional. We compare it with Model-Based Method and Direct Method Regression meanwhile validating the empirical behavior is in favor of our major theoretical claim. These two baselines are provided in the open-source implementation of the second benchmark. We show the result in two graphs,  [graph1](https://ibb.co/3BtQvzv), [graph 2](https://ibb.co/bB5v9G7 ).

From the second plot, we can see a polynomial dependence on horizon of the absolute error, which is consistent with the result in our theoretical claim.

Reviewer comment (Experiments):
The experimental evaluation is too limited, and the existing experimental evaluation is not in the authors' favor.

Author response:
We thank the reviewer for bringing it up. To further demonstrate the empirical behavior of our method, we carry out experiments investigating how error change with horizon and number of trajectories. The existing experiments is carried out on a benchmark proposed by Fu et.al.[1], which only provides settings for infinite dimensional MDP and the number of trajectories in the behavior dataset also cannot be changed. Therefore, we now utilize the benchmark in Voloshin et.al.[2] to verify our results.

This time, we do additional experiments on pixel Gridworld which is a pixel-based Grildworld with state being 64-dimensional. We compare it with Model-Based Method and Direct Method Regression meanwhile validating the empirical behavior is in favor of our major theoretical claim. These two baselines are provided in the open-source implementation of the second benchmark. We show the result in two graphs,  [graph1](https://ibb.co/3BtQvzv), [graph 2](https://ibb.co/bB5v9G7 ).

From the second plot, we can see a polynomial dependence on horizon of the absolute error, which is consistent with the result in our theoretical claim.

Reviewer comment (Experiments):
Certain obvious aspects of the experiments are missing — in particular, if theoretical results indicate favorable results with respect to horizon length, there should be some result showing this scaling, and how it improves compared to other methods.

Author response:
We thank the reviewer for bringing it up. To further demonstrate the empirical behavior of our method, we carry out experiments investigating how error change with horizon and number of trajectories. The existing experiments is carried out on a benchmark proposed by Fu et.al.[1], which only provides settings for infinite dimensional MDP and the number of trajectories in the behavior dataset also cannot be changed. Therefore, we now utilize the benchmark in Voloshin et.al.[2] to verify our results.

This time, we do additional experiments on pixel Gridworld which is a pixel-based Grildworld with state being 64-dimensional. We compare it with Model-Based Method and Direct Method Regression meanwhile validating the empirical behavior is in favor of our major theoretical claim. These two baselines are provided in the open-source implementation of the second benchmark. We show the result in two graphs,  [graph1](https://ibb.co/3BtQvzv), [graph 2](https://ibb.co/bB5v9G7 ).

From the second plot, we can see a polynomial dependence on horizon of the absolute error, which is consistent with the result in our theoretical claim.

Reviewer comment (Experiments):
The same is true for action and state space dimension — while the authors demonstrate favorable scaling, no thorough investigation is done to confirm this, or show that other methods do not have this favorable scaling.

Author response:
We thank the reviewer for bringing it up. To further demonstrate the empirical behavior of our method, we carry out experiments investigating how error change with horizon and number of trajectories. The existing experiments is carried out on a benchmark proposed by Fu et.al.[1], which only provides settings for infinite dimensional MDP and the number of trajectories in the behavior dataset also cannot be changed. Therefore, we now utilize the benchmark in Voloshin et.al.[2] to verify our results.

This time, we do additional experiments on pixel Gridworld which is a pixel-based Grildworld with state being 64-dimensional. We compare it with Model-Based Method and Direct Method Regression meanwhile validating the empirical behavior is in favor of our major theoretical claim. These two baselines are provided in the open-source implementation of the second benchmark. We show the result in two graphs,  [graph1](https://ibb.co/3BtQvzv), [graph 2](https://ibb.co/bB5v9G7 ).

From the second plot, we can see a polynomial dependence on horizon of the absolute error, which is consistent with the result in our theoretical claim.

Reviewer comment (Theory):
Assumption 6 in Section 4 is very strong because it requires very good exploration properties of the behavior policy: there should be enough data for any state-action pair to learn the model uniformly for all state-action pairs.

Author response:
We thank the reviewer for pointing it out, for it does need clarification. In fact, although we don't have direct assumptions on $\mu$, assumption 6 requires good properties of $\mu$ to achieve the small estimation error of the scores. For example, if $\mu$ is degenerate and the target policy $\pi$ is non-degenerate, there will be data points $(s,a)$ that are included in $\pi$'s support range but don't appear in the dataset $\mu$ generated which we use to estimate our score. As a result, the estimation error 

$\mathbb{E}_{q}[||Z(kh,X,(\theta,t,s,a))-Z||^2]$ 

on these $(s,a)$ will be a constant instead of $\epsilon^2$, and assumption 6 does not hold. By the above argument, it requires that the support of $\pi$ to be a subset of the support of $\mu$ (that is, $\sup_{s,a}\frac{\pi(a|s)}{\mu(a|s)}<\infty$) for assumption 6 to hold. In conclusion, the potential assumptions on $\mu$ for our theorem to hold is contained in assumption 6 which we already stated in our paper.

Reviewer comment (Theory):
No direct assumptions on $\mu$ looks very strange for off-policy evaluation algorithms since in the worst case $\mu$ may be a degenerate distribution whereas the goal is to evaluate a non-degenerate one; it requires additional comments.

Author response:
We thank the reviewer for pointing it out, for it does need clarification. In fact, although we don't have direct assumptions on $\mu$, assumption 6 requires good properties of $\mu$ to achieve the small estimation error of the scores. For example, if $\mu$ is degenerate and the target policy $\pi$ is non-degenerate, there will be data points $(s,a)$ that are included in $\pi$'s support range but don't appear in the dataset $\mu$ generated which we use to estimate our score. As a result, the estimation error 

$\mathbb{E}_{q}[||Z(kh,X,(\theta,t,s,a))-Z||^2]$ 

on these $(s,a)$ will be a constant instead of $\epsilon^2$, and assumption 6 does not hold. By the above argument, it requires that the support of $\pi$ to be a subset of the support of $\mu$ (that is, $\sup_{s,a}\frac{\pi(a|s)}{\mu(a|s)}<\infty$) for assumption 6 to hold. In conclusion, the potential assumptions on $\mu$ for our theorem to hold is contained in assumption 6 which we already stated in our paper.

Reviewer comment (Reproducibility):
Details on neural network, training procedure, inference baseline models and evaluation are not in appendix as it is written in the main text; there is only a link to a code. Is it possible to provide these details?

Author response:
We thank the reviewer for bringing it up. Those details will be added to the revised paper. 

* For the neural network, we use a modified U-Net with timestep embedding for the forward diffusion direction $Z(r, \mathbf{X}_r,\mathbf{M},\phi)$ and a transformer for the backward direction $\hat Z(r, \mathbf{X}_r,\mathbf{M},\phi)$. We refer to section c.4 in the appendix of the work by Chen et al.[1] for details.
* For the training procedure, we will add all the hyperparameters, including batch size, discretization, and optimizer in the revised paper.
* For the baseline and evaluation, we have fully followed their implementation as described in the work of Fu et al.[2], where an open-source implementation is available.

Reviewer comment (Theory):
The guarantees claimed in the abstract are very confusing since the policy error is not decreasing with any parameters of the method, whereas Theorem 4.1 shows that error decreases with approximation error and discretization step of SDE.

Author response:
We thank the reviewer for bringing up their confusion. The conclusion of our theoretical analysis is that when the approximation error is small, good OPE result can be obtained using our method without considering the dimension and horizon factor, which avoids the case where the OPE error grows exponentially with respect to dimension and horizon length which many other methods suffer from. We will rephrase our abstract and introduction more adequately to avoid confusion.

Reviewer comment (Theory):
Definition of transition kernel in the beginning of Section 2: current definition by probability of transition implies that a support of $T(s,a)$ is at most countable since the probability should be summable to $1$. Should it be defined as a probability density function?

Author response:
The transition kernel $T(s,a)$ is indeed defined as a probability density function, we are sorry for the confusion of definition. The expectation and conditional expectation of an event are samely defined as in probability theory. We will clear the typo and confusing notation in our revised paper.

Reviewer comment (Writing):
What is expectation of the event? What is conditional expectation of the event?

Author response:
The transition kernel $T(s,a)$ is indeed defined as a probability density function, we are sorry for the confusion of definition. The expectation and conditional expectation of an event are samely defined as in probability theory. We will clear the typo and confusing notation in our revised paper.