Score-Based Diffusion Policy Compatible with Reinforcement Learning via Optimal Transport
Authors: Mingyang Sun
Pengxiang Ding
Weinan Zhang
Donglin Wang
Venue: ICML 2025 poster (2025)
Keywords: Diffusion Policy
Reinforcement Learning
Optimal Transport
Primary area: reinforcement_learning->deep_rl

Abstract: Diffusion policies have shown promise in learning complex behaviors from demonstrations, particularly for tasks requiring precise control and long-term planning. However, they face challenges in robustness when encountering distribution shifts. This paper explores improving diffusion-based imitation learning models through online interactions with the environment. We propose OTPR (Optimal Transport-guided score-based diffusion Policy for Reinforcement learning fine-tuning), a novel method that integrates diffusion policies with RL using optimal transport theory. OTPR leverages the Q-function as a transport cost and views the policy as an optimal transport map, enabling efficient and stable fine-tuning. Moreover, we introduce masked optimal transport to guide state-action matching using expert keypoints and a compatibility-based resampling strategy to enhance training stability. Experiments on three simulation tasks demonstrate OTPR's superior performance and robustness compared to existing methods, especially in complex and sparse-reward environments. In sum, OTPR provides an effective framework for combining IL and RL, achieving versatile and reliable policy learning.

Score-Based Diffusion Policy
Compatible with Reinforcement Learning via Optimal Transport
Mingyang Sun 123
Pengxiang Ding 22
Weinan Zhang 43 Donglin Wang
Abstract
ity to exccute complex sequences, but also the adaptability to handle uncertainties and disturbances. Imitation Learn- ing (IL) has emerged as a popular data-driven approach for training robots by imitating demonstration data, with ad- vancements of Behavior Cloning (BC) like diffusion models (Chi et al., 2023; Ajay et al.) and action chunking (Zhao et al., 2023) enhancing its ability to learn complex, long- horizon behaviors. Notably, Diffusion Policy (DP) (Chi et al., 2023) has shown promise due to the capacity to handle multi-modal action distributions, excel in high- digh-dimmnsional spaces, and achieve stable training through techniques like denoising and score matching. However, these advance- ments still fail to address the fundamental flaws of BC, which remains highly susceptible to distributional shifts, where the policy encounters states outside its training data, leading to compounding emrors (Ross & Bagnell, 2010).
Diffusion policies have shown promise in learning complex behaviors from demonstrations, particu- larly for tasks requiring precise control and long- term planning. However, they face challenges in robustness when encountering distribution shifts. This paper explores improving diffusion-based imitation learning models through online interac- tions with the environment. We propose OTPR (Optimal
score-based diffusion Policy for Reinforcement learning fine-tuning! a novel method that integrates diffusion policies with RL using optimal transport theory. OTPR leverages the Q-function as a transport cost and views the policy as an optimal transport map. en- abling efficient and stable fine-tuning. Moreover, we introduce masked optimal transport to guide state-action matching using expert keypoints and a compatibility-basel resampling strategy to en- hance training stability. Experiments on three simulation tasks demonstrate OTPR's superior performance and robustness compared exist- ing methods, especially in complex and sparse- reward environments. In sum, OTPR provides an effective framework for combining II. and RI., achieving versatile and reliable policy learning.
Reinforcement Learning (RL) offers a powerful framework for autonomous learning through trial-and-erro interactions guided by reward signals, making particularly effective in training reactive controllers that adapt to noise, disturbances, and unforeseen states (Kober al., 2013). RL learns correc- tive behaviors directly from experience, enabling policies to recover from errors and handle states beyond the training distribution. Its ability to optimize over long time hori- zons can also refine action sequences, enhancing robustness and precision. Unlike IL, which benefits from leveraging demonstration data to jump start learning, RL enhances generalization by exploring diverse scenarios and adapting dynamically to environmental changes. However, RL also faces significant challenges, including the need for carefully designed reward functions and vast interaction data, which is costly to collect, particularly in real-world settings (Park et al., 2024).
1. Introduction
Robotic manipulation is an intricate endeavor, where the delicate interplay of long-term planning and instantaneous control poses a captivating challenge - the quest to develop policies that can seamlessly navigate this balance lies at the forefront of modern robotics (Heo et al., 2023; Mu et al., 2021; Chen et al., 2023). Tasks demand not only the abil-
These strengths and weaknesses suggest that an inte- grated approach, combining RL's adaptability with DP's demonstration-ddrivnn learning, holds promise for achieving reliable, scalable, and versatile robotic manipulation. The most common approach is to pretrain a imitation policy with human data and then finetune with RL (Black et al., 2024). Some methods apply additional regularization (Rajeswaran et al., 2018) or seperated policy network (Ankile al., 2024) to ensure that the knowledge from demonstrations
'Zhejiang University. Hangzhou, China "Westlake University, Hangzhou, China 'Shanghai Innovation Institute. Shanghai, China "Shanghai Jiao Tong University, Shanghai. China. Correspondence to: Donglin Wang <wangdonglin @ restlake..du.cn...
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Score-Based Diffusion Policy Compatible with Reinforcement Learning via Optimal Transport
does not get washed out quickly by the randomly initialized critics, which may suffer from hyper-parameter tuning issue. Additionally, the structure of diffusion models (iterative re- finement) inherently complicate the application of standard RL algorithms, often leading to low-efficiencyy instability or requiring significant architectural modifications (Ren et al., 2024; Mark et al., 2024; Psenka et al., 2024).
2022), a planning algorithm with diffusion models for of- fline reinforcement learning. This framework is extended to other tasks in the context of offline reinforcement learn- ing (Wang et al., 2022), where the training dataset includes reward values. Most typically, diffusion based policies are trained from human demonstrations through a supervised objective, and enjoy both high training stability and strong performance in modeling complex and multi-modal trajec- tory distributions. The application of DDPM (Ho et al., 2020) and DDIM (Song et al., 2020a) on visuomotor policy learning for physical robots (Chi et al., 2023; Gong et al., 2024; Jia et al., 2024) outperforms counterparts like Be- havioral Cloning. While these techniques effectively learn from multi-modal data, they often create models that are non-trivial to fine-tune using RL. Even if they were compati- ble with RL, the fine-tuning process can be computationally prohibitive due to the large number of parameters in modern policy models.
In this paper, we integrate insights from the optimal trans- port theory (Gu et al., 2023; Montesuma et al., 2024) to refine the diffusion policy optimization process, leverag- ing knowledge gained from expert trajectories to improve learning efficiency and policy performance in subsequent re- inforcement learning tasks. By utilizing the Q-function as a transport cost and viewing the policy as an optimal transport map, we establish an equivalent relationship between the optimal transport map and the optimal policy, which opens avenues for applying recent advantages of RL to diffusion policy directly. Our key contributions are as follows:
Training diffusion models with reinforcement learning. As demonstration data are often limited, there have been many approaches proposed to improve the performance of diffusion-based policies. One straightforward approach (Black et al., 2024; Fan al., 2024) involves framing diffu- sion denoising as a Markov Decision Process (MDP), which facilitates aefference-alignee generation with policy gradi- ent reinforcement learning. However, this approach often suffers from instability, limiting its practical applicability. (Ren et al., 2024) introduced policy gradient loss on a two- layer MDP for direct diffusion policy fine-tuning, which mitigates this instability, but the method is architecture- specific and does not introduce closed- control. Alter- native approaches integrating diffusion architectures with reinforcement learning (RL) include leveraging Q-function- based importance sampling (Hansen- et al., 2023), employing advantage- advantage-weighted regression (Goo & Niekum, 2022), or reformulating the objective as supervised learn- ing problem with return conditioning (Chen et al., 2021; Janner et al., 2022; Ajay et al.). Additionally, researchers have explored enhancing the denoising training objective by incorporating Q-function maximization (Wang et al., 2022) and iteratively refining the dataset using Q-functions (Yang et al., 2023). Another promising direction involves augment- ing a frozen, chunked diffusion policy model with a resid- ual policy trained through online RL, enabling improved performance without modifying the pre-trained diffusion model (Ankile et al., 2024).
We proposed an Optimal Transport guided score- based diffusion Policy for Reinforcement Learning fine-tuning (OTPR), which is the first work to sys- tematically combine optimal transport theory with dif- fusion policies for reinforcement learning fine-tuning. OTPR's 's core lies in solving the Le-regularized OT dual problem, thereby deriving a compatibility funetion that establishes soft coupling relationship between states and actions, effectively integrating imitation learning and reinforcement learning.
To enhance efficiency and accuracy, we introduce the Masked Optimal Transport to leverage the paired state-action from expert data as keypoinns to guide the matching of the other state-action points from replay buffer, which seamlessly integrates imitation learning and reinforcement learning objective.
To address the sub-optimal performance of the con- ditional score-based model when trained with stan- dard algorithms on mini- data, we introduce a Compatibility-basee Resampling strategy to selects action with high compatibility scores to guide the train- ing process, thereby enhancing performance.
We conduct extensive experiments on 3 simulation tasks spanning various difficulty levels. The results demonstrate that OTPR consistently matches or outperforms existing state-of- of-te--rtthe-rt methods in all tasks. with particularly no- table improvements in more challenging scenarios.
3. Background
2. Related Work
In this section, we offer fundamental definitions and theories to lay the groundwork for our framework, which will be thoroughly analyzed afterwards.
Diffusion based policies. Diffusion- policies have shown recent success in robotics and decision-making ap- plications. In a pioneering work, "Diffuser" (Janner et al.,
2

---

Score-Based Diffusion Policy Compatible with Reinforcement Learning via Optimal Transport
(a) Estimation Optimal Transport Plan
(c) Online Interaction
ww v
(ss, a) A
S
Hss, a)
Resample
a
(s, a)
Q
Diffusion Policy
VU (a)
(b) Training
QΦ So
T iter
Q-Function Training
Se (S,
a
Score Matching
a
Environments
Transitions
Expert
Replay Buffer
Demos
S
Figure 1. Overview of OTPR. (a) Estimation Optimal Transport Plan: A stochastic dual approach with two parametrized dual variables introduced to estimate optimal transport plan with Q-cost from state distribution and action distribution. (b) Training: OTPR pre-trains diffusion model from the expert's ' data. It then iteratively performs RL optimize Q-function and trains diffusion models by score matching. (c) Online Interaction: In the inference step, the policy makes action inference by iteratively denoising random noise, conditioned current state. OTPR then employs the compatibility function H to reweight cach action before resampling
3.1. Optimal Transport
In this work, we consider the L2 regularization introduced by (Dessein et al., 2018), whose computation is found to be more stable since there is no exponential term causing overflow. For all x E X' and y V
Given two probability spaces (X', (),,v) and a cost func- tion c X' x y + R, the Monge problem (Villani et al., 2009) is solving optimal map T y such that
Rςa(1) fr
dy(x, v)
inf ((T)) = Em-a [c(2, T(x)))
1T - (1)
du(z)dd(v). (4)
(dp(x)dr(g))
where the random variables
and Tæp is push for- ward of μ subject to (T*) == ()') for any measurable set 3" c V. Instead of finding the map T in the original Monge problem, the relaxed Kantorovich optimal scheme (>) is obtained by y realizing
where dr(ar.) is the density, i..., the Radon-Nikodym du(x))d(s) derivative of y w.r.t. μ X v.
Regularized OT Dual We refer the objective Kx(7) as the primal objective, and will use [N, u) to refer the associated dual objective, with dual variables u, The dual of the regularized OT problems can be obtained through the duality theorem.
inf [(+)= = Exxyo b) T(,,)).
(2)
where I'(H,, v) is the space composed of all joint probability measures y on X' x y with marginals p and v.
sup E (a.))~uxx [u(x)) v (((()),e())) u,6
(5)
Regularized OT Regularization was introduced in (Cuturi, 2013) to speed up the computation of problem, which is achieved by incorporating a negative-entropy penalty the primal variable of Problem
1 where (((()) 4X (u(x) + v(y) c(æ, y))
is concave w.r.t (u, v) and a. max (a, 0). In order to recover the solution 7> of the regularized primal prob- lem 3, we can use the first-order optimality conditions of the Monchee-Rockkaellarr 's duality theorem,
inf (K(()) = Exxy- (c(x, y) + XR(7) 7 E Te v)
(3) As highlighted by (Daniels et al., 2021), adding regular- ization term with a-strong convexity (such as entropy or squared L2 norm) to the problem is sufficient condition for Aa-strong convexity (x()) ( in L2 norm, which makes the dual problem unconstrained maximization problem.
d> y) Hx y)dp(())vv()))
(6)
where (, 2A (u(x) v(y) v)).
called compatibility function.
3

