
\paragraph{Notations} $\Omega$ is the bounded domain of interest, $\partial \Omega$ denotes the boundary, and $D$ is the radius of a ball centered at the origin that covers $\Omega$. $\nabla \cdot$ and $\nabla$ denote the divergence and gradient operator (with respect to $\bx$). $\varepsilon$ is the entropic regularizer; $\epsilon$ controls the perturbation of the marginals. $c_{\varepsilon}$ and $c$ are cost functions in EOT, and $c_{\varepsilon}=c/\varepsilon$.  $\psi_{\star}$ and ${\varphi_{\star}}$ are the Schr\"{o}dinger potentials; $\nabla \log \overrightarrow\psi(\bx_t, t)$ and $\nabla\log \overleftarrow\varphi(\bx_t, t)$ are the forward-backward score functions in reflected FB-SDE. $\mathcal{D}(\mu_{\star}, \nu_{\star})\subset C(\Omega, [0, T])$ is the path space with marginals $\mu_{\star}$ (data distribution) at time $t=0$ and $\nu_{\star}$ (prior distribution) at $t=T$, $\Pi(\mu_{\star}, \nu_{\star})\subset \Omega^2$ is the product space containing couplings with the first marginal $\mu_{\star}$ and second marginal $\nu_{\star}$.



\section{Reflected Forward-Backward SDE}





