Schrodinger Bridge to Bridge Generative Diffusion Method to Off-Policy Evaluation

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.

Could the authors explain how in the algorithm is chosen in general?

$T$ in the implementation is always chosen to be 1, and discretized into 50 or 100 time steps, i.e. the stepsize $h=0.02$ or $h=0.01$.

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.

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.

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.

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.

The core of the contribution seems to be section 3.3, but there are a few elements that are not quite clear

• Why does the loss need to be masked?
• 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.
• Overall, this section could do with a careful rewrite, in particular since this section seems to be the core part of the methodological innovation in the paper

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.

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? Imortance sampling? SMC?

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.

Is the data in figure 1 identical to that in table 1? ....The authors demonstrate favorable scaling, but no thorough investigation is done to confirm this, or show that other methods do not have this favorable scaling.

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, graph 2.

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.

While the paper is very rigorous in defining the assumptions of the theoretical analysis, 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?

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.

[1] Fu, Justin, et al. "Benchmarks for Deep Off-Policy Evaluation." International Conference on Learning Representations. 2020.

[2] Voloshin, Cameron, et al. "Empirical study of off-policy policy evaluation for reinforcement learning." arXiv preprint arXiv:1911.06854 (2019).

[3]Sitan Chen, Sinho Chewi, Jerry Li, Yuanzhi Li, Adil Salim, and Anru R. Zhang. Sampling is as easy as learning the score: theory for diffusion models with minimal data assumptions, 2023a.

[4]Chen, H., Lee, H. & Lu, J.. (2023). Improved Analysis of Score-based Generative Modeling: User-Friendly Bounds under Minimal Smoothness Assumptions. Proceedings of the 40th International Conference on Machine Learning,

[5]Block Adam, Mroueh Youssef, Rakhlin Alexander. (2020). Generative Modeling with Denoising Auto-Encoders and Langevin Sampling.

[6]Xiuyuan Cheng, Jianfeng Lu, Yixin Tan, and Yao Xie. Convergence of flow-based generative models via proximal gradient descent in wasserstein space, 2023.

[7]Gen Li, Yuting Wei, Yuxin Chen, and Yuejie Chi. Towards faster non-asymptotic convergence for diffusion-based generative models, 2023.

No direct assumptions on $\mu$ looks very strange for off-policy evaluation algorithms since in the worst-case may be degenerate distribution whereas the goal is to evaluation non-degenerate one; It requires additional comments.

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.

Details on neural network, training procedure, inference baseline models and evaluation are not in appendix at it is written in the main text, there is only a link to a code. Is it possible to provide these details?

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.

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

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.

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 be it defined as a probability density function? What is expectation of the event? What is conditional expectation of the event?

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.

[1] Chen, Yu, et al. "Provably Convergent Schr" odinger Bridge with Applications to Probabilistic Time Series Imputation." arXiv preprint arXiv:2305.07247 (2023).

[2] Fu, Justin, et al. "Benchmarks for Deep Off-Policy Evaluation." International Conference on Learning Representations. 2020.