---

Score-Based Diffusion Policy Compatible with Reinforcement Learning via Optimal Transport
3.2. Reinforcement Learning and Imitation Learning
4. Method
Reinfrocement Learning We consider a standard Markov decision process (MDP) consisting of state space s E S, continuous action space E A. deterministic state tran- sition function P S x A + S, reward function r : S + I and discount factor k. T ~ π denotes the dis- tribution of trajectory (so, ao, 81, a1....) given the policy T(lss). The action-state value function is Q (8, a) = E Ei=nk'relar
4.1. An Optimal Transport View of Policy Learning
We approach the policy optimization problem from the per- spective of optimal transport. Considering Eq. 1, by sub- stituting the cost function c(2; a) with the critic Q(s, a) and viewing our policy я map that moves mass from the stationary state distribution pu(s,, to the corresponding distribution of actions v(() given by an optimal behavior policy x (.((11)) (-1s), we formulate the following primal state- conditioned Monge OT problem:
a, So = s]. The goal of RL is to learn the policy that maximizes the discounted expected cumulative reward over trajectory T, defined as (TT) E k= r-l.
inf { M(x) = Eg-p [-Q" (s))
T#/ =v] (8)
Imitation Learning We assume access to a dataset 'D of demonstrations collected by expert human operators (often assumed to be optimal), Each trajectory T E consists of a sequence of transitions {(s0, ao)..... (sK, ag)) The most common IL method is behavior cloning (BC) which trains parameterized policy Re to minimize the negative log- likelihood of data, i..., L(0) = - -B(..a)- llog To (a) (ss.]. In this work, we assume ne follows an isotropic Gaussian as its action distribution for simplicity. With the isotropic as- sumption, the BC training objective can be formulated as the following squared loss: (0) = (s)aa)> (T) a 1l2
The objective is to minimize the expectation of the negative critic function Q" while mapping exclusively to the distri- bution of actions given by the behavior policy л², concept also explored in offline RL work (Asadulaev et al., 2025).
Proposition 4.1. Given an optimal behavior policy TO and a critic-based cost function = -Q², let TT is the solution to Eq. 8 with the Q3 cost function. Then it holds that: JRL = (x)
The proof is given in Appendix B.I. Proposition 4.1 of- fers valuable insights into the connection between optimal transport theory and RL by establishing an equivalent rela- tionship between the optimal transport map and the optimal policy. Meanwhile, given the paired state-action data de- rived from expert policy, the IL problem can be reframed as achieving conditional optimal transport map (i.e., dif- fusion policy). This link indicates that the transformations defined by the optimal transport map can effectively inte- grate RL with IL. We will next demonstrate how to use an estimated optimal transport plan (as intuitively illustrated in Fig. of Appendix) to serve as guide to utilize reinforce- ment learning to optimize the pre-trained diffusion policy via imitation learning.
3.3. Conditional Score Based Diffusion Policy
The conditional Score Based Diffusion Models (SBDMs) (Song et al., 2021; Batzolis et al., 2021) aim to generate a tar- get sample y from the distribution y of target training data given a condition data æ. In imitation learning, diffusion pol- icy regard state 8 as condition x and use a forward stochastic differential equation (SDE) to add Gaussian noises to the tar- get training data a for training the conditional score- score-based model. The forward SDE is da; = f(a,, t)dt + g())dw with ag ~ V, where w E RDD is standard Wiener process, f(,tt): _ + IR is the drift coefficient, and g(t) E R is the diffusion coefficient. Let 1/jo be the conditional distribu- tion of a given the initial state ag, and v4 be the marginal distribution of a, The conditional score-based model is trained by denoising score-matchin loss:
4.2. OT-Guided Conditional Denoising Score Matching
In IL setting, denote the all condition states as Sound (a) for a target action and f is the measure by push- forwarding v using sond, i.c., m(s) = Σ (a=) v(a) over the paired training dataset D. Section 3.3 provides a explicit reformulation for the conditional score- based diffu- sion policy with the paired training data. Proposition 4.2. Let C(s, is the generalized Dirac delta function, 5(s) = 0, if s & Scomd(a), then JusM(0) (0) in Eq. 7 can be reformulated as
(os Ee we ||so(ar; Scona( Va, log u )l
(7)
n. & smmiia) (8) where 8
where w; is the weight for time t. In this paper, is uni- formly sampled from [0, T). i.e., 1 ~ U(00 T)). With the trained sy(a; s, s,,), given a condition data 8, the tar- get sample ag is generated by the reverse SDE as da,- = [f( - gt)=((t)ssp)))))))))) (s))ss((aa)) : s, t)]dt -adt-gop gt)) where w is a stan- dard Wiener process in the reverse-time direction.
(CDSM(0) =E,w, E.wp Bann Fap B) / (a)la) ||so(a (a: - Va. log 1 (9)
Furthermore, v(s, a) = C(s, ())(()v(()) is a joint distri- bution for marginal distributions p and
4

---

Score-Based Diffusion Policy Compatible with Reinforcement Learning via Optimal Transport
The proof is given in Appendix B.2. In Proposition 4.2, the coupling relationship of condition state and target action is explicitly modeled in C(s, a). Nevertheless, in contrast to IL, the definition of C(s, a) in RL is not explicit due to the absence of an optimal paired relationship between 8, a. Fortunately, the joint distribution v exhibits a similar formulation to the transport plan y in Eq. 6. We therefore use Legreullarzeed OT to model the coupling relationship between state and action for unpaired settings. Specifi- cally, given a Q network learnt by RL, the L2-regularized OT is applied to the distributions/ jr, // to approximately con- struct a conditional transport plan { 8) = H (s, a)v(a), and the coupling relationship of the condition data s and target data a built by the compatibility function H(s, a). We then extend the formulation for paired setting in Eq. 9 by replacing C with H/ develop the training objective for unpaired setting, which is given by
mizing the optimal tranport problem gives u, v, which is often computationally challenging because OT needs trans- port all the mass of state to exactly match the mass of action distribution, which presents computational challenges. For- tunately, in imitation learning, expert demonstrations De have provide matched pairs of state and action data points (called "keypoints") X = {(8;, a,)) ) These keypoints are not only valuable but also crucial for investigating how to leverage them to guide the correct matching in OT. We introduce masked OT (Gu et al., 2022) to leverage the given matched keypoints to guide the correct transport in OT by preserving the relation of each data point to the keypoints:
inf { {K()) == Bexan (9ss a)] | ((..rm)) (12)
where the transport plan m @ 9 is (m O (s, = m(s, (s, a), and m is a binary mask function. Given pair of keypoints (si,
Baww (s, Ea (0) (s - Va, (og)))(as)a)) 2 (10)
K, then m(s;, ai) 1, m(s;, a) = 0 and m(s, a) 1 if s, a do not coin- cide with any keypoint. The mask-based modeling of the transport plan ensures that the keypoint pairs are always matched in the derived transport plan. g in Eq. 12 is de- fined as g(s, a) = d(R" R(),, where Ry Ri E (0, 1)N model the vector of relation of s, a to each of the paired keypoints in state and action space respectively, and d is the Jensen- divergence. The i-th elements of Rg and Ra are respectively defined by
In Eq. 10, H is a "soft" coupling relationship of state data and action data, because there may exist multiple a satis- fying (s, a) > 0 for each 8. We minimize JUDSM(() to train the conditional score-based model se(ar; 8, t).
Theorem 4.3. For da, = f(,,,t)dtt g() with ag A--ss) and E [0, Let (s)()))) be the distribution of at and (Csm(O) Bannaaanpaanntalll
μ. consider the forward SDE
8, Va, log (ggv((a)(s13 then we have (nnnss = (f)
exp(-c(8,s:)/p)
R Σラ= exp( exp(_c((s 8;)pp) (xp(-c(a,,)/pp) Rd, Σ5= exp(- kz)/p)
(13)
We give the proof in Appendix B.3. Theorem 4.3 indi- cates that the trained so(a;; 8, () using Eq. 10 approximates Va, log (a(aa)s) (s). Based on this, we can interpret our ap- proach as follows. Given a condition data we sample action data ag from the conditional transport plan ((ao) s), produce at by the forward SDE solver (examples given in Appendix A.2), and train (a, 8, f) to approximate Va, log v/ (s))s.)
where p is a commonly used temperature in the softmax function. Further, the masked matrix is introduced into the duality of the 1,2 reguarized OT problem, and the penalty term Fx is updated as:
sup E ( (u(s) + v(a) + Fx(u(s), (14) 0,00
Sample Generation We denote the trained conditional score-based model as "a( 8, t) where 0 is the value of 9 after training. Given the condition state s, we generate action samples by the following SDE:
1 m m() + u(()) g(s, a)) a 4X
(x(us),
The dual 5 and 14 are unconstrained concave, which can be maximized through stochastic gradient methods by sam- pling batches from j v. Following (Seguy et al., 2018). we use deep neural networks for their ability to approximate Uw Uae with the parameters w and the estimate of OT plan is
da,- = [/(a,,t) 2(1)21g (a; 8, 0)] di + g())dw. (11)
Numerical solvers such as the Euler-Maruyama method, DDIM, or DPM-Solver can be employed to efficiently solve this reverse SDE, enabling the generation of high-quality action samples (Song et al., 2020a; Lu et al., 2022).
(s, = -H(ss a)du(s)dv(().
(15)
4.3. Expert Data Masked Optimal Transport
where (ss 2X (uv(s) + =a(a) - g(s, a))+.
For the computation of H, a value- based reinforcement learning can provide an estimated Q-network, while opti-
The pseudo-codeis given in Appendix A.I.
5

---

Score-Based Diffusion Policy Compatible with Reinforcement Learning via Optimal Transport
Algorithm 1 Online Score-Based Diffusion Policy Training
To be specific, we investigate the upper bound of the ex- pected Wasserstein distance Wa (+|ss) Since (-(---) is proper metric, we can conveniently leverage the triangle inequality to derive an upper bound for this expectation: Expp (- (1 (-1s)) S
The pre-trained imitation policy, expert demon- strations D. initialzed Q-network and replay buffer B. Output: Trained conditional score-based policy s0- for iteration = ....... do
Input:
Learn Mw: t/ by optimizing the dual problem 14. while no done with episode do
+ Bgn W2 (-1s), (-1s)) where 4 is the estimated OT plan depending on uas Va- This inequality provides a means to assess the upper bound by breaking it down into two more manageable comparisons.
Observe current state 8: Sample 0; by the SDE 11, 1 1,.... L;
Compute (s, 01) using Eq. 15;
To bound the first term, we denote the Lagrange function for Le-regularized OTs in Eq. 5 as L(y, u, v) with dual variables u, v as follows:
Normalize: Pt H(s,a))
Select a as categorical from Pt
Store transition in B.
C(y u.) / (T(s,a) a(,,a)²) ((s)v(s)) ds
end while
Learn Q-network by employed RL algorithm.
Update 0 of score model so by fitting noise.
+ (u(s) (f ((8 (((8aa)da-116 ds (16)
end for
+ ffola) (f7ls,a)ds - v(a) da.
4.4. OT-Guided Training
Because L(7,, u, v) is sum of Kx(7) and linear terms, the Lagrangian inherits X-strong convexity in L1 Lorm.. Given the trained ug and no which are c-approximate maximizers of [u, v), the pseudo-plan ? = H(s, u2 (a)p(s)v(()) satisfies:
To implement (O) in Eq. 10 using training sam- ples to optimize 6, we can sample mini-batch data from replay buffer, and then compute H(s, a) and Js = (|l [|s0|| (ay; s, V log ( over the pairs of (8, a) in S and A. However, such strategy is sub-optimal. This is because given mini-batch of samples and a, for each source sample s, there may not exist tar- get sample a in the mini-batch with a higher value of H (s, a) that matches condition data s. Therefore, few or even no samples in a mini- batch contribute to the loss function, lead- ing to a large bias of the computed loss and instability of the training. To tackle this challenge, we generate L samples from policy model, and then use the compatibility function to reweight these actions, ultimately forming the intended policy when resampled. This approach is summarized in Algorithm 1. In implementation, our approach can be used to replace the policy improvement step in multiple RL algo- rithms, while keeping the critic training as At evaluation time, we simply taking the action by setting L 1 to reduce computational requirements.
^ A - >- S [(9, us, v2) L(y", u², u²) S (17)
Since the strong convexity of L implies Polyak- Lojasiewicz (PL) inequality, we have,
|1 17-7" 1h-7"11 LV us, us
(18)
Consequently, we can derive an upper bound for the ex- pected Wasserstein distance as follows:
R (2((((()()) S ! ll VV us, uol mo)l (19)
x
where 7 maxa, a' 12
For the bound of Eg (L (-)s),))))()(111)) (-1s)) to get without explicit f and g given, but from the existing convergence guarantees for general class of score- based generative models, we get 1E (. 4-11s) S e, which can be easily interpreted as two terms (1) the initial- ization of the algorithm at êr (-|s) instead of (((((s))), (2) the discretization and score-matchie moteenn errors in running the algorithm (Kwon et al., 2022; Gao et al., 2023).
it is difficult
5. Analysis
The proposed OTPR essentially aims to develop a condi- tional score-based diffusion policy for data transport from state space to action space in OT. To generate samples from conditional OT plan (-ss) the algorithm involves two key module learning: the dual term (M vw) and the score model 8p. In this section, we will provide an analysis from the perspective of optimal transport, illustrating how the two aforementioned processes establish the upper bound of the distance between the distribution ,AD (als) of generated samples and the conditional optimal transport plan 7 (a]s).
6. Experiments
In this section, we evaluate OTPR and several prior ap- proaches, in a number of benchmark domains that require learning policies from static offline expert data and then
6

---

Score-Based Diffusion Policy Compatible with Reinforcement Learning via Optimal Transport
kilchen-complete-v0
kicchen-mixed-00
Kilchen-pardall-0
100
100
100
60
Normalized Score
Nomalize Score
Nommatzd Scort
-0
40
-0
DFPO
3PPG
DPFO
DQ_
(DCL
IDQL
DQL
DGL.
DQ.
Interacton Stepe (x10K)
Interaction Steps (-10K)
Interaction Steps (010))
CALVIN
Can
Square
192
1.0
DPPO
IDGL
OTPR
Nomalize Score
Normalzzd 02
0.8
0.2
Normallzdd
0.4
OPPO
DP
02
DOL
D0.
OTFR
GTPR
0.0
0.
Intoraction Staps (H1K)
Interaction Staps (-200k)
Intoraction Stopa (=1M)
Figure 2. Learning curves of online fine-tuning with various methods. Observe that OTPR largely always dominates or attains similar performance to the next best method. Other methods for fine-tuning diffusion policies (IDQL, DQL, DPPO) are bit unstable, and perform substantially worse.
fine-tune them with limited online interaction in the MDP (offlinett-onllnee fine-tuning). We also study the hybrid RL problem setting (i.e., online RL with offline data put in the replay buffer) for some experiments. More results are shown on Appendix D. Finally, we perform ablation experiments to understand the utility of different components of OTPR.
and best practices for running OTPR in Appendix. We instantiate OTPR using the popular IQL. with keeping the critic training as is. For the image-based do- main, we use a ResNet 18 encoder and store features in the replay buffer to facilitate the estimation of the dual terms.
6.2. Results
6.1. Experimental Setup
Comparisons with Other Online Fine- Fine-Tuning Methods. We conduct a comprehensive comparison of OTPR against a range of reinforcement learning (RI.) methods designed for fine-tuning diffusion-based policies. Specifically, we evaluate the following approaches: (1) Implicit Diffusion Q-Learning (IDQL) MansennFFstuuhh et al., 2023), which extends Implicit Q-Learning (IQL) to incorporate diffusion policies through critic-based re-ranking; (2) Diffusion Pol- icy Optimization (DPPO) (Ren et al., 2024), which fine- tunes diffusion policies initially learned via imitation learn ing by optimizing a two-layer Markov Decision Process (MDP) loss; and (3) Diffusion Q-Learning (DQL) (Wang et al., 2022), which trains diffusion policies using a repa- rameterized policy gradient estimator similar to the Soft Actor-Critic (SAC) framework (Haarnoja et al., 201 8).
Environments and tasks. We study: (1) Robomimic tasks (Mandlekar et al., 2021), which is a commonly used benchmark designed to study imitation learning for robot manipulation. The evaluation score represents the success rate. (2) Franka-Kitchen tasks (Gupta et al., 2019), which require solving sequence of four manipulation tasks in kitchen environment with a 9-Dof Franka robot; and (3) the CALVIN benchmark (Mees et al., 2022), an eval- uation benchmark designed for long-horizon, language- conditioned manipulation, which requires solving a se- quence of four manipulation tasks in a tabletop environment. The evaluation score for a trajectory is the maximum number of sub-tasks completed simultaneously at any single point in the trajectory. The CALVIN task is significantly challeng- ing, as policies must be learned directly from pixels using offline play data obtained through human teleoperation.
Overall, OTPR performs consistently and significantly im- proves fine-tuning efficiency and asymptotic performance of diffusion policies. Notably, OTPR consistently main-
Implementation details. We provide a detailed list of
7

---

Score-Based Diffusion Policy Compatible with Reinforcement Learning via Optimal Transport
Table Comparison of OTPR with other demo-augmented RL algorithms. OTPR outperforms every other approach, both in terms of the offline performance (left of -+) and performance after fine-tuning -tnnin (right of ++)
Franka-Kitchent Mixehen-Mixed-v0
RoboMimic
Kitphen-Complete-v0
Kitchen-Partial-v)
Can-State
Square-State
RLPD
0 18
0 14
0 - 34
0 - 0
0 3
Cal-QL
19 + 57
37 + 72
59 + 84
0 - 0
0 - 0
IBRL
0 + 25
0 - 13
0 -> 15
0 -> 64
0 - 50
OTPR
61 - 92
59 + 79
42 + 93
63 - 99
40 - 98
tains high normalized scores in the kicchen-complet--o,, CALVIN and Can task, while other methods exhibit relative instability, especially DQL and IDQL, which show consider- able fluctuations in performance across different interaction steps. This may be attributed to both DQL and IDQL per- forming off- off-iy updates and propagating gradients from
Norma lood Broee
the imperfect Q function to the actor, which results in even
OTPE CTPRU
greater training instability in sparse-reward tasks given the continuous action space and large action chunk sizes. In contrast, OTPR can quickly mitigate the adverse effects brought about by this issue by leveraging the guidance of the compatible function. This analysis suggests that OTPR is robust and effective approach for online fine-tuning in diffusion policy tasks, consistently outperforming the other methods in terms of stability and overall performance.
Intenacioon Stape (-203K)
Interaction Siepa (x2D0K)
Figure 3. (left) Comparison between OTPR with different guidance (H. Q and A). (right) Comparison between OTPR with (OTPR-M) and without (OTPR-U) the expert data mask.
over other diffusion-based fin--tuning methods that rely on Q-values. Now, to spotlight the pivotal role of our method's core component-the guidance from the compatibility func- tion H, we replace it with Q and advantages A within the same training framework. The experimental results on the Robomimic-Cao task are illustrated in Fig. 3(left). Clearly, compared using Q-value and advantage, OT-guided train- ing demonstrates significantly faster convergence and supe- rior evaluation performance.
Comparisons with demo-augmented RL. Next, we com- pare OTPR with recently proposed RL methods for train- ing robot policies (not necessarily diffusion based) lever- aging offline data, including RLPD (Ball et al., 2023), Cal- QL (Nakamoto et al., 2024), and IBRL (Hu et al., 2023). These methods add expert data in the replay buffer and per- forms off-policy updates. We evaluate these methods on Konka-Kiithenn and RoboMimic environents. IBRL and Cal-QL are also pretrained with behavior cloning and of- fline RI. objectives, respectively. All of results are shown on Table 1. In the Franka-Kitchen domains, while Cal- QL demonstrates competitive performance, OTPR shows more impressive score improvements, rising from 61 to 92 in Kitchen-Come Ottchen-Complete-v) and from 59 to 79 in Kitchen- Mixed-vO. In contrast, other methods such as RLPD, IQL, and IBRL perform significantly worse, particularly in the -aattal-aaatil-- task, where OTPR leads with final score of 93. In the RoboMimic environment, OTPR contin- ues to excel, achieving high scores of 99 in Can-State and 98 in Square-State, showcasing its robustness across diverse scenarios. Although IBRL performs the best among the competitors, there remains a significant gap in performance.
Effect of the masked OT. OTPR incorporates masked Opti- mal Transport (OT) to utilize expert data as keypoints, guid- ing accurate distribution transport. As depicted in Fig. 3, OTPR-U, which lacks the mask matrix, exhibits instability and reduced efficiency, despite outperforming other main- stream methods. Notably, even without the mask, OTPR can still operate as fully functional offline RL algorithm by leveraging the compatibility function without reward.
7. Conclusion
This paper introduced OTPR, a novel method integrating optimal transport theory with diffusion policies to enhance the efficiency and adaptability of reinforcement learning fine-tuning. OTPR leverages the Q-function as transport cost and uses masked optimal transport to guide state-action matching, improving learning stability and performance. Experiments demonstrated OTPR's s superior performance across multiple tasks, especially in complex environments.
6.3. Ablation Experiments
Effect of the compatibility function. In the previous sec- tion, we have already demonstrated the advantages of OTPR
8

---

Score-Based Diffusion Policy Compatible with Reinforcement Learning via Optimal Transport
Future work will focus on scaling OTPR larger state- action spaces, and exploring its integration with other ad- vanced policy architectures.
Chen, S.-F. Wang, H.-C., Hsu, M.-H., Lai, C.-M., and Sun, S.-H. Diffusion model-augmented behavioral cloning. arXiv preprint 2xXx::332211333: 2023.
Chi, C., Xu, Z., Feng, S., Cousineau, E., Du, Y., Burchfiel, B., Tedrake, R., and Song, S. Diffusion policy: Visuomo- tor policy learning via action diffusion. The International Journal of Robotics Research, pp. 02783649241273666, 2023.
Acknowledgements
This work was supported by the National Science and Technology Innovation 2030 Major Project (Grant No. 20222ID2208800, and NSFC General Program (Grant No. 62176215).
Cuturi, M. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in neural information processing systems, 26, 2013.
Impact Statement
This paper presents work whose goal is to advance the field of reinforcement learning for robotic application. There are many potential societal consequences of our work, none of which we feel must be specifically highlighted here.
Daniels, M., Maunu, T., and Hand, P. Score-based gener- ative neural networks for large-scale optimal transport. Advances neural information processing systems, 34: 13955-12966 2021.
References
Dessein, A., Papadakis, N., and Rouas, J.-L. Regularized optimal transport and the rot mover's 's distance. Journal of Machine Learning Research, 19(15): 2018.
Ajay, A., Du, Y., Gupta, A., Tenenbaum, J. B., Jaakkola, T. S., and Agrawal, P. Is conditional generative modeling all you need for decision making? In The Eleventh International Conference on Learning Representations.
Fan, Y., Watkins, O., Du, Y., Liu, H., Ryu, M., Boutilier, C., Abbeel, P., Ghavamzadeh, M., Lee, K., and Lee, K. Rein- forcement learning for fine-tuning text-to-image diffusion models. Advances in Neural Information Processing Sys- tems, 36, 2024.
Ankile, L., Simeonov, A., Shenfeld, I, Torne, M., and Agrawal, P. From imitation to refinement-residuall rl for precise assembly, arXiv preprint arXiv:2407.16677, 2024.
Gao, X., Nguyen, H. M., and Zhu, L. Wasserstein con- vergence guarantees for a general class of score-based generative models. CoRR, abs/2311. 1000,, 2023. doi:
Asadulaev, A., Korst, R., Korotin, A., Egiazarian, V., Filchenkov, A., and Burnaev, E. Rethinking optimal transport in offline reinforcement learning. Advances in Neural Information Processing Systems, 37:123592- 123607, 2025.
1044555AARXI2231.10000
Gong, Z., Ding, P., Lyu, S., Huang, S., Sun, M., Zhao, W., Fan, Z., and Wang, D. Carp: Visuomotor policy learn- ing via coarse-to-fing fine autoregressive prediction. arXiv preprint arXiv:2412.06772. 2024.
Ball, P. J., Smith, I.,, Kostrikov, I,, and Levine, S. Ef-
ficient online reinforcement learning with offline data. In International Conference on Machine Learning, pp. 1577-1594. PMLR, 2023.
Goo, W. and Nickum, S. Know your boundaries: The necessity of explicit behavioral cloning in offline rl. arXiv preprint 22N6::220000066 2022.
Batzolis, G., Stanezuk, J.. Schönlieb, C.-B., and Etmann, C. Conditional image generation with based diffusion models. arXiv preprint arXiv:2111..336 2021.
Gu, X., Yang. Y., Zeng, W., Sun, J., and Xu, Z. Keypoint- guided optimal transport with applications in heteroge- neous domain adaptation. Advances in Neural Informa- tion Processing Systems, 35:14972-14499 2022.
Black, K., Janner, M., Du, Y., Kostrikov, I., and Levine, S. Training diffusion models with reinforcement learning. In The Twelfth International Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenReview.net, 2024.
Gu, X., Yang, L.,, Sun, J., and Xu, Z. Optimal transport- guided conditional score-based diffusion model. Ad- vances in Neural Information Processing Systems, 36: 36540-36552, 2023.
Chen, L., Lu, K., Rajeswaran, A., Lee, K., Grover, A., Laskin, M., Abbeel, P., Srinivas, A., and Mordatch, I. De- cision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems, 34:1 5084-15097, 2021.
Gupta, A., Kumar, V., Lynch, C., Levine, S., and Hausman, K. Relay policy learning: Solving long-horizon tasks via imitation and reinforcement learning. arXiv preprint urXiv:/9100 1956, 2019.

---

Reore-Based Diffusion Policy Compatible with Reinforcement Learning via Optimal Transport
Haarnoja, T., Zhou, A., Abbeel, P., and Levine, S. Soft actor-critic: Off-policy maximum entropy deep reinforce- ment learning with stochastic actor. In International conference on machine learning, pp. 1861-1870. PMLR, 2018.
Lu, C., Zhou, Y., Bao, E,, Chen, J., Li, C., and Zhu, J. Dpm-solver: A fast ode solver for diffusion probabilistic model sampling in around 10 steps. Advances in Neural Information Processing Systems, 14:5775-5787 2022.
Mandlekar, A., Xu, D., Wong, J., Nasiriany, S., Wang. C., Kulkarni, R., Fei-Fei, L., Savarese, S., Zhu, Y., and Martin-Martin, R. What matters in learning from offiin human demonstrations for robot manipulation. arXiv preprint
Hansen-Estruch, P., Kostrikov, L., Janner, M., Kuba, J. G., and Levine, S. Idql: Implicit q-learning as an actor- critic method with diffusion policies. arXiv preprint arXiv:2304. 10573, 2023.
2021.
Heo, M., Lee, Y., Lee, D., and Lim, J. J. Furniturebench: Re- produeible real-world benchmark for long-horizon com- plex manipulation. CoRR, abs/2305.1288 2023. doi: 10.48550AARXI.22355.222
Mark, M. S., Gao, T., Sampaio, G. G., Srirama, M. K., Sharma, A., Finn, C., and Kumar, Policy agnostic rl: Offline and online fine-tuning of any class and backbone. arXiv preprint arXh::241200666 2024.
Mees, O., Hermann, L.,, Rosete- Beate, E., and Burgard, W.
Ho, J., Jain, A., and Abbeel, P. Denoising diffusion proba- bilistic models. Advances in neural information process- ing systems, 33:6840-688 6851, 2020.
Calvin: A benchmark for Ianguage-conditione policy learning for long-horizon robot manipulation tasks. IEEE Robotics and Aulomation Letters, 7(3):7327-7334, 2022.
Hu, H., Mirchandani, S., and Sadigh, D. Imitation bootstrapped reinforcement learning, arXiv preprint arXiv:23 1,02/9,, 2023.
Montesuma, E. F., Mboula, F. M. N., and Souloumiac, A. Recent advances in optimal transport for machine learn- ing. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2024.
Janner, M., Du, Y., Tenenbaum, J., and Levine, S. Planning with diffusion for flexible behavior synthesis. In Interna- tional Conference on Machine Learning, pp. 9902-9915. PMLR, 2022.
Mu, T., Ling, Z., Xiang, F., Yang, D., Li, Tao, S., Huang, Z., Jia, Z., and Su, H. Maniskill: Generalizable ma- nipulation skill benchmark with large-scale demonstra- tions. In Vanschoren, J. and Yeung, S. (,ds.), Proceed- ings of the Neural Information Processing Systems Truck on Datasets and Benchmarks 1. NeurIPS Datasets and Benchmarks 2021, December 2021, virtual, 2021.
Jia, B., Ding, P., Cui, C., Sun, M., Qian, P., Huang, S., Fan, Z., and Wang, D. Score and distribution matching policy: Advanced accelerated visuomotor policies via matched distillation. arXiv preprint arXiv:2412.09225, 2024.
Nakamoto, M., Zhai, S., Singh, A., Sobol Mark, M., Ma, Y., Finn, C., Kumar, A., and Levine, S. Cal-ql: Calibrated offline rl pre-training for efficient online fine-tuning. Ad- vances Neural Information Processing Systems, 36, 2024.
Kakade, S. and Langford, J. Approximately optimal ap- proximate reinforcement learning. In Proceedings of the Nineteenth International Conference on Machine Learn- ing, pp. 267-274, 2002.
Kober, J., Bagnell, J. A., and Peters, J. Reinforcement learning in robotics: A survey. The International Journal of Robotics Research, 32(11):1238- 1274, 2013.
Park, Y., Margolis, G. B., and Agrawal, P. Position: Auto- matic environment shaping is the next frontier in RL. In Forty-first International Conference on Machine Learn- ing, ICML 2024, Vienna, Austria, July 21-27, 2024. Open- Review.net, 2024.
Kwon, D., Fan, Y., and Lee, K. Score-based generative modeling secretly minimizes the wasserstein distance. In Koyejo, S., Mohamed, S., Agarwal, A., Belgrave, D., Cho, K., and Oh, A. (eds.), Advances in Neural Infor- mation Processing Systems 35: Annual Conference on Neural Information Processing Systems 2022, NeurIPS 2022, New Orleans, LA, USA, November 28- December 9, 2022, 2022.
Psenka, M., Escontrela, A., Abbeel, P., and Ma, Y. Learning a diffusion model policy from rewards via q-score match- ing. In International Conference on Machine Learning, pp. 41163-41182, PMLR, 2024.
Rajeswaran, A., Kumar, V., Gupta, A., Vezzani, G., Schul- man, J., Todorov, E., and Levine, S. Learning complex dexterous manipulation with deep reinforcement learn- ing and demonstrations. In Kress-Gazit, H., Srinivasa, S. S., Howard, T., and Atanasov, N. (eds.), Robotics: Science and Systems XIV, Carnegie Mellon University,
Langley, P. Crafting papers on machine learning. In Langley, P. (ed.), Proceedings of the 17th International Conference on Machine Learning (ICML 2000), pp. 1207-1216, Stan- ford, CA, 2000. Morgan Kaufmann.
10

---

Score-Basee Based Diffusion Policy Compatible with Reinforcement Learning via Optimal Transport
Pitlsburgh, Pennsylvania, USA, June 26-30, 2018, 2018. doi: 10. 10.15677/SSS
Ren, A. Z., Lidard, J., Ankile, L. L.,, Simeonov, A., Agrawal, P., Majumdar, A., Burchfiel, B., Dai, H., and Simchowitz, M. Diffusion policy policy optimization. arXiv preprint arXiv:2409. 00588848 2024.
Ross, S. and Bagnell, D. Efficient reductions for imitation learning. In Proceedings of the thirteenth international conference on artificial intelligence and statistics, pp. 661-668. JMLR Workshop and Conference Proceedings, 2010.
Seguy, V., Damodaran, B., Flamary, R., Courty, N., Rolet, A., and Blondel, M. Large-scale optimal transport and mapping estimation. In ICLR 201 2018-Inme 80ttteernttt Confer- on Learning Representations, pp. 1-15, 2018.
Song, J., Meng, C., and Ermon, S. Denoising diffusion im- plicit models. arXiv preprint arXiv:2010.02502 2020a.
Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Er- mon, S., and Poole, B. Score-hased generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13455 2020b.
Song. Y., Durkan, C., Murray, I., and Ermon, S. Maxi- mum likelihood training of score-based diffusion models. Advances in neural information processing systems, 34: 1415 1428, 2021.
Villani, C. et al. Optimal transport: old and new, volume 338. Springer, 2009.
Vincent, P. A connection between matching and de- noising autoencoders. Neural Compulation, 23(7):1661- 1674, 2011. doi: 10.1162/NECO a 00142.
Wang, Z., Hunt, and Zhou, M. Diffusion policies as an expressive policy class for offline reinforcement learning. arXiv preprint ar arXiv:2208.0611 06193, 2022.
Yang, L., Huang, Z., Lei, F., Zhong, Y., Yang, Y., Fang, C., Wen. S., Zhou, B., and Lin, Z. Policy representation via diffusion probability model for reinforcement learning. arXiv preprint #rXiv:2355.33122 2023.
Zhao, Z., Kumar, V., Levine, S., and Finn, C. Learn- ing lnn-ggaann bimanual manipulation with low-cost hardware. In Bekris, K. E., Hauser, K., Herbert, S. L.., and Yu, J. (eds.), Robotics: Science and Systems XIX, Daegu, Republic of Korea, July 10-14, 2023, 2023. doi: 10.15607/RSS.200 2TT.SS222......1101 XIX.016.

---

Score-Based Diffusion Policy Compatible with Reinforcement Learning via Optimal Transport
Q-value Matrix Visualization
OT Plan Visualization
0.0
1.0
0.0
0.005
0.2
0.8
0.2
State Space
State Space
0.004
0.6
0.4
0.6
0.4
0.6
0.002
0.8
0.2
0.8
0.001
0.6
1.0
Action Space
Action Space
Figure This example demonstrates a clear and concise visualization Q-value matrix alongside corresponding estimated optimal transport plan.
A. Additional Details for Algorithm
A.1. Pseudo-codes algorithm for training twe Max
The pseudo-codes of the algorithm to learn the dual terms Way Mar a.k.a., potentials, are given in Algorithm (Seguy et al., 2018: Gu al., 2023).
Algorithm Algorithm for estimating potentials tuy Many
Input: Q-network and replay buffer B, nets Uea Vwn batch size b, learning rate expert demonstrations D (if available).
Output: Learned potential Mwi Vwg-
for iteration = 1,2.... do
Calculate the cost between & and a in
Sample a state batch {s1, s2,..., Sb} from B.
Sample action batch {a,, ; ) from B. if Expert data D is available then
Update + w1 + 0: Ei, Vay tay (s:) + 0, Fa(uuuu (mwn (8i), Uwa (a))) V.n Muyy (s:)
Update W2 + w2 + a Ei,
Var Baz (a,) + Fx (si); Dwrs (a,)) )) Vw Uwny (am)
else
Update W1 + w1 + a E
Vaz twr (s:) D. (Uws (5:), Uwo (a,) Uwss (s:)
Update W2 w2 + a Zi.j
Var Twny (az) Fx (Man (8i); Uwg (a;)) We Uang (am)
end if
end for
A.2. Training by fitting noise.
We consider the VE-SDE and the VP-SDE examples of forward SDEs. In the VE-SDE, f(a, t) and g
[1) where a > 0 is an increasing function of For the VP-SDE, f(a, 1) 3(1) and g(() v 3(), with
dt
((t) Bmin + (Bmax Btin)t. The distribution P:jo(a:lao) for a; given do is defined as:
[N(allas, (t)T), Pejo Pejo (a; ao) [a(arlase h(1), (1 (()))), for VP-SDE,
for VE-SDE,
(20)
where h(t)= -3 Bmin) t.Bmins and I the identity matrix.
12

---

Ccore-Based Diffusion Policy Compatible with Reinforcement Learning via Optimal Transport
We define the ofI the standard variation of Pejo(a, (a), specifically, of o²) for VE-SDE and of el(()) for VP- SDE. Using the reparameterization trick, given sampling(s, 8, a), we have a a+E for VE-SDE, and a₁ eth(t) a oye for VP-SDE, where e ~ ~N, I). Further, Va, log Pog written as
oe Therefore, the loss Ja,a for fitting noise can be
Js.a = E(,.)N(0.1)
Wt
(s + t,c + +el
(21)
of
For VE-SDE, (1)aa) = a, while for VP-SDE, (1) = e2 the scaled noise oz
Equation 21 indicates that so (363 2, t) is trained to match
B. Proofs
B.1. Proof of Proposition 4.1
-Q², let n is the solution to Proposition B.1. Given an optimal behavior policy T and a critic-based cost function
Eq. 8 with the Q cost function. Then it holds that: [RL(TT) = )
Proof. We use (u) and Supp(y) to refer to the support of μs and v. two subsets of S and A. respectively, which are also the set of values which s ~ μ and a ~ y can take. Given a point E Supp(u), the Monge problem would send the whole mass at x to unique location a E Supp(v). The primal state-conditioned Monge OT problem discripted with can be formulated as:
inf Esnn (s, subject (s) Supp(v) for all E Supp(μ)..
(22)
According to Lemma 6.1 (Kakade & Langford, 2002), we can use the performance difference lemma to compare performance of the two policies n" and T
J( = 1 - K
Eanp [A x*]]
(23)
Eany
(Q, T (s)) - (s)]
(24)
k
Egny (",(s, aESupp(3(11a))
- (s)) -V
(25)
k
In the setting of Proposition, T an optimal expert policy, V*(s) = max, Q3 (s. a), thus Supp(r"(-Is)) indicates the optimal actions a" from 7 which maximize Q Then have: J(=") [[x 0.
B.2. Proof of Proposition 4.2
1 8(s (a)m where 8 the Dirac delta function, then (DSM Eq. 7 can be Proposition B.2. Let C(s, a) u(e) reformulated as
(o EepprEe-pp a) (a (s)(ar;8,)) Va, log u [=
(26)
Furthermore,
u(s, a) = =({
(27)
is a joint distribution for marginal distributions and v.
Proof. We first prove Eq. 26, and then demonstrate that v(s, a) serves as joint distribution for the marginal distributions and v. (1) The right side of Eq. 26
13

---

Score-Based Diffusion Policy Compatible with Reinforcement Learning via Optimal Transport
Benn a (lllaarr |l 8,() Va, log (op(aaa)))) Earw (a (a;;s, s, Va. log
(28)
(29)
Earw / ö(s Scond
Ve, log (og
(30)
( 8cond (a), Va. log log
(31)
which is the definition of (0) in Eq. 7. (2) We demonstrate that the marginal distributions of v(8, a) are and follows. Firstly,
/ a(s.a) a) ds / 6(s Scond (a))(() ds v((a) / B(s Scoud (a) ds v(a)
(32)
Next, from the definition of 8(-), obtain 8(s Scand Scand(a)) Ea (anlaaa) δ(a a'). Then, we have
/ a((sa))da da
/
b(s- Scond Scond(a)) 6(s-Scona(a))
(a)) da
(33)
/
Σ
8(a' - a)v(() da
(34)
(al (a') ==a)
Σ
8(a' a)v(a) da
(35)
(a daaaaaa ==?
Σ
v(aa)
(36)
(a':8om((a"==
= -p(s)
(37)
B.3. Proof of Theorem 4.3
Theorem B.3. For p. consider the forward SDE da, f(a,, t)dt g(t)dw with ag ((1ss) and E [0, T Let ve (ar)s) be the distribution of an and Jcsu(() = E.,, (a) (500) Va, log (a, ())] then we have VoJUNSM((O) (oJesu(O)).
Proof. To establish the equivalence between (HD)M(O) and (O), we start by examining the difference between the two objective functions:
(!) (f) Ew; Eap H an) (a (at; s, t) Va, log (a (a)) JCsM(()
Eau.EE.,pBann Ba || (a;; s, t) - Va, log vejo (a, (an)) JcsM( (38)
Since ag d is Markov Chain in the forward SDE process, the distribution Pajo ag s) of a simplifies to vaj(aalaao 8) (ap(aa)) which is the distribution of a, by the forward SDE da, [(a,, l)d/ g(()dw with initial state ag. Then, we have
(O) JCSM (8) =Ew,
llo s,t) Va. log ag s) E.W,E los 8, t) - Va. log lo
(39)
According to (Vincent, 2011), given any s and f, we have
Bao-i(au (x) Bar (a.. (lsollls - Va. log (appoaraaa s))
= -Eg vn (a: so (a 8, Va. Vog (a Co,t.
(40)
14

---

Score-Based Diffusion Policy Compatible with Reinforcement Learning via Optimal Transport
where Ca. a constant 0 depending and l. Substituting this result into the previous Eq. 39, we get
JHDSM (0) JcsM (0) = E.-n,B
(41)
Since the right-hand side a constant have conclude that
V VoTIDSM(()) (B).
(42)
C. Details for Experiments
All experiments are conducted on an NVIDIA Tesla A100 80GB GPU, and all fine-tuning methods use the same pre-trained policy. In the pretraining, the observations and actions are normalized to [0, 1] using min/max statistics from the pre-training dataset. No history observation (pixel, proprioception, or ground truth object states) is used. The deffusion policy is trained with learning rate 1e 4 decayed to le with a cosine schedule. weight decay le-6 and 50 parallelized. For Kitchen and Robomimic tasks, epochs is 8000 and batch size is 128; for CALVIN tasks, epochs is 5000 and batch size is 512.
C.1. Details and Hyper-parameters for OTPR
Details for training Uur Vu- The architecture of both uhe and th a two MLP. À is set 1e 5. The batch size set 64. We employ the Adam algorithm update the parameters with le 6 learning rate.
Details for training sg. We take the VP-SDE (Song et al., 2020b) the forward SDE. In inference, we take the sampling method in DDIM (Song et al., 2020a) to perform the reverse SDE to generate action. The observations and actions are normalized to jo, 1] using min/max statistics from the pre-training dataset. For diffusion-based policies, we use MLP with two-layer residual connection similar to DPPO.
Table 2. Hyper- Seppee-arramters for OTPR
Task
Parameter
Franka-Keticher
CALVIN
Robomimic-
Robomimic-
Buffer size Actor Learning Rate
1000000
250000
250000
250000
1.00E-05
1.00E-05
1.00E-05
1.00E-05
Discount Optimizer L
0.99
0.99
0.999
0.999
Adam
8
8
T
20
20
20
20
0.7
0.7
0.7
0.7
Actor Batch Size Critic (Q and V) Hidden Layer Sizes
1024
1024
1024
1024
[512, 512, 512]
[512, 512, 512]
[256, 256, 256]
[256, 256, 256]
Critic (Q and V) Batch Size
256
256
256
256
C.2. Details and Hyper-parameters for Baselines
DPPO For the state-based tasks Robomimic and FrankaKitchen, we trained DPPO-MLP following the original paper's specifications, using an action chunking size of 4 for Robomimic and for FrankaKitchen. For the pixel-based task CALVIN, trained DPPO-ViT-MIL ViT-MLP with an action chunking size of 4.
IQLI We employ the IDQL-Imp version of IDQL., wherein the Q-function, value function, and diffusion policy are refined through new experiences. For Robomimic tasks, we employ the same network architecture as OTPR, while the original IDQL architectures are preserved for Franka-Kitchen and CALVIN. For the IQL 7 expectile, we set to 0.7 for each task.
DQL We set the weighting coefficient 0.5 for Robomimic, 0.005 for
and 0.01 for CALVIN.
15

---

Score-Basee Boore-Based Diffusion Policy Compatible with Reinforcement Learning via Optimal Transport
IBRL We adhere the original implementations" hyperparameters, with wider (1024) MLP layers and dropout during pre-training.
Cal-QL We set the mixing ratio to 0.25 for Tranka-Kitchen and 0.5 for CALVIN and Robomimic.
D. Additional Experimental Results
D.I. Pixel-based Tasks in Robomimic
Can-Pixel
Square-Poxel
Transpoor-PP
10
IDGL
mai
0.0
0.5
DQ OTFR
05
DCL
Nomellzze Score
01P1
Normakzed
Normalizee
0-1
0.4
DPRO
02
IDGL
0.2
02
DCL
O1P1
Interaction Steps (K203K)
Interaction Stops (=1MM
Interaction Stepa (x500K)
Figure 5. Learning curves of online fine-tuning with various methods on pixel-based Robomimic Tasks, Ohserve that OTPR largely always dominates attains similar performance the next best method. Other methods for fine-tuning diffusion policies (IDQL, DQL. DPPO) are bit unstable, and perform substantially worse.
D.2. LIBERO-Long Tasks
Table 3.
Success rate of online fine-tuning on task in Libero-Lone
Libero-Long
Task
Task 2
Task 3
Task 4
Task 5
Task 6
Task 7
Task 8
Task 9
Task 10
DP
0.65
0.5
0.85
0.7
0.6
0.8
0.6
0.65
0.1
0.7
OTPR
0.85
0.65
0.85
0.9
0.8
0.9
0.7
0.9
0.35
0.85
D.3. Qualitative Toy Experiments
Compatibility Function Validation
Objective: Validate the accuracy of the compatibility function to verify whether Algorithm can effectively learn the dual term.
Setup: We conducted experiments in 2D space, where Algorithm 2 applied to a random dataset consisting of Gaussian source distribution and multi-modal target distribution (8 Gaussian), with the Euclidean distance serving the cost function.
Fig. 6 left: We visualize the source distribution as colored level sets and the target distribution as randomly sampled points.
Result (Fig. middle): We constructed 200 sample points from the source distribution and 2000 samples from the target distribution, then paired them using the compatibility function The compatibility function H successfully matches source samples (yellow) target samples (green) with low cost, confirming its effectiveness.
Result 2 (Fig. right): Generated samples obtained by sampling from the source distribution. We learn optimal map as neural network by approximating the barycentric projection of the OT plan from Algorithm 2. Generated samples closely match the target distribution.
OT-Guided Diffusion Policy Verification
Objective: Assess diffusion policy's ' ability to recover target distributions under OT guidance.
16

---

Score-Based Diffusion Policy Compatible with Reinforcement Learning via Optimal Transport
Setup: We leverage 2 synthetic 2D datasets (8Gaussian and swissroll) further verify the effectiveness of OT guided diffusion policy. Each dataset contains data points paired with specific Q values (Figure Fig. 7 left).
Results (Fig. right); The ultimate samples generated by the diffusion model, which closely match the ground-truth target distribution.
larget [Action] and Source (State) Distribution
larget (Actioni and Source (State) Distribution
Generated Samples
Figure Left: A continuous Gaussian source distribution (colored level sets) and modti- target samples from 8 Gaussian distribution (yellow). Middle: Samples from source distribution (yellow), paired samples (green) and unpired samples (blue) from target distribution, Right: Generated samples by approximating the barycentric projection the OT plan.
04
Figure 7. Illustration of the 8gaussians and swissroll dataset (left) and the generation result of diffusion model (right)
17


Reviews:
[{"review_id": "sVPJtm7mtM", "reviewer": "Reviewer_L6dM", "summary": "This paper proposes to reformulate offline-to-online diffusion policy training with optimal transport. It views policy as a transport from the state distribution to action distribution, using the (negative) Q-function as a a transport cost and treating the policy as an optimal transport map. The authors show that the score matching objective of the diffusion policy training can be augmented with a weighting function that constitutes the joint distribution of state and action pairs. They further show that this weighting function can be relaxed with a compatibility function that involves the Q function as well as some dual variables $u_w(\\mathbf{s})$ and $v_w(\\mathbf{a})$ derived from the dual form of regularized optimal transport. As this relaxed weighting function gives zero weight to state-action pairs where the state is in expert demonstrations while the action is not, they coin their objective Masked Optimal Transport. The authors also did an analysis to illustrate that the the proposed training algorithm is optimizing an upper bound of the distance between the diffusion policy and the optimal transport plan. Their experiments demonstrate clear offline-to-online improvement.\n\n### After rebuttal ###\nThe authors addressed my major concerns. I updated my rating accordingly.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer, \n\nThanks for your thorough review and constructive feedback on our submission. \n\n## C1. Claims: \n*The connection between OT and RL appears superficial, and the role of Q-learning methods requires clarification.*\n\n**R1:** While the cliam has already garnered recognition from other reviewers (e.g. *OTPR unifies insights from generative modeling, reinforcement learning...* by Reviewer ZmNb), we provide these focused responses to your concerns:\n\n**Theoretical Equivalence:** \nTo support this claim, we have provided proof in the Appendix B.1 demonstrating the equivalence between the OT plan and the optimal policy. This theoretical result forms a solid foundation for our approach and shows that the connection is not merely superficial. \n\n**The Relevance of Q as cost:** The use of the Q-function as a transport cost has built on prior static offline RL works (e.g., [1]). Our work uniquely integrates this idea with diffusion policy to enable smooth offline-to-online finetuning. \n\n**Q-Learning Compatibility:** OTPR is an algorithm-agnostic framework incorporating any existing Q learning algorithms.\n\n[1] Asadulaev, Arip, et al. \"Rethinking Optimal Transport in Offline Reinforcement Learning.\" Advances in Neural Information Processing Systems (2024)\n\n## C2. Generalization: \n*The mask appears to be a restriction to generalization.* \n\n**R2:** We appreciate the reviewer’s attention to generalization capabilities. We believe the authors have understood that our algorithm consists of two parts: (1) estimating the OT plan to provide a compatibility function H(s,a), and (2) using H to guide the  diffusion policy optimization. \n\nThe mask is introduced during (1) and designed to fully leverage the paired state-action data from expert demonstrations to improve the accuracy of $H$. The effectiveness of this key-point guidance approach has been validated in some domain adaptation studies.  Crucially, the mask does not constrain the policy’s action space. The mask focuses only on actions that are known to be effective in specific seen states, as demonstrated by expert data, and does not directly influence the learning of the policy and inference. During online fine-tuning, the policy can freely explore novel actions in both seen and unseen states.\n\n## C3. Additional Experiments: \n\n**R3:**  We additionally added two 2D qualitative toy experiments and three pixel-based Robomimic tasks results in https://anonymous.4open.science/r/OTPR_Supplementary-5F55/README.md. See response to reviewer iuX1 for a detailed description. \n\n## C4. Ablation Results of Masked OT:  \n\n**R4:** The masked OT mechanism is designed to refine the estimation of the OT plan by leveraging expert demonstrations as high-confidence priors. Its contributions are twofold: (1) reducing the occurrence of poor solutions, and (2) improving overall stability across multiple fine-tuning evaluations. As shown in Figure 3, these improvements are captured by the consistent trend in mean performance and the reduction in variance (tighter confidence intervals). We also supplement additional ablation experiments (Figure 4 in https://anonymous.4open.science/r/OTPR_Supplementary-5F55/README.md). \n\n## C5. Notation Suggestion  \n\n**R5:** Our intention was to use **bold font** to denote variables, while regular font indicates a specific sampled instance of that variable, following [1]. This distinction also helps differentiate between actions at each timestep and those at each denoising step. We acknowledge that this notation might have caused confusion, and we will unify the font style and fix typo errors in the revised manuscript. \n\n[1] Ada, et al. \"Diffusion policies for out-of-distribution generalization in offline reinforcement learning.\" IEEE Robotics and Automation Letters\n\n## C6. Questions for Authors \n\n**R6:**\n\n**Masking Scheme and Action Discovery**  As we mentioned in our response to the Generalization Concerns (R2), the Mask mechanism is introduced to assist in estimating the OT plan. Nevertheless, it is the compatibility function H(s,a) that truly evaluates the sampled actions and guides the optimization of the diffusion policy. The computation of H is a comprehensive process that includes both the potential functions u and v, along with the Q-values, ensuring robustness. Actions with equally high Q-values are also permitted to achieve correspondingly high H-values.Thus, this does not preclude the model from discovering other good actions. \n\n**Demo-augmented RL Algorithms Performance** Similar reproduction results were also observed in the DPPO paper. RLPD is an online RL algorithm leveraging offline data. Since it does not involve a pre-train process, we set its offline performance to 0. For IBRL, we strictly adhered to the original implementation protocol and its behavioral cloning objective during the offline training stage. Its offline performance is bad, which may be attributed to the presence of noise and multi-modality of data.", "claims_and_evidence": "The major claim is that the OT perspectives help bridge diffusion policy with RL, which seems to be only weakly supported by showing the OT is using the (negative) Q function as a cost function. However, this connection appears to be fairly superficial since which Q learning method is used and what policy induces its Q value is not discussed.", "methods_and_evaluation_criteria": "Viewing policy learning as an optimal transport from state distribution to action distribution makes sense conceptually. However, how to define the transportation cost is tricky. The proposed method has an implicit assumption that the key states, i.e. states covered by expert demonstrations, can only be paired with the associated actions in demonstrations. This appears to be a restriction to generalization. \n\nThe evaluation criteria, i.e. the learning curves and the final performance, are standard in the literature.", "theoretical_claims": "The proofs make sense to me at a high level. But I didn't check the details.", "experimental_designs_or_analyses": "The authors conducted experiments in online funetuning on 6 RL tasks, which seems not very sufficient. Fig 2 shows the proposed method has significant effect in most of them. \n\nThey also compare with methods that use both offline data that are not necessarily optimal and expert data in offline training and online finetuning. The proposed method appears to be the only effective one in offline training, and performance is boosted after online finetuning.", "relation_to_broader_scientific_literature": "The problem of offline-to-online learning is crucial for robotics. This paper is an attempt in this direction.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "\\+ The ablation of different compatibility method clearly illustrate the effectiveness of the proposed H function. \n\n\\- The formulation appears to be complicated. I appreciate the authors' effort in introducing the perspectives of optimal transport to diffusion policy. But the exposition of the paper involves non-intuitive formal notions, making it less accessible. The proposed algorithm also look complicated. At each iteration, the dual variables, which are MLPs need to be \"optimized\" as the first step of learning. \n\\- The proposed masked OT does not introduce significant gains according to FIg3.", "other_comments_or_suggestions": "The notations appear to be very unorganized. For example, $\\mathbf{a}$ and $a$ are used interchangeably in Section 3.2. \n\nThere is a typo in Eq 11.", "questions_for_authors": "The mask scheme that only allows states in the expert demonstration to be matched with the actions appears to be pretty restrictive. Does that mean the model won't discover any actions that are equivalently good to actions that appear in the demonstration? \n\nHow is the offline learning in Table 1 performed? Why does the proposed method performs well in offline learning while most of the baseline methods failed?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "gqMtUw8xC5", "reviewer": "Reviewer_iuX1", "summary": "The paper proposes OTPR that leverages optimal transport for fine-tuning diffusion policy in RL. Q function is treated as the transport cost and the policy is considered the transport map. Masked OT with resampling is also applied to improve training stability. Experiment results show generally improved performance compared to other diffusion RL fine-tuning methods and demo-augmented RL methods.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer iuX1, \n\nThank you for your positive review and constructive feedback on our submission. We appreciate your recognition of our work and are glad to hear that you found our approach and experimental results valuable. Below, we address your specific comments and questions. \n\n## C1: Qualitative Experiments in Claims And Evidence\n\n**R1:** As suggested, we designed two illustrative 2D toy experiments to visually validate our method’s core components. Full visualizations are available in [Supplementary Section 1] (https://anonymous.4open.science/r/OTPR_Supplementary-5F55/README.md).\n\n### 1. Compatibility Function Validation\n\n**Objective:** Validate the accuracy of the compatibility function to verify whether Algorithm 2 can effectively learn the dual term. \n\n**Setup:** We conducted experiments in a 2D space, where Algorithm 2 was applied to a random dataset consisting of a Gaussian source distribution and a multi-modal target distribution (8Gaussian), with the Euclidean distance serving as the cost function. \n\n**Figure 1.1 left:** We visualize the source distribution as colored level sets and the target distribution as randomly sampled points. \n\n**Result 1 (Figure 1.1 middle):** We constructed 200 sample points from the source distribution and 2000 samples from the target distribution, then paired them using the compatibility function H. The compatibility function H successfully **matches source samples (yellow) to target samples (green) with low cost**, confirming its effectiveness. \n\n**Result 2 (Figure 1.1 right):** \nGenerated samples obtained by sampling from the source distribution. We learn an optimal map as a neural network by approximating the barycentric projection [1] of the OT plan from Algorithm 2. Generated samples closely match the target distribution.\n\n### 2. OT-Guided Diffusion Policy Verification\n**Objective:** Assess diffusion policy’s ability to recover target distributions under OT guidance.\n\n**Setup:** We leverage 2 synthetic 2D datasets (8Gaussian and swissroll) used in [2] to further verify the effectiveness of OT guided diffusion policy. Each dataset contains data points paired with specific Q values (**Figure 1.2 left**). \n\n**Results (Figure 1.2 right):** The ultimate samples generated by the diffusion model, which closely match the ground-truth target distribution.  \n\n[1] Seguy, Vivien, et al. \"Large-Scale Optimal Transport and Mapping Estimation.\" ICLR 2018-International Conference on Learning Representations. 2018.\n\n[2] Lu, Cheng, et al. \"Contrastive energy prediction for exact energy-guided diffusion sampling in offline reinforcement learning.\" International Conference on Machine Learning. PMLR, 2023.\n \n\n \n\n## C2. Experimental Details \n\n**R2:**  Thanks for the reviewer’s suggestion, we will supplement the more relevant details of the pre-trained diffusion policy in Section 6 or Appendix C if the page number is limited. \n\n*In the pretraining, the observations and actions are normalized to [0, 1] using min/max statistics from the pre-training dataset. No history observation (pixel, proprioception, or ground-truth object states) is used. The deffusion policy is trained with learning rate 1e-4 decayed to 1e-5 with a cosine schedule, weight decay 1e-6 and 50 parallelized. For Franka-Kitchen and Robomimic tasks, epochs is 8000 and batch size is 128; for CALVIN tasks, epochs is 5000 and batch size is 512.*  \n\n \n\n## C3. Missing Reference \n\n**R3:** We appreciate your suggestion to include **QSM** in our manuscript, which is closely related to our approach and offers valuable insights into the field. We will reference this paper appropriately in the Related Work section. \n\n \n\n## C4. Additional Robomimic Task \n\n**R4:** We sincerely appreciate this constructive feedback. To address the request, we have conducted additional experiments on pixel-based tasks in Robomimic, including the Transport task highlighted by the reviewers. We used ResNet as the visual encoder, similar to the setup in CALVIN. The learning curves and comparisons with baseline methods can be found at (https://anonymous.4open.science/r/OTPR_Supplementary-5F55/README.md).\n\nAs shown in Figure 2 from the linked results, our method still dominates or attains similar performance to the next best method. The results align with our original pixel-based CALVIN experiments, further validating our framework’s capability to handle visual inputs while preserving offline-to-online finetuning strengths.", "claims_and_evidence": "The claim of improved performance is supported by results on the Franka-Kitchen, CALVIN, and Robomimic tasks, which are manipulation tasks and generally more challenging than dense-reward mujoco locomotion tasks.\n\nHowever, the paper lacks qualitative discussion on its effectiveness. I would vote for strong accept if there are qualitative demonstrations of improved performance from leveraging optimal transport, e.g., in a carefully designed toy problem.", "methods_and_evaluation_criteria": "I appreciate the author considering more challenging RL benchmarks including the vision-based CALVIN.", "theoretical_claims": "I skimmed over the proofs in Appendix B and did not see any egregious error.", "experimental_designs_or_analyses": "I do not spot any particular issue with the experimental design.", "relation_to_broader_scientific_literature": "Diffusion RL fine-tuning is a very important area that requires active research as we have seen many successes in training diffusion policy with behavior cloning, and RL fine-tuning is critical to further improve the robustness. This paper bridges the gap between Ren et al. in online diffusion-based RL and many existing offline diffusion-based RL methods.", "essential_references_not_discussed": "A missing reference is Psenka et al., Learning a diffusion model policy from rewards via q-score matching, which also considers online diffusion-based RL with Q function.", "other_strengths_and_weaknesses": "The paper is well-written and figures (especially ones in the experiment section) are nicely done.", "other_comments_or_suggestions": "I suggest putting more experimental details, such as how the diffusion policies are pre-trained, in the beginning of experiment section.", "questions_for_authors": "Did you try pixel-based robomimic tasks? Would also be interesting to try Transport from robomimic, which is significantly more challenging than the tasks considered.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "gG3SLGZqrJ", "reviewer": "Reviewer_ZmNb", "summary": "This paper introduces OTPR, a novel method that integrates optimal transport theory with diffusion policies to enhance the robustness and performance of imitation learning models through online interactions with the environment . The core algorithmic idea involves leveraging the Q-function as a transport cost and viewing the policy as an optimal transport map to establish a connection between optimal transport and RL . OTPR also introduces masked optimal transport to guide state-action matching using expert data as keypoints and a compatibility-based resampling strategy to improve training stability. The paper's main findings from experiments on three simulation tasks demonstrate that OTPR consistently matches or outperforms existing state-of-the-art methods, especially in complex and sparse-reward scenarios, highlighting its effectiveness in combining imitation learning and reinforcement learning for versatile and reliable policy learning.\n\n## update after rebuttal\nI confirm my score. Authors addressed comments and added clarity and results to the original submission.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer ZmNb, \n\nThank you for your positive review and valuable feedback on our submission. Your comments have provided us with clear directions for improvement, and we are committed to addressing them in the revised version of our paper. Below, we address your main concerns and questions. \n\n## C1: Questions For Runtime \n\n**R1:** We thank the reviewers for their suggestion. Since all experiments were conducted under the same computational resource configuration, we directly provide the wall-clock time statistics for comparison in the Table of Supplementary Section 3 (https://anonymous.4open.science/r/OTPR_Supplementary-5F55/README.md).\n\nOur algorithm incurs approximately **8%–11%** more runtime compared to other baselines. This additional overhead stems from learning the dual term, but it is negligible considering the performance improvements achieved. \n\n## C2: Weakness (Real World Experiments)\n\n**R2:** We sincerely appreciate the reviewer’s valuable feedback regarding the importance of real-world validation. We fully acknowledge that performance in simulated environments may not directly translate to real-world scenarios. However, due to the time constraints of the rebuttal period and the current resource limitations, we are unable to conduct real-world experiments at this stage. As demonstrated in DPPO deployments, policies trained in high-fidelity simulations can achieve zero-shot transfer to physical hardware without real-data fine-tuning. So we have justifiable confidence that our approach can achieve comparable or better zero-shot transfer performance under identical experimental settings. We are currently actively developing a simulation-to-real evaluation platform and commit to publishing empirical validation results by the Camera-Ready deadline.\n \n\n## C3: Typos and Terminology \n\n**R3:** Thank you for catching the typo in Section 3.2. We will correct \"Reinfrocement Learning\" to \"Reinforcement Learning\" in the final version of the paper. We appreciate your attention to detail. \n\n\nThank you again for your time and insightful comments.", "claims_and_evidence": "The paper’s claims that OTPR (1) integrates optimal transport with diffusion policies to stabilize fine-tuning, (2) achieves notable performance gains over baseline methods, and (3) remains robust in sparse-reward environments—are generally well-supported. Multiple experiments on robotic tasks with varying difficulty, along with comparisons to several recent diffusion-based and demo-augmented RL baselines, underscore OTPR’s improvements. The authors also include ablations (e.g., masked vs. unmasked OT) to illustrate how each component contributes to the final performance.", "methods_and_evaluation_criteria": "The paper takes Robomimic, Franka-kitchen and CALVIN as the evaluation benchmark, which are suitable for assessing multi-step tasks under distribution shifts, making the evaluation well-matched to the paper’s goals.", "theoretical_claims": "There are no apparent flaws in the theoretical arguments as presented.", "experimental_designs_or_analyses": "The paper’s ablation studies that remove “masked OT” or the compatibility-based resampling strategy appear valid for isolating each component’s effect on performance. The comparisons to both diffusion-based and non-diffusion-based RL methods also lend credibility to the authors’ main performance claims. However, finer details like runtime comparisons are treated less thoroughly, so it’s hard to assess how sensitive or resource-intensive the method might be in broader settings. Overall, the key experiments are logically consistent, and no major flaws are evident in their design or analysis.", "relation_to_broader_scientific_literature": "OTPR unifies insights from generative modeling, reinforcement learning, and OT to tackle distribution mismatch more effectively.", "essential_references_not_discussed": "No", "other_strengths_and_weaknesses": "**Strengths:**\n\n- OTPR is the first method that combines diffusion based policy, RL and OT.\n- Propose Mask Optimal Transport in RL finetuning.\n\n**Weakness:**\n\n- Still need to test in Real-world environment like DPPO. Because performance achieved in a simulated environment may not directly translate to equivalent real-world performance.", "other_comments_or_suggestions": "Typos in Section 3.2: “Reinfrocement Learning” should be “Reinforcement Learning”.", "questions_for_authors": "How long is fine-tuning the diffusion policy with OTPR?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "LDqnkY6U1Z", "reviewer": "Reviewer_an3h", "summary": "The method proposed in this paper is quite complex, but I will do my best to summarize it:\n* The authors attempt to integrate diffusion policies (which are typically learned via imitation on expert data) with online environment interactions.\n* To this end, the authors make the following key observation (Prop 4.1): minimizing the optimal transport problem from expert state distribution -> expert action distribution, with the cost being the negative of the expert Q function, yields a policy of equivalent performance to the expert.\n* So our algorithm is roughly repeats the following steps:\n\na) Learn the optimal transport problem solution (as dual functions, although we can recover the primal), given the current learned Q-function and data from the replay buffer.\n\nb) For each inference step, sample a bunch of state-conditioned actions from the diffusion policy. Weight the probabilities state-action pairs by the coupling function H(s, a) from the OT problem solution; intuitively, we are doing weighted sampling based on how \"likely\" it is that an optimal policy following the current Q-function would take each action. Select an action according to the reweighted probabilities and execute in the environment.\n\nc) After doing a rollout, update the Q-function based on the observed rewards.\n\nd) Update the diffusion policy using equation (10); this is similar to advantage-weighted regression, where we are upweighting the state-action pairs where H(s, a) is high.\n\n* Basically, if my understanding is correct, we solve the OT problem transform a Q function into this coupling object H(s, a), which tells us how \"good\" a state-action pair is. We then use H to both weight (s,a) pairs in training the diffusion policy and guide sampling at inference time.\n* None of the above strictly requires expert demonstrations. But if we do have expert demonstrations, these come in the OT solving step via a masking function which gives keypoints to guide the OT solution.\n* The authors show convincing performance of this method across a range of problems.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer an3h, \n\nThank you for your detailed review and constructive feedback on our submission. Below, we address your main concerns and questions. \n\n## C1: Claims And Evidence: \n\n**R1:** We appreciate the feedback regarding evidence granularity. To substantiate our claims, we have conducted additional experiments on the Robomimic-Square environment (see Figure 4 in  https://anonymous.4open.science/r/OTPR_Supplementary-5F55/README.md). The observed \"instability\" manifests as OTPR-U producing worse fine-tuning outcomes more frequently across multiple trials, which is visually reflected in the increased variance of success rates. We agree the original terminology could be misinterpreted and will revise \"instability\" to \"statistically significant variance increase\".\n\n## C2: Weaknesses \n\n**R2:** We acknowledge that the confusing math exposition are weaknesses. We'll work to improve the clarity of the mathematical derivations and the overall exposition. Below, we address your main confusion: \n\n**1.** Here, we assume the existence of an stationary optimal behavior policy, $\\pi^\\beta$ in standard MDP. In this context, the \"state distribution\" is formally defined as the idealized stationary distribution induced by the expert policy. This theoretical abstraction serves to formulate the imitation learning problem as a Monge optimal transport (OT) problem between distributions. We will clarify it.\n\nHowever, we explicitly recognize that: The true stationary  state/action distribution is intractable to compute directly in practice, especially for complex environments. Real-world expert datasets only provide finite samples from this distribution. So we introduce a neural network-based approximation in 4.3 to estimate the underlying OT plan on expert dataset.\n\n**2.** The reviewer is correct; we did not assume a one-to-one mapping between state and action here. If the term \"condition data\" has caused confusion, we are willing to revise it to \"all conditional states.\" \n\n**3.** Thank you for the reminder. Indeed, there was an error on our part, and we will correct it. \n\n**4.** The \"initialized Q-network\" refers to a newly instantiated Q-network within the deep RL that has not undergone any training iterations.  To enhance algorithmic transparency, we propose adding an explicit initialization declaration step in the pseudocode. \n\n**5.** The reviewer is correct. Thank you for the reminder and we will correct it. \n\n**6.** We directly employed **Lemma 6.1** from (Kakade & Langford, 2002). As suggested by the reviewers, \nwe will explicitly clarify the equivalence between the advantage function $A$ and $R$ to enhance the readability of the derivation. \n\n**7.** Thanks to the reviewer for the suggestion. In order to facilitate a clear understanding of Regularized OT Dual for the reader, we will additionally provide an introduction to Kantorovich duality and provide a proof in the appendix, the core idea is rewrite the constrained infimum problem as an **inf** **sup** **problem**, and exchange the two operations by formally applying a minimax principle,i.e. replacing an “inf sup” by a “sup inf”. \n\n## C3: Question 1 For Authors \n*H(s, a) as a Better Version of the Q Function?* \n\n**R3:** First, we would like to clarify the reviewer's understanding. As described in Equation 15, **the compatibility function $H$ incorporates two estimated dual variables $u$ and $v$, rather than the distributions $\\mu$ and $\\nu$**. We sincerely apologize for the unintended ambiguity caused by the visual similarity between English and Greek letters in our notation. In the revised manuscript, we will adopt more distinctive notation to better differentiate these variables. \n\nSecondly, imitation learning aims to acquire a deterministic state-action mapping from expert demonstrations, while reinforcement learning focuses on learning a Q-function to evaluate action sampling, even for suboptimal actions. Our approach bridges these two objectives by introducing a compatibility function $H$ that establishes a **soft** state-action coupling relationship from a data distribution perspective. Specifically, the proposed compatibility function offers two advantages: (1) For clearly advantageous state-action pairs (e.g., those from expert demonstrations), $H$ provides precise guidance; (2) For novel state-action pairs that may emerge during RL exploration (particularly those absent from the training data), where Q-value estimates could be unreliable, the potential functions derived from the optimal transport plan serve as a corrective mechanism to adjust these estimates. \n\n## C4:  Question 2 For Authors\n*In figure 4, why isn't optimal transport assigning a unique action for every state? Am I right that this is the relaxed scheme (eq (2)) instead of the Monge formulation (eq (1))?* \n\n**R4:** The reviewer is correct. As indicated in the title of Figure 4, we indeed visualize the Optimal Transport Plan in this figure.", "claims_and_evidence": "Claim: the OTPR policy outperforms competing baselines.\n- Evidence: strong. Across all experiments in Figure 2, OTPR is either clearly best or tied. It is also clearly the most consistent--no other baseline performs close to OTPR for all scenarios. The performance improvement compared to other demo-augmented RL algorithms in Talbe 1 is quite substantial.\n\nClaim: the guidance from coupling function H (for training and sampling at inference) is better than that from the Q function or the advantage function.\n- Evidence: moderate to strong. The authors compare these for a single robomimic-can task (Figure 3 left), where the difference is substantial, but I'd like to see this for other environments as well.\n\nClaim: OTPR without the expert-demonstration masking \"exhibits instability and reduced efficiency\" (420).\n- Evidence: weak to moderate. The distinction between the masked and unmasked versions of OTPR in Figure 3 right are marginal, and again only for a single environment. I don't see what suggests that the unmasked version \"exhibits instability\".", "methods_and_evaluation_criteria": "Yes, and they are thorough.", "theoretical_claims": "I did have a look at the proof of Proposition B.1 and it seems reasonable, although I have some confusions about the problem setup which means I wasn't able to confidently check it (see below).", "experimental_designs_or_analyses": "The experiments are reasonable. I like that the hyperparameters for the OTPR are consistent across all the tasks in Table 2 (Appendix C.1), meaning that the method doesn't require extensive hyperparameter tuning.", "relation_to_broader_scientific_literature": "This paper is addressing a really important topic in robotics right now: how can we improve diffusion policies with online data? This is a challenging problem for diffusion policies in particular. Typically, in online RL you need to take the gradient of the action likelihood with respect to the policy parameters; but for diffusion policies, there is no closed-form action likelihood as actions produced via an iterative sampling scheme.\n\nA natural idea is to weight the regression targets in the diffusion loss by the Q function or the advantage. I'm mentally understanding the submission as proposing a better version of this using an OT coupling function H(s, a) instead. The experiments are pretty compelling, even though I don't fully understand the derivation.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "My main reservation is that the paper is pretty hard to understand. There's a lot going on, the math exposition is pretty confusing, and I think there are quite a few typos. Here are some of my confusions:\n1. In line 166 RHS, the authors say that the policy \"moves mass from the state distribution $\\mu(s)$ to the distribution of actions $\\nu(a)$.\" What are these state and action distributions? Are they the stationary distributions under the expert policy? Effectively, the marginals distributions of the (s,a) pairs from the expert dataset?\n2. In section 4.2, the authors denote the \"condition data\" of an action as $s_{cond}(a)$. This makes very little sense to me, as the same action can be taken in multiple different states, so there isn't one unique state that produces a particular action.\n3. I think in the RHS of 13 the i index should appear in the numberator, instead of k?\n4. The authors say that the input to Algorithm 1 is an \"initailized Q-network\"-- what does this mean? Is it pretrained?\n5. In (22), should it read $\\pi(s) \\in \\text{Supp}(\\nu)$?\n6. In the proof of Prop B.1, the relevant result from Kakade & Langford should be reproduced to make everything self-contained.\n7. When introducing the Regularized OT Dual, it should be explicitly stated what kind of objects u and v are. I think it would be also nice if the derivation of the dual could be provided in the appendix because it's not obvious to me how that happened.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "1. What you show is that H(s, a) is effectively just a \"better\" version of the Q function (Figure 3 left) for guiding and training diffusion policies. It seems that H(s, a) comes from three main ingredients: the Q function, the marginal distribution over states, and the marginal distribution over actions. What, intuitively, do these last two ingredients \"add in\" to make H better than Q? It seems in algorithm 2 that you are just learning H from the data in the replay buffer, which is collected from a suboptimal policy (the one currently being trained). Why would that help?\n2. In figure 4, why isn't optimal transport assigning a unique action for every state? Am I right that this is the relaxed scheme (eq (2)) instead of the Monge formulation (eq (1))?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper develops an optimal transport method for improving diffusion policy optimization. This is an important and timely topic for machine learning. From what I understand, it uses the Q function to define a cost for a transport problem that maps from states to actions, masking to guide matching based on demonstrated state-action pairs, and a resampling strategy to better guide the training process using online environment interactions.  As the reviewers note, the approach is difficult to fully understand, and the authors deserve some responsibility for not more clearly explaining it. There are also numerous typos in the paper (that should mostly be caught by spellcheck), and the authors' rebuttals to reviewer questions often increased rather than decreased reviewer confusion. However, on the positive side, the experimental results often outperform the baselines and the method appears to provide better stability. Since all reviewers are in favor of acceptance, I also recommend acceptance, but hope substantial improvements can be made to the paper's clarity before publication.
Decision: Accept (poster)