\subsection{From Reflected Schr\"{o}dinger Bridge to Reflected FB-SDE}

\label{rFB-SDE_derive}
We consider the stochastic control of the reflected SBP 
\begin{align}
    &\inf_{\bu\in \mathcal{U}} \E\bigg\{\int_0^T \frac{1}{2}\|\bu(\bx,t)\|^2_2 \dd t \bigg\} \notag\\
    \text{s.t.} &\ \ \dd  \bx_t=\left[\bbf(\bx, t)+g(t)\bu(\bx,t)\right]\dd t+\sqrt{2\varepsilon} g(t)\dd  \mathbf{w}_t + \bn(\bx)\dd \mathbf{L}_t  \label{reflected_control_diffusion}\\
    &\ \ \bx_0\sim  \mu_{\star} ,\ \  \bx_T\sim  \nu_{\star}, \ \ \bx_t \in\Omega \notag
    ,
\end{align}
where $\Omega$ is the state-space of $\bx$ and $\bu:\Omega\times [0,T]\rightarrow \mathbb{R}^d$ is the control variable in the space of $\mathcal{U}$; $\bbf: \Omega\times [0,T]\rightarrow \mathbb{R}^d$ is the vector field; $\mathbf{w}_t$ denotes the Brownian motion; The expectation is evaluated w.r.t the PDF $\rho(\bx, t)$ of \eqref{reflected_control_diffusion}; $\varepsilon$ is the diffusion term and also the entropic regularizer; $\mathbf{L}$ is the local time supported on $\{t\in [0, T]|\bx_t\in \partial \Omega\}$ and forces the particle to go back to ${\Omega}$. More precisely, $\mathbf{L}$ is a continuous non-decreasing process with $\mathbf{L}_0=0$ and it increases only when $\mathbf{x}_t$ hits the boundary $\partial \Omega$, that is, 
\begin{align}
\mathbf{L}_t=\int_0^t\mathbf{1}_{\{\mathbf{x}_s\in \partial \Omega\}}d\mathbf{L}_s. \label{local-time}
\end{align}

The existence and uniqueness of SDE \eqref{reflected_control_diffusion} can be addressed through the so-called \emph{Skorokhod problem} \citep{Skorokhod61} which amounts to finding the decomposition for any given continuous path $\mathbf{w}_t\in C(\mathbb{R}^d, [0, T])$, there exists a pair $(\mathbf{y}_t,\mathbf{L}_t)$ such that 
\begin{align}    \mathbf{w}_t=\mathbf{y}_t+\mathbf{L}_t,\notag
\end{align}
where $\mathbf{y}_t\in C({\Omega}, [0,T])$ and $\mathbf{L}_t$ satisfies \eqref{local-time} \citep{Lions_Sznitman_1984}.

Rewrite reflected SBP into a variational form \citep{Chen21}
\begin{align}
    &\inf_{\bu\in \mathcal{U}, \rho} \int_0^T \int_{\Omega} \frac{1}{2}\|\bu(\bx,t)\|^2_2 \rho(\bx, t)\dd \bx \dd t \label{variational_form_supp} \\
     \text{s.t.} &\ \ \frac{\partial \rho}{\partial t}+\nabla \cdot  \mathbf{J}|_{\bx\in\Omega}=0, \ \ \big\langle \mathbf{J},  \bn \big\rangle|_{\bx \in\partial \Omega}=0,  \label{FKP_eqn_supp}
\end{align}
where $\mathbf{J}$ is the probability flux of continuity equation \citep{Pavliotis14} given by %, a variant of Fokker Planck equation.
\begin{align}\label{prob_flux}
    \mathbf{J}\equiv \rho (\bbf+g \bu)-\varepsilon g^2 \nabla \rho.
\end{align}



Explore the Lagrangian of \eqref{variational_form_supp} and incorporate a multiplier: $\phi(\bx, t): \Omega\times [0,T]\rightarrow \mathbb{R}$
\begin{align}
\footnotesize
    \mathcal{L}(\rho, \bu, \phi)&=\int_0^T \int_{\Omega} \frac{1}{2}\|\bu\|^2_2\rho + \phi \bigg(\frac{\partial \rho}{\partial t}+\nabla \cdot \mathbf{J} \bigg) \dd \bx \dd t \label{un_contrained_eqs}\\
    &=\int_0^T \int_{\Omega} \bigg(\frac{1}{2}\rho\|\bu\|^2_2\ - \rho\frac{\partial \phi}{\partial t}-\langle \nabla \phi, \mathbf{J}\rangle \bigg)\dd \bx \dd t \notag + \underbrace{\int_{\Omega} \phi \rho |_{t=0}^T \dd \bx}_{\text{constant term}} + \underbrace{\int_0^T \int_{\partial \Omega} \big\langle \mathbf{J}, \bn\big\rangle     \dd \sigma(\bx) \dd t}_{:=0 \text{ by Eq.} \eqref{FKP_eqn_supp}},\notag
\end{align}
where the second equation follows by Stokes’ theorem.

Plugging \eqref{prob_flux} into \eqref{un_contrained_eqs} and ignoring constant terms, we have
\begin{align}
    \overline{\mathcal{L}}(\rho, \bu, \phi)=\int_0^T \int_{\Omega} \bigg(\frac{1}{2}\rho\|\bu\|^2_2\ - \rho\frac{\partial \phi}{\partial t}-\big\langle \nabla \phi, \rho (\bbf+g \bu)-\varepsilon g^2 \nabla \rho\big\rangle \bigg)\dd \bx \dd t.\label{un_contrained_eqs_v2}
\end{align}






The optimal control $\bu^{\star}$ follows by taking gradient with respect to $\bu$ 
\begin{equation}\label{score_primal_optimal}
    \bu^{\star}(\bx, t)=g(t)\nabla\phi(\bx, t).
\end{equation}

 
Plugging $\bu^{\star}(\bx, t)$ into \eqref{un_contrained_eqs_v2} and setting $\overline{\mathcal{L}}(\rho, \bu^{\star}, \phi)\equiv 0$, we apply integration by parts and derive
\begin{align*}
    0&=-\int_0^T \int_{\Omega} \bigg(\frac{1}{2} \rho g^2 \|\nabla \phi\|^2_2 \ + \rho\frac{\partial \phi}{\partial t}+\rho\big\langle \nabla \phi, \bbf\big\rangle -\varepsilon g^2 \big\langle \nabla \phi,  \nabla \rho\big\rangle \bigg)\dd \bx \dd t\\
    &=-\int_0^T \int_{\Omega} \rho\bigg(\frac{1}{2} g^2 \|\nabla \phi\|^2_2 \ + \frac{\partial \phi}{\partial t}+\big\langle \nabla \phi, \bbf\big\rangle +\varepsilon g^2 \Delta \phi  \bigg)\dd \bx \dd t+\int_0^T\int_{\partial \Omega} \varepsilon g^2 \rho \langle \nabla\phi, \bn\rangle \dd\sigma(\bx)\dd t.
\end{align*}

This yields the following constrained \emph{Hamilton–Jacobi–Bellman} (HJB) PDE:
\begin{align*}
\begin{cases}
&\frac{\partial \phi}{\partial t}+\varepsilon g^2\Delta\phi + \langle \nabla \phi, \bbf \rangle=-\frac{1}{2}\|g(t)\nabla\phi(\bx, t)\|^2_2 \quad \text{   in }\Omega\\
&\langle \nabla\phi, \bn\rangle=0  \qquad\qquad\qquad\qquad\qquad\qquad\qquad\quad \text{   on }\partial\Omega.
\end{cases}
\end{align*}

Applying the Cole-Hopf transformation:
\begin{align}\label{CH_transform}
    \overrightarrow\psi(\bx, t)&=\exp\bigg(\frac{\phi(\bx, t)}{2\varepsilon}\bigg), \ \ 
    \phi(\bx, t)=2\varepsilon \log \overrightarrow\psi(\bx, t),
\end{align}
then we see that $\overrightarrow\psi$ satisfies a backward Kolmogorov equation with Neumann boundary condition
\begin{align}
\begin{cases}
&\frac{\partial \overrightarrow\psi}{\partial t}+\varepsilon g^2\Delta\overrightarrow\psi + \langle \nabla \overrightarrow\psi, \bbf \rangle=0 \qquad \text{   in }\Omega\\
&\langle \nabla\overrightarrow\psi, \bn\rangle=0  \qquad\qquad\qquad\qquad\quad \text{   on }\partial\Omega.\notag
\end{cases}
\end{align}

On the other hand, we set 
\begin{align}\label{FB-score-properties}
    \overleftarrow\varphi(\bx, t)=\rho^{\star}(\bx, t)/\overrightarrow\psi(\bx, t),
\end{align}



where $\rho^{\star}(\bx, t)$ is the probability density of Eq.\eqref{variational_form_supp} given the optimal control variable $\bu^{\star}$. Then from $\rho^{\star}=\overleftarrow\varphi \overrightarrow\psi$, Eq.\eqref{FKP_eqn_supp} can be further simplified to
\begin{align*}
    0&=\partial_t\rho^{\star}+\nabla\cdot \bigg[\rho^{\star} \big(\bbf+g\bu^{\star}\big)-\varepsilon g^2 \nabla\rho^{\star}\bigg]\\
    &=\partial_t (\overleftarrow\varphi \overrightarrow\psi)+\nabla\cdot \bigg[\overleftarrow\varphi \overrightarrow\psi \big(\bbf+g^2 \nabla\phi\big)-\varepsilon g^2 \nabla (\overleftarrow\varphi \overrightarrow\psi)\bigg]\\
    &=(\partial_t \overleftarrow \varphi)\overrightarrow\psi  + \overleftarrow\varphi(\partial_t \overrightarrow\psi)+\nabla\cdot \bigg[\overleftarrow\varphi \overrightarrow\psi \big(\bbf+2\varepsilon g^2 \nabla \log \overrightarrow\psi \big)\bigg]-\varepsilon g^2 \Delta (\overleftarrow\varphi \overrightarrow\psi)\\
    &=\cdots=\overrightarrow\psi \bigg(\partial_t \overleftarrow \varphi +\nabla\cdot \big(\overleftarrow \varphi \bbf - \varepsilon g^2 \nabla \overleftarrow \varphi \big)\bigg),
\end{align*}
where we use the identify $\Delta (\overrightarrow\psi \overleftarrow \varphi)=\overleftarrow \varphi \Delta \overrightarrow\psi + \Delta\overleftarrow \varphi  \overrightarrow\psi+2\langle \nabla \overrightarrow\psi, \nabla
\overleftarrow \varphi \rangle$. Then we arrive at the forward Kolmogorov equation with the Robin boundary condition
\begin{align*}
\begin{cases}
&\partial_t \overleftarrow \varphi +\nabla\cdot \big(\overleftarrow \varphi \bbf - \varepsilon g^2 \nabla \overleftarrow \varphi \big)=0 \qquad \text{ in }\Omega\\
&\langle \overleftarrow \varphi \bbf - \varepsilon g^2 \nabla \overleftarrow \varphi, \bn\rangle =0 \ \qquad\qquad\quad\ \text{ on } \partial \Omega,
\end{cases}
\end{align*}
where the second boundary condition follows by invoking the Stokes’ theorem for the first equation.


Plugging Eq.\eqref{score_primal_optimal} and Eq.\eqref{CH_transform} into Eq.\eqref{reflected_control_diffusion}, the backward PDE corresponds to the forward SDE
\begin{align*}
    \dd  \bx_t=\left[\bbf(\bx_t, t) +  2\varepsilon g(t)^2\nabla\log\overrightarrow\psi(\bx_t, t)\right]\dd t+ \sqrt{2\varepsilon} g(t) \dd  \mathbf{w}_t+\bn(\bx)\dd {\mathbf{L}}_t, \ \  \bx_0\sim \mu_{\star}.
\end{align*}
Reversing the forward SDE \citep{reversal_reflected_BM, Cattiaux_1988} with $\log \overrightarrow\psi(\cdot, t) + \log\overleftarrow\varphi(\cdot, t)=\log\rho^{\star}(\cdot, t)$ based on Eq.\eqref{FB-score-properties}, we arrive at the backward SDE
\begin{align*}
    \dd  \bx_t=\left[\bbf(\bx_t, t) - 2\varepsilon g(t)^2 \nabla\log\overleftarrow\varphi(\bx_t, t)\right]\dd t+ \sqrt{2\varepsilon} g(t) \dd  \overline{\mathbf{w}}_t+\bn(\bx)\dd \overline{\mathbf{L}}_t,\ \  \bx_T \sim \nu_{\star}. 
\end{align*}


Our derivation is in a spirit similar to \citet{Caluya_reflected_SB}. The difference is that the proof is derived from the perspective of probability flux and enables us to derive the Neunman and Robin boundaries more explicitly.

\begin{remark} Regarding the scores $\nabla\log\overrightarrow\psi$ and $\nabla\log\overleftarrow\varphi$ at $t=0$ and $T$, we follow the standard truncation techniques \citep{DDPM, Diffusion_constrained} and fix them to $\bm{0}$. We refer readers to Appendix C of \citet{score_sde} and Appendix B of \citet{song_likelihood_training} for more discussions.
\end{remark}

\subsection{Connections between reflected FB-SDEs and flow-based models}\label{prob-flow-ode}

Similar to \citet{Diffusion_constrained, reflected_diffusion_model}, our flow representation in Eq.\eqref{FKP_eqn_supp} together with \eqref{FB-score-properties} naturally yields
\begin{proposition}[Probability Flow ODE]
    Consider the reflected FB-SDEs \eqref{FB-SDE} with Neumann and Robin boundary conditions. The corresponding probability flow ODE is given by
\begin{align*} 
\dd  \bx_t&=\left[\bbf(\bx_t, t) + \varepsilon g(t)^2\big(\nabla\log\overrightarrow\psi(\bx_t, t)- \nabla\log\overleftarrow\varphi(\bx_t, t)\big)\right]\dd t. 
\end{align*}
\end{proposition}

The result is the same as in \citet{forward_backward_SDE} and provides a stable alternative to compute the log-likelihood of the data.